diff --git a/.claude/skills/mooncake-api/SKILL.md b/.claude/skills/mooncake-api/SKILL.md index d0edb73af5..f59b189623 100644 --- a/.claude/skills/mooncake-api/SKILL.md +++ b/.claude/skills/mooncake-api/SKILL.md @@ -19,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) @@ -307,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 @@ -465,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 @@ -510,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-troubleshoot/SKILL.md b/.claude/skills/mooncake-troubleshoot/SKILL.md index 734bb2ed02..c95410f035 100644 --- a/.claude/skills/mooncake-troubleshoot/SKILL.md +++ b/.claude/skills/mooncake-troubleshoot/SKILL.md @@ -1,6 +1,6 @@ --- name: mooncake-troubleshoot -description: Automatically diagnose Mooncake deployment and runtime issues. Checks services (mooncake_master, metadata server), RDMA devices, environment variables, connectivity, memory limits, and analyzes logs for common error patterns. Use when Mooncake deployment fails, services won't start, connections fail, or you encounter runtime errors like "Error from etcd client", "No matched device found", "Failed to register memory", "NO_AVAILABLE_HANDLE", or any RDMA/networking issues. Also use when user asks to troubleshoot, debug, diagnose, or fix Mooncake problems. +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 @@ -62,8 +62,8 @@ echo "https_proxy: $https_proxy" Verify critical environment variables are set correctly: ```bash -# Display all MC_* variables -env | grep ^MC_ +# Display Mooncake variables +env | grep -E '^(MC_|MOONCAKE_STORE_CHECKSUM=)' # Key variables to check: echo "MC_METADATA_SERVER: $MC_METADATA_SERVER" @@ -76,6 +76,7 @@ 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:** @@ -88,6 +89,7 @@ echo "MC_ENABLE_DEST_DEVICE_AFFINITY: $MC_ENABLE_DEST_DEVICE_AFFINITY" - `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 @@ -216,6 +218,8 @@ Search logs for common error patterns and their meanings: - **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 @@ -244,6 +248,18 @@ env | grep MC_FORCE_TCP - 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 @@ -266,6 +282,7 @@ env | grep MC_FORCE_TCP | -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 | diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile index a5729ef936..60b39f8118 100644 --- a/.devcontainer/Dockerfile +++ b/.devcontainer/Dockerfile @@ -8,7 +8,6 @@ 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 \ diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index eabf4b29fd..24dcbb753e 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -6,19 +6,19 @@ # EP: @UNIDY2002 UNIDY2002@outlook.com # PG: @UNIDY2002 UNIDY2002@outlook.com -.github @stmatengss @ykwd @Ann-1024 @luketong777 -/docs @ShangmingCai @stmatengss @ykwd +.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 +/mooncake-integration/store @ykwd @stmatengss @zxpdemonio /mooncake-pg @UNIDY2002 @ympcMark @yuechen-sys /mooncake-store @ykwd @stmatengss @XucSh @YiXR -/mooncake-store/*/ha/ @Libotry @YiXR @00fish0 -/mooncake-transfer-engine @alogfans @doujiang24 @chestnut-Q -/mooncake-transfer-engine/tent @alogfans @doujiang24 @chestnut-Q @staryxchen @00fish0 @dtcccc +/mooncake-store/*/ha/ @Libotry @YiXR @00fish0 @Icedcoco +/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-transfer-engine/*/transport/efa_transport/ @alogfans @whn09 -/mooncake-wheel @ShangmingCai @stmatengss +/mooncake-wheel @ShangmingCai @stmatengss @zxpdemonio /scripts/tone_tests @luketong777 /scripts/ascend/ @ascend-direct-dev @VNightMare @MingYang119 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/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..11216c6c7d --- /dev/null +++ b/.github/workflows/_build-wheel.yaml @@ -0,0 +1,262 @@ +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 -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' + 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 -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|cuda13:arm64) + variant_flag='' + if [ "$BUILD_VARIANT" = cuda13 ]; then + variant_flag=CU13_BUILD + fi + generator=Ninja + arch_list='9.0' + ep_versions='' + build_nvlink=false + cmake_args='-DBUILD_UNIT_TESTS=OFF -DUSE_HTTP=ON -DUSE_CUDA=ON -DUSE_MNNVL=ON -DWITH_EP=OFF -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 + 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"])') + master="$site_packages/mooncake/mooncake_master" + if readelf -d "$master" | grep -Eq \ + 'Shared library: \[(libcuda|libcudart)\.so'; then + echo "Non-CUDA mooncake_master depends on CUDA" + readelf -d "$master" | grep 'Shared library:' + 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 index 59e76321db..b04eeae3ba 100644 --- a/.github/workflows/assistant.yml +++ b/.github/workflows/assistant.yml @@ -58,3 +58,32 @@ jobs: 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.yml b/.github/workflows/ci.yml index 7a6642965b..9aeabbb2d7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -2,9 +2,13 @@ name: 'Build & Test (Linux)' on: push: - branches: [ "main" ] + branches: + - "main" + - "release/**" pull_request: - branches: [ "main" ] + branches: + - "main" + - "release/**" types: [opened, synchronize, reopened, labeled] workflow_dispatch: {} @@ -18,7 +22,7 @@ concurrency: jobs: build: needs: [spell-check, clang-format, check-paths] - if: >- + if: &run-ci-for-source-changes >- (needs.check-paths.outputs.should-run-downstream == 'true' || github.event_name == 'workflow_dispatch') && (github.event_name == 'push' || @@ -57,10 +61,7 @@ jobs: 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 + uses: ./.github/actions/free-disk-space - name: Install CUDA Toolkit uses: Jimver/cuda-toolkit@v0.2.24 @@ -70,24 +71,16 @@ jobs: method: 'network' sub-packages: '["nvcc"]' - - name: Install coverage tools and build utilities + - name: Install build utilities run: | sudo apt-get update - sudo apt-get install -y lcov gcovr ninja-build + sudo apt-get install -y ninja-build - name: Test HugeTLB sizing helper run: | python3 scripts/test_hicache_hugepage_requirements.py shell: bash - - 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 @@ -98,7 +91,7 @@ jobs: 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 + - name: Configure project run: | sudo apt update -y sudo bash -x dependencies.sh -y @@ -135,58 +128,21 @@ jobs: shell: bash - name: Run Mooncake Store Rust smoke test and benchmark - run: | - $GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master \ - --eviction_high_watermark_ratio=0.95 \ - --cluster_id=ci_rust_test_cluster \ - --port 50051 & - MASTER_PID=$! - sleep 3 - cd mooncake-store/rust - 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:$LD_LIBRARY_PATH - export MOONCAKE_BUILD_DIR=$GITHUB_WORKSPACE/build - export MOONCAKE_STORE_LIB_DIR=$GITHUB_WORKSPACE/build/mooncake-store/src - export MOONCAKE_STORE_INCLUDE_DIR=$GITHUB_WORKSPACE/mooncake-store/include - # This job builds Mooncake with -DENABLE_ASAN=ON, so the C++ libraries - # the Rust crate links against carry undefined __asan_* references. Opt - # in to linking the ASan runtime; build.rs emits -lasan first, which - # keeps libasan first in the initial library list as ASan requires. - # Non-sanitized builds leave this unset and link without ASan. - export MOONCAKE_LINK_ASAN=1 - export MC_METADATA_SERVER=http://127.0.0.1:8080/metadata - export MC_RUST_STORE_RUN_INTEGRATION=true - export MC_RUST_STORE_MASTER_ADDR=127.0.0.1:50051 - export MC_RUST_STORE_LOCAL_HOSTNAME=127.0.0.1 - export MC_RUST_STORE_PROTOCOL=tcp - export MC_RUST_STORE_DEVICE_NAME= - cargo test --test minimal_smoke -- --nocapture - MC_RUST_BENCH_ITERATIONS=4 \ - MC_RUST_BENCH_VALUE_SIZE=4096 \ - MC_RUST_BENCH_WARMUP=1 \ - cargo run --release --example store_benchmark - kill $MASTER_PID 2>/dev/null || true + env: + MOONCAKE_STORE_CLUSTER_ID: ci_rust_test_cluster + MOONCAKE_STORE_RUST_LINK_ASAN: "1" + run: ./scripts/ci/run_store_rust_smoke.sh shell: bash - name: Run Go store binding integration tests - run: | - $GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master \ - --eviction_high_watermark_ratio=0.95 \ - --cluster_id=ci_go_test_cluster \ - --port 50051 & - MASTER_PID=$! - sleep 3 - cd mooncake-store/go - export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/build/mooncake-common:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd - export CGO_ENABLED=1 - export CGO_CFLAGS="-I$GITHUB_WORKSPACE/mooncake-store/include -I$GITHUB_WORKSPACE/mooncake-transfer-engine/include" - export CGO_LDFLAGS="-L$GITHUB_WORKSPACE/build/mooncake-store/src -L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src -L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base -L$GITHUB_WORKSPACE/build/mooncake-common -L$GITHUB_WORKSPACE/build/mooncake-common/etcd -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -lmlx5 -ljsoncpp -lzstd -lcurl -luring -lasan -lm -lgcov -lxxhash" - # Link cudart if CUDA is available (needed for D2H staging in mooncake_store) - if [ -d /usr/local/cuda/lib64 ]; then export CGO_LDFLAGS="$CGO_LDFLAGS -L/usr/local/cuda/lib64 -lcudart"; fi - ASAN_OPTIONS=detect_leaks=0:verify_asan_link_order=0 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata go test -v ./tests/... - kill $MASTER_PID 2>/dev/null || true + env: + MOONCAKE_STORE_CLUSTER_ID: ci_go_test_cluster + MOONCAKE_STORE_GO_LINK_COMMON: "1" + MOONCAKE_STORE_GO_SANITIZED: "1" + run: ./scripts/ci/run_store_go_integration.sh shell: bash - - name: Test (in build env) with coverage + - name: Test (in build env) run: | cd build export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib @@ -205,120 +161,21 @@ jobs: DEFAULT_KV_LEASE_TTL=500 ./mooncake-store/tests/task_integration_test --gtest_filter='TaskExecutorIntegrationTest.DrainJobCompleteFlow' shell: bash - - 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 - } - - 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" - fi - - - name: Generate Python version tag - id: generate_tag_build - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - 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 - 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: + # 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: >- - (needs.check-paths.outputs.should-run-downstream == 'true' || - github.event_name == 'workflow_dispatch') && - (github.event_name == 'push' || - github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) - runs-on: ubuntu-22.04 - container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64 - steps: - - uses: actions/checkout@v4 - with: - persist-credentials: false - - - name: Mark repository as safe - run: git config --global --add safe.directory $GITHUB_WORKSPACE - 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 -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DENABLE_DEBUG_SYMBOLS=OFF - shell: bash - - - name: Build project - run: | - cd build - source ~/.bashrc - cmake --build . - cmake --install . - shell: bash + if: *run-ci-for-source-changes + 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: [spell-check, clang-format, build-flags] + needs: [spell-check, clang-format, build-wheel] if: >- - needs.build-flags.result == 'success' && + needs.build-wheel.result == 'success' && (github.event_name == 'push' || github.event_name == 'workflow_dispatch' || github.event.action == 'opened' || @@ -347,7 +204,7 @@ jobs: - name: Download wheel artifact uses: actions/download-artifact@v4 with: - name: mooncake-wheel-ubuntu-py${{ steps.generate_tag_test.outputs.python_version_tag }} + name: mooncake-wheel-ci-py${{ steps.generate_tag_test.outputs.python_version_tag }} path: mooncake-wheel/dist - name: Verify wheel file exists @@ -360,12 +217,7 @@ jobs: 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 @@ -379,12 +231,6 @@ jobs: bash scripts/test_installation.sh shell: bash - - name: Start metadata server - run: | - source test_env/bin/activate - mooncake_http_metadata_server --port 8080 & - shell: bash - - name: Run tests with ssd run: | # Reserve port 50052 (mooncake_client RPC port) so the kernel never @@ -408,7 +254,8 @@ jobs: mooncake_master \ --eviction_high_watermark_ratio=0.95 \ --cluster_id=ci_test_cluster \ - --port 50051 & + --port 50051 \ + --enable_http_metadata_server=true & sleep 3 shell: bash @@ -475,6 +322,20 @@ jobs: python mooncake-pg/tests/test_pg_collectives.py shell: bash + - name: Test PyTorch 2.13 Single-Buffer Collectives (CPU Only) + if: matrix.ubuntu-version == 'ubuntu-22.04' && matrix.python-version == '3.10' + env: + MC_FORCE_TCP: "true" + run: | + source test_env/bin/activate + python -m pip install --force-reinstall "torch==2.13.0" \ + --index-url https://download.pytorch.org/whl/cu126 \ + --extra-index-url https://pypi.org/simple + python mooncake-pg/tests/test_pg_collectives.py \ + TestMooncakePGCollectivesCPU.test_all_gather_into_tensor \ + TestMooncakePGCollectivesCPU.test_reduce_scatter_sum + shell: bash + - name: Test Safetensor Functions run: | source test_env/bin/activate @@ -484,22 +345,19 @@ jobs: build-flags: needs: [spell-check, clang-format, check-paths] - if: >- - (needs.check-paths.outputs.should-run-downstream == 'true' || - github.event_name == 'workflow_dispatch') && - (github.event_name == 'push' || - github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) + 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" + MAX_JOBS: "2" + EP_TORCH_VERSIONS: "2.11.0;2.12.0;2.12.1;2.13.0" + CMAKE_RELWITHDEBINFO_FLAGS: "-O2 -DNDEBUG" steps: - uses: actions/checkout@v4 @@ -512,12 +370,7 @@ jobs: 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 + uses: ./.github/actions/free-disk-space - name: Install CUDA Toolkit uses: Jimver/cuda-toolkit@v0.2.24 @@ -556,17 +409,53 @@ jobs: cd build export PATH=/usr/local/nvidia/bin:/usr/local/nvidia/lib64:$PATH export LD_LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LD_LIBRARY_PATH - cmake -G Ninja .. -DUSE_ETCD=OFF -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=OFF -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs" -DENABLE_DEBUG_SYMBOLS=OFF + 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 + cd .. + rm -rf build + df -h shell: bash - name: Configure project with all settings are ON run: | mkdir build cd build - cmake -G Ninja .. -DUSE_ETCD=ON -DUSE_CXL=ON -DUSE_REDIS=ON -DUSE_HTTP=ON -DWITH_STORE=ON -DWITH_P2P_STORE=ON -DWITH_METRICS=ON -DBUILD_UNIT_TESTS=ON -DBUILD_EXAMPLES=ON -DENABLE_SCCACHE=ON -DUSE_CUDA=ON -DUSE_MNNVL=OFF -DUSE_UB=OFF -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs" -DENABLE_DEBUG_SYMBOLS=OFF + # ENABLE_DEBUG_SYMBOLS=OFF alone still leaves CMake's RelWithDebInfo -g. + cmake -G Ninja .. \ + -DUSE_ETCD=ON \ + -DUSE_CXL=ON \ + -DUSE_REDIS=ON \ + -DUSE_HTTP=ON \ + -DWITH_STORE=ON \ + -DWITH_P2P_STORE=ON \ + -DWITH_METRICS=ON \ + -DBUILD_UNIT_TESTS=ON \ + -DBUILD_EXAMPLES=ON \ + -DENABLE_SCCACHE=ON \ + -DUSE_CUDA=ON \ + -DUSE_MNNVL=OFF \ + -DUSE_UB=OFF \ + -DCMAKE_EXE_LINKER_FLAGS="-L/usr/local/cuda/lib64/stubs" \ + "-DCMAKE_C_FLAGS_RELWITHDEBINFO=${CMAKE_RELWITHDEBINFO_FLAGS}" \ + "-DCMAKE_CXX_FLAGS_RELWITHDEBINFO=${CMAKE_RELWITHDEBINFO_FLAGS}" \ + -DENABLE_DEBUG_SYMBOLS=OFF shell: bash # TODO: lack USE_NVMEOF,USE_MNNVL @@ -610,13 +499,37 @@ jobs: MOONCAKE_STORE_LIB_DIR=$GITHUB_WORKSPACE/build/mooncake-store/src \ MOONCAKE_STORE_INCLUDE_DIR=$GITHUB_WORKSPACE/mooncake-store/include \ cargo test --examples --tests --no-run + cargo clean + shell: bash + + - name: Verify Mooncake Store Rust dlopen bindings and packaging + run: | + cd mooncake-store/rust + # 1. Committed dlopen bindings must stay in sync with store_c.h. + cargo run --locked --example generate_dlopen_bindings + git diff --exit-code -- src/generated/ffi_dlopen_bindings.rs + # 2. The published dlopen crate must build with no header/bindgen: package + # it, extract, and check the dlopen feature against the packaged files. + 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) shell: bash - name: Configure project run: | cd build rm -r */tests - cmake -G Ninja .. -DBUILD_UNIT_TESTS=OFF -DBUILD_EXAMPLES=OFF -DUSE_HTTP=ON -DENABLE_SCCACHE=ON -DUSE_CXL=ON -DWITH_EP=ON -DEP_TORCH_VERSIONS="2.9.1;2.10.0;2.11.0;2.12.0;2.12.1" -DENABLE_DEBUG_SYMBOLS=OFF + cmake -G Ninja .. \ + -DBUILD_UNIT_TESTS=OFF \ + -DBUILD_EXAMPLES=OFF \ + -DUSE_HTTP=ON \ + -DENABLE_SCCACHE=ON \ + -DUSE_CXL=ON \ + -DWITH_EP=ON \ + "-DEP_TORCH_VERSIONS=${EP_TORCH_VERSIONS}" \ + -DENABLE_DEBUG_SYMBOLS=OFF shell: bash - name: Build project @@ -628,26 +541,6 @@ jobs: sudo cmake --install . shell: bash - - 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 - shell: bash - - - name: Build project with TENT - run: | - cd build-tent - cmake --build . - sudo cmake --install . - shell: bash - - - name: Test (TENT) - run: | - cd build-tent - ctest --test-dir mooncake-transfer-engine/tent/tests -j --output-on-failure - shell: bash - - name: Build nvlink_allocator.so run: | mkdir -p build/mooncake-transfer-engine/nvlink-allocator @@ -663,53 +556,10 @@ jobs: shell: bash run: ${SCCACHE_PATH} --show-stats - - name: Generate Python version tag - id: generate_tag_flags - run: | - echo "python_version_tag=$(echo ${{ matrix.python-version }} | tr -d '.')" >> $GITHUB_OUTPUT - 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 - run: | - PYTHON_VERSION=${{ matrix.python-version }} OUTPUT_DIR=dist-py${{ steps.generate_tag_flags.outputs.python_version_tag }} ./scripts/build_wheel.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 - needs: [spell-check, clang-format, check-paths] - if: >- - (needs.check-paths.outputs.should-run-downstream == 'true' || - github.event_name == 'workflow_dispatch') && - (github.event_name == 'push' || - github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) - 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: Build Docker image - 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: >- + if: &run-ci >- (github.event_name == 'push' || github.event_name == 'workflow_dispatch' || github.event.action == 'opened' || @@ -725,11 +575,7 @@ jobs: clang-format: name: Check code format - if: >- - (github.event_name == 'push' || - github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci')) + if: *run-ci runs-on: ubuntu-22.04 steps: - name: Checkout Actions Repository @@ -769,17 +615,13 @@ 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 docs-check: name: Check Sphinx docs build - if: >- - github.event_name == 'push' || - github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci') + if: *run-ci runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v4 @@ -793,7 +635,7 @@ jobs: filters: | docs: - 'docs/**' - - 'requirements_docs.txt' + - 'docs/requirements-docs.txt' - name: Set up Python if: steps.filter.outputs.docs == 'true' @@ -805,7 +647,7 @@ jobs: if: steps.filter.outputs.docs == 'true' run: | python -m pip install --upgrade pip - pip install -r requirements_docs.txt + pip install -r docs/requirements-docs.txt - name: Build docs with strict mode if: steps.filter.outputs.docs == 'true' @@ -816,20 +658,19 @@ jobs: check-paths: - if: >- - github.event_name == 'push' || - github.event_name == 'workflow_dispatch' || - github.event.action == 'opened' || - contains(github.event.pull_request.labels.*.name, 'run-ci') + 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 + run: | + echo "src=true" >> $GITHUB_OUTPUT + echo "tent=true" >> $GITHUB_OUTPUT - uses: actions/checkout@v4 if: github.event_name != 'workflow_dispatch' with: @@ -846,9 +687,32 @@ jobs: - '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.10", "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' || @@ -857,21 +721,123 @@ jobs: github.event_name == 'workflow_dispatch' || github.event.action == 'opened' || contains(github.event.pull_request.labels.*.name, 'run-ci')) - uses: ./.github/workflows/ci_cu13.yml - secrets: inherit - - ascend-test: - needs: [build, check-paths] - if: needs.check-paths.outputs.should-run-downstream == 'true' - uses: ./.github/workflows/ci_ascend.yml + uses: ./.github/workflows/ci_rocm.yml secrets: inherit integration-test: - needs: [build, check-paths] + needs: [build, build-wheel-cu13, check-paths] if: needs.check-paths.outputs.should-run-downstream == 'true' uses: ./.github/workflows/integration-test.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.action == 'opened' || + 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: 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 + + # 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) + if: '!matrix.need_cuda' + run: | + cd build-tent + ctest --test-dir mooncake-transfer-engine/tent/tests -j --output-on-failure + 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() @@ -880,12 +846,13 @@ jobs: - clang-format - docs-check - build - - build-musa + - build-wheel - build-flags - - build-docker - test-wheel-ubuntu - build-wheel-cu13 - - ascend-test + - build-wheel-efa + - build-wheel-rocm + - tent-ci - integration-test runs-on: ubuntu-latest steps: diff --git a/.github/workflows/ci_cu13.yml b/.github/workflows/ci_cu13.yml deleted file mode 100644 index 828491f3f8..0000000000 --- a/.github/workflows/ci_cu13.yml +++ /dev/null @@ -1,126 +0,0 @@ -name: 'Build Wheel (CUDA 13)' - -on: - workflow_call: {} - -jobs: - build-wheel-cu13: - 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 - with: - persist-credentials: false - - - 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: Install dependencies - run: | - sudo apt update -y - sudo apt install -y ninja-build - sudo bash -x dependencies.sh -y - df -h - shell: bash - - - name: Configure project - run: | - mkdir build - cd build - cmake -G Ninja .. \ - -DUSE_ETCD=ON \ - -DUSE_REDIS=ON \ - -DUSE_HTTP=ON \ - -DWITH_STORE=ON \ - -DWITH_P2P_STORE=ON \ - -DWITH_EP=ON \ - -DEP_TORCH_VERSIONS="2.9.1;2.10.0;2.11.0;2.12.0;2.12.1" \ - -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" \ - -DENABLE_DEBUG_SYMBOLS=OFF - 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 - cmake --build . - sudo cmake --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: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - 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/code-review.yml b/.github/workflows/code-review.yml deleted file mode 100644 index 94622c0896..0000000000 --- a/.github/workflows/code-review.yml +++ /dev/null @@ -1,29 +0,0 @@ -name: Qoder Auto Code Review - -on: - pull_request: - types: [opened, synchronize, reopened] - -jobs: - qoder-review: - # Skip fork and cross-repo PRs (head repo must match this repository) - if: github.event.pull_request.head.repo.full_name == github.repository - runs-on: ubuntu-latest - permissions: - contents: read - pull-requests: write - id-token: write - - steps: - - name: Checkout repository - uses: actions/checkout@v4 - with: - fetch-depth: 0 - - - name: Run Qoder Code Review - uses: QoderAI/qoder-action@v0 - with: - qoder_personal_access_token: ${{ secrets.QODER_PERSONAL_ACCESS_TOKEN }} - prompt: | - /review-pr - REPO:${{ github.repository }} PR_NUMBER:${{ github.event.pull_request.number }} diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 34dcd9ea74..6862480bae 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -6,7 +6,7 @@ on: branches: ["main"] paths: - 'docs/**' - - 'requirements_docs.txt' + - 'docs/requirements-docs.txt' - '.github/workflows/deploy.yml' # Allows you to run this workflow manually from the Actions tab workflow_dispatch: @@ -42,26 +42,26 @@ jobs: 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/integration-test.yml b/.github/workflows/integration-test.yml index 0d976f5624..4d5cb8ad52 100644 --- a/.github/workflows/integration-test.yml +++ b/.github/workflows/integration-test.yml @@ -34,12 +34,34 @@ jobs: attempt=1 while [ $attempt -le $max_attempts ]; do echo "Attempt $attempt: Fetching artifact..." - if curl -L -fs -o artifact.json -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" https://api.github.com/repos/${{ github.repository }}/actions/artifacts?per_page=100; then - artifact_id="" - if jq empty artifact.json >/dev/null 2>&1; then - artifact_id=$(jq -r ".artifacts[] | select(.name | contains(\"py312\") ) | select(.name | contains(\"mooncake\") ) | select(.name | contains(\"cu130\") | not) | select(.workflow_run.head_sha == \"$SHA\" ) | .id" artifact.json | head -n 1) + echo "Target SHA=${SHA}" + artifact_id="" + run_id="" + 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 artifact list. Retrying..." + echo "Failed to download workflow run list. Retrying..." + 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 + artifact_id=$(jq -r '.artifacts[] | select(.name | contains("py312") ) | select(.name | contains("mooncake") ) | select(.name | contains("cu130") ) | .id' artifact.json | head -n 1) + 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 find Build & Test workflow run for SHA $SHA. Retrying..." + echo "Available workflow runs for SHA:" + jq -r '.workflow_runs[] | "\(.id) \(.name) \(.path) \(.status) \(.conclusion)"' runs.json || true fi if [ -n "$artifact_id" ]; then echo "Successfully fetched expected artifact id $artifact_id" @@ -51,7 +73,7 @@ jobs: fi fi else - echo "Failed to fetch artifacts. Retrying..." + echo "Failed to fetch workflow runs. Retrying..." if [ $attempt -lt $max_attempts ]; then sleep $((attempt * 60)) fi diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml new file mode 100644 index 0000000000..ec31a8082c --- /dev/null +++ b/.github/workflows/nightly.yml @@ -0,0 +1,536 @@ +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 + uses: ./.github/workflows/_build-wheel.yaml + with: + variant: ${{ matrix.variant }} + architecture: ${{ matrix.architecture }} + python-versions: '["3.10", "3.12"]' + artifact-prefix: ${{ matrix.artifact-prefix }} + version-override: ${{ needs.version-stamp.outputs.nightly_version }} + + 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 + cmake --build . --parallel 1 + 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 \ + -DSTORE_USE_ETCD=ON -DCMAKE_BUILD_TYPE=Release \ + -DBUILD_UNIT_TESTS=ON -DENABLE_SCCACHE=ON + + - name: Build project + run: | + cd build + export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:${LIBRARY_PATH:-} + 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 + export LIBRARY_PATH=/usr/local/cuda/lib64/stubs:$LIBRARY_PATH + 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 & + sleep 2 + + - name: Run CTest unit tests + 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 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 Go store binding integration tests + env: + MOONCAKE_STORE_CLUSTER_ID: nightly_go_cluster + MOONCAKE_STORE_GO_LINK_COMMON: "0" + MOONCAKE_STORE_GO_SANITIZED: "0" + run: ./scripts/ci/run_store_go_integration.sh + + - name: Build and install Python wheel for integration tests + run: | + 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 + 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_SSD_OFFLOAD_IN_EVICT: "1" + TEST_PROMOTION_ON_HIT: "1" + TEST_CXL: "1" + run: | + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}:/usr/local/lib + bash scripts/run_tests.sh + + 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: 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: 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 + - 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/pre-release.yaml b/.github/workflows/pre-release.yaml index b500fcf6c4..2a07c4b4ad 100644 --- a/.github/workflows/pre-release.yaml +++ b/.github/workflows/pre-release.yaml @@ -1,7 +1,8 @@ name: Pre-Release -# Dry-run of the release pipelines: build wheels like Release / Release Non-CUDA / -# Release CUDA 13, validate artifacts, but do not create a GitHub Release or publish to PyPI. +# Dry run of the release pipelines: build wheels through the same _build-wheel.yaml +# the Release / Release Non-CUDA / Release CUDA 13 workflows use, validate the +# artifacts, but do not create a GitHub Release or publish to PyPI. # # Trigger by pushing a pre-release tag, for example: # git tag v1.0.0-rc1 && git push origin v1.0.0-rc1 @@ -13,295 +14,39 @@ on: - 'v*-beta*' - 'v*-pre*' -env: - SCCACHE_GHA_ENABLED: "true" - jobs: - build-cuda: - name: Build (CUDA 12) - runs-on: ubuntu-22.04 - permissions: - contents: read - strategy: - 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 version from tag - run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV - - - 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: 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.1;2.10.0;2.11.0;2.12.0;2.12.1" -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 -E 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: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - 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: | - 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: pre-release-cuda-py${{ steps.generate_tag_release.outputs.python_version_tag }} - path: mooncake-wheel/dist-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl - - build-non-cuda: - name: Build (Non-CUDA) - runs-on: ubuntu-22.04 - permissions: - contents: read - strategy: - 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 version from tag - run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV - - - 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: 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: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - 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: | - 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: pre-release-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 - - build-cuda13: - name: Build (CUDA 13) - runs-on: ubuntu-22.04 - permissions: - contents: read + build: 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 version from tag - run: echo "VERSION=${GITHUB_REF_NAME#v}" >> $GITHUB_ENV - - - 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: 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.1;2.10.0;2.11.0;2.12.0;2.12.1" -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: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - 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: | - 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: pre-release-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: + - 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 }} validate-release: name: Validate release artifacts - needs: [build-cuda, build-non-cuda, build-cuda13] + needs: build runs-on: ubuntu-22.04 permissions: contents: read @@ -313,7 +58,7 @@ jobs: uses: actions/download-artifact@v4 with: path: mooncake-wheel/dist-all - pattern: pre-release-* + pattern: mooncake-wheel*pre-release* - name: Prepare wheels for validation run: | @@ -324,8 +69,8 @@ jobs: ls -la mooncake-wheel/dist-release/ wheel_count=$(find mooncake-wheel/dist-release -name "*.whl" | wc -l) echo "wheel_count=${wheel_count}" >> "$GITHUB_ENV" - if [ "${wheel_count}" -lt 12 ]; then - echo "Expected at least 12 wheels (4 Python versions x 3 variants), found ${wheel_count}" + if [ "${wheel_count}" -lt 24 ]; then + echo "Expected at least 24 wheels (4 Python versions x 3 variants x 2 architectures), found ${wheel_count}" exit 1 fi @@ -342,7 +87,7 @@ jobs: - name: Upload validated wheels as workflow artifacts uses: actions/upload-artifact@v4 with: - name: pre-release-wheels-${{ github.ref_name }} + name: mooncake-wheels-pre-release-${{ github.ref_name }} path: mooncake-wheel/dist-release/*.whl retention-days: 14 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 8b08a8ed0a..4e45f477a3 100644 --- a/.github/workflows/release-cuda13.yaml +++ b/.github/workflows/release-cuda13.yaml @@ -5,143 +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: 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.1;2.10.0;2.11.0;2.12.0;2.12.1" -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 -E 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: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - 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..dcbfe7434f --- /dev/null +++ b/.github/workflows/release-efa-cuda13.yaml @@ -0,0 +1,35 @@ +name: Release EFA CUDA 13 + +on: + push: + tags: + - 'v*' + +# Publishes mooncake-transfer-engine-efa-cuda13, the CUDA 13-aware AWS EFA wheel. +jobs: + build: + 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..6ea27b00d6 --- /dev/null +++ b/.github/workflows/release-efa-non-cuda.yaml @@ -0,0 +1,34 @@ +name: Release EFA Non-CUDA + +on: + push: + tags: + - 'v*' + +# Publishes mooncake-transfer-engine-efa-non-cuda, the CPU/DRAM AWS EFA wheel. +jobs: + build: + 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..c10d46bc89 --- /dev/null +++ b/.github/workflows/release-efa.yaml @@ -0,0 +1,35 @@ +name: Release EFA + +on: + push: + tags: + - 'v*' + +# Publishes mooncake-transfer-engine-efa, the CUDA 12-aware AWS EFA wheel. +jobs: + build: + 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 1852717366..7d52069fb2 100644 --- a/.github/workflows/release-non-cuda.yaml +++ b/.github/workflows/release-non-cuda.yaml @@ -5,120 +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: 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: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - 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 index 7f665e19fe..9ef99593d3 100644 --- a/.github/workflows/release-npu.yaml +++ b/.github/workflows/release-npu.yaml @@ -4,13 +4,16 @@ on: push: tags: - 'v*' - -env: - CANN_VERSION: "9.0.0" + workflow_dispatch: + inputs: + tag: + description: 'Release tag (e.g. v0.1.0)' + required: true + type: string jobs: build: - if: ${{ !contains(github.ref_name, '-') }} + 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 @@ -31,13 +34,13 @@ jobs: contents: write env: - BUILD_WITH_EP: "0" 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 }} @@ -58,8 +61,11 @@ jobs: shell: bash run: | set -euo pipefail - CANN_URL="https://ascend-repo.obs.cn-east-2.myhuaweicloud.com/CANN/CANN%20${{ env.CANN_VERSION }}/Ascend-cann-toolkit_${{ env.CANN_VERSION }}_linux-${{ matrix.cann_arch }}.run" - echo "Downloading CANN ${{ env.CANN_VERSION }} from ${CANN_URL}" + 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 @@ -148,7 +154,7 @@ jobs: path: mooncake-wheel/dist-npu-py${{ steps.generate_tag_release.outputs.python_version_tag }}/*.whl publish-release: - if: ${{ !contains(github.ref_name, '-') }} + 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 @@ -160,6 +166,8 @@ jobs: 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 @@ -177,6 +185,7 @@ jobs: - 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 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 4bd194d96b..f5cb1c5ff4 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -5,142 +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: - # Skip semver pre-release tags (e.g. v1.0.0-rc1); those are handled by pre-release.yaml. if: ${{ !contains(github.ref_name, '-') }} - runs-on: ubuntu-22.04 - permissions: - contents: write 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: 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.1;2.10.0;2.11.0;2.12.0;2.12.1" -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 -E 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: Run sccache stat for check - if: ${{ env.SCCACHE_PATH != '' }} - shell: bash - run: ${SCCACHE_PATH} --show-stats - - - 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/.gitignore b/.gitignore index 7dfe7163cc..35a253a732 100644 --- a/.gitignore +++ b/.gitignore @@ -198,14 +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 3fefbf6af2..e29bb9c9be 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,7 @@ # Pre-commit hooks configuration for Mooncake # Install: pip install -r requirements-dev.txt && pre-commit install -# Run manually: pre-commit run --all-files +# Run manually on all files: 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 +10,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 @@ -26,12 +30,21 @@ repos: - repo: local hooks: - id: mooncake-code-format - name: Run Mooncake code format script - entry: ./scripts/code_format.sh + 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 - always_run: true 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 @@ -47,14 +60,7 @@ repos: hooks: - id: codespell exclude: '^(extern/|FAST25-release/)' - args: ['--ignore-words-list=te,mooncake,KVCache,cann'] - - - 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'] - repo: https://github.com/cheshirekow/cmake-format-precommit rev: v0.6.13 diff --git a/.typos.toml b/.typos.toml index 5ba9fb610e..4843e506f2 100644 --- a/.typos.toml +++ b/.typos.toml @@ -1,5 +1,5 @@ [default] -extend-ignore-words = ["CANN", "ASO", "fre", "wqs", "hsa"] +extend-ignore-words = ["CANN", "ASO", "fre", "wqs", "hsa", "ue"] [default.extend-words] CANN = "CANN" @@ -9,8 +9,13 @@ 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..f7332c2d0a --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,25 @@ +# 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. +- Run pre-commit locally on the files touched by the change before handoff when + the toolchain is available. If broader hooks or `pre-commit run --all-files` + rewrite unrelated files, do not include those unrelated edits in 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 2872bd1ad4..a82f333c77 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -18,8 +18,16 @@ option(WITH_STORE_GO "build Go bindings for mooncake store" OFF) option(WITH_P2P_STORE "build p2p store library and sample code" OFF) option(WITH_RUST_EXAMPLE "build the Rust interface and sample code for the transfer engine" OFF) option(WITH_STORE_RUST "build the Rust bindings for the Mooncake Store" ON) +option(WITH_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) include(${CMAKE_CURRENT_SOURCE_DIR}/mooncake-common/SetupPython.cmake) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extern/pybind11) @@ -70,6 +78,7 @@ add_compile_definitions(ASIO_SEPARATE_COMPILATION ASIO_DYN_LINK) add_subdirectory(mooncake-common) include_directories(mooncake-common/etcd) +include_directories(mooncake-common/k8s-lease) include_directories(mooncake-common/include) if (WITH_TE) @@ -93,12 +102,24 @@ endif() 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_subdirectory(mooncake-pg) - include_directories(mooncake-pg/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 and PG Python extensions") if(USE_CUDA) @@ -145,7 +166,9 @@ if (WITH_EP) "-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}" "-DSTAGING_DIR=${EP_PG_STAGING_DIR}" "-DENGINE_SO_PATH=$" + "-DPython3_EXECUTABLE=${Python3_EXECUTABLE}" "-DEP_USE_MUSA=$,1,0>" + "-DEP_USE_MACA=$,1,0>" -P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-ep/BuildEpExt.cmake" COMMENT "Building Mooncake EP Python extension(s)" DEPENDS engine @@ -155,17 +178,19 @@ if (WITH_EP) 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" + "-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}" - "-DTORCH_CUDA_ARCH_LIST=${_torch_cuda_arch_list_pipe}" "-DSTAGING_DIR=${EP_PG_STAGING_DIR}" - "-DENGINE_SO_PATH=$" + "-DPG_CORE_SO_PATH=$" + "-DPG_DEVICE_SO_PATH=$" + "-DPython3_EXECUTABLE=${Python3_EXECUTABLE}" "-DEP_USE_MUSA=$,1,0>" - -P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg/BuildPgExt.cmake" + "-DEP_USE_MACA=$,1,0>" + -P "${CMAKE_CURRENT_SOURCE_DIR}/mooncake-pg/torch/BuildPgExt.cmake" COMMENT "Building Mooncake PG Python extension(s)" - DEPENDS engine mooncake_ep_ext + DEPENDS mooncake_pg mooncake_pg_device mooncake_ep_ext VERBATIM ) endif () diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b18fa536b7..68d0741220 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,23 +47,27 @@ Mooncake uses [pre-commit](https://pre-commit.com/) to enforce consistent format | Type | Tool | Purpose | |------|------|---------| | Generic | trailing-whitespace / end-of-file-fixer | Basic hygiene | -| Project | `./scripts/code_format.sh` | Enforce Mooncake C/C++ formatting script before commit | +| 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 will run `./scripts/code_format.sh` automatically. If it rewrites files, re-stage the changes and commit again. +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): +Run hooks on all files (the first run installs hook environments). The C/C++ +hook remains limited to staged line ranges; use `./scripts/code_format.sh --all` +for an intentional whole-project C/C++ format: ```bash pre-commit run --all-files ``` @@ -68,7 +78,8 @@ 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-20 ``` @@ -80,7 +91,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..297963d94c 100644 --- a/MAINTAINERS.md +++ b/MAINTAINERS.md @@ -12,6 +12,6 @@ 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 58bc1222df..5d0cb1f2ee 100644 --- a/README.md +++ b/README.md @@ -16,26 +16,30 @@ [![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](https://img.shields.io/pypi/v/mooncake-transfer-engine)](https://pypi.org/project/mooncake-transfer-engine) [![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 CUDA 12](https://img.shields.io/static/v1?label=pypi&message=EFA%20%2B%20CUDA%2012&color=F59E0B)](https://pypi.org/project/mooncake-transfer-engine-efa/) + [![PyPI EFA CUDA 13](https://img.shields.io/static/v1?label=pypi&message=EFA%20%2B%20CUDA%2013&color=F59E0B)](https://pypi.org/project/mooncake-transfer-engine-efa-cuda13/) + [![PyPI EFA Non-CUDA](https://img.shields.io/static/v1?label=pypi&message=EFA%20non-CUDA&color=F59E0B)](https://pypi.org/project/mooncake-transfer-engine-efa-non-cuda/)
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

+- **Jul 2, 2026**: [DSpark](https://x.com/mgoin_/status/2072785822231728363) scales fully online training on a GB300 NVL72 system with Speculators and Mooncake: 9 vLLM nodes serve the GLM 5.2 FP8 verifier through Mooncake RDMA Store to 6 FSDP training nodes (DP=24), achieving 125k prefill tokens/s and 1.5 steps/s. - **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. -- **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. +- **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.
@@ -55,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. @@ -74,19 +78,17 @@ This repository also hosts its technical report and the open-sourced traces.

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

🔥 Show Cases

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

🔥 Show Cases

+ ### Transfer Engine (TE) 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. @@ -126,13 +128,13 @@ Mooncake Store is a high-performance distributed key-value cache storage engine - **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/getting_started/examples/lmcache-integration.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. +- **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.
### Mooncake EP and Process Group (PG) -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/python-api-reference/ep-backend.html) for details. +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 @@ -153,7 +155,7 @@ Mooncake EP and Mooncake PG extend Mooncake from high-performance data movement 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/getting_started/examples/sglang-integration/index.html)) +### 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. @@ -176,7 +178,7 @@ Mooncake is deeply integrated into [SGLang](https://github.com/sgl-project/sglan
-### vLLM Integration ([Guide](https://kvcache-ai.github.io/Mooncake/getting_started/examples/vllm-integration/index.html)) +### vLLM Integration ([Guide](https://kvcache-ai.github.io/Mooncake/deployment/integrations/vllm/index.html)) 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. @@ -193,46 +195,15 @@ Mooncake integrates with [vLLM](https://github.com/vllm-project/vllm) to acceler

🖥️ Supported Hardware

-Mooncake supports hardware backends across accelerator vendors, cloud fabrics, and standard datacenter interconnects. - -The following hardware partners and cloud platforms are supported by the Mooncake, covering GPUs, specialized AI accelerators, and cloud-native interconnects: - -
- - - - - - - - - - - - -
Huawei
Huawei
Cambricon
Cambricon
Moore Threads
Moore Threads
MetaX
MetaX
T-Head
T-Head
NVIDIA
NVIDIA
AMD
AMD
Alibaba Cloud
Alibaba Cloud
AWS
AWS
-
- -For complete protocol behavior, SDK requirements, and vendor-specific configuration, see the [supported protocols](https://kvcache-ai.github.io/Mooncake/getting_started/supported-protocols.html), [build guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html), and [Transfer Engine design docs](https://kvcache-ai.github.io/Mooncake/design/transfer-engine/index.html). - -

🚀 Quick Start

- -### Before using Mooncake - -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)*. -- Cambricon Neuware, if the package is built with `-DUSE_MLU`. By default Mooncake looks for Neuware under `NEUWARE_HOME` or `/usr/local/neuware`. -- Hygon DTK SDK, if the package is built with `-DUSE_HYGON`. By default Mooncake looks for DTK under `DTK_HOME` or `/opt/dtk`. -- Iluvatar CoreX SDK, if the package is built with `-DUSE_COREX`. By default Mooncake looks for CoreX under `COREX_HOME` or `/usr/local/corex`. +| 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 @@ -243,44 +214,19 @@ pip install mooncake-transfer-engine pip install mooncake-transfer-engine-cuda13 ``` -**For non-CUDA systems:** -```bash -pip install mooncake-transfer-engine-non-cuda -``` - -**For NPU systems:** -```bash -pip install mooncake-transfer-engine-npu -``` +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. -> - MLU support is currently available through source builds with `-DUSE_MLU=ON`; there is no dedicated prebuilt MLU wheel yet. -> - If users encounter problems such as missing `lib*.so`, they should uninstall the package they installed and build the binaries manually. +- [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) -### Build From Source -For the default source build, use the automatic dependency script and standard CMake flow: +### Skills for AI Assistants -```bash -git clone https://github.com/kvcache-ai/Mooncake.git -cd Mooncake - -sudo bash dependencies.sh - -mkdir build -cd build -cmake .. -make -j -sudo make install # optional, make it ready to be used by vLLM/SGLang -``` +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. -For custom accelerator backends, Docker deployment, NVMe-oF, EFA, CXL, Redis / HTTP metadata, Rust bindings, or other advanced build options, see the [Build Guide](https://kvcache-ai.github.io/Mooncake/getting_started/build.html). - -### Skills for AI coding assistants - -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: +
+Details | Skill | Description | |-------|-------------| @@ -299,29 +245,35 @@ Install them without cloning the repository via the [Claude Code plugin marketpl 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 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] -} -``` -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. +

📦 Open Source Traces and Tools

-**_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._** +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 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}, @@ -354,19 +306,6 @@ Please kindly cite our papers if you find the papers 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}, @@ -374,3 +313,5 @@ Please kindly cite our papers if you find the papers or the traces are useful: url = {https://arxiv.org/abs/2407.00079}, } ``` + +
diff --git a/benchmarks/storage_benchmark/storage_benchmark.py b/benchmarks/storage_benchmark/storage_benchmark.py deleted file mode 100644 index 689bc98e99..0000000000 --- a/benchmarks/storage_benchmark/storage_benchmark.py +++ /dev/null @@ -1,956 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 - -""" -Mooncake KVCache Storage Benchmark Tool -""" - -import argparse -import json -import time -import os -import statistics -import random -import errno -from pathlib import Path -from typing import Dict, List, Optional -from dataclasses import dataclass - -# ============================================================================ -# Constants -# ============================================================================ - -BLOCK_SIZE_TOKENS = 512 # Number of tokens per block -DEFAULT_BYTES_PER_TOKEN = 2048 # 7B model FP16 (2KB per token) -BLOCK_SIZE_BYTES = BLOCK_SIZE_TOKENS * DEFAULT_BYTES_PER_TOKEN # 1MB per block -MIN_LATENCY_MS = 0.001 # Minimum latency in milliseconds (1 microsecond) - -# Model KVCache sizes (bytes per token, based on LMCache calculator) -# Source: https://lmcache.ai/kv_cache_calculator.html -MODEL_BYTES_PER_TOKEN = { - "llama-3.1-405b": 327680, - "qwen3-32b": 81920, - "deepseek-v3": 1748992, - "glm-4.6": 157013, - "default": DEFAULT_BYTES_PER_TOKEN, -} - -# ============================================================================ -# Data Structures -# ============================================================================ - -@dataclass -class KVCacheRequest: - """KVCache request - - Attributes: - timestamp: Request timestamp in milliseconds - hash_ids: List of block IDs (each ID corresponds to a 512-token block) - input_length: Input token count - output_length: Output token count - """ - timestamp: float - hash_ids: List[int] - input_length: int - output_length: int - -# ============================================================================ -# Storage Layer: Offset Allocator -# ============================================================================ - -class OffsetAllocatorStorage: - """High-performance block storage based on Offset Allocator - - Architecture: - ----------- - 1. Single large file stores all blocks (avoids file explosion) - 2. Uses offset to manage file space (similar to Mooncake's OffsetAllocator) - 3. hash_id -> offset mapping stored in memory (fast lookup) - - Block Organization: - ----------- - Each block corresponds to 512 tokens, fixed size 1MB: - - hash_id[0] -> block_0 (tokens [0...511]) -> offset 0 - - hash_id[1] -> block_1 (tokens [512...1023]) -> offset 1 - - hash_id[i] -> block_i (tokens [i*512...(i+1)*512-1]) -> offset i - - Performance Advantages: - ----------- - - Only one file, no file explosion - - Offset reuse, reduces memory allocation - - pread/pwrite, thread-safe, no seek needed - - Keep fd open, reduces open/close overhead - - Metadata in memory, O(1) lookup - - Attributes: - storage_dir: Storage directory path - block_size_bytes: Block size in bytes - max_blocks: Maximum number of blocks - hash_id_to_offset: hash_id -> offset mapping - free_offsets: List of reusable offsets - next_offset: Next allocatable offset - """ - - def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN, - max_blocks: int = 100000, block_size_tokens: int = 512, - fsync_mode: str = 'batch', fsync_batch_size: int = 100): - """Initialize Offset Allocator storage - - Args: - storage_dir: Storage directory path - bytes_per_token: Bytes per token - max_blocks: Maximum number of blocks (determines file size) - block_size_tokens: Number of tokens per block - fsync_mode: When to fsync ('batch', 'always', 'end', 'none') - fsync_batch_size: Number of writes between fsync in batch mode - """ - self.storage_dir = Path(storage_dir) - self.bytes_per_token = bytes_per_token - self.block_size_tokens = block_size_tokens - self.block_size_bytes = self.block_size_tokens * self.bytes_per_token - self.max_blocks = max_blocks - - # Fsync configuration - self.fsync_mode = fsync_mode - self.fsync_batch_size = fsync_batch_size - self.pending_sync_count = 0 - - # Create storage directory - self.storage_dir.mkdir(parents=True, exist_ok=True) - - # Single large file - self.storage_file = self.storage_dir / "kvcache_storage.bin" - self.file_size = self.max_blocks * self.block_size_bytes - - # Initialize storage file - if not self.storage_file.exists(): - self._init_storage_file() - - # hash_id -> offset mapping (metadata, in memory) - self.hash_id_to_offset: Dict[int, int] = {} - - # Offset allocator (free list) - self.free_offsets: List[int] = [] - self.next_offset = 0 - - # File descriptor (keep open, avoid repeated open/close) - self.fd = None - - # Pre-allocated data buffer with pattern to avoid SSD compression artifacts - # Using a repeating pattern that looks like realistic data (not all zeros) - # Pattern: 64-byte repeated sequence mixed with some variation - pattern = bytes([(i & 0xFF) for i in range(256)]) # 0-255 byte pattern - pattern_repeats = (self.block_size_bytes // len(pattern)) + 1 - self._data_buffer = (pattern * pattern_repeats)[:self.block_size_bytes] - - # Statistics - self.stats = { - 'read_count': 0, - 'write_count': 0, - 'read_bytes': 0, - 'write_bytes': 0, - 'read_latencies_ms': [], - 'write_latencies_ms': [], - 'sync_count': 0, # Number of fsync operations performed - } - - # ======================================================================== - # Internal Methods - # ======================================================================== - - def _init_storage_file(self): - """Initialize storage file (pre-allocate space) - - Create sparse file to avoid actual disk space usage until data is written - """ - with open(self.storage_file, 'wb') as f: - f.seek(self.file_size - 1) - f.write(b'\0') - f.flush() - os.fsync(f.fileno()) - - def _get_fd(self): - """Get file descriptor (lazy open) - - Returns: - int: File descriptor - """ - if self.fd is None: - # Use O_RDWR | O_CREAT, no O_DIRECT (Python compatibility) - self.fd = os.open(self.storage_file, os.O_RDWR | os.O_CREAT) - return self.fd - - def _allocate_offset(self) -> int: - """Allocate a new offset - - Prioritize reusing freed offsets, otherwise allocate new offset - - Returns: - int: Allocated offset - """ - if self.free_offsets: - return self.free_offsets.pop() - offset = self.next_offset - self.next_offset += 1 - return offset - - def _free_offset(self, offset: int): - """Free offset for reuse - - Args: - offset: Offset to free - """ - self.free_offsets.append(offset) - - # ======================================================================== - # Public Interface - # ======================================================================== - - def block_exists(self, hash_id: int) -> bool: - """Check if block exists - - Args: - hash_id: Unique block identifier - - Returns: - bool: Whether block exists - """ - return hash_id in self.hash_id_to_offset - - def read_block(self, hash_id: int) -> float: - """Read block using pread - - Args: - hash_id: Unique block identifier - - Returns: - float: Read latency in milliseconds, or 0 if block doesn't exist - """ - if hash_id not in self.hash_id_to_offset: - return 0.0 # Block doesn't exist, no latency to measure - - offset = self.hash_id_to_offset[hash_id] - file_offset = offset * self.block_size_bytes - - start = time.perf_counter() - - try: - fd = self._get_fd() - data = os.pread(fd, self.block_size_bytes, file_offset) - latency_ms = (time.perf_counter() - start) * 1000.0 - - self.stats['read_count'] += 1 - self.stats['read_bytes'] += len(data) - self.stats['read_latencies_ms'].append(latency_ms) - return latency_ms - except OSError as e: - print(f"Error reading block {hash_id} at offset {file_offset}: {e}") - return 0.0 # Error case, don't pollute stats - - def write_block(self, hash_id: int) -> float: - """Write block using pwrite - - Args: - hash_id: Unique block identifier - - Returns: - float: Write latency in milliseconds - """ - # Allocate offset - offset = self._allocate_offset() - file_offset = offset * self.block_size_bytes - - # Use pre-allocated buffer (much faster than os.urandom) - data = self._data_buffer - - start = time.perf_counter() - - try: - fd = self._get_fd() - written = os.pwrite(fd, data, file_offset) - - write_done = time.perf_counter() - - # Conditional fsync based on mode - if self.fsync_mode == 'always': - # Include fsync in latency measurement - os.fsync(fd) - self.stats['sync_count'] += 1 - self.pending_sync_count = 0 - latency_ms = (time.perf_counter() - start) * 1000.0 - # Evict from page cache AFTER fsync to ensure reads measure actual SSD performance - os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED) - elif self.fsync_mode == 'batch': - # For batch mode, only measure write time (fsync is deferred) - self.pending_sync_count += 1 - if self.pending_sync_count >= self.fsync_batch_size: - os.fsync(fd) - self.stats['sync_count'] += 1 - self.pending_sync_count = 0 - latency_ms = (write_done - start) * 1000.0 # Only write time - # Evict from page cache after each write - os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED) - elif self.fsync_mode == 'none': - latency_ms = (write_done - start) * 1000.0 - # Evict from page cache even when not syncing - os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED) - else: # 'end' mode - latency_ms = (write_done - start) * 1000.0 - # Evict from page cache (fsync will happen at the end) - os.posix_fadvise(fd, file_offset, self.block_size_bytes, os.POSIX_FADV_DONTNEED) - - # Update mapping - self.hash_id_to_offset[hash_id] = offset - - self.stats['write_count'] += 1 - self.stats['write_bytes'] += written - self.stats['write_latencies_ms'].append(latency_ms) - return latency_ms - except OSError as e: - if e.errno == errno.ENOSPC: - print(f"Error: Disk full when writing block {hash_id} at offset {file_offset}") - else: - print(f"Error writing block {hash_id} at offset {file_offset}: {e}") - return 0.0 # Error case, don't pollute stats - - def __enter__(self): - """Context manager entry""" - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """Context manager exit - ensures cleanup""" - # Perform final fsync before closing for 'end' and 'batch' modes - self._finalize_sync() - self.close(force_sync=False) # Already synced above - return False - - def _finalize_sync(self): - """Perform final fsync before closing (for 'end' mode and pending batch writes)""" - if self.fd is not None: - if self.fsync_mode == 'end': - try: - os.fsync(self.fd) - self.stats['sync_count'] += 1 - except OSError: - pass - elif self.fsync_mode == 'batch' and self.pending_sync_count > 0: - # Flush remaining pending writes - try: - os.fsync(self.fd) - self.stats['sync_count'] += 1 - self.pending_sync_count = 0 - except OSError: - pass - - def close(self, force_sync: bool = True): - """Close file - - Args: - force_sync: Whether to force fsync before closing - """ - # For backward compatibility with non-context-manager usage - if force_sync: - self._finalize_sync() - - if self.fd is not None: - os.close(self.fd) - self.fd = None - - def get_stats(self) -> Dict: - """Get statistics - - Returns: - Dict: Dictionary containing read/write statistics - """ - def calc_stats(latencies): - """Calculate latency statistics""" - if not latencies: - return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} - return { - 'avg_ms': statistics.mean(latencies), - **calc_percentiles(latencies), - } - - return { - 'read': { - 'count': self.stats['read_count'], - 'mb': self.stats['read_bytes'] / 1024 / 1024, - **calc_stats(self.stats['read_latencies_ms']) - }, - 'write': { - 'count': self.stats['write_count'], - 'mb': self.stats['write_bytes'] / 1024 / 1024, - **calc_stats(self.stats['write_latencies_ms']) - }, - 'sync_count': self.stats['sync_count'], - 'total_blocks': len(self.hash_id_to_offset), - 'free_blocks': len(self.free_offsets), - } - - -# ============================================================================ -# Benchmark Layer -# ============================================================================ - -class StorageBenchmark: - """KVCache storage benchmark - - Based on Mooncake OffsetAllocator + vLLM PagedAttention implementation: - - Example: - ----- - Request A: [1, 2, 4] - -> hash_id 1 -> not exist, write block_1 (offset=0, 1MB) - -> hash_id 2 -> not exist, write block_2 (offset=1, 1MB) - -> hash_id 4 -> not exist, write block_4 (offset=2, 1MB) - - Request B: [1, 2, 4, 6] - -> hash_id 1 -> exists, read block_1 (offset=0) ✓ prefix reuse - -> hash_id 2 -> exists, read block_2 (offset=1) ✓ prefix reuse - -> hash_id 4 -> exists, read block_4 (offset=2) ✓ prefix reuse - -> hash_id 6 -> not exist, write block_6 (offset=3, 1MB) - - Performance Advantages: - --------- - - Single file operation, no file explosion - - Offset reuse, reduces memory allocation - - pread/pwrite, thread-safe - """ - - def __init__(self, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN, - max_blocks: int = 100000, block_size_tokens: int = 512, - fsync_mode: str = 'batch', fsync_batch_size: int = 100): - """Initialize benchmark - - Args: - storage_dir: Storage directory - bytes_per_token: Bytes per token - max_blocks: Maximum number of blocks - block_size_tokens: Number of tokens per block - fsync_mode: When to fsync ('batch', 'always', 'end', 'none') - fsync_batch_size: Number of writes between fsync in batch mode - """ - self.storage = OffsetAllocatorStorage( - storage_dir, bytes_per_token, max_blocks, - block_size_tokens, fsync_mode, fsync_batch_size - ) - self.bytes_per_token = bytes_per_token - self.block_size_tokens = block_size_tokens - - # Statistics - self.stats = { - 'total_requests': 0, - 'total_blocks': 0, - 'read_blocks': 0, - 'write_blocks': 0, - 'prefix_hit_blocks': 0, # Number of prefix hit blocks - 'request_latencies_ms': [], - } - - def process_request(self, req: KVCacheRequest) -> float: - """Process a KVCache request - - Based on vLLM's prefix caching mechanism: - - Each hash_id corresponds to an independent block - - Prefix reuse achieved through hash_id matching - - Args: - req: KVCache request - - Returns: - float: Request latency in milliseconds - """ - self.stats['total_requests'] += 1 - self.stats['total_blocks'] += len(req.hash_ids) - - start_time = time.perf_counter() - total_latency = 0.0 - - # Process each hash_id (in order) - for hash_id in req.hash_ids: - if self.storage.block_exists(hash_id): - # Block exists, read (reuse cached block) - total_latency += self.storage.read_block(hash_id) - self.stats['read_blocks'] += 1 - self.stats['prefix_hit_blocks'] += 1 # Count all cache hits as prefix reuse - else: - # Block doesn't exist, write (new block) - total_latency += self.storage.write_block(hash_id) - self.stats['write_blocks'] += 1 - - latency_ms = total_latency if total_latency > 0 else MIN_LATENCY_MS - self.stats['request_latencies_ms'].append(latency_ms) - - return latency_ms - - def get_stats(self) -> Dict: - """Get statistics - - Returns: - Dict: Statistics dictionary - """ - storage_stats = self.storage.get_stats() - - request_latencies = self.stats['request_latencies_ms'] - - if request_latencies: - latency_stats = { - 'avg_ms': statistics.mean(request_latencies), - **calc_percentiles(request_latencies), - } - else: - latency_stats = {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} - - total_blocks = self.stats['total_blocks'] - read_blocks = self.stats['read_blocks'] - write_blocks = self.stats['write_blocks'] - - return { - 'total_requests': self.stats['total_requests'], - 'total_blocks': total_blocks, - 'read_blocks': read_blocks, - 'write_blocks': write_blocks, - 'prefix_hit_blocks': self.stats['prefix_hit_blocks'], - 'block_hit_rate': read_blocks / total_blocks if total_blocks > 0 else 0, - 'write_ratio': write_blocks / total_blocks if total_blocks > 0 else 0, - 'tokens_per_block': self.block_size_tokens, # Configurable block size in tokens - 'latency': latency_stats, - 'storage': storage_stats, - } - - def __enter__(self): - """Context manager entry""" - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """Context manager exit - ensures cleanup""" - self.close() - return False - - def close(self, force_sync: bool = True): - """Close storage - - Args: - force_sync: Whether to force final sync before closing - """ - self.storage.close(force_sync=force_sync) - - -# ============================================================================ -# Utility Functions -# ============================================================================ - -def calc_percentiles(data: List[float]) -> Dict[str, float]: - """Calculate latency percentiles - - Uses linear interpolation for accurate percentile calculation. - This is more accurate than statistics.quantiles() for small datasets. - - Args: - data: List of latency values in milliseconds - - Returns: - Dict containing p50, p95, p99 percentiles - """ - if not data: - return {'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} - - # Sort data for percentile calculation - sorted_data = sorted(data) - n = len(sorted_data) - - def get_percentile(p: float) -> float: - """Get percentile using linear interpolation - - Args: - p: Percentile (0-100) - - Returns: - Value at percentile - """ - index = (n - 1) * p / 100 - lower = int(index) - upper = min(lower + 1, n - 1) - - if lower == upper: - return sorted_data[lower] - - # Linear interpolation - weight = index - lower - return sorted_data[lower] * (1 - weight) + sorted_data[upper] * weight - - return { - 'p50_ms': get_percentile(50), - 'p95_ms': get_percentile(95), - 'p99_ms': get_percentile(99), - } - - -# ============================================================================ -# Trace Loader -# ============================================================================ - -class TraceLoader: - """Load KVCache trace""" - - def __init__(self, trace_path: str): - """Initialize trace loader - - Args: - trace_path: Trace file path - """ - self.trace_path = trace_path - self.requests = [] - self._load_trace() - - def _load_trace(self): - """Load trace file with error handling""" - line_num = 0 - try: - with open(self.trace_path, 'r') as f: - for line in f: - line_num += 1 - line = line.strip() - if not line: - continue - try: - req = json.loads(line) - # Validate required fields - if not all(k in req for k in ['timestamp', 'hash_ids', 'input_length', 'output_length']): - print(f"Warning: Line {line_num} missing required fields, skipping") - continue - if not isinstance(req['hash_ids'], list): - print(f"Warning: Line {line_num} has invalid hash_ids (not a list), skipping") - continue - self.requests.append(KVCacheRequest( - timestamp=float(req['timestamp']), - hash_ids=req['hash_ids'], - input_length=int(req['input_length']), - output_length=int(req['output_length']) - )) - except (json.JSONDecodeError, ValueError, KeyError) as e: - print(f"Warning: Line {line_num} has invalid format: {e}, skipping") - continue - except FileNotFoundError: - raise FileNotFoundError(f"Trace file not found: {self.trace_path}") - except OSError as e: - raise OSError(f"Error reading trace file {self.trace_path}: {e}") - - def get_requests(self) -> List[KVCacheRequest]: - """Get request list - - Returns: - List[KVCacheRequest]: Request list - """ - return self.requests - - -# ============================================================================ -# Benchmark Runner -# ============================================================================ - -def run_benchmark(trace_path: str, storage_dir: str, bytes_per_token: int = DEFAULT_BYTES_PER_TOKEN, - max_requests: Optional[int] = None, max_blocks: int = 100000, - replay_timestamps: bool = False, time_scale: float = 1.0, - block_size_tokens: int = 512, - fsync_mode: str = 'batch', fsync_batch_size: int = 100) -> Dict: - """Run benchmark - - Args: - trace_path: Trace file path - storage_dir: Storage directory - bytes_per_token: Bytes per token - max_requests: Maximum number of requests (None = all) - max_blocks: Maximum number of blocks - replay_timestamps: Whether to replay timestamps from trace (simulate realistic timing) - time_scale: Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed) - block_size_tokens: Number of tokens per block - fsync_mode: When to fsync ('batch', 'always', 'end', 'none') - fsync_batch_size: Number of writes between fsync in batch mode - - Returns: - Dict: Benchmark results - """ - block_size_bytes = block_size_tokens * bytes_per_token - - print(f"\n{'='*80}") - print(f"Running: {Path(trace_path).name}") - print(f"Architecture: Offset Allocator (Mooncake style)") - print(f"Block size: {block_size_tokens} tokens/block ({block_size_bytes:,} bytes)") - print(f"Storage: Single large file with offset-based block management") - print(f"Bytes per token: {bytes_per_token}") - print(f"Max blocks: {max_blocks}") - print(f"Fsync mode: {fsync_mode}" + (f" (batch_size={fsync_batch_size})" if fsync_mode == 'batch' else '')) - print(f"Timestamp replay: {'Enabled' if replay_timestamps else 'Disabled'}") - if replay_timestamps: - scale_desc = 'real-time' if time_scale == 1.0 else f'{1/time_scale:.1f}x speed' if time_scale < 1.0 else f'{time_scale}x slower' - print(f"Time scale: {time_scale}x ({scale_desc})") - print(f"{'='*80}") - - # Load trace - loader = TraceLoader(trace_path) - requests = loader.get_requests() - - if max_requests: - requests = requests[:max_requests] - - print(f"Loaded {len(requests)} requests") - - # Show timestamp range - if replay_timestamps and requests: - timestamps = [req.timestamp for req in requests] - time_span_ms = max(timestamps) - min(timestamps) - print(f"Timestamp range: {min(timestamps):.1f} - {max(timestamps):.1f} ms (span: {time_span_ms:.1f} ms)") - - # Create benchmark instance with context manager for cleanup - with StorageBenchmark( - storage_dir, bytes_per_token, max_blocks, - block_size_tokens, fsync_mode, fsync_batch_size - ) as benchmark: - - # Run benchmark - start_time = time.perf_counter() - total_io_time = 0.0 # Actual I/O time (excluding sleep) - last_timestamp = None - base_time = time.time() # Use wall time for replay synchronization - - for i, req in enumerate(requests): - # Replay by timestamps - sleep_time = 0.0 - if replay_timestamps and last_timestamp is not None: - # Calculate time interval from previous request - delta_ms = req.timestamp - last_timestamp - sleep_time = delta_ms / 1000.0 / time_scale # Apply time scaling - - if sleep_time > 0: - time.sleep(sleep_time) - - # Process request (measure I/O time) - req_start = time.perf_counter() - benchmark.process_request(req) - req_io_time = time.perf_counter() - req_start - total_io_time += req_io_time - - # Record current request timestamp - last_timestamp = req.timestamp - - # Progress output - if (i + 1) % 100 == 0: - if replay_timestamps: - elapsed_wall_time = time.time() - base_time - simulated_time = (req.timestamp - requests[0].timestamp) / 1000.0 / time_scale - print(f" Processed {i + 1}/{len(requests)}... (wall: {elapsed_wall_time:.1f}s, simulated: {simulated_time:.1f}s, io: {total_io_time:.1f}s)") - else: - print(f" Processed {i + 1}/{len(requests)}...") - - elapsed = time.perf_counter() - start_time - - # Perform final sync to include it in stats - benchmark.storage._finalize_sync() - - # Get statistics (context manager will handle cleanup) - stats = benchmark.get_stats() - - # Calculate actual I/O time (excluding sleep) - io_time = total_io_time if replay_timestamps else elapsed - - return { - 'trace_file': Path(trace_path).name, - 'total_requests': len(requests), - 'simulation_time_s': elapsed, - 'io_time_s': io_time, # Actual I/O time - 'wall_time_s': elapsed, # Wall time (including sleep) - 'requests_per_second': len(requests) / io_time if io_time > 0 else 0, # Based on I/O time - 'timestamp_replay_enabled': replay_timestamps, - 'time_scale': time_scale, - 'bytes_per_token': bytes_per_token, - 'block_size_tokens': block_size_tokens, - 'fsync_mode': fsync_mode, - **stats, - } - - -# ============================================================================ -# Result Output -# ============================================================================ - -def print_results(results: List[Dict]): - """Print benchmark results - - Args: - results: List of benchmark results - """ - for i, r in enumerate(results, 1): - print(f"\n{'='*80}") - print(f" [{i}/{len(results)}] {r['trace_file']}") - print(f"{'='*80}") - - print(f"\n[Performance Overview]") - print(f" Total Requests: {r['total_requests']:,}") - print(f" Queries Per Second (QPS): {r['requests_per_second']:.2f}") - print(f" Cache Hit Rate: {r['block_hit_rate']:.2%}") - print(f" Write Ratio: {r['write_ratio']:.2%}") - print(f" Total Blocks: {r['total_blocks']:,}") - print(f" Read Blocks: {r['read_blocks']:,}") - print(f" Write Blocks: {r['write_blocks']:,}") - print(f" Prefix Hits: {r['prefix_hit_blocks']:,}") - - print(f"\n[Latency Analysis]") - req_lat = r['latency'] - print(f" Request Latency (End-to-End): Avg={req_lat['avg_ms']:.2f}ms, P50={req_lat['p50_ms']:.2f}ms, P95={req_lat['p95_ms']:.2f}ms, P99={req_lat['p99_ms']:.2f}ms") - read_lat = r['storage']['read'] - write_lat = r['storage']['write'] - print(f" Single I/O Operation (Per Block):") - print(f" Read: Avg={read_lat.get('avg_ms', 0):.3f}ms, P50={read_lat.get('p50_ms', 0):.3f}ms, P95={read_lat.get('p95_ms', 0):.3f}ms, P99={read_lat.get('p99_ms', 0):.3f}ms") - print(f" Write: Avg={write_lat.get('avg_ms', 0):.3f}ms, P50={write_lat.get('p50_ms', 0):.3f}ms, P95={write_lat.get('p95_ms', 0):.3f}ms, P99={write_lat.get('p99_ms', 0):.3f}ms") - - print(f"\n[I/O & Bandwidth]") - print(f" Total Read I/O: {r['storage']['read']['mb']:>10.1f} MB ({r['storage']['read']['count']:,} ops)") - print(f" Total Write I/O: {r['storage']['write']['mb']:>10.1f} MB ({r['storage']['write']['count']:,} ops)") - io_time = r['io_time_s'] - bandwidth = (r['storage']['read']['mb'] + r['storage']['write']['mb']) / io_time - print(f" Effective Bandwidth: {bandwidth:>10.1f} MB/s") - - print(f"\n[Storage Details]") - print(f" Blocks in Use: {r['storage']['total_blocks']:>10,}") - print(f" Free Blocks: {r['storage']['free_blocks']:>10,}") - print(f" Tokens per Block: {r['tokens_per_block']:>10,}") - print(f" Block Size: {r['tokens_per_block'] * r.get('bytes_per_token', 2048) / 1024 / 1024:>10.2f} MB") - if 'sync_count' in r['storage']: - print(f" Fsync Operations: {r['storage']['sync_count']:>10,}") - - print(f"\n[Execution Time]") - if r.get('timestamp_replay_enabled'): - print(f" Wall Time (Total): {r['wall_time_s']:>10.2f} s") - print(f" I/O Time (Actual): {r['io_time_s']:>10.2f} s") - print(f" Sleep Time (Replay): {r['wall_time_s'] - r['io_time_s']:>10.2f} s") - else: - print(f" Total Execution Time: {r['wall_time_s']:>10.2f} s") - - print(f"\n{'='*80}\n") - - -# ============================================================================ -# Main Program -# ============================================================================ - -def main(): - """Main entry point""" - parser = argparse.ArgumentParser( - description='Mooncake KVCache Storage Benchmark', - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=""" -Examples: - # Quick test (100 requests) - python storage_benchmark.py --scenario=toolagent --max-requests=100 - - # Test with large model preset (Llama-3.1-405B) - python storage_benchmark.py --scenario=toolagent --model=llama-3.1-405b --max-requests=100 - - # Test with Deepseek V3 (extra large model) - python storage_benchmark.py --scenario=toolagent --model=deepseek-v3 --max-requests=100 - - # Realistic replay (with timestamps, 10x speed) - python storage_benchmark.py --scenario=toolagent --max-requests=1000 \\ - --replay-timestamps --time-scale=0.1 - - # All scenarios with custom bytes_per_token - python storage_benchmark.py --scenario=all --bytes-per-token=512 - - # Test with different block sizes and fsync modes - python storage_benchmark.py --scenario=toolagent --block-size-tokens=256 --fsync-mode=always - - # Test with custom fsync batch size - python storage_benchmark.py --scenario=toolagent --fsync-mode=batch --fsync-batch-size=50 - -Performance Tuning: - --fsync-mode=batch (default): Balance between performance and safety - --fsync-mode=always: Safest but slowest, measures full persistence cost - --fsync-mode=end: Fastest, only measures write I/O (not persistence) - --fsync-mode=none: Testing only, no durability guarantees - -Available model presets: - llama-3.1-405b, qwen3-32b, deepseek-v3, glm-4.6, default - -For more information: tools/STORAGE_BENCHMARK_README.md - """ - ) - - parser.add_argument('--trace-dir', type=str, default='../../FAST25-release/traces', - help='Trace files directory') - parser.add_argument('--scenario', type=str, choices=['conversation', 'synthetic', 'toolagent', 'all'], - default='toolagent', help='Test scenario') - parser.add_argument('--storage-dir', type=str, default='/tmp/mooncake_bench', - help='Storage directory') - parser.add_argument('--model', type=str, choices=list(MODEL_BYTES_PER_TOKEN.keys()), - default='default', - help=f'Model preset (overrides --bytes-per-token). Available: {", ".join(MODEL_BYTES_PER_TOKEN.keys())}') - parser.add_argument('--bytes-per-token', type=int, default=DEFAULT_BYTES_PER_TOKEN, - help='Bytes per token (default %d, overridden by --model if specified)' % DEFAULT_BYTES_PER_TOKEN) - parser.add_argument('--max-requests', type=int, default=None, - help='Maximum number of requests (default: unlimited)') - parser.add_argument('--max-blocks', type=int, default=100000, - help='Maximum number of blocks in storage file (determines file size)') - parser.add_argument('--replay-timestamps', action='store_true', - help='Enable timestamp replay (simulate realistic request timing)') - parser.add_argument('--time-scale', type=float, default=1.0, - help='Time scaling factor (1.0=real-time, 0.1=10x speed, 10.0=0.1x speed)') - parser.add_argument('--block-size-tokens', type=int, default=512, - help='Number of tokens per block (default: 512)') - parser.add_argument('--fsync-mode', type=str, choices=['batch', 'always', 'end', 'none'], - default='batch', - help='When to fsync: batch=every N writes (default), always=after each write, end=only at close, none=never') - parser.add_argument('--fsync-batch-size', type=int, default=100, - help='Number of writes between fsync in batch mode (default: 100)') - - args = parser.parse_args() - - # Print benchmark header - print(f"\n{'='*80}") - print(f"{'Mooncake KVCache Storage Benchmark':^80}") - print(f"{'='*80}") - - # Determine bytes_per_token (model preset takes precedence) - bytes_per_token = MODEL_BYTES_PER_TOKEN.get(args.model, args.bytes_per_token) - if args.model != 'default': - print(f"Using model preset: {args.model} ({bytes_per_token} bytes/token, ~{bytes_per_token/1024:.1f} KB/token)") - else: - print(f"Using custom bytes_per_token: {bytes_per_token}") - - # Determine test scenarios - scenarios = ['conversation', 'synthetic', 'toolagent'] if args.scenario == 'all' else [args.scenario] - trace_files = { - 'conversation': 'conversation_trace.jsonl', - 'synthetic': 'synthetic_trace.jsonl', - 'toolagent': 'toolagent_trace.jsonl' - } - - # Run benchmarks - results = [] - - for scenario in scenarios: - trace_path = Path(args.trace_dir) / trace_files[scenario] - if trace_path.exists(): - result = run_benchmark( - str(trace_path), - str(Path(args.storage_dir) / scenario), - bytes_per_token, - args.max_requests, - args.max_blocks, - args.replay_timestamps, - args.time_scale, - args.block_size_tokens, - args.fsync_mode, - args.fsync_batch_size - ) - results.append(result) - else: - print(f"Warning: Trace file not found: {trace_path}") - - # Print results - if results: - print_results(results) - - -if __name__ == '__main__': - main() diff --git a/benchmarks/storage_benchmark_v1/benchmark.py b/benchmarks/storage_benchmark_v1/benchmark.py index 3f5ba5ed38..f840732393 100644 --- a/benchmarks/storage_benchmark_v1/benchmark.py +++ b/benchmarks/storage_benchmark_v1/benchmark.py @@ -10,10 +10,11 @@ import sys import time import statistics -import signal +from contextlib import ExitStack +from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from pathlib import Path -from typing import Iterator, List, Dict, Any +from typing import List, Dict, Any from storage import DiskHashTable from layout import get_model_config, create_layout @@ -103,7 +104,8 @@ def __init__(self, storage_dir: str, model_config: dict, 'read_pages': 0, 'write_pages': 0, 'page_hits': 0, - 'request_latencies_ms': [], + 'request_io_latencies_ms': [], + 'request_wall_latencies_ms': [], } def process_request(self, req: KVCacheRequest) -> float: @@ -118,13 +120,14 @@ def process_request(self, req: KVCacheRequest) -> float: self.stats['total_requests'] += 1 self.stats['total_tokens'] += req.input_length + req.output_length - total_latency = 0.0 + 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 - total_latency += self.storage.read( + io_latency_ms += self.storage.read( access.page_id, offset_in_page=access.offset_in_page, length=access.length @@ -133,39 +136,23 @@ def process_request(self, req: KVCacheRequest) -> float: self.stats['page_hits'] += 1 else: # Page doesn't exist, perform WRITE - total_latency += self.storage.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 - latency_ms = total_latency if total_latency > 0 else 0.0 - if latency_ms > 0: - self.stats['request_latencies_ms'].append(latency_ms) - return latency_ms + 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_latencies = self.stats['request_latencies_ms'] - - if request_latencies: - sorted_latencies = sorted(request_latencies) - n = len(sorted_latencies) - - def get_percentile(p: float) -> float: - idx = int(n * p) - return sorted_latencies[idx] if idx < n else sorted_latencies[-1] - - latency_stats = { - 'avg_ms': statistics.mean(request_latencies), - 'p50_ms': sorted_latencies[n // 2], - 'p95_ms': get_percentile(0.95), - 'p99_ms': get_percentile(0.99), - } - else: - latency_stats = {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} + 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'] @@ -178,7 +165,8 @@ def get_percentile(p: float) -> float: '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, - 'latency': latency_stats, + 'request_io_latency': latency_stats(request_io_latencies), + 'request_wall_latency': latency_stats(request_wall_latencies), 'storage': storage_stats, } @@ -205,10 +193,228 @@ def get_max_page_id(requests: List[KVCacheRequest]) -> int: 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) -> Dict: + 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: @@ -220,6 +426,9 @@ def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, 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 @@ -229,6 +438,8 @@ def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, 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 @@ -266,7 +477,11 @@ def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, 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 @@ -279,65 +494,69 @@ def run_benchmark(trace_path: str, storage_dir: str, model_config: dict, surplus_pct = (surplus / max_pages) * 100 if max_pages > 0 else 0 print(f" ✓ Direct mapping: all {max_pages_needed:,} logical pages uniquely mapped") - # Run benchmark - 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: - start_time = time.perf_counter() - try: - for i, req in enumerate(requests): - benchmark.process_request(req) - # Print progress for each request - elapsed = time.perf_counter() - start_time - qps = (i + 1) / elapsed if elapsed > 0 else 0 - stats = benchmark.get_stats() - storage = stats.get('storage', {}) - read_latency = storage.get('read', {}).get('avg_ms', 0) - write_latency = storage.get('write', {}).get('avg_ms', 0) - read_mb = storage.get('read', {}).get('mb', 0) - write_mb = storage.get('write', {}).get('mb', 0) - read_time = storage.get('read', {}).get('time_s', 0) - write_time = storage.get('write', {}).get('time_s', 0) - read_mbps = read_mb / read_time if read_time > 0 else 0 - write_mbps = write_mb / write_time if write_time > 0 else 0 - print(f" [{i+1:5d}/{len(requests)}] ids={len(req.hash_ids):3d} " - f"tokens={req.input_length+req.output_length:6d} | " - f"QPS={qps:7.2f} | " - f"R={stats['read_pages']:6d} ({read_latency:6.2f}ms, {read_mbps:6.1f}MB/s) | " - f"W={stats['write_pages']:6d} ({write_latency:6.2f}ms, {write_mbps:6.1f}MB/s)") - except KeyboardInterrupt: - print(f"\n\n{'='*80}") - print(f"Interrupted! Showing partial results:") - print(f"{'='*80}") - elapsed = time.perf_counter() - start_time - stats = benchmark.get_stats() - print_results([{ - 'trace_file': Path(trace_path).name, - 'total_requests': i + 1, - 'io_time_s': elapsed, - 'requests_per_second': (i + 1) / elapsed if elapsed > 0 else 0, - 'model': model_config['name'], - 'fsync_mode': fsync_mode, - **stats, - }]) - sys.exit(0) - - elapsed = time.perf_counter() - start_time - stats = benchmark.get_stats() + 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': len(requests), - 'io_time_s': elapsed, - 'requests_per_second': len(requests) / elapsed if elapsed > 0 else 0, + '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, - **stats, + 'threads': threads, + 'replay_scale': replay_scale, + **result['stats'], } @@ -350,6 +569,8 @@ def format_storage_stats(stats: Dict, title: str = "Storage"): 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}]") @@ -357,12 +578,28 @@ def format_storage_stats(stats: Dict, title: str = "Storage"): # 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):,}") @@ -438,8 +675,18 @@ def main(): 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}") @@ -447,6 +694,11 @@ def main(): 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] @@ -458,20 +710,28 @@ def main(): # 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(): - result = run_benchmark( - str(trace_path), - str(Path(args.storage_dir) / scenario), - model_config, - args.max_requests, - args.max_pages, - args.page_size_tokens, - args.fsync_mode, - args.fsync_batch_size - ) - results.append(result) + 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}") diff --git a/benchmarks/storage_benchmark_v1/doc/README.md b/benchmarks/storage_benchmark_v1/doc/README.md index 4104a6136e..8b53be660a 100644 --- a/benchmarks/storage_benchmark_v1/doc/README.md +++ b/benchmarks/storage_benchmark_v1/doc/README.md @@ -9,6 +9,7 @@ The KVCache Storage Benchmark is a tool for evaluating storage performance of KV ### 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 @@ -25,14 +26,49 @@ python benchmark.py --scenario conversation \ | `--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` | +| `--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, each request displays real-time statistics: +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) @@ -55,12 +91,26 @@ Fields: [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 @@ -89,6 +139,28 @@ Fields: 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: diff --git a/benchmarks/storage_benchmark_v1/storage/disk.py b/benchmarks/storage_benchmark_v1/storage/disk.py index 286ff7510f..b54de98a64 100644 --- a/benchmarks/storage_benchmark_v1/storage/disk.py +++ b/benchmarks/storage_benchmark_v1/storage/disk.py @@ -7,7 +7,7 @@ import os import time from pathlib import Path -from typing import Dict, List, Optional, Any +from typing import Dict, Any from .interface import Storage @@ -18,11 +18,16 @@ def calc_percentiles(data): return {'avg_ms': 0, 'p50_ms': 0, 'p95_ms': 0, 'p99_ms': 0} import statistics sorted_data = sorted(data) - n = len(sorted_data) + def get_percentile(p): - idx = int(n * p / 100) - if idx >= n: idx = n - 1 - return sorted_data[idx] + 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), diff --git a/dependencies.sh b/dependencies.sh index 7e064650e4..c1e42a772c 100755 --- a/dependencies.sh +++ b/dependencies.sh @@ -136,6 +136,8 @@ 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" @@ -156,7 +158,6 @@ if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then unzip \ libibverbs-dev \ libgoogle-glog-dev \ - libgtest-dev \ libjsoncpp-dev \ libunwind-dev \ libnuma-dev \ @@ -173,6 +174,7 @@ if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then liburing-dev \ libjemalloc-dev \ libmsgpack-dev \ + libzmq3-dev \ libzstd-dev \ libasio-dev \ libxxhash-dev \ @@ -187,25 +189,26 @@ if [ "$OS" = "ubuntu" ] || [ "$OS" = "debian" ]; then 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 \ - gtest-devel \ + gflags-devel \ jsoncpp-devel \ libunwind-devel \ numactl-devel \ python3-devel \ - boost-devel \ + boost1.78-devel \ openssl-devel \ - grpc-devel \ protobuf-devel \ yaml-cpp-devel \ - grpc-plugins \ libcurl-devel \ hiredis-devel \ liburing-devel \ jemalloc-devel \ + msgpack-devel \ + libzstd-devel \ pkgconf-pkg-config \ elfutils-libelf-devel \ patchelf \ 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 3dda2e6cc2..c0a9c8d760 100644 --- a/docker/mooncake.Dockerfile +++ b/docker/mooncake.Dockerfile @@ -9,16 +9,18 @@ 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 EP_TORCH_VERSIONS="2.13.0" ARG TORCH_CUDA_ARCH_LIST="8.0;9.0" +# 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}" @@ -50,7 +52,9 @@ COPY . /workspace # Install Mooncake dependencies (yalantinglibs, 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 -G Ninja .. \ @@ -63,18 +67,19 @@ RUN mkdir -p build && \ -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 . - -# 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 diff --git a/docker/musa.Dockerfile b/docker/musa.Dockerfile index d47564a321..d33f65d263 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, yalantinglibs, 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 95% rename from requirements_docs.txt rename to docs/requirements-docs.txt index 2aacb26838..66a4886911 100644 --- a/requirements_docs.txt +++ b/docs/requirements-docs.txt @@ -4,6 +4,7 @@ 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 @@ -17,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 index 35e77191e1..089c657c8b 100644 --- a/docs/source/api-reference/cpp/index.md +++ b/docs/source/api-reference/cpp/index.md @@ -2,15 +2,15 @@ | Module | Description | |--------|-------------| -| [Transfer Engine C++ API](../../design/transfer-engine/cpp-api) | `TransferEngine` class — memory registration, batch transfer, segment management, RDMA transport | -| [TENT C++ API](../../design/tent/cpp-api) | `mooncake::tent::TransferEngine` — next-gen transfer engine with automatic transport selection and fault tolerance | +| [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, `BufferAllocatorBase` interface | :::{toctree} :maxdepth: 1 :hidden: -../../design/transfer-engine/cpp-api -../../design/tent/cpp-api +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 index 41c3028a0f..779655148c 100644 --- a/docs/source/api-reference/cpp/mooncake-store.md +++ b/docs/source/api-reference/cpp/mooncake-store.md @@ -99,6 +99,11 @@ tl::expected CreateCopyTask( - 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++ @@ -118,6 +123,12 @@ tl::expected CreateMoveTask( - 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++ @@ -140,6 +151,17 @@ struct QueryTaskResponse { }; ``` +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++ diff --git a/docs/source/design/tent/cpp-api.md b/docs/source/api-reference/cpp/tent.md similarity index 95% rename from docs/source/design/tent/cpp-api.md rename to docs/source/api-reference/cpp/tent.md index fc68a22c62..22b56a0c68 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 | @@ -200,7 +201,7 @@ 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](qos.md). +- `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. @@ -274,6 +275,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 diff --git a/docs/source/design/transfer-engine/cpp-api.md b/docs/source/api-reference/cpp/transfer-engine.md similarity index 99% rename from docs/source/design/transfer-engine/cpp-api.md rename to docs/source/api-reference/cpp/transfer-engine.md index 44c9ff9d8d..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. diff --git a/docs/source/design/conductor/indexer-api-design.md b/docs/source/api-reference/http/conductor-indexer.md similarity index 61% rename from docs/source/design/conductor/indexer-api-design.md rename to docs/source/api-reference/http/conductor-indexer.md index e79367e6cd..0abd67e41c 100644 --- a/docs/source/design/conductor/indexer-api-design.md +++ b/docs/source/api-reference/http/conductor-indexer.md @@ -360,3 +360,134 @@ 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/http-api-reference/http-service.md b/docs/source/api-reference/http/http-service.md similarity index 93% rename from docs/source/http-api-reference/http-service.md rename to docs/source/api-reference/http/http-service.md index cbd79e007d..8c8ee90fdd 100644 --- a/docs/source/http-api-reference/http-service.md +++ b/docs/source/api-reference/http/http-service.md @@ -91,7 +91,7 @@ curl "http://localhost:8080/query_key?key=my_object" ``` #### `/batch_query_keys` -Retrieve replica information for multiple keys in a single request, including memory locations and transport endpoints for each key. +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) @@ -115,6 +115,25 @@ curl "http://localhost:8080/batch_query_keys?keys=key1,key2,key3" "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": { @@ -125,6 +144,8 @@ curl "http://localhost:8080/batch_query_keys?keys=key1,key2,key3" } ``` +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. diff --git a/docs/source/api-reference/http/index.md b/docs/source/api-reference/http/index.md index e946409f9d..c67a9fbede 100644 --- a/docs/source/api-reference/http/index.md +++ b/docs/source/api-reference/http/index.md @@ -2,13 +2,13 @@ | Module | Description | |--------|-------------| -| [HTTP Service](../../http-api-reference/http-service) | RESTful endpoints for cluster management, metrics, and data inspection | -| [Conductor Indexer API](../../design/conductor/indexer-api-design) | Cache-aware routing: service registration, KV event subscription, prefix cache-hit query APIs | +| [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-api-reference/http-service -../../design/conductor/indexer-api-design +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/python-api-reference/ep-backend.md b/docs/source/api-reference/python/ep-backend.md similarity index 97% rename from docs/source/python-api-reference/ep-backend.md rename to docs/source/api-reference/python/ep-backend.md index 5619349ff0..a3c12a38dd 100644 --- a/docs/source/python-api-reference/ep-backend.md +++ b/docs/source/api-reference/python/ep-backend.md @@ -17,8 +17,8 @@ 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 Backend (PG) design guide](../design/mooncake-backend-pg.md) -and the [Mooncake EP design guide](../design/mooncake-ep.md). +For implementation details, see the [Mooncake Backend (PG) design guide](../../design/mooncake-backend-pg.md) +and the [Mooncake EP design guide](../../design/mooncake-ep.md). ## Installation and build notes @@ -132,7 +132,6 @@ Arguments: | `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_preferred_hca(backend, location)` | Query topology-preferred HCA for a location. | Useful for topology-aware placement/debugging. | | `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 ranks synchronized by the backend. | Diagnostic helper. | | `pg.extend_group_size_to(backend, size)` | Reserve additional inactive ranks. | Newly extended ranks do not participate until recovered. | @@ -422,5 +421,5 @@ shape consistent with the process group world size or reserved `max_world_size`. - EP correctness and failure simulation: `mooncake-ep/tests/test_ep_grid.py` - Wheel-level EP example: `mooncake-wheel/tests/test_mooncake_ep.py` -See [PG/EP troubleshooting](../troubleshooting/pg-ep-troubleshooting.md) for +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 index 75ca5c5a6c..4c61577c59 100644 --- a/docs/source/api-reference/python/index.md +++ b/docs/source/api-reference/python/index.md @@ -2,15 +2,17 @@ | Module | Description | |--------|-------------| -| [Mooncake Store](../../python-api-reference/mooncake-store) | Distributed KV cache storage client — `put`/`get`/`remove`/`replicate` operations | -| [Transfer Engine](../../python-api-reference/transfer-engine) | High-performance RDMA/TCP data transfer between nodes | -| [EP Backend](../../python-api-reference/ep-backend) | Expert-parallel backend for large MoE model deployment | +| [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: -../../python-api-reference/mooncake-store -../../python-api-reference/transfer-engine -../../python-api-reference/ep-backend +mooncake-store +dataproto-structured-object-transfer +transfer-engine +ep-backend ::: diff --git a/docs/source/python-api-reference/mooncake-store.md b/docs/source/api-reference/python/mooncake-store.md similarity index 91% rename from docs/source/python-api-reference/mooncake-store.md rename to docs/source/api-reference/python/mooncake-store.md index 3125b72378..fe6d71d3da 100644 --- a/docs/source/python-api-reference/mooncake-store.md +++ b/docs/source/api-reference/python/mooncake-store.md @@ -12,9 +12,13 @@ pip install mooncake-transfer-engine 📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine/](https://pypi.org/project/mooncake-transfer-engine/) ### Required Service -Only one service is required now: +The only always-required service is: -- `mooncake_master` — Master service which now embeds the HTTP metadata server +- `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 @@ -60,13 +64,17 @@ print(data.decode()) # Output: Hello, Mooncake Store! store.close() ``` -**RDMA device selection**: Leave `rdma_devices` as `""` to auto-select RDMA NICs. Provide a comma-separated list (e.g. `"mlx5_0,mlx5_1"`) to pin to specific hardware. +**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 (preview) +#### P2P Hello World -The following setup uses the new P2P handshake and does not require an HTTP metadata server. This feature is not released yet; use only if you’re testing the latest code. +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 @@ -183,6 +191,27 @@ 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)`: @@ -937,11 +966,14 @@ print("Retrieved all keys successfully:", retrieved == values) ## Topology & Devices -- Auto-discovery: Disabled by default. For `protocol="rdma"`, you must specify RDMA devices. -- Enable auto-discovery (optional): - - `MC_MS_AUTO_DISC=1` enables auto-discovery; then `rdma_devices` is not required. - - Optionally restrict candidates with `MC_MS_FILTERS`, a comma-separated whitelist of NIC names, e.g. `MC_MS_FILTERS=mlx5_0,mlx5_2`. - - If `MC_MS_AUTO_DISC` is not set or set to `0`, auto-discovery remains disabled and `rdma_devices` is required for RDMA. +- 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: @@ -1027,17 +1059,35 @@ def setup( protocol: str = "tcp", rdma_devices: str = "", master_server_addr: str, + engine: Optional[TransferEngine] = None, + enable_ssd_offload: bool = False, + ssd_offload_path: str = "", + tenant_id: str = "default", ) -> int ``` **Parameters:** - `local_hostname` (str): **Required**. Local hostname and port (e.g., "localhost" or "localhost:12345") -- `metadata_server` (str): **Required**. Metadata server address (e.g., "http://localhost:8080/metadata") -- `global_segment_size` (int): Memory segment size in bytes for mounting (default: 16MB = 16777216) -- `local_buffer_size` (int): Local buffer size in bytes (default: 1GB = 1073741824) -- `protocol` (str): Network protocol - "tcp" or "rdma" (default: "tcp") -- `rdma_devices` (str): RDMA device name(s), e.g. `"mlx5_0"` or `"mlx5_0,mlx5_1"`. Leave empty to auto-select NICs. Provide device names to pin the NICs. Always empty for TCP. +- `metadata_server` (str): **Required**. Metadata connection string, e.g. `"P2PHANDSHAKE"` or `"http://localhost:8080/metadata"`. +- `global_segment_size` (int): Memory segment size in bytes for mounting. +- `local_buffer_size` (int): Local buffer size in bytes. +- `protocol` (str): Network protocol, usually `"tcp"`, `"rdma"`, `"efa"`, `"cxl"`, or `"ascend"` depending on the build. +- `rdma_devices` (str): 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**. Master server address (e.g., "localhost:50051") +- `engine` (Optional[TransferEngine]): Existing Transfer Engine instance to reuse. Defaults to `None`. +- `enable_ssd_offload` (bool): Enable client-side SSD offload support. Defaults to `False`. +- `ssd_offload_path` (str): SSD offload directory. When provided, overrides the storage path environment configuration. +- `tenant_id` (str): Tenant namespace for object keys. Defaults to `"default"`. + +**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) @@ -1082,6 +1132,22 @@ def setup_dummy(self, mem_pool_size: int, local_buffer_size: int, server_address 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() @@ -1559,7 +1625,8 @@ def batch_get_buffer(self, keys: List[str]) -> List[BufferHandle] **Returns:** - `List[BufferHandle]`: List of buffer objects, with None for keys not found -**Note:** This function is not supported for dummy client. +**Note:** This function is supported for dummy clients through the real +client-owned shared-memory staging buffer. **Example:** ```python @@ -1850,6 +1917,12 @@ def create_copy_task(self, key: str, targets: List[str]) -> Tuple[UUID, int] - 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 @@ -1884,6 +1957,12 @@ def create_move_task(self, key: str, source: str, target: str) -> Tuple[UUID, in - 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 @@ -1916,6 +1995,26 @@ def query_task(self, task_id: UUID) -> Tuple[QueryTaskResponse | None, int] - 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 @@ -1993,7 +2092,8 @@ def put_from_with_metadata(self, key: str, buffer_ptr: int, metadata_buffer_ptr: **Returns:** - `int`: Status code (0 = success, non-zero = error code) -**Note:** This function is not supported for dummy client. +**Note:** This function is supported for dummy clients through the real +client-owned shared-memory staging buffer. **Example:** ```python @@ -2393,7 +2493,8 @@ def upsert_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int **Returns:** - `int`: Status code (0 = success, non-zero = error code) -**Note:** This function is not supported for dummy client. +**Note:** This function is supported for dummy clients through the real +client-owned shared-memory staging buffer. #### batch_upsert_tensor_from() @@ -2427,7 +2528,9 @@ def batch_upsert_tensor(self, keys: List[str], tensors_list: List[torch.Tensor]) **Returns:** - `List[int]`: List of status codes for each tensor operation. -**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client. +**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() @@ -2445,7 +2548,9 @@ def upsert_pub_tensor(self, key: str, tensor: torch.Tensor, config: ReplicateCon **Returns:** - `int`: Status code (0 = success, non-zero = error code) -**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client. +**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 @@ -2479,7 +2584,9 @@ def batch_upsert_pub_tensor(self, keys: List[str], tensors_list: List[torch.Tens **Returns:** - `List[int]`: List of status codes for each tensor operation. -**Note:** This function requires `torch` to be installed and available in the environment. Not supported for dummy client. +**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. --- @@ -2897,7 +3004,7 @@ 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. +📋 **Complete Error Codes Reference**: See [Error Code Explanation](../../troubleshooting/error-code.md) for detailed descriptions of all error codes and their meanings. --- diff --git a/docs/source/python-api-reference/transfer-engine.md b/docs/source/api-reference/python/transfer-engine.md similarity index 98% rename from docs/source/python-api-reference/transfer-engine.md rename to docs/source/api-reference/python/transfer-engine.md index 59ab17933c..9958d9edf7 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 @@ -165,7 +167,7 @@ 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. +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() diff --git a/docs/source/api-reference/rust/index.md b/docs/source/api-reference/rust/index.md new file mode 100644 index 0000000000..30ab8e8325 --- /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) | Rust bindings for Transfer Engine (FFI wrapper used by the Rust example crate) | + +```{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..ed3d13e432 --- /dev/null +++ b/docs/source/api-reference/rust/transfer-engine.md @@ -0,0 +1,206 @@ +# Transfer Engine Rust API + +This page documents the Rust bindings living under `mooncake-transfer-engine/rust`. + +At the time of writing, the Rust package (`transfer_engine_rust`) is primarily used as a **Rust-side binding + example binary**. The public Rust types are implemented in `src/transfer_engine.rs` and wrap the Transfer Engine C API (`transfer_engine_c.h`). + +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` + +## Build & runtime prerequisites + +The Rust package uses bindgen + CMake to link against the Transfer Engine C/C++ build outputs. + +- **Build**: + - Requires Rust toolchain and libclang (bindgen). + - The package has `build.rs` that expects to find / build the native library via CMake. +- **Runtime**: + - Dynamic linker must find Transfer Engine shared libraries. + - You need a metadata server backend (commonly etcd) and a reachable peer segment registry. + +## Mental model + +The Transfer Engine operates on **segments** and **transfer batches**: + +- You create a `TransferEngine` bound to: + - `metadata_uri` (for example, etcd endpoint) + - `local_server_name` (this node's address/name) + - `rpc_port` (RPC listener port) +- You register local memory regions as RDMA-capable buffers. +- You open a remote segment to obtain a `segment_id` (an integer handle). +- You allocate a batch id for a fixed number of transfer requests. +- You submit a batch of `TransferRequest`. +- You poll status per task id inside the batch, then free the batch id. + +## API reference + +### Enums + +#### `OpcodeEnum` + +- `OpcodeEnum::Read` +- `OpcodeEnum::Write` + +Used by `TransferRequest.opcode`. + +#### `TransferStatusEnum` + +Status values returned by the C layer. Common values you will check for: + +- `Completed` +- `Failed` +- `Timeout` + +### Structs + +#### `TransferRequest` + +One transfer operation inside a batch. + +Fields: + +- `opcode: OpcodeEnum` +- `source: *mut c_void`: local source/destination pointer (depends on opcode). +- `target_id: i32`: segment id returned by `open_segment()`. +- `target_offset: u64`: byte offset inside the target segment. +- `length: u64`: transfer length in bytes. + +#### `BufferEntry` + +Used for batch memory registration: + +- `addr: *mut c_void` +- `length: u64` + +### `TransferEngine` + +#### `new(metadata_uri, local_server_name, rpc_port) -> anyhow::Result` + +Create a new engine handle. + +Notes: + +- `metadata_uri` and `local_server_name` are passed through `CString`; interior `\0` bytes will error. +- The wrapper currently disables `auto_discover` in the underlying C call. + +#### `discover_topology() -> anyhow::Result<()>` + +Trigger topology discovery. + +#### `install_transport(proto) -> anyhow::Result<()>` + +Install a transport by name (e.g. `"tcp"`, `"rdma"`, `"efa"` depending on build/runtime support). + +#### `register_local_memory(addr, length, location) -> anyhow::Result<()>` + +Register a local memory region for zero-copy transfers. + +- `addr`: pointer to the memory region. +- `length`: size in bytes. +- `location`: location string such as `"cpu:0"`. + +#### `unregister_local_memory(addr) -> anyhow::Result<()>` + +Unregister a previously registered memory region. + +#### `register_local_memory_batch(buffer_list, location) -> anyhow::Result<()>` + +Batch register multiple local buffers. + +- No-op when `buffer_list` is empty. + +#### `unregister_local_memory_batch(buffer_list) -> anyhow::Result<()>` + +Batch unregister multiple local buffers. + +#### `open_segment(name: String) -> anyhow::Result` + +Open a remote segment by name and get a segment id. + +#### `close_segment(segment_id: i32) -> anyhow::Result<()>` + +Close a previously opened segment. + +#### `warmup_efa_segment(name: &str) -> anyhow::Result<()>` + +Eagerly establish EFA endpoints so the first `submit_transfer()` does not pay the serial connection setup cost. + +- No-op on non-EFA transports. +- Call after `open_segment()` and after the metadata server has published the peer's NIC list. + +#### `sync_segment_cache() -> anyhow::Result<()>` + +Synchronize segment cache from metadata. + +#### `allocate_batch_id(batch_size) -> anyhow::Result` + +Allocate a batch id for `batch_size` transfer requests. + +You must call `free_batch_id(batch_id)` after all tasks are done. + +#### `submit_transfer(batch_id, requests) -> anyhow::Result<()>` + +Submit a batch transfer request list. + +- No-op when `requests` is empty. +- The wrapper converts each `TransferRequest` into the C representation (`transfer_request_t`). + +#### `get_transfer_status(batch_id, task_id) -> anyhow::Result<(i32, u64)>` + +Get status for one task in a batch. + +- `task_id` is an index inside the batch, typically `0..batch_size`. +- Returns `(status_code, transferred_bytes)`. + +The `status_code` maps to values in `TransferStatusEnum` (represented as `i32`). + +#### `free_batch_id(batch_id) -> anyhow::Result<()>` + +Free a previously allocated batch id. + +## Minimal usage example (pseudo-code) + +The crate's `src/main.rs` contains a full benchmark-style example. The following sketch shows the typical control flow: + +```rust +use std::ffi::c_void; +use transfer_engine_rust::transfer_engine::{OpcodeEnum, TransferEngine, TransferRequest}; + +fn main() -> anyhow::Result<()> { + let engine = TransferEngine::new("127.0.0.1:2379", "127.0.0.1", 12345)?; + engine.discover_topology()?; + engine.install_transport("tcp")?; + + // Register local memory (example only; you must allocate and pin memory appropriately). + let mut buffer = vec![0u8; 4096]; + engine.register_local_memory(buffer.as_mut_ptr() as *mut c_void, buffer.len(), "cpu:0")?; + + let seg_id = engine.open_segment("target-seg".to_string())?; + let batch_id = engine.allocate_batch_id(1)?; + + let mut reqs = [TransferRequest { + opcode: OpcodeEnum::Write, + source: buffer.as_mut_ptr() as *mut c_void, + target_id: seg_id, + target_offset: 0, + length: buffer.len() as u64, + }]; + + engine.submit_transfer(batch_id, &mut reqs)?; + let (status, bytes) = engine.get_transfer_status(batch_id, 0)?; + println!("status={status}, bytes={bytes}"); + + engine.free_batch_id(batch_id)?; + engine.close_segment(seg_id)?; + engine.unregister_local_memory(buffer.as_mut_ptr() as *mut c_void)?; + Ok(()) +} +``` + +## Safety & thread-safety + +- The wrapper marks `TransferEngine` as `Send + Sync`, but it owns an FFI handle (`transfer_engine_t`). +- All pointer-based arguments must satisfy Rust’s aliasing and lifetime rules. +- You must ensure registered memory remains valid until it is unregistered. diff --git a/docs/source/conf.py b/docs/source/conf.py index a505aa4f78..c10fad2189 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -45,6 +45,7 @@ "sphinxarg.ext", "sphinx_design", "sphinx_togglebutton", + "sphinx_reredirects", "sphinxcontrib.mermaid", ] myst_enable_extensions = [ @@ -85,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, } @@ -249,3 +251,87 @@ 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 = { + "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 97% rename from docs/source/getting_started/examples/lmcache-integration.md rename to docs/source/deployment/integrations/lmcache/index.md index 3df78327e7..8723f4a8bc 100644 --- a/docs/source/getting_started/examples/lmcache-integration.md +++ b/docs/source/deployment/integrations/lmcache/index.md @@ -28,8 +28,8 @@ For a complete deployment guide with step-by-step instructions, see: :::{toctree} :maxdepth: 1 -../examples/vllm-integration/vllmv1-lmcache-integration -../examples/vllm-integration/vllmv1-lmcache-mp-integration +vllmv1-lmcache-integration +vllmv1-lmcache-mp-integration ::: ## Performance Benchmarking and Results diff --git a/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-integration.md b/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-integration.md similarity index 98% rename from docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-integration.md rename to docs/source/deployment/integrations/lmcache/vllmv1-lmcache-integration.md index 024bd760a4..df7fdd5961 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-integration.md +++ b/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-integration.md @@ -28,7 +28,7 @@ your environment. Install Mooncake, vLLM, and LMCache on both machines. For installation details, refer to the official documentation of each project: -- [Mooncake build guide](../../build.md) +- [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/) @@ -186,6 +186,6 @@ together: ## Additional Resources -* [Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System](../lmcache-integration.md) +* [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/getting_started/examples/vllm-integration/vllmv1-lmcache-mp-integration.md b/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-mp-integration.md similarity index 98% rename from docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-mp-integration.md rename to docs/source/deployment/integrations/lmcache/vllmv1-lmcache-mp-integration.md index d728753a2b..7af01e6ad4 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllmv1-lmcache-mp-integration.md +++ b/docs/source/deployment/integrations/lmcache/vllmv1-lmcache-mp-integration.md @@ -35,7 +35,7 @@ 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. +`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` @@ -188,6 +188,6 @@ together: ## Additional Resources -* [Mooncake x LMCache: Unite to Pioneer KVCache-Centric LLM Serving System](../lmcache-integration.md) +* [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 91% rename from docs/source/getting_started/examples/lmdeploy-integration-v0.9.md rename to docs/source/deployment/integrations/lmdeploy.md index ac6fd6fbe3..50c9f9d7b5 100644 --- a/docs/source/getting_started/examples/lmdeploy-integration-v0.9.md +++ b/docs/source/deployment/integrations/lmdeploy.md @@ -16,7 +16,7 @@ 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](../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 @@ -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 ``` @@ -80,10 +80,10 @@ lmdeploy serve api_server Qwen/Qwen3-8B \ ```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 ``` 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 98% 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 4a6e011cb0..7078faa02e 100644 --- a/docs/source/getting_started/examples/sglang-integration/hicache-integration-v1.md +++ b/docs/source/deployment/integrations/sglang/hicache-integration-v1.md @@ -97,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):** @@ -326,7 +326,7 @@ python -m sglang.launch_server \ 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](../sglang-integration-v1) to set up the prefill, decode, and router workers. +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. 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 98% 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 a338ce2903..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. In benchmarks, pre-populated Mooncake achieves **best TTFT** across all tiers, maintaining high cache hit rates as conversation rounds grow ([details](../../../performance/sglang-hicache-benchmark-results-v1)). +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). diff --git a/docs/source/getting_started/examples/sglang-integration/index.md b/docs/source/deployment/integrations/sglang/index.md similarity index 88% rename from docs/source/getting_started/examples/sglang-integration/index.md rename to docs/source/deployment/integrations/sglang/index.md index 885c444f2d..8926ad3c69 100644 --- a/docs/source/getting_started/examples/sglang-integration/index.md +++ b/docs/source/deployment/integrations/sglang/index.md @@ -15,9 +15,9 @@ SGLang uses Mooncake's Transfer Engine for direct zero-copy KV cache transfer be +-----------+ +------------+ ``` -**Related:** [Full PD Disaggregation Guide](../sglang-integration-v1) — installation, cross-node/same-node setup, XpYd topology, EP backend for MoE models, and EPD backend for multimodal models. +**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-benchmark-results-v1) — compares 1P1D disaggregation with regular SGLang instances. +**Benchmark:** [PD Disaggregation Performance](../../../performance/sglang/sglang-benchmark-results-v1) — compares 1P1D disaggregation with regular SGLang instances. --- @@ -48,7 +48,7 @@ HiCache extends SGLang's RadixAttention with three memory tiers, using Mooncake :maxdepth: 1 :hidden: -../sglang-integration-v1 +pd-disaggregation hicache-quick-start hicache-integration-v1 :::: diff --git a/docs/source/getting_started/examples/sglang-integration-v1.md b/docs/source/deployment/integrations/sglang/pd-disaggregation.md similarity index 99% rename from docs/source/getting_started/examples/sglang-integration-v1.md rename to docs/source/deployment/integrations/sglang/pd-disaggregation.md index 4338001dda..443c976b43 100644 --- a/docs/source/getting_started/examples/sglang-integration-v1.md +++ b/docs/source/deployment/integrations/sglang/pd-disaggregation.md @@ -4,7 +4,7 @@ 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-benchmark-results-v1)). +In benchmarks, PD disaggregation with Mooncake achieves **~30% lower ITL** while maintaining comparable throughput ([details](../../../performance/sglang/sglang-benchmark-results-v1)). ``` +-----------+ Transfer Engine (RDMA) +-----------+ diff --git a/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md b/docs/source/deployment/integrations/vllm/disagg-prefill-decode.md similarity index 97% rename from docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md rename to docs/source/deployment/integrations/vllm/disagg-prefill-decode.md index c8f0ab246d..93453080f1 100644 --- a/docs/source/getting_started/examples/vllm-integration/disagg-prefill-decode.md +++ b/docs/source/deployment/integrations/vllm/disagg-prefill-decode.md @@ -37,7 +37,7 @@ 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](../../build.md). +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 @@ -135,7 +135,7 @@ vllm serve Qwen/Qwen2.5-7B-Instruct \ ### 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. --- @@ -146,7 +146,7 @@ For detailed performance benchmarks and results, see the [vLLM Benchmark](../../ 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). Preview benchmark results are available at [vLLM Benchmark Results V0.2](../../../performance/vllm-benchmark-results-v0.2.md). +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 @@ -157,7 +157,7 @@ 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](../../build.md). +- 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. ``` diff --git a/docs/source/getting_started/examples/vllm-integration/index.md b/docs/source/deployment/integrations/vllm/index.md similarity index 89% rename from docs/source/getting_started/examples/vllm-integration/index.md rename to docs/source/deployment/integrations/vllm/index.md index d6107e502e..3c6d443c4f 100644 --- a/docs/source/getting_started/examples/vllm-integration/index.md +++ b/docs/source/deployment/integrations/vllm/index.md @@ -5,7 +5,7 @@ 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. +- **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 | |----------|-------|-------------| diff --git a/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md b/docs/source/deployment/integrations/vllm/kv-cache-storage.md similarity index 97% rename from docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md rename to docs/source/deployment/integrations/vllm/kv-cache-storage.md index c25427adef..53b5f7a50c 100644 --- a/docs/source/getting_started/examples/vllm-integration/kv-cache-storage.md +++ b/docs/source/deployment/integrations/vllm/kv-cache-storage.md @@ -4,7 +4,7 @@ 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). See [benchmark results](../../../performance/vllm-benchmark-results-v1.md) for details. +Compared to Redis-based backends, MooncakeStore achieves significantly lower TTFT (e.g., **~32% improvement** in mean TTFT for 2P2D tp=2 under RDMA). --- @@ -168,7 +168,7 @@ 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](../../build.md). +- 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. ``` @@ -388,15 +388,6 @@ curl -s http://localhost:8000/v1/completions \ --- -## Performance - -| Scenario | Document | -|----------|----------| -| V1 MooncakeStoreConnector vs Redis | [Benchmark V1](../../../performance/vllm-benchmark-results-v1.md) | -| V0 MooncakeStore vs Redis | [Benchmark V0](../../../performance/vllm-benchmark-results-v0.2.md) | - ---- - ## Troubleshooting - If you encounter connection issues, check that: 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 97% 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 6a567dd732..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,3 +1,7 @@ +--- +orphan: true +--- + # vLLM V0 Disaggregated Serving Demo ```{admonition} Archived @@ -6,11 +10,11 @@ This page has been **consolidated** into the unified [Disaggregated Prefill-Deco ``` ## 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 @@ -21,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 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 98% 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 42e3410601..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,3 +1,7 @@ +--- +orphan: true +--- + # vLLM V0 Disaggregated Serving with MooncakeStore ```{admonition} Archived @@ -18,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 @@ -29,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 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 94% 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 6e7ad79eda..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,3 +1,7 @@ +--- +orphan: true +--- + # vLLM v1 backend Disaggregated Serving with MooncakeConnector ```{admonition} Archived @@ -21,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 @@ -50,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 @@ -123,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/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md b/docs/source/deployment/integrations/vllm/vllm-mooncakestoreconnector.md similarity index 96% rename from docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md rename to docs/source/deployment/integrations/vllm/vllm-mooncakestoreconnector.md index d76f5646a3..8517cb5de8 100644 --- a/docs/source/getting_started/examples/vllm-integration/vllm-mooncakestoreconnector.md +++ b/docs/source/deployment/integrations/vllm/vllm-mooncakestoreconnector.md @@ -1,3 +1,7 @@ +--- +orphan: true +--- + # Guide: vLLM MooncakeStoreConnector ```{admonition} Archived @@ -133,3 +137,8 @@ python examples/disaggregated/disaggregated_serving/mooncake_connector/mooncake_ > ``` > > 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..cb7706f166 --- /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 [Tenant Quota Management](mooncake-store-deployment-guide.md#tenant-quota-management) 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 dcd52722f1..640aec3c2e 100644 --- a/docs/source/deployment/mooncake-store-deployment-guide.md +++ b/docs/source/deployment/mooncake-store-deployment-guide.md @@ -141,11 +141,13 @@ Runs a cluster of master instances coordinated through etcd. If the leader fails # Start each master instance with: mooncake_master \ --enable_ha=true \ - --etcd_endpoints="10.0.0.1:2379;10.0.0.2:2379;10.0.0.3:2379" \ + --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`. The etcd cluster used for HA can be shared with or separate from the Transfer Engine's metadata etcd. +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;...`. @@ -163,7 +165,7 @@ mooncake_master \ --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`. +**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`. --- @@ -199,11 +201,12 @@ mooncake_master \ --offload_on_evict=true \ --promotion_on_hit=true \ --promotion_admission_threshold=2 \ - --root_fs_dir=/mnt/ssd_cache \ --enable_http_metadata_server=true \ --http_metadata_server_port=8080 ``` +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 @@ -234,9 +237,139 @@ mooncake_master \ 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`). + +For snapshot-based standby bootstrap, also configure: + +- `--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`. + +### 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 its state from the `PromotionContext`, populating metadata shards and the segment manager. +5. **Invalid Endpoint Filtering**: During restoration, any replica endpoints that correspond to segments no longer in the registry are automatically filtered out from `GetReplicaList` results. + +### 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 +export MOONCAKE_SNAPSHOT_LOCAL_PATH=/data/mooncake_snapshots +``` + +Start the cluster: + +```bash +# Start Primary +mooncake_master --config_path=primary.yaml + +# Start Standby +mooncake_master --config_path=standby.yaml +``` + +### Resetting a Legacy OpLog Namespace + +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. + +Reset is destructive: + +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`: @@ -249,6 +382,99 @@ curl -s http://:9003/metrics 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_used_bytes{tenant_id}` +- `mooncake_tenant_quota_reserved_bytes{tenant_id}` +- `mooncake_tenant_quota_committed_count{tenant_id}` +- `mooncake_tenant_quota_metadata_object_count{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 + +Tenant quota admission is disabled by default. Enable strict multi-tenant mode on the master when you want memory writes admitted against connector-managed per-tenant quota: + +```bash +mooncake_master \ + --enable_multi_tenants=true \ + --tenant_quota_connector_type=file \ + --tenant_quota_connector_uri=/etc/mooncake/tenant_quotas.yaml +``` + +You can also store the same YAML policy in etcd when Mooncake Store is built with `STORE_USE_ETCD=ON`: + +```bash +mooncake_master \ + --enable_multi_tenants=true \ + --cluster_id=mooncake_cluster \ + --tenant_quota_connector_type=etcd \ + --tenant_quota_connector_uri=127.0.0.1:2379 +``` + +The etcd connector stores the policy at `mooncake-store//tenant_quota_policy`. If the key does not exist, the master starts with an empty policy so the first tenant policy can be created through the admin API. It shares the process-wide store etcd client used by HA/oplog, so if HA or oplog also uses etcd, `tenant_quota_connector_uri` must match those etcd endpoints. The policy must use schema version `1`; tenant names must be non-empty, unique, must not start with `_`, and must not contain NUL or control characters; quotas must be positive integers with optional `B`, `KB`, `MB`, `GB`, or `TB` units: + +```yaml +version: 1 + +tenants: + - name: tenant-a + quota: 200GB + + - name: tenant-b + quota: 500GB +``` + +When strict multi-tenant mode is enabled, write requests must include a registered tenant. The `default` tenant is not special unless it is explicitly registered in the connector policy. + +The same HTTP port used for metrics exposes the tenant quota admin API: + +```bash +# List tenant quota snapshots +curl -s http://:9003/api/v1/tenant_quotas + +# Query one tenant +curl -s "http://:9003/api/v1/tenant_quotas?tenant_id=tenant-a" + +# Upsert an explicit policy. Explicit tenant policies must be positive. +curl -s -X PUT "http://:9003/api/v1/tenant_quotas?tenant_id=tenant-a" \ + -H 'Content-Type: application/json' \ + -d '{"requested_quota_bytes":2147483648}' + +# Delete an explicit policy. The tenant must not own objects or quota usage. +curl -s -X DELETE "http://:9003/api/v1/tenant_quotas?tenant_id=tenant-a" +``` + +Each tenant quota snapshot returns: + +```json +{ + "success": true, + "data": { + "tenant_id": "tenant-a", + "requested_quota_bytes": 2147483648, + "effective_quota_bytes": 2147483648, + "used_bytes": 0, + "reserved_bytes": 0, + "committed_count": 0, + "metadata_object_count": 0, + "over_quota": false, + "has_explicit_policy": true + } +} +``` + +In HA mode, quota admin requests are served only by the active master service. Standby, candidate, or inactive services return HTTP 503. If strict multi-tenant mode is disabled, the quota admin API returns HTTP 409 with `UNAVAILABLE_IN_CURRENT_MODE`. Deleting a non-empty tenant returns HTTP 409 with `TENANT_NOT_EMPTY`. + --- ## Quick Tips @@ -256,8 +482,8 @@ curl -s http://:9003/metrics/summary - 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-offload). -- For NVMe-oF SSD pool configuration see the [NVMe-oF SSD Pool Deployment Guide](nvmf-ssd-deployment-guide) +- 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 experimental 3FS (USRBIO) integration as a persistent storage backend, see the [3FS USRBIO Plugin guide](../getting_started/plugin-usage/3FS-USRBIO-Plugin). - For detailed monitoring and observation see [Observability](../getting_started/observability) @@ -265,8 +491,8 @@ curl -s http://:9003/metrics/summary :maxdepth: 1 :hidden: -ssd-offload -NvMe-Of SSD Pool +KV Cache Sharing and Isolation +SSD Storage HF3FS Plugin (Experimental)<../getting_started/plugin-usage/3FS-USRBIO-Plugin> ../getting_started/observability ::: @@ -299,6 +525,45 @@ glog's standard flags (`--log_dir`, `--max_log_size`, `--logtostderr`, ...) cont | `--enable_metric_reporting` | `true` | Periodically log master metrics | | `--metrics_port` | `9003` | HTTP port for `/metrics` endpoints | +### KV Cache Event Publisher + +The master can publish KV cache lifecycle events over a ZMQ PUB socket for +cache-aware indexers such as Mooncake Conductor. This feature is compiled out +by default. Install `libzmq3-dev` and configure Mooncake Store with +`-DENABLE_KV_EVENTS=ON` before enabling it at runtime. + +Both `--kv_events_bind_endpoint` and `--kv_events_backend_id` are required when +the publisher is enabled. If either value is empty, or the ZMQ socket cannot +bind, the master logs an error and continues with event publishing disabled. + +```bash +mooncake_master \ + --enable_kv_events=true \ + --kv_events_bind_endpoint=tcp://0.0.0.0:5557 \ + --kv_events_backend_id=store-node-1 +``` + +Register an address reachable by the indexer, rather than the wildcard bind +address, through the indexer's `POST /register` endpoint. For the event format, +registration fields, and object-key behavior, see the {ref}`Mooncake Store +master publisher ` reference. + +| Flag | Default | Description | +|------|---------|-------------| +| `--enable_kv_events` | `false` | Enable the ZMQ KV cache event publisher; requires a build with `ENABLE_KV_EVENTS=ON` | +| `--kv_events_bind_endpoint` | empty | ZMQ PUB bind endpoint, for example `tcp://0.0.0.0:5557`; required when enabled | +| `--kv_events_backend_id` | empty | Cache-owner identity emitted as `backend_id`; required when enabled | +| `--kv_events_emit_legacy_compat` | `true` | Include vLLM/SGLang-compatible aliases such as `type` and `block_hashes` | +| `--kv_events_emit_object_key` | `true` | Include the Mooncake `object_key`; unparsable sequence hashes are still published when this is enabled | +| `--kv_events_queue_capacity` | `65536` | Maximum pending events; the publisher drops the oldest event when the queue is full. Set to `0` for an unbounded queue | + +The legacy flags `--kv_events_model_name`, `--kv_events_tenant_id`, +`--kv_events_additional_salt`, `--kv_events_lora_name`, +`--kv_events_block_size`, and `--kv_events_dp_rank` are retained for config +compatibility but are not emitted in event payloads. Supply model, block-size, +hash-namespace, LoRA, and data-parallel metadata when registering the publisher +with the indexer; each event carries its object's tenant ID. + ### HTTP Metadata Server (Embedded) | Flag | Default | Description | @@ -306,6 +571,56 @@ glog's standard flags (`--log_dir`, `--max_log_size`, `--logtostderr`, ...) cont | `--enable_http_metadata_server` | `false` | Enable embedded HTTP metadata server | | `--http_metadata_server_host` | `0.0.0.0` | Metadata bind host | | `--http_metadata_server_port` | `8080` | Metadata TCP port | +| `--enable_metadata_cleanup_on_timeout` | `false` | Delete a client's stale HTTP metadata (`mooncake/[/]ram/` and `mooncake/[/]rpc_meta/`) when its heartbeat times out (see below) | + +### Stale Metadata Cleanup on Client Timeout + +When a client crashes or is force-killed (`kill -9`, OOM, node failure), it cannot +run its normal cleanup, leaving stale entries on the HTTP metadata server +(`mooncake/[/]ram/` and `mooncake/[/]rpc_meta/`). +The HTTP metadata server has no heartbeat of its own, so these entries linger and +can mislead nodes that later connect or restart with different RDMA parameters. + +With `--enable_metadata_cleanup_on_timeout=true`, the Master Service reuses its +existing client-heartbeat monitor: when a client's `--client_ttl` expires, in +addition to unmounting the segment it also removes that client's `ram/` and +`rpc_meta/` keys from the HTTP metadata server. It supports both deployment +topologies: + +- **Co-located** (`--enable_http_metadata_server=true`): the master removes the + keys via a direct in-process call (no network overhead). +- **Separately deployed** HTTP metadata server: the master derives the metadata + server address from the cluster's existing configuration and removes the keys + via HTTP `DELETE`. The address is read, in priority order, from: + 1. the `MOONCAKE_TE_META_DATA_SERVER` environment variable (the same Transfer + Engine metadata connection string the clients use, e.g. + `http://host:8080/metadata`), then + 2. the `metadata_server` field of the JSON file pointed to by + `MOONCAKE_CONFIG_PATH`. + +Notes: +- Only `http(s)` metadata servers are supported; `etcd`/`redis`/`P2PHANDSHAKE` + backends are not cleaned up (a warning is logged and cleanup stays disabled). +- The feature is opt-in and best-effort: if no co-located server is enabled and + no HTTP metadata address can be derived, the master logs a warning and + disables cleanup. Remote `DELETE` failures are logged but never block the + client-monitor thread or the main process. +- Respects `MC_METADATA_CLUSTER_ID` for custom key prefixes (matching the + Transfer Engine). + +```bash +# Co-located metadata server +mooncake_master \ + --enable_http_metadata_server=true \ + --enable_metadata_cleanup_on_timeout=true \ + --client_ttl=10 + +# Separately-deployed HTTP metadata server (address derived from the env var) +export MOONCAKE_TE_META_DATA_SERVER=http://metadata-host:8080/metadata +mooncake_master \ + --enable_metadata_cleanup_on_timeout=true \ + --client_ttl=10 +``` ### Memory Allocator @@ -317,7 +632,7 @@ glog's standard flags (`--log_dir`, `--max_log_size`, `--logtostderr`, ...) cont | Flag | Default | Description | |------|---------|-------------| -| `--allocation_strategy` | `random` | `random` (pure random, fastest), `free_ratio_first` (best load balance), or `cxl` (prefer CXL memory) | +| `--allocation_strategy` | `random` | Allocation strategy: `random` (pure random, fastest), `free_ratio_first` (best memory load balance), `ssd_free_ratio_first` (SSD-aware free-ratio-first), `cxl` (prefer CXL memory), or `local_first` (prefer local host memory segments before ordered remote fallback) | ### PutStart Timeouts @@ -330,13 +645,21 @@ glog's standard flags (`--log_dir`, `--max_log_size`, `--logtostderr`, ...) cont | Flag | Default | Description | |------|---------|-------------| -| `--default_kv_lease_ttl` | `5000` ms | Lease TTL for KV objects. Supports `5000ms`, `5s`, `30m`, `1h` | +| `--default_kv_lease_ttl` | `10000` ms | Lease TTL for KV objects. Supports `5000ms`, `5s`, `30m`, `1h` | | `--default_kv_soft_pin_ttl` | `1800000` ms | Soft pin TTL (30 min) | | `--allow_evict_soft_pinned_objects` | `true` | Allow evicting soft-pinned objects | | `--eviction_ratio` | `0.05` | Fraction evicted at high watermark | -| `--eviction_high_watermark_ratio` | `0.95` | Usage ratio triggering eviction | +| `--eviction_high_watermark_ratio` | `0.90` | Usage ratio triggering eviction | | `--client_ttl` | `10` s | Seconds before a silent client is considered disconnected | +### Tenant Quota + +| Flag | Default | Description | +|------|---------|-------------| +| `--enable_multi_tenants` | `false` | Enable strict tenant registration and per-tenant memory quota admission | +| `--tenant_quota_connector_type` | `file` | Tenant quota policy connector type: `file` or `etcd` when built with `STORE_USE_ETCD=ON` | +| `--tenant_quota_connector_uri` | empty | Connector URI; for `file`, the writable YAML policy path; for `etcd`, the endpoints string | + ### High Availability **Master Node High Availability** @@ -345,8 +668,12 @@ glog's standard flags (`--log_dir`, `--max_log_size`, `--logtostderr`, ...) cont | `--enable_ha` | `false` | Enable HA mode | | `--ha_backend_type` | `etcd` | HA backend: `etcd`, `redis`, or `k8s` | | `--ha_backend_connstring` | empty | HA backend connection string | -| `--etcd_endpoints` | empty | etcd endpoints, semicolon separated (when `--ha_backend_type=etcd`) | +| `--etcd_endpoints` | empty | Backward-compatible etcd HA endpoints, used only for `ha_backend_type=etcd` when `--ha_backend_connstring` is empty | | `--cluster_id` | `mooncake_cluster` | Cluster ID for HA persistence | +| `--enable_oplog` | `false` | Enable the primary OpLog writer and standby reader; currently requires `enable_ha=true` and `ha_backend_type=etcd` | +| `--oplog_poll_interval_ms` | `1000` | Base polling and retry delay for the batch standby, in milliseconds | +| `--oplog_batch_max_entries` | `1024` | Maximum number of entries admitted to an ordered batch | +| `--batch_oplog_retry_timeout_sec` | `180` | Maximum consecutive retryable batch-standby failure window in seconds | ```{caution} Metadata Snapshot And Restore is experimental feature. @@ -392,6 +719,8 @@ Flags for controlling data movement between DRAM and SSD. | `--enable_offload` | `false` | Enable offload from DRAM to SSD | | `--offload_on_evict` | `false` | Defer offload to eviction time rather than at `Put` | | `--offload_force_evict` | `false` | Force-evict objects exceeding capacity without offload | +| `--offloading_queue_limit` | `50000` | Max number of objects allowed in the offloading queue per local disk segment. Increase to allow more objects to be offloaded to SSD before force-eviction kicks in | +| `--offload_cap_ratio` | `0.5` | Per-cycle offload cap as a fraction of `offloading_queue_limit` (range `[0.0, 1.0]`). Controls how many objects can be queued for offload in a single eviction cycle before falling back to force-evict | | `--promotion_on_hit` | `false` | Promote SSD-resident keys to DRAM on read hit | | `--promotion_admission_threshold` | `2` | Min CountMinSketch count to allow promotion (`1` = disable gating) | | `--promotion_max_per_heartbeat` | `1` | Max promotion tasks handed to a single client per heartbeat. Each task is a synchronous SSD-read + RDMA-write on the client; serializing them avoids blocking past the client-liveness window | @@ -401,6 +730,10 @@ Flags for controlling data movement between DRAM and SSD. Start with `--enable_offload=true` for eager asynchronous SSD persistence after `Put` completion. Add `--offload_on_evict=true` when you want SSD writes to happen only when memory pressure selects an object for eviction. Add `--promotion_on_hit=true` to allow hot SSD-only data to be promoted back to DRAM, and tune `--promotion_admission_threshold` to control how many observed reads are required before promotion is queued. +For SSD offload, configure the disk path on each real client with `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`; the master tracks these objects as `LOCAL_DISK` replicas. Do not use the legacy `--root_fs_dir` parameter with `--enable_offload=true`. + +When `--offload_on_evict=true` is active, each `BatchEvict` cycle can queue at most `offloading_queue_limit * offload_cap_ratio` objects for SSD offload (default: `50000 * 0.5 = 25000`); objects exceeding this cap fall back to force-evict (discard) if `--offload_force_evict=true`, otherwise they remain in memory. For SSD-heavy workloads where NVMe bandwidth is underutilized while the KV-cache hit rate suffers, raise both `--offloading_queue_limit` and `--offload_cap_ratio` so more objects per cycle are actually persisted to SSD instead of discarded. Example: `--offloading_queue_limit=500000 --offload_cap_ratio=0.8` yields a per-cycle cap of `400000` (vs the default `25000`). + ### CXL Memory | Flag | Default | Description | @@ -415,21 +748,23 @@ When `--allocation_strategy=cxl` is set alongside `--enable_cxl=true`, the maste | Flag | Default | Description | |------|---------|-------------| -| `--root_fs_dir` | empty | DFS mount directory for multi-layer storage backend | +| `--root_fs_dir` | empty | Legacy DFS persistence directory; do not use with SSD offload | | `--global_file_segment_size` | `INT64_MAX` (unlimited) | Max available space for DFS segments; default does not cap DFS usage | +`--root_fs_dir` is a legacy persistence parameter and is expected to be replaced as the distributed filesystem path is refactored. For SSD offload, configure `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` on each real client instead. + ### NoF (NVMe-oF SSD Pool) ```{caution} NVMe-oF SSD Pool (NoF) is an experimental feature. ``` -Master-side flags for the NVMe-oF SSD pool. They control eviction within the NoF SSD tier and the heartbeat used to detect and unmount unresponsive NoF segments. For the client-side NoF I/O tuning (`MC_NOF_*`), see the [NVMe-oF SSD Pool Deployment Guide](nvmf-ssd-deployment-guide.md). +Master-side flags for the NVMe-oF SSD pool. They control eviction within the NoF SSD tier and the heartbeat used to detect and unmount unresponsive NoF segments. For the client-side NoF I/O tuning (`MC_NOF_*`), see the [NVMe-oF SSD Pool Deployment Guide](ssd/nvmf-ssd-deployment-guide.md). | Flag | Default | Description | |------|---------|-------------| | `--nof_eviction_ratio` | `0.05` | Fraction of objects evicted when NoF SSD space is full | -| `--nof_eviction_high_watermark_ratio` | `0.95` | Usage ratio that triggers eviction in the NoF SSD tier | +| `--nof_eviction_high_watermark_ratio` | `0.90` | Usage ratio that triggers eviction in the NoF SSD tier | | `--nof_heartbeat_interval_sec` | `10` | How often the master probes each mounted NoF segment | | `--nof_heartbeat_probe_timeout_ms` | `1000` | Timeout for a single NoF heartbeat probe | | `--nof_heartbeat_failures_threshold` | `3` | Consecutive NoF heartbeat failures before a segment is unmounted | @@ -447,6 +782,20 @@ rpc_interface: "eth0" rpc_port: 50051 ``` +### Local-first Allocation + +Mooncake can prefer memory segments on the writer's host before falling back to remote hosts. This is useful when colocating inference workers and store segments, because a store node failure only invalidates the KV cache written to that host instead of spreading one request's cache across the whole cluster. + +This feature is disabled by default. Enable it on the master by selecting the local-first allocation strategy: + +```yaml +allocation_strategy: "local_first" +``` + +When enabled, the master applies local-first allocation only for memory replicas with `replica_num == 1`. Explicit `preferred_segment` or `preferred_segments` are tried first; if they are unavailable or full, Mooncake falls back through active hosts in cyclic lexicographic host-id order, starting from the writer host when it has active segments, or otherwise from the next greater active host id. Within the same host, segment names are sorted and rotated by key hash so multiple segments on one host do not always receive the first allocation attempt. + +The client derives the host id from `local_hostname` by removing the port. For example, `host-a:50051` and `host-a:50052` map to the same host id, `host-a`. For local-first allocation to work correctly, all writer and store processes on the same physical or logical host must use the same stable, globally unique host part in `local_hostname`. In deployments with multiple NIC IPs, hostname aliases, or container/pod networking, choose one canonical host name or IP and use it consistently across processes on that host. Empty, loopback, and wildcard values such as `localhost`, `127.0.0.1`, `0.0.0.0`, `::1`, and `::` are treated as unknown and do not trigger automatic local-first placement for that client. + --- (reference-client-configuration-tuning)= @@ -471,16 +820,18 @@ Arguments of `MooncakeDistributedStore.setup(...)`: | `metadata_server` | str | required | `P2PHANDSHAKE` / `http://…:8080/metadata` / etcd address | | `global_segment_size` | int (bytes) | required | DRAM contributed to the cluster (the sample uses 3.2 GB) | | `local_buffer_size` | int (bytes) | required | Transfer Engine buffer | -| `protocol` | str | required | `tcp` / `rdma` / `cxl` / `ascend` | +| `protocol` | str | required | `tcp` / `rdma` / `efa` / `cxl` / `ascend` | | `rdma_devices` | str | required | RDMA NIC(s), comma-separated (pass `""` for non-RDMA). **Keyword is `rdma_devices`, not `device_name`** | | `master_server_addr` | str | required | Master `host:port`. **Keyword is `master_server_addr`, not `master_server_address`** | | `engine` | TransferEngine | `None` | *(advanced)* Reuse an existing Transfer Engine instance instead of creating one | | `enable_ssd_offload` | bool | `false` | *(advanced)* Enable client-side SSD offload | | `ssd_offload_path` | str | empty | *(advanced)* SSD offload directory | | `tenant_id` | str | `default` | *(advanced)* Tenant identifier | +| `enable_client_http_server` | bool | `false` | Enable the client-side HTTP `/health`, `/metrics`, and `/metrics/summary` endpoints | +| `client_http_port` | int | `9300` | Client-side HTTP endpoint port, used only when `enable_client_http_server=true` | ```{note} -The first seven arguments have **no Python default** — the C++ defaults are not exposed by the pybind binding, so they must all be supplied (a bare `setup(local_hostname, metadata_server)` raises `TypeError`). Only `engine` / `enable_ssd_offload` / `ssd_offload_path` / `tenant_id` are optional. Also, in Method A only the `MC_*` engine variables below have any effect — `MOONCAKE_*` are ignored. +The first seven arguments have **no Python default** — the C++ defaults are not exposed by the pybind binding, so they must all be supplied (a bare `setup(local_hostname, metadata_server)` raises `TypeError`). The later arguments (`engine`, SSD offload fields, `tenant_id`, and client HTTP endpoint fields) are optional. Also, in Method A the `MOONCAKE_*` variables used by `MooncakeConfig` are ignored; low-level runtime variables such as the `MC_*` engine variables below are still read by the C++ client. ``` ### Method B — Service / Integration (`MOONCAKE_*` + CLI) @@ -499,13 +850,16 @@ The store service CLI only accepts `--config`, `-D/--define`, `--port`, and `--m |----------|-------------------------|---------|-------------| | `MOONCAKE_MASTER` | `master_server_addr` | — (required unless `MOONCAKE_CONFIG_PATH`) | Master `host:port` | | `MOONCAKE_TE_META_DATA_SERVER` | `metadata_server` | `P2PHANDSHAKE` | `P2PHANDSHAKE` / `http://…:8080/metadata` / etcd address | -| `MOONCAKE_PROTOCOL` | `protocol` | `tcp` | `tcp` / `rdma` / `cxl` / `ascend` | -| `MOONCAKE_DEVICE` | `rdma_devices` | empty | RDMA device(s), comma-separated; `auto-discovery` supported | +| `MOONCAKE_PROTOCOL` | `protocol` | `tcp` | `tcp` / `rdma` / `efa` / `cxl` / `ascend` | +| `MOONCAKE_DEVICE` | `rdma_devices` | empty | RDMA/EFA device(s), comma-separated; `auto-discovery` supported | | `MOONCAKE_GLOBAL_SEGMENT_SIZE` | `global_segment_size` | `3355443200` (3.125 GiB) | DRAM contributed; accepts byte integer **or** suffixed form like `500gb` | | `MOONCAKE_LOCAL_BUFFER_SIZE` | `local_buffer_size` | `1073741824` (1 GiB) | Transfer Engine buffer; same parsing as above | | `MOONCAKE_LOCAL_HOSTNAME` | `local_hostname` | `localhost` | | | `MOONCAKE_OFFLOAD_ENABLED` | `enable_ssd_offload` | `false` | Client-side SSD offload | | `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` | `ssd_offload_path` | empty | Offload directory | +| `MOONCAKE_TENANT_ID` | `tenant_id` | `default` | Tenant identifier | +| `MOONCAKE_ENABLE_CLIENT_HTTP_SERVER` | `enable_client_http_server` | `false` | Enable client-side `/health`, `/metrics`, and `/metrics/summary` endpoints | +| `MOONCAKE_CLIENT_HTTP_PORT` | `client_http_port` | `9300` | Client-side HTTP endpoint port | | `MOONCAKE_CONFIG_PATH` | — | unset | Path to a JSON config file (takes precedence over the variables above) | ```{note} @@ -536,12 +890,16 @@ Or via a JSON config file. The service also exposes a lightweight HTTP API (on ` "local_buffer_size": 268435456, "protocol": "tcp", "device_name": "", - "master_server_address": "127.0.0.1:50051" + "master_server_address": "127.0.0.1:50051", + "tenant_id": "default", + "enable_client_http_server": false, + "client_http_port": 9300 } ``` ```bash python -m mooncake.mooncake_store_service --config= --port=8081 +python -m mooncake.mooncake_store_service --config= -Dtenant_id=tenant-a ``` ### Method C — Resource-owning Real Client (`mooncake_client`) @@ -552,21 +910,55 @@ Run the `mooncake_client` binary as a standalone RPC process that owns storage r mooncake_client \ --global_segment_size="4GB" \ --master_server_address="127.0.0.1:50051" \ - --metadata_server="http://127.0.0.1:8080/metadata" + --metadata_server="http://127.0.0.1:8080/metadata" \ + --tenant_id="default" ``` | Flag | Default | Description | |------|---------|-------------| -| `--host` | `0.0.0.0` | Client service bind host | -| `--port` | `50052` | Client service listen port | +| `--host` | `0.0.0.0` | Client service bind host. Accepts `ip:port` to specify the data plane port for TransferEngine | +| `--port` | `50052` | Client RPC listen port (dummy↔real client control plane) | | `--global_segment_size` | `4 GB` | Global segment size contributed by the client | | `--master_server_address` | `127.0.0.1:50051` | Master service address | | `--metadata_server` | `http://127.0.0.1:8080/metadata` | Transfer Engine metadata service | | `--protocol` | `tcp` | Transfer protocol | | `--device_names` | empty | Transfer device name(s), comma-separated | | `--threads` | `1` | Client worker thread count | +| `--tenant_id` | `default` | Tenant identifier | | `--enable_offload` | `false` | Enable client-side SSD offload | | `--start_offload_rpc_server` | `true` | Start the offload RPC server for dummy clients | +| `--enable_http_server` | `false` | Enable client-side `/health`, `/metrics`, and `/metrics/summary` endpoints | +| `--http_port` | `9300` | Client-side HTTP endpoint port | + +### Client HTTP Health and Metrics Endpoint + +Each real client can expose its own lightweight HTTP endpoint independently of the master admin HTTP server and the Python store REST API. This endpoint is disabled by default for programmatic clients and `mooncake_store_service`; enable it explicitly when you want to scrape client-local metrics: + +```python +store.setup( + local_hostname, + metadata_server, + global_segment_size, + local_buffer_size, + protocol, + rdma_devices, + master_server_addr, + enable_client_http_server=True, + client_http_port=9300, +) +``` + +For `mooncake_store_service`, use `MOONCAKE_ENABLE_CLIENT_HTTP_SERVER=true` and optionally `MOONCAKE_CLIENT_HTTP_PORT=`, or set the same fields in the JSON config. For `mooncake_client`, use `--enable_http_server=true --http_port=`. + +| Endpoint | Description | +|----------|-------------| +| `GET /health` | Client health check | +| `GET /metrics` | Prometheus-format client metrics | +| `GET /metrics/summary` | Human-readable client metrics summary | + +```{note} +`MC_STORE_CLIENT_METRIC` controls whether client metrics are collected. If the client HTTP server is enabled but `MC_STORE_CLIENT_METRIC=0`, `/metrics` and `/metrics/summary` return HTTP 503 with `metrics not available`. +``` ### Engine Runtime Tuning (`MC_*`) @@ -579,9 +971,17 @@ The following `MC_*` variables are read directly by the engine/client at runtime | `MC_RPC_PROTOCOL` | `tcp` | RPC transport protocol between master and clients: `tcp` or `rdma` | | `MC_RPC_TIMEOUT_MS` | `30000` | Per-request deadline (ms) for all client→master RPCs. Applies uniformly to every RPC method. A negative value disables the timeout. On expiry the call returns `RPC_TIMEOUT` | | `MC_RPC_CONNECT_TIMEOUT_MS` | `30000` | Connection-establishment timeout (ms) for the master RPC client | +| `MC_RPC_CLIENT_IO_THREADS` | `min(16, online CPU count)`, minimum `1` | Fallback number of threads and `io_context` instances for each component's RPC client I/O pool. A positive integer overrides the default; invalid values and `0` use the default | +| `MC_STORE_RPC_CLIENT_IO_THREADS` | `MC_RPC_CLIENT_IO_THREADS` | Store/Master client RPC I/O pool size. This pool is isolated from Transfer Engine traffic. Invalid values and `0` use the fallback | +| `MC_TE_RPC_CLIENT_IO_THREADS` | `MC_RPC_CLIENT_IO_THREADS` | Transfer Engine and TENT client RPC I/O pool size. This pool is isolated from Store/Master traffic. Invalid values and `0` use the fallback | | `MC_USE_TENT` / `MC_USE_TEV1` | unset | Set to any value to enable the TENT (next-gen) transfer engine | | `MC_STORE_CLUSTER_ID` | unset | Cluster ID label attached to client metrics | +RPC client I/O pool settings are read and resolved once when the process-wide +`Environ` singleton is initialized. Changes therefore require a process +restart. When Store and Transfer Engine run in the same process, each component +owns the configured number of threads and `io_context` instances. + #### Topology Discovery | Variable | Default | Description | @@ -618,6 +1018,14 @@ Local hot cache provides a DRAM read cache on top of SSD-resident objects for fa | `MC_STORE_LOCAL_HOT_CACHE_USE_SHM` | unset | Set `1` to use memfd-backed shared memory | | `MC_STORE_LOCAL_HOT_ADMISSION_THRESHOLD` | unset | Minimum CountMinSketch count before a key is admitted to hot cache | +#### Object-Level Checksum Diagnostics + +Set `MOONCAKE_STORE_CHECKSUM=1` on a Mooncake Store client process before the client is created to enable object-level CRC-64 checks. The client computes the checksum before `put`/`upsert`, stores it in master metadata, and verifies the logical `object_size` bytes returned by a full-object `get`. For complete diagnostic coverage, enable the switch on every writer and reader client. A client with the switch disabled does not generate or verify checksums; an enabled reader skips verification for objects whose metadata has no checksum. + +This switch is intended for corruption diagnosis, not normal production use. It adds a full data scan to writes and reads, performs device-to-host staging for GPU buffers, and disables the local hot cache. Range reads, including `get_into_ranges`, are intentionally not verified. + +Do not run binaries from before and after checksum support was introduced in the same deployment; Mooncake Store clients, the primary master, and the standby master must all use a checksum-capable version. Checksum-capable masters persist checksum metadata in new snapshots and can load snapshots created by older versions; objects restored from an older snapshot have no checksum and are read without verification. Snapshots containing checksum metadata cannot be restored by binaries that predate checksum support, so rolling back requires an older compatible snapshot or a fresh deployment. + #### Local Memory Optimization | Variable | Default | Description | @@ -635,6 +1043,21 @@ Local hot cache provides a DRAM read cache on top of SSD-resident objects for fa | `MC_MMAP_ARENA_POOL_SIZE` | unset | Pre-allocated arena pool size (e.g., `8gb`). Explicitly set to enable the arena | | `MC_DISABLE_MMAP_ARENA` | unset | Disable arena, fall back to per-call `mmap()`. Accepts `1`/`true`/`yes`/`on` (or `0`/`false`/`no`/`off`) | +RDMA Store segments backed by HugeTLB are populated in parallel immediately +before transfer-engine registration. No additional population-mode setting is +required: + +```bash +export MC_STORE_USE_HUGEPAGE=1 +export MC_STORE_HUGEPAGE_SIZE=2MB +``` + +For direct mappings, workers divide the mapping into page ranges. For +NUMA-segmented mappings, each worker is scheduled on the NUMA node associated +with its `mbind()` region before touching pages. The mmap arena retains its +eager `MAP_POPULATE` behavior for DMA safety; set `MC_DISABLE_MMAP_ARENA=1` if +the deferred direct-mmap path is desired while the arena is otherwise enabled. + #### yalantinglibs Log Level ```bash diff --git a/docs/source/deployment/ssd/index.md b/docs/source/deployment/ssd/index.md new file mode 100644 index 0000000000..110445391a --- /dev/null +++ b/docs/source/deployment/ssd/index.md @@ -0,0 +1,17 @@ +# SSD Storage + +Mooncake Store supports both node-local SSD offload and shared NVMe-over-Fabrics +storage pools. Choose the guide that matches the storage tier in your deployment. + +| Storage option | Description | +|----------------|-------------| +| [SSD Offload](ssd-offload) | Configure local SSD offload, eviction, and I/O behavior. | +| [NVMe-oF SSD Pool](nvmf-ssd-deployment-guide) | Configure a shared NVMe-over-Fabrics storage tier. | + +:::{toctree} +:maxdepth: 1 +:hidden: + +ssd-offload +nvmf-ssd-deployment-guide +::: diff --git a/docs/source/deployment/nvmf-ssd-deployment-guide.md b/docs/source/deployment/ssd/nvmf-ssd-deployment-guide.md similarity index 94% rename from docs/source/deployment/nvmf-ssd-deployment-guide.md rename to docs/source/deployment/ssd/nvmf-ssd-deployment-guide.md index eb8558383b..5b27e077b1 100644 --- a/docs/source/deployment/nvmf-ssd-deployment-guide.md +++ b/docs/source/deployment/ssd/nvmf-ssd-deployment-guide.md @@ -16,7 +16,7 @@ Mooncake Store. ## 1. Build Mooncake with NoF Support Follow the "Build with NVMe-oF SSD Pool" section in the -[Build Guide](../getting_started/build.md) to install SPDK dependencies and +[Build Guide](../../getting_started/build.md) to install SPDK dependencies and build Mooncake with `-DUSE_NOF=ON`. ## 2. Deploy Mooncake Services @@ -233,7 +233,7 @@ Enter the SPDK directory on the target node and run the following commands. ### 6.2 Use NoF with vLLM + LMCache For the general VLLM + LMCache + Mooncake deployment flow, see -[vLLM V1 Disaggregated Serving with Mooncake Store and LMCache](../getting_started/examples/vllm-integration/vllmv1-lmcache-integration.md). +[vLLM V1 Disaggregated Serving with Mooncake Store and LMCache](../integrations/lmcache/vllmv1-lmcache-integration.md). After the NVMe-oF SSD pool is registered with Mooncake, add the NoF-specific Mooncake configuration below. @@ -254,7 +254,6 @@ remote_url: "mooncakestore://192.168.65.81:50051/" remote_serde: "naive" local_cpu: True max_local_cpu_size: 8 -enable_mooncake_nof_pool: True extra_config: local_hostname: "localhost" @@ -268,8 +267,12 @@ extra_config: **Notes**: -- `enable_mooncake_nof_pool=True` enables writing KV cache objects to the - registered NoF pool. +- Writing KV cache objects to the registered NoF pool is controlled per put + by `nof_replica_num` in the store client's `ReplicateConfig` (exposed in + the Python binding), and the master must be built with `USE_NOF`. The + LMCache Mooncake connector does not currently parse or forward such a + setting, so enabling the NoF pool for LMCache writes needs a connector + change or a direct store client. - `global_segment_size: 0` means the inference process does not contribute a memory segment to the Mooncake cluster. - Keep `local_buffer_size` non-zero because the client still needs local diff --git a/docs/source/deployment/ssd-offload.md b/docs/source/deployment/ssd/ssd-offload.md similarity index 70% rename from docs/source/deployment/ssd-offload.md rename to docs/source/deployment/ssd/ssd-offload.md index 439f5f29b5..34d2dd28a9 100644 --- a/docs/source/deployment/ssd-offload.md +++ b/docs/source/deployment/ssd/ssd-offload.md @@ -2,9 +2,9 @@ ## Overview -Mooncake Store supports offloading KV cache objects from distributed memory to local SSD. When memory pressure is high, the master instructs clients to persist selected objects to disk. On a cache miss, the client automatically falls back to reading from SSD. +Mooncake Store supports offloading KV cache objects from distributed memory to a local filesystem path, typically backed by local SSDs. When memory pressure is high, the master instructs clients to persist selected objects to disk. On a cache miss, the client automatically falls back to reading from the local filesystem-backed offload path. -For measured TTFT and throughput impact in multi-turn workloads, see [Mooncake SSD Offload Benchmark](../performance/ssd-offload-benchmark-results.md). +For measured TTFT and throughput impact in multi-turn workloads, see [Mooncake SSD Offload Benchmark](../../performance/mooncake/ssd-offload-benchmark-results.md). SSD offload requires the **Real Client** and supports two deployment modes: @@ -13,6 +13,8 @@ SSD offload requires the **Real Client** and supports two deployment modes: In both modes, all SSD reads and writes happen within the Real Client (embedded or standalone). +SSD offload does not use the master's `--root_fs_dir` option. Configure the local disk path on each Real Client with `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`; the master tracks offloaded objects as `LOCAL_DISK` replicas. `--root_fs_dir` is a legacy parameter from an older persistence path and may cause issues when used with `--enable_offload=true`. + ## Startup Steps ### Step 1: Create the SSD storage directory @@ -25,8 +27,7 @@ mkdir -p /nvme/mooncake_offload ```bash mooncake_master \ - --rpc_port=50051 \ - --enable_offload=true + --rpc_port=50051 ``` ### Step 3A (Mode A): Start the application with embedded Real Client @@ -92,16 +93,37 @@ store.setup_dummy( |------|---------|-------------| | `--metadata_server` | `http://127.0.0.1:8080/metadata` | Metadata server connection string | | `--master_server_address` | `127.0.0.1:50051` | Master address | -| `--host` | `0.0.0.0` | This machine's externally reachable IP | +| `--host` | `0.0.0.0` | This machine's externally reachable IP. Accepts `ip:port` to specify the data plane port for TransferEngine | | `--port` | `50052` | Real client RPC listening port | | `--device_names` | ` ` | NIC name(s), e.g. `eth0` or `mlx5_0` | | `--protocol` | `tcp` | Transport protocol: `tcp` or `rdma` | | `--global_segment_size` | `4 GB` | Memory pool size allocated for this node | | `--enable_offload` | `false` | **Must be set to `true` to enable SSD offload** | +| `--start_offload_rpc_server` | `true` | Start the offload RPC server used by DummyClients for SSD reads. Effective only when `--enable_offload=true`; disable only for write-only owners | | `--threads` | `1` | Number of RPC server threads | --- +## Master Offload Parameters + +These flags control when the master asks clients to persist objects to SSD, whether SSD-only objects can be promoted back to DRAM, and how much offload work can be queued during eviction. + +| Flag | Default | Description | +|------|---------|-------------| +| `--enable_offload` | `false` | Enables the master-side SSD offload control plane. Set this on the master and on every real client that owns SSD storage | +| `--offload_on_evict` | `false` | Defer SSD persistence until memory eviction selects an object. When `false`, successful `Put` completion queues eager SSD persistence | +| `--offload_force_evict` | `false` | If offload-on-evict exceeds the per-cycle offload cap, force-evict excess objects instead of leaving them in DRAM | +| `--offloading_queue_limit` | `50000` | Maximum pending offload objects per local disk segment. Must be greater than `0` and at most `100000000` | +| `--offload_cap_ratio` | `0.5` | Per-eviction-cycle cap as a fraction of `offloading_queue_limit`; range `[0.0, 1.0]`. Default cap is `25000` objects per cycle | +| `--promotion_on_hit` | `false` | Promote SSD-resident objects back to DRAM after read hits | +| `--promotion_admission_threshold` | `2` | Minimum CountMinSketch count before promotion is admitted. Set `1` to disable second-touch gating | +| `--promotion_max_per_heartbeat` | `1` | Maximum promotion tasks returned to one client per heartbeat. Keep conservative for large objects because each task does SSD read plus memory write | +| `--promotion_queue_limit` | `50000` | Maximum in-flight promotion tasks tracked by the master | + +Start with `--enable_offload=true` for eager SSD persistence. Add `--offload_on_evict=true` when SSD writes should happen only under memory pressure. For SSD-heavy workloads where offload-on-evict is dropping too many objects, raise both `--offloading_queue_limit` and `--offload_cap_ratio`; for example, `--offloading_queue_limit=500000 --offload_cap_ratio=0.8` allows up to `400000` objects to be queued in one eviction cycle. + +--- + ## SSD Offload Configuration ### Core settings @@ -118,6 +140,11 @@ store.setup_dummy( | `MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_INTERVAL_SECONDS` | `1` | Interval for reclaiming expired offload buffers; defaults to the heartbeat interval in the current implementation | | `MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS` | `5000` | Lease time for buffers returned by `batch_get_offload_object` before GC reclaims them | | `MOONCAKE_OFFLOAD_USE_URING` | `false` | Enable io_uring for async file I/O | +| `MOONCAKE_OFFLOAD_ENABLE_DISK_WATERMARK_EVICTION` | `true` | Enable proactive local-disk eviction from the FileStorage heartbeat | +| `MOONCAKE_OFFLOAD_DISK_EVICTION_HIGH_WATERMARK_RATIO` | `0.90` | Trigger proactive disk eviction when backend usage exceeds this ratio of its quota | +| `MOONCAKE_OFFLOAD_DISK_EVICTION_LOW_WATERMARK_RATIO` | `0.80` | Target backend usage ratio for proactive disk eviction | + +The `MOONCAKE_OFFLOAD_*` watermark names are preferred. Short aliases `MOONCAKE_DISK_EVICTION_HIGH_WATERMARK_RATIO` and `MOONCAKE_DISK_EVICTION_LOW_WATERMARK_RATIO` are also accepted. The high watermark must be greater than the low watermark. ### Bucket backend settings @@ -137,7 +164,7 @@ Applies when `MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR=file_per_key_storage_b | Environment Variable | Default | Description | |---|---|---| | `MOONCAKE_OFFLOAD_FSDIR` | `file_per_key_dir` | Subdirectory name created under `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` | -| `ENABLE_EVICTION` | `true` | Enables local-storage eviction logic for this backend | +| `MOONCAKE_OFFLOAD_ENABLE_EVICTION` | `true` | Enables local-storage eviction logic for this backend. Short alias: `ENABLE_EVICTION` | --- @@ -165,7 +192,7 @@ Stores each object in an individual file. Simple and easy to inspect, but genera | Environment Variable | Default | Description | |---|---|---| | `MOONCAKE_OFFLOAD_FSDIR` | `file_per_key_dir` | Subdirectory name under `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH` where objects are stored | -| `MOONCAKE_OFFLOAD_ENABLE_EVICTION` | `true` | Enable disk eviction when the total size exceeds the quota | +| `MOONCAKE_OFFLOAD_ENABLE_EVICTION` | `true` | Enable disk eviction for this backend. Short alias: `ENABLE_EVICTION` | Best for: debugging or small-scale deployments. @@ -181,7 +208,9 @@ Best for: high-concurrency scenarios with many small objects where restart durab --- -## Eviction (Bucket Backend Only) +## Eviction + +### Write-time eviction When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend automatically evicts buckets before writing new ones if total disk usage would exceed the limit. @@ -193,6 +222,20 @@ When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend automatically Eviction is two-phase: the bucket is removed from metadata and master is notified first, then in-flight reads are drained before files are deleted. +### Proactive watermark eviction + +When `MOONCAKE_OFFLOAD_ENABLE_DISK_WATERMARK_EVICTION=true`, the FileStorage heartbeat asks the backend to check local-disk usage every `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` seconds. If usage exceeds `MOONCAKE_OFFLOAD_DISK_EVICTION_HIGH_WATERMARK_RATIO`, the backend evicts toward `MOONCAKE_OFFLOAD_DISK_EVICTION_LOW_WATERMARK_RATIO`. + +| Backend | Behavior | +|---------|----------| +| `bucket_storage_backend` | Reuses `MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY`; set it to `fifo` or `lru` because the default policy `none` makes watermark eviction a no-op | +| `file_per_key_storage_backend` | Requires `MOONCAKE_OFFLOAD_ENABLE_EVICTION=true`; reuses the FIFO file eviction queue and recovered keys from startup metadata scan | +| `offset_allocator_storage_backend` | No-op in this version | + +The watermark ratios apply to each backend's quota. For `bucket_storage_backend`, the quota is `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE`; when that value is `0`, the backend defaults it to 90% of the physical disk capacity. For `file_per_key_storage_backend`, the underlying file backend uses its storage quota, which defaults to 90% of the physical disk capacity in SSD offload mode. + +For watermark eviction, the real client notifies the master before deleting local files. If the notification fails, the selected files remain tracked locally and are retried by a later heartbeat. + --- ## Example @@ -211,7 +254,10 @@ The following example starts a master and a real client on a single machine. ```bash mooncake_master \ --rpc_port=50051 \ - --enable_offload=true + --enable_offload=true \ + --offload_on_evict=true \ + --promotion_on_hit=true \ + --promotion_admission_threshold=2 ``` ### Start the real client (new terminal) diff --git a/docs/source/design/conductor/conductor-architecture-design.md b/docs/source/design/conductor/conductor-architecture-design.md index e21cb74e6d..115156b5c4 100644 --- a/docs/source/design/conductor/conductor-architecture-design.md +++ b/docs/source/design/conductor/conductor-architecture-design.md @@ -178,5 +178,5 @@ mooncake-conductor/ +-- CMakeLists.txt ``` -See [Indexer API](./indexer-api-design.md) for the HTTP API and KV Events wire +See [Indexer API](../../api-reference/http/conductor-indexer.md) for the HTTP API and KV Events wire format. diff --git a/docs/source/design/hicache-design.md b/docs/source/design/hicache-design.md index db2dfa2b2e..a2a301e869 100644 --- a/docs/source/design/hicache-design.md +++ b/docs/source/design/hicache-design.md @@ -2,7 +2,7 @@ With the rapid development of tasks such as Agentic Coding, the length of request contexts continues to grow. Increasing the capacity of the KV Cache to improve its hit rate has become increasingly important for enhancing throughput and reducing TTFT. In this context, SGLang introduces **HiCache**, which extends the original RadixAttention (previously limited to GPU memory) by adding hierarchical caching support and integrating with distributed storage backends such as Mooncake. -Inspired by the classic three-level cache design of modern CPUs, HiCache organizes GPU memory as L1, host memory as L2, and distributed storage as L3. This hierarchy enables HiCache to fully exploit the "idle" storage space of GPUs and CPUs, while integrating distributed cache systems for global KV cache storage and scheduling. As a result, HiCache significantly expands KV cache capacity while maintaining strong read performance, especially in workloads such as multi-QA and long-context inference, where KV cache reuse is frequent. For detailed benchmark results, see [this document](https://kvcache-ai.github.io/Mooncake/performance/sglang-hicache-benchmark-results-v1.html). +Inspired by the classic three-level cache design of modern CPUs, HiCache organizes GPU memory as L1, host memory as L2, and distributed storage as L3. This hierarchy enables HiCache to fully exploit the "idle" storage space of GPUs and CPUs, while integrating distributed cache systems for global KV cache storage and scheduling. As a result, HiCache significantly expands KV cache capacity while maintaining strong read performance, especially in workloads such as multi-QA and long-context inference, where KV cache reuse is frequent. For detailed benchmark results, see [this document](https://kvcache-ai.github.io/Mooncake/performance/sglang/sglang-hicache-benchmark-results-v1.html). While HiCache supports multiple L3 backends, this document focuses primarily on the **Mooncake** backend. @@ -115,7 +115,7 @@ Furthermore, **Mooncake** supports efficient batch read and write operations and ## Integration with PD-Disaggregation Deployment Mode -SGLang supports a PD (Prefill-Decode) disaggregation deployment mode through the **Mooncake TransferEngine** (for details, see [this document](https://docs.sglang.ai/advanced_features/pd_disaggregation.html)). +SGLang supports a PD (Prefill-Decode) disaggregation deployment mode through the **Mooncake TransferEngine** (for details, see [this document](https://docs.sglang.ai/advanced_features/pd_disaggregation.html)). In the PD-disaggregation deployment mode, HiCache can be enabled on the Prefill nodes to optimize prefill performance. With the hierarchical caching mechanism provided by **HiCache + Mooncake Store**, prefill nodes can handle long-context and multi-turn dialogue scenarios more efficiently, significantly improving performance during the prefill phase. HiCache can also be enabled on the decode nodes to write computation results back to L3. diff --git a/docs/source/design/index.md b/docs/source/design/index.md new file mode 100644 index 0000000000..6b69421b4e --- /dev/null +++ b/docs/source/design/index.md @@ -0,0 +1,37 @@ +--- +orphan: true +--- + +# Design Documents + +Architecture and implementation details for Mooncake's storage, transfer, and +distributed execution components. + +## Core Architecture + +| Document | Description | +|----------|-------------| +| [Mooncake Architecture](architecture) | KVCache-centric disaggregated serving architecture. | +| [Mooncake Store](mooncake-store) | Distributed object and KV cache storage design. | +| [Transfer Engine](transfer-engine/index) | High-performance data movement architecture and transports. | +| [P2P Store](p2p-store) | Peer-to-peer checkpoint and object transfer design. | + +## Serving and Cache Systems + +| Document | Description | +|----------|-------------| +| [HiCache](hicache-design) | Hierarchical KV cache design. | +| [Engram](engram) | Distributed serving and cache architecture. | +| [Unified Parallel Tensor I/O](unified-parallel-tensor-io) | Parallel tensor storage and transfer model. | +| [SSD Offload](ssd-offload) | SSD-backed cache hierarchy design. | +| [SSD Free-Ratio-First Allocation](ssd-free-ratio-first-allocation) | Capacity-aware replica placement strategy. | + +## Distributed Execution and Routing + +| Document | Description | +|----------|-------------| +| [Mooncake Backend (PG)](mooncake-backend-pg) | Fault-tolerant PyTorch process-group backend. | +| [Mooncake EP](mooncake-ep) | Expert-parallel communication and recovery. | +| [TENT](tent/overview) | Next-generation transfer engine design. | +| [TENT Benchmark](tent/tebench) | TENT benchmark framework and methodology. | +| [Conductor](conductor/conductor-architecture-design) | Cache-aware request routing architecture. | diff --git a/docs/source/design/mooncake-backend-pg.md b/docs/source/design/mooncake-backend-pg.md index 3983560505..ed61f776a3 100644 --- a/docs/source/design/mooncake-backend-pg.md +++ b/docs/source/design/mooncake-backend-pg.md @@ -220,5 +220,5 @@ subgroups, `extend_group_size_to()`, `get_peer_state()`, `recover_ranks()`, or ## Related documentation - [Mooncake EP design](mooncake-ep.md) -- [Python API reference](../python-api-reference/ep-backend.md) +- [Python API reference](../api-reference/python/ep-backend.md) - [PG/EP troubleshooting](../troubleshooting/pg-ep-troubleshooting.md) diff --git a/docs/source/design/mooncake-ep.md b/docs/source/design/mooncake-ep.md index e0734e0b76..cbd977e0af 100644 --- a/docs/source/design/mooncake-ep.md +++ b/docs/source/design/mooncake-ep.md @@ -243,5 +243,5 @@ Adapt launch commands to the target environment and number of GPUs. ## Related documentation - [Mooncake Backend (PG) design](mooncake-backend-pg.md) -- [Python API reference](../python-api-reference/ep-backend.md) +- [Python API reference](../api-reference/python/ep-backend.md) - [PG/EP troubleshooting](../troubleshooting/pg-ep-troubleshooting.md) diff --git a/docs/source/design/mooncake-store.md b/docs/source/design/mooncake-store.md index 86f7ba72b5..aa68f03cf8 100644 --- a/docs/source/design/mooncake-store.md +++ b/docs/source/design/mooncake-store.md @@ -91,6 +91,46 @@ To reduce cache warm-up time after a master restart, the Master Service supports > > The snapshot storage location is **exclusively managed** by the Mooncake snapshot system. Old snapshots are automatically deleted during cleanup. **DO NOT store other files in this location.** Use a dedicated, isolated storage for snapshots. +### Tenant Quota + +The Master Service can optionally enforce strict multi-tenant memory quota admission. This feature is disabled by default. When `enable_multi_tenants=false`, request tenant IDs are ignored for object placement, all objects use the `default` namespace, and tenant quota management requests return `UNAVAILABLE_IN_CURRENT_MODE`. + +When strict multi-tenant mode is enabled, the tenant quota policy is loaded from the configured connector. Supported connector types are `file` and, when the store is built with `STORE_USE_ETCD=ON`, `etcd`. The `file` connector uses `tenant_quota_connector_uri=` as a writable YAML policy path. The `etcd` connector uses `tenant_quota_connector_uri=` as the etcd endpoints string and stores the same YAML policy in `mooncake-store//tenant_quota_policy`; if that key does not exist, the master starts with an empty policy so the first policy can be created through the admin API. The etcd connector shares the process-wide store etcd client used by HA/oplog, so deployments that enable both must configure matching etcd endpoints. Tenants must be explicitly present in that connector policy before they can write. Missing tenants, empty tenants, and an unregistered `default` tenant are rejected with `TENANT_NOT_REGISTERED`. + +The YAML policy uses schema version `1`: + +```yaml +version: 1 + +tenants: + - name: tenant-a + quota: 200GB +``` + +Tenant names must be non-empty, unique, must not start with `_`, and must not contain NUL or control characters. Quotas must be positive integers and may use `B`, `KB`, `MB`, `GB`, or `TB` units. + +Effective quota is recomputed from the current registered memory capacity: + +- If explicit tenant requests fit within the registered memory capacity, tenants receive their requested quotas and remaining capacity stays unallocated. +- If explicit tenant requests exceed registered memory capacity, explicit tenants receive quota scaled proportionally by request size. +- Remainders are assigned deterministically by tenant ID, so repeated recomputes produce stable results. +- Tenants present in restored metadata but missing from the connector policy become in-memory orphans with requested quota `0`, effective quota `0`, and `over_quota=true` while they still own metadata. Reads and removals are allowed so operators can clean them up; writes remain blocked until the tenant is re-registered or emptied. + +`PutStart` and size-changing `UpsertStart` charge quota before memory is allocated. If the first reservation fails, the master performs tenant-scoped memory eviction for the target tenant and retries the reservation. The retry is bounded to two eviction attempts. Tenant quota eviction scans only the target tenant, skips hard-pinned objects, honors soft-pin eviction configuration, and preserves grouped-object lease safety checks. + +Admin policy changes are persisted before the final in-memory policy is applied. `PUT` writes the connector first and then applies the policy in memory. `DELETE` first marks the tenant unregistered in memory to block concurrent writes, verifies the tenant is empty, writes the connector, and rolls back the in-memory mark if the connector write fails. The admin HTTP API exposes: + +| Method | Path | Description | +|--------|------|-------------| +| `GET` | `/api/v1/tenant_quotas` | List quota snapshots for active or explicit tenants | +| `GET` | `/api/v1/tenant_quotas?tenant_id=` | Query one tenant quota snapshot | +| `PUT` | `/api/v1/tenant_quotas?tenant_id=` | Create or update a tenant quota policy | +| `DELETE` | `/api/v1/tenant_quotas?tenant_id=` | Delete an empty tenant quota policy | + +Tenant quota snapshots include `tenant_id`, `requested_quota_bytes`, `effective_quota_bytes`, `used_bytes`, `reserved_bytes`, `committed_count`, `metadata_object_count`, `over_quota`, and `has_explicit_policy`. + +Snapshots restore object runtime state only. Tenant quota policy is always loaded from the connector after metadata restore, then usage and effective quota are rebuilt from restored metadata and current registered capacity. If the connector cannot be loaded in strict multi-tenant mode, startup fails. + ### Master Service APIs The protobuf definition between Master and Client is as follows: @@ -454,7 +494,7 @@ Mooncake Store provides two concrete implementations of `BufferAllocatorBase`: **OffsetBufferAllocator (default and recommended)**: This allocator is derived from [OffsetAllocator](https://github.com/sebbbi/OffsetAllocator), which uses a custom bin-based allocation strategy that supports fast hard realtime `O(1)` offset allocation with minimal fragmentation. Mooncake Store optimizes this allocator based on the specific memory usage characteristics of LLM inference workloads, thereby enhancing memory utilization in LLM scenarios. -For measured utilization and allocation latency across LLM-style workloads, see [Allocator Performance](../performance/allocator-benchmark-result.md). +For measured utilization and allocation latency across LLM-style workloads, see [Allocator Performance](../performance/mooncake/allocator-benchmark-result.md). **CachelibBufferAllocator (deprecated)**: This allocator leverages Facebook's [CacheLib](https://github.com/facebook/CacheLib) to manage memory using a slab-based allocation strategy. It provides efficient memory allocation with good fragmentation resistance and is well-suited for high-performance scenarios. However, in our modified version, it does not handle workloads with highly variable object sizes effectively, so it is currently marked as deprecated. @@ -524,7 +564,7 @@ Mooncake Store provides multiple built-in allocation strategies to control how s ./build/mooncake-store/src/mooncake_master --allocation_strategy=free_ratio_first ``` -Valid values are: `random` (default), `free_ratio_first`, `cxl` (case-sensitive). +Valid values are: `random` (default), `free_ratio_first`, `ssd_free_ratio_first`, `cxl`, `local_first` (case-sensitive). #### How to Choose @@ -532,7 +572,9 @@ Valid values are: `random` (default), `free_ratio_first`, `cxl` (case-sensitive) |---|---|---| | `random` | Maximum throughput, stable clusters | Limited load balancing; slow convergence when new segments join | | `free_ratio_first` | Balanced utilization, dynamic scaling | Slightly lower throughput due to sampling and sorting overhead | +| `ssd_free_ratio_first` | SSD-aware memory allocation when SSD offloading is enabled | Depends on SSD usage metrics; falls back to random allocation when needed | | `cxl` | CXL memory hardware | CXL-specific; single-replica only | +| `local_first` | Colocated inference workers and memory store segments | Requires stable host identity in `local_hostname`; single memory replica only | **Use `random`** (default) when your cluster is relatively stable (segments rarely join or leave) and you want the highest possible allocation throughput. @@ -540,9 +582,13 @@ Valid values are: `random` (default), `free_ratio_first`, `cxl` (case-sensitive) - Segments have different capacities and you want even utilization ratios. - New segments are dynamically added at runtime and you need them to absorb load quickly. With `random`, convergence to a well-balanced state can be slow on large or dynamic clusters; `free_ratio_first` accelerates this by preferentially filling emptier segments, substantially increasing the likelihood that newly joined segments are selected for allocations (see details below). +**Use `ssd_free_ratio_first`** when SSD offloading is enabled and you want memory allocation to prefer segments whose backing SSD still has more free capacity. + **Use `cxl`** only when your hardware includes CXL (Compute Express Link) memory devices and you want to allocate data exclusively on CXL segments. -For benchmark data comparing `random` and `free_ratio_first` across segment counts, replica counts, and skewed capacities, see [AllocationStrategy Performance](../performance/allocation-strategy-benchmark-result.md). +**Use `local_first`** when inference workers and Mooncake Store memory segments are colocated and you want writes to prefer the writer's host before falling back to other hosts. For this strategy to work correctly, all writer and store processes on the same physical or logical host must use the same stable, globally unique host part in `local_hostname`. + +For benchmark data comparing `random` and `free_ratio_first` across segment counts, replica counts, and skewed capacities, see [AllocationStrategy Performance](../performance/mooncake/allocation-strategy-benchmark-result.md). #### Strategy Details @@ -571,6 +617,16 @@ The overhead is minimal: sampling is `O(K)` and sorting is `O(K log K)`, where K The key insight behind Best-of-N is that if a new/empty segment is sampled, it will almost certainly be ranked first due to having the highest free ratio, which naturally accelerates convergence when new segments join the cluster. +**`ssd_free_ratio_first` — SsdFreeRatioFirstAllocationStrategy** + +An SSD-aware variant of the free-ratio-first strategy. It first tries preferred segments, then samples candidate segment names and sorts them by SSD free ratio reported by the local disk segment metrics provider. If it cannot satisfy all replicas from the sorted candidates, it falls back to random allocation for the remaining replicas. + +**`local_first` — Local-first allocation** + +Host-aware local-first allocation reuses the normal preferred-segment flow. The master derives the writer host id from the request's client host identity and builds an ordered preferred segment list: active hosts are visited in cyclic lexicographic host-id order, starting from the writer host when it has active segments, or otherwise from the next greater active host id. Within the same host, segment names are sorted and rotated by key hash so multiple local segments do not always receive the first allocation attempt. + +This strategy currently applies to memory allocation with `replica_num == 1`. Explicit `preferred_segment` or `preferred_segments` in `ReplicateConfig` are still tried first; if they are unavailable or full, allocation continues with the local-first ordered fallback list. + **`cxl` — CxlAllocationStrategy** Specialized for CXL (Compute Express Link) memory hardware. Unlike the other strategies, this one does not perform random or load-balanced selection — it always allocates from a specific CXL segment: @@ -583,7 +639,7 @@ Limitations: This strategy only supports single-replica allocation (does not dis ## Eviction Policy -When a `PutStart` request fails due to insufficient memory, or when the eviction thread detects that space usage has reached the configured high watermark (95% by default, configurable via `-eviction_high_watermark_ratio`), an eviction task is triggered to free up space by evicting a portion of objects (5% by default, configurable via `-eviction_ratio`). Similar to `Remove`, evicted objects are simply marked as deleted, with no data transfer required. +When a `PutStart` request fails due to insufficient memory, or when the eviction thread detects that space usage has reached the configured high watermark (90% by default, configurable via `-eviction_high_watermark_ratio`), an eviction task is triggered to free up space by evicting a portion of objects (5% by default, configurable via `-eviction_ratio`). Similar to `Remove`, evicted objects are simply marked as deleted, with no data transfer required. Currently, an approximate LRU policy is adopted, where the least recently used objects are preferred for eviction. To avoid data races and corruption, objects currently being read or written by clients should not be evicted. For this reason, objects that have leases or have not been marked as complete by `PutEnd` requests will be ignored by the eviction task. @@ -597,7 +653,7 @@ For grouped objects, a successful `ExistKey` or `GetReplicaList` refreshes the l However, if the lease expires before a `Get` operation finishes reading the data, the operation will be considered failed, and no data will be returned, in order to prevent potential data corruption. -The default lease TTL is 5 seconds and is configurable via a startup parameter of `master_service`. +The default lease TTL is 10 seconds and is configurable via a startup parameter of `master_service`. ## Soft Pin @@ -675,6 +731,8 @@ When the user specifies `--root_fs_dir=/path/to/dir` when starting the master, a ​Note​​: When enabling this feature, the user must ensure that the DFS-mounted directory (`root_fs_dir=/path/to/dir`) is valid and consistent across all client hosts. If some clients have invalid or incorrect mount paths, it may cause abnormal behavior in Mooncake Store. +This `root_fs_dir` path is a legacy persistence path. SSD offload uses `--enable_offload=true` on the master and real client, stores data under the real client's `MOONCAKE_OFFLOAD_FILE_STORAGE_PATH`, and records `LOCAL_DISK` replicas. Do not use `--root_fs_dir` with `--enable_offload=true`. + ### Persistent Storage Space Configuration​ Mooncake provides configurable DFS available space. Users can specify `--global_file_segment_size=1048576` when starting the master, indicating a maximum usable space of 1MB on DFS. The current default setting is the maximum value of int64 (as we generally do not restrict DFS storage usage), which is displayed as `infinite` in `mooncake_maseter`'s console logs. @@ -731,7 +789,7 @@ For detailed guidance on monitoring master metrics, Prometheus endpoints, and he ## Mooncake Store Python API -**Complete Python API Documentation**: [https://kvcache-ai.github.io/Mooncake/python-api-reference/mooncake-store.html](https://kvcache-ai.github.io/Mooncake/python-api-reference/mooncake-store.html) +**Complete Python API Documentation**: [https://kvcache-ai.github.io/Mooncake/api-reference/python/mooncake-store.html](https://kvcache-ai.github.io/Mooncake/api-reference/python/mooncake-store.html) ## Version Management Policy diff --git a/docs/source/design/ssd-free-ratio-first-allocation.md b/docs/source/design/ssd-free-ratio-first-allocation.md new file mode 100644 index 0000000000..337f171c72 --- /dev/null +++ b/docs/source/design/ssd-free-ratio-first-allocation.md @@ -0,0 +1,152 @@ +# SSD Free-Ratio-First Allocation Design + +## Overview + +Mooncake Store distributes KV cache objects across multiple memory segments hosted on different nodes. The master's allocation strategy decides which segment receives each new object replica. When using DDR-only strategies such as `random` or `free_ratio_first`, the allocator ignores SSD state entirely. In deployments where some segments have SSD offload enabled and others do not, this blind allocation can concentrate traffic on a small subset of segments whose SSD capacity is quickly exhausted, while segments with ample SSD headroom remain underutilized. + +This document describes `SsdFreeRatioFirstAllocationStrategy`, an allocation strategy that ranks candidate segments by their SSD free ratio and preferentially allocates to segments with the most available SSD space. It integrates with the existing allocation framework and adds SSD usage tracking to `LocalDiskSegment` so that the master can make informed placement decisions. + +--- + +## Architecture + +``` + Allocation Request + │ + ▼ + ┌───────────────────────┐ + │ Sample candidate │ + │ segments (up to │ + │ 6 * replica_num) │ + └───────────┬───────────┘ + │ + ▼ + ┌───────────────────────┐ + │ Compute SSD free │ + │ ratio per candidate │ + └───────────┬───────────┘ + │ + ▼ + ┌───────────────────────┐ + │ Sort by SSD free │ + │ ratio descending │ + └───────────┬───────────┘ + │ + ▼ + ┌───────────────────────┐ + │ Allocate from top │ + │ candidates │ + └───────────┬───────────┘ + │ + ┌──────────┴──────────┐ + │ Remaining replicas? │ + └──────┬──────┬───────┘ + │ │ + Yes ◄┘ └► No → Done + │ + ▼ + ┌───────────────────────┐ + │ Fallback: random │ + │ allocation for │ + │ remaining replicas │ + └───────────────────────┘ +``` + +The flow follows the same high-level structure as the existing `FreeRatioFirstAllocationStrategy`: sample a subset of candidates, compute a ranking metric, sort, and allocate from the top. The key difference is that the ranking metric is SSD free ratio rather than DRAM free ratio. + +`MasterService` only passes an `SsdMetricsProvider` when the effective allocation strategy is `SSD_FREE_RATIO_FIRST`. Non-SSD strategies receive `nullptr`, so they avoid unnecessary local-disk segment access. + +--- + +## Core Algorithm + +### Candidate sampling + +For each allocation request requesting `replica_num` replicas, the strategy samples `min(6 * replica_num, total_segments)` candidate segments. This bounded sampling keeps the sorting cost predictable regardless of cluster size while still providing a statistically diverse candidate set. + +### SSD free ratio + +For each sampled segment, the SSD free ratio is computed as: + +``` +ssd_free_ratio = (ssd_total_capacity - ssd_used_bytes) / ssd_total_capacity +``` + +A segment with 1 TB total SSD capacity and 200 GB used has an SSD free ratio of 0.80. A segment whose SSD is full has a ratio of 0.0. + +Before calculating the ratio, `ssd_used_bytes` is clamped to `[0, ssd_total_capacity]`. This keeps transient concurrent accounting drift from producing a negative free ratio or a value greater than 1.0. If no SSD metrics provider is available, or if the reported total capacity is not positive, the strategy treats the segment as fully free. + +### Sorting + +Candidates are sorted by SSD free ratio in descending order. Segments with more available SSD space appear first and are preferred for allocation. + +### Preferred segments + +As with other allocation strategies, segments marked as preferred by the caller are handled first. Preferred segments bypass the SSD free ratio ranking and are allocated immediately if they have sufficient capacity. + +### Fallback to random allocation + +After allocating from the SSD-ranked candidates, any remaining replicas that could not be satisfied are allocated using the standard random strategy as a fallback. This ensures that allocation succeeds even when SSD metrics are unavailable (for example, on segments without SSD offload configured). + +--- + +## SSD Usage Tracking + +### `ssd_used_bytes` counter + +`LocalDiskSegment` maintains `ssd_total_capacity_bytes`, updated by `ReportSsdCapacity`, and an atomic counter `ssd_used_bytes` that tracks the total number of bytes currently occupied by offloaded replicas on the segment's SSD. `ssd_used_bytes` is updated alongside metadata changes: + +- **Increment**: `NotifyOffloadSuccess` increments `ssd_used_bytes` by the object size only after the master successfully adds a `LOCAL_DISK` replica to the object entry. If the object has already disappeared from metadata, the notification is ignored and the counter is not changed. +- **Decrement**: The master decrements `ssd_used_bytes` via `ReleaseLocalDiskUsage` whenever a `LOCAL_DISK` replica is removed from metadata — on full object deletion (`EraseMetadata`), and on partial replica removal through the shared erase helper (`EraseReplicasWithCacheTotalAccounting`, used by batch clear, revoke, and stale-handle cleanup) and the local-disk eviction path. `ReleaseLocalDiskUsage` iterates only `LOCAL_DISK` replicas and is a no-op for other replica types, so it is safe to call on any replica set. + +The counter is atomic to allow concurrent updates from multiple RPC handler threads without requiring a separate lock. The allocation strategy treats it as an eventually consistent placement signal and clamps it before computing the free ratio. + +### `SsdMetricsProvider` interface + +`ScopedLocalDiskSegmentAccess` implements the `SsdMetricsProvider` interface, which exposes two methods: + +| Method | Return type | Description | +|--------|-------------|-------------| +| `getSsdTotalCapacity` | `int64_t` | Total SSD capacity configured for the segment, in bytes | +| `getSsdUsedBytes` | `int64_t` | Current SSD usage, read from `ssd_used_bytes` | + +The allocation strategy queries these methods through the `SsdMetricsProvider` interface, keeping the strategy decoupled from the concrete segment implementation. + +--- + +## Configuration Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `allocation_strategy` | string | `"random"` | Set to `"ssd_free_ratio_first"` to enable SSD-ratio-based load balancing | + +The parameter is passed as a gflag to the master process at startup. + +--- + +## Code Structure + +| File | Change | +|------|--------| +| `mooncake-store/include/types.h` | Add `SSD_FREE_RATIO_FIRST` enum value to the allocation strategy enum | +| `mooncake-store/include/allocation_strategy.h` | Add `SsdMetricsProvider` interface and `SsdFreeRatioFirstAllocationStrategy` class | +| `mooncake-store/include/segment.h` | Add `ssd_total_capacity_bytes` and `ssd_used_bytes` fields to `LocalDiskSegment`; inherit `SsdMetricsProvider` | +| `mooncake-store/src/segment.cpp` | Implement `getSsdTotalCapacity` and `getSsdUsedBytes` | +| `mooncake-store/src/master_service.cpp` | Pass SSD metrics only to `SSD_FREE_RATIO_FIRST`; update `ssd_used_bytes` after successful `NotifyOffloadSuccess` metadata insertion; release usage when `LOCAL_DISK` replicas are erased | + +--- + +## Usage Example + +Start the master with the SSD free-ratio-first strategy: + +```bash +./mooncake_master --allocation_strategy=ssd_free_ratio_first +``` + +With this configuration: + +1. The master samples up to `6 * replica_num` candidate segments for each allocation request. +2. Candidates are ranked by SSD free ratio (descending). +3. Allocation proceeds from the top-ranked candidates. +4. Any remaining replicas fall back to random allocation. diff --git a/docs/source/design/ssd-offload.md b/docs/source/design/ssd-offload.md index d9430f5dd1..1ca2607499 100644 --- a/docs/source/design/ssd-offload.md +++ b/docs/source/design/ssd-offload.md @@ -6,7 +6,7 @@ Mooncake Store supports offloading KV cache objects from distributed memory to l SSD offload is implemented as a background subsystem within the **real client** process. It is transparent to the application: a `Put` that would otherwise be evicted from memory is persisted to disk, and a `Get` that finds no memory replica automatically falls back to reading from SSD. -For multi-turn conversation benchmark results, see [Mooncake SSD Offload Benchmark](../performance/ssd-offload-benchmark-results.md). +For multi-turn conversation benchmark results, see [Mooncake SSD Offload Benchmark](../performance/mooncake/ssd-offload-benchmark-results.md). --- @@ -87,7 +87,7 @@ Step by step: 1. **Heartbeat**: The heartbeat thread wakes up every `MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS` seconds and calls `client_->OffloadObjectHeartbeat(enable_offloading_, offloading_objects)`. The master replies with a map of `{key → size}` for objects it has selected to evict from memory. 2. **Read slices from memory**: `FileStorage::OffloadObjects` groups the keys into buckets (for `BucketStorageBackend`) and calls `BatchQuerySegmentSlices` to obtain `{key → Slice}` from the local memory segment via `client_->BatchQuery`. -3. **Eviction** (if capacity limit is set): Before writing, `PrepareEviction` removes old buckets from metadata under the exclusive lock and collects their keys. The `eviction_handler` callback calls `client_->BatchEvictDiskReplica` to notify the master in a single RPC. `FinalizeEviction` then deletes the corresponding files. +3. **Eviction** (if capacity limit is set): Before writing, `PrepareEviction` removes old buckets from metadata under the exclusive lock and collects their keys. The `eviction_handler` callback calls `client_->BatchEvictDiskReplica` to notify the master in one batched RPC. `FinalizeEviction` then deletes the corresponding files. 4. **Write to SSD**: `StorageBackend::BatchOffload` serializes and writes the key-value data to disk. 5. **Notify master**: On success, the `complete_handler` calls `client_->NotifyOffloadSuccess(keys, metadatas)`. The master adds a `LOCAL_DISK` replica entry (carrying the real client's RPC address as `transport_endpoint`) to the object's replica list. @@ -150,20 +150,36 @@ Each object is stored as an individual file. The file path is derived from the k ### OffsetAllocatorStorageBackend -A single pre-allocated file (`kv_cache.data`) is shared by all objects. Space within the file is managed by an `OffsetAllocator`. Metadata is sharded across 1024 independent maps to reduce lock contention under high concurrency. Records follow the layout `[key_len: u32 | value_len: u32 | key | value]`. +A single pre-allocated file (`kv_cache.data`) is shared by all objects. Space within the file is managed by an `OffsetAllocator`. Metadata is sharded across 1024 independent maps to reduce lock contention under high concurrency. Records follow the layout `[key_len: u32 | value_len: u32 | seq: u64 | flags: u32 | crc32: u32 | key | zero padding | value]`. The value region always starts at a 4 KiB boundary within the record (padding is a pure function of `key_len`) so that DMA writers (e.g. GDS/cuFile) can use aligned file offsets. `seq` is a monotonic stamp used by restart recovery to drop post-checkpoint writes; `crc32` is a CRC-32C over header prefix + key + value and is present only when `flags & kFlagHasCrc` (see `enable_record_crc` in `OffsetAllocatorBackendConfig`). --- -## Eviction (BucketStorageBackend) +## Eviction + +### Write-time eviction (BucketStorageBackend) When `MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE` is set, the backend evicts existing buckets to make room before writing a new one. Eviction is disabled by default (`BucketEvictionPolicy::NONE`). +### Proactive watermark eviction + +`FileStorage::Heartbeat()` also calls the backend-level proactive disk watermark path when `MOONCAKE_OFFLOAD_ENABLE_DISK_WATERMARK_EVICTION=true`. The high watermark decides when eviction starts and the low watermark decides the target usage after eviction. This path is independent of write admission, so disk usage can move back toward the low watermark even when no new write arrives. + +| Backend | Candidate source | +|---------|------------------| +| `BucketStorageBackend` | The configured bucket eviction policy (`FIFO` or `LRU`) | +| `StorageBackendAdaptor` | The FIFO file queue rebuilt by `StorageBackend::Init()`; `ScanMeta()` restores object keys for recovered files | +| `OffsetAllocatorStorageBackend` | No-op in this version | + +The watermark path uses the same master-notification ordering as bucket write-time eviction: selected keys are reported through the eviction handler before files are deleted. If master notification fails, `BucketStorageBackend` restores the prepared metadata and `StorageBackendAdaptor` restores the selected file records to the FIFO queue. + +For `BucketStorageBackend`, write-time eviction and watermark eviction both use `PrepareEviction()` and `FinalizeEviction()`, so bucket candidate selection is serialized by the same metadata lock. For `StorageBackendAdaptor`, write-time reactive eviction and watermark eviction use the same FIFO queue and space-accounting locks. + ### Policies | Policy | Candidate selection | |--------|---------------------| | `FIFO` | `buckets_.begin()` — always the oldest bucket, since `buckets_` is ordered by bucket ID | -| `LRU` | `std::min_element` over `BucketMetadata::last_access_ns_` — the bucket with the smallest last-read timestamp | +| `LRU` | `lru_index_`, ordered by `{last_access_ns_, bucket_id}` — the bucket with the smallest last-read timestamp | `last_access_ns_` is an atomic `int64_t` updated on every `BatchLoad` with relaxed ordering. Buckets that have never been read have `last_access_ns_ == 0` and are therefore always evicted first under LRU, giving FIFO-among-unread semantics. @@ -179,20 +195,28 @@ Eviction is split into two phases to ensure that the master is notified before f **Between phases** — notify master: -The caller invokes the `eviction_handler` callback with the full list of evicted keys. The handler calls `MasterClient::BatchEvictDiskReplica`, which sends a single RPC to the master to remove the disk replicas for all evicted keys atomically. +The caller invokes the `eviction_handler` callback with the full list of evicted keys. The handler calls `MasterClient::BatchEvictDiskReplica`, which sends one batched RPC to the master and receives a per-key result. **Phase 2 — `FinalizeEviction(pending)`** (called after master notification): For each evicted bucket: -1. Spin-wait (with a 10-second timeout) until `inflight_reads_ == 0`. -2. Evict any stale file-handle cache entries. -3. Delete the `.bucket` and `.meta` files. - -This ordering guarantees: +1. Attempt to delete the `.meta` file first, reducing the chance that a later + restart recovers a bucket whose master replica has already been removed. +2. Spin-wait (with a 10-second timeout) until `inflight_reads_ == 0`. +3. Evict any stale file-handle cache entries. +4. Delete the `.bucket` file. + +This ordering provides the following behavior when the relevant file operations +succeed: - The master never serves a stale disk-replica location for a file that has already been deleted. - Ongoing reads complete successfully before their files are removed. - Freed disk space is available for the incoming write before `WriteBucket` is called. +Cleanup failures after the master update are logged and treated as best-effort. +A crash or filesystem failure that leaves both files intact still requires +cross-node reconciliation; local file ordering alone cannot make that window +transactional. + --- ## io_uring File I/O diff --git a/docs/source/design/tent/metrics.md b/docs/source/design/tent/metrics.md index 463756b9c1..65055fd74a 100644 --- a/docs/source/design/tent/metrics.md +++ b/docs/source/design/tent/metrics.md @@ -23,7 +23,7 @@ By default, metrics are **disabled** at compile time for maximum performance. To cmake -DTENT_METRICS_ENABLED=ON .. ``` -When disabled at compile time (`TENT_METRICS_ENABLED=OFF`, the default), all metrics macros expand to `((void)0)`, resulting in **zero runtime overhead**. +When disabled at compile time (`TENT_METRICS_ENABLED=OFF`, the default), all metrics macros expand to `((void)0)` and the `recordTaskCompletionMetrics` body is `#if`-gated out, resulting in **zero runtime overhead** on the transfer hot path. ### Runtime Disable (Minimal Overhead) @@ -65,11 +65,7 @@ TENT metrics configuration is integrated into the main `transfer-engine.json` co "http_port": 9100, "http_host": "0.0.0.0", "http_server_threads": 2, - "report_interval_seconds": 30, - "enable_prometheus": true, - "enable_json": true, - "latency_buckets": [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1.0], - "size_buckets": [1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864, 268435456, 1073741824] + "report_interval_seconds": 30 }, "transports": { // ... transport configuration @@ -77,10 +73,9 @@ TENT metrics configuration is integrated into the main `transfer-engine.json` co } ``` -**Note**: +**Note**: - `report_interval_seconds`: Set to 0 to disable periodic logging -- `latency_buckets`: Values are in **seconds** (e.g., 0.001 = 1ms). The system internally converts to microseconds for histogram storage. -- `size_buckets`: Values are in bytes +- Histogram buckets are fixed at compile time (see `kLatencyBuckets` / `kSizeBuckets` in `tent_metrics.h`) for reproducible observability across deployments. ### Environment Variables @@ -91,14 +86,6 @@ TENT_METRICS_HTTP_PORT=9100 TENT_METRICS_HTTP_HOST=0.0.0.0 TENT_METRICS_HTTP_SERVER_THREADS=2 TENT_METRICS_REPORT_INTERVAL=30 # Set to 0 to disable periodic logging - -# Output formats -TENT_METRICS_ENABLE_PROMETHEUS=true -TENT_METRICS_ENABLE_JSON=true - -# Custom buckets (comma-separated, latency in seconds, size in bytes) -TENT_METRICS_LATENCY_BUCKETS="0.0001,0.0005,0.001,0.005,0.01,0.05,0.1,0.5,1.0" -TENT_METRICS_SIZE_BUCKETS="1024,4096,16384,65536,262144,1048576" ``` ## Quick Start @@ -131,17 +118,28 @@ tent_metrics.initialize(config); ```cpp // Using convenience macros (recommended) -TENT_RECORD_READ_COMPLETED(1024*1024, 0.025); // 1MB read in 25ms -TENT_RECORD_WRITE_COMPLETED(512*1024, 0.015); // 512KB write in 15ms -TENT_RECORD_READ_FAILED(1024*1024); // 1MB read failed -TENT_RECORD_WRITE_FAILED(512*1024); // 512KB write failed +TENT_RECORD_READ_COMPLETED(RDMA, 1024*1024, 0.025); // 1MB read in 25ms +TENT_RECORD_WRITE_COMPLETED(RDMA, 512*1024, 0.015); // 512KB write in 15ms +TENT_RECORD_READ_FAILED(TCP); // read failed (no bytes recorded) +TENT_RECORD_WRITE_FAILED(TCP); // write failed (no bytes recorded) +TENT_RECORD_TRANSPORT_FAILOVER(RDMA, TCP); // cross-transport failover event // Direct API usage auto& tent_metrics = TentMetrics::instance(); -tent_metrics.recordReadCompleted(1024*1024, 0.025); -tent_metrics.recordWriteCompleted(512*1024, 0.015); -tent_metrics.recordReadFailed(1024*1024); -tent_metrics.recordWriteFailed(512*1024); +tent_metrics.recordReadCompleted(RDMA, 1024*1024, 0.025); +tent_metrics.recordWriteCompleted(RDMA, 512*1024, 0.015); +tent_metrics.recordReadFailed(TCP); +tent_metrics.recordWriteFailed(TCP); +tent_metrics.recordTransportFailover(RDMA, TCP); + +// Deadline feasibility (RFC #2519, observability only): +tent_metrics.recordDeadlineMLU(RDMA, 0.8); // MLU < 1 met the deadline +tent_metrics.recordDeadlineInfeasible(TCP); // deadline was in the past at submit + +// Causal-chain per-stage latency breakdown (microseconds): +tent_metrics.recordStageLatency(TentMetrics::Stage::QueueWait, RDMA, 12.0); +tent_metrics.recordStageLatency(TentMetrics::Stage::Dispatch, RDMA, 45.0); +tent_metrics.recordStageLatency(TentMetrics::Stage::Transport, RDMA, 130.0); ``` ### RAII Latency Measurement @@ -149,12 +147,12 @@ tent_metrics.recordWriteFailed(512*1024); ```cpp // Automatic latency measurement using RAII { - TENT_SCOPED_READ_LATENCY(1024 * 1024); // e.g. 1MB + TENT_SCOPED_READ_LATENCY(RDMA, 1024 * 1024); // e.g. 1MB // ... perform read operation ... } // latency automatically recorded when scope exits { - TENT_SCOPED_WRITE_LATENCY(512 * 1024); // e.g. 512KB + TENT_SCOPED_WRITE_LATENCY(RDMA, 512 * 1024); // e.g. 512KB // ... perform write operation ... } ``` @@ -174,32 +172,30 @@ The HTTP server provides multiple endpoints: ``` # HELP tent_read_bytes_total Total bytes read via TENT # TYPE tent_read_bytes_total counter -tent_read_bytes_total 1048576 +tent_read_bytes_total{transport="rdma"} 1048576 +tent_read_bytes_total{transport="tcp"} 524288 # HELP tent_write_bytes_total Total bytes written via TENT # TYPE tent_write_bytes_total counter -tent_write_bytes_total 524288 +tent_write_bytes_total{transport="rdma"} 524288 # HELP tent_read_requests_total Total read requests via TENT # TYPE tent_read_requests_total counter -tent_read_requests_total 100 - -# HELP tent_write_requests_total Total write requests via TENT -# TYPE tent_write_requests_total counter -tent_write_requests_total 50 +tent_read_requests_total{transport="rdma"} 100 +tent_read_requests_total{transport="tcp"} 50 # HELP tent_read_failures_total Total read failures via TENT # TYPE tent_read_failures_total counter -tent_read_failures_total 2 +tent_read_failures_total{transport="tcp"} 2 -# HELP tent_write_failures_total Total write failures via TENT -# TYPE tent_write_failures_total counter -tent_write_failures_total 1 +# HELP tent_transport_failover_total Total cross-transport failover events +# TYPE tent_transport_failover_total counter +tent_transport_failover_total{from="rdma",to="tcp"} 1 # HELP tent_read_latency_us Read latency distribution in microseconds # TYPE tent_read_latency_us histogram -tent_read_latency_us_bucket{le="100"} 10 -tent_read_latency_us_bucket{le="500"} 50 +tent_read_latency_us_bucket{transport="rdma",le="100"} 10 +tent_read_latency_us_bucket{transport="rdma",le="500"} 50 ... ``` @@ -222,18 +218,59 @@ Read: 1.00 MB (100 reqs, 2 fails) | Write: 512.00 KB (50 reqs, 1 fails) ## Available Metrics -| Metric Name | Type | Description | -|-------------|------|-------------| -| `tent_read_bytes_total` | Counter | Total bytes read via TENT | -| `tent_write_bytes_total` | Counter | Total bytes written via TENT | -| `tent_read_requests_total` | Counter | Total read requests via TENT | -| `tent_write_requests_total` | Counter | Total write requests via TENT | -| `tent_read_failures_total` | Counter | Total read failures via TENT | -| `tent_write_failures_total` | Counter | Total write failures via TENT | -| `tent_read_latency_us` | Histogram | Read latency distribution in microseconds | -| `tent_write_latency_us` | Histogram | Write latency distribution in microseconds | -| `tent_read_size_bytes` | Histogram | Read request size distribution in bytes | -| `tent_write_size_bytes` | Histogram | Write request size distribution in bytes | +| Metric Name | Type | Labels | Description | +|-------------|------|--------|-------------| +| `tent_read_bytes_total` | Counter | `transport` | Total bytes read via TENT (success only; failures record no bytes) | +| `tent_write_bytes_total` | Counter | `transport` | Total bytes written via TENT (success only) | +| `tent_read_requests_total` | Counter | `transport` | Total read requests via TENT (success + failure) | +| `tent_write_requests_total` | Counter | `transport` | Total write requests via TENT (success + failure) | +| `tent_read_failures_total` | Counter | `transport` | Total read failures via TENT | +| `tent_write_failures_total` | Counter | `transport` | Total write failures via TENT | +| `tent_transport_failover_total` | Counter | `from`, `to` | Total cross-transport failover events | +| `tent_transport_attempts_total` | Counter | `transport`, `operation` | Physical transport attempts submitted for execution | +| `tent_transport_attempt_failures_total` | Counter | `transport`, `operation` | Physical transport attempts that terminated with `FAILED` | +| `tent_deadline_infeasible_total` | Counter | `transport` | Transfers whose deadline was already in the past at submit time | +| `tent_read_latency_us` | Histogram | `transport` | Read latency distribution in microseconds | +| `tent_write_latency_us` | Histogram | `transport` | Write latency distribution in microseconds | +| `tent_read_size_bytes` | Histogram | `transport` | Read request size distribution in bytes | +| `tent_write_size_bytes` | Histogram | `transport` | Write request size distribution in bytes | +| `tent_deadline_mlu_permille` | Histogram | `transport` | Deadline feasibility ratio (MLU x 1000); 1000 = MLU 1.0 (the met/missed boundary) | +| `tent_stage_queue_wait_us` | Histogram | `transport` | Causal chain: queue wait latency in microseconds | +| `tent_stage_dispatch_us` | Histogram | `transport` | Causal chain: dispatch latency in microseconds | +| `tent_stage_transport_us` | Histogram | `transport` | Causal chain: transport execution latency in microseconds | +| `tent_transport_attempt_latency_us` | Histogram | `transport`, `operation` | Observed latency of each physical transport attempt | + +**Notes**: +- `*_requests_total` counts terminal logical/merged-transfer outcomes. Its `transport` label is the **final transport**: a request recovered by RDMA→TCP failover is counted once under `tcp`. +- `tent_transport_attempts_total` and `tent_transport_attempt_failures_total` measure physical transport reliability. A recovered RDMA→TCP request contributes one failed RDMA attempt and one successful TCP attempt. +- Attempt latency currently ends when polling or the progress worker observes the terminal status, so it may include completion-observation delay. +- `*_failures_total` does not record bytes; failed transfers transfer no bytes. +- `tent_deadline_infeasible_total` is a dedicated counter (not a histogram sentinel) so infeasible-at-submit cases are distinguishable from genuine high-MLU samples. +- yalantinglibs omits zero-valued counters/histograms from the Prometheus output, so a metric only appears once it has been observed at least once. + +### Labels + +Transfer and attempt metrics carry a `transport` label (the +`tent_transport_failover_total` counter uses `from` and `to` instead) so they +can be sliced by transport without grepping logs. Attempt metrics also carry +an `operation` label. Label values come exclusively from the `TransportType` +enum closed set — no arbitrary transport strings are accepted. + +| Label | Values | Description | +|-------|--------|-------------| +| `transport` | `unspec`, `rdma`, `mnnvl`, `shm`, `nvlink`, `gds`, `io_uring`, `tcp`, `ascend`, `sunrise_link`, `tpu` | The transport that handled the transfer | +| `operation` | `read`, `write` | Attempt operation | +| `from` | (same set) | Transport that failed before failover | +| `to` | (same set) | Transport that the failover switched to | + +Transport label values come from the shared `transportTypeName()` mapping. +`unspec` covers transfers that failed before a transport was selected. + +**Cardinality**: the `transport` label has 11 values; the failover +`from`/`to` pair has at most 11x11 = 121 combinations (in practice only a +few pairs ever occur), and each attempt metric has at most 11x2 = 22 +transport/operation combinations. Total series across all metrics is bounded +at ~1500. ## Integration with TransferEngine @@ -254,8 +291,22 @@ Metrics are automatically recorded at the TENT layer: - **Latency tracking**: Start time is recorded when `submitTransfer` is called - **Metrics recording**: When `getTransferStatus` detects task completion, latency is calculated and metrics are recorded +- **Attempt tracking**: Each concrete `Transport::submitTransferTasks()` call is counted as one attempt. A failed attempt is closed before failover changes the task's current transport, and the replacement attempt gets a fresh attempt timestamp. Synchronous submit failures are also closed as failed attempts. The transport is captured when the attempt starts, so it is attributed correctly even if failover overwrites the task's current transport afterwards. Staging is an orchestration step, not a transport attempt: `ProxyManager` chunks the transfer and issues the real transport submissions, which are the ones counted, so a staged transfer is not double-counted. + +This provides two complementary views: + +- Logical request latency and outcome, attributed to the final transport for backward compatibility. +- Physical attempt count, failure count, and latency, attributed to the transport that actually executed that attempt. -This provides end-to-end latency measurement across all transport types (RDMA, TCP, NVLink, etc.). +For a recovered RDMA→TCP request, request metrics record one successful TCP +outcome, while attempt metrics record one failed RDMA attempt and one successful +TCP attempt. The causal-chain stage metrics (`tent_stage_queue_wait_us`, +`tent_stage_dispatch_us`, `tent_stage_transport_us`) are unchanged by this +addition: they remain attributed to the final transport and +`tent_stage_transport_us` still spans the whole request, so existing dashboards +keep their meaning. Use `tent_transport_attempt_latency_us` (labeled by the +transport that actually ran each attempt) to inspect per-attempt latency in a +multi-attempt request. **Note**: Remember to build with `-DTENT_METRICS_ENABLED=ON` to enable metrics collection. @@ -289,11 +340,11 @@ void TentMetrics::registerMetrics() { // ... existing counters ... &new_counter_, // Add new counter here }; - + histograms_ = { - &read_latency_, + {&read_latency_, &kLatencyBuckets}, // ... existing histograms ... - &new_histogram_, // Add new histogram here + {&new_histogram_, &kNewBuckets}, // Add new histogram + its buckets here }; } ``` @@ -325,22 +376,21 @@ No changes to `getPrometheusMetrics()` or `getJsonMetrics()` are required. ## Advanced Configuration -### Custom Buckets +### Histogram Buckets -Define custom histogram buckets for specific use cases: +Histogram bucket boundaries are fixed at compile time (defined as `static inline const std::vector` members in `tent_metrics.h`: `kLatencyBuckets`, `kSizeBuckets`, `kMluPerMilleBuckets`, `kStageBuckets`). They are intentionally not runtime-configurable so that observability is reproducible across deployments. To change buckets, edit the constant in `tent_metrics.h` and rebuild. -```cpp -// Latency buckets (in seconds, converted to microseconds internally) -std::vector rdma_latency_buckets = { - 0.000001, 0.000005, 0.00001, 0.00005, 0.0001, // 1-100μs - 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1 // 0.5-100ms -}; - -// Size buckets for different data patterns (in bytes) -std::vector message_size_buckets = { - 64, 256, 1024, 4096, 16384, 65536, 262144 // 64B to 256KB -}; -``` +### Compatibility + +The `metrics/latency_buckets` and `metrics/size_buckets` config keys, and the `TENT_METRICS_LATENCY_BUCKETS` / `TENT_METRICS_SIZE_BUCKETS` environment variables, were removed and are now silently ignored. Histogram buckets are fixed at compile time (see Histogram Buckets above). + +**Migration:** remove these keys from your `transfer-engine.json` and unset the environment variables. If custom bucket boundaries are required, edit `kLatencyBuckets` / `kSizeBuckets` in `tent/metrics/tent_metrics.h` and rebuild. + +> **Note:** the previous runtime-configurable implementation had a latent bug — the JSON output labeled custom buckets with the compile-time default boundaries, so `/metrics/json` was already mislabeled for custom-bucket deployments. Removing the knob fixes this rather than preserving a buggy code path. + +The `metrics/enable_prometheus` and `metrics/enable_json` config keys, and the `TENT_METRICS_ENABLE_PROMETHEUS` / `TENT_METRICS_ENABLE_JSON` environment variables, were removed and are now silently ignored. The `/metrics` and `/metrics/json` HTTP endpoints are now registered unconditionally and cannot be toggled independently — both are read-only and served on the same port, so there was no real scenario where a deployment wanted one disabled. + +**Migration:** remove these keys from your `transfer-engine.json` and unset the environment variables. To fully disable the metrics subsystem (no HTTP server, no periodic summary logging, no recording), set `metrics/enabled` to `false` (or `TENT_METRICS_ENABLED=false`). Note that "log-only mode" — where periodic summaries are still logged but `/metrics` is unavailable — only occurs when metrics are enabled but the HTTP port cannot be bound (e.g. port conflict); it is not triggered by `metrics/enabled=false`. ### Validation @@ -370,19 +420,46 @@ scrape_configs: ### Grafana Queries ```promql -# Transfer throughput (MB/s) +# Transfer throughput (MB/s) — all transports rate(tent_read_bytes_total[5m]) / 1024 / 1024 rate(tent_write_bytes_total[5m]) / 1024 / 1024 +# Transfer throughput (MB/s) — per transport +rate(tent_read_bytes_total{transport="rdma"}[5m]) / 1024 / 1024 +rate(tent_read_bytes_total{transport="tcp"}[5m]) / 1024 / 1024 + # Request rate rate(tent_read_requests_total[5m]) rate(tent_write_requests_total[5m]) -# Failure rate +# Failure rate (all transports) rate(tent_read_failures_total[5m]) / rate(tent_read_requests_total[5m]) +# Failure rate (per transport) +rate(tent_read_failures_total{transport="tcp"}[5m]) / rate(tent_read_requests_total{transport="tcp"}[5m]) + # P99 latency (note: latency is in microseconds, convert to seconds for display) histogram_quantile(0.99, rate(tent_read_latency_us_bucket[5m])) / 1000000 histogram_quantile(0.99, rate(tent_write_latency_us_bucket[5m])) / 1000000 -``` +# P99 latency per transport +histogram_quantile(0.99, rate(tent_read_latency_us_bucket{transport="rdma"}[5m])) / 1000000 + +# Failover rate by transport pair +rate(tent_transport_failover_total{from="rdma",to="tcp"}[5m]) + +# RDMA physical-attempt failure rate +sum(rate(tent_transport_attempt_failures_total{transport="rdma"}[5m])) +/ +sum(rate(tent_transport_attempts_total{transport="rdma"}[5m])) + +# P99 latency of RDMA write attempts +histogram_quantile( + 0.99, + sum by (le) ( + rate(tent_transport_attempt_latency_us_bucket{ + transport="rdma",operation="write" + }[5m]) + ) +) / 1000000 +``` diff --git a/docs/source/design/tent/overview.md b/docs/source/design/tent/overview.md index 3288743588..9559622614 100644 --- a/docs/source/design/tent/overview.md +++ b/docs/source/design/tent/overview.md @@ -73,9 +73,9 @@ TENT extends the classic Mooncake Transfer Engine by moving transport selection, The design favors predictable behavior and operational simplicity over manual tuning and static configuration. -## TENT C++ API Reference +## TENT C++ API Reference -[TENT C++ API Reference](cpp-api.md) +[TENT C++ API Reference](../../api-reference/cpp/tent.md) ## TENT Metrics System diff --git a/docs/source/design/tent/qos.md b/docs/source/design/tent/qos.md index 06ecfc27db..e5fc0e0952 100644 --- a/docs/source/design/tent/qos.md +++ b/docs/source/design/tent/qos.md @@ -62,6 +62,20 @@ This ensures that: - High-priority requests normally never wait behind lower-priority work - Low-priority requests eventually get serviced even under continuous high-priority load +**Promotion configuration**: + +| Config key | Default | Behavior | +|---|---|---| +| `transports/rdma/priority_promotion_timeout_us` | `10000` (10ms) | How long an entry may wait before it is eligible for promotion. | +| `transports/rdma/priority_promotion_per_entry` | `false` | Selects the promotion policy (see below). | + +`priority_promotion_per_entry` controls *which* entries a promotion pass moves up: + +- **`false` (default, head-only)**: a pass inspects only the queue head; if the head has timed out, the whole queue is promoted one level. This is the original, lowest-overhead "flush the tier" behavior — coarse, but it never scans the queue. +- **`true` (per-entry)**: a pass promotes exactly the entries that have themselves timed out, leaving freshly enqueued entries in place, and considers both MEDIUM→HIGH and LOW→MEDIUM in the same tick. This avoids promoting non-starving requests and avoids stalling a starving LOW entry behind an unrelated MEDIUM promotion, at the cost of scanning the drained queue. Behavior is identical to head-only for the all-timed-out / empty cases. + +The default is byte-for-byte the historical behavior; set the flag to `true` to opt into finer-grained, fairer promotion. + ### Global Slot Coordination For multi-process environments, TENT implements global time-sliced coordination using shared memory: @@ -70,7 +84,7 @@ For multi-process environments, TENT implements global time-sliced coordination Time slices rotate every N milliseconds: Slot 0 (0-Nms): Only HIGH priority requests allowed -Slot 1 (N-2Nms): MEDIUM + HIGH priority requests allowed +Slot 1 (N-2Nms): MEDIUM + HIGH priority requests allowed Slot 2 (2N-3Nms): All priorities allowed ...repeats... ``` @@ -393,4 +407,4 @@ Monitor queue depths for each priority level (requires instrumentation). - [TENT Overview](overview.md) - [TENT Slice Spraying](slice-spraying.md) -- [TENT C++ API](cpp-api.md) +- [TENT C++ API](../../api-reference/cpp/tent.md) diff --git a/docs/source/design/tent/slice-spraying.md b/docs/source/design/tent/slice-spraying.md index 5cf6d3bc75..09355e6f19 100644 --- a/docs/source/design/tent/slice-spraying.md +++ b/docs/source/design/tent/slice-spraying.md @@ -395,4 +395,4 @@ device_selector_->printTrafficStats(); - [TENT Overview](overview.md) - [TENT QoS](qos.md) -- [TENT C++ API](cpp-api.md) +- [TENT C++ API](../../api-reference/cpp/tent.md) diff --git a/docs/source/design/tent/tebench.md b/docs/source/design/tent/tebench.md index c25745b5bd..4e9ce606f7 100644 --- a/docs/source/design/tent/tebench.md +++ b/docs/source/design/tent/tebench.md @@ -91,6 +91,102 @@ Each output row corresponds to one benchmark configuration. A short (~1 second) warmup phase is executed before measurements begin. +### 4.1 QoS Metrics Baseline + +Use `--qos_classes` to partition a fixed number of worker threads into QoS +classes: + +```text +name:threads:slo_us:weight[:isolated_gbps],... +``` + +For readability, the same contract can be supplied as JSON with +`--qos_classes_json`: + +```json +[ + {"name": "foreground", "threads": 4, "slo_us": 1000, "weight": 4, "isolated_gbps": 12.5}, + {"name": "checkpoint", "threads": 12, "slo_us": 0, "weight": 1, "isolated_gbps": 10.0} +] +``` + +Use only one of `--qos_classes` and `--qos_classes_json`. + +For example, the following closed-loop mixed workload assigns four workers to +an SLO-constrained foreground class and twelve workers to a best-effort +checkpoint class: + +```bash +./tebench \ + --target_seg_name= \ + --backend=tent \ + --start_num_threads=16 \ + --max_num_threads=16 \ + --qos_classes=foreground:4:1000:4:12.5,checkpoint:12:0:1:10.0 \ + --qos_link_capacity_gbps=25 \ + --qos_output_jsonl=qos-results.jsonl +``` + +The class thread counts must add up to the fixed `start_num_threads` value. +An `slo_us` of zero marks a best-effort class. The SLO is a reporting threshold: +QoS baseline mode measures whether each completed transfer meets it, without +changing request scheduling policy on either backend. + +The human-readable summary and the optional versioned JSONL record report: + +| Metric | Definition | +| ------ | ---------- | +| `slo_attainment` | Fraction of completed batches whose measured end-to-end transfer time is at most `slo_us` | +| `p99_us` | P99 end-to-end batch transfer latency for the class | +| `goodput_gbps` | Class throughput multiplied by SLO attainment; best-effort classes use attainment 1 | +| `weighted_goodput_gbps` | Sum of `weight × goodput_gbps` | +| `jain_fairness` | Jain index over per-class `throughput_gbps / weight` | +| `isolation_leakage` | `max(0, 1 - mixed_throughput / isolated_throughput)` | +| `total_utilization` | Aggregate measured throughput divided by `qos_link_capacity_gbps` | + +Isolation leakage requires a matching class-only baseline, supplied as the +optional fifth class field. Total utilization requires +`--qos_link_capacity_gbps`. Missing baselines are emitted as `N/A` in text and +`null` in JSON rather than being inferred from the mixed run. Run isolated and +mixed cases with the same block size, batch size, transport, memory type, and +host pair. JSONL records retain `isolated_throughput_gbps` and +`link_capacity_gbps` alongside the derived values so every metric can be +recomputed from one record. + +### 4.2 Mixed Foreground/Background Traffic + +The QoS class flags above use one block size, batch size, and request intent for +all workers. Use `--workload_classes_json` when classes must generate different +transfer shapes concurrently, for example latency-sensitive reads competing +with a bulk KV migration: + +```bash +./tebench \ + --target_seg_name= \ + --backend=tent \ + --start_num_threads=8 \ + --max_num_threads=8 \ + --workload_classes_json='[ + {"name":"foreground","threads":2,"block_size":4096,"batch_size":1, + "intent_type":"foreground_get","deadline_us":250,"slo_us":300,"weight":4}, + {"name":"migration","threads":6,"block_size":4194304,"batch_size":2, + "intent_type":"migration","slo_us":0,"weight":1} + ]' \ + --qos_link_capacity_gbps=25 \ + --qos_output_jsonl=mixed-traffic.jsonl +``` + +This mode runs one fixed workload instead of sweeping the global block and +batch size flags. Class thread counts must add up to the fixed benchmark thread +count. `deadline_us` is optional and controls TENT deadline tagging; +`slo_us` remains the independent reporting threshold. Per-class +`transferred_bytes` is included in text and JSONL output so throughput is +computed from each class's actual transfer size. + +`--workload_classes_json` is mutually exclusive with `--qos_classes`, +`--qos_classes_json`, and the global `--tent_intent_type`. Non-default +per-class intents and deadlines require the TENT backend. + ## 5. Runtime Configuration This section summarizes the key runtime options that control workload behavior, @@ -147,6 +243,13 @@ Example: * `--seg_name` : local segment name (typically left empty) * `--seg_type` : `DRAM | VRAM` (default: `DRAM`) +* `--seg_type_mix` : comma-separated segment types for mixed DRAM+VRAM runs, + e.g. `dram,vram`. When set, the target registers buffers of each listed + type in one segment, and initiator threads round-robin across them so a + single tebench process drives traffic over multiple memory types (and + thus multiple transports — SHM for DRAM, NVLink for VRAM) concurrently. + Empty falls back to `--seg_type` (single type, existing behavior). See + Section 5.8 for usage and the multi-transport configuration it requires. * `--target_seg_name` : target segment name (empty → Target mode) **Scan ranges** @@ -186,9 +289,85 @@ gpu_id + thread_id **Transport (TENT only)** -* `--xport_type` : `rdma | shm | mnnvl | gds | iouring` +* `--xport_type` : `rdma | shm | mnnvl | gds | iouring`. Selects a single + transport to enable (all others are disabled). Empty means no transport + is explicitly enabled or disabled by tebench — the engine reads the + transport enable list from the `MC_TENT_CONF` config file (see Section + 5.8 for multi-transport scenarios). +* `--tent_intent_type` : attach a standard transfer intent to every request, + such as `foreground_get`, `background_prefetch`, or `checkpoint`. This is + useful for validating intent-specific transport and QoS policy selection. **Metadata service** * `--metadata_type` : `p2p | etcd | redis | http` (default: `p2p`) * `--metadata_url_list` : comma-separated URLs (ignored in `p2p` mode) + +### 5.7 QoS Reporting + +* `--qos_classes` : class/thread/SLO/weight contract described in Section 4.1 +* `--workload_classes_json` : class-specific transfer shape, intent, deadline, + and QoS contract described in Section 4.2 +* `--qos_link_capacity_gbps` : measured usable link capacity in decimal GB/s +* `--qos_output_jsonl` : append one schema-versioned JSON object per benchmark + configuration + +QoS mode intentionally requires a fixed thread count. Sweep offered load by +running explicit cases with different class thread allocations so every output +record has an unambiguous workload contract. + +### 5.8 Mixed DRAM+VRAM and Multi-Transport + +By default, tebench allocates buffers of a single `--seg_type` and enables a +single `--xport_type` per run. To exercise a mixed-transport workload where +the engine's transport selector naturally picks SHM for DRAM→DRAM transfers +and NVLink for VRAM→VRAM transfers within one process, combine +`--seg_type_mix` with the `MC_TENT_CONF` environment variable: + +1. **Enable multiple transports via `MC_TENT_CONF`** — point it at a JSON + config file (or inline JSON string) that enables the transports you want + active: + + ```json + { + "transports": { + "shm": {"enable": true}, + "nvlink": {"enable": true} + } + } + ``` + + Leave `--xport_type` empty so tebench does not override the config's + enable list. `MC_TENT_CONF` is read by the engine's `ConfigHelper` and + applies to both target and initiator. + +2. **Register a mixed DRAM+VRAM segment via `--seg_type_mix`**: + + ```bash + # Target: mixed DRAM+VRAM segment, all GPUs + MC_TENT_CONF=/path/to/config.json ./tebench --backend=tent \ + --metadata_type=p2p --seg_name=tgtmix --seg_type_mix=dram,vram \ + --local_gpu_id=-1 & + + # Initiator: mixed DRAM+VRAM, 8 threads (4 DRAM→SHM, 4 VRAM→NVLink) + MC_TENT_CONF=/path/to/config.json ./tebench --backend=tent \ + --metadata_type=p2p --target_seg_name= --seg_type_mix=dram,vram \ + --local_gpu_id=-1 --op_type=write --duration=30 \ + --start_num_threads=8 --max_num_threads=8 + ``` + + Initiator threads round-robin across the listed seg_types by `thread_id` + (even→DRAM, odd→VRAM). The engine selects transport based on the target + buffer's location — SHM for `cpu:N` buffers, NVLink for `cuda:N` buffers — + so `/metrics` will show both `transport="shm"` and `transport="nvlink"` + labels with non-zero counts from a single metrics endpoint. + +3. **`--local_gpu_id=-1`** is recommended for VRAM so each worker thread + gets its own GPU buffer. With the default `--local_gpu_id=0`, all VRAM + threads share one GPU buffer and contend, which underrepresents NVLink + throughput. + +Mixed mode requires `--start_num_threads` ≥ 2 (at least one thread per +seg_type) and a build with `USE_CUDA=ON` (for VRAM/NVLink support). When +`--seg_type_mix` is empty, tebench falls back to `--seg_type` (single type) +and `--xport_type` (single transport) — the existing behavior. diff --git a/docs/source/design/tent/transport-selector.md b/docs/source/design/tent/transport-selector.md index d1a90e0877..e4ba13e2ec 100644 --- a/docs/source/design/tent/transport-selector.md +++ b/docs/source/design/tent/transport-selector.md @@ -23,8 +23,9 @@ Transport selection is driven by configuration with pattern-based rules. { "policy": [ { - "name": "high_prio_fast", + "name": "foreground_get", "segment_type": "memory", + "intent_type": "foreground_get", "priority": "high", "devices": ["mlx5_0", "mlx5_1", "mlx5_2"], "transports": ["nvlink", "rdma", "shm"] @@ -52,9 +53,46 @@ Transport selection is driven by configuration with pattern-based rules. | `name` | string | Yes | Policy identifier (for logging) | | `segment_type` | string | Yes | `"memory"` or `"file"` | | `priority` | string or int | No | Match only requests with this priority: `"high"` (0), `"medium"` (1), `"low"` (2) | +| `intent_type` | string or int | No | Match a standard transfer intent such as `"foreground_get"`, `"background_prefetch"`, `"migration"`, `"checkpoint"`, `"weight_loading"`, or `"staging_internal"` | | `devices` | array[string] | No | List of allowed device names (empty = all devices) | | `transports` | array[string] | No | Transport preference list (evaluated in order) | +### Intent-Based Policy Binding + +`Request::intent_type` can select an intent-specific policy before transport, +device, QP-pool, and SL/TC resolution: + +```json +{ + "policy": [ + { + "name": "foreground-kv", + "segment_type": "memory", + "intent_type": "foreground_get", + "qp_pool": "foreground", + "service_level": 3, + "traffic_class": 96, + "transports": ["rdma"] + }, + { + "name": "memory-fallback", + "segment_type": "memory", + "transports": ["rdma", "tcp"] + } + ] +} +``` + +Policies are evaluated in JSON order, so intent-specific entries should appear +before a catch-all entry. A policy without `intent_type` retains the historical +behavior and matches any intent. `INTENT_UNSPEC` therefore behaves exactly as +before with existing configurations. + +An explicit `Request::policy_name` remains the strongest per-request override: +it selects the named policy by segment type and bypasses the policy's other +match filters, including `intent_type`. Invalid intent values cause that policy +entry to be skipped rather than silently converted into a catch-all rule. + ### Memory Type Filters For `memory` segments, you can filter by source/destination memory type: @@ -228,6 +266,7 @@ TransportSelector.select(context, transports, transport_index) ↓ Match policy by: - segment_type (file/memory) + - intent_type (exact match if specified in policy) - priority (exact match if specified in policy) - location constraints - size constraints diff --git a/docs/source/design/transfer-engine/efa_transport.md b/docs/source/design/transfer-engine/efa_transport.md index cb9bd96f0f..b624485646 100644 --- a/docs/source/design/transfer-engine/efa_transport.md +++ b/docs/source/design/transfer-engine/efa_transport.md @@ -4,6 +4,7 @@ This document describes how to build and use Mooncake with AWS Elastic Fabric Ad ## Prerequisites +(efa-prerequisites-driver)= ### 1. AWS EFA Driver and libfabric EFA driver and libfabric should be pre-installed on AWS instances with EFA support (e.g., p6-b300.48xlarge, p6-b200.48xlarge, p5en.48xlarge, p5e.48xlarge, p5.48xlarge). @@ -34,6 +35,27 @@ This installs all system packages, git submodules (including pybind11 and yalant > **Note:** The EFA driver and libfabric are **not** installed by `dependencies.sh`. They must be pre-installed on the instance (see section 1 above). +## Installing from PyPI (recommended) + +Pre-built EFA wheels are published to PyPI by the official release pipeline, so most users do not need to build from source. The EFA transport's memory path is CUDA-aware, so separate CUDA 12, CUDA 13, and non-CUDA variants are published: + +```bash +# GPU memory transfers with CUDA 12 — built with USE_CUDA=ON +pip install mooncake-transfer-engine-efa + +# GPU memory transfers with CUDA 13 — built with USE_CUDA=ON +pip install mooncake-transfer-engine-efa-cuda13 + +# CPU/DRAM-only transfers — built with USE_CUDA=OFF +pip install mooncake-transfer-engine-efa-non-cuda +``` + +The CUDA 13 wheel requires an NVIDIA 580-series or newer driver, following the [CUDA compatibility requirements](https://docs.nvidia.com/deploy/cuda-compatibility/minor-version-compatibility.html). + +> **Note:** These wheels deliberately do **not** bundle `libfabric`/`libefa` (see the runtime note in [Building a Distributable Wheel](#efa-distributable-wheel)). They resolve to the system AWS EFA installation at runtime, so the EFA driver and libfabric from the [Prerequisites](#efa-prerequisites-driver) must still be present on the instance. Make sure `/opt/amazon/efa/lib` is on `LD_LIBRARY_PATH`. + +To build from source instead (for development, an unreleased revision, or a custom configuration), follow the sections below. + ## Building Mooncake with EFA Support ### 1. Build with EFA Enabled @@ -80,11 +102,10 @@ cp mooncake-common/libasio.so ../mooncake-wheel/mooncake/ pip install -e ../mooncake-wheel --no-build-isolation ``` +(efa-distributable-wheel)= ### 3. Building a Distributable Wheel (optional) -To produce a relocatable wheel for distribution (instead of the editable -install above), use `scripts/build_wheel.sh`, which runs `auditwheel -repair` to bundle non-system dependencies: +To produce a relocatable wheel for distribution (instead of the editable install above), use `scripts/build_wheel.sh`, which runs `auditwheel repair` to bundle non-system dependencies: ```bash # After the cmake/make build above completes: @@ -92,21 +113,22 @@ PYTHON_VERSION=3.13 BUILD_DIR=build bash scripts/build_wheel.sh 3.13 dist pip install dist/mooncake_transfer_engine-*.whl ``` -> **Important (EFA builds):** `auditwheel repair` excludes `libfabric` -> and `libefa` from the wheel so they resolve to the system EFA -> installation (`/opt/amazon/efa/lib`) at runtime. This is required -> because the in-process `aws-ofi-nccl` plugin (loaded by NCCL) links the -> **same** system `libfabric`. If the wheel bundled its own copy, the -> process would load two independent libfabric instances — Mooncake's -> bundled one and NCCL's system one — and whichever initializes first -> claims the EFA device, leaving the other with an empty provider list -> (`fi_getinfo: provider efa output empty list`). NCCL then silently -> falls back to the TCP provider and cross-node collectives such as -> `all_gather_object` hang. Excluding libfabric/libefa (see -> `scripts/build_wheel.sh`) keeps a single shared libfabric in the -> process. If you are on an older Mooncake build whose wheel still bundles -> libfabric, force the system copy with -> `export LD_PRELOAD=/opt/amazon/efa/lib/libfabric.so.1` as a workaround. +To produce a wheel whose package name matches one of the published variants, set the corresponding build-variant environment variable — this is exactly what the release pipeline does: + +```bash +# GPU build (cmake was configured with USE_CUDA=ON): +EFA_BUILD=1 PYTHON_VERSION=3.13 BUILD_DIR=build bash scripts/build_wheel.sh 3.13 dist + +# CUDA 13 GPU build (cmake was configured with CUDA 13 and USE_CUDA=ON): +EFA_CU13_BUILD=1 PYTHON_VERSION=3.13 BUILD_DIR=build bash scripts/build_wheel.sh 3.13 dist + +# CPU build (cmake was configured with USE_CUDA=OFF): +EFA_NON_CUDA_BUILD=1 PYTHON_VERSION=3.13 BUILD_DIR=build bash scripts/build_wheel.sh 3.13 dist +``` + +> **CI/CD:** EFA wheels are built and published automatically — see `.github/workflows/ci_efa.yml` (per-PR build validation), `.github/workflows/release-efa.yaml` (CUDA 12 release), `.github/workflows/release-efa-cuda13.yaml` (CUDA 13 release), and `.github/workflows/release-efa-non-cuda.yaml` (non-CUDA release). No EFA hardware is required to *build* the wheel: only the libfabric headers/library are needed to compile and link, which the CI runner obtains from the distro `libfabric-dev` package. + +> **Important (EFA builds):** `auditwheel repair` excludes `libfabric` and `libefa` from the wheel so they resolve to the system EFA installation (`/opt/amazon/efa/lib`) at runtime. This is required because the in-process `aws-ofi-nccl` plugin (loaded by NCCL) links the **same** system `libfabric`. If the wheel bundled its own copy, the process would load two independent libfabric instances — Mooncake's bundled one and NCCL's system one — and whichever initializes first claims the EFA device, leaving the other with an empty provider list (`fi_getinfo: provider efa output empty list`). NCCL then silently falls back to the TCP provider and cross-node collectives such as `all_gather_object` hang. Excluding libfabric/libefa (see `scripts/build_wheel.sh`) keeps a single shared libfabric in the process. If you are on an older Mooncake build whose wheel still bundles libfabric, force the system copy with `export LD_PRELOAD=/opt/amazon/efa/lib/libfabric.so.1` as a workaround. ## Verification @@ -123,6 +145,13 @@ print(f'Initialize result: {result}') # Should be 0 # EFA device (libfabric): rdmap79s0, domain: rdmap79s0-rdm, fabric: efa, provider: efa ``` +Mooncake Store also auto-discovers topology when `protocol="efa"` and no +device names are supplied. The Store passes the requested protocol to the +Transfer Engine, which uses its normal topology discovery and installs the EFA +transport instead of RDMA. Use `MC_MS_FILTERS` to restrict discovery to a +comma-separated device whitelist. To disable discovery, set +`MC_MS_AUTO_DISC=0` and provide device names explicitly. + ## Unit Tests Run the EFA transport unit tests (requires EFA hardware): @@ -158,18 +187,10 @@ cd build && ctest --output-on-failure -R 'efa' Use `transfer_engine_bench` to measure EFA transport throughput between two nodes. -The following commands are the GPU-to-GPU configuration that produces -the headline numbers in the [Benchmark Results](#benchmark-results) -tables (≈ 350 GB/s write on a p5en.48xlarge pair, ≈ 302 GB/s on -p6-b200.48xlarge). Two things matter the most: +The following commands are the GPU-to-GPU configuration that produces the headline numbers in the [Benchmark Results](#benchmark-results) tables (≈ 350 GB/s write on a p5en.48xlarge pair, ≈ 302 GB/s on p6-b200.48xlarge). Two things matter the most: -- `--gpu_id=-1` on **both** sides — this fans buffers across every GPU, - which in turn lets both NUMA nodes' NICs saturate. Pinning a single - GPU (the default `--gpu_id=0`) halves throughput because half the - NICs end up cross-NUMA. -- `--block_size=1048576` (1MB, not the 64 KB default) — each block - becomes one `fi_write` / `fi_read`, so larger blocks amortize - per-op overhead and are the main knob for hitting line rate. +- `--gpu_id=-1` on **both** sides — this fans buffers across every GPU, which in turn lets both NUMA nodes' NICs saturate. Pinning a single GPU (the default `--gpu_id=0`) halves throughput because half the NICs end up cross-NUMA. +- `--block_size=1048576` (1MB, not the 64 KB default) — each block becomes one `fi_write` / `fi_read`, so larger blocks amortize per-op overhead and are the main knob for hitting line rate. ### 1. Target Node (receiver) @@ -182,9 +203,7 @@ p6-b200.48xlarge). Two things matter the most: --gpu_id=-1 ``` -`--buffer_size` must be at least as large as the initiator's -`--buffer_size` — the initiator writes into offsets `[0, buffer_size)` -on the target, so keep these in sync. +`--buffer_size` must be at least as large as the initiator's `--buffer_size` — the initiator writes into offsets `[0, buffer_size)` on the target, so keep these in sync. ### 2. Initiator Node (sender) @@ -204,12 +223,11 @@ on the target, so keep these in sync. --report_unit=GB ``` -Replace `:` with the target node's -address shown in the target's startup log (e.g., `ip-172-31-29-226:12345`). +Replace `:` with the target node's address shown in the target's startup log (e.g., `ip-172-31-29-226:12345`). > **CPU-to-CPU** (no GPUs): build with `-DUSE_CUDA=OFF`, **or** pass `--use_vram=false` to a CUDA-enabled binary. Drop `--gpu_id=-1` in that case — the bench will spread buffers across NUMA nodes instead. -> **Why `threads=16` and not 32:** the SRD shared endpoint caps outstanding WRs per NIC (default 256 — see `MC_MAX_WR`). With `threads × batch ≤ NICs × max_wr` the CQ never saturates; going higher triggers backoff and times out. 32 threads × 128 batch = 4096 slices chasing 16 × 256 = 4096 WRs has no headroom, so the steady-state config settles at 16 threads. +> **Why `threads=16` and not 32:** measured, on 16-NIC hosts. It used to also be forced by the WR cap — when that cap was a fixed 256, 32 threads × 128 batch = 4096 slices chased 16 × 256 = 4096 WRs with no headroom and runs timed out. The cap is now the provider's real transmit depth (see the WR-cap tip under **Tuning Tips** below), which is far higher, so `threads=16` stands on the measurement alone. ### Key Parameters @@ -231,11 +249,7 @@ address shown in the target's startup log (e.g., `ip-172-31-29-226:12345`). | `--init_mem` | true | Zero-fill the allocated buffer; rarely needs to change | | `--auto_discovery` | false | Auto-discover topology on init; off for reproducible runs | -> **Note on EFA slicing:** EFA transport does not split each transfer -> into fixed-size slices the way RDMA transport does — each transfer -> is sent as a single `fi_write` / `fi_read` whose size equals -> `block_size`, round-robin'd across NICs per request. **`block_size` -> is the key tuning parameter** for EFA throughput. +> **Note on EFA slicing:** EFA transport does not split each transfer into fixed-size slices the way RDMA transport does — each transfer is sent as a single `fi_write` / `fi_read` whose size equals `block_size`, round-robin'd across NICs per request. **`block_size` is the key tuning parameter** for EFA throughput. > **Note:** `buffer_size` must be >= `block_size * batch_size * threads`. The benchmark auto-adjusts if too small. @@ -271,10 +285,7 @@ Tested on two p6-b300.48xlarge instances (Intel Xeon Platinum 8559C, 8× B300, 1 Tested on two p6-b200.48xlarge instances in the same AWS placement group. -> **Note:** numbers below predate the SRD shared-endpoint refactor (#1944) and -> current EFA tuning work. They are a lower bound for the current -> code; we will re-sweep and update when a B200 pair is available -> again. +> **Note:** numbers below predate the SRD shared-endpoint refactor (#1944) and current EFA tuning work. They are a lower bound for the current code; we will re-sweep and update when a B200 pair is available again. **GPU-to-GPU** (build with `-DUSE_CUDA=ON`, `--gpu_id=-1` for all 8 GPUs): @@ -374,20 +385,10 @@ Tested on two p5.48xlarge instances (AMD EPYC 7R13, 8× H100 80GB, 32 EFA device ### Single-host loopback -EFA NICs have no hardware loopback short-circuit: when a transfer's source and -destination resolve to the same host, the data does not go out on the wire as -GPUDirect/device RDMA. libfabric handles the same-host case in software, and -there are **two distinct provider knobs** that select how: - -- **`FI_EFA_ENABLE_SHM_TRANSFER`** (default `1`, on): when on, the EFA - provider routes same-host peers through the **`shm` provider** — verifiable - at runtime, where libfabric reports `Opened fabric: shm` alongside - `Opened fabric: efa` even on a default (device-RDMA-enabled) configuration. - This SHM path is the one that supplies the same-host memcpy fast path; it is - active **by default**, independent of `FI_EFA_USE_DEVICE_RDMA`. -- **`FI_EFA_USE_DEVICE_RDMA`** (default `1` after #2041): controls whether the - EFA RDM data path uses device RDMA vs libfabric's emulated RDM path. It is a - provider-level flag resolved at `fi_getinfo` time; Mooncake does not wrap it. +EFA NICs have no hardware loopback short-circuit: when a transfer's source and destination resolve to the same host, the data does not go out on the wire as GPUDirect/device RDMA. libfabric handles the same-host case in software, and there are **two distinct provider knobs** that select how: + +- **`FI_EFA_ENABLE_SHM_TRANSFER`** (default `1`, on): when on, the EFA provider routes same-host peers through the **`shm` provider** — verifiable at runtime, where libfabric reports `Opened fabric: shm` alongside `Opened fabric: efa` even on a default (device-RDMA-enabled) configuration. This SHM path is the one that supplies the same-host memcpy fast path; it is active **by default**, independent of `FI_EFA_USE_DEVICE_RDMA`. +- **`FI_EFA_USE_DEVICE_RDMA`** (default `1` after #2041): controls whether the EFA RDM data path uses device RDMA vs libfabric's emulated RDM path. It is a provider-level flag resolved at `fi_getinfo` time; Mooncake does not wrap it. ```{warning} **GPU (FI_HMEM_CUDA) buffers — known segfault.** The default same-host **SHM** @@ -411,34 +412,25 @@ transfer — `__memcpy_avx_unaligned` ← `ofi_copy_to_mr_iov` ← `smr_copy_fro then fall back to device RDMA, which is GPU-aware and correct. ``` -For **host (DRAM) buffers** the SHM memcpy path is safe (host→host copy) and is -the same-host fast path the measurements below exercise. +For **host (DRAM) buffers** the SHM memcpy path is safe (host→host copy) and is the same-host fast path the measurements below exercise. -Measured on p5.48xlarge (1 NIC, ~1.2 GiB per `put_from` call, host DRAM buffer, -same-host producer/consumer in **separate processes**): +Measured on p5.48xlarge (1 NIC, ~1.2 GiB per `put_from` call, host DRAM buffer, same-host producer/consumer in **separate processes**): | same-host path | per-write latency | |---|---:| | device RDMA (NIC round-trip, no fast-path for loopback) | ~830 ms | | SHM memcpy fast path (default) | ~390 ms | -For reference, a cross-host `put_from` of the same payload (device RDMA, 1 NIC) -is ~340 ms — i.e., driving a same-host loopback through the NIC is *slower* than -going over the wire to another host, because the NIC has no fast-path for -loopback. Cross-host transfers always use device RDMA and are unaffected by -`FI_EFA_ENABLE_SHM_TRANSFER`: leave it at its default on any process that also -talks to remote peers. +For reference, a cross-host `put_from` of the same payload (device RDMA, 1 NIC) is ~340 ms — i.e., driving a same-host loopback through the NIC is *slower* than going over the wire to another host, because the NIC has no fast-path for loopback. Cross-host transfers always use device RDMA and are unaffected by `FI_EFA_ENABLE_SHM_TRANSFER`: leave it at its default on any process that also talks to remote peers. ### Tuning Tips - **Use `--block_size=1048576` (1MB)** — the single most important knob. The 64 KB default reaches only ~26% of peak. 1 MB is within a few percent of the 2 MB plateau while leaving headroom for `batch_size` under the shared-endpoint WR cap. -- **Keep `threads × batch_size ≤ num_nics × max_wr`** — under the SRD shared endpoint each NIC carries one `fid_ep` with a 256 WR cap (`MC_MAX_WR`), giving `16 NICs × 256 = 4096` in-flight slots on a 16-NIC host (b300, b200, p5en) and `32 NICs × 256 = 8192` on p5. Exceeding this trips "timed out waiting for CQ drain". `threads=16, batch=128` is a solid baseline on 16-NIC hosts; on 32-NIC p5, `threads=32, batch=128` works the same way. +- **Keep `threads × batch_size ≤ num_nics × max_wr`** — under the SRD shared endpoint each NIC carries one `fid_ep` whose WR cap is the provider's transmit queue depth (`fi_tx_attr: size`: 4096 on p5, 2048 on p6-b300). Exceeding it trips "timed out waiting for CQ drain". In practice this is no longer the binding constraint — a 16-NIC b300 allows `16 × 2048 = 32768` in-flight slices, far more than any sane `threads × batch` — so pick `threads` and `batch_size` from the measurements above (`threads=16, batch=128` on 16-NIC hosts; `threads=32, batch=128` on 32-NIC p5) and treat this as a sanity check. Read your host's depth from the `provider tx queue=` field in Mooncake's per-device startup log, or with `fi_info -p efa -t FI_EP_RDM -v | grep -A9 fi_tx_attr`. - **Write vs read:** write benefits from larger batches (peak at `batch=128`); on 16-NIC p5en read prefers smaller queues (peak at `batch=32`), but on 32-NIC p5 reads scale up to `batch=128` because the wider fabric absorbs larger in-flight queues. - For **GPU-to-GPU**: pass `--gpu_id=-1` on **both** sides so buffers fan out across every GPU. Pinning a single GPU halves throughput because half the NICs end up cross-NUMA. - For **CPU-to-CPU**: DRAM bandwidth is the ceiling. NUMA-split (separate initiator/target instances per NUMA node) can help reduce contention when one instance can't saturate both nodes. -- `--buffer_size` only needs `≥ block × batch × threads`; larger - values do not improve throughput. The example commands use 4 GB - because that is safe for any reasonable config. +- `--buffer_size` only needs `≥ block × batch × threads`; larger values do not improve throughput. The example commands use 4 GB because that is safe for any reasonable config. ### First-request latency @@ -461,7 +453,7 @@ The SRD shared-endpoint refactor (#1944) speeds up first-request latency two dif - Rust: `TransferEngine::warmup_efa_segment(name: &str)` - Python: `engine.warmup_efa_segment(segment_name)` - Call once per peer right after `openSegment`. The call is idempotent. Under this refactor `warmupSegment` itself is ~15× faster than the pre-#1944 code (1.1 s vs 17 s), bounded by the peer's single-threaded handshake RPC daemon (`accept` + JSON parse serialized on one thread), so it scales linearly with the number of fresh NIC pairs. +Call once per peer right after `openSegment`. The call is idempotent. Under this refactor `warmupSegment` itself is ~15× faster than the pre-#1944 code (1.1 s vs 17 s), bounded by the peer's single-threaded handshake RPC daemon (`accept` + JSON parse serialized on one thread), so it scales linearly with the number of fresh NIC pairs. vLLM and SGLang do not currently call `warmupSegment` — they go through the generic `TransferEngine` interface and pick up the 4× cold-submit speedup automatically. The API is there for direct Mooncake callers that want the larger win. @@ -556,12 +548,7 @@ vllm-router --policy round_robin \ --host 0.0.0.0 --port 30000 ``` -> **Do not add `--intra-node-data-parallel-size` here.** The prefill / decode -> instances above are launched with `-tp 8` (pure tensor parallelism, data -> parallel size = 1), so there is no intra-node DP to advertise. Only pass -> `--intra-node-data-parallel-size N` when your instances actually run `N`-way -> data parallelism per node (e.g. you launched them with `--data-parallel-size N`); -> setting it to match `-tp 8` is wrong and will misroute requests. +> **Do not add `--intra-node-data-parallel-size` here.** The prefill / decode instances above are launched with `-tp 8` (pure tensor parallelism, data parallel size = 1), so there is no intra-node DP to advertise. Only pass `--intra-node-data-parallel-size N` when your instances actually run `N`-way data parallelism per node (e.g. you launched them with `--data-parallel-size N`); setting it to match `-tp 8` is wrong and will misroute requests. ## Usage with SGLang @@ -595,16 +582,27 @@ export FI_EFA_USE_DEVICE_RDMA=1 export LD_LIBRARY_PATH=/opt/amazon/efa/lib:$LD_LIBRARY_PATH ``` -> **Note on additional `MC_*` knobs:** `MC_NUM_CQ_PER_CTX`, `MC_MAX_WR`, `MC_MAX_CQE_PER_CTX`, `MC_SLICE_SIZE`, and `MC_EFA_STRIPING_THRESHOLD` are **not** required at typical PD-disagg loads — the SRD shared-endpoint refactor (#1944) makes them redundant up to high concurrency on 1k/1k traffic. Treat them as emergency switches for CQ-overflow or long-running drift symptoms. +> **Do not set `MC_MAX_WR` on EFA — the default is already the correct value.** It caps a counter that paces submission against the endpoint's transmit queue, whose depth the *provider* chooses per device (`fi_tx_attr: size` — **4096** on p5.48xlarge, **2048** on p6-b300.48xlarge, derived from the device's `max_sq_wr`), so the transport reads that depth back instead of using a compiled-in number. +> +> Overriding it desynchronizes the counter from the queue it paces, and **both directions of error are real transfer failures, not just slowdowns** — EFA has no transport-level retransmit, so the affected slices are reported up as `FAILED`: +> +> - **Too small** (e.g. the old fixed default of 256): the counter saturates while the queue is mostly empty, so submitters spin in the credit-wait loop and give up — `timed out waiting for CQ drain (wr_depth=256, max=256)` — with thousands of slots the NIC would have accepted. +> - **Too large** (e.g. `MC_MAX_WR=16384`): the counter hands out credit the queue cannot honor, so `fi_write` refuses with `-FI_EAGAIN` while `wr_depth` sits pinned at the provider's real depth and the rest of the credit is nominally free — `1024 consecutive FI_EAGAIN waves posted nothing`. +> +> A value **below** the provider's depth is still honored, as a deliberate per-NIC throttle; a value above it is clamped with a warning. Verify what took effect from the per-device startup line — no probe needed: +> +> ``` +> EFA device (libfabric): rdmap79s0, ... (shared endpoint, max_wr=4096, provider tx queue=4096, max_cqe=12288) +> ``` +> +> If you genuinely need a deeper queue, raise it at the provider with **`FI_EFA_TX_SIZE`** rather than with `MC_MAX_WR`, and Mooncake will track the new depth on its own — measured on p5, `FI_EFA_TX_SIZE=16384` moves `tx_attr->size` to 16384 and the CQ to 24576, and both counters follow. `MC_MAX_WR` cannot do this: it only moves Mooncake's counter, leaving the hardware queue where it was. -> **`MC_EFA_CQ_THREADS`** — caps the number of CQ polling threads spawned by the EFA transport. Default is `1`, which reaches 99.93% of peak GPU-to-GPU throughput while saving CPU for other workloads. Set to `0` to disable the cap (one poller per EFA context — the legacy behavior). Higher values (e.g., `MC_EFA_CQ_THREADS=4`) are available as an escape hatch for throughput tuning but rarely help in practice. +> **`MC_EFA_CQ_THREADS`** — caps the number of CQ polling threads. The default of `1` reaches 99.93% of peak GPU-to-GPU throughput while leaving CPU for other work, so **leave it alone**. Pollers busy-wait (the CQs are opened `FI_WAIT_NONE`, so there is no descriptor to block on), which means each extra thread burns a full core whether or not completions are arriving — the opposite of what you want in PD-disagg, where the same cores serve the inference loop. Raising it also cannot fix `FI_EAGAIN` symptoms: those mean the provider is refusing new work, not that completions are going unreaped. Set `0` to lift the cap entirely (one poller per EFA device — the pre-#2113 behavior). Values above the device count are ignored; no excess threads are created. > > ```bash -> export MC_EFA_CQ_THREADS=1 # default: single CQ poller (recommended) -> export MC_EFA_CQ_THREADS=0 # disable cap: one poller per EFA context (legacy) +> export MC_EFA_CQ_THREADS=1 # default: single CQ poller +> export MC_EFA_CQ_THREADS=0 # lift cap: one poller per EFA device (legacy) > ``` -> -> If the value exceeds the number of EFA contexts, it is safely ignored (no excess threads are created). ### 3. Prefill Instance @@ -651,18 +649,9 @@ The trailing `8998` after `--prefill` must match the prefill's `--disaggregation ### Why libfabric instead of ibverbs? -AWS EFA exposes an RDMA-capable device through the ibverbs interface, but it does -**not** implement the full ibverbs API. In particular, EFA only supports -**SRD** (Scalable Reliable Datagram) and **UD** (Unreliable Datagram) queue -pairs — it does **not** support the **RC** (Reliable Connection) queue pairs -that Mooncake's RDMA (`rdma`) transport is built on. Attempting to create an RC -QP on an EFA device fails (`EOPNOTSUPP`), and SRD has no one-sided RC-style -`ibv_post_send(RDMA_WRITE)` verb in the public ibverbs API. +AWS EFA exposes an RDMA-capable device through the ibverbs interface, but it does **not** implement the full ibverbs API. In particular, EFA only supports **SRD** (Scalable Reliable Datagram) and **UD** (Unreliable Datagram) queue pairs — it does **not** support the **RC** (Reliable Connection) queue pairs that Mooncake's RDMA (`rdma`) transport is built on. Attempting to create an RC QP on an EFA device fails (`EOPNOTSUPP`), and SRD has no one-sided RC-style `ibv_post_send(RDMA_WRITE)` verb in the public ibverbs API. -The portable way to drive EFA's SRD transport is libfabric, whose EFA provider -exposes SRD through the `FI_EP_RDM` (Reliable Datagram Message) endpoint type and -implements `fi_write` / `fi_read` (one-sided RMA) on top of it. Mooncake's EFA -transport therefore targets libfabric directly rather than ibverbs. +The portable way to drive EFA's SRD transport is libfabric, whose EFA provider exposes SRD through the `FI_EP_RDM` (Reliable Datagram Message) endpoint type and implements `fi_write` / `fi_read` (one-sided RMA) on top of it. Mooncake's EFA transport therefore targets libfabric directly rather than ibverbs. ### EFA Transport Architecture @@ -691,15 +680,7 @@ Under the SRD shared-endpoint model every peer is addressed through one `fid_ep` └───────────────────────────────────────────────────────────┘ ``` -> **Peer-map keying.** `peer_map_` is keyed by the **full** `host:port@nic` -> path, *not* a port-stripped form. Under SGLang DP > 1 each DP worker on a -> peer host is a separate process with its own Mooncake `TransferEngine` and -> its own P2PHANDSHAKE RPC port; they share host + NIC but have distinct EFA -> addresses. Normalizing the port away would collapse every DP worker on that -> host onto one `EfaEndPoint`, so each arriving handshake would look like a -> "peer reconnected" to the previous holder and trigger `fi_av_remove` + -> `fi_av_insert` churn on every KV transfer. Keeping the port in the key costs -> nothing in steady state (the port is stable for a worker's lifetime). +> **Peer-map keying.** `peer_map_` is keyed by the **full** `host:port@nic` path, *not* a port-stripped form. Under SGLang DP > 1 each DP worker on a peer host is a separate process with its own Mooncake `TransferEngine` and its own P2PHANDSHAKE RPC port; they share host + NIC but have distinct EFA addresses. Normalizing the port away would collapse every DP worker on that host onto one `EfaEndPoint`, so each arriving handshake would look like a "peer reconnected" to the previous holder and trigger `fi_av_remove` + `fi_av_insert` churn on every KV transfer. Keeping the port in the key costs nothing in steady state (the port is stable for a worker's lifetime). ### Thread Safety @@ -708,12 +689,7 @@ The EFA transport requests `FI_THREAD_SAFE` at the domain level and guards the s - Multiple submission threads may route slices through the same shared endpoint concurrently. - libfabric's EFA RDM endpoints are not thread-safe for concurrent `fi_write`/`fi_read` even under `FI_THREAD_SAFE` at the domain level — concurrent posts corrupt provider internals and completions silently vanish. -CQ completion queues are polled by dedicated worker threads that run -independently of submission threads. The poller count is `min(MC_EFA_CQ_THREADS, -num_EFA_devices)`; `MC_EFA_CQ_THREADS` defaults to `1`, so a single poller -round-robins every context's CQ (which already reaches ~99.9% of peak — see the -SGLang env-var note above). Set `MC_EFA_CQ_THREADS=0` to lift the cap and spawn -one poller per EFA device (the legacy behavior). +CQs are polled by dedicated worker threads that run independently of submission threads. The poller count is `min(MC_EFA_CQ_THREADS, num_EFA_devices)`, and `MC_EFA_CQ_THREADS` defaults to `1`, so by default a single thread walks every device's CQ in a busy-wait loop (yielding only when a full pass reaped nothing). That already reaches ~99.9% of peak — see the env-var note under *Usage with SGLang* before changing it. ### EFA vs RoCE RDMA @@ -728,12 +704,7 @@ one poller per EFA device (the legacy behavior). | Throughput CPU-to-CPU (16×200G, p5en) | 213 GB/s (tuned) | — | | AWS availability | All EFA-enabled instances | Not available on AWS | -> Mooncake requests libfabric API ≥ 1.18 at `fi_getinfo`, which makes -> `FI_EFA_USE_DEVICE_RDMA=1` the default on every supported EFA generation -> (p5/p5e included). On this path `fi_write` / `fi_read` are hardware-offloaded -> one-sided RMA over SRD — the host CPU is not in the data path. The -> software-emulated RMA path only applies if you explicitly set -> `FI_EFA_USE_DEVICE_RDMA=0`. +> Mooncake requests libfabric API ≥ 1.18 at `fi_getinfo`, which makes `FI_EFA_USE_DEVICE_RDMA=1` the default on every supported EFA generation (p5/p5e included). On this path `fi_write` / `fi_read` are hardware-offloaded one-sided RMA over SRD — the host CPU is not in the data path. The software-emulated RMA path only applies if you explicitly set `FI_EFA_USE_DEVICE_RDMA=0`. ### Supported AWS Instance Types diff --git a/docs/source/design/transfer-engine/index.md b/docs/source/design/transfer-engine/index.md index 9ffbff0272..93d6d40125 100644 --- a/docs/source/design/transfer-engine/index.md +++ b/docs/source/design/transfer-engine/index.md @@ -161,7 +161,7 @@ The following video shows a normal run as described above, with the Target on th ## Transfer Engine C/C++ API Transfer Engine provides interfaces through the `TransferEngine` class (located in `mooncake-transfer-engine/include/transfer_engine.h`), where the specific data transfer functions for different backends are implemented by the `Transport` class, currently supporting `TcpTransport`, `RdmaTransport`, `EfaTransport` (for AWS EFA), `NVMeoFTransport`, `NvlinkTransport` (for NVIDIA GPUs), `IntraNodeNvlinkTransport` (for NVIDIA GPUs), and `HipTransport` (for AMD GPUs). -For a complete C++ API reference, see [Transfer Engine C++ API Reference](cpp-api.md). +For a complete C++ API reference, see [Transfer Engine C++ API Reference](../../api-reference/cpp/transfer-engine.md). ### Data Transfer Transfer Engine provides batch-based read/write transfers between segments (DRAM/VRAM/NVMeof). A typical flow is: register local memory, open a target segment, submit a batch, and poll status. Detailed function signatures and usage are documented in the C++ API reference. @@ -256,7 +256,186 @@ The HTTP server should implement three following RESTful APIs, while the metadat 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). -## Using Transfer Engine to Your Projects +## Using Transfer Engine in Your Projects + +(using-python-interface)= +### Using Python Interface + +For Python applications, use `mooncake.engine.TransferEngine` directly. Most +serving-framework users should configure Mooncake through vLLM or SGLang +instead; this example is for developers who need to call the low-level Transfer +Engine API. + +Install Mooncake from PyPI before running this example. The package selected in +the [Quick Start installation step](../../getting_started/quick-start.md#installation) +also provides the Transfer Engine Python bindings and runtime components. + +#### Runtime prerequisites + +The Python scripts below only import the Python standard library and +`mooncake.engine`, but the Mooncake package still depends on the system runtime +libraries used by the transfer stack. On Ubuntu, install them with: + +```bash +sudo apt-get update && sudo apt-get install -y libcurl4 libibverbs1 rdma-core librdmacm1 libnuma1 liburing2 +``` + +The local TCP example uses `P2PHANDSHAKE`, so it does not require a separate +metadata service. When using RDMA, make sure RDMA drivers, devices, and device +permissions are configured; you may need to run with `sudo` or adjust device +permissions. + +The following two-process TCP example keeps the Python-side dependencies minimal +by using only the Python standard library plus the Mooncake package. + +::::{dropdown} Receiver (`receiver.py`) + +Save the following as `receiver.py`, then run it in the first terminal: + +```python +import ctypes +import json +import socket + +from mooncake.engine import TransferEngine + + +HOSTNAME = "localhost" +METADATA_SERVER = "P2PHANDSHAKE" +PROTOCOL = "tcp" +DEVICE_NAME = "" +BUFFER_SIZE = 1024 * 1024 + + +def main(): + engine = TransferEngine() + engine.initialize(HOSTNAME, METADATA_SERVER, PROTOCOL, DEVICE_NAME) + session_id = f"{HOSTNAME}:{engine.get_rpc_port()}" + + server_buffer = (ctypes.c_uint8 * BUFFER_SIZE)() + server_ptr = ctypes.addressof(server_buffer) + server_len = ctypes.sizeof(server_buffer) + + ret = engine.register_memory(server_ptr, server_len) + if ret != 0: + raise RuntimeError("Mooncake memory registration failed.") + + print(f"Receiver session ID: {session_id}") + print(f"Receiver buffer address: {server_ptr}, length: {server_len}") + + listener = None + try: + listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + listener.bind(("0.0.0.0", 5555)) + listener.listen(1) + + print("Waiting for sender to connect on port 5555...") + conn, _ = listener.accept() + with conn: + payload = { + "session_id": session_id, + "ptr": server_ptr, + "len": server_len, + } + conn.sendall(json.dumps(payload).encode("utf-8") + b"\n") + print("Buffer information sent to sender.") + + input("Press Enter after the sender finishes...") + print(f"First byte in receiver buffer: {server_buffer[0]}") + finally: + ret = engine.unregister_memory(server_ptr) + if ret != 0: + raise RuntimeError("Mooncake memory deregistration failed.") + if listener is not None: + listener.close() + + +if __name__ == "__main__": + main() +``` + +:::: + +::::{dropdown} Sender (`sender.py`) + +Save the following as `sender.py`, then run it in a second terminal: + +```python +import ctypes +import json +import socket + +from mooncake.engine import TransferEngine + + +HOSTNAME = "localhost" +METADATA_SERVER = "P2PHANDSHAKE" +PROTOCOL = "tcp" +DEVICE_NAME = "" +BUFFER_SIZE = 1024 * 1024 + + +def recv_json_line(sock): + chunks = [] + while True: + chunk = sock.recv(4096) + if not chunk: + break + chunks.append(chunk) + if b"\n" in chunk: + break + return json.loads(b"".join(chunks).split(b"\n", 1)[0].decode("utf-8")) + + +def main(): + with socket.create_connection(("localhost", 5555)) as sock: + buffer_info = recv_json_line(sock) + + server_session_id = buffer_info["session_id"] + server_ptr = buffer_info["ptr"] + server_len = buffer_info["len"] + print(f"Receiver session ID: {server_session_id}") + print(f"Receiver buffer address: {server_ptr}, length: {server_len}") + + engine = TransferEngine() + engine.initialize(HOSTNAME, METADATA_SERVER, PROTOCOL, DEVICE_NAME) + session_id = f"{HOSTNAME}:{engine.get_rpc_port()}" + + client_buffer = (ctypes.c_uint8 * BUFFER_SIZE)() + client_ptr = ctypes.addressof(client_buffer) + client_len = ctypes.sizeof(client_buffer) + ctypes.memset(client_ptr, 1, client_len) + + ret = engine.register_memory(client_ptr, client_len) + if ret != 0: + raise RuntimeError("Mooncake memory registration failed.") + + try: + print(f"Sender session ID: {session_id}") + print("Transferring data to receiver...") + ret = engine.transfer_sync_write( + server_session_id, + client_ptr, + server_ptr, + min(client_len, server_len), + ) + if ret < 0: + raise RuntimeError("Transfer failed.") + print("Transfer successful.") + finally: + ret = engine.unregister_memory(client_ptr) + if ret != 0: + raise RuntimeError("Mooncake memory deregistration failed.") + + +if __name__ == "__main__": + main() +``` + +:::: + +For more Python APIs, see [Transfer Engine Python API](../../api-reference/python/transfer-engine.md). ### Using C/C++ Interface After compiling Mooncake Store, you can move the compiled static library file `libtransfer_engine.a` and the C header file `transfer_engine_c.h` into your own project. There is no need to reference other files under `src/transfer_engine`. @@ -276,7 +455,8 @@ For advanced users, TransferEngine provides the following advanced runtime optio - `MC_NUM_COMP_CHANNELS_PER_CTX` The number of Completion Channel created per device instance, default value 1 - `MC_IB_PORT` The IB port number used per device instance, default value 1 - `MC_IB_TC` Adjust RDMA NIC Traffic Class when switch/NIC defaults differ or for traffic planning. Default value -1 -- `MC_IB_PCI_RELAXED_ORDERING` Setting the PCIe ordering to relaxed for the network adapter sometimes results in better performance. Can set 1 to enable RO function. Default value 0 +- `MC_IB_SL` Set the InfiniBand Service Level (0-15) of RDMA QPs. The switch maps SL to a Virtual Lane for QoS isolation, e.g. to steer KV-cache traffic into a different VL than Expert-Parallel all-to-all traffic that shares the same NIC. -1 keeps the default (0). Default value -1 +- `MC_IB_PCI_RELAXED_ORDERING` Controls PCIe Relaxed Ordering (RO) for RDMA memory regions. `0`: disabled, `1`: enabled if supported by hardware (default), `2`: auto. Requires `ibv_reg_mr_iova2` (libibverbs ≥ 1.8); falls back to strict ordering if unavailable. - `MC_MLX5_QP_UDP_SPORTS` Comma-separated list of UDP source ports (0-65535) used to override the RoCEv2 UDP source port of each QP, for spreading traffic across different ECMP/LAG paths. QP at index *i* uses `list[i % size]`. Default empty (driver chooses). **Requires** an mlx5 NIC + RoCEv2, and the binary built with `-DUSE_MLX5DV=ON`. Recommend ports in the dynamic range 49152-65535. Example: `MC_MLX5_QP_UDP_SPORTS="49152,49153,49154,49155"` - `MC_MLX5_QP_LAG_PORT_BALANCE` Set to `1` or `true` to enable automatic LAG port balancing across bonded physical ports. QP at index *i* is pinned to port `(i % num_lag_ports) + 1`; the number of LAG ports is queried from hardware via `mlx5dv_query_device` at startup and printed in the device log. If the device is not in LAG mode the setting is a no-op. Default: disabled. **Requires** the binary built with `-DUSE_MLX5DV=ON`. Example: `MC_MLX5_QP_LAG_PORT_BALANCE=1` - `MC_GID_INDEX` The GID index used per device instance, default value 3 (or the maximum value supported by the platform) @@ -285,15 +465,18 @@ For advanced users, TransferEngine provides the following advanced runtime optio - `MC_MAX_EP_PER_CTX` The maximum number of active EndPoint per device instance, default value 65536. **Note:** For versions prior to 0.3.7.post1, the default value is 256, and it cannot be manually set to 65536. The maximum supported value is 65535! - `MC_NUM_QP_PER_EP` The number of QPs per EndPoint, the more the number, the better the fine-grained I/O performance, default value 2 - `MC_MAX_SGE` The maximum number of SGEs supported per QP, default value 4 (or the highest value supported by the platform) -- `MC_MAX_WR` The maximum number of Work Request supported per QP, default value 256 (or the highest value supported by the platform) +- `MC_MAX_WR` The maximum number of Work Request supported per QP, default value 256 (or the highest value supported by the platform). **On EFA this should be left unset:** the transport derives the depth from the libfabric provider's transmit queue (a per-device attribute, so no fixed value is right on every instance type), and an override that exceeds it is clamped. See the EFA transport page. - `MC_MAX_INLINE` The maximum Inline write data volume (bytes) supported per QP, default value 64 (or the highest value supported by the platform) - `MC_MTU` The MTU length used per device instance, can be 512, 1024, 2048, 4096, default value 4096 (or the maximum length supported by the platform) - `MC_WORKERS_PER_CTX` The number of asynchronous worker threads corresponding to each device instance - `MC_SLICE_SIZE` The segmentation granularity of user requests in Transfer Engine - `MC_RETRY_CNT` The maximum number of retries in Transfer Engine +- `MC_TE_FILTERS` Restrict which RDMA NICs the engine discovers and uses, as a comma-separated allow-list of device names (e.g. `mlx5_bond_0,mlx5_bond_1`). Only the listed NICs are kept; all others are ignored. Unset (default) discovers all NICs. This is the **same env var and semantics as the legacy Transfer Engine's device whitelist** (see below), so a single variable scopes NICs across both engines. Useful on multi-NIC / multi-NUMA hosts to keep the engine (and its rail selection) off NICs that are not routable to the peer. +- `MC_TE_FILTERS_EXCLUDE` The deny-list counterpart of `MC_TE_FILTERS`: a comma-separated list of device names to exclude from discovery. Ignored if `MC_TE_FILTERS` is set (allow-list takes precedence). Unset (default) excludes nothing. (New; the legacy engine has an allow-list only.) - `MC_AUTO_GID_MAX_RETRIES` The maximum number of automatic local GID reprobe retries during classic RDMA handshake recovery. Default value 2. Set to 0 to disable automatic GID retry. - `MC_LOG_LEVEL` This option can be set as `TRACE`/`INFO`/`WARNING`/`ERROR` (see [glog doc](https://github.com/google/glog/blob/master/docs/logging.md)), and more detailed logs will be output during runtime - `MC_DISABLE_METACACHE` Disable local meta cache to prevent transfer failure due to dynamic memory registrations, which may downgrades the performance +- `MC_TE_METADATA_REFRESH_INTERVAL_SECONDS` Periodically refresh Transfer Engine metadata-derived local caches. Currently refreshes cached remote segment descriptors from the metadata service. Default value 0 disables background polling; callers may still manually invoke `syncSegmentCache()`. Set a positive interval in seconds when peers may re-register the same segment name after restart and cached descriptors must converge automatically - `MC_HANDSHAKE_LISTEN_BACKLOG` The backlog size of socket listening for handshaking, default value is 128 - `MC_HANDSHAKE_CONNECT_TIMEOUT` Connect timeout in seconds for outbound handshake-port requests (QP handshake, probe, notify, metadata exchange), default value is 5. Bounds the stall when the peer address is unreachable; without it, a connect to an unroutable address (e.g. a removed node) blocks for the kernel's full TCP SYN retry cycle, which can take minutes - `MC_HANDSHAKE_MAX_LENGTH` The maximum handshake message length in bytes for P2P mode. Valid range: 1MB to 128MB. Default value is 1MB (1048576 bytes). Increase this value when using a single RDMA instance with many registered memory buffers (>10,000) to avoid handshake failures. Example: set to 10485760 for 10MB @@ -302,7 +485,11 @@ For advanced users, TransferEngine provides the following advanced runtime optio - `MC_REDIS_DB_INDEX` The database index for Redis storage plugin, must be an integer between 0 and 255. Only takes effect when Redis is specified as the metadata server. If not set or invalid, the default value is 0. - `MC_FRAGMENT_RATIO ` In RdmaTransport::submitTransferTask, if the last data piece after division is ≤ 1/MC_FRAGMENT_RATIO of the block size, it merges with the previous block to reduce overhead. The default value is 4 - `MC_ENABLE_DEST_DEVICE_AFFINITY` Enable device affinity for RDMA performance optimization. When enabled, Transfer Engine will prioritize communication with remote NICs that have the same name as local NICs to reduce QP count and improve network performance in rail-optimized topologies. The default value is false +- `MC_TRACK_RDMA_POSTED_SLICES` Enable RDMA posted-slice tracking for timeout diagnostics. When enabled, CQ timeout logs include stuck transfer groups by peer NIC path, slice count, bytes, oldest post age, and sample addresses. This adds synchronization on the RDMA post and poll hot paths, so it is disabled by default and should be enabled only while diagnosing stuck completions. - `MC_ENABLE_PARALLEL_REG_MR` Control parallel memory region registration across multiple RDMA NICs. Valid values: -1 (auto, default), 0 (disabled), 1 (enabled). When set to -1, parallel registration is automatically enabled when multiple RNICs exist and memory has been pre-touched. Note: If memory hasn't been touched before registration, parallel registration can be slower than sequential registration +- `MC_MAX_CONCURRENT_REG_MR` Cap on how many buffers `registerLocalMemoryBatch` registers concurrently (EFA transport). The default 0 means unbounded — one thread per buffer, the historical behavior. Note the cap is **per process**, so a framework running one `TransferEngine` per TP rank multiplies it by the rank count. Capping can cut registration time substantially when a batch holds many large GPU buffers. Registration is CPU-bound, so a reasonable value is `cores / processes-per-node` — on a 192-core node running 8 ranks, around 16. Oversubscribing costs more than undersubscribing, and the result also depends on the order the caller passes buffers in, so a poorly chosen cap can be slower than unbounded — hence opt-in. +- `MC_EFA_NIC_SELECTION` Which NICs the EFA transport registers a buffer on. `all` (the default) registers every buffer on every NIC. `local` restricts **device** memory to the NICs the topology reports as closest to that GPU, which on p5.48xlarge is the 4 EFA devices sharing the GPU's PCIe root complex. Because EFA charges device-memory registration in proportion to the device bytes already registered on the same libfabric domain, narrowing the NIC set cuts registration time by close to the fan-out ratio. Set it when registering many GPU buffers is a startup bottleneck; it is opt-in because fewer NICs can serve a transfer touching that buffer, so a job whose working set sits behind a single GPU is capped at that rail group's bandwidth rather than the node's. Host memory is unaffected. Combine with `MC_MAX_CONCURRENT_REG_MR`, which bounds a different variable: this reduces the cost per registration, that one reduces how many run at once. +- `MC_EFA_CQ_THREADS` Cap on the number of CQ polling threads in the EFA transport, default value 1 (which already reaches ~99.9% of peak throughput). Pollers busy-wait, so each extra thread costs a full core. Set 0 to lift the cap (one poller per EFA device). Values above the device count are ignored - `MC_FORCE_HCA` Force to use RDMA as the active transport, return error if no HCA has been found. - `MC_FORCE_MNNVL` Force to use Multi-Node NVLink as the active transport regardless whether RDMA devices are installed. - `MC_INTRA_NVLINK` Enable intra-node NVLINK transport, and cannot be used together with MC_FORCE_MNNVL. @@ -315,6 +502,7 @@ For advanced users, TransferEngine provides the following advanced runtime optio - `MC_ENDPOINT_STORE_TYPE` Choose FIFO Endpoint Store (`FIFO`) or Sieve Endpoint Store (`SIEVE`), default is `SIEVE`. - `MC_TCP_ENABLE_CONNECTION_POOL` Enable TCP Connection Pool to avoid excessive sockets. - `MC_TCP_SLICE_SIZE` The segmentation granularity (in bytes) of TCP transport for splitting large transfers into socket read/write operations. Corresponds to `MC_SLICE_SIZE` for RDMA. Default value 65536 (64KB). +- `MC_TCP_PROTO` When set to `1`, TCP initiators use the legacy unacknowledged framing even against servers that support acknowledged framing (protocol v2). Under v2 (the default against v2-capable servers), a WRITE completes only after the receiver confirms the payload has been applied to destination memory, and server-side rejections surface as failed transfers instead of silent data loss. Use this variable only as a rollback escape hatch during mixed-version upgrades. ## C++ API Reference diff --git a/docs/source/getting_started/build.md b/docs/source/getting_started/build.md index 47c31bcb15..bab4632376 100644 --- a/docs/source/getting_started/build.md +++ b/docs/source/getting_started/build.md @@ -2,56 +2,68 @@ This document describes how to build Mooncake. -## PyPI Package -Install the Mooncake Transfer Engine package from PyPI, which includes both Mooncake Transfer Engine and Mooncake Store Python bindings: +## Build From Source + +### Recommended Version +- OS: Ubuntu 22.04 LTS+ +- cmake: 3.20.x +- gcc: 9.4+ + +### Default Build + +Install common build dependencies first. A stable Internet connection is +required because the script installs system packages, initializes submodules, +installs Go, and builds/installs yalantinglibs. -**For CUDA-enabled systems:** ```bash -pip install mooncake-transfer-engine +sudo bash dependencies.sh ``` -📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine/](https://pypi.org/project/mooncake-transfer-engine/) -**For non-CUDA systems:** +Then build and install Mooncake: + ```bash -pip install mooncake-transfer-engine-non-cuda +mkdir build +cd build +cmake .. +make -j +sudo make install ``` -📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine-non-cuda/](https://pypi.org/project/mooncake-transfer-engine-non-cuda/) -> **Note**: The CUDA version includes Mooncake-EP and GPU topology detection, requiring CUDA 12.1+. The non-CUDA version is for environments without CUDA dependencies. -> **Note**: MLU support is currently source-build only. If you need Cambricon MLU memory support, install Neuware and build with `-DUSE_MLU=ON`. +### Build with VRAM Segment -## Automatic Build +To enable VRAM Segment, install CUDA toolkit and build Mooncake with +`USE_VRAM_SEGMENT` enabled: -### Recommended Version -- OS: Ubuntu 22.04 LTS+ -- cmake: 3.20.x -- gcc: 9.4+ +```bash +sudo bash dependencies.sh + +mkdir build +cd build +cmake .. -DUSE_VRAM_SEGMENT=ON +make -j +sudo make install +``` + +If NVLink is available in your environment, you can also enable it +with `-DUSE_INTRA_NVLINK=ON`: + +```bash +sudo bash dependencies.sh + +mkdir build +cd build +cmake .. -DUSE_VRAM_SEGMENT=ON -DUSE_INTRA_NVLINK=ON +make -j +sudo make install +``` -### Steps -1. Install dependencies, stable Internet connection is required: - ```bash - bash dependencies.sh - ``` - -2. In the root directory of this project, run the following commands: - ```bash - mkdir build - cd build - cmake .. - make -j - ``` -3. Install Mooncake python package and mooncake_master executable - ```bash - sudo make install - ``` - -**Build with NVMe-oF SSD Pool** +### Build with NVMe-oF SSD Pool To enable the NVMe-oF SSD pool, install the SPDK dependencies and build Mooncake with `USE_NOF` enabled: ```bash -bash dependencies.sh --with-spdk +sudo bash dependencies.sh --with-spdk mkdir build cd build @@ -63,148 +75,60 @@ sudo make install `-DUSE_NOF=ON` builds the NoF registration APIs and deployment tools. Use `-DUSE_NOF=OFF` or omit the option when the NVMe-oF SSD pool is not needed. -## Manual Build +### Hardware Backend Setup + +Run `sudo bash dependencies.sh` before using any of these backend-specific build +options. The script handles common dependencies; vendor SDKs and runtime +environment setup must be prepared separately. + +| Hardware / backend | Build option | External SDK / setup | Environment and notes | +| --- | --- | --- | --- | +| NVIDIA CUDA / GPUDirect | `-DUSE_CUDA=ON` | Install CUDA 12.1+ and enable `nvidia-fs` for cuFile builds. | Add CUDA libraries to `LIBRARY_PATH` and `LD_LIBRARY_PATH`, for example `/usr/local/cuda/lib64`. | +| NVIDIA NCCL DeviceTransport | `-DUSE_NCCL_DEVICE=ON` | Install NCCL 2.30.4+ with `nccl_device.h`. Requires CUDA. | Set `NCCL_ROOT` when NCCL is outside the standard search paths. NCCL Device API device code must be rebuilt or re-JITed with headers that exactly match the runtime `libnccl`. | +| NVIDIA NCCL host RMA (WRITE only) | `-DUSE_NCCL_HOST=ON` | Install NCCL 2.30.4+. Requires CUDA. | Set `NCCL_ROOT` when NCCL is outside the standard search paths. Install NCCL as the only transport in a `TransferEngine(false)` instance before registering buffers. Peers must register matching buffer sizes in the same order. It has no multi-transport fallback and supports WRITE requests only because NCCL 2.30 has no public host Get operation. | +| NVIDIA Multi-Node NVLink | `-DUSE_MNNVL=ON` | Requires CUDA. | Also set `-DUSE_CUDA=ON`. Not used with MUSA, HIP, or MACA builds. | +| Moore Threads MUSA | `-DUSE_MUSA=ON` | Install MUSA SDK and `mthreads-peermem` for GPUDirect RDMA. | Add `/usr/local/musa/lib` to `LIBRARY_PATH` and `LD_LIBRARY_PATH`. | +| Cambricon MLU | `-DUSE_MLU=ON` | Install Cambricon Neuware SDK. | Set `NEUWARE_HOME`, or pass `-DNEUWARE_ROOT=/path/to/neuware`. Use `-DMLU_INCLUDE_DIR` and `-DMLU_LIB_DIR` for custom layouts. | +| MetaX MACA | `-DUSE_MACA=ON` | Install MACA SDK. | Set `MACA_HOME`, or pass `-DMACA_ROOT=/path/to/maca`. Use `-DMACA_INCLUDE_DIR`, `-DMACA_LIB_DIR`, and `-DMACA_RUNTIME_LIBS` for custom layouts. | +| Huawei Ascend Direct | `-DUSE_ASCEND_DIRECT=ON` | Install Ascend CANN Toolkit and ADXL dependencies. | Source `/usr/local/Ascend/cann/set_env.sh` before configuring CMake. This is the recommended Ascend path. | +| Huawei Ascend UBSHMEM | `-DUSE_UBSHMEM=ON` | Install Ascend CANN Toolkit. Requires CANN >= 9.0.0, driver >= 26.0.0, Lingqu >= 1.5. | Source the CANN `set_env.sh` before configuring CMake. | +| AMD HIP / ROCm | `-DUSE_HIP=ON` | Install ROCm/HIP SDK. | Ensure HIP compiler, headers, and runtime libraries are visible to CMake. | +| Hygon DCU | `-DUSE_HYGON=ON` | Install DTK SDK. | Set `DTK_HOME`, or pass `-DDTK_ROOT=/path/to/dtk`. Use `-DDTK_INCLUDE_DIR` and `-DDTK_LIB_DIR` for custom layouts. | +| Iluvatar CoreX | `-DUSE_COREX=ON` | Install CoreX SDK. | Set `COREX_HOME`, or pass `-DCOREX_ROOT=/path/to/corex`. Use `-DCOREX_INCLUDE_DIR` and `-DCOREX_LIB_DIR` for custom layouts. | + +```{admonition} NCCL host RMA constraints +:class: important +The first valid NCCL host WRITE freezes the ordered CUDA-buffer catalog before +bootstrap. Registration and unregistration are not allowed afterward, even if +bootstrap fails. Session initialization is attempted once for each +endpoint/device pair. If the session reaches a terminal failure, it retains the +error and subsequent transfers fail without retrying bootstrap. Recovery +requires destroying and recreating the NCCL-only `TransferEngine` on both +peers, then registering the buffers again. A one-sided restart is unsupported. +Same-engine targets remain unsupported and should use +the intra-node NVLink/P2P transport. +``` -### Recommended Version -- cmake: 3.22.x -- boost-devel: 1.66.x -- googletest: 1.12.x -- gcc: 10.2.1 -- go: 1.22+ -- hiredis -- curl - -### Steps - -1. Install dependencies from system software repository: - ```bash - # For debian/ubuntu - apt-get install -y build-essential \ - cmake \ - libibverbs-dev \ - libgoogle-glog-dev \ - libgtest-dev \ - libjsoncpp-dev \ - libnuma-dev \ - libunwind-dev \ - libpython3-dev \ - libboost-dev \ - libssl-dev \ - pybind11-dev \ - libcurl4-openssl-dev \ - libhiredis-dev \ - pkg-config \ - patchelf - - # For centos/alibaba linux os - yum install cmake \ - gflags-devel \ - glog-devel \ - libibverbs-devel \ - numactl-devel \ - gtest \ - gtest-devel \ - boost-devel \ - openssl-devel \ - hiredis-devel \ - libcurl-devel - ``` - - NOTE: You may need to install gtest, glog, gflags from source code: - ```bash - git clone https://github.com/gflags/gflags - git clone https://github.com/google/glog - git clone https://github.com/abseil/googletest.git - ``` - -2. If you want to compile the GPUDirect support module, first follow the instructions in https://docs.nvidia.com/cuda/cuda-installation-guide-linux/ to install CUDA (ensure to enable `nvidia-fs` for proper `cuFile` module compilation). After that: - 1) Configure `LIBRARY_PATH` and `LD_LIBRARY_PATH` to ensure linking of `cuFile`, `cudart`, and other libraries during compilation: - ```bash - export LIBRARY_PATH=$LIBRARY_PATH:/usr/local/cuda/lib64 - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64 - ``` - ```{admonition} GPU-Direct RDMA - :class: note - Mooncake could use the DMA-BUF path for GPU-Direct RDMA, which does **not** require the `nvidia-peermem` kernel module. If you prefer the DMA-BUF path, please set the runtime environment variable `WITH_NVIDIA_PEERMEM=0` before starting Mooncake. If you prefer the legacy `ibv_reg_mr` path (which requires `nvidia-peermem`), set the runtime environment variable `WITH_NVIDIA_PEERMEM=1`. See Section 3.7 of https://docs.nvidia.com/cuda/gpudirect-rdma/ for instructions on installing `nvidia-peermem`. - ``` - -3. If you want to compile the Moore Mthreads GPUDirect support module, first follow the instructions in https://docs.mthreads.com/musa-sdk/musa-sdk-doc-online/install_guide to install MUSA. After that: - 1) Install `mthreads-peermem` for enabling GPU-Direct RDMA - 2) Configure `LIBRARY_PATH` and `LD_LIBRARY_PATH` to ensure linking of `musart`, and other libraries during compilation: - ```bash - export LIBRARY_PATH=$LIBRARY_PATH:/usr/local/musa/lib - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/musa/lib - ``` - -4. If you want to compile Cambricon MLU support, first install the Cambricon Neuware SDK. After that: - 1) Export `NEUWARE_HOME` or pass `-DNEUWARE_ROOT=/path/to/neuware` to CMake - 2) Configure `LIBRARY_PATH` and `LD_LIBRARY_PATH` to ensure linking of `cnrt`, `cndrv`, and other Neuware libraries during compilation: - ```bash - export NEUWARE_HOME=/usr/local/neuware - export LIBRARY_PATH=$LIBRARY_PATH:${NEUWARE_HOME}/lib64 - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${NEUWARE_HOME}/lib64 - ``` - - If your Neuware installation lives outside the default include/library layout, you can also pass: - ```bash - cmake .. -DUSE_MLU=ON \ - -DMLU_INCLUDE_DIR=/path/to/neuware/include \ - -DMLU_LIB_DIR=/path/to/neuware/lib64 - ``` - - For Cambricon MLU builds, enable the MLU backend explicitly: - ```bash - cmake .. -DUSE_MLU=ON -DNEUWARE_ROOT=${NEUWARE_HOME:-/usr/local/neuware} - make -j - ``` - -5. If you want to compile MetaX (Muxi) MACA support (e.g. C500), install the MACA SDK so headers and libraries are available under `MACA_ROOT` (defaults to `MACA_HOME` env var if set, otherwise `/opt/maca`). SDK layouts vary; include both `lib` and `lib64` in runtime paths when needed: - ```bash - export MACA_HOME=/opt/maca - export LIBRARY_PATH=$LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64 - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${MACA_HOME}/lib:${MACA_HOME}/lib64 - ``` - Build with `-DUSE_MACA=ON`. Optional overrides: - - `-DMACA_ROOT=/path/to/maca` - - `-DMACA_INCLUDE_DIR=/path/to/maca/include` - - `-DMACA_LIB_DIR=/path/to/maca/lib64` - - `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"` (semicolon-separated CMake list) - -6. If you want to compile Huawei Ascend NPU support, first install the Ascend CANN Toolkit following the instructions at https://www.hiascend.com/document. After that: - 1) Source `set_env.sh` in the CANN installation directory to configure the build environment (no need to manually set `ASCEND_HOME_PATH` or other related environment variables). - 2) Mooncake provides two Ascend NPU transport paths, choose one as needed: - - `-DUSE_ASCEND_DIRECT=ON` (**recommended**): Ascend Direct transport based on the ADXL engine. (refer to [Version Compatibility Guide](https://gitcode.com/cann/hixl/wiki/Mooncake%20+%20HIXL%20%E5%BF%AB%E9%80%9F%E4%B8%8A%E6%89%8B%E6%8C%87%E5%8D%97.md) for details). - - `-DUSE_UBSHMEM=ON`: Shared memory transport based on CANN VMM APIs (requires CANN >= 9.0.0, driver >= 26.0.0, Lingqu >= 1.5). - - Example for building with Ascend NPU: - ```bash - source /usr/local/Ascend/cann/set_env.sh - cmake .. -DUSE_ASCEND_DIRECT=ON - make -j - ``` - -7. Install yalantinglibs - ```bash - git clone https://github.com/alibaba/yalantinglibs.git - cd yalantinglibs - mkdir build && cd build - cmake .. -DBUILD_EXAMPLES=OFF -DBUILD_BENCHMARK=OFF -DBUILD_UNIT_TESTS=OFF - make -j$(nproc) - make install - ``` - -8. In the root directory of this project, run the following commands: - ```bash - mkdir build - cd build - cmake .. - make -j - ``` - -9. Install Mooncake python package and mooncake_master executable - ```bash - make install - ``` +```{admonition} NCCL DeviceTransport version contract +:class: important +Mooncake currently requires NCCL Device API device code, whether AOT-compiled +or JIT-compiled, to use NCCL headers that exactly match the loaded runtime +`libnccl.so`. `NcclTransport::initialize()` rejects a mismatch. After upgrading +NCCL, rebuild Mooncake and every AOT CUDA kernel that includes +`transport/device/nccl_device.cuh`. Invalidate and regenerate any cached NCCL +Device API JIT kernels before running. This is required for the current GIN +device-code model, which is not cross-version compatible. +``` + +```{admonition} GPU-Direct RDMA +:class: note +Mooncake can use the DMA-BUF path for GPU-Direct RDMA, which does **not** +require the `nvidia-peermem` kernel module. Set `WITH_NVIDIA_PEERMEM=0` before +starting Mooncake to use DMA-BUF. Set `WITH_NVIDIA_PEERMEM=1` to use the legacy +`ibv_reg_mr` path, which requires `nvidia-peermem`. See Section 3.7 of +https://docs.nvidia.com/cuda/gpudirect-rdma/ for `nvidia-peermem` installation +instructions. +``` ## Use Mooncake in Docker Containers Mooncake supports Docker-based deployment. You can either build the image from @@ -251,44 +175,91 @@ sudo docker run --gpus all \ The `64gb` / `56gb` values above are tuned examples for large HiCache deployments, not defaults. The arena remains disabled unless you explicitly enable it, and if you enable it via gflag without an env override the default pool size is `8gb`. On smaller hosts, start with `8gb` or `16gb` and size upward with the helper. When you want the baseline direct-`mmap()` path instead of the arena, set `MC_DISABLE_MMAP_ARENA=1` (also accepts `true`, `yes`, or `on`) and omit `MC_MMAP_ARENA_POOL_SIZE`. Set it before the first Mooncake mmap-buffer allocation in the process. If you build the image from source with `docker/mooncake.Dockerfile`, that source-built image also installs the helper as `mooncake-hicache-sizing`. Without `MC_STORE_USE_HUGEPAGE=1`, the arena may opportunistically try hugepages and then retry on regular pages if HugeTLB is unavailable. When `MC_STORE_USE_HUGEPAGE=1` is set, both the arena path and the direct-`mmap()` fallback path require HugeTLB pages. Mooncake will not silently degrade that explicit hugepage request to regular pages. +For RDMA Store segments backed by HugeTLB, page population is automatically +deferred until immediately before transfer-engine registration and +parallelized across CPU threads: + +```bash +export MC_STORE_USE_HUGEPAGE=1 +export MC_STORE_HUGEPAGE_SIZE=2MB +``` + +Direct mappings use a generic worker pool. NUMA-segmented mappings bind each +worker to the node associated with its memory region. The mmap arena keeps its +eager population behavior; set `MC_DISABLE_MMAP_ARENA=1` if an arena was +otherwise enabled and deferred direct-mmap population is desired. + ## Advanced Compile Options -The following options can be used during `cmake ..` to specify whether to compile certain components of Mooncake. -- `-DUSE_CUDA=[ON|OFF]`: Enable GPU memory support (GPUDirect RDMA, NVMe-oF, and GPU-aware TCP transport). **Default: OFF.** Required when transferring GPU memory (e.g., KV cache in vLLM disaggregated serving), even when using TCP protocol. -- `-DUSE_MNNVL=[ON|OFF]`: Enable Multi-Node NVLink transport support, default is OFF. **Note:** `-DUSE_CUDA` is required when `-DUSE_MNNVL` is on (not used when building with `-DUSE_MUSA=ON`, `-DUSE_HIP=ON`, or `-DUSE_MACA=ON`). -- `-DUSE_MUSA=[ON|OFF]`: Enable Moore Threads GPU support via MUSA -- `-DUSE_MACA=[ON|OFF]`: Enable MetaX (Muxi) GPU support via MACA. -- `-DMACA_ROOT=/path/to/maca`: Override the MACA SDK root (`MACA_HOME` env var is also honored; default `/opt/maca`). -- `-DMACA_INCLUDE_DIR=/path/to/include`: Override MACA include directory when `-DUSE_MACA=ON`. -- `-DMACA_LIB_DIR=/path/to/lib64`: Override MACA library directory when `-DUSE_MACA=ON`. -- `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"`: Override MACA runtime libraries linked by `transfer_engine`. -- `-DUSE_HIP=[ON|OFF]`: Enable AMD GPU support via HIP/ROCm -- `-DUSE_HYGON=[ON|OFF]`: Enable Hygon DCU support via DTK SDK. **Default: OFF.** Uses CUDA-compatible runtime. -- `-DDTK_ROOT=/path/to/dtk`: Override the default DTK SDK root used when `-DUSE_HYGON=ON`. If unset, Mooncake uses `DTK_HOME` or `/opt/dtk`. -- `-DDTK_INCLUDE_DIR=/path/to/include`: Override the DTK include directory when `-DUSE_HYGON=ON`. -- `-DDTK_LIB_DIR=/path/to/lib64`: Override the DTK library directory when `-DUSE_HYGON=ON`. -- `-DUSE_COREX=[ON|OFF]`: Enable Iluvatar CoreX GPU support. **Default: OFF.** Uses CUDA-compatible runtime. -- `-DCOREX_ROOT=/path/to/corex`: Override the default CoreX SDK root used when `-DUSE_COREX=ON`. If unset, Mooncake uses `COREX_HOME` or `/usr/local/corex`. -- `-DCOREX_INCLUDE_DIR=/path/to/include`: Override the CoreX include directory when `-DUSE_COREX=ON`. -- `-DCOREX_LIB_DIR=/path/to/lib`: Override the CoreX library directory when `-DUSE_COREX=ON`. -- `-DUSE_MLU=[ON|OFF]`: Enable Cambricon MLU memory support via Neuware. **Default: OFF.** Supports MLU memory detection, topology discovery, and RDMA registration for Transfer Engine. -- `-DNEUWARE_ROOT=/path/to/neuware`: Override the default Neuware SDK root used when `-DUSE_MLU=ON`. If unset, Mooncake uses `NEUWARE_HOME` or `/usr/local/neuware`. -- `-DMLU_INCLUDE_DIR=/path/to/include`: Override the Neuware include directory when `-DUSE_MLU=ON`. -- `-DMLU_LIB_DIR=/path/to/lib64`: Override the Neuware library directory when `-DUSE_MLU=ON`. -- `-DUSE_EFA=[ON|OFF]`: Enable AWS Elastic Fabric Adapter transport via libfabric. **Default: OFF.** See [EFA Transport](../design/transfer-engine/efa_transport.md) for details. -- `-DUSE_INTRA_NVLINK=[ON|OFF]`: Enable intranode nvlink transport -- `-DUSE_CXL=[ON|OFF]`: Enable CXL support -- `-DWITH_STORE=[ON|OFF]`: Build Mooncake Store component -- `-DWITH_P2P_STORE=[ON|OFF]`: Enable Golang support and build P2P Store component, require go 1.23+ -- `-DWITH_RUST_EXAMPLE=[ON|OFF]`: Build the Transfer Engine Rust interface and sample code. **Default: OFF.** -- `-DWITH_STORE_RUST=[ON|OFF]`: Build Mooncake Store Rust bindings and CMake Rust targets. **Default: ON.** -- `-DWITH_EP=[ON|OFF]`: Build the EP (Expert Parallelism) and PG Python extensions for CUDA. Requires CUDA toolkit and PyTorch. Use `-DEP_TORCH_VERSIONS="2.9.1"` (semicolon-separated) to build for specific PyTorch versions, or leave empty to use the currently-installed torch. The CUDA version is detected automatically. **Default: OFF.** -- `-DUSE_REDIS=[ON|OFF]`: Enable Redis-based metadata service for the Transfer Engine, require hiredis -- `-DUSE_HTTP=[ON|OFF]`: Enable Http-based metadata service -- `-DUSE_ETCD=[ON|OFF]`: Enable etcd-based metadata service, require go 1.23+ -- `-DSTORE_USE_ETCD=[ON|OFF]`: Enable etcd-based failover for Mooncake Store, require go 1.23+. **Note:** `-DUSE_ETCD` and `-DSTORE_USE_ETCD` are two independent options. Enabling `-DSTORE_USE_ETCD` does **not** depend on `-DUSE_ETCD` -- `-DSTORE_USE_REDIS=[ON|OFF]`: Enable Redis-based failover for Mooncake Store, require hiredis. **Default: OFF.** **Note:** `-DUSE_REDIS` and `-DSTORE_USE_REDIS` are two independent options. Enabling `-DSTORE_USE_REDIS` does **not** depend on `-DUSE_REDIS`. -- `-DBUILD_SHARED_LIBS=[ON|OFF]`: Build Transfer Engine as shared library, default is OFF -- `-DBUILD_UNIT_TESTS=[ON|OFF]`: Build unit tests, default is ON -- `-DBUILD_EXAMPLES=[ON|OFF]`: Build examples, default is ON -- `-DUSE_ASCEND_DIRECT=[ON|OFF]`: Enable Ascend Direct transport and HCCS support via the ADXL engine (**recommended**). -- `-DUSE_UBSHMEM=[ON|OFF]`: Enable Huawei Ascend NPU shared memory transport via CANN VMM APIs. +The following options can be passed to `cmake ..`. + +### Accelerator and Hardware Options + +| Option | Default | Description | +| --- | --- | --- | +| `-DUSE_CUDA=ON/OFF` | `OFF` | Enable GPU memory support, including GPUDirect RDMA, NVMe-oF, and GPU-aware TCP transport. Required when transferring GPU memory, even when using TCP. | +| `-DUSE_NCCL_DEVICE=ON/OFF` | `OFF` | Enable the experimental NCCL DeviceTransport backend. Requires CUDA and NCCL 2.30.4+; AOT- and JIT-compiled NCCL device code must use headers that exactly match the runtime `libnccl`. | +| `-DUSE_NCCL_HOST=ON/OFF` | `OFF` | Enable the experimental, WRITE-only NCCL host RMA transport. Requires CUDA and NCCL 2.30.4+, must be installed before its buffers are registered, and must be the engine's only installed transport. | +| `-DUSE_MNNVL=ON/OFF` | `OFF` | Enable Multi-Node NVLink transport. Requires `-DUSE_CUDA=ON`; not used with MUSA, HIP, or MACA builds. | +| `-DUSE_MUSA=ON/OFF` | `OFF` | Enable Moore Threads GPU support via MUSA. | +| `-DUSE_MACA=ON/OFF` | `OFF` | Enable MetaX (Muxi) GPU support via MACA. | +| `-DUSE_HIP=ON/OFF` | `OFF` | Enable AMD GPU support via HIP/ROCm. | +| `-DUSE_HYGON=ON/OFF` | `OFF` | Enable Hygon DCU support via DTK SDK. Uses a CUDA-compatible runtime. | +| `-DUSE_COREX=ON/OFF` | `OFF` | Enable Iluvatar CoreX GPU support. Uses a CUDA-compatible runtime. | +| `-DUSE_MLU=ON/OFF` | `OFF` | Enable Cambricon MLU memory support via Neuware, including memory detection, topology discovery, and RDMA registration. | +| `-DUSE_ASCEND_DIRECT=ON/OFF` | `OFF` | Enable Ascend Direct transport and HCCS support via the ADXL engine. Recommended for Ascend builds. | +| `-DUSE_UBSHMEM=ON/OFF` | `OFF` | Enable Huawei Ascend NPU shared memory transport via CANN VMM APIs. | +| `-DUSE_INTRA_NVLINK=ON/OFF` | `OFF` | Enable intranode NVLink transport. | +| `-DUSE_VRAM_SEGMENT=ON/OFF` | `OFF` | Enable create VRAM Segment instead of (default) DRAM Segment. | +| `-DUSE_CXL=ON/OFF` | `OFF` | Enable CXL support. | + +### Vendor SDK Path Overrides + +| Option | Applies when | Description | +| --- | --- | --- | +| `-DMACA_ROOT=/path/to/maca` | `-DUSE_MACA=ON` | Override the MACA SDK root. `MACA_HOME` is also honored; default is `/opt/maca`. | +| `-DMACA_INCLUDE_DIR=/path/to/include` | `-DUSE_MACA=ON` | Override the MACA include directory. | +| `-DMACA_LIB_DIR=/path/to/lib64` | `-DUSE_MACA=ON` | Override the MACA library directory. | +| `-DMACA_RUNTIME_LIBS="mcruntime;mxc-runtime64;rt"` | `-DUSE_MACA=ON` | Override MACA runtime libraries linked by `transfer_engine`. | +| `-DDTK_ROOT=/path/to/dtk` | `-DUSE_HYGON=ON` | Override the DTK SDK root. `DTK_HOME` is also honored; default is `/opt/dtk`. | +| `-DDTK_INCLUDE_DIR=/path/to/include` | `-DUSE_HYGON=ON` | Override the DTK include directory. | +| `-DDTK_LIB_DIR=/path/to/lib64` | `-DUSE_HYGON=ON` | Override the DTK library directory. | +| `-DCOREX_ROOT=/path/to/corex` | `-DUSE_COREX=ON` | Override the CoreX SDK root. `COREX_HOME` is also honored; default is `/usr/local/corex`. | +| `-DCOREX_INCLUDE_DIR=/path/to/include` | `-DUSE_COREX=ON` | Override the CoreX include directory. | +| `-DCOREX_LIB_DIR=/path/to/lib` | `-DUSE_COREX=ON` | Override the CoreX library directory. | +| `-DNEUWARE_ROOT=/path/to/neuware` | `-DUSE_MLU=ON` | Override the Neuware SDK root. `NEUWARE_HOME` is also honored; default is `/usr/local/neuware`. | +| `-DMLU_INCLUDE_DIR=/path/to/include` | `-DUSE_MLU=ON` | Override the Neuware include directory. | +| `-DMLU_LIB_DIR=/path/to/lib64` | `-DUSE_MLU=ON` | Override the Neuware library directory. | + +### Transport and Metadata Options + +| Option | Default | Description | +| --- | --- | --- | +| `-DUSE_EFA=ON/OFF` | `OFF` | Enable AWS Elastic Fabric Adapter transport via libfabric. See [EFA Transport](../design/transfer-engine/efa_transport.md). | +| `-DUSE_NOF=ON/OFF` | `OFF` | Build Mooncake Store with NVMe-oF SSD pool support. Use `sudo bash dependencies.sh --with-spdk` before enabling it. | +| `-DUSE_REDIS=ON/OFF` | `OFF` | Enable Redis-based metadata service for Transfer Engine. Requires hiredis. | +| `-DUSE_HTTP=ON/OFF` | `ON` | Enable HTTP-based metadata service. | +| `-DUSE_ETCD=ON/OFF` | `OFF` | Enable etcd-based metadata service. Requires Go 1.23+. | +| `-DSTORE_USE_ETCD=ON/OFF` | `OFF` | Enable etcd-based failover for Mooncake Store. Independent from `-DUSE_ETCD`. Requires Go 1.23+. | +| `-DSTORE_USE_REDIS=ON/OFF` | `OFF` | Enable Redis-based failover for Mooncake Store. Independent from `-DUSE_REDIS`. Requires hiredis. | +| `-DSTORE_USE_K8S_LEASE=ON/OFF` | `OFF` | Enable Kubernetes Lease-based failover for Mooncake Store. Cannot be enabled with `-DSTORE_USE_ETCD=ON` or non-legacy `-DUSE_ETCD=ON`. | + +### Component Options + +| Option | Default | Description | +| --- | --- | --- | +| `-DWITH_TE=ON/OFF` | `ON` | Build the Mooncake Transfer Engine component and sample code. | +| `-DWITH_STORE=ON/OFF` | `ON` | Build the Mooncake Store component. | +| `-DWITH_STORE_GO=ON/OFF` | `OFF` | Build Go bindings for Mooncake Store when `-DWITH_STORE=ON`. | +| `-DWITH_P2P_STORE=ON/OFF` | `OFF` | Enable Golang support and build the P2P Store component. Requires Go 1.23+. | +| `-DWITH_RUST_EXAMPLE=ON/OFF` | `OFF` | Build the Transfer Engine Rust interface and sample code. | +| `-DWITH_STORE_RUST=ON/OFF` | `ON` | Build Mooncake Store Rust bindings and CMake Rust targets. | +| `-DWITH_EP=ON/OFF` | `OFF` | Build the EP and PG Python extensions for CUDA. Requires CUDA toolkit and PyTorch. Use `-DEP_TORCH_VERSIONS="2.13.0"` to build for specific PyTorch versions, or leave empty to use the currently installed torch. The CUDA version is detected automatically. | + +### Build Behavior Options + +| Option | Default | Description | +| --- | --- | --- | +| `-DBUILD_SHARED_LIBS=ON/OFF` | `OFF` | Build Transfer Engine as a shared library. | +| `-DBUILD_UNIT_TESTS=ON/OFF` | `ON` | Build unit tests. | +| `-DBUILD_EXAMPLES=ON/OFF` | `ON` | Build examples. | +| `-DSTORE_USE_JEMALLOC=ON/OFF` | `OFF` | Use jemalloc in the Mooncake Store master. | diff --git a/docs/source/getting_started/observability.md b/docs/source/getting_started/observability.md index fd0e3922b1..ffc6b007bb 100644 --- a/docs/source/getting_started/observability.md +++ b/docs/source/getting_started/observability.md @@ -165,3 +165,46 @@ The admin HTTP server is configured in the master config file (`master.json` or ``` Set `enable_metric_reporting` to `false` to disable the periodic metrics log. HTTP endpoints (`/metrics`, `/health`, etc.) remain available regardless of this setting. + +## Client Metrics Endpoint + +Mooncake clients can also expose a client-local HTTP endpoint for health checks +and client metrics. This is separate from the master admin endpoint above and is +disabled by default for Python/programmatic clients. + +Enable it through the Python setup arguments: + +```python +store.setup( + local_hostname, + metadata_server, + global_segment_size, + local_buffer_size, + protocol, + rdma_devices, + master_server_addr, + enable_client_http_server=True, + client_http_port=9300, +) +``` + +For `mooncake.mooncake_store_service`, set +`MOONCAKE_ENABLE_CLIENT_HTTP_SERVER=true` and optionally +`MOONCAKE_CLIENT_HTTP_PORT=`. For the standalone `mooncake_client`, use +`--enable_http_server=true --http_port=`. + +| Endpoint | Content-Type | Description | +|----------|--------------|-------------| +| `GET /health` | `application/json` | Client health check | +| `GET /metrics` | `text/plain; version=0.0.4` | Prometheus-format client metrics | +| `GET /metrics/summary` | `text/plain` | Human-readable client metrics summary | + +```bash +curl http://:9300/health +curl http://:9300/metrics +curl http://:9300/metrics/summary +``` + +Set `MC_STORE_CLIENT_METRIC=0` to disable client metric collection. If the +client HTTP server remains enabled while metrics are disabled, `/metrics` and +`/metrics/summary` return HTTP 503 with `metrics not available`. diff --git a/docs/source/getting_started/quick-start.md b/docs/source/getting_started/quick-start.md index d9f0379761..90d17c6f30 100644 --- a/docs/source/getting_started/quick-start.md +++ b/docs/source/getting_started/quick-start.md @@ -2,263 +2,122 @@ This document describes how to quickly start using Mooncake Transfer Engine and Mooncake Store. +## Before using Mooncake + +Install the following prerequisites before running any Mooncake component: +- Python 3.10 or later; a virtual environment is recommended. +- RDMA driver and SDK (for example, Mellanox OFED), if you plan to use RDMA for data transfer. +- CUDA 12.1 or later, if the package is built with `-DUSE_CUDA` (disabled by default). For most CUDA-enabled use cases, such as RDMA-based KV cache transfer between GPUs or between GPU and DRAM, NVIDIA GPUDirect support is also required. *You may install them from [here](https://developer.nvidia.com/cuda-downloads)*. +- Cambricon Neuware, if the package is built with `-DUSE_MLU`. By default Mooncake looks for Neuware under `NEUWARE_HOME` or `/usr/local/neuware`. +- Hygon DTK SDK, if the package is built with `-DUSE_HYGON`. By default Mooncake looks for DTK under `DTK_HOME` or `/opt/dtk`. +- Iluvatar CoreX SDK, if the package is built with `-DUSE_COREX`. By default Mooncake looks for CoreX under `COREX_HOME` or `/usr/local/corex`. + ## Installation -Install the Mooncake Transfer Engine package from PyPI, which includes both Mooncake Transfer Engine and Mooncake Store Python bindings: +Install the Mooncake package from PyPI. The same package provides: + +- Mooncake Store Python bindings for vLLM and SGLang HiCache integrations. +- Transfer Engine Python bindings and runtime components for direct + `mooncake.engine.TransferEngine` usage. **For CUDA-enabled systems:** + +- CUDA < 13.0 ```bash -pip install mooncake-transfer-engine numpy pyzmq +pip install mooncake-transfer-engine ``` -📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine/](https://pypi.org/project/mooncake-transfer-engine/) -**For non-CUDA systems:** +- CUDA >= 13.0 ```bash -pip install mooncake-transfer-engine-non-cuda numpy pyzmq +pip install mooncake-transfer-engine-cuda13 ``` -📦 **Package Details**: [https://pypi.org/project/mooncake-transfer-engine-non-cuda/](https://pypi.org/project/mooncake-transfer-engine-non-cuda/) - -> **Note**: The CUDA version includes Mooncake-EP and GPU topology detection, requiring CUDA 12.1+. The non-CUDA version is for environments without CUDA dependencies. - -## Transfer Engine Quick Start - -> **Note**: When using RDMA protocol, you may need to run with `sudo` for proper permissions. - -### Start Transfer Engine Receiver (Server) - -```python - -import numpy as np -import zmq -from mooncake.engine import TransferEngine - -def main(): - # Initialize ZMQ context and socket - context = zmq.Context() - socket = context.socket(zmq.PUSH) - socket.bind("tcp://*:5555") # Bind to port 5555 for buffer info - - HOSTNAME = "localhost" # localhost for simple demo - METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...] - PROTOCOL = "rdma" # [rdma, tcp, ...] - DEVICE_NAME = "" # auto discovery if empty - - # Initialize server engine - server_engine = TransferEngine() - server_engine.initialize( - HOSTNAME, - METADATA_SERVER, - PROTOCOL, - DEVICE_NAME - ) - session_id = f"{HOSTNAME}:{server_engine.get_rpc_port()}" - - # Allocate memory on server side (1MB buffer) - server_buffer = np.zeros(1024 * 1024, dtype=np.uint8) - server_ptr = server_buffer.ctypes.data - server_len = server_buffer.nbytes - - # Register memory with Mooncake - if PROTOCOL == "rdma": - ret_value = server_engine.register_memory(server_ptr, server_len) - if ret_value != 0: - print("Mooncake memory registration failed.") - raise RuntimeError("Mooncake memory registration failed.") - - print(f"Server initialized with session ID: {session_id}") - print(f"Server buffer address: {server_ptr}, length: {server_len}") - - # Send buffer info to client - buffer_info = { - "session_id": session_id, - "ptr": server_ptr, - "len": server_len - } - socket.send_json(buffer_info) - print("Buffer information sent to client") - - # Keep server running - try: - while True: - input("Press Ctrl+C to exit...") - except KeyboardInterrupt: - print("\nShutting down server...") - finally: - # Cleanup - if PROTOCOL == "rdma": - ret_value = server_engine.unregister_memory(server_ptr) - if ret_value != 0: - print("Mooncake memory deregistration failed.") - raise RuntimeError("Mooncake memory deregistration failed.") - - socket.close() - context.term() - -if __name__ == "__main__": - main() - +**For non-CUDA systems:** +```bash +pip install mooncake-transfer-engine-non-cuda ``` -### Start Transfer Engine Sender (Client) - -```python - - -import numpy as np -import zmq -from mooncake.engine import TransferEngine - -def main(): - # Initialize ZMQ context and socket - context = zmq.Context() - socket = context.socket(zmq.PULL) - socket.connect(f"tcp://localhost:5555") - - # Wait for buffer info from server - print("Waiting for server buffer information...") - buffer_info = socket.recv_json() - server_session_id = buffer_info["session_id"] - server_ptr = buffer_info["ptr"] - server_len = buffer_info["len"] - print(f"Received server info - Session ID: {server_session_id}") - print(f"Server buffer address: {server_ptr}, length: {server_len}") - - # Initialize client engine - HOSTNAME = "localhost" # localhost for simple demo - METADATA_SERVER = "P2PHANDSHAKE" # [ETCD_SERVER_URL, P2PHANDSHAKE, ...] - PROTOCOL = "rdma" # [rdma, tcp, ...] - DEVICE_NAME = "" # auto discovery if empty - - client_engine = TransferEngine() - client_engine.initialize( - HOSTNAME, - METADATA_SERVER, - PROTOCOL, - DEVICE_NAME - ) - session_id = f"{HOSTNAME}:{client_engine.get_rpc_port()}" - - # Allocate and initialize client buffer (1MB) - client_buffer = np.ones(1024 * 1024, dtype=np.uint8) # Fill with ones - client_ptr = client_buffer.ctypes.data - client_len = client_buffer.nbytes - - # Register memory with Mooncake - if PROTOCOL == "rdma": - ret_value = client_engine.register_memory(client_ptr, client_len) - if ret_value != 0: - print("Mooncake memory registration failed.") - raise RuntimeError("Mooncake memory registration failed.") - - print(f"Client initialized with session ID: {session_id}") - - # Transfer data from client to server - print("Transferring data to server...") - for _ in range(10): - ret = client_engine.transfer_sync_write( - server_session_id, - client_ptr, - server_ptr, - min(client_len, server_len) # Transfer minimum of both lengths - ) - - if ret >= 0: - print("Transfer successful!") - else: - print("Transfer failed!") - - # Cleanup - if PROTOCOL == "rdma": - ret_value = client_engine.unregister_memory(client_ptr) - if ret_value != 0: - print("Mooncake memory deregistration failed.") - raise RuntimeError("Mooncake memory deregistration failed.") - - socket.close() - context.term() - -if __name__ == "__main__": - main() - +**For NPU systems:** +```bash +pip install mooncake-transfer-engine-npu ``` -### More Examples and Documentation +> **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, but it still needs system runtime libraries such as `libcurl4`, `libibverbs1`, `rdma-core`, `librdmacm1`, `libnuma1`, and `liburing2` on Ubuntu. In a fresh environment, run `sudo apt-get update` before installing them: +> ```bash +> sudo apt-get update && sudo apt-get install -y libcurl4 libibverbs1 rdma-core librdmacm1 libnuma1 liburing2 +> ``` +> - MLU support is currently available through source builds with `-DUSE_MLU=ON`; there is no dedicated prebuilt MLU wheel yet. +> - If users encounter problems such as missing `lib*.so`, first install the corresponding system runtime libraries. If the issue persists, uninstall the package and build the binaries manually. -Please refer to the [Transfer Engine Python API](../python-api-reference/transfer-engine.md) and [Transfer Engine](../design/transfer-engine/index.md) for more examples and documentation. +## Connect vLLM or SGLang -## Mooncake Store Quick Start +Choose the integration path that matches your serving deployment. -### Start Master (with HTTP enabled) +### PD Disaggregation -Enable the built-in HTTP metadata server when starting the master: +PD disaggregation paths use Mooncake Transfer Engine for direct KV transfer +between prefill and decode workers. Configure these paths through the serving +framework guides, not by calling Transfer Engine APIs directly: -```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`. +- [SGLang Integration Overview](../deployment/integrations/sglang/index.md) +- [vLLM Integration Overview](../deployment/integrations/vllm/index.md) -If the master runs in a container and its IP is dynamic, set `--rpc_interface=` such as `--rpc_interface=eth0`. Mooncake Master will resolve the current IPv4 address from that interface at startup instead of relying on a fixed `--rpc_address`. +### Mooncake Store -Optional: Use the free-ratio-first allocation strategy for better load balancing across segments with different sizes or utilization: +Mooncake Store provides distributed KV cache storage for vLLM and SGLang +HiCache: -```bash -mooncake_master \ - --allocation_strategy=free_ratio_first \ - --enable_http_metadata_server=true \ - --http_metadata_server_port=8080 -``` +| Framework | Use case | Setup guide | +|-----------|----------|-------------| +| SGLang | HiCache L3 storage backend with Mooncake Store | [SGLang HiCache Quick Start](../deployment/integrations/sglang/hicache-quick-start.md) | +| vLLM | KV cache storage and sharing with `MooncakeStoreConnector` | [vLLM KV Cache Storage & Sharing](../deployment/integrations/vllm/kv-cache-storage.md) | + +The serving framework guides include the required Mooncake Store service +startup and connector configuration for each path. -The free-ratio-first strategy balances memory utilization ratio across segments by sampling multiple candidates and preferentially allocating to those with higher free space ratios, leading to more even utilization. +## Optional Python Smoke Test -### Hello World Example +If you want to verify the Store Python API without a serving framework, run this +single-node `put`/`get` example after starting `mooncake_master`. It uses +`P2PHANDSHAKE`, so no separate Transfer Engine metadata service is required. ```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 + local_hostname="localhost", + metadata_server="P2PHANDSHAKE", + global_segment_size=512 * 1024 * 1024, + local_buffer_size=128 * 1024 * 1024, + protocol="tcp", + rdma_devices="", + master_server_addr="127.0.0.1:50051", ) -# 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() ``` -### More Examples and Documentation - -Please refer to the [Mooncake Store Python API](../python-api-reference/mooncake-store.md), [Mooncake Store](../design/mooncake-store.md) and [Mooncake Store Deployment & Tuning Guide](../deployment/mooncake-store-deployment-guide.md) for more examples and documentation. +## AI Coding Assistant Skills -## Skills for AI Coding Assistants +If you use Claude Code or another coding assistant that supports reusable +skills, Mooncake provides built-in playbooks for common development tasks: -Mooncake ships a set of **built-in skills** under [`.claude/skills`](https://github.com/kvcache-ai/Mooncake/tree/main/.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: +| Skill | Use it for | +|-------|------------| +| `/mooncake-troubleshoot` | Diagnose services, RDMA, environment variables, and runtime logs. | +| `/mooncake-ci-local` | Run pre-PR local validation with Mooncake's CI script. | +| `/mooncake-api` | Work with Mooncake Store, Transfer Engine, and EP/Backend Python APIs. | -| 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. | - -Install them without cloning the repository via the [Claude Code plugin marketplace](https://code.claude.com/docs/en/plugin-marketplaces): +Install them from the Claude Code plugin marketplace without cloning the full +repository: ```text /plugin marketplace add kvcache-ai/Mooncake --sparse .claude-plugin @@ -267,4 +126,11 @@ Install them without cloning the repository via the [Claude Code plugin marketpl /plugin install mooncake-api@mooncake ``` -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. +## Next Steps + +For production deployment, standalone store services, high availability, +allocation strategies, SSD offload, and runtime tuning, continue to the +[Mooncake Store Deployment & Tuning Guide](../deployment/mooncake-store-deployment-guide.md). + +For API details, see the [Mooncake Store Python API](../api-reference/python/mooncake-store.md) +and [Mooncake Store design](../design/mooncake-store.md). diff --git a/docs/source/getting_started/supported-protocols.md b/docs/source/getting_started/supported-protocols.md index 6ebbb9a796..baa8a6ebbd 100644 --- a/docs/source/getting_started/supported-protocols.md +++ b/docs/source/getting_started/supported-protocols.md @@ -16,6 +16,7 @@ Mooncake Transfer Engine supports multiple communication protocols for data tran | **barex** | RDMA-capable NIC | Bare-metal RDMA extension | ⚠️ Advanced | | **cxl** | CXL-capable hardware | Memory pooling and sharing | ⚠️ Advanced | | **ascend** | Huawei Ascend NPU | Ascend NPU communication | ⚠️ Advanced | +| **tpu** | Google TPU (PJRT) | TPU KV-cache transfer via host-DRAM staging | 🧪 Experimental (TENT) | ## Commonly Used Protocols (Python API) @@ -266,6 +267,39 @@ export MC_FORCE_MNNVL=true - [Heterogeneous Ascend](../design/transfer-engine/heterogeneous_ascend.md) - [Ascend Transport](../design/transfer-engine/ascend_transport.md) +### TPU Transport (tpu) — Experimental + +**Description:** Google TPU support in the TENT runtime. Because TPU HBM is not +directly addressable by the NIC, transfers touching TPU memory are staged +through host DRAM: the HBM ↔ host-DRAM hop is performed by a PJRT device-copy +adapter, and the host ↔ host hop is carried by an existing transport (RDMA/TCP). +The two stages are chained automatically by the TENT staging pipeline +(`ProxyManager`), so no separate networked TPU transport is required. + +**Status:** Experimental. The C++/TENT data path is gated behind `-DUSE_TPU=ON` +(OFF by default). A serving-framework (JAX / PyTorch-XLA) integration layer is +planned as a follow-up. + +**Use When:** +- Disaggregated prefill/decode serving on TPU hosts +- KV-cache transfer between TPU nodes over RDMA/TCP + +**Requirements:** +- Built with `-DUSE_TPU=ON -DUSE_TENT=ON` +- A PJRT device-copy adapter shared library exposing the `mc_tpu_pjrt_*` C ABI + (see `tpu_pjrt_abi.h`). The adapter is resolved at runtime via `dlopen`; its + path defaults to `libmooncake_tpu_pjrt.so` and can be overridden with the + `MC_TPU_PJRT_LIB` environment variable. No PJRT/TPU SDK is required at build + time. +- An RDMA (or TCP) transport enabled for the host ↔ host hop. + +**Design notes:** +- TPU memory is reported as a distinct memory type (`tpu:N` locations); the + staging policy routes the local HBM ↔ host copy to the TPU device-copy + transport and the cross-node hop to RDMA/TCP. +- DMA-mapped (pinned) staging buffers for true async device DMA are a planned + performance follow-up. + ## Configuration Examples ### Configuration File (JSON) @@ -363,8 +397,8 @@ If a protocol fails to initialize: ## See Also -- [Quick Start Guide](quick-start.md) - Getting started with Mooncake +- [Quick Start](quick-start.md) - Start with Mooncake integrations for serving frameworks - [Transfer Engine Design](../design/transfer-engine/index.md) - Detailed architecture - [Transfer Engine Benchmark](../design/transfer-engine/transfer-engine-bench-tuning.md) - Performance tuning -- [Python API Reference](../python-api-reference/transfer-engine.md) - API documentation +- [Python API Reference](../api-reference/python/transfer-engine.md) - API documentation - [Deployment Guide](../deployment/mooncake-store-deployment-guide.md) - Production deployment diff --git a/docs/source/image/sglang_pd_qwen3_235b_bandwidth.png b/docs/source/image/sglang_pd_qwen3_235b_bandwidth.png new file mode 100644 index 0000000000..63dfde7504 Binary files /dev/null and b/docs/source/image/sglang_pd_qwen3_235b_bandwidth.png differ diff --git a/docs/source/image/sglang_pd_qwen3_235b_transfer_time.png b/docs/source/image/sglang_pd_qwen3_235b_transfer_time.png new file mode 100644 index 0000000000..6073a85439 Binary files /dev/null and b/docs/source/image/sglang_pd_qwen3_235b_transfer_time.png differ diff --git a/docs/source/image/sglang_pd_qwen3_235b_ttft_breakdown.png b/docs/source/image/sglang_pd_qwen3_235b_ttft_breakdown.png new file mode 100644 index 0000000000..b36c5b3de3 Binary files /dev/null and b/docs/source/image/sglang_pd_qwen3_235b_ttft_breakdown.png differ diff --git a/docs/source/index.md b/docs/source/index.md index 47204a2535..4f980bcc30 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -62,9 +62,9 @@ At the center of Mooncake is a KVCache-centric scheduler that balances effective - **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](getting_started/examples/lmdeploy-integration-v0.9) for LMDeploy. +- **Jun 20, 2025**: Mooncake becomes a PD disaggregation [backend](deployment/integrations/lmdeploy) 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](getting_started/examples/lmcache-integration) unite to pioneer KVCache-centric LLM serving system. +- **May 8, 2025**: [Mooncake x LMCache](deployment/integrations/lmcache/index) 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. @@ -85,7 +85,7 @@ At the center of Mooncake is a KVCache-centric scheduler that balances effective :::{toctree} :caption: Getting Started -:maxdepth: 2 +:maxdepth: 1 getting_started/build getting_started/quick-start @@ -96,13 +96,14 @@ getting_started/quick-start :::{toctree} :caption: Deployment -:maxdepth: 2 +:maxdepth: 1 deployment/mooncake-store-deployment-guide -getting_started/examples/sglang-integration/index -getting_started/examples/vllm-integration/index -Mooncake x LMCache Integration -Mooncake x LMDeploy Integration +deployment/kubernetes-deployment-guide/index +deployment/integrations/sglang/index +deployment/integrations/vllm/index +Mooncake x LMCache Integration +Mooncake x LMDeploy Integration ::: @@ -111,17 +112,16 @@ Mooncake x LMDeploy Integration \ +MOONCAKE_PROTOCOL=rdma \ +WITH_NVIDIA_PEERMEM=1 \ +python -m sglang.launch_server \ + --model-path /path/to/Qwen3-235B-A22B-Instruct-2507 \ + --host 0.0.0.0 --port 30000 \ + --tp-size 8 \ + --dtype bfloat16 \ + --trust-remote-code \ + --stream-interval 1 \ + --chunked-prefill-size -1 \ + --disable-radix-cache \ + --disaggregation-mode prefill \ + --disaggregation-bootstrap-port 8998 \ + --disaggregation-transfer-backend mooncake \ + --disaggregation-ib-device /path/to/ib_devices.json \ + --enable-request-time-stats-logging +``` + +**Decoder:** + +```bash +SGLANG_HOST_IP= \ +MOONCAKE_PROTOCOL=rdma \ +WITH_NVIDIA_PEERMEM=1 \ +python -m sglang.launch_server \ + --model-path /path/to/Qwen3-235B-A22B-Instruct-2507 \ + --host 0.0.0.0 --port 30001 \ + --tp-size 8 \ + --dtype bfloat16 \ + --trust-remote-code \ + --stream-interval 1 \ + --chunked-prefill-size -1 \ + --disable-radix-cache \ + --disaggregation-mode decode \ + --disaggregation-transfer-backend mooncake \ + --disaggregation-ib-device /path/to/ib_devices.json \ + --enable-request-time-stats-logging +``` + +**Router (Decoder Node):** + +```bash +python -m sglang_router.launch_router \ + --host 0.0.0.0 --port 8000 \ + --pd-disaggregation \ + --prefill http://:30000 8998 \ + --decode http://:30001 \ + --policy round_robin +``` + +### Benchmark Script + +We used the official `sglang.benchmark.serving` module to generate traffic with varying prompt lengths. Each prompt length contains 50 requests with an output length of 128 tokens. + +```bash +for prompt_len in 128 256 512 1024 2048 4096 8192 16384 32768; do + python -m sglang.benchmark.serving \ + --backend sglang \ + --base-url http://127.0.0.1:8000 \ + --model /path/to/Qwen3-235B-A22B-Instruct-2507 \ + --served-model-name Qwen3-235B-A22B-Instruct-2507 \ + --tokenizer /path/to/Qwen3-235B-A22B-Instruct-2507 \ + --dataset-name random \ + --dataset-path /path/to/ShareGPT_V3_unfiltered_cleaned_split.json \ + --num-prompts 50 \ + --random-input-len "${prompt_len}" \ + --random-output-len 128 \ + --random-range-ratio 1.0 \ + --request-rate inf \ + --max-concurrency 1 \ + --seed 42 \ + --flush-cache \ + --warmup-requests 1 \ + --tokenize-prompt \ + --output-details \ + --output-file "result-${prompt_len}.jsonl" +done +``` + +## PD Disaggregation vs. Regular SGLang + +We evaluated the earlier implementation on two A10 servers. By comparing the performance of a 1P1D configuration with that of two regular (non-disaggregated) instances, we observed that P/D disaggregation achieves approximately 30% lower ITL while maintaining comparable total throughput. This aligns with findings from the Mooncake paper, which highlighted that P/D disaggregation is effective in reducing TBT/ITL under similar throughput conditions—or conversely, in enabling higher throughput under stricter ITL/TBT SLOs. + +Moreover, we anticipate even greater benefits in larger-scale clusters where both the number of prefill and decode nodes (x and y in xPyD configurations) increase, offering enhanced scheduling flexibility and resource efficiency. + +### Traffic Request Rate: 1.0 + +- model: Qwen2.5-7B-Instruct-GPTQ-Int4 +- TP: 4 +- random_input_len=8192, random_output_len=512 +- num prompt=50 + +| Configuration | Output Token Throughput (tok/s) | Mean E2E Latency (ms) |Total Token Throughput (tok/s) | Mean TTFT (ms) | P99 TTFT (ms) | Mean ITL (ms) | P99 ITL (ms) | +|----------------|----------------------------------|-----------------------|----------------|---------------|---------------|--------------|---------------------------------| +| 1P1D | 407.59 | 3413.86 |7084.46 | 732.54 | 2952.57 | 7.23 | 10.76 | +| 2 Regular | 427.65 | 4586.54 |7433.27 | 767.18 | 1264.88 | 10.30 | 12.73 | + +### Traffic Request Rate: 4.0 + +- model: Qwen2.5-7B-Instruct-GPTQ-Int4 +- TP: 2 +- random_input_len=2048, random_output_len=512 +- num prompt=200 + +| Configuration | Output Token Throughput (tok/s) | Mean E2E Latency (ms) | Total Token Throughput (tok/s) | Mean TTFT (ms) |P99 TTFT (ms) | Mean ITL (ms) | P99 ITL (ms) | +|---------------|---------------------------------|-------------|--------------------------------|----------------|---------------|----------------|--------------| +| 1P1D | 1215.17 | 11519.24 | 6161.43 | 1111.94 | 2725.89 | 17.06 | 19.72 | +| 2 Regular | 1223.03 | 11683.15 | 6201.29 | 310.01 | 720.91 | 25.74 | 294.89 | + +By the Mooncake Team + +© Copyright 2026, Mooncake Team. diff --git a/docs/source/performance/sglang-hicache-benchmark-results-v1.md b/docs/source/performance/sglang/sglang-hicache-benchmark-results-v1.md similarity index 96% rename from docs/source/performance/sglang-hicache-benchmark-results-v1.md rename to docs/source/performance/sglang/sglang-hicache-benchmark-results-v1.md index a127845d88..bc8ec053fd 100644 --- a/docs/source/performance/sglang-hicache-benchmark-results-v1.md +++ b/docs/source/performance/sglang/sglang-hicache-benchmark-results-v1.md @@ -1,4 +1,4 @@ -# SGLang HiCache with Mooncake Backend Benchmark +# SGLang HiCache x Mooncake Store Performance We evaluated the performance of SGLang HiCache using a multi-turn conversation benchmark designed to simulate realistic user interactions. The benchmark spawns concurrent clients, each engaging in multi-round conversations. For every client, starting from the second round, the input consists of the concatenation of the input and output from the preceding round. @@ -13,13 +13,13 @@ Since HiCache is currently used to accelerate the prefill stage, this benchmark ## Benchmark Result -![overall performance](../image/hicache_multi_turn_overall.png) +![overall performance](../../image/hicache_multi_turn_overall.png) We first evaluated the overall performance on a cluster consisting 3 servers, each has 2 NVIDIA A10 GPUs and 2 100Gbps eRDMA NICs. As shown in the figure, in terms of prefill performance, `populated Mooncake` achieves the best results, followed by `+Mooncake`, `+L2`, and finally `GPU only`. -![overall performance](../image/hicache_multi_turn_per_turn.png) +![overall performance](../../image/hicache_multi_turn_per_turn.png) Next, we take a closer look at the sources of performance differences. We recorded the TTFT and cache hit rate for each conversation round. The experiment was conducted on a server equipped with 8 × H800 GPUs and 8 × mlx5 RDMA NICs. To minimize interference from the decode stage and highlight prefill performance differences, we set the output length to 1. @@ -43,7 +43,7 @@ In practical deployment, Mooncake aggregates memory across the entire cluster in **Cluster Deployment:** The benchmark was conducted across a 3-node cluster connected via eRDMA: - Machine A: SGLang server, Mooncake master service, Mooncake client (30GB memory) -- Machine B: Mooncake client (60GB memory) +- Machine B: Mooncake client (60GB memory) - Machine C: Mooncake client (60GB memory) - Total Distributed Memory Pool: 150GB @@ -178,4 +178,4 @@ python3 -m sglang.launch_server \ --hicache-ratio 2 \ --hicache-storage-prefetch-policy timeout \ --hicache-storage-backend mooncake -``` \ No newline at end of file +``` diff --git a/docs/source/performance/vllm-benchmark-results-v0.2.md b/docs/source/performance/vllm-benchmark-results-v0.2.md deleted file mode 100644 index 4f3bd612f0..0000000000 --- a/docs/source/performance/vllm-benchmark-results-v0.2.md +++ /dev/null @@ -1,49 +0,0 @@ -# Benchmark performance on NVIDIA A10 -Here are some preview mooncake benchmark results on A10 with up to 2 RDMA NICs. We are currently having some trouble benchmarking `PyNcclConnector` now. For some unknown reasons, it crashes a lot for inter-node disaggregated scenarios. So the benchmark results haven't included the `PyNcclConnector` yet. - -In addition, we are also coordinating resources to integrate some machines with more RDMA NICs and more advanced GPUs. The official benchmark results will be released in due time. - -## Varying tp (input length = 1024, qps = 2, output length =6) -| Setting | num_rdma_nic | Successful Requests | Duration (s) | Total Input Tokens | Total Generated Tokens | Req Throughput (req/s) | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms) | Mean TPOT (ms) | Median TPOT (ms) | P99 TPOT (ms) | Mean ITL (ms) | Median ITL (ms) | P99 ITL (ms) | -|-----------------|--------------|---------------------|--------------|--------------------|------------------------|------------------------|---------------------------------|-------------------------------|----------------|-----------------|--------------|---------------|------------------|--------------|--------------|----------------|-------------| -| tp = 1 | 2 | 200 | 99.47 | 201995 | 1200 | 2.01 | 12.06 | 2042.74 | 1056.76 | 635.00 | 4006.59 | 97.08 | 26.94 | 781.91 | 97.01 | 14.05 | 2205.51 | -| tp = 2 | 2 | 200 | 98.98 | 201995 | 1200 | 2.02 | 12.12 | 2052.95 | 314.87 | 231.20 | 949.40 | 25.65 | 15.56 | 129.60 | 25.62 | 15.48 | 288.06 | -| tp = 4 | 2 | 200 | 98.76 | 201995 | 1200 | 2.03 | 12.15 | 2057.44 | 198.10 | 160.03 | 461.61 | 23.52 | 18.93 | 94.38 | 23.50 | 18.01 | 187.79 | -| tp = 1 | 1 | 200 | 99.44 | 201995 | 1200 | 2.01 | 12.07 | 2043.39 | 1071.12 | 631.56 | 4361.02 | 83.93 | 26.93 | 794.75 | 83.86 | 14.13 | 1932.66 | -| tp = 2 | 1 | 200 | 98.96 | 201995 | 1200 | 2.02 | 12.13 | 2053.35 | 335.26 | 258.30 | 997.93 | 28.84 | 15.56 | 144.82 | 28.80 | 15.42 | 397.56 | -| tp = 4 | 1 | 200 | 98.78 | 201995 | 1200 | 2.02 | 12.15 | 2057.03 | 201.68 | 162.85 | 456.33 | 22.31 | 16.74 | 94.76 | 22.29 | 16.73 | 189.13 | -| tp = 1 | TCP | 200 | 99.55 | 201995 | 1200 | 2.01 | 12.05 | 2041.13 | 1414.05 | 766.23 | 6035.36 | 155.01 | 35.28 | 1191.24 | 154.91 | 14.32 | 3148.99 | -| tp = 2 | TCP | 200 | 98.97 | 201995 | 1200 | 2.02 | 12.12 | 2053.03 | 333.74 | 251.32 | 954.63 | 28.74 | 15.49 | 161.24 | 28.70 | 15.35 | 393.52 | -| tp = 4 | TCP | 200 | 98.78 | 201995 | 1200 | 2.02 | 12.15 | 2056.94 | 205.37 | 162.92 | 463.70 | 21.54 | 16.51 | 94.04 | 21.51 | 16.56 | 170.54 | - -## Varying qps (length = 1024, tp = 4, output length =6) -|Setting | num_rdma_nic | Successful Requests | Duration (s) | Total Input Tokens | Total Generated Tokens | Req Throughput (req/s) | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms) | Mean TPOT (ms) | Median TPOT (ms) | P99 TPOT (ms) | Mean ITL (ms) | Median ITL (ms) | P99 ITL (ms) | -|-----------------|--------------|---------------------|--------------|--------------------|------------------------|------------------------|---------------------------------|-------------------------------|----------------|-----------------|--------------|---------------|------------------|--------------|--------------|----------------|-------------| -| qps = 2 | 2 | 200 | 98.77 | 201995 | 1200 | 2.02 | 12.15 | 2057.33 | 200.64 | 156.62 | 478.22 | 22.63 | 17.35 | 99.61 | 22.60 | 17.08 | 186.25 | -| qps = 4 | 2 | 200 | 49.75 | 201995 | 1200 | 4.02 | 24.12 | 4084.03 | 341.88 | 240.68 | 1430.54 | 38.36 | 18.39 | 313.45 | 38.31 | 17.17 | 588.80 | -| qps = 6 | 2 | 200 | 33.44 | 201995 | 1200 | 5.98 | 35.88 | 6075.54 | 851.15 | 501.59 | 3239.89 | 102.51 | 47.67 | 606.77 | 102.34 | 18.35 | 1704.79 | -| qps = 8 | 2 | 200 | 27.16 | 201995 | 1200 | 7.36 | 44.19 | 7482.52 | 4835.08 | 5733.45 | 8846.27 | 1276.59 | 1150.11 | 4401.23 | 1274.43 | 48.34 | 20682.35 | -| qps = 2 | 1 | 200 | 98.77 | 201995 | 1200 | 2.02 | 12.15 | 2057.31 | 201.77 | 161.53 | 473.44 | 22.13 | 16.52 | 96.18 | 22.11 | 16.51 | 190.40 | -| qps = 4 | 1 | 200 | 49.76 | 201995 | 1200 | 4.02 | 24.12 | 4083.83 | 337.31 | 243.38 | 1395.85 | 39.95 | 17.61 | 325.39 | 39.88 | 17.06 | 838.68 | -| qps = 6 | 1 | 200 | 33.44 | 201995 | 1200 | 5.98 | 35.88 | 6075.99 | 820.53 | 458.84 | 3169.52 | 83.92 | 30.50 | 663.07 | 83.78 | 17.85 | 1306.32 | -| qps = 8 | 1 | 200 | 27.19 | 201995 | 1200 | 7.36 | 44.14 | 7473.44 | 5291.91 | 6160.55 | 9596.56 | 1190.36 | 1040.63 | 4418.66 | 1188.33 | 47.61 | 20815.23 | -| qps = 2 | TCP | 200 | 98.76 | 201995 | 1200 | 2.03 | 12.15 | 2057.42 | 207.22 | 160.81 | 511.01 | 22.17 | 16.59 | 94.96 | 22.15 | 16.59 | 181.82 | -| qps = 4 | TCP | 200 | 49.79 | 201995 | 1200 | 4.02 | 24.10 | 4081.06 | 355.43 | 252.63 | 1554.91 | 40.15 | 16.92 | 314.28 | 40.09 | 16.66 | 708.50 | -| qps = 6 | TCP | 200 | 33.49 | 201995 | 1200 | 5.97 | 35.83 | 6067.71 | 907.74 | 514.85 | 3253.93 | 122.75 | 45.51 | 648.40 | 122.56 | 18.09 | 2282.92 | -| qps = 8 | TCP | 200 | 28.39 | 201995 | 1200 | 7.04 | 42.26 | 7156.09 | 6714.57 | 7885.09 | 11787.51 | 1116.06 | 408.32 | 4645.25 | 1114.29 | 46.87 | 21898.03 | - -## Varying input length (tp = 4, qps = 2, output length =6) -| Setting | num_rdma_nic | Successful Requests | Duration (s) | Total Input Tokens | Total Generated Tokens | Req Throughput (req/s) | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms) | Mean TPOT (ms) | Median TPOT (ms) | P99 TPOT (ms) | Mean ITL (ms) | Median ITL (ms) | P99 ITL (ms) | -|-----------------|--------------|---------------------|--------------|--------------------|------------------------|------------------------|---------------------------------|-------------------------------|----------------|-----------------|--------------|---------------|------------------|--------------|--------------|----------------|-------------| -| 1024 | 2 | 200 | 98.77 | 201995 | 1200 | 2.02 | 12.15 | 2057.32 | 195.47 | 151.55 | 482.84 | 22.83 | 19.27 | 96.55 | 22.81 | 18.12 | 158.16 | -| 2048 | 2 | 200 | 99.22 | 406707 | 1200 | 2.02 | 12.09 | 4110.95 | 723.76 | 488.67 | 2941.96 | 67.25 | 18.93 | 632.73 | 67.20 | 17.49 | 1209.54 | -| 4096 | 2 | 200 | 117.42 | 818415 | 1200 | 1.70 | 10.22 | 6979.90 | 14616.48 | 18323.82 | 23191.04 | 8042.84 | 7593.16 | 19851.11 | 8040.02 | 65.43 | 93511.26 | -| 8192 | 2 | 200 | 247.77 | 1636065 | 1200 | 0.81 | 4.84 | 6608.10 | 75783.36 | 79331.60 | 147544.42 | 16961.27 | 15140.11 | 39278.98 | 16958.32 | 90.01 | 186151.61 | -| 1024 | 1 | 200 | 98.77 | 201995 | 1200 | 2.02 | 12.15 | 2057.31 | 201.77 | 161.53 | 473.44 | 22.13 | 16.52 | 96.18 | 22.11 | 16.51 | 190.40 | -| 2048 | 1 | 200 | 99.25 | 406707 | 1200 | 2.02 | 12.09 | 4109.96 | 719.43 | 482.02 | 3208.13 | 61.92 | 17.64 | 681.26 | 61.86 | 16.83 | 978.90 | -| 4096 | 1 | 200 | 111.88 | 818415 | 1200 | 1.79 | 10.73 | 7326.16 | 20362.10 | 22807.05 | 31853.55 | 5915.16 | 4521.51 | 18739.12 | 5913.18 | 67.03 | 81600.29 | -| 8192 | 1 | 200 | 270.01 | 1636065 | 1200 | 0.74 | 4.44 | 6063.79 | 103355.40 | 106546.65 | 172025.11 | 12894.35 | 11027.66 | 35110.13 | 12892.85 | 64.84 | 151774.68 | -| 1024 | TCP | 200 | 98.81 | 201995 | 1200 | 2.02 | 12.14 | 2056.44 | 203.32 | 160.83 | 460.90 | 21.81 | 16.96 | 95.27 | 21.78 | 16.91 | 171.80 | -| 2048 | TCP | 200 | 99.27 | 406707 | 1200 | 2.01 | 12.09 | 4108.98 | 731.60 | 484.78 | 3213.69 | 68.55 | 17.88 | 639.93 | 68.49 | 17.33 | 1257.45 | -| 4096 | TCP | 200 | 118.37 | 818415 | 1200 | 1.69 | 10.14 | 6923.89 | 23735.69 | 27101.97 | 36573.47 | 6386.62 | 5102.00 | 20032.26 | 6384.71 | 69.57 | 92811.27 | -| 8192 | TCP | 200 | 278.12 | 1636065 | 1200 | 0.72 | 4.31 | 5886.95 | 106873.23 | 109941.33 | 179781.64 | 13360.87 | 12155.24 | 36022.96 | 13359.20 | 68.01 | 156716.38 | diff --git a/docs/source/performance/vllm-benchmark-results-v1.md b/docs/source/performance/vllm-benchmark-results-v1.md deleted file mode 100644 index 05032748b6..0000000000 --- a/docs/source/performance/vllm-benchmark-results-v1.md +++ /dev/null @@ -1,22 +0,0 @@ -# Benchmark performance on NVIDIA A10 -Here are some preview MooncakeStore benchmark results on A10 with "Qwen/Qwen2.5-7B-Instruct-GPTQ-Int4". - - -## Varying PD ratio (input length = 1024, qps = 2, output length =6, num of requests = 200) -| Configuration | Backend | Duration (s) | Output Token Throughput (tok/s) | Total Token Throughput (tok/s) | Mean TTFT (ms) | Median TTFT (ms) | P99 TTFT (ms) | Mean TPOT (ms) | Median TPOT (ms) | P99 TPOT (ms) | Mean ITL (ms) | Median ITL (ms) | P99 ITL (ms) | -|----------------------------|----------------------|--------------|---------------------------------|-------------------------------|----------------|-----------------|--------------|---------------|------------------|--------------|--------------|----------------|-------------| -| 2P2D tp = 1 | Redis | 99.47 | 12.06 | 2042.75 | 844.28 | 666.84 | 2270.91 | 16.88 | 11.57 | 104.83 | 16.84 | 11.56 | 239.67 | -| | MooncakeStore (TCP) | 99.44 | 12.07 | 2043.30 | 817.43 | 639.48 | 1969.89 | 12.49 | 11.55 | 45.52 | 12.46 | 11.55 | 15.31 | -| | MooncakeStore (RDMA) | 99.33 | 12.08 | 2045.57 | 763.58 | 604.22 | 2030.34 | 12.43 | 11.53 | 43.02 | 12.39 | 11.52 | 15.40 | -| 2P2D tp = 2 | Redis | 98.92 | 12.13 | 2054.12 | 397.20 | 352.37 | 782.44 | 9.00 | 8.05 | 36.06 | 8.97 | 8.03 | 13.94 | -| | MooncakeStore (TCP) | 98.81 | 12.14 | 2056.43 | 327.91 | 309.38 | 573.36 | 8.33 | 8.04 | 17.62 | 8.30 | 8.03 | 11.23 | -| | MooncakeStore (RDMA) | 98.74 | 12.15 | 2057.79 | 271.25 | 250.11 | 532.00 | 8.34 | 8.12 | 14.70 | 8.31 | 8.10 | 11.12 | -| 3P3D (1 remote P, 1 remote D) tp = 2 qps = 2 | Redis | 98.89 | 12.13 | 2054.80 | 382.73 | 358.18 | 659.31 | 8.07 | 8.02 | 8.86 | 8.04 | 7.99 | 10.14 | -| | MooncakeStore (TCP) | 98.71 | 12.16 | 2058.47 | 298.71 | 302.74 | 512.84 | 8.06 | 8.04 | 8.54 | 8.03 | 8.02 | 8.88 | -| | MooncakeStore (RDMA) | 98.69 | 12.16 | 2058.88 | 269.73 | 252.96 | 543.38 | 8.13 | 8.04 | 10.49 | 8.10 | 8.02 | 11.29 | -| 4P2D (2 remote P) tp = 2 | Redis | 98.85 | 12.14 | 2055.66 | 350.39 | 339.15 | 506.78 | 8.54 | 8.01 | 26.42 | 8.51 | 7.99 | 11.43 | -| | MooncakeStore (TCP) | 98.76 | 12.15 | 2057.56 | 312.32 | 307.50 | 475.79 | 8.29 | 8.03 | 19.87 | 8.25 | 8.01 | 9.51 | -| | MooncakeStore (RDMA) | 98.71 | 12.16 | 2058.59 | 259.87 | 251.23 | 461.96 | 8.20 | 8.05 | 10.20 | 8.17 | 8.03 | 11.50 | -| 2P4D (2 remote D) tp = 2 | Redis | 98.88 | 12.14 | 2054.90 | 381.91 | 338.25 | 722.00 | 8.07 | 8.05 | 8.55 | 8.04 | 8.02 | 9.15 | -| | MooncakeStore (TCP) | 98.78 | 12.15 | 2057.11 | 317.42 | 304.53 | 521.66 | 8.07 | 8.03 | 8.75 | 8.04 | 8.02 | 9.62 | -| | MooncakeStore (RDMA) | 98.73 | 12.15 | 2058.02 | 275.13 | 251.57 | 487.43 | 8.18 | 8.06 | 9.19 | 8.15 | 8.05 | 10.53 | diff --git a/docs/source/performance/vllm/index.md b/docs/source/performance/vllm/index.md index b32561c7dc..f6df6cafdf 100644 --- a/docs/source/performance/vllm/index.md +++ b/docs/source/performance/vllm/index.md @@ -1,18 +1,17 @@ -# vLLM Integration Performance Benchmarks +# vLLM Integration Performance Benchmarks evaluating Mooncake's integration with vLLM across different backends and scenarios. -| Document | Backend | Key Findings | -|----------|---------|---------------| -| [vLLM V1 + MooncakeConnector](../vllm-v1-support-benchmark) | vLLM V1 | 1P1D PD disaggregation on H800 with 8x RoCE: **142.25 GB/s** peak transfer bandwidth (71.1% of theoretical), KV transfer overhead just **4.2%** of total TTFT at 32K tokens | -| [vLLM V1 + MooncakeStore vs Redis](../vllm-benchmark-results-v1) | vLLM V1 | MooncakeStore RDMA consistently outperforms Redis across all XpYd topologies — e.g., **~32% lower** mean TTFT in 2P2D tp=2 | -| [vLLM V0 + MooncakeConnector (Legacy)](../vllm-benchmark-results-v0.2) | vLLM V0 | TP=4 reduces TTFT by ~80% vs TP=1; RDMA provides significant latency advantage over TCP across varying QPS and input lengths | +| Document | Scenario | Highlights | +|----------|----------|---------------| +| [PD Disaggregation Performance](vllm-v1-pd-performance) | PD disaggregation with Mooncake Connector | 1P1D PD disaggregation on H800 with 8x RoCE: **142.25 GB/s** peak transfer bandwidth (71.1% of theoretical), KV transfer overhead just **4.2%** of total TTFT at 32K tokens | +| [vLLM x Mooncake Store Performance](vllm-v1-mooncake-store) | distributed KV cache pool with Mooncake Store | Distributed KV cache pool improves throughput by **3.8x**, reduces P50 TTFT and E2E latency by **46x** and **8.6x**, and scales to **60 GPUs** with >95% cache hit rate | :::{toctree} :maxdepth: 1 :hidden: -../vllm-v1-support-benchmark -../vllm-benchmark-results-v1 -../vllm-benchmark-results-v0.2 +vllm-v1-pd-performance +vllm-v1-mooncake-store + ::: diff --git a/docs/source/performance/vllm/vllm-v1-mooncake-store.md b/docs/source/performance/vllm/vllm-v1-mooncake-store.md new file mode 100644 index 0000000000..df8a78c3db --- /dev/null +++ b/docs/source/performance/vllm/vllm-v1-mooncake-store.md @@ -0,0 +1,44 @@ +# vLLM x Mooncake Store Performance +Mooncake leverages the `MooncakeStoreConnector` in vLLM V1 to enable a distributed KV cache pool, supporting cross-instance sharing and reuse of KV caches. Furthermore, vLLM's `MultiConnector` can be configured to orchestrate both the `MooncakeConnector` (for peer-to-peer KV transfer) and the `MooncakeStoreConnector` (for the shared pool), enabling prefill-decode (PD) disaggregation. + +![Overall Performance](https://vllm.ai/blog-assets/figures/2026-05-06-mooncake-store/hero_vllm_mooncake.svg) + +We thank the vLLM team for conducting the performance evaluation. The detailed results are presented below. + +> The original blog is available at https://vllm.ai/blog/2026-05-06-mooncake-store. + + +## Speeding up real agentic traces + +Setup: Kimi-2.5 NVFP4 model on GB200 nodes with PD disaggregation + +In this experiment, the model was deployed with a 1P1D configuration across 12 GPUs in total. + +![Throughput and Latency Comparison on Agentic Traces](https://vllm.ai/blog-assets/figures/2026-05-06-mooncake-store/pd_compare_mooncake_vs_nixl.png) + +The distributed KV cache pool improves vLLM throughput by 3.8x and reduces P50 TTFT and E2E latency by 46x and 8.6x, respectively. These gains are driven by a dramatic increase in cache hit rate: from 1.7%, where only the system prompt is cached, to 92.2%, where nearly the entire prefix is cached. + +## Scaling out to multiple nodes + +Experiment settings: + +* 20K common tokens (system instructions) +* 10K tokens first input +* 2,048 tokens per-turn input length +* 900 output tokens +* 30 turns total +* Number of sessions scaled with number of GPUs: 75 → 150 → 225 → 300 → 375 +* Parameters were chosen to roughly align with the original Codex workload and keep the total output/input ratio ~1.3% + +![Scaling Performance across Multiple Nodes](https://vllm.ai/blog-assets/figures/2026-05-06-mooncake-store/pd_scaling.png) + +To stress-test the datapath under cross-node traffic, we used round-robin routing. As a result, requests could be scheduled on different nodes across turns and often needed to fetch KV caches from a previous node. + +Without a distributed KV cache pool, this routing pattern would cause massive cache misses and severe throughput degradation. With Mooncake Store, vLLM consistently achieves a cache hit rate above 95%, and the system scales nearly linearly to 60 GPUs. + +This result shows that the distributed KV cache pool substantially improves cache hit rate while maintaining an efficient datapath as the cluster grows. + + +## Benchmark Scripts + +The benchmark scripts are provided in the artifact repository [here](https://github.com/ivanium/vllm/tree/feat/mooncake-store-int/scripts/mooncake/artifacts). diff --git a/docs/source/performance/vllm-v1-support-benchmark.md b/docs/source/performance/vllm/vllm-v1-pd-performance.md similarity index 90% rename from docs/source/performance/vllm-v1-support-benchmark.md rename to docs/source/performance/vllm/vllm-v1-pd-performance.md index d6cc7c578f..8780b3a4e5 100644 --- a/docs/source/performance/vllm-v1-support-benchmark.md +++ b/docs/source/performance/vllm/vllm-v1-pd-performance.md @@ -1,6 +1,6 @@ -# vLLM with Mooncake Transfer Engine Benchmark +# vLLM PD Disaggregation Performance -Mooncake has now implemented a vLLM connector, enabling direct support for the Prefill-Decode (PD) separation architecture in vLLM v1. We evaluated the performance of this integration, focusing on the efficiency of cross-node KV cache transfer using RDMA. +Mooncake has now implemented a vLLM connector, enabling direct support for the Prefill-Decode (PD) disaggregation architecture in vLLM v1. We evaluated the performance of this integration, focusing on the efficiency of cross-node KV cache transfer using RDMA. ## Benchmark Result @@ -8,7 +8,7 @@ Mooncake has now implemented a vLLM connector, enabling direct support for the P We measured the actual transfer bandwidth during the execution of requests with varying prompt lengths. -![KV Transfer Bandwidth (Actual)](../image/vllm_benchmark_actual_bandwidth.png) +![KV Transfer Bandwidth (Actual)](../../image/vllm_benchmark_actual_bandwidth.png) In a 1P1D (1 Prefiller, 1 Decoder) configuration using the Qwen3-8B model, Mooncake achieved a peak actual transfer bandwidth of **142.25 GB/s**. Given the theoretical maximum bandwidth of approximately 200 GB/s for the 8x RoCE connections, this represents a **71.1% bandwidth utilization rate**. This efficiency demonstrates that the custom transfer protocol and GPU Direct RDMA capabilities can effectively saturate high-performance networks. @@ -16,9 +16,9 @@ In a 1P1D (1 Prefiller, 1 Decoder) configuration using the Qwen3-8B model, Moonc We analyzed the Time To First Token (TTFT) to understand the impact of KV transfer overhead on end-to-end latency. -![TTFT Breakdown](../image/vllm_benchmark_ttft_breakdown.png) +![TTFT Breakdown](../../image/vllm_benchmark_ttft_breakdown.png) -![Transfer Time vs KV Size](../image/vllm_benchmark_transfer_time.png) +![Transfer Time vs KV Size](../../image/vllm_benchmark_transfer_time.png) The results show that Mooncake's high-speed transfer ensures that the overhead of moving KV cache is negligible compared to the computation time. For a prompt length of 32,768 tokens (transferring 4.50 GB of data), the actual KV transfer took only **31.65 ms**, accounting for merely **4.2%** of the total TTFT. diff --git a/docs/source/troubleshooting/error-code.md b/docs/source/troubleshooting/error-code.md index c4ee8e4349..2b8b19f5d3 100644 --- a/docs/source/troubleshooting/error-code.md +++ b/docs/source/troubleshooting/error-code.md @@ -49,6 +49,7 @@ Mooncake Store may generate various types of errors during execution. For most A | | OBJECT_HAS_LEASE (-706) | Object has lease | | | LEASE_EXPIRED (-707) | Lease expired before data transfer completed | | Transfer | TRANSFER_FAIL (-800) | Transfer operation failed | +| Checksum | CHECKSUM_MISMATCH (-801) | Retrieved object data does not match its stored checksum | | RPC | RPC_FAIL (-900) | RPC operation failed | | High Availability | ETCD_OPERATION_ERROR (-1000) | etcd operation failed | | | ETCD_KEY_NOT_EXIST (-1001) | Key not found in etcd | @@ -62,4 +63,7 @@ Mooncake Store may generate various types of errors during execution. For most A | | FILE_WRITE_FAIL (-1103) | Error writing file | | | FILE_INVALID_BUFFER (-1104) | File buffer is wrong | | | FILE_LOCK_FAIL (-1105) | File lock operation failed | -| | FILE_INVALID_HANDLE (-1106) | Invalid file handle | \ No newline at end of file +| | FILE_INVALID_HANDLE (-1106) | Invalid file handle | +| Task / Job | TASK_NOT_FOUND (-1400) | Task ID not found, or a completed task has already been pruned from the master's in-memory history | +| | TASK_PENDING_LIMIT_EXCEEDED (-1401) | The master-side pending task queue is full and cannot accept another task | +| | JOB_NOT_FOUND (-1402) | Job ID not found | diff --git a/docs/source/troubleshooting/index.md b/docs/source/troubleshooting/index.md new file mode 100644 index 0000000000..ba51937a5c --- /dev/null +++ b/docs/source/troubleshooting/index.md @@ -0,0 +1,13 @@ +--- +orphan: true +--- + +# Troubleshooting + +Diagnose common Mooncake deployment, runtime, and distributed execution issues. + +| Guide | Description | +|-------|-------------| +| [Troubleshooting Guide](troubleshooting) | Common setup and runtime problems. | +| [Error Codes](error-code) | Error-code meanings and recommended actions. | +| [PG/EP Troubleshooting](pg-ep-troubleshooting) | Diagnose Mooncake PG and EP failures. | diff --git a/docs/source/troubleshooting/troubleshooting.md b/docs/source/troubleshooting/troubleshooting.md index 3bf11c2e4c..9ddf6ce0e9 100644 --- a/docs/source/troubleshooting/troubleshooting.md +++ b/docs/source/troubleshooting/troubleshooting.md @@ -8,6 +8,21 @@ This document lists common errors that may occur when using Mooncake Store and p > - [ ] Incorrect RDMA device name and connection status is not active. > - [ ] etcd is not started normally and it is not bind with `0.0.0.0`. +## Corrupted Data or Garbled Output + +Use object-level checksum diagnostics when a full-object Mooncake Store read returns corrupted data or the application produces garbled output that may originate from stored data. Deploy checksum-capable Mooncake Store client, primary master, and standby master binaries from the same version, then set `MOONCAKE_STORE_CHECKSUM=1` before starting every writer and reader client process: + +```bash +export MOONCAKE_STORE_CHECKSUM=1 +``` + +Reproduce the issue with full-object `put`/`upsert` and `get` operations. The writer computes a CRC-64 checksum over the source object before transfer, and the reader verifies the logical `object_size` bytes returned by `get`. Range reads, including `get_into_ranges`, are not verified. + +`CHECKSUM_MISMATCH` (-801) proves that the bytes returned by the covered `get` differ from the bytes checksummed before the corresponding write. Treat the read as failed and do not use the destination buffer. It does not identify whether the corruption occurred during transfer, storage, or another covered Store stage. + +An enabled reader skips verification when the object's metadata has no checksum, such as an object written by a client with the switch disabled or restored from an older snapshot. This case is logged as `object_checksum_absent` at VLOG(1). A successful read without a mismatch therefore does not prove that checksum verification occurred, and a verified Store read does not rule out corruption introduced elsewhere in the application. + +Checksum diagnostics add a full data scan to writes and reads, perform device-to-host staging for GPU buffers, and disable the local hot cache. Disable the switch after diagnosis. See the [Mooncake Store Deployment and Tuning Guide](../deployment/mooncake-store-deployment-guide.md) for deployment and snapshot compatibility details. ## Metadata and Out-of-Band Communication 1. At startup, a `TransferMetadata` object is constructed according to the incoming `metadata_server` parameter. During program execution, this object is used to communicate with the etcd server to maintain internal data required for connection. diff --git a/docs/source/zh_archive/mooncake-store.md b/docs/source/zh_archive/mooncake-store.md index 768b310a3d..aade0b725d 100644 --- a/docs/source/zh_archive/mooncake-store.md +++ b/docs/source/zh_archive/mooncake-store.md @@ -612,7 +612,7 @@ virtual tl::expected, ErrorCode> Allocate( ### 替换策略 -当 `PutStart` 请求因内存不足而失败,或者当后台线程检测到空间使用率达到配置的高水位线(默认 95%,可通过 `-eviction_high_watermark_ratio` 配置)时,会触发一次替换任务,通过换出一部分对象来释放空间(默认 5%,可通过 `-eviction_ratio` 配置)。与 `Remove` 类似,被换出的对象仅仅会被标记为已删除,不需要进行数据传输。 +当 `PutStart` 请求因内存不足而失败,或者当后台线程检测到空间使用率达到配置的高水位线(默认 90%,可通过 `-eviction_high_watermark_ratio` 配置)时,会触发一次替换任务,通过换出一部分对象来释放空间(默认 5%,可通过 `-eviction_ratio` 配置)。与 `Remove` 类似,被换出的对象仅仅会被标记为已删除,不需要进行数据传输。 目前采用的是一种近似的 LRU 策略,即尽可能优先换出最近最少被访问的对象。为了避免数据竞争和数据损坏,正在被客户端读取或写入的对象不会被换出。因此,拥有租约或尚未被 `PutEnd` 请求标记为 complete 的对象不会被换出。 @@ -622,7 +622,7 @@ virtual tl::expected, ErrorCode> Allocate( 然而,如果在 `Get` 操作完成读取数据之前租约已过期,该操作将被视为失败,并且不会返回任何数据,以防止潜在的数据损坏。 -默认的租约时间为 5 秒,并可通过 `master_service` 的启动参数进行配置。 +默认的租约时间为 10 秒,并可通过 `master_service` 的启动参数进行配置。 ### 软固定机制 @@ -747,7 +747,7 @@ HTTP 元数据服务器可通过以下参数进行配置: ## Mooncake Store Python API -**完整的 Python API 文档**: [https://kvcache-ai.github.io/Mooncake/python-api-reference/mooncake-store.html](https://kvcache-ai.github.io/Mooncake/python-api-reference/mooncake-store.html) +**完整的 Python API 文档**: [https://kvcache-ai.github.io/Mooncake/api-reference/python/mooncake-store.html](https://kvcache-ai.github.io/Mooncake/api-reference/python/mooncake-store.html) ## 编译及使用方法 diff --git a/docs/source/zh_archive/transfer-engine.md b/docs/source/zh_archive/transfer-engine.md index 180fa44f2c..cf3b156308 100644 --- a/docs/source/zh_archive/transfer-engine.md +++ b/docs/source/zh_archive/transfer-engine.md @@ -394,7 +394,7 @@ int init(const std::string &metadata_conn_string, - `MC_NUM_COMP_CHANNELS_PER_CTX` 每个设备实例创建的 Completion Channel 数量,默认值 1 - `MC_IB_PORT` 每个设备实例使用的 IB 端口号,默认值 1 - `MC_IB_TC` 当使用`RDMA`通信协议时,在交换机和网卡默认配置不一致场景/需要流量规划场景下,可能需要修改 RDMA 网卡的 Traffic Class 配置,默认值 -1 -- `MC_IB_PCI_RELAXED_ORDERING` 将网络适配器的PCIe顺序设置为放宽有时会带来更好的性能。可设置 1 以启用RO功能,默认值 0 +- `MC_IB_PCI_RELAXED_ORDERING` 控制 RDMA 内存区域的 PCIe Relaxed Ordering(RO)功能。`0`:禁用,`1`:硬件支持时启用(默认),`2`:自动。需要 `ibv_reg_mr_iova2`(libibverbs ≥ 1.8),不支持时自动回退到严格顺序模式。 - `MC_GID_INDEX` 每个设备实例使用的 GID 序号,默认值 3(或平台支持的最大值) - `MC_PKEY_INDEX` QP 转换到 INIT 状态时使用的 `pkey_index`(partition key 表索引)。有效范围:0 到 65535,默认值 0。当 fabric 所需的 partition key 不在 HCA pkey 表的 0 号位置时需要设置该值 - `MC_MAX_CQE_PER_CTX` 每个设备实例中 CQ 缓冲区大小,默认值 4096 diff --git a/extern/yalantinglibs b/extern/yalantinglibs index 6a0e067d9a..7801bc9ad9 160000 --- a/extern/yalantinglibs +++ b/extern/yalantinglibs @@ -1 +1 @@ -Subproject commit 6a0e067d9a43492cf8e4e280b531924fbd724dbd +Subproject commit 7801bc9ad9021781f15217552214e325a1cf7373 diff --git a/image/hardwares/MetaX_logo.png b/image/hardwares/MetaX_logo.png new file mode 100755 index 0000000000..307e2344f7 Binary files /dev/null and b/image/hardwares/MetaX_logo.png differ diff --git a/image/hardwares/T-Head_logo.png b/image/hardwares/T-Head_logo.png new file mode 100755 index 0000000000..0d001c9ec3 Binary files /dev/null and b/image/hardwares/T-Head_logo.png differ diff --git a/image/hardwares/biren_logo.png b/image/hardwares/biren_logo.png new file mode 100644 index 0000000000..329580e556 Binary files /dev/null and b/image/hardwares/biren_logo.png differ diff --git a/image/hardwares/cambricon_logo.png b/image/hardwares/cambricon_logo.png new file mode 100755 index 0000000000..56588e74c6 Binary files /dev/null and b/image/hardwares/cambricon_logo.png differ diff --git a/image/partners/hygon_logo.png b/image/partners/hygon_logo.png new file mode 100644 index 0000000000..4d418152a1 Binary files /dev/null and b/image/partners/hygon_logo.png differ diff --git a/image/partners/nvidia_logo.png b/image/partners/nvidia_logo.png index 5d51569efd..70ec4de270 100644 Binary files a/image/partners/nvidia_logo.png and b/image/partners/nvidia_logo.png differ diff --git a/image/partners/sunrise_logo.png b/image/partners/sunrise_logo.png new file mode 100644 index 0000000000..c4b1f9215d Binary files /dev/null and b/image/partners/sunrise_logo.png differ diff --git a/monitoring/grafana/dashboards/mooncake.json b/monitoring/grafana/dashboards/mooncake.json index 2c497d1c56..77933094f5 100644 --- a/monitoring/grafana/dashboards/mooncake.json +++ b/monitoring/grafana/dashboards/mooncake.json @@ -1,67 +1,909 @@ { - "annotations": { - "list": [ - { - "builtIn": 1, - "datasource": "-- Grafana --", - "enable": true, - "hide": true, - "iconColor": "rgba(0, 211, 255, 1)", - "name": "Annotations & Alerts", - "type": "dashboard" - } - ] - }, + "annotations": { "list": [ { "builtIn": 1, "datasource": "-- Grafana --", "enable": true, "hide": true, "iconColor": "rgba(0, 211, 255, 1)", "name": "Annotations & Alerts", "type": "dashboard" } ] }, "editable": true, "gnetId": null, + "refresh": "5s", "graphTooltip": 0, "links": [], "panels": [ { - "title": "RPC Requests", - "type": "graph", + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 0 }, + "id": 100, + "panels": [], + "title": "Cluster Overview", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 4, "x": 0, "y": 1 }, + "id": 101, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_active_clients", "legendFormat": "", "refId": "A" }], + "title": "Active Clients", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null }, { "color": "orange", "value": 1000000 }, { "color": "red", "value": 10000000 } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 4, "x": 4, "y": 1 }, + "id": 102, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_key_count", "legendFormat": "Keys", "refId": "A" }], + "title": "Total Keys", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 4, "x": 8, "y": 1 }, + "id": 103, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_soft_pin_key_count", "legendFormat": "Soft Pinned", "refId": "A" }], + "title": "Soft Pinned Keys", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "bytes" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 1 }, + "id": 104, + "options": { "legend": { "calcs": ["mean", "lastNotNull", "max"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "master_allocated_bytes", "legendFormat": "Memory Allocated", "refId": "A" }, + { "expr": "master_total_capacity_bytes", "legendFormat": "Memory Capacity", "refId": "B" } + ], + "title": "Memory Storage", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "max": 100, "min": 0, "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null }, { "color": "orange", "value": 70 }, { "color": "red", "value": 90 } ] }, "unit": "percent" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 4, "x": 0, "y": 5 }, + "id": 105, + "options": { "colorMode": "background", "graphMode": "none", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "100 * master_allocated_bytes / master_total_capacity_bytes", "legendFormat": "", "refId": "A" }], + "title": "Memory Usage %", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "max": 100, "min": 0, "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null }, { "color": "orange", "value": 70 }, { "color": "red", "value": 90 } ] }, "unit": "percent" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 4, "x": 4, "y": 5 }, + "id": 106, + "options": { "colorMode": "background", "graphMode": "none", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "100 * master_nof_allocated_bytes / master_total_nof_capacity_bytes", "legendFormat": "", "refId": "A" }], + "title": "NoF SSD Usage %", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "max": 100, "min": 0, "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null }, { "color": "orange", "value": 70 }, { "color": "red", "value": 90 } ] }, "unit": "percent" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 4, "x": 8, "y": 5 }, + "id": 107, + "options": { "colorMode": "background", "graphMode": "none", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "100 * master_allocated_file_size_bytes / master_total_file_capacity_bytes", "legendFormat": "", "refId": "A" }], + "title": "File Storage Usage %", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "bytes" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 9 }, + "id": 108, + "options": { "legend": { "calcs": ["mean", "lastNotNull", "max"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "master_nof_allocated_bytes", "legendFormat": "NoF SSD Allocated", "refId": "A" }, + { "expr": "master_total_nof_capacity_bytes", "legendFormat": "NoF SSD Capacity", "refId": "B" } + ], + "title": "NoF SSD Storage", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 17 }, + "id": 200, + "panels": [], + "title": "RPC QPS — Core Operations", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 0, "y": 18 }, + "id": 201, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_put_start_requests_total[5m])", "legendFormat": "PutStart", "refId": "A" }, + { "expr": "rate(master_put_end_requests_total[5m])", "legendFormat": "PutEnd", "refId": "B" }, + { "expr": "rate(master_put_revoke_requests_total[5m])", "legendFormat": "PutRevoke", "refId": "C" } + ], + "title": "Put Operations QPS", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 8, "y": 18 }, + "id": 202, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_get_replica_list_requests_total[5m])", "legendFormat": "GetReplicaList", "refId": "A" }, + { "expr": "rate(master_get_replica_list_by_regex_requests_total[5m])", "legendFormat": "GetReplicaListByRegex", "refId": "B" }, + { "expr": "rate(master_exist_key_requests_total[5m])", "legendFormat": "ExistKey", "refId": "C" } + ], + "title": "Get / Query Operations QPS", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 16, "y": 18 }, + "id": 203, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_remove_requests_total[5m])", "legendFormat": "Remove", "refId": "A" }, + { "expr": "rate(master_remove_by_regex_requests_total[5m])", "legendFormat": "RemoveByRegex", "refId": "B" }, + { "expr": "rate(master_remove_all_requests_total[5m])", "legendFormat": "RemoveAll", "refId": "C" }, + { "expr": "rate(master_ping_requests_total[5m])", "legendFormat": "Ping", "refId": "D" } + ], + "title": "Remove / Ping Operations QPS", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 0, "y": 26 }, + "id": 204, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_put_start_failures_total[5m])", "legendFormat": "PutStart Fail", "refId": "A" }, + { "expr": "rate(master_put_start_alloc_failures_total[5m])", "legendFormat": "PutStart Alloc Fail", "refId": "B" }, + { "expr": "rate(master_put_end_failures_total[5m])", "legendFormat": "PutEnd Fail", "refId": "C" }, + { "expr": "rate(master_put_revoke_failures_total[5m])", "legendFormat": "PutRevoke Fail", "refId": "D" } + ], + "title": "Put Failure Rates", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 8, "y": 26 }, + "id": 205, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_get_replica_list_failures_total[5m])", "legendFormat": "GetReplicaList Fail", "refId": "A" }, + { "expr": "rate(master_get_replica_list_by_regex_failures_total[5m])", "legendFormat": "GetByRegex Fail", "refId": "B" }, + { "expr": "rate(master_exist_key_failures_total[5m])", "legendFormat": "ExistKey Fail", "refId": "C" }, + { "expr": "rate(master_remove_failures_total[5m])", "legendFormat": "Remove Fail", "refId": "D" }, + { "expr": "rate(master_ping_failures_total[5m])", "legendFormat": "Ping Fail", "refId": "E" } + ], + "title": "Get / Remove / Ping Failure Rates", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 16, "y": 26 }, + "id": 206, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_remove_by_regex_failures_total[5m])", "legendFormat": "RemoveByRegex Fail", "refId": "A" }, + { "expr": "rate(master_remove_all_failures_total[5m])", "legendFormat": "RemoveAll Fail", "refId": "B" } + ], + "title": "Other Failure Rates", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 34 }, + "id": 300, + "panels": [], + "title": "Segment Lifecycle Ops", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 35 }, + "id": 301, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_mount_segment_requests_total[5m])", "legendFormat": "Mount MEM", "refId": "A" }, + { "expr": "rate(master_unmount_segment_requests_total[5m])", "legendFormat": "Unmount MEM", "refId": "B" }, + { "expr": "rate(master_remount_segment_requests_total[5m])", "legendFormat": "Remount MEM", "refId": "C" } + ], + "title": "Memory Segment Mount / Unmount / Remount Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 35 }, + "id": 302, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_mount_nof_segment_requests_total[5m])", "legendFormat": "Mount NoF", "refId": "A" }, + { "expr": "rate(master_unmount_nof_segment_requests_total[5m])", "legendFormat": "Unmount NoF", "refId": "B" }, + { "expr": "rate(master_remount_nof_segment_requests_total[5m])", "legendFormat": "Remount NoF", "refId": "C" } + ], + "title": "NoF SSD Segment Mount / Unmount / Remount Rate", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 43 }, + "id": 400, + "panels": [], + "title": "Copy / Move / Task Operations", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 0, "y": 44 }, + "id": 401, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_copy_start_requests_total[5m])", "legendFormat": "CopyStart", "refId": "A" }, + { "expr": "rate(master_copy_end_requests_total[5m])", "legendFormat": "CopyEnd", "refId": "B" }, + { "expr": "rate(master_copy_revoke_requests_total[5m])", "legendFormat": "CopyRevoke", "refId": "C" } + ], + "title": "Copy Protocol Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 8, "y": 44 }, + "id": 402, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_move_start_requests_total[5m])", "legendFormat": "MoveStart", "refId": "A" }, + { "expr": "rate(master_move_end_requests_total[5m])", "legendFormat": "MoveEnd", "refId": "B" }, + { "expr": "rate(master_move_revoke_requests_total[5m])", "legendFormat": "MoveRevoke", "refId": "C" } + ], + "title": "Move Protocol Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 16, "y": 44 }, + "id": 403, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_create_copy_task_requests_total[5m])", "legendFormat": "CreateCopyTask", "refId": "A" }, + { "expr": "rate(master_create_move_task_requests_total[5m])", "legendFormat": "CreateMoveTask", "refId": "B" }, + { "expr": "rate(master_query_task_requests_total[5m])", "legendFormat": "QueryTask", "refId": "C" }, + { "expr": "rate(master_fetch_tasks_requests_total[5m])", "legendFormat": "FetchTasks", "refId": "D" }, + { "expr": "rate(master_update_task_requests_total[5m])", "legendFormat": "MarkTaskComplete", "refId": "E" }, + { "expr": "rate(master_evict_disk_replica_requests_total[5m])", "legendFormat": "EvictDiskReplica", "refId": "F" } + ], + "title": "Task Management Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 24, "x": 0, "y": 52 }, + "id": 404, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_copy_start_failures_total[5m])", "legendFormat": "CopyStart Fail", "refId": "A" }, + { "expr": "rate(master_copy_end_failures_total[5m])", "legendFormat": "CopyEnd Fail", "refId": "B" }, + { "expr": "rate(master_move_start_failures_total[5m])", "legendFormat": "MoveStart Fail", "refId": "C" }, + { "expr": "rate(master_move_end_failures_total[5m])", "legendFormat": "MoveEnd Fail", "refId": "D" }, + { "expr": "rate(master_create_copy_task_failures_total[5m])", "legendFormat": "CreateCopyTask Fail", "refId": "E" }, + { "expr": "rate(master_create_move_task_failures_total[5m])", "legendFormat": "CreateMoveTask Fail", "refId": "F" } + ], + "title": "Copy / Move / Task Failure Rates", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 60 }, + "id": 500, + "panels": [], + "title": "Eviction Statistics", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 0, "y": 61 }, + "id": 501, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_attempted_evictions_total[5m])", "legendFormat": "Total Attempts", "refId": "A" }, + { "expr": "rate(master_successful_evictions_total[5m])", "legendFormat": "Total Success", "refId": "B" } + ], + "title": "Total Eviction Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 8, "y": 61 }, + "id": 502, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_attempted_evictions_mem[5m])", "legendFormat": "MEM Attempts", "refId": "A" }, + { "expr": "rate(master_successful_evictions_mem[5m])", "legendFormat": "MEM Success", "refId": "B" }, + { "expr": "rate(master_attempted_evictions_nof[5m])", "legendFormat": "NoF Attempts", "refId": "C" }, + { "expr": "rate(master_successful_evictions_nof[5m])", "legendFormat": "NoF Success", "refId": "D" } + ], + "title": "MEM vs NoF Eviction Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 16, "y": 61 }, + "id": 503, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_evicted_key_count[5m])", "legendFormat": "Total Keys", "refId": "A" }, + { "expr": "rate(master_evicted_key_count_mem[5m])", "legendFormat": "MEM Keys", "refId": "B" }, + { "expr": "rate(master_evicted_key_count_nof[5m])", "legendFormat": "NoF Keys", "refId": "C" } + ], + "title": "Evicted Keys Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "Bps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 24, "x": 0, "y": 69 }, + "id": 504, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_evicted_size_bytes[5m])", "legendFormat": "Total Bytes", "refId": "A" }, + { "expr": "rate(master_evicted_size_bytes_mem[5m])", "legendFormat": "MEM Bytes", "refId": "B" }, + { "expr": "rate(master_evicted_size_bytes_nof[5m])", "legendFormat": "NoF Bytes", "refId": "C" } + ], + "title": "Evicted Bytes Rate", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 77 }, + "id": 600, + "panels": [], + "title": "Promotion-on-Hit", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null }, { "color": "orange", "value": 100 }, { "color": "red", "value": 1000 } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 3, "x": 0, "y": 78 }, + "id": 601, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_promotion_in_flight", "legendFormat": "", "refId": "A" }], + "title": "In-Flight", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 3, "x": 3, "y": 78 }, + "id": 602, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_promotion_admitted_total", "legendFormat": "", "refId": "A" }], + "title": "Admitted", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 3, "x": 6, "y": 78 }, + "id": 603, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_promotion_completed_total", "legendFormat": "", "refId": "A" }], + "title": "Completed", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "red", "value": null }, { "color": "green", "value": 0 } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 3, "x": 9, "y": 78 }, + "id": 604, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_promotion_failed_total", "legendFormat": "", "refId": "A" }], + "title": "Failed", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "orange", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 3, "x": 12, "y": 78 }, + "id": 605, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_promotion_cancelled_total", "legendFormat": "", "refId": "A" }], + "title": "Cancelled", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "orange", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 3, "x": 15, "y": 78 }, + "id": 606, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_promotion_expired_total", "legendFormat": "", "refId": "A" }], + "title": "Expired", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "blue", "value": null } ] }, "unit": "decbytes" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 3, "x": 18, "y": 78 }, + "id": 607, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_promotion_completed_bytes_total", "legendFormat": "", "refId": "A" }], + "title": "Completed Bytes", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "purple", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 3, "x": 21, "y": 78 }, + "id": 608, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [ + { "expr": "master_promotion_rejected_frequency_total + master_promotion_rejected_watermark_total + master_promotion_rejected_cap_total", "legendFormat": "", "refId": "A" } + ], + "title": "Total Rejected", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 82 }, + "id": 609, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_promotion_admitted_total[5m])", "legendFormat": "Admitted", "refId": "A" }, + { "expr": "rate(master_promotion_completed_total[5m])", "legendFormat": "Completed", "refId": "B" }, + { "expr": "rate(master_promotion_failed_total[5m])", "legendFormat": "Failed", "refId": "C" }, + { "expr": "rate(master_promotion_cancelled_total[5m])", "legendFormat": "Cancelled", "refId": "D" }, + { "expr": "rate(master_promotion_expired_total[5m])", "legendFormat": "Expired", "refId": "E" } + ], + "title": "Promotion Lifecycle Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 82 }, + "id": 610, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_promotion_rejected_frequency_total[5m])", "legendFormat": "Rejected: Freq", "refId": "A" }, + { "expr": "rate(master_promotion_rejected_watermark_total[5m])", "legendFormat": "Rejected: Watermark", "refId": "B" }, + { "expr": "rate(master_promotion_rejected_cap_total[5m])", "legendFormat": "Rejected: Cap", "refId": "C" } + ], + "title": "Promotion Rejection Rate", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 90 }, + "id": 700, + "panels": [], + "title": "Cache Hit Statistics", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 0, "y": 91 }, + "id": 701, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(mem_cache_hit_nums_[5m])", "legendFormat": "MEM Hits", "refId": "A" }, + { "expr": "rate(file_cache_hit_nums_[5m])", "legendFormat": "SSD Hits", "refId": "B" } + ], + "title": "Cache Hit Count Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "Bps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 8, "y": 91 }, + "id": 702, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(mem_cache_hit_bytes_total[5m])", "legendFormat": "MEM Bytes", "refId": "A" }, + { "expr": "rate(file_cache_hit_bytes_total[5m])", "legendFormat": "SSD Bytes", "refId": "B" } + ], + "title": "Cache Hit Bytes Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 16, "y": 91 }, + "id": 703, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "mem_cache_nums_", "legendFormat": "MEM Objects", "refId": "A" }, + { "expr": "file_cache_nums_", "legendFormat": "SSD Objects", "refId": "B" } + ], + "title": "Current Cached Objects", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "max": 100, "min": 0, "thresholds": { "mode": "absolute", "steps": [ { "color": "red", "value": null }, { "color": "orange", "value": 50 }, { "color": "green", "value": 80 } ] }, "unit": "percent" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 6, "x": 0, "y": 99 }, + "id": 704, + "options": { "colorMode": "background", "graphMode": "none", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "100 * rate(mem_cache_hit_nums_[5m]) / clamp_min(rate(mem_cache_hit_nums_[5m]) + rate(file_cache_hit_nums_[5m]), 1e-9)", "legendFormat": "", "refId": "A" }], + "title": "MEM Hit Ratio % (Store-observed)", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "max": 100, "min": 0, "thresholds": { "mode": "absolute", "steps": [ { "color": "red", "value": null }, { "color": "orange", "value": 80 }, { "color": "green", "value": 95 } ] }, "unit": "percent" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 6, "x": 6, "y": 99 }, + "id": 705, + "options": { "colorMode": "background", "graphMode": "none", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "100 * rate(valid_get_nums_[5m]) / clamp_min(rate(total_get_nums_[5m]), 1e-9)", "legendFormat": "", "refId": "A" }], + "title": "Valid Get Rate %", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 99 }, + "id": 706, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(valid_get_nums_[5m])", "legendFormat": "Valid Gets", "refId": "A" }, + { "expr": "rate(total_get_nums_[5m])", "legendFormat": "Total Gets", "refId": "B" } + ], + "title": "Get Operations (Valid vs Total)", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 107 }, + "id": 800, + "panels": [], + "title": "NoF Heartbeat & Health", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 8, "x": 0, "y": 108 }, + "id": 801, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_nof_heartbeat_success_total[5m])", "legendFormat": "Success", "refId": "A" }, + { "expr": "rate(master_nof_heartbeat_failure_total[5m])", "legendFormat": "Failure", "refId": "B" }, + { "expr": "rate(master_nof_heartbeat_timeout_total[5m])", "legendFormat": "Timeout", "refId": "C" } + ], + "title": "NoF Heartbeat Rate", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null }, { "color": "red", "value": 1 } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 4, "x": 8, "y": 108 }, + "id": 802, + "options": { "colorMode": "background", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_nof_segments_unmounted_by_heartbeat_total", "legendFormat": "", "refId": "A" }], + "title": "Unmounted by Heartbeat", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "ms", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "ms" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 108 }, + "id": 803, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "histogram_quantile(0.95, sum by (le) (rate(master_nof_heartbeat_probe_latency_ms_bucket[5m])))", "legendFormat": "p95", "refId": "A" } + ], + "title": "NoF Heartbeat Probe Latency (Histogram)", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 116 }, + "id": 900, + "panels": [], + "title": "Snapshot Operations", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 4, "x": 0, "y": 117 }, + "id": 901, + "options": { "colorMode": "value", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_snapshot_success", "legendFormat": "", "refId": "A" }], + "title": "Snapshot Success", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null }, { "color": "red", "value": 1 } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 4, "x": 4, "y": 117 }, + "id": 902, + "options": { "colorMode": "background", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_snapshot_fail", "legendFormat": "", "refId": "A" }], + "title": "Snapshot Failures", + "type": "stat" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "ms", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "ms" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 16, "x": 8, "y": 117 }, + "id": 903, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "histogram_quantile(0.95, sum by (le) (rate(master_snapshot_duration_ms_bucket[5m])))", "legendFormat": "p95", "refId": "A" } + ], + "title": "Snapshot Duration Distribution (Histogram)", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 125 }, + "id": 1000, + "panels": [], + "title": "PutStart Discard / Staging", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 126 }, + "id": 1001, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "master_put_start_discard_cnt", "legendFormat": "Discarded", "refId": "A" }, + { "expr": "master_put_start_release_cnt", "legendFormat": "Released", "refId": "B" } + ], + "title": "PutStart Discard / Release Cumulative", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "thresholds" }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null }, { "color": "orange", "value": 1073741824 }, { "color": "red", "value": 10737418240 } ] }, "unit": "bytes" }, + "overrides": [] + }, + "gridPos": { "h": 4, "w": 6, "x": 12, "y": 126 }, + "id": 1002, + "options": { "colorMode": "background", "graphMode": "area", "justifyMode": "auto", "orientation": "auto", "reduceOptions": { "calcs": ["lastNotNull"], "fields": "", "values": false }, "textMode": "auto" }, + "targets": [{ "expr": "master_put_start_discarded_staging_size", "legendFormat": "", "refId": "A" }], + "title": "Staging Size (Discarded not Released)", + "type": "stat" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 134 }, + "id": 1100, + "panels": [], + "title": "Batch Operations (Overview)", + "type": "row" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "reqps" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 0, "y": 135 }, + "id": 1101, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_batch_put_start_requests_total[5m])", "legendFormat": "BatchPutStart", "refId": "A" }, + { "expr": "rate(master_batch_put_end_requests_total[5m])", "legendFormat": "BatchPutEnd", "refId": "B" }, + { "expr": "rate(master_batch_put_revoke_requests_total[5m])", "legendFormat": "BatchPutRevoke", "refId": "C" }, + { "expr": "rate(master_batch_get_replica_list_requests_total[5m])", "legendFormat": "BatchGetReplicaList", "refId": "D" }, + { "expr": "rate(master_batch_exist_key_requests_total[5m])", "legendFormat": "BatchExistKey", "refId": "E" }, + { "expr": "rate(master_batch_query_ip_requests_total[5m])", "legendFormat": "BatchQueryIp", "refId": "F" }, + { "expr": "rate(master_batch_replica_clear_requests_total[5m])", "legendFormat": "BatchReplicaClear", "refId": "G" } + ], + "title": "Batch Operation Request Rates", + "type": "timeseries" + }, + { + "datasource": "Prometheus", + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "short" }, + "overrides": [] + }, + "gridPos": { "h": 8, "w": 12, "x": 12, "y": 135 }, + "id": 1102, + "options": { "legend": { "calcs": ["mean", "lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, + "targets": [ + { "expr": "rate(master_batch_put_start_items_total[5m])", "legendFormat": "BatchPutStart Items", "refId": "A" }, + { "expr": "rate(master_batch_put_end_items_total[5m])", "legendFormat": "BatchPutEnd Items", "refId": "B" }, + { "expr": "rate(master_batch_get_replica_list_items_total[5m])", "legendFormat": "BatchGetRL Items", "refId": "C" }, + { "expr": "rate(master_batch_exist_key_items_total[5m])", "legendFormat": "BatchExistKey Items", "refId": "D" }, + { "expr": "rate(master_batch_query_ip_items_total[5m])", "legendFormat": "BatchQueryIp Items", "refId": "E" }, + { "expr": "rate(master_batch_replica_clear_items_total[5m])", "legendFormat": "BatchReplicaClear Items", "refId": "F" } + ], + "title": "Batch Operation Item Throughput", + "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { "h": 1, "w": 24, "x": 0, "y": 143 }, + "id": 1200, + "panels": [], + "title": "Value Size Distribution (Histogram)", + "type": "row" + }, + { "datasource": "Prometheus", - "gridPos": { - "h": 9, - "w": 12, - "x": 0, - "y": 0 + "fieldConfig": { + "defaults": { "color": { "mode": "palette-classic" }, "custom": { "axisCenteredZero": false, "axisColorMode": "text", "axisLabel": "bytes", "axisPlacement": "auto", "barAlignment": 0, "drawStyle": "line", "fillOpacity": 10, "gradientMode": "none", "hideFrom": { "legend": false, "tooltip": false, "viz": false }, "lineInterpolation": "linear", "lineWidth": 1, "pointSize": 5, "scaleDistribution": { "type": "linear" }, "showPoints": "never", "spanNulls": false, "stacking": { "group": "A", "mode": "none" }, "thresholdsStyle": { "mode": "off" } }, "mappings": [], "thresholds": { "mode": "absolute", "steps": [ { "color": "green", "value": null } ] }, "unit": "bytes" }, + "overrides": [] }, - "id": 2, + "gridPos": { "h": 8, "w": 24, "x": 0, "y": 144 }, + "id": 1201, + "options": { "legend": { "calcs": ["lastNotNull"], "displayMode": "table", "placement": "bottom" }, "tooltip": { "mode": "multi", "sort": "desc" } }, "targets": [ - { - "expr": "rate(mooncake_master_rpc_requests_total[5m])", - "legendFormat": "{{method}}", - "refId": "A" - } - ] + { "expr": "histogram_quantile(0.95, sum by (le) (rate(master_value_size_bytes_bucket[5m])))", "legendFormat": "p95 bytes", "refId": "A" } + ], + "title": "Object Value Size p95", + "type": "timeseries" } ], - "schemaVersion": 22, + "schemaVersion": 36, "style": "dark", - "tags": [], - "templating": { - "list": [] - }, + "tags": ["mooncake", "master", "storage"], + "templating": { "list": [] }, "title": "Mooncake Master", - "time": { - "from": "now-6h", - "to": "now" - }, - "timepicker": { - "refresh_intervals": [ - "5s", - "10s", - "30s", - "1m", - "5m", - "15m", - "30m", - "1h", - "2h", - "1d" - ] - }, + "time": { "from": "now-6h", "to": "now" }, + "timepicker": { "refresh_intervals": ["5s", "10s", "30s", "1m", "5m", "15m", "30m", "1h", "2h", "1d"] }, "timezone": "browser", "version": 1 } diff --git a/mooncake-common/FindGTest.cmake b/mooncake-common/FindGTest.cmake new file mode 100644 index 0000000000..5fe684101b --- /dev/null +++ b/mooncake-common/FindGTest.cmake @@ -0,0 +1,26 @@ +include_guard(GLOBAL) + +include(FetchContent) + +if(POLICY CMP0135) + cmake_policy(SET CMP0135 NEW) +endif() + +# Keep test builds independent of the host distribution and Python environment. +# In particular, manylinux images may expose headers from /opt/conda while +# dependencies.sh installs an incompatible system gtest library. +set(BUILD_GMOCK + OFF + CACHE BOOL "Build GoogleMock" FORCE) +set(INSTALL_GTEST + OFF + CACHE BOOL "Install GoogleTest" FORCE) + +set(GOOGLETEST_VERSION 1.17.0) +FetchContent_Declare( + googletest + URL https://github.com/google/googletest/archive/refs/tags/v${GOOGLETEST_VERSION}.tar.gz + URL_HASH + SHA256=65fab701d9829d38cb77c14acdc431d2108bfdbf8979e40eb8ae567edf10b27c) + +FetchContent_MakeAvailable(googletest) diff --git a/mooncake-common/FindNCCLDevice.cmake b/mooncake-common/FindNCCLDevice.cmake new file mode 100644 index 0000000000..7581d0e1bc --- /dev/null +++ b/mooncake-common/FindNCCLDevice.cmake @@ -0,0 +1,106 @@ +# Locate NCCL 2.30.4+. The host backend needs nccl.h; the device backend also +# requires the experimental nccl_device.h API. +# +# Prefer NCCL's exported package so version metadata, include paths, library +# variants, CUDA, and thread dependencies come from NCCL itself. Keep a manual +# fallback for tarball and pre-package installations. + +set(_NCCL_DEVICE_MIN_VERSION "2.30.4") +set(_NCCL_DEVICE_USING_CONFIG FALSE) + +find_package(NCCL ${_NCCL_DEVICE_MIN_VERSION} CONFIG QUIET) +if(NCCL_FOUND AND TARGET NCCL::nccl) + set(_NCCL_DEVICE_USING_CONFIG TRUE) + set(NCCLDevice_VERSION "${NCCL_VERSION}") + find_path(NCCL_DEVICE_INCLUDE_DIR + NAMES nccl.h + HINTS ${NCCL_INCLUDE_DIRS} + NO_DEFAULT_PATH) +endif() + +if(NOT _NCCL_DEVICE_USING_CONFIG) + find_path(NCCL_DEVICE_INCLUDE_DIR + NAMES nccl.h + HINTS + ${NCCL_ROOT} + $ENV{NCCL_ROOT} + $ENV{NCCL_HOME} + PATH_SUFFIXES include build/include + PATHS + /usr/local/cuda + /usr/local + /usr) + + find_library(NCCL_DEVICE_LIBRARY + NAMES nccl + HINTS + ${NCCL_ROOT} + $ENV{NCCL_ROOT} + $ENV{NCCL_HOME} + PATH_SUFFIXES lib lib64 build/lib build/lib64 + PATHS + /usr/local/cuda + /usr/local + /usr) + + if(NCCL_DEVICE_INCLUDE_DIR AND + EXISTS "${NCCL_DEVICE_INCLUDE_DIR}/nccl.h") + foreach(_component MAJOR MINOR PATCH) + file(STRINGS "${NCCL_DEVICE_INCLUDE_DIR}/nccl.h" + _nccl_${_component}_line + REGEX "^#define NCCL_${_component}[ \t]+[0-9]+" + LIMIT_COUNT 1) + string(REGEX MATCH "[0-9]+$" _nccl_${_component} + "${_nccl_${_component}_line}") + endforeach() + if(NOT "${_nccl_MAJOR}" STREQUAL "" AND + NOT "${_nccl_MINOR}" STREQUAL "" AND + NOT "${_nccl_PATCH}" STREQUAL "") + set(NCCLDevice_VERSION + "${_nccl_MAJOR}.${_nccl_MINOR}.${_nccl_PATCH}") + endif() + endif() +endif() + +if(USE_NCCL_DEVICE AND NCCL_DEVICE_INCLUDE_DIR) + find_path(NCCL_DEVICE_API_INCLUDE_DIR + NAMES nccl_device.h + HINTS ${NCCL_DEVICE_INCLUDE_DIR} + NO_DEFAULT_PATH) +endif() + +include(FindPackageHandleStandardArgs) +set(_NCCL_DEVICE_REQUIRED_VARS NCCL_DEVICE_INCLUDE_DIR) +if(NOT _NCCL_DEVICE_USING_CONFIG) + list(APPEND _NCCL_DEVICE_REQUIRED_VARS NCCL_DEVICE_LIBRARY) +endif() +if(USE_NCCL_DEVICE) + list(APPEND _NCCL_DEVICE_REQUIRED_VARS NCCL_DEVICE_API_INCLUDE_DIR) +endif() +find_package_handle_standard_args(NCCLDevice + REQUIRED_VARS ${_NCCL_DEVICE_REQUIRED_VARS} + VERSION_VAR NCCLDevice_VERSION) + +if(NCCLDevice_FOUND) + find_package(CUDAToolkit REQUIRED) + + if(NOT TARGET NCCL::nccl) + add_library(NCCL::nccl UNKNOWN IMPORTED) + set_target_properties(NCCL::nccl PROPERTIES + IMPORTED_LOCATION "${NCCL_DEVICE_LIBRARY}" + INTERFACE_INCLUDE_DIRECTORIES "${NCCL_DEVICE_INCLUDE_DIR}" + INTERFACE_LINK_LIBRARIES "CUDA::cudart") + endif() + + if(USE_NCCL_DEVICE) + set_property(TARGET NCCL::nccl APPEND PROPERTY + INTERFACE_COMPILE_DEFINITIONS + NCCL_DEVICE_PERMIT_EXPERIMENTAL_CODE=1) + endif() +endif() + +mark_as_advanced(NCCL_DEVICE_INCLUDE_DIR NCCL_DEVICE_API_INCLUDE_DIR + NCCL_DEVICE_LIBRARY) +unset(_NCCL_DEVICE_REQUIRED_VARS) +unset(_NCCL_DEVICE_MIN_VERSION) +unset(_NCCL_DEVICE_USING_CONFIG) diff --git a/mooncake-common/FindUrma.cmake b/mooncake-common/FindUrma.cmake index 0af8d1a7ca..3966437a31 100644 --- a/mooncake-common/FindUrma.cmake +++ b/mooncake-common/FindUrma.cmake @@ -1,20 +1,46 @@ include(FetchContent) -# UMDK 头文件库 -FetchContent_Declare( - urma - GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git - GIT_TAG v25.12.0.B081 -) +# Prefer a system UMDK installation so configure works without downloading a +# second copy and the headers match the liburma ABI used at runtime. +find_path( + URMA_SYSTEM_INCLUDE_DIR + NAMES urma_api.h + PATHS /usr/include /usr/local/include + PATH_SUFFIXES urma umdk src/urma/lib/urma/core/include) +find_library( + URMA_LIBRARY + NAMES urma + PATHS /usr/lib /usr/lib64 /usr/local/lib /usr/local/lib64) -FetchContent_MakeAvailable(urma) +if(URMA_SYSTEM_INCLUDE_DIR) + set(urma_INCLUDE_DIR "${URMA_SYSTEM_INCLUDE_DIR}") +else() + # The source fallback supplies headers only. Production TENT UB remains + # disabled at runtime when no real liburma is present; tests inject their own + # adapter instead of defining a second set of global urma_* mock symbols. + FetchContent_Declare( + urma + GIT_REPOSITORY https://atomgit.com/openeuler/umdk.git + GIT_TAG v25.12.0.B081) + FetchContent_MakeAvailable(urma) + set(urma_INCLUDE_DIR "${urma_SOURCE_DIR}/src/urma/lib/urma/core/include") +endif() -# 输出实际路径,确认位置 -message(STATUS "URMA source dir: ${urma_SOURCE_DIR}") -message(STATUS "URMA binary dir: ${urma_BINARY_DIR}") +if(NOT TARGET Urma::urma) + add_library(Urma::urma INTERFACE IMPORTED GLOBAL) + set_property(TARGET Urma::urma PROPERTY INTERFACE_INCLUDE_DIRECTORIES + "${urma_INCLUDE_DIR}") + if(URMA_LIBRARY) + set_property(TARGET Urma::urma PROPERTY INTERFACE_LINK_LIBRARIES + "${URMA_LIBRARY}") + endif() +endif() -# 假设 UMDK 头文件在其 include 目录下 -set(urma_INCLUDE_DIR ${urma_SOURCE_DIR}/src/urma/lib/urma/core/include) - -# 添加到需要的目标 -message(STATUS "urma_INCLUDE_DIR: ${urma_INCLUDE_DIR}") \ No newline at end of file +set(URMA_INCLUDE_DIR "${urma_INCLUDE_DIR}") +set(URMA_FOUND TRUE) +message(STATUS "URMA include directory: ${URMA_INCLUDE_DIR}") +if(URMA_LIBRARY) + message(STATUS "URMA library: ${URMA_LIBRARY}") +else() + message(STATUS "URMA library not found; real UB backends will be unavailable") +endif() diff --git a/mooncake-common/common.cmake b/mooncake-common/common.cmake index eefa33bdb4..73c3110811 100644 --- a/mooncake-common/common.cmake +++ b/mooncake-common/common.cmake @@ -4,7 +4,8 @@ set(CMAKE_CUDA_STANDARD 20) option(ENABLE_DEBUG_SYMBOLS "Include debug symbols (-g) in compilation" ON) -set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wextra -Wno-unused-parameter -fPIC") +set(CMAKE_CXX_FLAGS + "${CMAKE_CXX_FLAGS} -Wall -Wextra -Wno-unused-parameter -fPIC") set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wall -Wextra -Wno-unused-parameter -fPIC") if(ENABLE_DEBUG_SYMBOLS) @@ -14,7 +15,8 @@ endif() if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fcoroutines") - set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -fno-tree-slp-vectorize") + set(CMAKE_CXX_FLAGS_RELEASE + "${CMAKE_CXX_FLAGS_RELEASE} -fno-tree-slp-vectorize") endif() set(CMAKE_C_FLAGS_RELEASE "-O3") @@ -29,7 +31,7 @@ endif() option(ENABLE_ASAN "enable address sanitizer" OFF) -if (ENABLE_ASAN) +if(ENABLE_ASAN) set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fsanitize=leak") set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsanitize=leak") set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fsanitize=address") @@ -37,7 +39,7 @@ if (ENABLE_ASAN) endif() # keep debuginfo by default -if (NOT CMAKE_BUILD_TYPE) +if(NOT CMAKE_BUILD_TYPE) set(CMAKE_BUILD_TYPE "RelWithDebInfo") endif() @@ -52,7 +54,7 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON) include(${CMAKE_CURRENT_LIST_DIR}/limit_jobs.cmake) option(ENABLE_SCCACHE "Whether to open sccache" OFF) -if (ENABLE_SCCACHE) +if(ENABLE_SCCACHE) find_program(SCCACHE sccache REQUIRED) endif() if(SCCACHE AND ENABLE_SCCACHE) @@ -67,44 +69,61 @@ add_compile_options(-fno-tree-slp-vectorize) option(BUILD_EXAMPLES "Build examples" ON) option(BUILD_UNIT_TESTS "Build unit tests" ON) +if(BUILD_UNIT_TESTS) + include(${CMAKE_CURRENT_LIST_DIR}/FindGTest.cmake) +endif() option(BUILD_BENCHMARK "Build benchmarks" ON) option(USE_CUDA "option for enabling gpu features for NVIDIA GPU" OFF) +option(USE_NCCL_DEVICE "option for enabling the NCCL DeviceTransport backend" OFF) +option(USE_NCCL_HOST "option for enabling the NCCL host RMA transport" OFF) option(USE_MLU "option for enabling Cambricon MLU features" OFF) option(USE_MUSA "option for enabling gpu features for MTHREADS GPU" OFF) option(USE_MACA "option for enabling gpu features for MUXI GPU with MACA" OFF) option(USE_HIP "option for enabling gpu features for AMD GPU" OFF) option(USE_HYGON "option for enabling gpu features for Hygon DCU with DTK" OFF) option(USE_COREX "option for enabling gpu features for Iluvatar CoreX" OFF) +option(USE_SUPA "option for enabling gpu features for Biren GPU with SUPA" OFF) option(USE_NVMEOF "option for using NVMe over Fabric" OFF) option(USE_TCP "option for using TCP transport" ON) option(USE_BAREX "option for using accl-barex transport" OFF) option(USE_ASCEND "option for using npu with HCCL" OFF) option(USE_ASCEND_DIRECT "option for using ascend npu with adxl engine" OFF) option(USE_UBSHMEM "option for using ascend npu with shmem" OFF) -option(USE_ASCEND_HETEROGENEOUS "option for transferring between ascend npu and gpu" OFF) +option(USE_ASCEND_HETEROGENEOUS + "option for transferring between ascend npu and gpu" OFF) option(USE_MNNVL "option for using Multi-Node NVLink transport" OFF) option(USE_CXL "option for using CXL protocol" OFF) option(USE_EFA "option for using AWS EFA transport" OFF) option(USE_UB "option for using UB protocol transport" OFF) -option(USE_SUNRISE "option for enabling gpu features for Sunrise GPU with Tang runtime" OFF) - -if (USE_UB) +option(USE_SUNRISE + "option for enabling gpu features for Sunrise GPU with Tang runtime" OFF) +option(USE_TPU + "option for enabling TPU (PJRT) staging support in TENT; the PJRT adapter is loaded at runtime via dlopen, no build-time SDK required" + OFF) +option(USE_VRAM_SEGMENT "option for vram segment" OFF) + +if(USE_UB) add_compile_definitions(USE_UB) message(STATUS "ub transport is enabled") include(${CMAKE_CURRENT_LIST_DIR}/FindUrma.cmake) endif() -if (USE_EFA) +if(USE_EFA) # Find libfabric headers and library; default to AWS EFA installer path - find_path(LIBFABRIC_INCLUDE_DIR rdma/fabric.h + find_path( + LIBFABRIC_INCLUDE_DIR rdma/fabric.h HINTS /opt/amazon/efa/include PATH_SUFFIXES include) - find_library(LIBFABRIC_LIBRARY fabric + find_library( + LIBFABRIC_LIBRARY fabric HINTS /opt/amazon/efa/lib PATH_SUFFIXES lib lib64) - if (NOT LIBFABRIC_INCLUDE_DIR OR NOT LIBFABRIC_LIBRARY) - message(FATAL_ERROR "libfabric not found. Install AWS EFA or set LIBFABRIC_INCLUDE_DIR/LIBFABRIC_LIBRARY.") + if(NOT LIBFABRIC_INCLUDE_DIR OR NOT LIBFABRIC_LIBRARY) + message( + FATAL_ERROR + "libfabric not found. Install AWS EFA or set LIBFABRIC_INCLUDE_DIR/LIBFABRIC_LIBRARY." + ) endif() get_filename_component(LIBFABRIC_LIB_DIR ${LIBFABRIC_LIBRARY} DIRECTORY) @@ -115,190 +134,297 @@ if (USE_EFA) message(STATUS " libfabric include: ${LIBFABRIC_INCLUDE_DIR}") message(STATUS " libfabric library: ${LIBFABRIC_LIBRARY}") endif() +if(USE_CXI) + # Find libfabric headers and library; default to AWS EFA installer path + find_path(LIBFABRIC_INCLUDE_DIR rdma/fabric.h PATH_SUFFIXES include) + find_library(LIBFABRIC_LIBRARY fabric PATH_SUFFIXES lib lib64) + + if(NOT LIBFABRIC_INCLUDE_DIR OR NOT LIBFABRIC_LIBRARY) + message( + FATAL_ERROR + "libfabric not found. Install AWS EFA or set LIBFABRIC_INCLUDE_DIR/LIBFABRIC_LIBRARY." + ) + endif() + + get_filename_component(LIBFABRIC_LIB_DIR ${LIBFABRIC_LIBRARY} DIRECTORY) + include_directories(${LIBFABRIC_INCLUDE_DIR}) + link_directories(${LIBFABRIC_LIB_DIR}) + add_compile_definitions(USE_CXI) + message(STATUS "HPE CXI (libfabric) transport is enabled") + message(STATUS " libfabric include: ${LIBFABRIC_INCLUDE_DIR}") + message(STATUS " libfabric library: ${LIBFABRIC_LIBRARY}") +endif() option(USE_ETCD "option for enable etcd as metadata server" OFF) option(USE_ETCD_LEGACY "option for enable etcd based on etcd-cpp-api-v3" OFF) option(USE_REDIS "option for enable redis as metadata server" OFF) option(USE_HTTP "option for enable http as metadata server" ON) -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_METRICS "enable metrics and metrics reporting thread" ON) option(USE_3FS "option for using 3FS storage backend" OFF) -option(USE_EVENT_DRIVEN_COMPLETION "option for using event-driven completion (store & transfer engine)" OFF) +option(USE_EVENT_DRIVEN_COMPLETION + "option for using event-driven completion (store & transfer engine)" OFF) option(USE_TENT "option for building Mooncake TENT" OFF) -option(ENABLE_MULTI_PROTOCOL "option for enabling multi-protocol support in transfer engine" OFF) -if (ENABLE_MULTI_PROTOCOL) - add_compile_definitions(ENABLE_MULTI_PROTOCOL) - message(STATUS "Multi-protocol support is enabled") +option(ENABLE_MULTI_PROTOCOL + "option for enabling multi-protocol support in transfer engine" OFF) +if(ENABLE_MULTI_PROTOCOL) + add_compile_definitions(ENABLE_MULTI_PROTOCOL) + message(STATUS "Multi-protocol support is enabled") endif() option(USE_LRU_MASTER "option for using LRU in master service" OFF) option(USE_INTRA_NVLINK "option for using IntraNode nvlink transport" OFF) -option(USE_MLX5DV "enable mlx5 direct verbs (libmlx5) for QP UDP source port override" OFF) +option(USE_MLX5DV + "enable mlx5 direct verbs (libmlx5) for QP UDP source port override" OFF) set(LRU_MAX_CAPACITY 1000) -if (USE_LRU_MASTER) +if(USE_LRU_MASTER) add_compile_definitions(USE_LRU_MASTER) add_compile_definitions(LRU_MAX_CAPACITY) endif() -if (USE_EVENT_DRIVEN_COMPLETION) +if(USE_EVENT_DRIVEN_COMPLETION) add_compile_definitions(USE_EVENT_DRIVEN_COMPLETION) message(STATUS "Event-driven completion is enabled") else() message(STATUS "Event-driven completion is disabled") endif() -if (USE_NVMEOF) +if(USE_NVMEOF) set(USE_CUDA ON) add_compile_definitions(USE_NVMEOF) message(STATUS "NVMe-oF support is enabled") endif() -if (USE_MNNVL) - if (NOT USE_HIP AND NOT USE_MUSA AND NOT USE_MACA) +if(USE_MNNVL) + if(NOT USE_HIP + AND NOT USE_MUSA + AND NOT USE_MACA + AND NOT USE_SUPA) set(USE_CUDA ON) endif() add_compile_definitions(USE_MNNVL) message(STATUS "Multi-Node NVLink support is enabled") endif() -if (USE_CUDA) +if (USE_VRAM_SEGMENT) + set(USE_CUDA ON) + add_compile_definitions(USE_VRAM_SEGMENT) + message(STATUS "VRAM SEGMENT is ON") +endif() + +if(USE_CUDA) + find_package(CUDAToolkit REQUIRED) add_compile_definitions(USE_CUDA) message(STATUS "CUDA support is enabled") - include_directories(/usr/local/cuda/include) - link_directories( - /usr/local/cuda/lib - /usr/local/cuda/lib64 - ) + include_directories(${CUDAToolkit_INCLUDE_DIRS}) + # Include stubs directory so the linker can find libcuda.so on machines that + # have the CUDA toolkit but not a GPU driver (e.g. CI builders). On + # production systems with a driver the real libcuda.so in the system library + # path takes precedence at both link and runtime. + link_directories(${CUDAToolkit_LIBRARY_DIR} ${CUDAToolkit_LIBRARY_DIR}/stubs) +endif() + +if(USE_NCCL_DEVICE OR USE_NCCL_HOST) + if(NOT USE_CUDA) + message(FATAL_ERROR + "USE_NCCL_DEVICE and USE_NCCL_HOST require USE_CUDA=ON") + endif() + list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_LIST_DIR}) + find_package(NCCLDevice 2.30.4 REQUIRED MODULE) +endif() + +if(USE_NCCL_DEVICE) + add_compile_definitions(USE_NCCL_DEVICE) + message(STATUS + "NCCL DeviceTransport support is enabled (NCCL ${NCCLDevice_VERSION})") +endif() + +if(USE_NCCL_HOST) + add_compile_definitions(USE_NCCL_HOST) + message(STATUS + "NCCL host RMA transport is enabled (NCCL ${NCCLDevice_VERSION})") +endif() + +if(USE_SUPA) + add_compile_definitions(USE_SUPA) + message(STATUS "SUPA support is enabled") + if(NOT DEFINED BIREN_HOME OR BIREN_HOME STREQUAL "") + if(DEFINED ENV{BIREN_HOME} AND NOT "$ENV{BIREN_HOME}" STREQUAL "") + set(BIREN_HOME + "$ENV{BIREN_HOME}" + CACHE PATH "Biren SUPA SDK root") + else() + set(BIREN_HOME + "/usr/local/birensupa/all/latest" + CACHE PATH "Biren SUPA SDK root") + endif() + endif() + message(STATUS " BIREN_HOME: ${BIREN_HOME}") + include_directories(${BIREN_HOME}/supa/include) + link_directories(${BIREN_HOME}/supa/lib ${BIREN_HOME}/brumd/lib) +endif() + +if(USE_TPU) + # Every TPU source file lives under mooncake-transfer-engine/tent, which is + # only added when USE_TENT is ON. Without this guard -DUSE_TPU=ON configures + # and builds cleanly while compiling no TPU code at all. + if(NOT USE_TENT) + message( + FATAL_ERROR + "USE_TPU=ON requires USE_TENT=ON: all TPU support lives in TENT. Re-run cmake with -DUSE_TENT=ON." + ) + endif() + add_compile_definitions(USE_TPU) + message(STATUS "TPU (PJRT) staging support is enabled") endif() -if (NOT DEFINED NEUWARE_ROOT OR NEUWARE_ROOT STREQUAL "") - if (DEFINED ENV{NEUWARE_HOME} AND NOT "$ENV{NEUWARE_HOME}" STREQUAL "") - set(NEUWARE_ROOT "$ENV{NEUWARE_HOME}" CACHE PATH "Path to Cambricon Neuware SDK" FORCE) +if(NOT DEFINED NEUWARE_ROOT OR NEUWARE_ROOT STREQUAL "") + if(DEFINED ENV{NEUWARE_HOME} AND NOT "$ENV{NEUWARE_HOME}" STREQUAL "") + set(NEUWARE_ROOT + "$ENV{NEUWARE_HOME}" + CACHE PATH "Path to Cambricon Neuware SDK" FORCE) else() - set(NEUWARE_ROOT "/usr/local/neuware" CACHE PATH "Path to Cambricon Neuware SDK" FORCE) + set(NEUWARE_ROOT + "/usr/local/neuware" + CACHE PATH "Path to Cambricon Neuware SDK" FORCE) endif() endif() -if (NOT DEFINED MLU_INCLUDE_DIR OR MLU_INCLUDE_DIR STREQUAL "") +if(NOT DEFINED MLU_INCLUDE_DIR OR MLU_INCLUDE_DIR STREQUAL "") set(MLU_INCLUDE_DIR "${NEUWARE_ROOT}/include") endif() -if (NOT DEFINED MLU_LIB_DIR OR MLU_LIB_DIR STREQUAL "") +if(NOT DEFINED MLU_LIB_DIR OR MLU_LIB_DIR STREQUAL "") set(MLU_LIB_DIR "${NEUWARE_ROOT}/lib64") endif() -if (NOT DEFINED MACA_ROOT OR MACA_ROOT STREQUAL "") - if (DEFINED ENV{MACA_HOME} AND NOT "$ENV{MACA_HOME}" STREQUAL "") - set(MACA_ROOT "$ENV{MACA_HOME}" CACHE PATH "Path to MACA SDK" FORCE) +if(NOT DEFINED MACA_ROOT OR MACA_ROOT STREQUAL "") + if(DEFINED ENV{MACA_HOME} AND NOT "$ENV{MACA_HOME}" STREQUAL "") + set(MACA_ROOT + "$ENV{MACA_HOME}" + CACHE PATH "Path to MACA SDK" FORCE) else() - set(MACA_ROOT "/opt/maca" CACHE PATH "Path to MACA SDK" FORCE) + set(MACA_ROOT + "/opt/maca" + CACHE PATH "Path to MACA SDK" FORCE) endif() endif() -if (NOT DEFINED MACA_INCLUDE_DIR OR MACA_INCLUDE_DIR STREQUAL "") +if(NOT DEFINED MACA_INCLUDE_DIR OR MACA_INCLUDE_DIR STREQUAL "") set(MACA_INCLUDE_DIR "${MACA_ROOT}/include") endif() -if (NOT DEFINED MACA_LIB_DIR OR MACA_LIB_DIR STREQUAL "") - if (EXISTS "${MACA_ROOT}/lib64") +if(NOT DEFINED MACA_LIB_DIR OR MACA_LIB_DIR STREQUAL "") + if(EXISTS "${MACA_ROOT}/lib64") set(MACA_LIB_DIR "${MACA_ROOT}/lib64") else() set(MACA_LIB_DIR "${MACA_ROOT}/lib") endif() endif() -if (USE_MLU) +if(USE_MLU) add_compile_definitions(USE_MLU) message(STATUS "MLU support is enabled") include_directories(${MLU_INCLUDE_DIR}) - if (EXISTS "${MLU_LIB_DIR}") + if(EXISTS "${MLU_LIB_DIR}") link_directories(${MLU_LIB_DIR}) endif() endif() -if (USE_MACA) +if(USE_MACA) add_compile_definitions(USE_MACA) message(STATUS "MACA support is enabled") include_directories(${MACA_INCLUDE_DIR}) - if (EXISTS "${MACA_LIB_DIR}") + if(EXISTS "${MACA_LIB_DIR}") link_directories(${MACA_LIB_DIR}) endif() endif() -if (NOT DEFINED MC_TANGRT_ROOT OR MC_TANGRT_ROOT STREQUAL "") - if (DEFINED ENV{MC_TANGRT_ROOT} AND NOT "$ENV{MC_TANGRT_ROOT}" STREQUAL "") - set(MC_TANGRT_ROOT "$ENV{MC_TANGRT_ROOT}" CACHE PATH "Path to Tang runtime root" FORCE) +if(NOT DEFINED MC_TANGRT_ROOT OR MC_TANGRT_ROOT STREQUAL "") + if(DEFINED ENV{MC_TANGRT_ROOT} AND NOT "$ENV{MC_TANGRT_ROOT}" STREQUAL "") + set(MC_TANGRT_ROOT + "$ENV{MC_TANGRT_ROOT}" + CACHE PATH "Path to Tang runtime root" FORCE) else() - set(MC_TANGRT_ROOT "/usr/local/tangrt" CACHE PATH "Path to Tang runtime root" FORCE) + set(MC_TANGRT_ROOT + "/usr/local/tangrt" + CACHE PATH "Path to Tang runtime root" FORCE) endif() endif() -if (USE_SUNRISE) +if(USE_SUNRISE) add_compile_definitions(USE_SUNRISE) message(STATUS "Sunrise (Tang runtime) support is enabled") include_directories(${MC_TANGRT_ROOT}/include) endif() -if (USE_MUSA) +if(USE_MUSA) add_compile_definitions(USE_MUSA) message(STATUS "MUSA support is enabled") include_directories(/usr/local/musa/include) - link_directories( - /usr/local/musa/lib - ) + link_directories(/usr/local/musa/lib) endif() -if (USE_HYGON) - if (NOT DEFINED DTK_ROOT OR DTK_ROOT STREQUAL "") - if (DEFINED ENV{DTK_HOME} AND NOT "$ENV{DTK_HOME}" STREQUAL "") - set(DTK_ROOT "$ENV{DTK_HOME}" CACHE PATH "Path to Hygon DTK SDK" FORCE) +if(USE_HYGON) + if(NOT DEFINED DTK_ROOT OR DTK_ROOT STREQUAL "") + if(DEFINED ENV{DTK_HOME} AND NOT "$ENV{DTK_HOME}" STREQUAL "") + set(DTK_ROOT + "$ENV{DTK_HOME}" + CACHE PATH "Path to Hygon DTK SDK" FORCE) else() - set(DTK_ROOT "/opt/dtk" CACHE PATH "Path to Hygon DTK SDK" FORCE) + set(DTK_ROOT + "/opt/dtk" + CACHE PATH "Path to Hygon DTK SDK" FORCE) endif() endif() - if (NOT DEFINED DTK_INCLUDE_DIR OR DTK_INCLUDE_DIR STREQUAL "") + if(NOT DEFINED DTK_INCLUDE_DIR OR DTK_INCLUDE_DIR STREQUAL "") set(DTK_INCLUDE_DIR "${DTK_ROOT}/cuda/cuda-11/include") endif() - if (NOT DEFINED DTK_LIB_DIR OR DTK_LIB_DIR STREQUAL "") + if(NOT DEFINED DTK_LIB_DIR OR DTK_LIB_DIR STREQUAL "") set(DTK_LIB_DIR "${DTK_ROOT}/cuda/cuda-11/lib64") endif() add_compile_definitions(USE_HYGON) message(STATUS "Hygon DCU/DTK support is enabled") include_directories(${DTK_INCLUDE_DIR}) - if (EXISTS "${DTK_LIB_DIR}") + if(EXISTS "${DTK_LIB_DIR}") link_directories(${DTK_LIB_DIR}) endif() endif() -if (USE_COREX) - if (NOT DEFINED COREX_ROOT OR COREX_ROOT STREQUAL "") - if (DEFINED ENV{COREX_HOME} AND NOT "$ENV{COREX_HOME}" STREQUAL "") - set(COREX_ROOT "$ENV{COREX_HOME}" CACHE PATH "Path to Iluvatar CoreX SDK" FORCE) +if(USE_COREX) + if(NOT DEFINED COREX_ROOT OR COREX_ROOT STREQUAL "") + if(DEFINED ENV{COREX_HOME} AND NOT "$ENV{COREX_HOME}" STREQUAL "") + set(COREX_ROOT + "$ENV{COREX_HOME}" + CACHE PATH "Path to Iluvatar CoreX SDK" FORCE) else() - set(COREX_ROOT "/usr/local/corex" CACHE PATH "Path to Iluvatar CoreX SDK" FORCE) + set(COREX_ROOT + "/usr/local/corex" + CACHE PATH "Path to Iluvatar CoreX SDK" FORCE) endif() endif() - if (NOT DEFINED COREX_INCLUDE_DIR OR COREX_INCLUDE_DIR STREQUAL "") + if(NOT DEFINED COREX_INCLUDE_DIR OR COREX_INCLUDE_DIR STREQUAL "") set(COREX_INCLUDE_DIR "${COREX_ROOT}/include") endif() - if (NOT DEFINED COREX_LIB_DIR OR COREX_LIB_DIR STREQUAL "") + if(NOT DEFINED COREX_LIB_DIR OR COREX_LIB_DIR STREQUAL "") set(COREX_LIB_DIR "${COREX_ROOT}/lib") endif() add_compile_definitions(USE_COREX) message(STATUS "Iluvatar CoreX support is enabled") include_directories(${COREX_INCLUDE_DIR}) - if (EXISTS "${COREX_LIB_DIR}") + if(EXISTS "${COREX_LIB_DIR}") link_directories(${COREX_LIB_DIR}) endif() endif() -if (USE_HIP) +if(USE_HIP) list(APPEND CMAKE_PREFIX_PATH "/opt/rocm/lib/cmake") find_package(HIP REQUIRED) include_directories(${HIP_INCLUDE_DIRS}) @@ -307,52 +433,56 @@ if (USE_HIP) find_program(HIPIFY_PERL_EXECUTABLE hipify-perl) if(NOT HIPIFY_PERL_EXECUTABLE) - message(FATAL_ERROR - "hipify-perl not found.\n" - "Please ensure the ROCm or HIP SDK is installed and in your PATH.") + message( + FATAL_ERROR + "hipify-perl not found.\n" + "Please ensure the ROCm or HIP SDK is installed and in your PATH.") endif() endif() -# This function converts given CUDA source files into HIP-compatible -# files using hipify-perl, placing the outputs in the build directory for use in -# project compilation. The file path changes to a new location after hipify. +# This function converts given CUDA source files into HIP-compatible files using +# hipify-perl, placing the outputs in the build directory for use in project +# compilation. The file path changes to a new location after hipify. function(hipify_files input_var_name) - set(result_files) + set(result_files) - foreach(input_file IN LISTS ${input_var_name}) - file(RELATIVE_PATH rel_path ${CMAKE_SOURCE_DIR} ${input_file}) - set(output_file "${CMAKE_BINARY_DIR}/${rel_path}") + foreach(input_file IN LISTS ${input_var_name}) + file(RELATIVE_PATH rel_path ${CMAKE_SOURCE_DIR} ${input_file}) + set(output_file "${CMAKE_BINARY_DIR}/${rel_path}") - get_filename_component(output_dir ${output_file} DIRECTORY) - file(MAKE_DIRECTORY ${output_dir}) + get_filename_component(output_dir ${output_file} DIRECTORY) + file(MAKE_DIRECTORY ${output_dir}) - add_custom_command( - OUTPUT ${output_file} - COMMAND ${HIPIFY_PERL_EXECUTABLE} ${input_file} > ${output_file} - DEPENDS ${input_file} - COMMENT "HIPifying ${input_file} → ${output_file}" - ) + add_custom_command( + OUTPUT ${output_file} + COMMAND ${HIPIFY_PERL_EXECUTABLE} ${input_file} > ${output_file} + DEPENDS ${input_file} + COMMENT "HIPifying ${input_file} → ${output_file}") - list(APPEND result_files ${output_file}) - endforeach() + list(APPEND result_files ${output_file}) + endforeach() - set(${input_var_name} ${result_files} PARENT_SCOPE) + set(${input_var_name} + ${result_files} + PARENT_SCOPE) endfunction() -if (USE_CXL) +if(USE_CXL) add_compile_definitions(USE_CXL) message(STATUS "CXL support is enabled") endif() -if (USE_TCP) +if(USE_TCP) add_compile_definitions(USE_TCP) endif() -if (USE_BAREX) +if(USE_BAREX) add_compile_definitions(USE_BAREX) endif() -if (USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM) +if(USE_ASCEND + OR USE_ASCEND_DIRECT + OR USE_UBSHMEM) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DOPEN_BUILD_PROJECT ") set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -DOPEN_BUILD_PROJECT ") string(TOLOWER "${CMAKE_SYSTEM_PROCESSOR}" CURRENT_CPU) @@ -368,7 +498,8 @@ if (USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM) message(STATUS "Use env ASCEND_HOME_PATH") file(GLOB ASCEND_TOOLKIT_ROOT "$ENV{ASCEND_HOME_PATH}/${CPU_ARCH}-linux") else() - file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/${CPU_ARCH}-linux") + file(GLOB ASCEND_TOOLKIT_ROOT + "/usr/local/Ascend/ascend-toolkit/latest/${CPU_ARCH}-linux") endif() set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64") set(ASCEND_INCLUDE_DIR "${ASCEND_TOOLKIT_ROOT}/include") @@ -377,24 +508,25 @@ if (USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM) link_directories(${ASCEND_LIB_DIR}) endif() -if (USE_ASCEND) +if(USE_ASCEND) set(ASCEND_DEVLIB_DIR "${ASCEND_TOOLKIT_ROOT}/devlib") link_directories(${ASCEND_DEVLIB_DIR}) add_compile_definitions(USE_ASCEND) endif() -if (USE_ASCEND_DIRECT) +if(USE_ASCEND_DIRECT) set(BUILD_SHARED_LIBS ON) add_compile_definitions(USE_ASCEND_DIRECT) endif() -if (USE_UBSHMEM) +if(USE_UBSHMEM) set(BUILD_SHARED_LIBS ON) add_compile_definitions(USE_UBSHMEM) endif() -if (USE_ASCEND_HETEROGENEOUS) - file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux") +if(USE_ASCEND_HETEROGENEOUS) + file(GLOB ASCEND_TOOLKIT_ROOT + "/usr/local/Ascend/ascend-toolkit/latest/*-linux") set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64") set(ASCEND_INCLUDE_DIR "${ASCEND_TOOLKIT_ROOT}/include") add_compile_definitions(USE_ASCEND_HETEROGENEOUS) @@ -402,21 +534,26 @@ if (USE_ASCEND_HETEROGENEOUS) link_directories(${ASCEND_LIB_DIR}) endif() -if (USE_REDIS) +if(USE_REDIS) add_compile_definitions(USE_REDIS) message(STATUS "Redis as metadata server support is enabled") endif() -if (USE_HTTP) +if(USE_HTTP) add_compile_definitions(USE_HTTP) message(STATUS "Http as metadata server support is enabled") endif() -if (NOT USE_ETCD AND NOT USE_REDIS AND NOT USE_HTTP) - message(STATUS "None of USE_ETCD, USE_REDIS, USE_HTTP is selected, only \"P2PHANDSHAKE\" is supported as metadata server") +if(NOT USE_ETCD + AND NOT USE_REDIS + AND NOT USE_HTTP) + message( + STATUS + "None of USE_ETCD, USE_REDIS, USE_HTTP is selected, only \"P2PHANDSHAKE\" is supported as metadata server" + ) endif() -if (WITH_METRICS) +if(WITH_METRICS) add_compile_definitions(WITH_METRICS) message(STATUS "metrics is enabled") endif() @@ -426,8 +563,23 @@ if(USE_3FS) message(STATUS "3FS storage backend is enabled") endif() +if(EXISTS "/usr/include/boost1.78") + include_directories(SYSTEM "/usr/include/boost1.78") + link_directories("/usr/lib64/boost1.78") +endif() + set(GFLAGS_USE_TARGET_NAMESPACE "true") find_package(yaml-cpp REQUIRED) find_package(gflags REQUIRED) +if(NOT TARGET gflags::gflags) + foreach(_gflags_target gflags-shared gflags_shared gflags) + if(TARGET ${_gflags_target}) + add_library(gflags::gflags INTERFACE IMPORTED) + set_target_properties(gflags::gflags PROPERTIES INTERFACE_LINK_LIBRARIES + ${_gflags_target}) + break() + endif() + endforeach() +endif() find_package(yalantinglibs CONFIG REQUIRED) add_compile_definitions(YLT_ENABLE_IBV) diff --git a/mooncake-common/etcd/etcd_wrapper.go b/mooncake-common/etcd/etcd_wrapper.go index 6c1695d3d3..4ac718072c 100644 --- a/mooncake-common/etcd/etcd_wrapper.go +++ b/mooncake-common/etcd/etcd_wrapper.go @@ -97,7 +97,6 @@ func newStoreClientConfig(validEndpoints []string) clientv3.Config { DialTimeout: 5 * time.Second, DialKeepAliveTime: storeDialKeepAliveTime, DialKeepAliveTimeout: storeDialKeepAliveTimeout, - PermitWithoutStream: true, } } @@ -792,6 +791,66 @@ func EtcdStoreBatchCreateWrapper(keys **C.char, values **C.char, count C.int, er return 0 } +//export EtcdStoreTxnCompareAndPutWrapper +func EtcdStoreTxnCompareAndPutWrapper(compareKeys **C.char, compareKeySizes *C.int, compareKinds *C.int, compareValues **C.char, compareValueSizes *C.int, compareCount C.int, putKeys **C.char, putKeySizes *C.int, putValues **C.char, putValueSizes *C.int, putCount C.int, errMsg **C.char) int { + cli := getStoreClient() + if cli == nil { + *errMsg = C.CString("etcd client not initialized") + return -1 + } + + cmpN := int(compareCount) + putN := int(putCount) + + cmps := make([]clientv3.Cmp, 0, cmpN) + if cmpN > 0 { + compareKeyPtrs := (*[1 << 28]*C.char)(unsafe.Pointer(compareKeys))[:cmpN:cmpN] + compareKeySizeList := (*[1 << 28]C.int)(unsafe.Pointer(compareKeySizes))[:cmpN:cmpN] + compareKindList := (*[1 << 28]C.int)(unsafe.Pointer(compareKinds))[:cmpN:cmpN] + compareValuePtrs := (*[1 << 28]*C.char)(unsafe.Pointer(compareValues))[:cmpN:cmpN] + compareValueSizeList := (*[1 << 28]C.int)(unsafe.Pointer(compareValueSizes))[:cmpN:cmpN] + for i := 0; i < cmpN; i++ { + k := C.GoStringN(compareKeyPtrs[i], compareKeySizeList[i]) + switch int(compareKindList[i]) { + case 0: + v := C.GoStringN(compareValuePtrs[i], compareValueSizeList[i]) + cmps = append(cmps, clientv3.Compare(clientv3.Value(k), "=", v)) + case 1: + cmps = append(cmps, clientv3.Compare(clientv3.CreateRevision(k), "=", 0)) + default: + *errMsg = C.CString("unsupported compare kind") + return -1 + } + } + } + + ops := make([]clientv3.Op, 0, putN) + if putN > 0 { + putKeyPtrs := (*[1 << 28]*C.char)(unsafe.Pointer(putKeys))[:putN:putN] + putKeySizeList := (*[1 << 28]C.int)(unsafe.Pointer(putKeySizes))[:putN:putN] + putValuePtrs := (*[1 << 28]*C.char)(unsafe.Pointer(putValues))[:putN:putN] + putValueSizeList := (*[1 << 28]C.int)(unsafe.Pointer(putValueSizes))[:putN:putN] + for i := 0; i < putN; i++ { + k := C.GoStringN(putKeyPtrs[i], putKeySizeList[i]) + v := C.GoStringN(putValuePtrs[i], putValueSizeList[i]) + ops = append(ops, clientv3.OpPut(k, v)) + } + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + resp, err := cli.Txn(ctx).If(cmps...).Then(ops...).Commit() + if err != nil { + *errMsg = C.CString(err.Error()) + return -1 + } + if !resp.Succeeded { + *errMsg = C.CString("transaction compare failed") + return -2 + } + return 0 +} + //export EtcdStoreGetWithPrefixWrapper func EtcdStoreGetWithPrefixWrapper(prefix *C.char, prefixSize C.int, keys **C.char, keySizes **C.int, values **C.char, valueSizes **C.int, count *C.int, errMsg **C.char) int { cli := getStoreClient() diff --git a/mooncake-common/etcd/go.mod b/mooncake-common/etcd/go.mod index ef5b3503f1..0c5e2d57f9 100644 --- a/mooncake-common/etcd/go.mod +++ b/mooncake-common/etcd/go.mod @@ -19,10 +19,10 @@ require ( go.uber.org/multierr v1.6.0 // indirect go.uber.org/zap v1.17.0 // indirect golang.org/x/net v0.55.0 // indirect - golang.org/x/sys v0.39.0 // indirect - golang.org/x/text v0.32.0 // indirect - google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect - google.golang.org/grpc v1.79.3 // indirect - google.golang.org/protobuf v1.36.10 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/text v0.37.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20260414002931-afd174a4e478 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect + google.golang.org/grpc v1.82.1 // indirect + google.golang.org/protobuf v1.36.11 // indirect ) diff --git a/mooncake-common/etcd/go.sum b/mooncake-common/etcd/go.sum index 888b82dc64..ea7c850636 100644 --- a/mooncake-common/etcd/go.sum +++ b/mooncake-common/etcd/go.sum @@ -41,16 +41,16 @@ go.etcd.io/etcd/client/v3 v3.5.21 h1:T6b1Ow6fNjOLOtM0xSoKNQt1ASPCLWrF9XMHcH9pEyY go.etcd.io/etcd/client/v3 v3.5.21/go.mod h1:mFYy67IOqmbRf/kRUvsHixzo3iG+1OF2W2+jVIQRAnU= go.opentelemetry.io/auto/sdk v1.2.1 h1:jXsnJ4Lmnqd11kwkBV2LgLoFMZKizbCi5fNZ/ipaZ64= go.opentelemetry.io/auto/sdk v1.2.1/go.mod h1:KRTj+aOaElaLi+wW1kO/DZRXwkF4C5xPbEe3ZiIhN7Y= -go.opentelemetry.io/otel v1.39.0 h1:8yPrr/S0ND9QEfTfdP9V+SiwT4E0G7Y5MO7p85nis48= -go.opentelemetry.io/otel v1.39.0/go.mod h1:kLlFTywNWrFyEdH0oj2xK0bFYZtHRYUdv1NklR/tgc8= -go.opentelemetry.io/otel/metric v1.39.0 h1:d1UzonvEZriVfpNKEVmHXbdf909uGTOQjA0HF0Ls5Q0= -go.opentelemetry.io/otel/metric v1.39.0/go.mod h1:jrZSWL33sD7bBxg1xjrqyDjnuzTUB0x1nBERXd7Ftcs= -go.opentelemetry.io/otel/sdk v1.39.0 h1:nMLYcjVsvdui1B/4FRkwjzoRVsMK8uL/cj0OyhKzt18= -go.opentelemetry.io/otel/sdk v1.39.0/go.mod h1:vDojkC4/jsTJsE+kh+LXYQlbL8CgrEcwmt1ENZszdJE= -go.opentelemetry.io/otel/sdk/metric v1.39.0 h1:cXMVVFVgsIf2YL6QkRF4Urbr/aMInf+2WKg+sEJTtB8= -go.opentelemetry.io/otel/sdk/metric v1.39.0/go.mod h1:xq9HEVH7qeX69/JnwEfp6fVq5wosJsY1mt4lLfYdVew= -go.opentelemetry.io/otel/trace v1.39.0 h1:2d2vfpEDmCJ5zVYz7ijaJdOF59xLomrvj7bjt6/qCJI= -go.opentelemetry.io/otel/trace v1.39.0/go.mod h1:88w4/PnZSazkGzz/w84VHpQafiU4EtqqlVdxWy+rNOA= +go.opentelemetry.io/otel v1.43.0 h1:mYIM03dnh5zfN7HautFE4ieIig9amkNANT+xcVxAj9I= +go.opentelemetry.io/otel v1.43.0/go.mod h1:JuG+u74mvjvcm8vj8pI5XiHy1zDeoCS2LB1spIq7Ay0= +go.opentelemetry.io/otel/metric v1.43.0 h1:d7638QeInOnuwOONPp4JAOGfbCEpYb+K6DVWvdxGzgM= +go.opentelemetry.io/otel/metric v1.43.0/go.mod h1:RDnPtIxvqlgO8GRW18W6Z/4P462ldprJtfxHxyKd2PY= +go.opentelemetry.io/otel/sdk v1.43.0 h1:pi5mE86i5rTeLXqoF/hhiBtUNcrAGHLKQdhg4h4V9Dg= +go.opentelemetry.io/otel/sdk v1.43.0/go.mod h1:P+IkVU3iWukmiit/Yf9AWvpyRDlUeBaRg6Y+C58QHzg= +go.opentelemetry.io/otel/sdk/metric v1.43.0 h1:S88dyqXjJkuBNLeMcVPRFXpRw2fuwdvfCGLEo89fDkw= +go.opentelemetry.io/otel/sdk/metric v1.43.0/go.mod h1:C/RJtwSEJ5hzTiUz5pXF1kILHStzb9zFlIEe85bhj6A= +go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09nk+3A= +go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0= go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw= go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc= go.uber.org/multierr v1.6.0 h1:y6IPFStTAIT5Ytl7/XYmHvzXQ7S3g/IeZW9hyZ5thw4= @@ -66,20 +66,20 @@ golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= -golang.org/x/net v0.48.0 h1:zyQRTTrjc33Lhh0fBgT/H3oZq9WuvRR5gPC70xpDiQU= -golang.org/x/net v0.48.0/go.mod h1:+ndRgGjkh8FGtu1w1FGbEC31if4VrNVMuKTgcAAnQRY= +golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8= +golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= -golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk= -golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY= +golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= -golang.org/x/text v0.32.0 h1:ZD01bjUt1FQ9WJ0ClOL5vxgxOI/sVCNgX1YtKwcY0mU= -golang.org/x/text v0.32.0/go.mod h1:o/rUWzghvpD5TXrTIBuJU77MTaN0ljMWE47kxGJQ7jY= +golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= +golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= @@ -88,16 +88,16 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -gonum.org/v1/gonum v0.16.0 h1:5+ul4Swaf3ESvrOnidPp4GZbzf0mxVQpDCYUQE7OJfk= -gonum.org/v1/gonum v0.16.0/go.mod h1:fef3am4MQ93R2HHpKnLk4/Tbh/s0+wqD5nfa6Pnwy4E= -google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 h1:fCvbg86sFXwdrl5LgVcTEvNC+2txB5mgROGmRL5mrls= -google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:+rXWjjaukWZun3mLfjmVnQi18E1AsFbDN9QdJ5YXLto= -google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 h1:gRkg/vSppuSQoDjxyiGfN4Upv/h/DQmIR10ZU8dh4Ww= -google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217/go.mod h1:7i2o+ce6H/6BluujYR+kqX3GKH+dChPTQU19wjRPiGk= -google.golang.org/grpc v1.79.3 h1:sybAEdRIEtvcD68Gx7dmnwjZKlyfuc61Dyo9pGXXkKE= -google.golang.org/grpc v1.79.3/go.mod h1:KmT0Kjez+0dde/v2j9vzwoAScgEPx/Bw1CYChhHLrHQ= -google.golang.org/protobuf v1.36.10 h1:AYd7cD/uASjIL6Q9LiTjz8JLcrh/88q5UObnmY3aOOE= -google.golang.org/protobuf v1.36.10/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= +gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4= +gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E= +google.golang.org/genproto/googleapis/api v0.0.0-20260414002931-afd174a4e478 h1:yQugLulqltosq0B/f8l4w9VryjV+N/5gcW0jQ3N8Qec= +google.golang.org/genproto/googleapis/api v0.0.0-20260414002931-afd174a4e478/go.mod h1:C6ADNqOxbgdUUeRTU+LCHDPB9ttAMCTff6auwCVa4uc= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw= +google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8= +google.golang.org/grpc v1.82.1 h1:NnAxzGRA0677vCa4BUkOAnO5+FfQqVl9iUXeD0IqcGE= +google.golang.org/grpc v1.82.1/go.mod h1:yzTZ1TB1Z3SG+LIYaI+WiE8D5+PZ3ArnrSp8zF3+/ZA= +google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE= +google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= diff --git a/mooncake-common/include/ascii_string.h b/mooncake-common/include/ascii_string.h new file mode 100644 index 0000000000..2e92aa3b03 --- /dev/null +++ b/mooncake-common/include/ascii_string.h @@ -0,0 +1,50 @@ +#pragma once + +#include +#include +#include + +namespace mooncake { + +constexpr bool IsAsciiWhitespace(char ch) { + return ch == ' ' || ch == '\t' || ch == '\n' || ch == '\r' || ch == '\f' || + ch == '\v'; +} + +inline std::string_view TrimAsciiWhitespace(std::string_view value) { + while (!value.empty() && IsAsciiWhitespace(value.front())) { + value.remove_prefix(1); + } + while (!value.empty() && IsAsciiWhitespace(value.back())) { + value.remove_suffix(1); + } + return value; +} + +constexpr char AsciiToLower(char ch) { + return ch >= 'A' && ch <= 'Z' ? static_cast(ch + ('a' - 'A')) : ch; +} + +inline std::string AsciiToLower(std::string_view value) { + std::string normalized; + normalized.reserve(value.size()); + for (char ch : value) { + normalized.push_back(AsciiToLower(ch)); + } + return normalized; +} + +inline bool AsciiCaseInsensitiveEquals(std::string_view lhs, + std::string_view rhs) { + if (lhs.size() != rhs.size()) { + return false; + } + for (size_t i = 0; i < lhs.size(); ++i) { + if (AsciiToLower(lhs[i]) != AsciiToLower(rhs[i])) { + return false; + } + } + return true; +} + +} // namespace mooncake diff --git a/mooncake-common/include/bool_parser.h b/mooncake-common/include/bool_parser.h new file mode 100644 index 0000000000..1139aaa75f --- /dev/null +++ b/mooncake-common/include/bool_parser.h @@ -0,0 +1,62 @@ +#pragma once + +#include +#include + +#include "ascii_string.h" + +namespace mooncake { + +enum class BoolTokenSet { + kTrueFalse, + kCanonical, +}; + +struct BoolParseOptions { + BoolTokenSet token_set{BoolTokenSet::kCanonical}; + bool trim_ascii_whitespace{true}; +}; + +inline std::optional TryParseBool(std::string_view value, + BoolParseOptions options = {}) { + if (options.trim_ascii_whitespace) { + value = TrimAsciiWhitespace(value); + } + if (value.empty()) { + return std::nullopt; + } + + if (value == "1") { + return true; + } + if (value == "0") { + return false; + } + if (AsciiCaseInsensitiveEquals(value, "true")) { + return true; + } + if (AsciiCaseInsensitiveEquals(value, "false")) { + return false; + } + if (options.token_set == BoolTokenSet::kTrueFalse) { + return std::nullopt; + } + + if (AsciiCaseInsensitiveEquals(value, "yes")) { + return true; + } + if (AsciiCaseInsensitiveEquals(value, "no")) { + return false; + } + if (AsciiCaseInsensitiveEquals(value, "on") || + AsciiCaseInsensitiveEquals(value, "enable")) { + return true; + } + if (AsciiCaseInsensitiveEquals(value, "off") || + AsciiCaseInsensitiveEquals(value, "disable")) { + return false; + } + return std::nullopt; +} + +} // namespace mooncake diff --git a/mooncake-common/include/crc_checksum.h b/mooncake-common/include/crc_checksum.h new file mode 100644 index 0000000000..1b2c7d1c47 --- /dev/null +++ b/mooncake-common/include/crc_checksum.h @@ -0,0 +1,19 @@ +#pragma once + +#include +#include + +namespace mooncake { + +class CrcChecksum { + public: + void Update(const void* data, size_t size); + uint64_t Finalize() const { return crc_; } + + private: + uint64_t crc_{0}; +}; + +uint64_t ComputeCrcChecksum(const void* data, size_t size); + +} // namespace mooncake diff --git a/mooncake-common/include/default_config.h b/mooncake-common/include/default_config.h index f072baf212..ae5395ac27 100644 --- a/mooncake-common/include/default_config.h +++ b/mooncake-common/include/default_config.h @@ -12,6 +12,8 @@ #include #include +#include "ascii_string.h" + namespace mooncake { class DefaultConfig { public: @@ -156,9 +158,7 @@ inline void init_ylt_log_level() { easylog::set_min_severity(easylog::Severity::WARN); return; } - std::string level_str(env_level); - std::transform(level_str.begin(), level_str.end(), level_str.begin(), - [](unsigned char c) { return std::tolower(c); }); + const std::string level_str = AsciiToLower(env_level); easylog::Severity severity; if (level_str == "trace") { severity = easylog::Severity::TRACE; diff --git a/mooncake-common/include/duration_utils.h b/mooncake-common/include/duration_utils.h index 521922b6b5..f913772e46 100644 --- a/mooncake-common/include/duration_utils.h +++ b/mooncake-common/include/duration_utils.h @@ -1,24 +1,14 @@ #pragma once -#include #include #include #include #include -namespace mooncake { +#include "ascii_string.h" +#include "integer_parser.h" -inline std::string_view TrimAsciiWhitespace(std::string_view value) { - while (!value.empty() && - std::isspace(static_cast(value.front()))) { - value.remove_prefix(1); - } - while (!value.empty() && - std::isspace(static_cast(value.back()))) { - value.remove_suffix(1); - } - return value; -} +namespace mooncake { inline bool ParseDurationMs(std::string_view value, uint64_t* result, std::string* error = nullptr) { @@ -41,8 +31,8 @@ inline bool ParseDurationMs(std::string_view value, uint64_t* result, } size_t number_end = 0; - while (number_end < trimmed.size() && - std::isdigit(static_cast(trimmed[number_end]))) { + while (number_end < trimmed.size() && trimmed[number_end] >= '0' && + trimmed[number_end] <= '9') { ++number_end; } @@ -52,23 +42,14 @@ inline bool ParseDurationMs(std::string_view value, uint64_t* result, "s, m, or h as the unit suffix"); } - uint64_t numeric_value = 0; - for (size_t i = 0; i < number_end; ++i) { - const uint64_t digit = static_cast(trimmed[i] - '0'); - if (numeric_value > - (std::numeric_limits::max() - digit) / 10) { - return set_error("duration value is too large"); - } - numeric_value = numeric_value * 10 + digit; + const auto numeric_value = + TryParseInteger(trimmed.substr(0, number_end)); + if (!numeric_value.has_value()) { + return set_error("duration value is too large"); } std::string_view suffix = TrimAsciiWhitespace(trimmed.substr(number_end)); - std::string normalized_suffix; - normalized_suffix.reserve(suffix.size()); - for (char ch : suffix) { - normalized_suffix.push_back( - static_cast(std::tolower(static_cast(ch)))); - } + const std::string normalized_suffix = AsciiToLower(suffix); uint64_t multiplier = 1; if (normalized_suffix.empty() || normalized_suffix == "ms") { @@ -84,11 +65,11 @@ inline bool ParseDurationMs(std::string_view value, uint64_t* result, "'; supported units are ms, s, m, and h"); } - if (numeric_value > std::numeric_limits::max() / multiplier) { + if (*numeric_value > std::numeric_limits::max() / multiplier) { return set_error("duration value is too large after unit conversion"); } - *result = numeric_value * multiplier; + *result = *numeric_value * multiplier; return true; } diff --git a/mooncake-common/include/environ.h b/mooncake-common/include/environ.h index 78f8f45028..f6224696dd 100644 --- a/mooncake-common/include/environ.h +++ b/mooncake-common/include/environ.h @@ -6,11 +6,21 @@ namespace mooncake { +class EnvironSource { + public: + virtual ~EnvironSource() = default; + virtual const char* Get(const char* name) const = 0; +}; + class Environ { public: // Singleton access static Environ& Get(); + // Construct from an injected source. Production code should use Get(); + // this constructor allows tests to provide deterministic environment data. + explicit Environ(const EnvironSource& source); + // Getters for Environment Variables int GetNumCqPerCtx() const { return num_cq_per_ctx_; } int GetNumCompChannelsPerCtx() const { return num_comp_channels_per_ctx_; } @@ -51,20 +61,51 @@ class Environ { bool GetPathRoundrobin() const { return path_roundrobin_; } bool GetWithNvidiaPeermem() const { return with_nvidia_peermem_; } int GetEfaCqThreads() const { return efa_cq_threads_; } + bool GetStoreChecksumEnabled() const { return store_checksum_enabled_; } + + // AWS / S3 client configuration + std::string GetAwsRegion() const { return aws_region_; } + std::string GetAwsS3Endpoint() const { return aws_s3_endpoint_; } + std::string GetAwsBucketName() const { return aws_bucket_name_; } + std::string GetAwsAccessKeyId() const { return aws_access_key_id_; } + std::string GetAwsSecretAccessKey() const { return aws_secret_access_key_; } + bool GetAwsUseVirtualAddressing() const { + return aws_use_virtual_addressing_; + } + bool GetAwsUseHttps() const { return aws_use_https_; } + // Empty string means "unset" — s3_helper keeps the AWS SDK default in + // that case. Parsing to AWS enums is done by the consumer. + std::string GetAwsRequestChecksumCalculation() const { + return aws_request_checksum_calculation_; + } + std::string GetAwsResponseChecksumValidation() const { + return aws_response_checksum_validation_; + } + int64_t GetAwsConnectTimeoutMs() const { return aws_connect_timeout_ms_; } + int64_t GetAwsRequestTimeoutMs() const { return aws_request_timeout_ms_; } + uint32_t GetRpcClientIoThreads() const { return rpc_client_io_threads_; } + uint32_t GetStoreRpcClientIoThreads() const { + return store_rpc_client_io_threads_; + } + uint32_t GetTransferEngineRpcClientIoThreads() const { + return transfer_engine_rpc_client_io_threads_; + } // Helper method to get int from env static int GetInt(const char* name, int default_value); + static int64_t GetInt64(const char* name, int64_t default_value); + static uint32_t GetUInt32(const char* name, uint32_t default_value); + static uint64_t GetUInt64(const char* name, uint64_t default_value); // Helper method to get size_t from env static size_t GetSizeT(const char* name, size_t default_value); - // Helper method to get bool from env (checks for "1", "true", "TRUE") + // Helper method to get a canonical boolean from env. Invalid values use the + // caller-provided default. static bool GetBool(const char* name, bool default_value); // Helper method to get string from env static std::string GetString(const char* name, const std::string& default_value); private: - Environ(); - // Member variables int num_cq_per_ctx_; int num_comp_channels_per_ctx_; @@ -103,6 +144,23 @@ class Environ { bool path_roundrobin_; bool with_nvidia_peermem_; int efa_cq_threads_; + bool store_checksum_enabled_; + uint32_t rpc_client_io_threads_; + uint32_t store_rpc_client_io_threads_; + uint32_t transfer_engine_rpc_client_io_threads_; + + // AWS / S3 client configuration + std::string aws_region_; + std::string aws_s3_endpoint_; + std::string aws_bucket_name_; + std::string aws_access_key_id_; + std::string aws_secret_access_key_; + bool aws_use_virtual_addressing_; + bool aws_use_https_; + std::string aws_request_checksum_calculation_; + std::string aws_response_checksum_validation_; + int64_t aws_connect_timeout_ms_; + int64_t aws_request_timeout_ms_; }; } // namespace mooncake diff --git a/mooncake-common/include/integer_parser.h b/mooncake-common/include/integer_parser.h new file mode 100644 index 0000000000..da430c71d8 --- /dev/null +++ b/mooncake-common/include/integer_parser.h @@ -0,0 +1,48 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "ascii_string.h" + +namespace mooncake { + +struct IntegerParseOptions { + bool trim_ascii_whitespace{false}; + bool allow_leading_plus{false}; +}; + +template +std::optional TryParseInteger(std::string_view value, + IntegerParseOptions options = {}) { + static_assert(std::is_integral_v && + !std::is_same_v, bool>, + "TryParseInteger requires a non-bool integral type"); + + if (options.trim_ascii_whitespace) { + value = TrimAsciiWhitespace(value); + } + if (options.allow_leading_plus && !value.empty() && value.front() == '+') { + value.remove_prefix(1); + if (!value.empty() && (value.front() == '+' || value.front() == '-')) { + return std::nullopt; + } + } + if (value.empty()) { + return std::nullopt; + } + + Integer parsed{}; + const char* begin = value.data(); + const char* end = begin + value.size(); + const auto result = std::from_chars(begin, end, parsed); + if (result.ec != std::errc{} || result.ptr != end) { + return std::nullopt; + } + return parsed; +} + +} // namespace mooncake diff --git a/mooncake-common/include/rpc_client_io_context.h b/mooncake-common/include/rpc_client_io_context.h new file mode 100644 index 0000000000..b9842f0c99 --- /dev/null +++ b/mooncake-common/include/rpc_client_io_context.h @@ -0,0 +1,69 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace mooncake { + +std::shared_ptr CreateRpcClientIoContextPool( + uint32_t thread_count); + +template +coro_io::io_context_pool& GetRpcClientIoContextPool(uint32_t thread_count) { + static const auto io_pool = CreateRpcClientIoContextPool(thread_count); + return *io_pool; +} + +/** + * A replaceable client pool for callers that communicate with one target at a + * time. Requests retain a shared_ptr to the old pool while they are in flight; + * after an address switch the old pool is destroyed when those requests end. + */ +class RpcClientPool { + public: + using ClientPool = coro_io::client_pool; + using PoolConfig = ClientPool::pool_config; + + explicit RpcClientPool(coro_io::io_context_pool& io_context_pool, + PoolConfig config = {}) + : io_context_pool_(io_context_pool), config_(std::move(config)) { + // Address replacement supersedes background recovery of the old host. + config_.host_alive_detect_duration = std::chrono::seconds(0); + } + + std::shared_ptr GetOrCreateClientPool( + std::string_view address) { + std::lock_guard lock(mutex_); + if (!client_pool_ || address_ != address) { + client_pool_ = + ClientPool::create(address, config_, io_context_pool_); + address_ = address; + } + return client_pool_; + } + + std::shared_ptr GetClientPool() const { + std::shared_lock lock(mutex_); + return client_pool_; + } + + private: + mutable std::shared_mutex mutex_; + coro_io::io_context_pool& io_context_pool_; + PoolConfig config_; + std::string address_; + std::shared_ptr client_pool_; +}; + +} // namespace mooncake diff --git a/mooncake-common/k8s-lease/CMakeLists.txt b/mooncake-common/k8s-lease/CMakeLists.txt index 7ccc2ac62a..ed555f17b0 100644 --- a/mooncake-common/k8s-lease/CMakeLists.txt +++ b/mooncake-common/k8s-lease/CMakeLists.txt @@ -1,5 +1,6 @@ add_custom_command( OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so + ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.h COMMAND bash -c "go mod tidy" && bash -c "go build -buildmode=c-shared -o ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.so k8s_lease_wrapper.go" && cp ${CMAKE_CURRENT_BINARY_DIR}/libk8s_lease_wrapper.h ${CMAKE_CURRENT_SOURCE_DIR} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} COMMENT "Building K8s Lease Go shared library" diff --git a/mooncake-common/k8s-lease/go.mod b/mooncake-common/k8s-lease/go.mod index 4bfc205264..b3c7f5cdf2 100644 --- a/mooncake-common/k8s-lease/go.mod +++ b/mooncake-common/k8s-lease/go.mod @@ -1,6 +1,6 @@ module github.com/kvcache-ai/Mooncake/mooncake-common/k8s-lease -go 1.24.0 +go 1.25.0 require ( k8s.io/api v0.34.3 @@ -40,11 +40,11 @@ require ( github.com/x448/float16 v0.8.4 // indirect go.yaml.in/yaml/v2 v2.4.3 // indirect go.yaml.in/yaml/v3 v3.0.4 // indirect - golang.org/x/net v0.47.0 // indirect + golang.org/x/net v0.55.0 // indirect golang.org/x/oauth2 v0.30.0 // indirect - golang.org/x/sys v0.38.0 // indirect - golang.org/x/term v0.37.0 // indirect - golang.org/x/text v0.31.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/term v0.43.0 // indirect + golang.org/x/text v0.37.0 // indirect golang.org/x/time v0.9.0 // indirect google.golang.org/protobuf v1.36.8 // indirect gopkg.in/evanphx/json-patch.v4 v4.13.0 // indirect diff --git a/mooncake-common/k8s-lease/k8s_lease_wrapper.go b/mooncake-common/k8s-lease/k8s_lease_wrapper.go index 72f907318f..b67caca52a 100644 --- a/mooncake-common/k8s-lease/k8s_lease_wrapper.go +++ b/mooncake-common/k8s-lease/k8s_lease_wrapper.go @@ -20,6 +20,7 @@ import "C" import ( "context" + "encoding/json" "fmt" "os" "sync" @@ -29,6 +30,7 @@ import ( coordinationv1 "k8s.io/api/coordination/v1" apierrors "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/types" "k8s.io/apimachinery/pkg/watch" "k8s.io/client-go/kubernetes" "k8s.io/client-go/rest" @@ -486,4 +488,62 @@ func K8sLeaseCancelWatch( return 0 } +// patchPodLabel sets or removes a label on a pod using a JSON merge patch. +// If value is non-nil, the label is set; if nil, the label is removed. +func patchPodLabel(namespace, podName, labelKey string, value interface{}) error { + if err := ensureClientInitialized(); err != nil { + return err + } + + patch := map[string]interface{}{ + "metadata": map[string]interface{}{ + "labels": map[string]interface{}{ + labelKey: value, + }, + }, + } + patchBytes, err := json.Marshal(patch) + if err != nil { + return fmt.Errorf("failed to marshal label patch: %w", err) + } + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + _, err = globalClient.CoreV1().Pods(namespace).Patch( + ctx, podName, types.MergePatchType, patchBytes, metav1.PatchOptions{}) + if err != nil { + return fmt.Errorf("failed to patch pod label: %w", err) + } + return nil +} + +//export K8sPatchPodLabel +func K8sPatchPodLabel( + ns, podName, labelKey, labelValue *C.char, + errMsg **C.char, +) C.int { + err := patchPodLabel(C.GoString(ns), C.GoString(podName), + C.GoString(labelKey), C.GoString(labelValue)) + if err != nil { + *errMsg = C.CString(err.Error()) + return -1 + } + return 0 +} + +//export K8sRemovePodLabel +func K8sRemovePodLabel( + ns, podName, labelKey *C.char, + errMsg **C.char, +) C.int { + err := patchPodLabel(C.GoString(ns), C.GoString(podName), + C.GoString(labelKey), nil) + if err != nil { + *errMsg = C.CString(err.Error()) + return -1 + } + return 0 +} + func main() {} diff --git a/mooncake-common/k8s-lease/k8s_lease_wrapper_test.go b/mooncake-common/k8s-lease/k8s_lease_wrapper_test.go index 705b579a52..47b50e0276 100644 --- a/mooncake-common/k8s-lease/k8s_lease_wrapper_test.go +++ b/mooncake-common/k8s-lease/k8s_lease_wrapper_test.go @@ -7,6 +7,7 @@ import ( "time" coordinationv1 "k8s.io/api/coordination/v1" + corev1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/client-go/kubernetes" "k8s.io/client-go/kubernetes/fake" @@ -22,6 +23,15 @@ func swapClient(newClient kubernetes.Interface) kubernetes.Interface { return old } +// swapInitClientFn replaces initClientFn and returns the old one. +func swapInitClientFn(newFn func() error) func() error { + clientMutex.Lock() + defer clientMutex.Unlock() + old := initClientFn + initClientFn = newFn + return old +} + // TestGetHolderWithFakeClient tests getHolder using a fake K8s clientset. func TestGetHolderWithFakeClient(t *testing.T) { holderID := "node-1:8080" @@ -53,6 +63,61 @@ func TestGetHolderWithFakeClient(t *testing.T) { } } +func TestPatchPodLabelAutoInitializesClient(t *testing.T) { + fakeClient := fake.NewSimpleClientset(&corev1.Pod{ + ObjectMeta: metav1.ObjectMeta{ + Name: "test-pod", + Namespace: "default", + }, + }) + + oldClient := swapClient(nil) + defer swapClient(oldClient) + + initCalls := 0 + oldInitFn := swapInitClientFn(func() error { + initCalls++ + clientMutex.Lock() + globalClient = fakeClient + clientMutex.Unlock() + return nil + }) + defer swapInitClientFn(oldInitFn) + + const labelKey = "mooncake.io/store-role" + if err := patchPodLabel("default", "test-pod", labelKey, "leader"); err != nil { + t.Fatalf("patchPodLabel(set) failed: %v", err) + } + if initCalls != 1 { + t.Fatalf("initClientFn calls = %d, want 1", initCalls) + } + + pod, err := fakeClient.CoreV1().Pods("default").Get( + context.Background(), "test-pod", metav1.GetOptions{}) + if err != nil { + t.Fatalf("get pod after label set failed: %v", err) + } + if got := pod.Labels[labelKey]; got != "leader" { + t.Fatalf("label %q = %q, want %q", labelKey, got, "leader") + } + + if err := patchPodLabel("default", "test-pod", labelKey, nil); err != nil { + t.Fatalf("patchPodLabel(remove) failed: %v", err) + } + if initCalls != 1 { + t.Fatalf("initClientFn calls after second patch = %d, want 1", initCalls) + } + + pod, err = fakeClient.CoreV1().Pods("default").Get( + context.Background(), "test-pod", metav1.GetOptions{}) + if err != nil { + t.Fatalf("get pod after label removal failed: %v", err) + } + if _, exists := pod.Labels[labelKey]; exists { + t.Fatalf("label %q still present after removal", labelKey) + } +} + // TestGetHolderNotFound tests getHolder when the Lease does not exist. func TestGetHolderNotFound(t *testing.T) { fakeClient := fake.NewSimpleClientset() diff --git a/mooncake-common/src/CMakeLists.txt b/mooncake-common/src/CMakeLists.txt index 1f231775ce..2730485cf1 100644 --- a/mooncake-common/src/CMakeLists.txt +++ b/mooncake-common/src/CMakeLists.txt @@ -1,70 +1,48 @@ find_package(yaml-cpp REQUIRED) -find_package(asio QUIET) - -if(asio_FOUND) - message(STATUS "Found ASIO via find_package") - set(ASIO_INCLUDE_DIR ${asio_INCLUDE_DIR}) -else() - find_path(ASIO_INCLUDE_DIR - NAMES asio.hpp - PATHS - /usr/local/include - /usr/include - ${CMAKE_INSTALL_PREFIX}/include - DOC "Path to ASIO headers" - ) - - if(NOT ASIO_INCLUDE_DIR) - message(FATAL_ERROR "ASIO not found. Please install ASIO or set ASIO_INCLUDE_DIR manually.") - endif() - - message(STATUS "Found ASIO at: ${ASIO_INCLUDE_DIR}") -endif() - -set(MOONCAKE_COMMON_SOURCES - default_config.cpp - environ.cpp -) +set(MOONCAKE_COMMON_SOURCES crc_checksum.cpp default_config.cpp environ.cpp + rpc_client_io_context.cpp) add_library(asio_shared SHARED asio_impl.cpp) -target_compile_definitions(asio_shared - PUBLIC - ASIO_SEPARATE_COMPILATION - ASIO_DYN_LINK -) - -target_include_directories(asio_shared - PUBLIC - ${ASIO_INCLUDE_DIR} -) - -set_target_properties(asio_shared PROPERTIES - POSITION_INDEPENDENT_CODE ON - INSTALL_RPATH "$ORIGIN" - BUILD_WITH_INSTALL_RPATH TRUE - OUTPUT_NAME "asio" - LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/mooncake-common" -) - -target_link_libraries(asio_shared PUBLIC pthread) +target_compile_definitions(asio_shared PUBLIC ASIO_SEPARATE_COMPILATION + ASIO_DYN_LINK) + +set_target_properties( + asio_shared + PROPERTIES POSITION_INDEPENDENT_CODE ON + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH TRUE + OUTPUT_NAME "asio" + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/mooncake-common") + +target_link_libraries(asio_shared PUBLIC pthread yalantinglibs::yalantinglibs) + +# Static, PIC twin of asio_shared for consumers that absorb asio into their own +# shared object instead of taking a DT_NEEDED on libasio.so (see +# WITH_STORE_C_SHARED). Same TU and definitions as asio_shared so the symbols +# match what the rest of the tree was compiled against; only the linkage +# differs. +if(WITH_STORE_C_SHARED) + add_library(asio_static STATIC asio_impl.cpp) + target_compile_definitions(asio_static PUBLIC ASIO_SEPARATE_COMPILATION + ASIO_DYN_LINK) + target_include_directories(asio_static PUBLIC ${ASIO_INCLUDE_DIR}) + set_target_properties(asio_static PROPERTIES POSITION_INDEPENDENT_CODE ON) + target_link_libraries(asio_static PUBLIC pthread) +endif() -add_library(mooncake_common - ${MOONCAKE_COMMON_SOURCES} -) +add_library(mooncake_common ${MOONCAKE_COMMON_SOURCES}) -target_include_directories(mooncake_common PUBLIC - $ - $ -) +target_include_directories( + mooncake_common + PUBLIC $ + $) -target_link_libraries(mooncake_common PUBLIC - yaml-cpp - jsoncpp -) +target_link_libraries(mooncake_common PUBLIC asio_shared yaml-cpp jsoncpp + yalantinglibs::yalantinglibs) -if (BUILD_SHARED_LIBS) +if(BUILD_SHARED_LIBS) install(TARGETS mooncake_common DESTINATION lib) endif() diff --git a/mooncake-common/src/crc_checksum.cpp b/mooncake-common/src/crc_checksum.cpp new file mode 100644 index 0000000000..37c330201b --- /dev/null +++ b/mooncake-common/src/crc_checksum.cpp @@ -0,0 +1,42 @@ +#include "crc_checksum.h" + +#include + +namespace mooncake { + +namespace { + +constexpr uint64_t kCrc64EcmaPolynomial = 0x42F0E1EBA9EA3693ULL; + +constexpr std::array MakeCrc64EcmaTable() { + std::array table{}; + for (size_t i = 0; i < table.size(); ++i) { + uint64_t crc = static_cast(i) << 56; + for (int bit = 0; bit < 8; ++bit) { + crc = (crc & (1ULL << 63)) != 0 ? (crc << 1) ^ kCrc64EcmaPolynomial + : crc << 1; + } + table[i] = crc; + } + return table; +} + +constexpr auto kCrc64EcmaTable = MakeCrc64EcmaTable(); + +} // namespace + +void CrcChecksum::Update(const void* data, size_t size) { + const auto* bytes = static_cast(data); + for (size_t i = 0; i < size; ++i) { + const auto index = static_cast((crc_ >> 56) ^ bytes[i]); + crc_ = kCrc64EcmaTable[index] ^ (crc_ << 8); + } +} + +uint64_t ComputeCrcChecksum(const void* data, size_t size) { + CrcChecksum checksum; + checksum.Update(data, size); + return checksum.Finalize(); +} + +} // namespace mooncake diff --git a/mooncake-common/src/environ.cpp b/mooncake-common/src/environ.cpp index 46aa1e9a9c..c736823171 100644 --- a/mooncake-common/src/environ.cpp +++ b/mooncake-common/src/environ.cpp @@ -1,113 +1,213 @@ #include "environ.h" -#include -#include -#include + #include -#include #include +#include +#include + +#include "bool_parser.h" +#include "integer_parser.h" namespace mooncake { -Environ& Environ::Get() { - static Environ instance; - return instance; +namespace { + +constexpr char kRpcClientIoThreadsEnv[] = "MC_RPC_CLIENT_IO_THREADS"; +constexpr char kStoreRpcClientIoThreadsEnv[] = "MC_STORE_RPC_CLIENT_IO_THREADS"; +constexpr char kTransferEngineRpcClientIoThreadsEnv[] = + "MC_TE_RPC_CLIENT_IO_THREADS"; +constexpr uint32_t kDefaultRpcClientIoThreads = 16; + +class OsEnvironSource final : public EnvironSource { + public: + const char* Get(const char* name) const override { + return std::getenv(name); + } +}; + +const EnvironSource& GetOsEnvironSource() { + static const OsEnvironSource source; + return source; } -int Environ::GetInt(const char* name, int default_value) { - const char* val = std::getenv(name); - if (val) { - char* endptr = nullptr; - errno = 0; - long result = std::strtol(val, &endptr, 10); - if (endptr == val || *endptr != '\0' || errno == ERANGE || - result < INT_MIN || result > INT_MAX) { - std::cerr << "[Mooncake] Warning: invalid value '" << val - << "' for env " << name << ", using default " - << default_value << std::endl; - return default_value; - } - return static_cast(result); +template +Integer ReadInteger(const EnvironSource& source, const char* name, + Integer default_value) { + const char* value = source.Get(name); + if (value == nullptr) { + return default_value; + } + + const auto parsed = TryParseInteger( + std::string_view(value), + {.trim_ascii_whitespace = true, .allow_leading_plus = true}); + if (parsed.has_value()) { + return *parsed; } + + std::cerr << "[Mooncake] Warning: invalid value '" << value << "' for env " + << name << ", using default " << default_value << std::endl; return default_value; } -size_t Environ::GetSizeT(const char* name, size_t default_value) { - const char* val = std::getenv(name); - if (val) { - char* endptr = nullptr; - errno = 0; - long long result = std::strtoll(val, &endptr, 10); - if (endptr == val || *endptr != '\0' || errno == ERANGE || result < 0 || - static_cast(result) > SIZE_MAX) { - std::cerr << "[Mooncake] Warning: invalid value '" << val - << "' for env " << name << ", using default " - << default_value << std::endl; - return default_value; - } - return static_cast(result); +int ReadInt(const EnvironSource& source, const char* name, int default_value) { + return ReadInteger(source, name, default_value); +} + +int64_t ReadInt64(const EnvironSource& source, const char* name, + int64_t default_value) { + return ReadInteger(source, name, default_value); +} + +size_t ReadSizeT(const EnvironSource& source, const char* name, + size_t default_value) { + return ReadInteger(source, name, default_value); +} + +bool ReadBool(const EnvironSource& source, const char* name, + bool default_value) { + const char* value = source.Get(name); + if (value == nullptr) { + return default_value; } + + const auto parsed = TryParseBool(value); + if (parsed.has_value()) { + return *parsed; + } + + std::cerr << "[Mooncake] Warning: invalid value '" << value << "' for env " + << name << ", using default " << default_value << std::endl; return default_value; } +std::string ReadString(const EnvironSource& source, const char* name, + const std::string& default_value) { + const char* val = source.Get(name); + return val ? std::string(val) : default_value; +} + +uint32_t ResolveRpcClientIoThreads(const EnvironSource& source, + const char* env_name, uint32_t fallback) { + const int configured = + ReadInt(source, env_name, static_cast(fallback)); + return configured > 0 ? static_cast(configured) : fallback; +} + +} // namespace + +Environ& Environ::Get() { + static Environ instance(GetOsEnvironSource()); + return instance; +} + +int Environ::GetInt(const char* name, int default_value) { + return ReadInt(GetOsEnvironSource(), name, default_value); +} + +int64_t Environ::GetInt64(const char* name, int64_t default_value) { + return ReadInt64(GetOsEnvironSource(), name, default_value); +} + +uint32_t Environ::GetUInt32(const char* name, uint32_t default_value) { + return ReadInteger(GetOsEnvironSource(), name, default_value); +} + +uint64_t Environ::GetUInt64(const char* name, uint64_t default_value) { + return ReadInteger(GetOsEnvironSource(), name, default_value); +} + +size_t Environ::GetSizeT(const char* name, size_t default_value) { + return ReadSizeT(GetOsEnvironSource(), name, default_value); +} + bool Environ::GetBool(const char* name, bool default_value) { - const char* val = std::getenv(name); - if (val) { - std::string s(val); - std::transform(s.begin(), s.end(), s.begin(), - [](unsigned char c) { return std::tolower(c); }); - return s == "1" || s == "true" || s == "on" || s == "yes"; - } - return default_value; + return ReadBool(GetOsEnvironSource(), name, default_value); } std::string Environ::GetString(const char* name, const std::string& default_value) { - const char* val = std::getenv(name); - if (val) { - return std::string(val); - } - return default_value; + return ReadString(GetOsEnvironSource(), name, default_value); } -Environ::Environ() { - num_cq_per_ctx_ = GetInt("MC_NUM_CQ_PER_CTX", 1); - num_comp_channels_per_ctx_ = GetInt("MC_NUM_COMP_CHANNELS_PER_CTX", 1); - ib_port_ = GetInt("MC_IB_PORT", 1); - ib_tc_ = GetInt("MC_IB_TC", -1); - ib_pci_relaxed_ordering_ = GetInt("MC_IB_PCI_RELAXED_ORDERING", 0); - gid_index_ = GetInt("MC_GID_INDEX", 3); - max_cqe_per_ctx_ = GetInt("MC_MAX_CQE_PER_CTX", 4096); - max_ep_per_ctx_ = GetInt("MC_MAX_EP_PER_CTX", 65536); - num_qp_per_ep_ = GetInt("MC_NUM_QP_PER_EP", 2); - max_sge_ = GetInt("MC_MAX_SGE", 4); - max_wr_ = GetInt("MC_MAX_WR", 256); - max_inline_ = GetInt("MC_MAX_INLINE", 64); - mtu_ = GetInt("MC_MTU", 4096); - workers_per_ctx_ = GetInt("MC_WORKERS_PER_CTX", 2); - slice_size_ = GetSizeT("MC_SLICE_SIZE", 65536); - retry_cnt_ = GetInt("MC_RETRY_CNT", 9); - log_level_ = GetString("MC_LOG_LEVEL", "INFO"); - disable_metacache_ = GetBool("MC_DISABLE_METACACHE", false); - handshake_listen_backlog_ = GetInt("MC_HANDSHAKE_LISTEN_BACKLOG", 128); - handshake_max_length_ = GetInt("MC_HANDSHAKE_MAX_LENGTH", 1048576); - log_dir_ = GetString("MC_LOG_DIR", ""); - redis_password_ = GetString("MC_REDIS_PASSWORD", ""); - redis_db_index_ = GetInt("MC_REDIS_DB_INDEX", 0); - fragment_ratio_ = GetInt("MC_FRAGMENT_RATIO", 4); +Environ::Environ(const EnvironSource& source) { + const uint32_t hardware_threads = + static_cast(std::thread::hardware_concurrency()); + const uint32_t default_rpc_client_io_threads = std::min( + kDefaultRpcClientIoThreads, std::max(uint32_t{1}, hardware_threads)); + rpc_client_io_threads_ = ResolveRpcClientIoThreads( + source, kRpcClientIoThreadsEnv, default_rpc_client_io_threads); + store_rpc_client_io_threads_ = ResolveRpcClientIoThreads( + source, kStoreRpcClientIoThreadsEnv, rpc_client_io_threads_); + transfer_engine_rpc_client_io_threads_ = ResolveRpcClientIoThreads( + source, kTransferEngineRpcClientIoThreadsEnv, rpc_client_io_threads_); + + num_cq_per_ctx_ = ReadInt(source, "MC_NUM_CQ_PER_CTX", 1); + num_comp_channels_per_ctx_ = + ReadInt(source, "MC_NUM_COMP_CHANNELS_PER_CTX", 1); + ib_port_ = ReadInt(source, "MC_IB_PORT", 1); + ib_tc_ = ReadInt(source, "MC_IB_TC", -1); + ib_pci_relaxed_ordering_ = ReadInt(source, "MC_IB_PCI_RELAXED_ORDERING", 0); + gid_index_ = ReadInt(source, "MC_GID_INDEX", 3); + max_cqe_per_ctx_ = ReadInt(source, "MC_MAX_CQE_PER_CTX", 4096); + max_ep_per_ctx_ = ReadInt(source, "MC_MAX_EP_PER_CTX", 65536); + num_qp_per_ep_ = ReadInt(source, "MC_NUM_QP_PER_EP", 2); + max_sge_ = ReadInt(source, "MC_MAX_SGE", 4); + max_wr_ = ReadInt(source, "MC_MAX_WR", 256); + max_inline_ = ReadInt(source, "MC_MAX_INLINE", 64); + mtu_ = ReadInt(source, "MC_MTU", 4096); + workers_per_ctx_ = ReadInt(source, "MC_WORKERS_PER_CTX", 2); + slice_size_ = ReadSizeT(source, "MC_SLICE_SIZE", 65536); + retry_cnt_ = ReadInt(source, "MC_RETRY_CNT", 9); + log_level_ = ReadString(source, "MC_LOG_LEVEL", "INFO"); + disable_metacache_ = ReadBool(source, "MC_DISABLE_METACACHE", false); + handshake_listen_backlog_ = + ReadInt(source, "MC_HANDSHAKE_LISTEN_BACKLOG", 128); + handshake_max_length_ = ReadInt(source, "MC_HANDSHAKE_MAX_LENGTH", 1048576); + log_dir_ = ReadString(source, "MC_LOG_DIR", ""); + redis_password_ = ReadString(source, "MC_REDIS_PASSWORD", ""); + redis_db_index_ = ReadInt(source, "MC_REDIS_DB_INDEX", 0); + fragment_ratio_ = ReadInt(source, "MC_FRAGMENT_RATIO", 4); enable_dest_device_affinity_ = - GetBool("MC_ENABLE_DEST_DEVICE_AFFINITY", false); - use_ipv6_ = GetBool("MC_USE_IPV6", false); - min_rpc_port_ = GetInt("MC_MIN_RPC_PORT", GetInt("MC_MIN_PRC_PORT", 15000)); - max_rpc_port_ = GetInt("MC_MAX_RPC_PORT", GetInt("MC_MAX_PRC_PORT", 17000)); - enable_parallel_reg_mr_ = GetInt("MC_ENABLE_PARALLEL_REG_MR", -1); - endpoint_store_type_ = GetString("MC_ENDPOINT_STORE_TYPE", "SIEVE"); - force_tcp_ = GetBool("MC_FORCE_TCP", false); - force_hca_ = GetBool("MC_FORCE_HCA", false); - force_mnnvl_ = GetBool("MC_FORCE_MNNVL", false); - intra_nvlink_ = GetBool("MC_INTRA_NVLINK", false); - path_roundrobin_ = GetBool("MC_PATH_ROUNDROBIN", false); - with_nvidia_peermem_ = GetBool("WITH_NVIDIA_PEERMEM", true); - efa_cq_threads_ = GetInt("MC_EFA_CQ_THREADS", 1); + ReadBool(source, "MC_ENABLE_DEST_DEVICE_AFFINITY", false); + use_ipv6_ = ReadBool(source, "MC_USE_IPV6", false); + min_rpc_port_ = ReadInt(source, "MC_MIN_RPC_PORT", + ReadInt(source, "MC_MIN_PRC_PORT", 15000)); + max_rpc_port_ = ReadInt(source, "MC_MAX_RPC_PORT", + ReadInt(source, "MC_MAX_PRC_PORT", 17000)); + enable_parallel_reg_mr_ = ReadInt(source, "MC_ENABLE_PARALLEL_REG_MR", -1); + endpoint_store_type_ = + ReadString(source, "MC_ENDPOINT_STORE_TYPE", "SIEVE"); + force_tcp_ = ReadBool(source, "MC_FORCE_TCP", false); + force_hca_ = ReadBool(source, "MC_FORCE_HCA", false); + force_mnnvl_ = ReadBool(source, "MC_FORCE_MNNVL", false); + intra_nvlink_ = ReadBool(source, "MC_INTRA_NVLINK", false); + path_roundrobin_ = ReadBool(source, "MC_PATH_ROUNDROBIN", false); + with_nvidia_peermem_ = ReadBool(source, "WITH_NVIDIA_PEERMEM", true); + efa_cq_threads_ = ReadInt(source, "MC_EFA_CQ_THREADS", 1); + store_checksum_enabled_ = + ReadBool(source, "MOONCAKE_STORE_CHECKSUM", false); + + // AWS / S3 client configuration (consumed by s3_helper.cpp) + aws_region_ = ReadString(source, "MOONCAKE_AWS_REGION", ""); + aws_s3_endpoint_ = ReadString(source, "MOONCAKE_AWS_S3_ENDPOINT", ""); + aws_bucket_name_ = ReadString(source, "MOONCAKE_AWS_BUCKET_NAME", ""); + aws_access_key_id_ = ReadString(source, "MOONCAKE_AWS_ACCESS_KEY_ID", ""); + aws_secret_access_key_ = + ReadString(source, "MOONCAKE_AWS_SECRET_ACCESS_KEY", ""); + aws_use_virtual_addressing_ = + ReadBool(source, "MOONCAKE_AWS_USE_VIRTUAL_ADDRESSING", true); + aws_use_https_ = ReadBool(source, "MOONCAKE_AWS_USE_HTTPS", true); + // Empty string preserves "unset" semantics — s3_helper keeps the AWS SDK + // default in that case rather than forcing a value. + aws_request_checksum_calculation_ = + ReadString(source, "MOONCAKE_AWS_REQUEST_CHECKSUM_CALCULATION", ""); + aws_response_checksum_validation_ = + ReadString(source, "MOONCAKE_AWS_RESPONSE_CHECKSUM_VALIDATION", ""); + aws_connect_timeout_ms_ = + ReadInt64(source, "MOONCAKE_AWS_CONNECT_TIMEOUT_MS", 10000); + aws_request_timeout_ms_ = + ReadInt64(source, "MOONCAKE_AWS_REQUEST_TIMEOUT_MS", 30000); } } // namespace mooncake diff --git a/mooncake-common/src/rpc_client_io_context.cpp b/mooncake-common/src/rpc_client_io_context.cpp new file mode 100644 index 0000000000..cb15f5e8dc --- /dev/null +++ b/mooncake-common/src/rpc_client_io_context.cpp @@ -0,0 +1,10 @@ +#include "rpc_client_io_context.h" + +namespace mooncake { + +std::shared_ptr CreateRpcClientIoContextPool( + uint32_t thread_count) { + return coro_io::create_io_context_pool(thread_count); +} + +} // namespace mooncake diff --git a/mooncake-common/tests/CMakeLists.txt b/mooncake-common/tests/CMakeLists.txt index ec43cfa22d..ee2d458547 100644 --- a/mooncake-common/tests/CMakeLists.txt +++ b/mooncake-common/tests/CMakeLists.txt @@ -1,11 +1,29 @@ add_executable(default_config_test default_config_test.cpp) -target_link_libraries(default_config_test PUBLIC - mooncake_common - gtest - pthread -) +target_link_libraries(default_config_test PUBLIC mooncake_common gtest pthread) add_test(NAME default_config_test COMMAND default_config_test) add_executable(environ_test environ_test.cpp) target_link_libraries(environ_test PUBLIC mooncake_common gtest pthread) add_test(NAME environ_test COMMAND environ_test) + +add_executable(crc_checksum_test crc_checksum_test.cpp) +target_link_libraries(crc_checksum_test PUBLIC mooncake_common gtest gtest_main + pthread) +add_test(NAME crc_checksum_test COMMAND crc_checksum_test) + +add_executable(rpc_client_io_threads_test rpc_client_io_threads_test.cpp) +target_link_libraries(rpc_client_io_threads_test PUBLIC mooncake_common gtest + pthread) +add_test(NAME rpc_client_io_threads_test COMMAND rpc_client_io_threads_test) + +add_executable(rpc_client_io_context_test rpc_client_io_context_test.cpp) +target_link_libraries(rpc_client_io_context_test PUBLIC mooncake_common gtest + ibverbs pthread) +add_test(NAME rpc_client_io_context_test COMMAND rpc_client_io_context_test) + +foreach(parser_test ascii_string bool_parser integer_parser) + add_executable(${parser_test}_test ${parser_test}_test.cpp) + target_link_libraries(${parser_test}_test PUBLIC mooncake_common gtest + gtest_main pthread) + add_test(NAME ${parser_test}_test COMMAND ${parser_test}_test) +endforeach() diff --git a/mooncake-common/tests/ascii_string_test.cpp b/mooncake-common/tests/ascii_string_test.cpp new file mode 100644 index 0000000000..5cce0c1386 --- /dev/null +++ b/mooncake-common/tests/ascii_string_test.cpp @@ -0,0 +1,26 @@ +#include "ascii_string.h" + +#include + +#include + +namespace mooncake { +namespace { + +TEST(AsciiStringTest, TrimsOnlyAsciiWhitespace) { + EXPECT_EQ(TrimAsciiWhitespace(" \t\n\r\f\vvalue \t\n\r\f\v"), "value"); + EXPECT_TRUE(TrimAsciiWhitespace(" \t\n\r\f\v").empty()); + + const std::string non_ascii = + std::string("\xC2\xA0") + "value" + std::string("\xC2\xA0"); + EXPECT_EQ(TrimAsciiWhitespace(non_ascii), non_ascii); +} + +TEST(AsciiStringTest, NormalizesAndComparesWithoutLocale) { + EXPECT_EQ(AsciiToLower("AbC-123"), "abc-123"); + EXPECT_TRUE(AsciiCaseInsensitiveEquals("EnAbLe", "enable")); + EXPECT_FALSE(AsciiCaseInsensitiveEquals("enabled", "enable")); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-common/tests/bool_parser_test.cpp b/mooncake-common/tests/bool_parser_test.cpp new file mode 100644 index 0000000000..9a0c59c774 --- /dev/null +++ b/mooncake-common/tests/bool_parser_test.cpp @@ -0,0 +1,33 @@ +#include "bool_parser.h" + +#include + +#include + +namespace mooncake { +namespace { + +TEST(BoolParserTest, ParsesCanonicalGrammar) { + for (std::string_view value : {"1", "true", "YES", "On", "enable"}) { + ASSERT_EQ(TryParseBool(value), true) << value; + } + for (std::string_view value : {"0", "false", "NO", "Off", "disable"}) { + ASSERT_EQ(TryParseBool(value), false) << value; + } + EXPECT_EQ(TryParseBool(" \tTrUe\r\n"), true); + EXPECT_EQ(TryParseBool(""), std::nullopt); + EXPECT_EQ(TryParseBool("maybe"), std::nullopt); +} + +TEST(BoolParserTest, SupportsRestrictedTokenSets) { + const BoolParseOptions true_false_only{ + .token_set = BoolTokenSet::kTrueFalse, .trim_ascii_whitespace = false}; + EXPECT_EQ(TryParseBool("1", true_false_only), true); + EXPECT_EQ(TryParseBool("0", true_false_only), false); + EXPECT_EQ(TryParseBool("TRUE", true_false_only), true); + EXPECT_EQ(TryParseBool("yes", true_false_only), std::nullopt); + EXPECT_EQ(TryParseBool(" true ", true_false_only), std::nullopt); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-common/tests/crc_checksum_test.cpp b/mooncake-common/tests/crc_checksum_test.cpp new file mode 100644 index 0000000000..69da3a8291 --- /dev/null +++ b/mooncake-common/tests/crc_checksum_test.cpp @@ -0,0 +1,32 @@ +#include "crc_checksum.h" + +#include + +#include +#include +#include + +namespace mooncake { + +TEST(CrcChecksumTest, MatchesCrc64EcmaKnownVector) { + constexpr std::string_view value = "123456789"; + EXPECT_EQ(ComputeCrcChecksum(value.data(), value.size()), + 0x6C40DF5F0B497347ULL); +} + +TEST(CrcChecksumTest, StreamingMatchesContiguousForArbitraryLengths) { + const std::array value = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, + 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, + 0x0C, 0x0D, 0x0E, 0x0F, 0x10}; + + CrcChecksum streaming; + streaming.Update(value.data(), 3); + streaming.Update(value.data() + 3, 7); + streaming.Update(value.data() + 10, value.size() - 10); + + EXPECT_EQ(streaming.Finalize(), + ComputeCrcChecksum(value.data(), value.size())); + EXPECT_EQ(ComputeCrcChecksum(nullptr, 0), 0); +} + +} // namespace mooncake diff --git a/mooncake-common/tests/environ_test.cpp b/mooncake-common/tests/environ_test.cpp index 9a0b534063..4701f6e78e 100644 --- a/mooncake-common/tests/environ_test.cpp +++ b/mooncake-common/tests/environ_test.cpp @@ -28,9 +28,25 @@ class EnvironTest : public ::testing::Test { void clearTestEnvVars() { unsetenv("MC_TEST_INT"); + unsetenv("MC_TEST_INT64"); + unsetenv("MC_TEST_UINT32"); + unsetenv("MC_TEST_UINT64"); unsetenv("MC_TEST_SIZET"); unsetenv("MC_TEST_BOOL"); unsetenv("MC_TEST_STRING"); + // Make sure AWS vars don't leak in from the test runner's env. + unsetenv("MOONCAKE_AWS_REGION"); + unsetenv("MOONCAKE_AWS_S3_ENDPOINT"); + unsetenv("MOONCAKE_AWS_BUCKET_NAME"); + unsetenv("MOONCAKE_AWS_ACCESS_KEY_ID"); + unsetenv("MOONCAKE_AWS_SECRET_ACCESS_KEY"); + unsetenv("MOONCAKE_AWS_USE_VIRTUAL_ADDRESSING"); + unsetenv("MOONCAKE_AWS_USE_HTTPS"); + unsetenv("MOONCAKE_AWS_REQUEST_CHECKSUM_CALCULATION"); + unsetenv("MOONCAKE_AWS_RESPONSE_CHECKSUM_VALIDATION"); + unsetenv("MOONCAKE_AWS_CONNECT_TIMEOUT_MS"); + unsetenv("MOONCAKE_AWS_REQUEST_TIMEOUT_MS"); + unsetenv("MOONCAKE_STORE_CHECKSUM"); } }; @@ -85,6 +101,84 @@ TEST_F(EnvironTest, GetIntMinValue) { EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 0), INT_MIN); } +TEST_F(EnvironTest, GetIntSupportsTrimmedLeadingPlus) { + setenv("MC_TEST_INT", " \t+42\r\n", 1); + EXPECT_EQ(Environ::GetInt("MC_TEST_INT", 0), 42); +} + +// --- GetInt64 --- + +TEST_F(EnvironTest, GetInt64ValidValue) { + setenv("MC_TEST_INT64", "123456789012", 1); + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 0), 123456789012LL); +} + +TEST_F(EnvironTest, GetInt64Missing) { + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 9999), 9999); +} + +TEST_F(EnvironTest, GetInt64Empty) { + setenv("MC_TEST_INT64", "", 1); + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 555), 555); +} + +TEST_F(EnvironTest, GetInt64NonNumeric) { + setenv("MC_TEST_INT64", "abc", 1); + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 555), 555); +} + +TEST_F(EnvironTest, GetInt64Overflow) { + setenv("MC_TEST_INT64", "99999999999999999999999999", 1); + EXPECT_EQ(Environ::GetInt64("MC_TEST_INT64", 555), 555); +} + +TEST_F(EnvironTest, UnsignedGettersUseRequestedDefaultForInvalidValues) { + setenv("MC_TEST_UINT32", "4294967296", 1); + setenv("MC_TEST_UINT64", "-1", 1); + EXPECT_EQ(Environ::GetUInt32("MC_TEST_UINT32", 17), 17U); + EXPECT_EQ(Environ::GetUInt64("MC_TEST_UINT64", 23), 23U); +} + +// --- AWS / S3 fields --- +// +// NOTE: Environ is a singleton whose constructor caches every value the +// first time Get() is called. So all AWS env vars must be set BEFORE the +// first Environ::Get() in this process. We therefore cover the populate +// path in a single test that takes the singleton's "first call" for +// itself; the default-path behavior is implicitly covered by Environ's +// constructor defaults (any earlier test would lock the cache to defaults +// and prevent us from observing populated values here). + +TEST_F(EnvironTest, AwsFieldsPopulateFromEnv) { + setenv("MOONCAKE_AWS_REGION", "us-east-1", 1); + setenv("MOONCAKE_AWS_S3_ENDPOINT", "https://s3.example.com", 1); + setenv("MOONCAKE_AWS_BUCKET_NAME", "my-bucket", 1); + setenv("MOONCAKE_AWS_ACCESS_KEY_ID", "AKIA-test", 1); + setenv("MOONCAKE_AWS_SECRET_ACCESS_KEY", "secret", 1); + setenv("MOONCAKE_AWS_USE_VIRTUAL_ADDRESSING", "0", 1); + setenv("MOONCAKE_AWS_USE_HTTPS", "0", 1); + setenv("MOONCAKE_AWS_REQUEST_CHECKSUM_CALCULATION", "when_required", 1); + setenv("MOONCAKE_AWS_RESPONSE_CHECKSUM_VALIDATION", "when_supported", 1); + setenv("MOONCAKE_AWS_CONNECT_TIMEOUT_MS", "5000", 1); + // Bogus request timeout should fall back to the registered default. + setenv("MOONCAKE_AWS_REQUEST_TIMEOUT_MS", "bogus", 1); + setenv("MOONCAKE_STORE_CHECKSUM", "1", 1); + + const auto& e = Environ::Get(); + EXPECT_EQ(e.GetAwsRegion(), "us-east-1"); + EXPECT_EQ(e.GetAwsS3Endpoint(), "https://s3.example.com"); + EXPECT_EQ(e.GetAwsBucketName(), "my-bucket"); + EXPECT_EQ(e.GetAwsAccessKeyId(), "AKIA-test"); + EXPECT_EQ(e.GetAwsSecretAccessKey(), "secret"); + EXPECT_FALSE(e.GetAwsUseVirtualAddressing()); + EXPECT_FALSE(e.GetAwsUseHttps()); + EXPECT_EQ(e.GetAwsRequestChecksumCalculation(), "when_required"); + EXPECT_EQ(e.GetAwsResponseChecksumValidation(), "when_supported"); + EXPECT_EQ(e.GetAwsConnectTimeoutMs(), 5000); + EXPECT_EQ(e.GetAwsRequestTimeoutMs(), 30000); + EXPECT_TRUE(e.GetStoreChecksumEnabled()); +} + // --- GetSizeT --- TEST_F(EnvironTest, GetSizeTValidValue) { @@ -134,20 +228,27 @@ TEST_F(EnvironTest, GetSizeTOverflow) { // --- GetBool --- TEST_F(EnvironTest, GetBoolTrue) { - for (const char* v : - {"1", "true", "TRUE", "True", "on", "ON", "yes", "YES"}) { + for (const char* v : {"1", "true", "TRUE", "True", "on", "ON", "yes", "YES", + "enable", "EnAbLe", " true "}) { setenv("MC_TEST_BOOL", v, 1); EXPECT_TRUE(Environ::GetBool("MC_TEST_BOOL", false)) << "for: " << v; } } TEST_F(EnvironTest, GetBoolFalse) { - for (const char* v : {"0", "false", "FALSE", "off", "no", "whatever"}) { + for (const char* v : + {"0", "false", "FALSE", "off", "no", "disable", "DiSaBlE"}) { setenv("MC_TEST_BOOL", v, 1); EXPECT_FALSE(Environ::GetBool("MC_TEST_BOOL", false)) << "for: " << v; } } +TEST_F(EnvironTest, GetBoolInvalidUsesRequestedDefault) { + setenv("MC_TEST_BOOL", "whatever", 1); + EXPECT_TRUE(Environ::GetBool("MC_TEST_BOOL", true)); + EXPECT_FALSE(Environ::GetBool("MC_TEST_BOOL", false)); +} + TEST_F(EnvironTest, GetBoolMissing) { EXPECT_TRUE(Environ::GetBool("MC_TEST_BOOL", true)); EXPECT_FALSE(Environ::GetBool("MC_TEST_BOOL", false)); @@ -155,7 +256,8 @@ TEST_F(EnvironTest, GetBoolMissing) { TEST_F(EnvironTest, GetBoolEmpty) { setenv("MC_TEST_BOOL", "", 1); - EXPECT_FALSE(Environ::GetBool("MC_TEST_BOOL", true)); + EXPECT_TRUE(Environ::GetBool("MC_TEST_BOOL", true)); + EXPECT_FALSE(Environ::GetBool("MC_TEST_BOOL", false)); } // --- GetString --- diff --git a/mooncake-common/tests/integer_parser_test.cpp b/mooncake-common/tests/integer_parser_test.cpp new file mode 100644 index 0000000000..74c3661d14 --- /dev/null +++ b/mooncake-common/tests/integer_parser_test.cpp @@ -0,0 +1,28 @@ +#include "integer_parser.h" + +#include +#include + +#include + +namespace mooncake { +namespace { + +TEST(IntegerParserTest, ParsesStrictlyWithRangeChecks) { + EXPECT_EQ(TryParseInteger("-42"), -42); + EXPECT_EQ(TryParseInteger("18446744073709551615"), + std::numeric_limits::max()); + EXPECT_EQ(TryParseInteger("18446744073709551616"), std::nullopt); + EXPECT_EQ(TryParseInteger("4294967296"), std::nullopt); + EXPECT_EQ(TryParseInteger("12suffix"), std::nullopt); + EXPECT_EQ(TryParseInteger(" 12 "), std::nullopt); + EXPECT_EQ(TryParseInteger("+12"), std::nullopt); + EXPECT_EQ(TryParseInteger("+-12", {.allow_leading_plus = true}), + std::nullopt); + EXPECT_EQ(TryParseInteger(" +12 ", {.trim_ascii_whitespace = true, + .allow_leading_plus = true}), + 12); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-common/tests/rpc_client_io_context_test.cpp b/mooncake-common/tests/rpc_client_io_context_test.cpp new file mode 100644 index 0000000000..c3b18c7127 --- /dev/null +++ b/mooncake-common/tests/rpc_client_io_context_test.cpp @@ -0,0 +1,103 @@ +#include "rpc_client_io_context.h" + +#include +#include +#include + +#include + +namespace mooncake { +namespace { + +class AddressSwitchService { + public: + explicit AddressSwitchService(int id) : id_(id) {} + + int Identify() const { return id_; } + + private: + int id_; +}; + +struct FirstTestRpcClientIoContextPoolTag {}; +struct SecondTestRpcClientIoContextPoolTag {}; + +coro_io::io_context_pool& GetFirstTestRpcClientIoContextPool() { + return GetRpcClientIoContextPool(2); +} + +coro_io::io_context_pool& GetSecondTestRpcClientIoContextPool() { + return GetRpcClientIoContextPool(3); +} + +TEST(RpcClientIoContextPoolTest, UsesConfiguredSizeAndReusesPool) { + auto& first_pool = GetFirstTestRpcClientIoContextPool(); + auto& second_pool = GetSecondTestRpcClientIoContextPool(); + ASSERT_EQ(first_pool.pool_size(), 2U); + ASSERT_EQ(second_pool.pool_size(), 3U); + + EXPECT_EQ(&GetFirstTestRpcClientIoContextPool(), &first_pool); + EXPECT_EQ(&GetSecondTestRpcClientIoContextPool(), &second_pool); + EXPECT_NE(&first_pool, &second_pool); +} + +TEST(RpcClientIoContextPoolTest, ReplacesPoolWhenTargetChanges) { + RpcClientPool pools(GetFirstTestRpcClientIoContextPool()); + + auto first = pools.GetOrCreateClientPool("127.0.0.1:10001"); + std::weak_ptr old_pool = first; + EXPECT_EQ(pools.GetOrCreateClientPool("127.0.0.1:10001"), first); + + auto second = pools.GetOrCreateClientPool("127.0.0.1:10002"); + EXPECT_NE(first, second); + first.reset(); + EXPECT_TRUE(old_pool.expired()); + EXPECT_EQ(pools.GetClientPool(), second); +} + +TEST(RpcClientIoContextPoolTest, SendsToNewAddressAfterSwitch) { + AddressSwitchService first_service(1); + AddressSwitchService second_service(2); + coro_rpc::coro_rpc_server first_server(1, 0, "127.0.0.1"); + coro_rpc::coro_rpc_server second_server(1, 0, "127.0.0.1"); + first_server.register_handler<&AddressSwitchService::Identify>( + &first_service); + second_server.register_handler<&AddressSwitchService::Identify>( + &second_service); + ASSERT_FALSE(first_server.async_start().hasResult()); + ASSERT_FALSE(second_server.async_start().hasResult()); + + RpcClientPool pools(GetFirstTestRpcClientIoContextPool()); + + const auto call = [&](uint16_t port) { + auto pool = + pools.GetOrCreateClientPool("127.0.0.1:" + std::to_string(port)); + return async_simple::coro::syncAwait(pool->send_request( + [](coro_rpc::coro_rpc_client& client) + -> async_simple::coro::Lazy> { + co_return co_await client + .call<&AddressSwitchService::Identify>(); + })); + }; + + auto first_result = call(first_server.port()); + ASSERT_TRUE(first_result); + ASSERT_TRUE(first_result.value()); + EXPECT_EQ(first_result.value().value(), 1); + + auto second_result = call(second_server.port()); + ASSERT_TRUE(second_result); + ASSERT_TRUE(second_result.value()); + EXPECT_EQ(second_result.value().value(), 2); + + first_server.stop(); + second_server.stop(); +} + +} // namespace +} // namespace mooncake + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-common/tests/rpc_client_io_threads_test.cpp b/mooncake-common/tests/rpc_client_io_threads_test.cpp new file mode 100644 index 0000000000..02e93c1806 --- /dev/null +++ b/mooncake-common/tests/rpc_client_io_threads_test.cpp @@ -0,0 +1,67 @@ +#include "environ.h" + +#include + +#include +#include +#include +#include + +namespace mooncake { +namespace { + +class MockEnvironSource : public EnvironSource { + public: + MockEnvironSource( + std::initializer_list> values) + : values_(values) {} + + const char* Get(const char* name) const override { + const auto it = values_.find(name); + return it == values_.end() ? nullptr : it->second.c_str(); + } + + void Set(std::string name, std::string value) { + values_[std::move(name)] = std::move(value); + } + + private: + std::unordered_map values_; +}; + +TEST(RpcClientIoThreadsTest, UsesComponentOverridesAndFreezesValues) { + MockEnvironSource source{{"MC_RPC_CLIENT_IO_THREADS", "8"}, + {"MC_STORE_RPC_CLIENT_IO_THREADS", "4"}, + {"MC_TE_RPC_CLIENT_IO_THREADS", "6"}}; + const Environ env(source); + + EXPECT_EQ(env.GetRpcClientIoThreads(), 8U); + EXPECT_EQ(env.GetStoreRpcClientIoThreads(), 4U); + EXPECT_EQ(env.GetTransferEngineRpcClientIoThreads(), 6U); + + source.Set("MC_RPC_CLIENT_IO_THREADS", "12"); + source.Set("MC_STORE_RPC_CLIENT_IO_THREADS", "10"); + source.Set("MC_TE_RPC_CLIENT_IO_THREADS", "11"); + EXPECT_EQ(env.GetRpcClientIoThreads(), 8U); + EXPECT_EQ(env.GetStoreRpcClientIoThreads(), 4U); + EXPECT_EQ(env.GetTransferEngineRpcClientIoThreads(), 6U); +} + +TEST(RpcClientIoThreadsTest, ComponentValuesUseCommonFallback) { + const MockEnvironSource source{{"MC_RPC_CLIENT_IO_THREADS", "8"}, + {"MC_STORE_RPC_CLIENT_IO_THREADS", "0"}, + {"MC_TE_RPC_CLIENT_IO_THREADS", "invalid"}}; + const Environ env(source); + + EXPECT_EQ(env.GetRpcClientIoThreads(), 8U); + EXPECT_EQ(env.GetStoreRpcClientIoThreads(), 8U); + EXPECT_EQ(env.GetTransferEngineRpcClientIoThreads(), 8U); +} + +} // namespace +} // namespace mooncake + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-ep/BuildEpExt.cmake b/mooncake-ep/BuildEpExt.cmake index c94529cbb7..4a5a661ed4 100644 --- a/mooncake-ep/BuildEpExt.cmake +++ b/mooncake-ep/BuildEpExt.cmake @@ -11,6 +11,7 @@ # STAGING_DIR - destination directory for the built .so files # ENGINE_SO_PATH - absolute path to the built engine.cpython-XYZ.so # EP_USE_MUSA - set to "1" when building for MUSA (MTLink path) +# EP_USE_MACA - set to "1" when building for MACA (MTLink path) cmake_minimum_required(VERSION 3.16) @@ -40,6 +41,16 @@ if(EP_USE_MUSA) else() unset(ENV{MOONCAKE_EP_USE_MUSA}) endif() +if(EP_USE_MACA) + set(ENV{MOONCAKE_EP_USE_MACA} "1") + if(DEFINED ENV{MACA_PATH}) + set(ENV{MACA_HOME} "$ENV{MACA_PATH}") + elseif(DEFINED ENV{MACA_HOME}) + set(ENV{MACA_PATH} "$ENV{MACA_HOME}") + endif() +else() + unset(ENV{MOONCAKE_EP_USE_MACA}) +endif() # --------------------------------------------------------------------------- # 2. Ensure engine.so exists in mooncake-wheel/mooncake/ for setup.py linking. diff --git a/mooncake-ep/benchmarks/elastic_buffer_perf.py b/mooncake-ep/benchmarks/elastic_buffer_perf.py new file mode 100644 index 0000000000..0bd2d10d13 --- /dev/null +++ b/mooncake-ep/benchmarks/elastic_buffer_perf.py @@ -0,0 +1,307 @@ +#!/usr/bin/env python3 +"""Performance smoke for Mooncake ElasticBuffer dispatch/combine. + +The benchmark intentionally keeps the workload simple and reproducible. It is +not a full system benchmark; it provides a reviewer-friendly way to verify that +the new elastic path runs repeatedly, supports cached handles, and reports +per-rank effective payload bandwidth. + +Typical single-node usage: + + MOONCAKE_EP_NUM_LOCAL_RANKS=8 \ + torchrun --standalone --nproc_per_node=8 \ + mooncake-ep/benchmarks/elastic_buffer_perf.py --route alltoall +""" + +from __future__ import annotations + +import argparse +import os +import time +from dataclasses import dataclass + +import torch +import torch.distributed as dist +import torch.testing as testing + +from mooncake.mooncake_elastic_buffer import ElasticBuffer + + +@dataclass(frozen=True) +class RoutePlan: + topk_idx: torch.Tensor + expected_recv_tokens: int + expected_combine_factor: int + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Benchmark Mooncake ElasticBuffer") + parser.add_argument("--num-tokens", type=int, default=128) + parser.add_argument("--max-tokens", type=int, default=0) + parser.add_argument("--hidden", type=int, default=4096) + parser.add_argument("--num-experts", type=int, default=256) + parser.add_argument("--num-topk", type=int, default=8) + parser.add_argument("--num-sms", type=int, default=24) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument( + "--route", + choices=("alltoall", "local", "cross"), + default="alltoall", + help="Expert routing pattern to generate.", + ) + parser.add_argument( + "--reuse-handle", + action=argparse.BooleanOptionalAction, + default=True, + help="Reuse the first dispatch handle for later iterations.", + ) + parser.add_argument( + "--check-correctness", + action=argparse.BooleanOptionalAction, + default=True, + help="Check combine output on each iteration.", + ) + parser.add_argument( + "--sync-actual-count", + action="store_true", + help="Synchronize and verify GPU-side received-token count each iteration.", + ) + parser.add_argument("--seed", type=int, default=2026) + return parser.parse_args() + + +def init_distributed(seed: int) -> tuple[int, int]: + if not dist.is_initialized(): + dist.init_process_group("nccl") + rank = dist.get_rank() + local_rank = int(os.environ.get("LOCAL_RANK", rank % torch.cuda.device_count())) + torch.cuda.set_device(local_rank) + torch.set_default_device("cuda") + torch.set_default_dtype(torch.bfloat16) + torch.manual_seed(seed + rank) + return rank, dist.get_world_size() + + +def make_route_plan( + *, + rank: int, + world_size: int, + buffer: ElasticBuffer, + num_tokens: int, + num_topk: int, + num_experts: int, + route: str, +) -> RoutePlan: + local_experts = num_experts // world_size + if local_experts <= 0: + raise ValueError("num_experts must be at least world_size") + expert_offsets = torch.arange(num_topk, device="cuda", dtype=torch.long) % local_experts + + if route == "cross" and buffer.num_scaleout_ranks > 1: + dst_scaleout = (buffer.scaleout_rank_idx + 1) % buffer.num_scaleout_ranks + dst_rank = dst_scaleout * buffer.num_scaleup_ranks + buffer.scaleup_rank_idx + choices = dst_rank * local_experts + expert_offsets + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens, + 1, + ) + + if route == "local" or (route == "cross" and buffer.num_scaleout_ranks == 1): + choices = rank * local_experts + expert_offsets + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens, + 1, + ) + + dst_ranks = (rank + torch.arange(num_topk, device="cuda", dtype=torch.long)) % world_size + choices = dst_ranks * local_experts + expert_offsets + unique_dst_ranks = int(torch.unique(dst_ranks).numel()) + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens * unique_dst_ranks, + unique_dst_ranks, + ) + + +def make_input(rank: int, iteration: int, num_tokens: int, hidden: int) -> torch.Tensor: + base = torch.arange(num_tokens * hidden, device="cuda", dtype=torch.float32) + base = base.view(num_tokens, hidden) + return (base + rank * 1_000_000 + iteration * 17).to(torch.bfloat16).contiguous() + + +def check_output( + *, + rank: int, + route: str, + combined: torch.Tensor, + expected: torch.Tensor, +) -> None: + if route == "local": + if not torch.equal(combined, expected): + diff = (combined.float() - expected.float()).abs().max().item() + raise AssertionError(f"rank={rank}: local-route mismatch, max_diff={diff}") + return + + testing.assert_close( + combined, + expected, + rtol=1e-2, + atol=1e-3, + msg=lambda msg: f"rank={rank}: {route} combine mismatch: {msg}", + ) + + +def main() -> None: + args = parse_args() + rank, world_size = init_distributed(args.seed) + max_tokens = args.max_tokens or max(128, args.num_tokens) + num_experts = args.num_experts + if num_experts % world_size != 0: + raise ValueError("num_experts must be divisible by world_size") + + buffer = ElasticBuffer( + dist.group.WORLD, + num_max_tokens_per_rank=max_tokens, + hidden=args.hidden, + num_topk=args.num_topk, + use_fp8_dispatch=False, + deterministic=False, + allow_hybrid_mode=True, + allow_multiple_reduction=True, + num_gpu_timeout_secs=10, + ) + route_plan = make_route_plan( + rank=rank, + world_size=world_size, + buffer=buffer, + num_tokens=args.num_tokens, + num_topk=args.num_topk, + num_experts=num_experts, + route=args.route, + ) + weights = torch.ones((args.num_tokens, args.num_topk), device="cuda", dtype=torch.float32) + + def run_one(iteration: int, cached_handle): + x = make_input(rank, iteration, args.num_tokens, args.hidden) + dispatch_start = torch.cuda.Event(enable_timing=True) + dispatch_end = torch.cuda.Event(enable_timing=True) + combine_end = torch.cuda.Event(enable_timing=True) + + use_cached = args.reuse_handle and cached_handle is not None + dispatch_start.record() + recv_x, _recv_idx, recv_weights, handle, _ = buffer.dispatch( + x, + topk_idx=None if use_cached else route_plan.topk_idx, + topk_weights=None if use_cached else weights, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + expert_alignment=1, + handle=cached_handle if use_cached else None, + do_cpu_sync=True if not use_cached else None, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + dispatch_end.record() + + actual_recv_tokens = route_plan.expected_recv_tokens + if args.sync_actual_count: + actual_recv_tokens = int(handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()) + if actual_recv_tokens != route_plan.expected_recv_tokens: + raise AssertionError( + f"rank={rank}: got {actual_recv_tokens} received tokens, " + f"expected {route_plan.expected_recv_tokens}" + ) + + combined, _combined_weights, _ = buffer.combine( + recv_x[:actual_recv_tokens].contiguous(), + handle, + topk_weights=( + recv_weights[:actual_recv_tokens].contiguous() + if recv_weights is not None + else None + ), + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + combine_end.record() + torch.cuda.synchronize() + + if args.check_correctness: + expected = (x.float() * route_plan.expected_combine_factor).to(torch.bfloat16) + check_output(rank=rank, route=args.route, combined=combined, expected=expected) + + return ( + handle, + dispatch_start.elapsed_time(dispatch_end), + dispatch_end.elapsed_time(combine_end), + actual_recv_tokens, + ) + + cached_handle = None + for i in range(args.warmup): + cached_handle, _dispatch_ms, _combine_ms, _actual = run_one(i, cached_handle) + + dist.barrier() + torch.cuda.synchronize() + dispatch_ms = [] + combine_ms = [] + recv_tokens = [] + wall_start = time.time() + for i in range(args.iters): + cached_handle, d_ms, c_ms, actual = run_one(args.warmup + i, cached_handle) + dispatch_ms.append(d_ms) + combine_ms.append(c_ms) + recv_tokens.append(actual) + torch.cuda.synchronize() + dist.barrier() + wall_seconds = time.time() - wall_start + + stats = torch.tensor( + [ + sum(dispatch_ms) / len(dispatch_ms), + sum(combine_ms) / len(combine_ms), + min(dispatch_ms), + max(dispatch_ms), + min(combine_ms), + max(combine_ms), + sum(recv_tokens) / len(recv_tokens), + wall_seconds, + ], + device="cuda", + dtype=torch.float64, + ) + gathered = [torch.empty_like(stats) for _ in range(world_size)] + dist.all_gather(gathered, stats) + + if rank == 0: + table = torch.stack(gathered).cpu() + payload_bytes = table[:, 6].mean().item() * args.hidden * 2 + dispatch_avg_ms = table[:, 0].mean().item() + combine_avg_ms = table[:, 1].mean().item() + print( + "MOONCAKE_ELASTIC_PERF_OK", + f"world={world_size}", + f"route={args.route}", + f"reuse_handle={int(args.reuse_handle)}", + f"tokens={args.num_tokens}", + f"hidden={args.hidden}", + f"topk={args.num_topk}", + f"scaleout={buffer.num_scaleout_ranks}", + f"scaleup={buffer.num_scaleup_ranks}", + f"dispatch_avg_ms={dispatch_avg_ms:.3f}", + f"combine_avg_ms={combine_avg_ms:.3f}", + f"recv_tokens_avg={table[:, 6].mean().item():.1f}", + f"effective_payload_MB_per_rank={payload_bytes / 1e6:.1f}", + f"dispatch_effective_GBps={payload_bytes / dispatch_avg_ms / 1e6:.2f}", + f"combine_effective_GBps={payload_bytes / combine_avg_ms / 1e6:.2f}", + flush=True, + ) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/mooncake-ep/benchmarks/ep_benchmark/README.md b/mooncake-ep/benchmarks/ep_benchmark/README.md new file mode 100644 index 0000000000..3ff35c0d14 --- /dev/null +++ b/mooncake-ep/benchmarks/ep_benchmark/README.md @@ -0,0 +1,109 @@ +# Mooncake EP Dispatch/Combine Benchmark + +Measures dispatch/combine throughput and tail latency of the Mooncake EP Buffer +under uniform, k-hot incast, and Zipfian routing patterns. + +## Usage + +Single-node (uses `mp.spawn` internally): + +```bash +python mooncake-ep/benchmarks/ep_benchmark/run_ep_benchmark.py \ + --num-ranks 8 --num-experts 256 --hidden-size 7168 \ + --top-k 8 --num-tokens 1024 --dtype bf16 \ + --routing-mode k_hot --hot-experts 32 --hot-fraction 0.9 \ + --zero-copy --async-finish \ + --warmup-iters 20 --iters 100 \ + --pg-backend nccl \ + --json-output results/khot_8rank.json +``` + +Multi-node via `torchrun` (`--num-ranks` is ignored, `WORLD_SIZE` from torchrun is used): + +```bash +torchrun --nnodes=2 --nproc_per_node=4 --rdzv_backend=c10d \ + --rdzv_endpoint=$HEAD_NODE \ + mooncake-ep/benchmarks/ep_benchmark/run_ep_benchmark.py \ + --num-experts 256 --hidden-size 7168 \ + --top-k 8 --num-tokens 1024 --dtype bf16 \ + --routing-mode k_hot --hot-experts 32 --hot-fraction 0.9 \ + --zero-copy --async-finish \ + --pg-backend nccl \ + --json-output results/khot_8node.json +``` + +Using a config file (CLI flags override config values): + +```bash +python mooncake-ep/benchmarks/ep_benchmark/run_ep_benchmark.py \ + --config mooncake-ep/benchmarks/ep_benchmark/configs/cuda_zipfian.json \ + --num-ranks 8 +``` + +## Parameters + +| Flag | Default | Description | +|------|---------|-------------| +| `--config` | None | Path to JSON config file (CLI flags override) | +| `--num-ranks` | `8` | Number of EP ranks / GPUs | +| `--num-experts` | `256` | Total experts across all ranks | +| `--hidden-size` | `7168` | Hidden dimension | +| `--top-k` | `8` | Experts per token | +| `--num-tokens` | `1024` | Tokens per rank | +| `--dtype` | `bf16` | Dispatch data type (`bf16` or `fp8`) | +| `--routing-mode` | `uniform` | `uniform`, `k_hot`, or `zipf` | +| `--hot-experts` | `32` | Hot expert count (k_hot mode) | +| `--hot-fraction` | `0.9` | Fraction of tokens to hot experts (k_hot mode) | +| `--zipf-alpha` | `1.0` | Zipf distribution alpha (zipf mode) | +| `--zero-copy` | off | Zero-copy combine via get_next_combine_buffer | +| `--async-finish` | off | Event-based async sync | +| `--return-recv-hook` | off | Hook-based sync (mutually exclusive with `--async-finish`) | +| `--pg-backend` | `nccl` | `nccl` or `mooncake` (mooncake requires RDMA) | +| `--warmup-iters` | `20` | Warmup iterations (not timed) | +| `--iters` | `100` | Measured iterations | +| `--seed` | `0` | Base random seed (each rank uses `seed + rank`) | +| `--json-output` | stdout | Output JSON file path | +| `--master-addr` | `127.0.0.1` | Process group master address | +| `--master-port` | `29500` | Process group master port | + +`num_experts` must be divisible by `num_ranks`. + +`end_to_end_latency_ms` measures the full dispatch → mock expert forward → combine cycle, not dispatch + combine alone. + +## Output + +```json +{ + "benchmark": "mooncake_ep", + "world_size": 8, + "num_experts": 256, + "hidden_size": 7168, + "routing_mode": "k_hot", + "metrics": { + "dispatch_latency_ms": {"p50": 0, "p90": 0, "p99": 0, "p999": 0, "mean": 0}, + "combine_latency_ms": {"p50": 0, "p90": 0, "p99": 0, "p999": 0, "mean": 0}, + "end_to_end_latency_ms": {"p50": 0, "p90": 0, "p99": 0, "p999": 0, "mean": 0}, + "tokens_per_second": 0, + "expert_load": { + "max_tokens_per_expert": 0, + "min_tokens_per_expert": 0, + "mean_tokens_per_expert": 0, + "imbalance_ratio": 0 + }, + "per_rank_stats": { + "dispatch_latency_ms": { + "rank0": {"p50": 0, "p90": 0, "p99": 0, "p999": 0, "mean": 0}, + "rank1": {"p50": 0, "p90": 0, "p99": 0, "p999": 0, "mean": 0} + }, + "combine_latency_ms": { + "rank0": {"p50": 0, "p90": 0, "p99": 0, "p999": 0, "mean": 0}, + "rank1": {"p50": 0, "p90": 0, "p99": 0, "p999": 0, "mean": 0} + }, + "e2e_latency_ms": { + "rank0": {"p50": 0, "p90": 0, "p99": 0, "p999": 0, "mean": 0}, + "rank1": {"p50": 0, "p90": 0, "p99": 0, "p999": 0, "mean": 0} + } + } + } +} +``` diff --git a/mooncake-ep/benchmarks/ep_benchmark/configs/cuda_incast_k_hot.json b/mooncake-ep/benchmarks/ep_benchmark/configs/cuda_incast_k_hot.json new file mode 100644 index 0000000000..f81ade6de3 --- /dev/null +++ b/mooncake-ep/benchmarks/ep_benchmark/configs/cuda_incast_k_hot.json @@ -0,0 +1,18 @@ +{ + "backend": "cuda", + "num_ranks": 8, + "num_experts": 256, + "hidden_size": 7168, + "top_k": 8, + "num_tokens": 1024, + "dtype": "bf16", + "routing_mode": "k_hot", + "hot_experts": 32, + "hot_fraction": 0.9, + "zero_copy": true, + "async_finish": true, + "return_recv_hook": false, + "pg_backend": "nccl", + "warmup_iters": 20, + "iters": 100 +} diff --git a/mooncake-ep/benchmarks/ep_benchmark/configs/cuda_uniform.json b/mooncake-ep/benchmarks/ep_benchmark/configs/cuda_uniform.json new file mode 100644 index 0000000000..9c5c251706 --- /dev/null +++ b/mooncake-ep/benchmarks/ep_benchmark/configs/cuda_uniform.json @@ -0,0 +1,16 @@ +{ + "backend": "cuda", + "num_ranks": 8, + "num_experts": 256, + "hidden_size": 7168, + "top_k": 8, + "num_tokens": 1024, + "dtype": "bf16", + "routing_mode": "uniform", + "zero_copy": true, + "async_finish": true, + "return_recv_hook": false, + "pg_backend": "nccl", + "warmup_iters": 20, + "iters": 100 +} diff --git a/mooncake-ep/benchmarks/ep_benchmark/configs/cuda_zipfian.json b/mooncake-ep/benchmarks/ep_benchmark/configs/cuda_zipfian.json new file mode 100644 index 0000000000..9884d9cb36 --- /dev/null +++ b/mooncake-ep/benchmarks/ep_benchmark/configs/cuda_zipfian.json @@ -0,0 +1,17 @@ +{ + "backend": "cuda", + "num_ranks": 8, + "num_experts": 256, + "hidden_size": 7168, + "top_k": 8, + "num_tokens": 1024, + "dtype": "bf16", + "routing_mode": "zipf", + "zipf_alpha": 1.0, + "zero_copy": true, + "async_finish": true, + "return_recv_hook": false, + "pg_backend": "nccl", + "warmup_iters": 20, + "iters": 100 +} diff --git a/mooncake-ep/benchmarks/ep_benchmark/metrics.py b/mooncake-ep/benchmarks/ep_benchmark/metrics.py new file mode 100644 index 0000000000..31578b9afb --- /dev/null +++ b/mooncake-ep/benchmarks/ep_benchmark/metrics.py @@ -0,0 +1,143 @@ +"""Metrics computation and JSON output assembly for the EP benchmark.""" + +import json +from typing import Any, Optional + +import numpy as np +import torch + + +def compute_percentiles(values) -> dict[str, float]: + """Compute p50, p90, p99, p999, and mean from latency values (in ms).""" + if len(values) == 0: + return {"p50": 0.0, "p90": 0.0, "p99": 0.0, "p999": 0.0, "mean": 0.0} + arr = np.asarray(values, dtype=np.float64) + return { + "p50": float(np.percentile(arr, 50)), + "p90": float(np.percentile(arr, 90)), + "p99": float(np.percentile(arr, 99)), + "p999": float(np.percentile(arr, 99.9)), + "mean": float(arr.mean()), + } + + +def compute_global_expert_load( + topk_idx: torch.Tensor, + num_experts: int, + group: Optional["torch.distributed.ProcessGroup"] = None, +) -> dict[str, float]: + """ + Compute global expert load by all-reducing per-rank bincount. + + Each rank computes its local expert assignment counts via bincount, + then all_reduce(sum) to get global counts. + + Returns: + Dict with max/min/mean_tokens_per_expert and imbalance_ratio (= max/mean). + """ + import torch.distributed as dist + + flat = topk_idx.flatten().to(torch.int64) + local_counts = torch.bincount(flat, minlength=num_experts).to(torch.float64) + + if group is not None and dist.is_initialized(): + global_counts = local_counts.clone() + dist.all_reduce(global_counts, op=dist.ReduceOp.SUM, group=group) + else: + global_counts = local_counts + + max_val = float(global_counts.max().item()) + min_val = float(global_counts.min().item()) + mean_val = float(global_counts.mean().item()) + imbalance_ratio = max_val / mean_val if mean_val > 0 else 0.0 + + return { + "max_tokens_per_expert": max_val, + "min_tokens_per_expert": min_val, + "mean_tokens_per_expert": mean_val, + "imbalance_ratio": imbalance_ratio, + } + + +def assemble_json_output( + backend: str, + world_size: int, + num_experts: int, + hidden_size: int, + top_k: int, + num_tokens: int, + dtype: str, + routing_mode: str, + zero_copy: bool, + async_finish: bool, + return_recv_hook: bool, + warmup_iters: int, + iters: int, + hot_experts: int | None = None, + hot_fraction: float | None = None, + zipf_alpha: float | None = None, + dispatch_latencies_ms: list[float] | None = None, + combine_latencies_ms: list[float] | None = None, + e2e_latencies_ms: list[float] | None = None, + expert_load: dict[str, float] | None = None, + pg_backend: str = "nccl", + *, + per_rank_stats: dict[str, Any], +) -> dict[str, Any]: + """Assemble the final JSON-serializable output dict matching the RFC schema.""" + result = { + "benchmark": "mooncake_ep", + "backend": backend, + "world_size": world_size, + "num_experts": num_experts, + "hidden_size": hidden_size, + "top_k": top_k, + "num_tokens": num_tokens, + "dtype": dtype, + "routing_mode": routing_mode, + "zero_copy": zero_copy, + "async_finish": async_finish, + "return_recv_hook": return_recv_hook, + "warmup_iters": warmup_iters, + "iters": iters, + "pg_backend": pg_backend, + } + + if routing_mode == "k_hot": + result["hot_experts"] = hot_experts + result["hot_fraction"] = hot_fraction + elif routing_mode == "zipf": + result["zipf_alpha"] = zipf_alpha + + dispatch_pct = compute_percentiles(dispatch_latencies_ms or []) + combine_pct = compute_percentiles(combine_latencies_ms or []) + e2e_pct = compute_percentiles(e2e_latencies_ms or []) + + mean_e2e_ms = e2e_pct["mean"] + total_tokens = num_tokens * world_size + tokens_per_second = ( + total_tokens / (mean_e2e_ms / 1000.0) if mean_e2e_ms > 0 else 0.0 + ) + + result["metrics"] = { + "dispatch_latency_ms": dispatch_pct, + "combine_latency_ms": combine_pct, + "end_to_end_latency_ms": e2e_pct, + "tokens_per_second": tokens_per_second, + "expert_load": expert_load + or { + "max_tokens_per_expert": 0, + "min_tokens_per_expert": 0, + "mean_tokens_per_expert": 0, + "imbalance_ratio": 0, + }, + "per_rank_stats": per_rank_stats, + } + + return result + + +def write_json_output(data: dict[str, Any], path: str) -> None: + """Write the JSON output to a file with pretty formatting.""" + with open(path, "w") as f: + json.dump(data, f, indent=2) diff --git a/mooncake-ep/benchmarks/ep_benchmark/results/.gitignore b/mooncake-ep/benchmarks/ep_benchmark/results/.gitignore new file mode 100644 index 0000000000..d6b7ef32c8 --- /dev/null +++ b/mooncake-ep/benchmarks/ep_benchmark/results/.gitignore @@ -0,0 +1,2 @@ +* +!.gitignore diff --git a/mooncake-ep/benchmarks/ep_benchmark/routing.py b/mooncake-ep/benchmarks/ep_benchmark/routing.py new file mode 100644 index 0000000000..aa52aa6fb6 --- /dev/null +++ b/mooncake-ep/benchmarks/ep_benchmark/routing.py @@ -0,0 +1,153 @@ +""" +Expert-parallel routing generators for the EP benchmark. + +""" + +import torch + + +def generate_topk_weights( + num_tokens: int, top_k: int, device: torch.device, generator: torch.Generator +) -> torch.Tensor: + """Generate softmax-normalized weights.""" + raw = torch.rand( + num_tokens, top_k, dtype=torch.float32, device=device, generator=generator + ) + return torch.softmax(raw, dim=-1) + + +def uniform_routing( + num_tokens: int, + num_experts: int, + top_k: int, + device: torch.device | None = None, + seed: int = 0, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Uniform routing: each token picks top-k from the full expert set with + approximately uniform probability. + """ + if device is None: + device = torch.device("cpu") + generator = torch.Generator(device=device).manual_seed(seed) + scores = torch.randn( + num_tokens, + num_experts, + dtype=torch.float32, + device=device, + generator=generator, + ) + topk_idx = torch.topk(scores, top_k, dim=-1)[1].to(torch.int64) + topk_weights = generate_topk_weights(num_tokens, top_k, device, generator) + return topk_idx, topk_weights + + +def k_hot_routing( + num_tokens: int, + num_experts: int, + top_k: int, + device: torch.device | None = None, + seed: int = 0, + hot_experts: int = 4, + hot_fraction: float = 0.9, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + """Incast / k-hot routing: *hot_fraction* of tokens pick ALL top-k from + the first *hot_experts* set; the rest route uniformly. + + Raises: + ValueError: if hot_experts < top_k (cannot pick top_k distinct experts + from a set smaller than top_k). + ValueError: if hot_experts > num_experts. + ValueError: if hot_fraction not in (0, 1). + """ + if device is None: + device = torch.device("cpu") + if hot_experts > num_experts: + raise ValueError(f"hot_experts ({hot_experts}) > num_experts ({num_experts})") + if hot_experts < top_k: + raise ValueError( + f"hot_experts ({hot_experts}) < top_k ({top_k}); " + f"cannot pick {top_k} distinct experts from {hot_experts} hot experts" + ) + if not (0.0 < hot_fraction < 1.0): + raise ValueError(f"hot_fraction must be in (0, 1), got {hot_fraction}") + + generator = torch.Generator(device=device).manual_seed(seed) + + num_hot = int(num_tokens * hot_fraction) + num_cold = num_tokens - num_hot + + hot_scores = torch.randn( + num_hot, hot_experts, dtype=torch.float32, device=device, generator=generator + ) + hot_topk_idx = torch.topk(hot_scores, top_k, dim=-1)[1].to(torch.int64) + + cold_scores = torch.randn( + num_cold, + num_experts, + dtype=torch.float32, + device=device, + generator=generator, + ) + cold_topk_idx = torch.topk(cold_scores, top_k, dim=-1)[1].to(torch.int64) + + topk_idx = torch.cat([hot_topk_idx, cold_topk_idx], dim=0) + perm = torch.randperm(num_tokens, generator=generator, device=device) + topk_idx = topk_idx[perm] + + topk_weights = generate_topk_weights(num_tokens, top_k, device, generator) + return topk_idx, topk_weights + + +def zipfian_routing( + num_tokens: int, + num_experts: int, + top_k: int, + device: torch.device | None = None, + seed: int = 0, + zipf_alpha: float = 1.0, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor]: + """Zipfian routing: expert selection probability follows Zipf(alpha). + + P(expert_i) proportional to 1 / (i + 1)^alpha. + + Uses the Gumbel-max trick for top-k sampling without replacement: + for each token, add Gumbel(0,1) noise to log-probabilities, then take top-k. + """ + if device is None: + device = torch.device("cpu") + if zipf_alpha <= 0: + raise ValueError(f"zipf_alpha must be > 0, got {zipf_alpha}") + + generator = torch.Generator(device=device).manual_seed(seed) + + expert_ids = torch.arange(1, num_experts + 1, dtype=torch.float32, device=device) + log_probs = -zipf_alpha * torch.log(expert_ids) + log_probs = log_probs - log_probs.logsumexp(dim=0) + + gumbel = -torch.log( + -torch.log( + torch.rand( + num_tokens, + num_experts, + dtype=torch.float32, + device=device, + generator=generator, + ).clamp(min=1e-20) + ) + ) + noisy_scores = log_probs.unsqueeze(0) + gumbel + topk_idx = torch.topk(noisy_scores, top_k, dim=-1)[1].to(torch.int64) + + topk_weights = generate_topk_weights(num_tokens, top_k, device, generator) + return topk_idx, topk_weights + + +ROUTING_MODES = { + "uniform": uniform_routing, + "k_hot": k_hot_routing, + "zipf": zipfian_routing, +} diff --git a/mooncake-ep/benchmarks/ep_benchmark/run_ep_benchmark.py b/mooncake-ep/benchmarks/ep_benchmark/run_ep_benchmark.py new file mode 100755 index 0000000000..923765d964 --- /dev/null +++ b/mooncake-ep/benchmarks/ep_benchmark/run_ep_benchmark.py @@ -0,0 +1,467 @@ +#!/usr/bin/env python3 +"""Mooncake EP (Expert Parallel) dispatch/combine benchmark. + +Measures throughput and tail latency of the Mooncake EP Buffer's +dispatch/combine cycle under uniform, k-hot incast, and Zipfian routing +patterns. +""" + +import argparse +import json +import os +import sys + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from metrics import ( + assemble_json_output, + compute_global_expert_load, + compute_percentiles, + write_json_output, +) +from routing import ROUTING_MODES + +from mooncake.mooncake_ep_buffer import Buffer + + +def parse_args(): + # Pre-parse --config to load JSON defaults before full parsing + pre_parser = argparse.ArgumentParser(add_help=False) + pre_parser.add_argument("--config", default=None) + pre_args, _ = pre_parser.parse_known_args() + + config_defaults = {} + if pre_args.config: + with open(pre_args.config) as f: + config_defaults = json.load(f) + + parser = argparse.ArgumentParser( + description="Mooncake EP dispatch/combine benchmark" + ) + parser.add_argument( + "--config", + type=str, + default=None, + help="Path to JSON config file (CLI flags override config values)", + ) + parser.add_argument( + "--backend", + type=str, + default="cuda", + help="Device backend label for output JSON (default: cuda)", + ) + parser.add_argument( + "--num-ranks", + type=int, + default=8, + help="Number of EP ranks / GPUs", + ) + parser.add_argument( + "--num-experts", + type=int, + default=256, + help="Total number of experts across all ranks", + ) + parser.add_argument( + "--hidden-size", + type=int, + default=7168, + help="Hidden dimension", + ) + parser.add_argument( + "--top-k", + type=int, + default=8, + help="Number of experts each token routes to", + ) + parser.add_argument( + "--num-tokens", + type=int, + default=1024, + help="Number of tokens per rank", + ) + parser.add_argument( + "--dtype", + type=str, + default="bf16", + choices=["bf16", "fp8"], + help="Data type for dispatch (default: bf16)", + ) + parser.add_argument( + "--routing-mode", + type=str, + default="uniform", + choices=["uniform", "k_hot", "zipf"], + help="Routing pattern (default: uniform)", + ) + parser.add_argument( + "--hot-experts", + type=int, + default=32, + help="Number of hot experts for k_hot mode (default: 32)", + ) + parser.add_argument( + "--hot-fraction", + type=float, + default=0.9, + help="Fraction of tokens routed to hot experts in k_hot mode (default: 0.9)", + ) + parser.add_argument( + "--zipf-alpha", + type=float, + default=1.0, + help="Zipf alpha parameter for zipf mode (default: 1.0)", + ) + parser.add_argument( + "--zero-copy", + action="store_true", + help="Use zero-copy combine via get_next_combine_buffer", + ) + parser.add_argument( + "--async-finish", + action="store_true", + help="Use async_finish for dispatch/combine (event-based sync)", + ) + parser.add_argument( + "--return-recv-hook", + action="store_true", + help="Use return_recv_hook for dispatch/combine (hook-based sync)", + ) + parser.add_argument( + "--pg-backend", + type=str, + default="nccl", + choices=["nccl", "mooncake"], + help="Process group backend (default: nccl; mooncake requires RDMA)", + ) + parser.add_argument( + "--warmup-iters", + type=int, + default=20, + help="Number of warmup iterations (default: 20)", + ) + parser.add_argument( + "--iters", + type=int, + default=100, + help="Number of measured iterations (default: 100)", + ) + parser.add_argument( + "--seed", type=int, default=0, help="Base random seed (default: 0)" + ) + parser.add_argument( + "--json-output", + type=str, + default=None, + help="Path to write JSON output (default: stdout on rank 0)", + ) + parser.add_argument( + "--master-addr", + type=str, + default="127.0.0.1", + help="Master address for process group (default: 127.0.0.1)", + ) + parser.add_argument( + "--master-port", + type=int, + default=29500, + help="Master port for process group (default: 29500)", + ) + + if config_defaults: + parser.set_defaults(**config_defaults) + + return parser.parse_args() + + +def validate_args(args): + if args.num_experts % args.num_ranks != 0: + raise ValueError( + f"num_experts ({args.num_experts}) must be divisible by " + f"num_ranks ({args.num_ranks})" + ) + if args.async_finish and args.return_recv_hook: + raise ValueError("async_finish and return_recv_hook are mutually exclusive") + if args.dtype == "fp8" and args.hidden_size % 128 != 0: + raise ValueError( + f"hidden_size ({args.hidden_size}) must be divisible by 128 for fp8 dtype" + ) + + +def dequantize_fp8(x_fp8, scales): + """Dequantize FP8 packed data to bfloat16 (pattern from test_ep_grid.py:43-48).""" + hidden = x_fp8.shape[-1] + x_view = x_fp8.reshape(-1, hidden // 128, 128).float() + scales_view = scales.reshape(-1, hidden // 128, 1).float() + dequantized = (x_view * scales_view).reshape(x_fp8.shape) + return dequantized.to(torch.bfloat16) + + +class EPBenchmarkWorker: + """Per-rank worker that sets up the EP buffer, runs warmup and measured + iterations, then aggregates and outputs results.""" + + def __init__(self, rank, local_rank, args): + self.rank = rank + self.args = args + self.use_fp8 = args.dtype == "fp8" + self.num_local_experts = args.num_experts // args.num_ranks + self.timeout_us = -1 + + torch.cuda.set_device(local_rank) + + if args.pg_backend == "mooncake": + try: + import mooncake.pg # noqa: F401 + except Exception as exc: + raise RuntimeError( + "Failed to import mooncake.pg; required for --pg-backend=mooncake" + ) from exc + + dist.init_process_group( + backend=args.pg_backend, rank=rank, world_size=args.num_ranks + ) + self.group = dist.group.WORLD + + num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint( + args.num_tokens, args.hidden_size, args.num_ranks, args.num_experts + ) + self.buf = Buffer(self.group, num_ep_buffer_bytes) + + self.topk_idx, self.topk_weights = ROUTING_MODES[args.routing_mode]( + num_tokens=args.num_tokens, + num_experts=args.num_experts, + top_k=args.top_k, + device=torch.device("cuda"), + seed=args.seed + rank, + hot_experts=args.hot_experts, + hot_fraction=args.hot_fraction, + zipf_alpha=args.zipf_alpha, + ) + + self.x = torch.randn( + args.num_tokens, args.hidden_size, dtype=torch.bfloat16, device="cuda" + ) + self.active_ranks = torch.ones( + (args.num_ranks,), dtype=torch.int32, device="cuda" + ) + self.out_tensor = torch.zeros_like(self.x) + + def run_dispatch(self): + recv_x, recv_count, handle, event, hook = self.buf.dispatch( + self.x, + self.topk_idx, + self.active_ranks, + num_max_dispatch_tokens_per_rank=self.args.num_tokens, + num_experts=self.args.num_experts, + timeout_us=self.timeout_us, + use_fp8=self.use_fp8, + async_finish=self.args.async_finish, + return_recv_hook=self.args.return_recv_hook, + ) + if self.args.return_recv_hook: + hook() + elif self.args.async_finish: + event.current_stream_wait() + return recv_x, recv_count, handle + + def prepare_combine(self, expert_out, handle): + if self.args.zero_copy: + cb_buf = self.buf.get_next_combine_buffer(handle) + cb_buf.copy_(expert_out) + return cb_buf.contiguous(), handle + else: + return expert_out.contiguous(), handle + + def run_combine(self, expert_to_pass, handle): + combined_x, event, hook = self.buf.combine( + expert_to_pass, + self.topk_idx, + self.topk_weights, + self.active_ranks, + timeout_us=self.timeout_us, + handle=handle, + zero_copy=self.args.zero_copy, + async_finish=self.args.async_finish, + return_recv_hook=self.args.return_recv_hook, + out=self.out_tensor, + ) + if self.args.return_recv_hook: + hook() + elif self.args.async_finish: + event.current_stream_wait() + return combined_x + + def mock_expert_forward(self, recv_x): + if self.use_fp8: + recv_bf16 = dequantize_fp8(recv_x[0], recv_x[1]) + else: + recv_bf16 = recv_x + + expert_out = torch.empty_like(recv_bf16) + for le in range(self.num_local_experts): + expert_id = self.rank * self.num_local_experts + le + expert_out[le] = recv_bf16[le] * (expert_id * 0.1 + 1.0) + return expert_out.to(torch.bfloat16) + + def warmup(self): + for _ in range(self.args.warmup_iters): + recv_x, _, handle = self.run_dispatch() + expert_out = self.mock_expert_forward(recv_x) + expert_to_pass, handle = self.prepare_combine(expert_out, handle) + self.run_combine(expert_to_pass, handle) + torch.cuda.synchronize() + + def run_measured(self): + n = self.args.iters + dispatch_starts = [torch.cuda.Event(enable_timing=True) for _ in range(n)] + dispatch_ends = [torch.cuda.Event(enable_timing=True) for _ in range(n)] + combine_starts = [torch.cuda.Event(enable_timing=True) for _ in range(n)] + combine_ends = [torch.cuda.Event(enable_timing=True) for _ in range(n)] + + for i in range(n): + dispatch_starts[i].record() + + recv_x, _, handle = self.run_dispatch() + + dispatch_ends[i].record() + + expert_out = self.mock_expert_forward(recv_x) + expert_to_pass, handle = self.prepare_combine(expert_out, handle) + + combine_starts[i].record() + + self.run_combine(expert_to_pass, handle) + + combine_ends[i].record() + + torch.cuda.synchronize() + + dispatch_latencies = [ + dispatch_starts[i].elapsed_time(dispatch_ends[i]) for i in range(n) + ] + combine_latencies = [ + combine_starts[i].elapsed_time(combine_ends[i]) for i in range(n) + ] + # e2e includes dispatch + mock_expert_forward + combine (full cycle) + e2e_latencies = [ + dispatch_starts[i].elapsed_time(combine_ends[i]) for i in range(n) + ] + + return dispatch_latencies, combine_latencies, e2e_latencies + + def aggregate_and_output( + self, dispatch_latencies, combine_latencies, e2e_latencies + ): + dispatch_tensor = torch.tensor( + dispatch_latencies, dtype=torch.float64, device="cuda" + ) + combine_tensor = torch.tensor( + combine_latencies, dtype=torch.float64, device="cuda" + ) + e2e_tensor = torch.tensor(e2e_latencies, dtype=torch.float64, device="cuda") + + all_dispatch = [ + torch.zeros_like(dispatch_tensor) for _ in range(self.args.num_ranks) + ] + all_combine = [ + torch.zeros_like(combine_tensor) for _ in range(self.args.num_ranks) + ] + all_e2e = [torch.zeros_like(e2e_tensor) for _ in range(self.args.num_ranks)] + + dist.all_gather(all_dispatch, dispatch_tensor, group=self.group) + dist.all_gather(all_combine, combine_tensor, group=self.group) + dist.all_gather(all_e2e, e2e_tensor, group=self.group) + + expert_load = compute_global_expert_load( + self.topk_idx, self.args.num_experts, group=self.group + ) + + if self.rank == 0: + # A step completes only when the slowest rank finishes (critical path). + dispatch_critical = torch.stack(all_dispatch).max(dim=0).values.cpu().tolist() + combine_critical = torch.stack(all_combine).max(dim=0).values.cpu().tolist() + e2e_critical = torch.stack(all_e2e).max(dim=0).values.cpu().tolist() + + per_rank_stats = {} + for metric, tensors in ( + ("dispatch", all_dispatch), + ("combine", all_combine), + ("e2e", all_e2e), + ): + per_rank_stats[f"{metric}_latency_ms"] = { + f"rank{r}": compute_percentiles(latencies.cpu()) + for r, latencies in enumerate(tensors) + } + + result = assemble_json_output( + backend=self.args.backend, + world_size=self.args.num_ranks, + num_experts=self.args.num_experts, + hidden_size=self.args.hidden_size, + top_k=self.args.top_k, + num_tokens=self.args.num_tokens, + dtype=self.args.dtype, + routing_mode=self.args.routing_mode, + zero_copy=self.args.zero_copy, + async_finish=self.args.async_finish, + return_recv_hook=self.args.return_recv_hook, + warmup_iters=self.args.warmup_iters, + iters=self.args.iters, + hot_experts=( + self.args.hot_experts if self.args.routing_mode == "k_hot" else None + ), + hot_fraction=( + self.args.hot_fraction + if self.args.routing_mode == "k_hot" + else None + ), + zipf_alpha=( + self.args.zipf_alpha if self.args.routing_mode == "zipf" else None + ), + dispatch_latencies_ms=dispatch_critical, + combine_latencies_ms=combine_critical, + e2e_latencies_ms=e2e_critical, + expert_load=expert_load, + pg_backend=self.args.pg_backend, + per_rank_stats=per_rank_stats, + ) + + if self.args.json_output: + write_json_output(result, self.args.json_output) + print(f"Results written to {self.args.json_output}", file=sys.stderr) + else: + print(json.dumps(result, indent=2)) + + def run(self): + self.warmup() + dispatch_latencies, combine_latencies, e2e_latencies = self.run_measured() + self.aggregate_and_output(dispatch_latencies, combine_latencies, e2e_latencies) + dist.barrier() + dist.destroy_process_group() + + +def _worker_entry(rank, args): + EPBenchmarkWorker(rank, rank, args).run() + + +def main(): + args = parse_args() + + if "RANK" in os.environ: + args.num_ranks = int(os.environ["WORLD_SIZE"]) + validate_args(args) + rank = int(os.environ["RANK"]) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + EPBenchmarkWorker(rank, local_rank, args).run() + return + + validate_args(args) + os.environ.setdefault("MASTER_ADDR", args.master_addr) + os.environ.setdefault("MASTER_PORT", str(args.master_port)) + mp.spawn(_worker_entry, args=(args,), nprocs=args.num_ranks) + + +if __name__ == "__main__": + main() diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_api.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_api.cuh new file mode 100644 index 0000000000..aa99b8e0c9 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_api.cuh @@ -0,0 +1,20 @@ +#pragma once + +// Official DeepEP elastic source import surface for Mooncake. +// +// This umbrella intentionally lives under include/elastic, not in the legacy EP +// include root. It keeps the imported elastic implementation discoverable +// while allowing the host launch/runtime glue to opt in file-by-file without +// perturbing legacy Buffer dispatch/combine symbols. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_buffer.h b/mooncake-ep/include/elastic/mooncake_ep_elastic_buffer.h new file mode 100644 index 0000000000..6b5dacdd07 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_buffer.h @@ -0,0 +1,170 @@ +#ifndef MOONCAKE_EP_ELASTIC_BUFFER_H +#define MOONCAKE_EP_ELASTIC_BUFFER_H + +#include +#include +#include +#include +#include +#include + +#include + +namespace mooncake { + +struct ElasticLaunchContext; + +struct ElasticTopology { + int rank_idx = 0; + int num_ranks = 1; + int num_rdma_ranks = 1; + int num_nvlink_ranks = 1; + int num_scaleout_ranks = 1; + int num_scaleup_ranks = 1; + int scaleout_rank_idx = 0; + int scaleup_rank_idx = 0; + bool hybrid_enabled = false; +}; + +struct ElasticConfig { + int64_t num_max_tokens_per_rank = 0; + int64_t hidden = 0; + int64_t num_topk = 0; + bool use_fp8_dispatch = false; + bool deterministic = false; + bool allow_hybrid_mode = true; + bool allow_multiple_reduction = true; + bool prefer_overlap_with_compute = true; + int sl_idx = 3; + int num_allocated_qps = 0; + int num_cpu_timeout_secs = 300; + int num_gpu_timeout_secs = 100; +}; + +struct ElasticNativeHandle { + bool do_expand = false; + int num_experts = 0; + int expert_alignment = 1; + int num_max_tokens_per_rank = 0; + int num_sms = 0; + torch::Tensor topk_idx; + torch::Tensor psum_num_recv_tokens_per_scaleup_rank; + torch::Tensor psum_num_recv_tokens_per_expert; + torch::Tensor recv_src_metadata; + torch::Tensor recv_layout_range; + torch::Tensor dst_buffer_slot_idx; + std::optional token_metadata_at_forward; + std::optional channel_linked_list; + std::vector num_recv_tokens_per_expert_list; +}; + +struct ElasticDispatchOutput { + torch::Tensor recv_x; + std::optional recv_x_scales; + std::optional recv_topk_idx; + std::optional recv_topk_weights; + ElasticNativeHandle handle; + std::optional event; +}; + +struct ElasticCombineOutput { + torch::Tensor combined_x; + std::optional combined_topk_weights; + std::optional event; +}; + +class MooncakeElasticBuffer { + public: + MooncakeElasticBuffer(int rank, int num_ranks, int64_t num_buffer_bytes, + int64_t num_max_tokens_per_rank, int64_t hidden, + int64_t num_topk, bool use_fp8_dispatch, + bool deterministic, bool allow_hybrid_mode, + bool allow_multiple_reduction, + bool prefer_overlap_with_compute, int sl_idx, + int num_allocated_qps, int num_cpu_timeout_secs, + int num_gpu_timeout_secs); + + ~MooncakeElasticBuffer(); + + static int64_t calculate_buffer_size(int num_ranks, + int64_t num_max_tokens_per_rank, + int64_t hidden, int64_t num_topk, + bool use_fp8_dispatch, + bool allow_hybrid_mode, + bool allow_multiple_reduction); + + std::tuple get_physical_domain_size() const; + std::tuple get_logical_domain_size() const; + int get_theoretical_num_sms(int num_experts, int num_topk) const; + + ElasticDispatchOutput dispatch( + const torch::Tensor& x, const std::optional& sf, + const torch::Tensor& topk_idx, + const std::optional& topk_weights, + torch::Tensor& active_ranks, int num_experts, + int num_max_tokens_per_rank, int expert_alignment, int num_sms, + bool do_expand, bool do_cpu_sync, bool async_with_compute_stream, + const std::optional& cached_handle = std::nullopt); + + ElasticCombineOutput combine( + const torch::Tensor& x, const ElasticNativeHandle& handle, + const std::optional& topk_weights, + torch::Tensor& active_ranks, int num_sms, + bool async_with_compute_stream, + const std::optional& out); + + MooncakeEpBuffer& native_buffer() { return *native_buffer_; } + + bool ibgda_disabled() const { return native_buffer_->ibgda_disabled(); } + bool use_fast_path() { return native_buffer_->use_fast_path(); } + void update_local_qpns() { native_buffer_->update_local_qpns(); } + bool is_roce() const { return native_buffer_->is_roce(); } + void sync_ibgda_peers(const std::vector& remote_addrs, + const std::vector& remote_keys, + const std::vector>& peer_qpns, + const std::vector>& peer_lids, + const std::vector& subnet_prefixes, + const std::vector& interface_ids, + const std::vector& active_ranks_mask) { + native_buffer_->sync_ibgda_peers(remote_addrs, remote_keys, peer_qpns, + peer_lids, subnet_prefixes, + interface_ids, active_ranks_mask); + } + std::tuple get_mr_info() { + return native_buffer_->get_mr_info(); + } + std::tuple get_gid() { return native_buffer_->get_gid(); } + std::vector get_local_qpns() { + return native_buffer_->get_local_qpns(); + } + std::vector get_local_lids() { + return native_buffer_->get_local_lids(); + } + std::vector get_ipc_handle() { + return native_buffer_->get_ipc_handle(); + } + void sync_nvlink_ipc_handles( + const std::vector>& remote_handles, + const std::vector& active_ranks_mask) { + native_buffer_->sync_nvlink_ipc_handles(remote_handles, + active_ranks_mask); + } + + private: + ElasticConfig config_; + ElasticTopology topology_; + std::unique_ptr native_buffer_; + int64_t host_workspace_bytes_ = 0; + void* host_workspace_ = nullptr; + void* mapped_host_workspace_ = nullptr; + + static ElasticLaunchContext make_launch_context( + MooncakeEpBuffer& buffer, const ElasticTopology& topology, + void* mapped_host_workspace, int64_t timeout_cycles); + static ElasticTopology discover_topology(int rank, int num_ranks, + bool allow_hybrid_mode); +}; + +} // namespace mooncake + +#endif // MOONCAKE_EP_ELASTIC_BUFFER_H diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh new file mode 100644 index 0000000000..9e36b344f7 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_official.cuh @@ -0,0 +1,365 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include +#include + +#include + +namespace mooncake::elastic { + +template (), + int kNumTokensInLayout = get_num_tokens_in_layout< + kAllowMultipleReduction, kNumRanks, kNumTopk>(), + typename team_t = std::conditional_t< + kIsScaleupNVLink, transport::ScaleupTeam, transport::WorldTeam>> +__global__ void __launch_bounds__(kNumThreads, 1) + combine_impl(nv_bfloat16* x, float* topk_weights, int* src_metadata, + int* psum_num_recv_tokens_per_scaleup_rank, + const device::CommCtx comm_ctx, void* buffer, void* workspace, + const int rank_idx, int num_reduced_tokens) { + // Utils + const auto sm_idx = static_cast(blockIdx.x); + const auto thread_idx = static_cast(threadIdx.x); + const auto warp_idx = (ptx::get_warp_idx() + rank_idx) % kNumWarps; + const auto lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; + constexpr bool kDoExpandedSend = + not kAllowMultipleReduction and kUseExpandedLayout; + + // We should assign the real number of received tokens if without CPU sync + if (num_reduced_tokens == kNumMaxTokensPerRank * kNumRanks) + num_reduced_tokens = + __ldg(psum_num_recv_tokens_per_scaleup_rank + kNumRanks - 1); + + // Buffer layouts + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto token_layout = + layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false); + const auto tma_buffer = + layout::BufferLayout(token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx) + .get_token_buffer(0); + const auto recv_buffer = layout::BufferLayout( + token_layout, kNumTokensInLayout, kNumMaxTokensPerRank, buffer); + const auto send_buffer = layout::BufferLayout( + token_layout, kNumRanks, + kNumMaxTokensPerRank * (kDoExpandedSend ? kNumTopk : 1), + recv_buffer.get_buffer_end_ptr()); + + // Init TMA + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Expanding mode must not be backward + if constexpr (kUseExpandedLayout) EP_DEVICE_ASSERT(topk_weights == nullptr); + + // Gin handle + // We treat each warp as a "channel" + const auto [qp_idx, sharing_mode] = + comm::get_qp_mode(sm_idx, warp_idx); + const auto gin = transport::MooncakeGin(comm_ctx, qp_idx, sharing_mode, + kNumQPs, 0, 0, 0, kNumRanks); + + // Full barrier to ensure the remote buffer is available + const auto workspace_layout = + layout::WorkspaceLayout(workspace, 1, kNumRanks, kNumExperts); + comm::gpu_barrier(gin, workspace_layout, 0, rank_idx, sm_idx, + thread_idx); + + // Do TMA writes into the remote buffers + int num_tokens_per_warp = + math::ceil_div(num_reduced_tokens, kNumSMs * kNumWarps); + const int token_start_idx = num_tokens_per_warp * global_warp_idx; + const int token_end_idx = + min(token_start_idx + num_tokens_per_warp, num_reduced_tokens); + for (int i = token_start_idx; i < token_end_idx; ++i) { + // The master slot index during dispatch + constexpr int kMetadataStride = 2 + kNumTopk; + const int src_token_idx = + __ldg(src_metadata + i * kMetadataStride) % kNumMaxTokensPerRank; + const int src_rank_topk_idx = + __ldg(src_metadata + i * kMetadataStride + 1); + const int src_rank_idx = src_rank_topk_idx / kNumTopk; + const int src_topk_idx = src_rank_topk_idx % kNumTopk; + + // Directly to the remote or via RDMA + const bool nvlink_bypass = + gin.is_nvlink_accessible(src_rank_idx); + layout::TokenLayout master_token_buffer = [=]() { + // NVLink bypass + if (nvlink_bypass) { + auto token_buffer = + recv_buffer + .get_rank_buffer(kUseRankLayout ? rank_idx + : src_topk_idx) + .get_token_buffer(src_token_idx); + token_buffer.set_base_ptr(gin.get_sym_ptr( + token_buffer.get_base_ptr(), src_rank_idx)); + return token_buffer; + } + + // Use RDMA + return send_buffer.get_rank_buffer(src_rank_idx) + .get_token_buffer(src_token_idx); + }(); + + // Hidden requirements + EP_STATIC_ASSERT( + kHidden % (32 * sizeof(int4) / sizeof(nv_bfloat16)) == 0, + "Invalid hidden"); + using combine_vec_t = + typename CombineVecTraits::vec_t; + constexpr int kHiddenVec = + kHidden * sizeof(nv_bfloat16) / sizeof(combine_vec_t); + + // Read source indices for expand mode + int stored_topk_slot_idx = -1; + if constexpr (kUseExpandedLayout) { + if (lane_idx < kNumTopk) + stored_topk_slot_idx = + __ldg(src_metadata + i * kMetadataStride + (2 + lane_idx)); + __syncwarp(); + } + + // 3 cases: + // - no expand + no reduce, or expand + no reduce + // - expand + reduce + // - expand + send all + auto reduce_valid_mask = ptx::gather(stored_topk_slot_idx >= 0); + auto no_local_reduce = + not kUseExpandedLayout or + (kAllowMultipleReduction and __popc(reduce_valid_mask) == 1); + if (no_local_reduce) { + int token_idx_in_tensor = i; + if constexpr (kUseExpandedLayout) + token_idx_in_tensor = + ptx::exchange(stored_topk_slot_idx, + ptx::get_master_lane_idx(reduce_valid_mask)); + + // No reduce +#ifdef MOONCAKE_EP_USE_MUSA + { + const auto src_ptr = math::advance_ptr( + x, static_cast(token_idx_in_tensor) * + kNumHiddenBytes); + auto* dst_ptr = static_cast( + master_token_buffer.get_base_ptr()); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; + vec_idx += 32) { + ptx::st_na(dst_ptr + vec_idx, src_ptr[vec_idx]); + } + __syncwarp(); + __threadfence_system(); + } +#else + if (ptx::elect_one_sync()) { + const auto load_ptr = math::advance_ptr( + x, static_cast(token_idx_in_tensor) * + kNumHiddenBytes); + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, + mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + ptx::tma_store_1d(master_token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); +#endif + } else if constexpr (kAllowMultipleReduction) { + // Do local reduction + // Sort valid top-k indices to front + int topk_slot_idx[kNumTopk]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, [=](const int& idx) { + return ptx::exchange(stored_topk_slot_idx, idx); + }); + + // Reduce into shared memory + constexpr int kUnrollFactor = + get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, + static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ + [=](const int& slot_idx) { + return math::advance_ptr( + x, slot_idx * static_cast(kNumHiddenBytes)); + }, + /* Wait buffer release */ + [=]() { + ptx::tma_store_wait(); + __syncwarp(); + }); + ptx::tma_store_fence(); + __syncwarp(); + + // Issue TMA stores +#ifdef MOONCAKE_EP_USE_MUSA + { + const auto* src_ptr = static_cast( + tma_buffer.get_base_ptr()); + auto* dst_ptr = static_cast( + master_token_buffer.get_base_ptr()); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; + vec_idx += 32) { + ptx::st_na(dst_ptr + vec_idx, src_ptr[vec_idx]); + } + __syncwarp(); + __threadfence_system(); + } +#else + if (ptx::elect_one_sync()) { + ptx::tma_store_1d(master_token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); +#endif + } else { +// No local reduction, send all data (expanded send) +#pragma unroll + for (int k = 0; k < kNumTopk; ++k) { + const auto slot_idx = ptx::exchange(stored_topk_slot_idx, k); + if (slot_idx >= 0) { + const auto src_token_ptr = math::advance_ptr( + x, slot_idx * static_cast(kNumHiddenBytes)); + const auto token_buffer = + recv_buffer.get_rank_buffer(k).get_token_buffer( + src_token_idx); +#ifdef MOONCAKE_EP_USE_MUSA + if (nvlink_bypass) { + auto* dst_ptr = + static_cast(gin.get_sym_ptr( + token_buffer.get_base_ptr(), src_rank_idx)); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; + vec_idx += 32) { + ptx::st_na(dst_ptr + vec_idx, + src_token_ptr[vec_idx]); + } + } else { + const auto send_token_buffer = + send_buffer.get_rank_buffer(src_rank_idx) + .get_token_buffer(src_token_idx * kNumTopk + k); + auto* dst_ptr = static_cast( + send_token_buffer.get_base_ptr()); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; + vec_idx += 32) { + ptx::st_na(dst_ptr + vec_idx, + src_token_ptr[vec_idx]); + } + __syncwarp(); + if (ptx::elect_one_sync()) { + gin.put(token_buffer.get_base_ptr(), + send_token_buffer.get_base_ptr(), + kNumHiddenBytes, src_rank_idx); + } + } + __syncwarp(); + __threadfence_system(); +#else + if (ptx::elect_one_sync()) { + // Load + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + src_token_ptr, mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + + if (nvlink_bypass) { + // Write into the same position + ptx::tma_store_1d( + gin.get_sym_ptr( + token_buffer.get_base_ptr(), src_rank_idx), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } else { + // Write to the RDMA send buffer + const auto send_token_buffer = + send_buffer.get_rank_buffer(src_rank_idx) + .get_token_buffer(src_token_idx * kNumTopk + + k); + ptx::tma_store_1d(send_token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), + kNumHiddenBytes); + ptx::tma_store_commit(); + ptx::tma_store_wait(); + + // Issue RDMA + gin.put(token_buffer.get_base_ptr(), + send_token_buffer.get_base_ptr(), + kNumHiddenBytes, src_rank_idx); + } + } + __syncwarp(); +#endif + } + } + } + + // Write topk weights + if (not kUseExpandedLayout and topk_weights != nullptr and + lane_idx < kNumTopk) { + const float value = __ldg(topk_weights + (i * kNumTopk + lane_idx)); +#ifdef MOONCAKE_EP_USE_MUSA + ptx::st_relaxed_sys( + master_token_buffer.get_topk_weights_ptr() + lane_idx, value); +#else + master_token_buffer.get_topk_weights_ptr()[lane_idx] = value; +#endif + } + __syncwarp(); +#ifdef MOONCAKE_EP_USE_MUSA + __threadfence_system(); +#endif + + // Wait send buffer's TMA store and issue RDMA send + // NOTES: `kDoExpandedSend` mode has already issued + if (not kDoExpandedSend and not nvlink_bypass and + ptx::elect_one_sync()) { + ptx::tma_store_wait(); + const auto dst_ptr = + recv_buffer + .get_rank_buffer(kUseRankLayout ? rank_idx : src_topk_idx) + .get_token_buffer(src_token_idx) + .get_base_ptr(); + gin.put(dst_ptr, master_token_buffer.get_base_ptr(), + master_token_buffer.get_num_bytes(), + src_rank_idx); + } + } + + // Final barrier to ensure data arrival + comm::gpu_barrier(gin, workspace_layout, 0, rank_idx, sm_idx, + thread_idx); +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh new file mode 100644 index 0000000000..6e41aebab8 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_reduce_epilogue.cuh @@ -0,0 +1,212 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include +#include + +#include + +namespace mooncake::elastic { + +template (), + int kNumTokensInLayout = get_num_tokens_in_layout< + kAllowMultipleReduction, kNumRanks, kNumTopk>()> +__global__ void __launch_bounds__(kNumThreads, 1) + combine_reduce_epilogue_impl(nv_bfloat16* combined_x, + float* combined_topk_weights, + topk_idx_t* combined_topk_idx, + void* recv_buffer, void* bias_0, void* bias_1, + const int num_combined_tokens, + const int scaleout_rank_idx, + const int scaleup_rank_idx) { + constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks; + constexpr int kNumExpertsPerRank = + kNumExperts / (kNumScaleupRanks * kNumScaleoutRanks); + EP_STATIC_ASSERT(kNumExperts % (kNumScaleupRanks * kNumScaleoutRanks) == 0, + "Invalid number of experts or ranks"); + + // Utils + const auto sm_idx = static_cast(blockIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = + warp_idx * kNumSMs + + sm_idx; // NOTES: Here we prioritize distributing tasks to different + // SMs to ensure that the last wave is evenly concentrated on + // each SM. + + // Load buffers from scale-out or scale-up ranks + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto comm_token_layout = + layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false); + const auto comm_buffer = + layout::BufferLayout(comm_token_layout, kNumTokensInLayout, + kNumMaxTokensPerRank, recv_buffer); + + // Store buffers + const auto output_token_layout = + layout::TokenLayout(kNumHiddenBytes, 0, 0, false); + const auto output_buffer = layout::BufferLayout( + output_token_layout, 1, num_combined_tokens, combined_x); + const auto tma_buffer = + layout::BufferLayout(output_token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx) + .get_token_buffer(0); + + // Bias layout + const auto bias_0_buffer = layout::BufferLayout( + output_token_layout, 1, num_combined_tokens, bias_0); + const auto bias_1_buffer = layout::BufferLayout( + output_token_layout, 1, num_combined_tokens, bias_1); + + // Will block until the main combine kernel has finished and all data are + // visible NOTES: PDL is used, please do not use `__ldg` +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + // Read from buffers and do reduction + for (int token_idx = global_warp_idx; token_idx < num_combined_tokens; + token_idx += kNumWarps * kNumSMs) { + // Preprocess all indices + int stored_dst_rank_idx = -1, stored_dst_expert_idx = -1; + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections"); + if (lane_idx < kNumTopk) { + stored_dst_expert_idx = static_cast( + combined_topk_idx[token_idx * kNumTopk + lane_idx]); + stored_dst_rank_idx = + stored_dst_expert_idx >= 0 + ? stored_dst_expert_idx / (kNumScaleoutRanks == 1 + ? kNumExpertsPerRank + : kNumExpertsPerScaleout) + : -1; + } + __syncwarp(); + + // Sort valid top-k indices to front + const auto [should_deduplicate, + deduplicate_key] = [&]() -> std::pair { + if constexpr (kUseExpandedLayout and not kAllowMultipleReduction) { + // Activations are never reduced before + return {false, 0}; + } else if constexpr (kNumScaleoutRanks != 1 and + not kUseExpandedLayout and + not kAllowMultipleReduction) { + // Hybrid mode without expanded layout and multiple reduction. + // Should deduplicate on a per-rank basis + return {true, stored_dst_expert_idx >= 0 + ? stored_dst_expert_idx / kNumExpertsPerRank + : -1}; + } else { + // Should deduplicate on a per-rank (for non-hybrid mode) or a + // per-scale-rank (for hybrid mode) basis + return {true, stored_dst_rank_idx}; + } + }(); + auto reduce_valid_mask = + should_deduplicate + ? ptx::gather(ptx::deduplicate(deduplicate_key, lane_idx) and + stored_dst_rank_idx >= 0) + : ptx::gather(stored_dst_rank_idx >= 0); + int topk_slot_idx[kNumTokensInLayout]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, [=](const int& idx) { + return kUseRankLayout ? ptx::exchange(stored_dst_rank_idx, idx) + : idx; + }); + + // Iterate over per-hidden-chunk stage + using combine_vec_t = + typename CombineVecTraits::vec_t; + constexpr int kHiddenVec = + kHidden * sizeof(nv_bfloat16) / sizeof(combine_vec_t); + constexpr int kUnrollFactor = get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, + static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ + [=](const int& slot_idx) { + return static_cast( + comm_buffer.get_rank_buffer(slot_idx) + .get_token_buffer(token_idx) + .get_base_ptr()); + }, + /* Wait buffer release */ + [=]() { + ptx::tma_store_wait(); + __syncwarp(); + }, + /* Bias 0 */ bias_0 == nullptr + ? nullptr + : static_cast( + bias_0_buffer.get_token_buffer(token_idx).get_base_ptr()), + /* Bias 1 */ bias_1 == nullptr + ? nullptr + : static_cast( + bias_1_buffer.get_token_buffer(token_idx) + .get_base_ptr())); + ptx::tma_store_fence(); + __syncwarp(); + + // Issue TMA copy +#ifdef MOONCAKE_EP_USE_MUSA + { + const auto* src_ptr = + static_cast(tma_buffer.get_base_ptr()); + auto* dst_ptr = static_cast( + output_buffer.get_token_buffer(token_idx).get_base_ptr()); +#pragma unroll 1 + for (int vec_idx = lane_idx; vec_idx < kHiddenVec; vec_idx += 32) { + dst_ptr[vec_idx] = src_ptr[vec_idx]; + } + __syncwarp(); + } +#else + if (ptx::elect_one_sync()) { + ptx::tma_store_1d( + output_buffer.get_token_buffer(token_idx).get_base_ptr(), + tma_buffer.get_base_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); +#endif + + // Write top-k weights + if (combined_topk_weights != nullptr) { + const auto master_lane_idx = + ptx::get_master_lane_idx(ptx::match(stored_dst_rank_idx)); + if (lane_idx < kNumTopk) { + float value = 0; + if (stored_dst_rank_idx >= 0) { + const auto dst_ptr = + comm_buffer + .get_rank_buffer(kUseRankLayout + ? stored_dst_rank_idx + : master_lane_idx) + .get_token_buffer(token_idx) + .get_topk_weights_ptr() + + lane_idx; + value = *dst_ptr; + } + combined_topk_weights[token_idx * kNumTopk + lane_idx] = value; + } + __syncwarp(); + } + } +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_utils.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_utils.cuh new file mode 100644 index 0000000000..b2e3623f2a --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_combine_utils.cuh @@ -0,0 +1,209 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +namespace mooncake::elastic { + +template +constexpr bool use_rank_layout() { + if constexpr (not kAllowMultipleReduction) return false; + return kNumRanks <= kNumTopk; +} + +template +constexpr int get_num_tokens_in_layout() { + return use_rank_layout() + ? kNumRanks + : kNumTopk; +} + +template +constexpr int get_max_unroll_factor() { + for (int i = kMaxUnrollFactor; i >= 1; --i) + if (kLength % (kWarpSize * i) == 0) return i; +#ifdef MOONCAKE_EP_USE_MUSA + return 1; +#else + throw std::logic_error("Invalid length, cannot find unrolling factor"); +#endif +} + +// Determine the vector type for combine loads/stores based on arch and hidden +// size alignment +template +struct CombineVecTraits { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 1000) + // On SM100+, use longlong4_t (32 bytes) if hidden is aligned, otherwise + // fall back to int4 (16 bytes) + static constexpr bool kUseLonglong4 = + (kHiddenBytes % sizeof(longlong4_t) == 0) and + ((kHiddenBytes / sizeof(longlong4_t)) % 32 == 0); + using vec_t = std::conditional_t; +#else + using vec_t = int4; +#endif +}; + +template +__device__ __forceinline__ void compute_topk_slots( + int (&topk_slot_idx)[kNumValidTopk], uint32_t mask, + const fetch_func_t& fetch_func) { +#pragma unroll + for (int k = 0; k < kNumValidTopk; ++k) { + const int lowest_idx = __ffs(mask) - 1; + // Here we perform the exchange unconditionally to avoid `BRA.DIV` + const auto fetched = fetch_func(lowest_idx); + mask &= mask - 1; + topk_slot_idx[k] = lowest_idx >= 0 ? fetched : -1; + } +} + +template +__device__ __forceinline__ void combine_reduce( + const int& lane_idx, int (&topk_slot_idx)[kNumValidTopk], + vec_t* dst_buffer_ptr, + const get_src_buffer_ptr_func_t& get_src_buffer_ptr_func, + const wait_buffer_func_t& wait_buffer_func, vec_t* bias_0 = nullptr, + vec_t* bias_1 = nullptr) { + constexpr int kNumElemsPerVec = sizeof(vec_t) / sizeof(nv_bfloat16); + EP_STATIC_ASSERT(kNumElemsPerVec % 2 == 0, "Invalid number of elements"); + EP_STATIC_ASSERT(kHiddenVec % (kUnrollFactor * 32) == 0, + "Invalid unrolling"); + + // We use BF16 add as much as possible, as casting is slow + const bool enable_hadd_bypass = + (bias_0 == nullptr and bias_1 == nullptr) and + (kNumValidTopk <= 2 or topk_slot_idx[2] < 0); + EP_STATIC_ASSERT(kNumValidTopk > 0, "Invalid top-k"); + + if (enable_hadd_bypass) { +#pragma unroll 1 + for (int i = 0; i < kHiddenVec / (kUnrollFactor * 32); ++i) { + // Read values 0 + const auto slot_0 = topk_slot_idx[0]; + const auto src_base_ptr_0 = get_src_buffer_ptr_func(slot_0); + vec_t values_0[kUnrollFactor] = {}; +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { + values_0[j] = ptx::ldg_with_gez_pred( + src_base_ptr_0 + + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), + slot_0); + } + + // Read values 1 + vec_t values_1[kUnrollFactor] = {}; + const auto slot_1 = kNumValidTopk == 1 ? -1 : topk_slot_idx[1]; + const auto src_base_ptr_1 = get_src_buffer_ptr_func(slot_1); +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { + values_1[j] = ptx::ldg_with_gez_pred( + src_base_ptr_1 + + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), + slot_1); + } + + // Wait buffer releases for the first write + if (i == 0) wait_buffer_func(); + + // Reduce into shared memory + const auto bf162_view_0 = reinterpret_cast(values_0); + const auto bf162_view_1 = reinterpret_cast(values_1); +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { +#pragma unroll + for (int l = 0; l < kNumElemsPerVec / 2; ++l) { + const int idx = j * (kNumElemsPerVec / 2) + l; +#ifdef MOONCAKE_EP_USE_MUSA + bf162_view_0[idx] = __floats2bfloat162_rn( + __low2float(bf162_view_0[idx]) + + __low2float(bf162_view_1[idx]), + __high2float(bf162_view_0[idx]) + + __high2float(bf162_view_1[idx])); +#else + bf162_view_0[idx] += bf162_view_1[idx]; +#endif + } + dst_buffer_ptr[i * (kUnrollFactor * 32) + j * 32 + lane_idx] = + values_0[j]; + } + } + } else { +#pragma unroll 1 + for (int i = 0; i < kHiddenVec / (kUnrollFactor * 32); ++i) { + // Add bias + float2 reduced[kUnrollFactor * kNumElemsPerVec / 2] = {}; + const auto add_bias = [&](const vec_t* base_ptr) { + // Read + vec_t values[kUnrollFactor]; +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) + values[j] = ptx::ldg(base_ptr + i * (kUnrollFactor * 32) + + j * 32 + lane_idx); + + // Reduce + const auto bf162_view = reinterpret_cast(values); +#pragma unroll + for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++j) + ptx::accumulate(reduced[j], bf162_view[j]); + }; + bias_0 != nullptr ? add_bias(bias_0) : void(); + bias_1 != nullptr ? add_bias(bias_1) : void(); + +#pragma unroll + for (int k = 0; k < kNumValidTopk; ++k) { + // We have a limitation on `k` to reduce the branch instruction + // count + if (k >= kNumExpectedTopk and topk_slot_idx[k] < 0) break; + + // Read values + const auto src_base_ptr = + get_src_buffer_ptr_func(topk_slot_idx[k]); + vec_t values[kUnrollFactor] = {}; +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { + values[j] = ptx::ldg_with_gez_pred( + src_base_ptr + + (i * (kUnrollFactor * 32) + j * 32 + lane_idx), + topk_slot_idx[k]); + } + + // Reduce + const auto bf162_view = reinterpret_cast(values); +#pragma unroll + for (int j = 0; j < kUnrollFactor * kNumElemsPerVec / 2; ++j) + ptx::accumulate(reduced[j], bf162_view[j]); + } + + // Wait buffer releases for the first write + if (i == 0) wait_buffer_func(); + +// Cast into shared memory +#pragma unroll + for (int j = 0; j < kUnrollFactor; ++j) { + vec_t casted_value; + auto bf162_view = + reinterpret_cast(&casted_value); +#pragma unroll + for (int l = 0; l < kNumElemsPerVec / 2; ++l) { + const auto value = reduced[j * (kNumElemsPerVec / 2) + l]; +#ifdef MOONCAKE_EP_USE_MUSA + bf162_view[l] = __floats2bfloat162_rn(value.x, value.y); +#else + bf162_view[l] = __float22bfloat162_rn(value); +#endif + } + dst_buffer_ptr[i * (kUnrollFactor * 32) + j * 32 + lane_idx] = + casted_value; + } + } + } +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_comm.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_comm.cuh new file mode 100644 index 0000000000..886c2f12d8 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_comm.cuh @@ -0,0 +1,190 @@ +#pragma once + +#include +#include + +#include +#include +#include + +namespace mooncake::elastic::comm { + +static constexpr int64_t kNumOneSecCycles = 2000000000; + +static constexpr int kDeviceBarrierTag = 0; +static constexpr int kKernelBarrierTag = 1; +static constexpr int kDispatchTag0 = 2; +static constexpr int kDispatchTag1 = 3; +static constexpr int kCombineTag0 = 4; +static constexpr int kCombineTag1 = 5; +static constexpr int kHybridDispatchTag0 = 6; +static constexpr int kHybridDispatchTag1 = 7; +static constexpr int kHybridCombineTag0 = 8; +static constexpr int kHybridCombineTag1 = 9; + +static constexpr int kFlushAllAllocatedQPs = -1; + +template +__device__ __forceinline__ void timeout_while(const bool& condition, + const func_t& func, + int64_t start_clock = 0) { + if (start_clock == 0) start_clock = clock64(); + while (condition) { + const bool timeout = kNumTimeoutCycles >= 0 && + (clock64() - start_clock >= kNumTimeoutCycles); + if (func(timeout)) break; + if (timeout) { + const auto timeout_start = clock64(); + while (clock64() - timeout_start < kNumOneSecCycles) { + } + ptx::trap(); + } + } +} + +template +__device__ __forceinline__ void timeout_while(const func_t& func, + const int64_t& start_clock = 0) { + timeout_while(true, func, start_clock); +} + +template +__forceinline__ __device__ void local_grid_sync( + const layout::WorkspaceLayout& workspace, const int& thread_idx) { +#ifdef MOONCAKE_EP_USE_MUSA + (void)kNumThreads; + __shared__ unsigned long long ticket; + __syncthreads(); + if (thread_idx == 0) { + ticket = atomicAdd( + workspace.get_nvl_barrier_counter_ptr(kKernelBarrierTag), 1ULL); + } + __syncthreads(); + const auto target = ((ticket / kNumSMs) + 1ULL) * kNumSMs; + timeout_while(thread_idx == 0, [=](const bool&) { + return ptx::ld_volatile( + workspace.get_nvl_barrier_counter_ptr(kKernelBarrierTag)) >= + target; + }); + __syncthreads(); +#else + (void)workspace; + (void)thread_idx; + (gridDim.x > 1) ? cooperative_groups::this_grid().sync() : __syncthreads(); +#endif +} + +template +__device__ __forceinline__ std::pair get_qp_mode( + const int& sm_idx, const int& channel_in_sm_idx, + const bool& is_notify_warp = false) { + if constexpr (kNumQPs == 1) return {0, 1}; + if (is_notify_warp) return {0, 0}; + + constexpr int kQPStartIdx = static_cast(kWithNotifyWarps); + constexpr int kNumAvailableQPs = kNumQPs - kQPStartIdx; + if constexpr (kNumSMs <= kNumAvailableQPs) { + const int num_qps_in_sm = (kNumAvailableQPs / kNumSMs) + + (sm_idx < (kNumAvailableQPs % kNumSMs)); + return {kQPStartIdx + sm_idx + + (channel_in_sm_idx % max(1, num_qps_in_sm)) * kNumSMs, + 0}; + } else { + const auto global_channel_idx = + sm_idx * kNumChannelsPerSM + channel_in_sm_idx; + return {kQPStartIdx + (global_channel_idx % max(1, kNumAvailableQPs)), + 1}; + } +} + +template +__forceinline__ __device__ void mooncake_barrier_wo_local_sync( + const transport::MooncakeGin& gin, const layout::WorkspaceLayout& workspace, + const int& rank_idx, const int& sm_idx, const int& thread_idx) { + if (kNumSMs > 1 && sm_idx > 0) return; + + const int status = + static_cast((*workspace.get_nvl_barrier_counter_ptr(kTag)) & 3); + const int phase = status & 1; + const int sign = status >> 1; + const int* base_signal = workspace.get_nvl_barrier_signal_ptr(kTag, phase); + + if (thread_idx < kNumRanks) { + auto* dst_ptr = const_cast(base_signal) + rank_idx; + gin.red_add_rel(dst_ptr, sign ? -1 : 1, thread_idx); + } + __syncthreads(); + + if (thread_idx == 0) + atomicAdd(workspace.get_nvl_barrier_counter_ptr(kTag), 1ULL); + + timeout_while( + thread_idx == 0, [=](const bool& is_last_check) { + int sum = 0; +#pragma unroll + for (int i = 0; i < kNumRanks; ++i) { + sum += + ptx::ld_acquire_sys(const_cast(base_signal) + i); + } + // Mooncake's portable barrier uses one additive slot per source + // rank. Each positive phase adds +1 into a zeroed phase slot; the + // matching negative phase later adds -1 into the same phase slot. + // This matches RDMA atomic-add semantics and avoids relying on a + // remote store primitive for non-P2P peers. + const auto target = sign ? 0 : kNumRanks; + if (sum == target) return true; + if (is_last_check) { + printf( + "Mooncake elastic barrier timeout, tag: %d, rank: %d, " + "signal-sum: %d, target: %d\n", + kTag, rank_idx, sum, target); + } + return false; + }); +} + +template +__forceinline__ __device__ void gpu_barrier( + const transport::MooncakeGin& gin, const layout::WorkspaceLayout& workspace, + const int& scaleout_rank_idx, const int& scaleup_rank_idx, + const int& sm_idx, const int& thread_idx, bool do_scaleout = true, + bool do_scaleup = true) { + if constexpr (kFlushStores) gin.flush(); + if constexpr (kSyncAtStart) { + local_grid_sync(workspace, + thread_idx); + } + + do_scaleout &= kNumScaleoutRanks > 1; + do_scaleup &= kNumScaleupRanks > 1; + if (do_scaleup && !do_scaleout) { + mooncake_barrier_wo_local_sync(gin, workspace, scaleup_rank_idx, + sm_idx, thread_idx); + } else if (do_scaleout && !do_scaleup) { + mooncake_barrier_wo_local_sync( + gin, workspace, scaleout_rank_idx, sm_idx, thread_idx); + } else { + const int global_rank = + scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx; + mooncake_barrier_wo_local_sync< + transport::WorldTeam, kNumScaleoutRanks * kNumScaleupRanks, kNumSMs, + kNumThreads, kNumTimeoutCycles, kTag>(gin, workspace, global_rank, + sm_idx, thread_idx); + } + + if constexpr (kSyncAtEnd) { + local_grid_sync(workspace, + thread_idx); + } +} + +} // namespace mooncake::elastic::comm diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh new file mode 100644 index 0000000000..1850361106 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_compiled.cuh @@ -0,0 +1,115 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +// Make CLion CUDA indexing work +#ifdef __CLION_IDE__ +#define __CUDA_ARCH__ 900 +#define __CUDACC_RDC__ +#define __CUDACC__ +#endif + +// Remove Torch restrictions +#ifdef __CUDA_NO_HALF_CONVERSIONS__ +#undef __CUDA_NO_HALF_CONVERSIONS__ +#endif +#ifdef __CUDA_NO_HALF_OPERATORS__ +#undef __CUDA_NO_HALF_OPERATORS__ +#endif +#ifdef __CUDA_NO_HALF2_OPERATORS__ +#undef __CUDA_NO_HALF2_OPERATORS__ +#endif +#ifdef __CUDA_NO_BFLOAT16_CONVERSIONS__ +#undef __CUDA_NO_BFLOAT16_CONVERSIONS__ +#endif +#ifdef __CUDA_NO_BFLOAT162_OPERATORS__ +#undef __CUDA_NO_BFLOAT162_OPERATORS__ +#endif + +#include +#include +#include + +#if defined(MOONCAKE_EP_USE_MUSA) && defined(__MCC__) && \ + !defined(MOONCAKE_EP_MUSA_LDG_DEFINED) +#define MOONCAKE_EP_MUSA_LDG_DEFINED +template +__device__ __forceinline__ dtype_t __ldg(const dtype_t* ptr) { + return *ptr; +} +#endif + +#ifndef DISABLE_SM90_FEATURES +#include +#elif !defined(MOONCAKE_EP_USE_MUSA) +// Ampere does not support FP8 features +#define __NV_E4M3 0 +#define __NV_E5M2 1 +typedef int __nv_fp8_interpretation_t; +typedef int __nv_fp8x4_e4m3; +typedef uint8_t __nv_fp8_storage_t; +#endif + +// Compatibility: 256 bits LD/ST instructions +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(CUDART_VERSION) and \ + CUDART_VERSION >= 13000 +using longlong4_t = longlong4_32a; +#define make_longlong4_t make_longlong4_32a +#else +struct alignas(32) longlong4_t { + long long x, y, z, w; +}; +__device__ __forceinline__ longlong4_t make_longlong4_t(const long long& x, + const long long& y, + const long long& z, + const long long& w) { + return {x, y, z, w}; +} +#endif + +#ifndef EP_NUM_TOPK_IDX_BITS +#define EP_NUM_TOPK_IDX_BITS 64 +#endif + +namespace mooncake { + +#ifndef DISABLE_SM90_FEATURES +constexpr bool kEnableSM90Features = true; +#else +constexpr bool kEnableSM90Features = false; +#endif + +template +struct int_with_bits; +template <> +struct int_with_bits<8> { + using type = int8_t; +}; +template <> +struct int_with_bits<16> { + using type = int16_t; +}; +template <> +struct int_with_bits<32> { + using type = int32_t; +}; +template <> +struct int_with_bits<64> { + using type = int64_t; +}; + +using topk_idx_t = int_with_bits::type; + +union sf_pack_t { + float fp32; + int ue8m0x4; +}; + +constexpr int kNumTMAAlignedBytes = 16; +constexpr int kNumAlignedSFPacks = 16 / sizeof(sf_pack_t); + +// Some communication channel settings +constexpr int kNumMaxChannels = 1024; + +} // namespace mooncake diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh new file mode 100644 index 0000000000..47cfca8401 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_copy_epilogue.cuh @@ -0,0 +1,277 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include +#include +#include + +namespace mooncake::elastic { + +template < + bool kDoExpand, bool kCachedMode, + // NOTES: this channel concept only applies for scale-out ranks + int kNumSMs, int kNumChannels, int kNumWarps, int kNumScaleoutRanks, + int kNumScaleupRanks, int kNumHiddenBytes, int kNumSFPacks, + int kNumMaxTokensPerRank, int kNumExperts, int kNumTopk, + int kNumRanks = kNumScaleoutRanks * kNumScaleupRanks, + int kNumThreads = kNumWarps * 32, + int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, + kNumChannels), + bool kDoCreateLinkedList = (kNumScaleoutRanks > 1 and not kCachedMode)> +__global__ void __launch_bounds__(kNumThreads, 1) dispatch_copy_epilogue_impl( + void* buffer, void* workspace, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, void* recv_x, sf_pack_t* recv_sf, + topk_idx_t* recv_topk_idx, float* recv_topk_weights, int* recv_src_metadata, + int* channel_linked_list, int num_recv_tokens, + const int recv_sf_token_stride, const int recv_sf_hidden_stride, + const int scaleout_rank_idx, const int scaleup_rank_idx) { + // Utils + const auto sm_idx = static_cast(blockIdx.x), + thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; + + // For top-k index transformations + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; + const auto rank_idx = + scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx; + const auto expert_start_idx = kNumExpertsPerRank * rank_idx, + expert_end_idx = kNumExpertsPerRank * (rank_idx + 1); + + // Buffer layouts + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto token_layout = layout::TokenLayout( + kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true); + const auto tma_buffer = + layout::BufferLayout(token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx) + .get_token_buffer(0); + const auto scaleup_buffer = layout::BufferLayout( + token_layout, kNumScaleupRanks, + kNumScaleoutRanks * kNumMaxTokensPerRank, buffer); + + // Init TMA + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Will block until the main dispatch kernel has finished and all data are + // visible NOTES: PDL is used, please do not use `__ldg` +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + // For no CPU sync case, the number of received tokens should be read from + // the GPU tensor + if (num_recv_tokens == kNumMaxTokensPerRank * kNumRanks) + num_recv_tokens = + psum_num_recv_tokens_per_scaleup_rank[kNumScaleupRanks - 1]; + + // Current rank indices should be maintained + int current_rank_idx = -1, stored_psum_num_recv_tokens; + int current_rank_start = 0, current_rank_end = 0; +#pragma unroll + for (int i = global_warp_idx; i < num_recv_tokens; + i += kNumWarps * kNumSMs) { + // Calculate token index in the buffer + while (i >= current_rank_end) { + current_rank_idx += 1; + EP_DEVICE_ASSERT(current_rank_idx < kNumScaleupRanks); + const auto stored_lane_idx = current_rank_idx % 32; + if (stored_lane_idx == 0 and + current_rank_idx + lane_idx < kNumScaleupRanks) + stored_psum_num_recv_tokens = + psum_num_recv_tokens_per_scaleup_rank[current_rank_idx + + lane_idx]; + current_rank_start = current_rank_end; + current_rank_end = + ptx::exchange(stored_psum_num_recv_tokens, stored_lane_idx); + } + const auto buffer_token = + scaleup_buffer.get_rank_buffer(current_rank_idx) + .get_token_buffer(i - current_rank_start); + + // Wait buffer releases + ptx::tma_store_wait(); + __syncwarp(); + + // Issue TMA loads + // Including all stuffs: data, SF, top-k metadata + if (ptx::elect_one_sync()) { + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + buffer_token.get_base_ptr(), mbarrier_ptr, + tma_buffer.get_num_bytes()); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, + tma_buffer.get_num_bytes()); + } + __syncwarp(); + + // Load target expert indices separately to tolerate TMA load latency + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections"); + int dst_expert_idx = -1; + if (lane_idx < kNumTopk) + dst_expert_idx = buffer_token.get_topk_idx_ptr()[lane_idx]; + __syncwarp(); + + // Validate target expert indices and store for non-expand mode + const auto in_range = expert_start_idx <= dst_expert_idx and + dst_expert_idx < expert_end_idx; + const auto master_src_topk_idx = + ptx::get_master_lane_idx(ptx::gather(in_range)); + dst_expert_idx = in_range ? dst_expert_idx - expert_start_idx : -1; + EP_DEVICE_ASSERT(ptx::deduplicate(dst_expert_idx, lane_idx) or + dst_expert_idx == -1); + if (not kDoExpand and lane_idx < kNumTopk) + recv_topk_idx[i * kNumTopk + lane_idx] = + static_cast(dst_expert_idx); + __syncwarp(); + + // Calculate target indices in the tensor + int dst_tensor_idx = -1; + if (not kDoExpand and ptx::elect_one_sync()) { + dst_tensor_idx = i; + } else if (kDoExpand and dst_expert_idx >= 0) { + dst_tensor_idx = + atomicAdd(psum_num_recv_tokens_per_expert + dst_expert_idx, 1); + } + __syncwarp(); + + // Wait for TMA arrival + if (ptx::elect_one_sync()) + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + __syncwarp(); + + // Maintain linked list + if constexpr (kDoCreateLinkedList) { + if (ptx::elect_one_sync()) + channel_linked_list[tma_buffer.get_linked_list_idx_ptr() + [master_src_topk_idx]] = i; + __syncwarp(); + } + + // Issue TMA stores for data + if (kDoExpand ? (dst_tensor_idx >= 0) : ptx::elect_one_sync()) { + ptx::tma_store_1d( + math::advance_ptr(recv_x, static_cast(dst_tensor_idx) * + kNumHiddenBytes), + tma_buffer.get_hidden_ptr(), kNumHiddenBytes); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Store SF + if constexpr (kNumSFPacks > 0) { + constexpr auto kNumFullIters = kNumSFPacks / 32; + const bool do_last_iter = + (kNumSFPacks % 32 != 0) and + (kNumFullIters * 32 + lane_idx < kNumSFPacks); + EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, + "Unaligned SF element type"); + + // Load into registers + const auto smem_src_ptr = tma_buffer.get_sf_ptr(); + sf_pack_t reg_src[kNumFullIters + 1]; +#pragma unroll + for (int k = 0; k < kNumFullIters; ++k) + reg_src[k] = smem_src_ptr[k * 32 + lane_idx]; + if (do_last_iter) + reg_src[kNumFullIters] = + smem_src_ptr[kNumFullIters * 32 + lane_idx]; + + // Prepare strides + const auto recv_sf_token_stride_i64 = + static_cast(recv_sf_token_stride); + const auto recv_sf_hidden_stride_i64 = + static_cast(recv_sf_hidden_stride); + + // Iterate through all valid indices and store into output buffer + auto mask = kDoExpand ? ptx::gather(dst_tensor_idx >= 0) : 1; + while (mask) { + const int valid_lane_idx = __ffs(mask) - 1; + const auto gmem_dst = math::advance_ptr( + recv_sf, + ptx::exchange(dst_tensor_idx, valid_lane_idx) * + (recv_sf_token_stride_i64 * sizeof(sf_pack_t))); +#pragma unroll + for (int k = 0; k < kNumFullIters; ++k) + gmem_dst[(k * 32 + lane_idx) * recv_sf_hidden_stride_i64] = + reg_src[k]; + if (do_last_iter) + gmem_dst[(kNumFullIters * 32 + lane_idx) * + recv_sf_hidden_stride_i64] = + reg_src[kNumFullIters]; + mask ^= 1 << valid_lane_idx; + } + } + + // Store the top-k weights + if (kDoExpand and recv_topk_weights != nullptr and + dst_tensor_idx >= 0) { + recv_topk_weights[dst_tensor_idx] = + tma_buffer.get_topk_weights_ptr()[lane_idx]; + } else if (not kDoExpand and recv_topk_weights != nullptr and + lane_idx < kNumTopk) { + // For backward, weights are optional + recv_topk_weights[i * kNumTopk + lane_idx] = + tma_buffer.get_topk_weights_ptr()[lane_idx]; + } + __syncwarp(); + + // Write source token index + // And: + // - Non-hybrid mode: the source scaleup peer rank index and master + // top-k lane index + // - Hybrid mode: the slot index and master top-k lane index + constexpr int kMetadataStride = 2 + kNumTopk; + if (ptx::elect_one_sync()) { + recv_src_metadata[i * kMetadataStride + 0] = + *tma_buffer.get_src_token_global_idx_ptr(); + if constexpr (kNumScaleoutRanks == 1) { + recv_src_metadata[i * kMetadataStride + 1] = + current_rank_idx * kNumTopk + master_src_topk_idx; + } else { + recv_src_metadata[i * kMetadataStride + 1] = + (i - current_rank_start) * kNumTopk + master_src_topk_idx; + } + } + __syncwarp(); + + // Write reduction source indices + if (kDoExpand and lane_idx < kNumTopk) + recv_src_metadata[i * kMetadataStride + 2 + lane_idx] = + dst_tensor_idx; + __syncwarp(); + } + + // Maintain linked list's ending + // Or you can understand it as writing the tail at once + if constexpr (kDoCreateLinkedList) { + constexpr int kNumScaleupRanksPerLane = + math::constexpr_ceil_div(kNumScaleupRanks, 32); + const auto workspace_layout = layout::WorkspaceLayout( + workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + for (int i = global_warp_idx; i < kNumChannels; + i += kNumSMs * kNumWarps) { +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + if (const auto k = j * 32 + lane_idx; + j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks) { + channel_linked_list + [*workspace_layout.get_channel_scaleup_tail_ptr(i, k)] = + -1; + + // Clean for combine usages + *workspace_layout.get_channel_scaleup_tail_ptr(i, k) = 0; + } + } + __syncwarp(); + } + } +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh new file mode 100644 index 0000000000..323f81d5fb --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_deterministic_prologue.cuh @@ -0,0 +1,171 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include + +namespace mooncake::elastic { + +// Slot preassignment runs in the active scale-up domain. Hybrid scale-out +// forwarding is handled by the hybrid dispatch kernel. +template +__global__ void __launch_bounds__(kNumThreads, 1) + dispatch_deterministic_prologue_impl(topk_idx_t* topk_idx, + int* rank_count_buffer, + int* dst_buffer_slot_idx, + const int num_tokens, + const int scaleup_rank_idx) { + constexpr int kNumExpertsPerRank = kNumExperts / kNumScaleupRanks; + EP_STATIC_ASSERT(kNumExperts % kNumScaleupRanks == 0, + "Invalid number of experts or ranks"); + + // Utils + const auto sm_idx = static_cast(blockIdx.x), + thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto global_warp_idx = sm_idx * kNumWarps + warp_idx; + + // Token region the current warp is responsible for + const auto num_tokens_per_warp = + math::ceil_div(num_tokens, kNumSMs * kNumWarps); + const auto start_token_idx = global_warp_idx * num_tokens_per_warp; + const auto end_token_idx = + min(start_token_idx + num_tokens_per_warp, num_tokens); + + // Group configs + // NOTES: Group refers to the tokens that each warp handles concurrently + constexpr int kNumTokensPerGroup = 32 / kNumTopk; + const auto token_idx_offset = lane_idx / kNumTopk; + const unsigned token_mask = ((1u << kNumTopk) - 1) + << (token_idx_offset * kNumTopk); + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k"); + + // Shared memory for reduction + // NOTES: Each warp owns separate shared memory region for separate sum. + extern __shared__ int8_t smem[]; + const auto rank_count_global_psum = math::advance_ptr(smem, 0); + const auto rank_count_warp_sum = math::advance_ptr( + rank_count_global_psum, + (kNumScaleupRanks + warp_idx * kNumScaleupRanks) * sizeof(int)); + const auto rank_count_warp_psum = math::advance_ptr( + rank_count_warp_sum, kNumWarps * kNumScaleupRanks * sizeof(int)); + + // Initialize to zero before reduce + for (int i = thread_idx; i < kNumScaleupRanks * (1 + 2 * kNumWarps); + i += kNumThreads) + reinterpret_cast(smem)[i] = 0; + __syncthreads(); + + // Util functions + const auto map_expert_to_rank_idx = [&](const int& expert_idx) { + return expert_idx >= 0 ? expert_idx / kNumExpertsPerRank : -1; + }; + const auto is_unique = [&](const int& rank_idx) { + return ((ptx::match(rank_idx) & token_mask) >> lane_idx) == 1; + }; + const auto count_ones_before = [&](const unsigned& mask, + const int& bit_idx) { + return __popc(mask & ((1u << bit_idx) - 1)); + }; + const auto get_other_rank_count_warp_sum = [&](const int& other_warp_idx) { + // NOTES: pass negative num_bytes to advance pointer + return math::advance_ptr( + rank_count_warp_sum, + (other_warp_idx - warp_idx) * kNumScaleupRanks * sizeof(int)); + }; + + // Each warp scan the tokens separately + for (int i = start_token_idx; i < end_token_idx; i += kNumTokensPerGroup) { + const auto token_idx = i + token_idx_offset; + const auto is_active_thread = + lane_idx < kNumTopk * kNumTokensPerGroup and + token_idx < end_token_idx; + const int expert_idx = + is_active_thread + ? static_cast(__ldg(topk_idx + i * kNumTopk + lane_idx)) + : -1; + const auto rank_idx = map_expert_to_rank_idx(expert_idx); + + // Avoid duplicate messages to a single rank + const auto deduped_rank_idx = is_unique(rank_idx) ? rank_idx : -1; + const auto rank_idx_mask = ptx::match(deduped_rank_idx); + + // Let the one with the largest lane index send the count + if ((rank_idx_mask >> lane_idx) == 1 and deduped_rank_idx >= 0) + rank_count_warp_sum[deduped_rank_idx] += __popc(rank_idx_mask); + } + __syncthreads(); + + // Get block sum and store to global + for (int rank_idx = thread_idx; rank_idx < kNumScaleupRanks; + rank_idx += kNumThreads) { + int rank_count_block_sum = 0; + for (int i = 0; i < kNumWarps; i++) + rank_count_block_sum += get_other_rank_count_warp_sum(i)[rank_idx]; + rank_count_buffer[sm_idx * kNumScaleupRanks + rank_idx] = + rank_count_block_sum; + } + cooperative_groups::this_grid().sync(); + + // Get the prefix sum before the current SM + for (int rank_idx = lane_idx; rank_idx < kNumScaleupRanks; rank_idx += 32) { + int rank_count = 0; + for (int i = warp_idx; i < sm_idx; i += kNumWarps) + rank_count += rank_count_buffer[i * kNumScaleupRanks + rank_idx]; + atomicAdd_block(rank_count_global_psum + rank_idx, rank_count); + } + __syncthreads(); + + // Get each warp's prefix sum + for (int rank_idx = lane_idx; rank_idx < kNumScaleupRanks; rank_idx += 32) { + int rank_count = rank_count_global_psum[rank_idx]; + for (int i = 0; i < warp_idx; i++) + rank_count += get_other_rank_count_warp_sum(i)[rank_idx]; + rank_count_warp_psum[rank_idx] = rank_count; + } + __syncwarp(); + + // Each warp scan the tokens separately + for (int i = start_token_idx; i < end_token_idx; i += kNumTokensPerGroup) { + const auto token_idx = i + token_idx_offset; + const auto is_active_thread = + lane_idx < kNumTopk * kNumTokensPerGroup and + token_idx < end_token_idx; + const auto expert_idx = + is_active_thread + ? static_cast(__ldg(topk_idx + i * kNumTopk + lane_idx)) + : -1; + const auto rank_idx = map_expert_to_rank_idx(expert_idx); + + // Avoid duplicate messages to a single rank + const auto deduped_rank_idx = is_unique(rank_idx) ? rank_idx : -1; + const auto rank_idx_mask = ptx::match(deduped_rank_idx); + + // Store to target buffer + const auto stored_dst_slot_idx = + deduped_rank_idx >= 0 + ? rank_count_warp_psum[deduped_rank_idx] + + count_ones_before(rank_idx_mask, lane_idx) + : -1; + const auto value = + stored_dst_slot_idx >= 0 + ? scaleup_rank_idx * kNumMaxTokensPerRank + stored_dst_slot_idx + : -1; + if (is_active_thread) + dst_buffer_slot_idx[i * kNumTopk + lane_idx] = value; + + // Let the one with the largest lane index send the count + if ((rank_idx_mask >> lane_idx) == 1 and deduped_rank_idx >= 0) + rank_count_warp_psum[deduped_rank_idx] += __popc(rank_idx_mask); + __syncwarp(); + } +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh new file mode 100644 index 0000000000..761392dde4 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_dispatch_official.cuh @@ -0,0 +1,512 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace mooncake::elastic { + +template > +__global__ void __launch_bounds__(kNumThreads, 1) + dispatch_impl(void* x, sf_pack_t* sf, topk_idx_t* topk_idx, + float* topk_weights, topk_idx_t* copied_topk_idx, + int* cumulative_local_expert_recv_stats, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, + int* dst_buffer_slot_idx, const int num_tokens, + const int sf_token_stride, const int sf_hidden_stride, + const device::CommCtx comm_ctx, void* buffer, void* workspace, + void* mapped_host_workspace, const int rank_idx) { + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; + EP_STATIC_ASSERT(kNumExperts % kNumRanks == 0, + "Invalid number of experts or ranks"); + EP_STATIC_ASSERT(kNumNotifyWarps % 4 == 0, "Invalid warpgroup size"); + + // Utils + const auto sm_idx = static_cast(blockIdx.x), + thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + + // Workspaces + const auto workspace_layout = + layout::WorkspaceLayout(workspace, 1, kNumRanks, kNumExperts); + const auto host_workspace_layout = layout::WorkspaceLayout( + mapped_host_workspace, 1, kNumRanks, kNumExperts); + + // The kernel uses a fixed space of dynamic shared memory (no static shared + // memory) + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + constexpr int kNumSmemBytesForNotify = + kNumNotifyThreads > 0 ? math::constexpr_align(kNumRanks + kNumExperts, + kNumNotifyThreads) * + sizeof(int) + : 0; + EP_STATIC_ASSERT(kNumSmemBytesForNotify % ptx::kNumTMAAlignBytes == 0, + "Invalid TMA alignment"); + + // Named barrier indices + constexpr int kNotifyBarrierIndex = 1; + + // Gin handle + // We treat each warp as a "channel" + const auto [qp_idx, sharing_mode] = + comm::get_qp_mode 0)>( + sm_idx, warp_idx - kNumNotifyWarps, warp_idx < kNumNotifyWarps); + const auto gin = transport::MooncakeGin(comm_ctx, qp_idx, sharing_mode, + kNumQPs, 0, 0, 0, kNumRanks); + + // Barrier without TMA store flush, without prologue grid sync + comm::gpu_barrier(gin, workspace_layout, 0, rank_idx, sm_idx, + thread_idx); + + // Different warp roles + if (warp_idx < kNumNotifyWarps) { + // Assign shared memory + constexpr int kNumAlignedElems = kNumSmemBytesForNotify / sizeof(int); + const auto rank_expert_count = math::advance_ptr(smem, 0); + + // Clean initial counts + // NOTES: if you want to change the order of different warp roles, + // please take care of the `thread_idx` + int *rank_count = rank_expert_count, + *expert_count = rank_expert_count + kNumRanks; +#pragma unroll + for (int i = 0; i < kNumAlignedElems / kNumNotifyThreads; ++i) + rank_expert_count[i * kNumNotifyThreads + thread_idx] = 0; + ptx::named_barrier(kNotifyBarrierIndex); + + // Atomic add on shared memory + EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes"); + const auto global_warp_idx = warp_idx * kNumSMs + sm_idx; + for (int i = global_warp_idx; i < num_tokens; + i += kNumNotifyWarps * kNumSMs) { + // Expert choice can not be redundant + // NOTES: no assertions here as they are expensive + const auto dst_expert_idx = + lane_idx < kNumTopk ? static_cast(__ldg( + topk_idx + i * kNumTopk + lane_idx)) + : -1; + if (dst_expert_idx >= 0) + atomicAdd_block(expert_count + dst_expert_idx, 1); + + // Rank choice should do deduplication here + const auto dst_rank_idx = + dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1; + if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0) + atomicAdd_block(rank_count + dst_rank_idx, 1); + } + ptx::named_barrier(kNotifyBarrierIndex); + +// Do full-grid reduction +#pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + const int64_t counter = (1ll << 32ll) | rank_expert_count[i]; + ptx::red_add( + workspace_layout.get_notify_reduction_workspace_ptr() + i, + counter); + } + + // Do the remaining work by SM 0 + if (sm_idx == 0) { +// Reduce all SM's count +// Wait all SMs' arrival +#pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + comm::timeout_while< + kNumTimeoutCycles>(true, [=](const bool& is_last_check) { + const auto status = ptx::ld_volatile( + workspace_layout.get_notify_reduction_workspace_ptr() + + i); + if ((status >> 32) == kNumSMs) { + // Write into shared memory + // Write into send buffer if with RDMA + const auto encoded = math::encode_decode_positive( + static_cast(status & 0xffffffffll)); + rank_expert_count[i] = encoded; + if constexpr (not kIsScaleupNVLink) + workspace_layout + .get_scaleup_rank_expert_count_ptr()[i] = + encoded; + + // Clean for the next usage + workspace_layout + .get_notify_reduction_workspace_ptr()[i] = 0; + return true; + } + + if (is_last_check) { + printf( + "DeepEP notify (GPU reduction) timeout, rank: " + "%d/%d, " + "thread: %d, status: %d | %d, expected: %d\n", + rank_idx, kNumRanks, thread_idx, + static_cast(status >> 32), + static_cast(status & 0xffffffff), kNumSMs); + } + return false; + }); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // TODO: for further optimization, we can fuse rank and expert + // counters Issue scaleup rank count writes to peers + for (int i = thread_idx; i < kNumRanks; i += kNumNotifyThreads) { + // Rank counters + const auto dst_rank_counter = + workspace_layout.get_scaleup_rank_count_ptr() + + rank_idx; + gin.put_value(dst_rank_counter, + static_cast(rank_count[i]), i, + 0); + } + __syncwarp(); + + // Issue scaleup expert count writes to peers + if constexpr (kIsScaleupNVLink) { + // NVLink per-element copy + // We don't use TMA as the dtype of shared memory and global is + // different + for (int i = thread_idx; i < kNumExperts; + i += kNumNotifyThreads) { + const auto idx = kNumExpertsPerRank * rank_idx + + (i % kNumExpertsPerRank); + gin.put_value( + workspace_layout.get_scaleup_expert_count_ptr() + + idx, + static_cast(expert_count[i]), + i / kNumExpertsPerRank); + } + } else { + // RDMA bulk copy + for (int i = thread_idx; i < kNumRanks; + i += kNumNotifyThreads) { + const auto src_ptr = + workspace_layout.get_scaleup_expert_count_ptr() + + kNumExpertsPerRank * i; + const auto dst_ptr = + workspace_layout.get_scaleup_expert_count_ptr() + + kNumExpertsPerRank * rank_idx; + gin.put(dst_ptr, src_ptr, + kNumExpertsPerRank * sizeof(int64_t), i); + } + } + + // This is necessary, as the waited results will rewrite the shared + // memory + ptx::named_barrier(kNotifyBarrierIndex); + + // Wait for rank and expert count + const auto start_clock = clock64(); + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + comm::timeout_while( + [=](const bool& is_last_check) { + // NOTES: the global memory type has 64 bits + const auto count = static_cast< + int>(ptx::ld_volatile( + workspace_layout + .get_scaleup_rank_expert_count_ptr() + + i)); + const auto decoded = + math::encode_decode_positive(count); + if (math::is_decoded_positive_ready(decoded)) { + workspace_layout + .get_scaleup_rank_expert_count_ptr()[i] = + 0; + rank_expert_count[i] = decoded; + return true; + } + + if (is_last_check) + printf( + "DeepEP notify timeout, rank: %d, thread: %d, " + "count: %d\n", + rank_idx, i, decoded); + return false; + }, + start_clock); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Reduce expert count and add stats + for (int i = thread_idx; i < kNumExpertsPerRank; + i += kNumNotifyThreads) { + int sum = 0; +#pragma unroll + for (int j = 0; j < kNumRanks; ++j) + sum += expert_count[j * kNumExpertsPerRank + i]; + expert_count[i] = math::align(sum, kExpertAlignment); + + // Update statistics counters + if (cumulative_local_expert_recv_stats != nullptr) + atomicAdd(cumulative_local_expert_recv_stats + i, sum); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Write host workspace + if constexpr (kDoCPUSync) { + for (int i = thread_idx; i < kNumRanks + kNumExpertsPerRank; + i += kNumNotifyThreads) { + host_workspace_layout + .get_scaleup_rank_expert_count_ptr()[i] = + math::encode_decode_positive(rank_expert_count[i]); + } + __syncwarp(); + } + + // Do prefix sum by the warps + // NOTES: we may have fast implementation with `cub::BlockScan`, but + // it is too heavy to use + const auto do_psum = [=](const int* count, int* out, const int n, + const int is_exclusive) { + int psum = 0; +#pragma unroll + for (int i = 0; i < math::ceil_div(n + is_exclusive, 32); ++i) { + const auto idx = i * 32 + lane_idx; + const auto mem_idx = idx - is_exclusive; + const auto value = + (0 <= mem_idx and mem_idx < n) ? count[mem_idx] : 0; + const auto sum = + psum + ptx::warp_inclusive_sum(value, lane_idx); + + // Store into global memory + if (idx < n + is_exclusive) out[idx] = sum; + + // Update `psum` by using the last lane's value + psum = ptx::exchange(sum, 31); + } + }; + if (warp_idx == 0) { + // Inclusive prefix sum + do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, + kNumRanks, 0); + } else if (warp_idx == 1) { + // Exclusive prefix sum for later expanding + do_psum(expert_count, psum_num_recv_tokens_per_expert, + kNumExpertsPerRank, 1); + } + } + } else { + const int dispatch_warp_idx = warp_idx - kNumNotifyWarps; + + // Buffer layouts + const auto token_layout = layout::TokenLayout( + kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true); + const auto tma_buffer = + layout::BufferLayout( + token_layout, kNumDispatchWarps, 1, + math::advance_ptr(smem, kNumSmemBytesForNotify)) + .get_rank_buffer(dispatch_warp_idx) + .get_token_buffer(0); + auto recv_buffer = layout::BufferLayout( + token_layout, kNumRanks, kNumMaxTokensPerRank, buffer); + auto send_buffer = + layout::BufferLayout(token_layout, 1, kNumMaxTokensPerRank, + recv_buffer.get_buffer_end_ptr()); + recv_buffer = recv_buffer.get_rank_buffer(rank_idx); + + // Init TMA + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) + ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Iterate all tokens + const auto token_start = dispatch_warp_idx * kNumSMs + sm_idx; + const auto token_stride = kNumDispatchWarps * kNumSMs; + for (int token_idx = token_start; token_idx < num_tokens; + token_idx += token_stride) { + const auto token_i64_idx = static_cast(token_idx); + + // Wait TMA store arrivals + ptx::tma_store_wait(); + __syncwarp(); + + // Issue data TMA + ptx::tma_load_1d_warp( + tma_buffer.get_hidden_ptr(), + math::advance_ptr(x, token_i64_idx * kNumHiddenBytes), + mbarrier_ptr, kNumHiddenBytes, lane_idx); + __syncwarp(); + + // Issue SF TMA or cp.async + if constexpr (kNumSFPacks > 0) { + EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, + "Unaligned SF element type"); + const auto gmem_src_ptr = math::advance_ptr( + sf, token_i64_idx * sf_token_stride * sizeof(sf_pack_t)); + const auto smem_dst_ptr = tma_buffer.get_sf_ptr(); + + constexpr auto kNumFullIters = kNumSFPacks / 32; +#pragma unroll + for (int k = 0; k < kNumFullIters; ++k) { + ptx::cp_async_ca( + gmem_src_ptr + (k * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + k * 32 + lane_idx); + } + if (kNumFullIters * 32 + lane_idx < kNumSFPacks) { + ptx::cp_async_ca( + gmem_src_ptr + + (kNumFullIters * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + kNumFullIters * 32 + lane_idx); + } + ptx::cp_async_mbarrier_arrive(mbarrier_ptr); + __syncwarp(); + } + + // Load top-k indices and weights + EP_STATIC_ASSERT(kNumTopk <= 32, + "Insufficient lanes for loading top-k indices"); + int stored_dst_rank_idx = -1; + if (lane_idx < kNumTopk) { + const auto uncasted_dst_expert_idx = + __ldg(topk_idx + token_idx * kNumTopk + lane_idx); + const auto dst_expert_idx = + static_cast(uncasted_dst_expert_idx); + stored_dst_rank_idx = dst_expert_idx >= 0 + ? dst_expert_idx / kNumExpertsPerRank + : -1; + tma_buffer.get_topk_idx_ptr()[lane_idx] = dst_expert_idx; + if (topk_weights != nullptr) + tma_buffer.get_topk_weights_ptr()[lane_idx] = + __ldg(topk_weights + token_idx * kNumTopk + lane_idx); + if (copied_topk_idx != nullptr) + copied_topk_idx[token_idx * kNumTopk + lane_idx] = + uncasted_dst_expert_idx; + } + __syncwarp(); + + // Add source metadata (rank index and token index) + // Please ensure no TMA buffer shared memory writes after this part + if (ptx::elect_one_sync()) + *tma_buffer.get_src_token_global_idx_ptr() = + rank_idx * kNumMaxTokensPerRank + token_idx; + ptx::tma_store_fence(); + __syncwarp(); + + // Deduplicate ranks and assign slots + int stored_dst_slot_idx = -1; + if constexpr (kReuseSlotIndices) { + if (lane_idx < kNumTopk) + stored_dst_slot_idx = __ldg( + dst_buffer_slot_idx + token_idx * kNumTopk + lane_idx); + stored_dst_slot_idx = stored_dst_slot_idx >= 0 + ? (stored_dst_slot_idx - + rank_idx * kNumMaxTokensPerRank) + : -1; + } else { + if (ptx::deduplicate(stored_dst_rank_idx, lane_idx) and + stored_dst_rank_idx >= 0) + stored_dst_slot_idx = atomicAdd( + workspace_layout.get_scaleup_atomic_sender_counter() + + stored_dst_rank_idx, + 1); + if (lane_idx < kNumTopk) { + const auto value = stored_dst_slot_idx >= 0 + ? rank_idx * kNumMaxTokensPerRank + + stored_dst_slot_idx + : -1; + dst_buffer_slot_idx[token_idx * kNumTopk + lane_idx] = + value; + } + } + __syncwarp(); + + // Wait TMA load arrival + // NOTES: this arrive must be after the + // `ptx::cp_async_mbarrier_arrive` + if (ptx::elect_one_sync()) { + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + } + __syncwarp(); + + // TMA store to send buffer + auto send_buffer_ptr = + send_buffer.get_token_buffer(token_idx).get_base_ptr(); + if constexpr (not kIsScaleupNVLink) { + if (ptx::elect_one_sync()) + ptx::tma_store_1d(send_buffer_ptr, + tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + ptx::tma_store_commit(); + __syncwarp(); + } + + // Issue TMA NVLink stores + EP_STATIC_ASSERT(kNumTopk <= 32, "Invalid top-k selection"); + const auto dst_ptr = + stored_dst_slot_idx >= 0 + ? gin.get_sym_ptr( + recv_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + stored_dst_rank_idx) + : nullptr; + if (dst_ptr != nullptr) + ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + ptx::tma_store_commit(); + __syncwarp(); + + // Issue RDMA put + if constexpr (not kIsScaleupNVLink) { + // Wait the send buffer store to arrive + ptx::tma_store_wait<1>(); + __syncwarp(); + + // NOTES: we should skip the NVLink accessible ranks + if (stored_dst_slot_idx >= 0 and dst_ptr == nullptr) { + gin.put( + recv_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + send_buffer_ptr, tma_buffer.get_num_bytes(), + stored_dst_rank_idx); + } + __syncwarp(); + } + } + } + + // Barrier to ensure data arrival + comm::gpu_barrier(gin, workspace_layout, 0, rank_idx, sm_idx, + thread_idx); + + // Trigger the copy epilogue kernel +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + // Clean atomic counters + EP_STATIC_ASSERT(kNumRanks <= kNumThreads, "Insufficient threads"); + if (not kReuseSlotIndices and sm_idx == 0 and thread_idx < kNumRanks) + workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0; +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_exception.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_exception.cuh new file mode 100644 index 0000000000..26e1c8b984 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_exception.cuh @@ -0,0 +1,81 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include + +#ifndef EP_STATIC_ASSERT +#define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason) +#endif + +#define EPExceptionWithLineInfo(name, message) \ + EPException(name, __FILE__, __LINE__, message) + +#ifndef EP_HOST_ASSERT +#define EP_HOST_ASSERT(cond) \ + do { \ + if (not(cond)) { \ + throw EPException("Assertion", __FILE__, __LINE__, #cond); \ + } \ + } while (0) +#endif + +#ifndef EP_HOST_UNREACHABLE +#define EP_HOST_UNREACHABLE(reason) \ + (throw EPException("Assertion", __FILE__, __LINE__, reason)) +#endif + +#ifndef EP_DEVICE_ASSERT +#define EP_DEVICE_ASSERT(cond) \ + do { \ + if (not(cond)) { \ + printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, \ + __LINE__, #cond); \ + asm("trap;"); \ + } \ + } while (0) +#endif + +#ifndef EP_UNIFIED_ASSERT +#ifdef __CUDA_ARCH__ +#define EP_UNIFIED_ASSERT(cond) EP_DEVICE_ASSERT(cond) +#else +#define EP_UNIFIED_ASSERT(cond) EP_HOST_ASSERT(cond) +#endif +#endif + +#ifndef CUDA_RUNTIME_CHECK +#define CUDA_RUNTIME_CHECK(cmd) \ + do { \ + const auto e = (cmd); \ + if (e != cudaSuccess) { \ + std::stringstream ss; \ + ss << static_cast(e) << " (" << cudaGetErrorName(e) << ", " \ + << cudaGetErrorString(e) << ")"; \ + throw EPException("CUDA runtime", __FILE__, __LINE__, ss.str()); \ + } \ + } while (0) +#endif + +#ifndef CUDA_DRIVER_CHECK +#define CUDA_DRIVER_CHECK(cmd) \ + do { \ + const auto e = (cmd); \ + if (e != CUDA_SUCCESS) { \ + std::stringstream ss; \ + const char *name, *info; \ + lazy_cuGetErrorName(e, &name), lazy_cuGetErrorString(e, &info); \ + ss << static_cast(e) << " (" << name << ", " << info << ")"; \ + throw EPException("CUDA driver", __FILE__, __LINE__, ss.str()); \ + } \ + } while (0) +#endif + +#ifndef NCCL_CHECK +#define NCCL_CHECK(cmd) \ + do { \ + (void)(cmd); \ + } while (0) +#endif diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh new file mode 100644 index 0000000000..1689e8f01f --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_combine_official.cuh @@ -0,0 +1,787 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace mooncake::elastic { + +template < + bool kUseExpandedLayout, bool kAllowMultipleReduction, int kNumSMs, + int kNumScaleupWarps, int kNumForwardWarps, int kNumScaleoutRanks, + int kNumScaleupRanks, int kHidden, int kNumMaxTokensPerRank, + int kNumExperts, int kNumTopk, int kNumQPs, int64_t kNumTimeoutCycles, + int kNumScaleupRanksPerLane = math::constexpr_ceil_div(kNumScaleupRanks, + 32), + int kNumScaleupUpdateInterval = 3, int kNumChannelsPerSM = kNumForwardWarps, + int kNumChannels = kNumChannelsPerSM * kNumSMs, + int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, + kNumChannels), + int kNumRanks = kNumScaleoutRanks * kNumScaleupRanks, + int kNumWarps = kNumScaleupWarps + kNumForwardWarps, + int kNumThreads = kNumWarps * 32, + int kNumHiddenBytes = kHidden * sizeof(nv_bfloat16), + bool kUseScaleoutRankLayout = + use_rank_layout(), + bool kUseScaleupRankLayout = + use_rank_layout(), + int kNumTokensInScaleoutLayout = get_num_tokens_in_layout< + kAllowMultipleReduction, kNumScaleoutRanks, kNumTopk>(), + int kNumTokensInScaleupLayout = get_num_tokens_in_layout< + kAllowMultipleReduction, kNumScaleupRanks, kNumTopk>()> +__global__ void __launch_bounds__(kNumThreads, 1) + hybrid_combine_impl(nv_bfloat16* x, float* topk_weights, int* src_metadata, + int* psum_num_recv_tokens_per_scaleup_rank, + int* token_metadata_at_forward, + int* channel_linked_list, + const device::CommCtx comm_ctx, void* buffer, + void* workspace, const int scaleout_rank_idx, + const int scaleup_rank_idx, int num_reduced_tokens) { + // Utils + const auto sm_idx = static_cast(blockIdx.x); + const auto thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(); + const auto lane_idx = ptx::get_lane_idx(); + constexpr bool kDoExpandedSend = + not kAllowMultipleReduction and kUseExpandedLayout; + + // Combine vector type selection + using combine_vec_t = typename CombineVecTraits::vec_t; + constexpr int kHiddenVec = kNumHiddenBytes / sizeof(combine_vec_t); + + // Workspaces + const auto workspace_layout = layout::WorkspaceLayout( + workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + + // We should assign the real number of received tokens if without CPU sync + if (num_reduced_tokens == kNumMaxTokensPerRank * kNumRanks) + num_reduced_tokens = + __ldg(psum_num_recv_tokens_per_scaleup_rank + kNumScaleupRanks - 1); + + // Token layouts + const auto token_layout = + layout::TokenLayout(kNumHiddenBytes, 0, kNumTopk, false); + + // TMA buffers + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + const auto tma_buffer = + layout::BufferLayout(token_layout, kNumWarps, 1, smem) + .get_rank_buffer(warp_idx) + .get_token_buffer(0); + + // All the buffer layouts + auto scaleup_buffer = layout::BufferLayout( + token_layout, kNumTokensInScaleupLayout, + kNumScaleoutRanks * kNumMaxTokensPerRank, buffer); + auto scaleout_recv_buffer = layout::BufferLayout( + token_layout, kNumTokensInScaleoutLayout, kNumMaxTokensPerRank, + scaleup_buffer.get_buffer_end_ptr()); + auto scaleout_send_buffer = layout::BufferLayout( + token_layout, kAllowMultipleReduction ? 1 : kNumTopk, + kNumChannels * (kNumScaleoutRanks * kNumMaxTokensPerChannel), + scaleout_recv_buffer.get_buffer_end_ptr()); + + // Init TMA for scale-up and forward warps + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (ptx::elect_one_sync()) ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Mooncake Gin handle + // Each warp is a channel + const auto [qp_idx, sharing_mode] = + comm::get_qp_mode( + sm_idx, warp_idx % kNumChannelsPerSM); + const auto gin = transport::MooncakeGin( + comm_ctx, qp_idx, sharing_mode, kNumQPs, scaleout_rank_idx, + scaleup_rank_idx, kNumScaleupRanks, kNumRanks); + + // Global parallel barriers for scale-out subteam and scale-up subteam + // NOTES: this barrier needs a grid sync, as there are channel scale-up tail + // cleaning before + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, + thread_idx); + + // Adjust register count at certain cases + // TODO: support more cases, or try to make channel count more aligned + // DeepEP's register redistribution uses setmaxnreg, which is not accepted + // by ptxas for the SM90 target used by current Mooncake NV validation. + // Keep the official role split but disable this SM100-only optimization. + constexpr bool kAdjustRegisters = false; + constexpr int kNumRegistersForScaleupWarps = 40; + constexpr int kNumRegistersForForwardWarps = + 256 - kNumRegistersForScaleupWarps; + + // Different warp roles + if (warp_idx < kNumScaleupWarps) { + const auto channel_idx = sm_idx * kNumChannelsPerSM + warp_idx; + + // Adjust registers + if constexpr (kAdjustRegisters) + ptx::warpgroup_reg_dealloc(); + + // Shift into the right buffer if using rank layout + if constexpr (kUseScaleupRankLayout) + scaleup_buffer = scaleup_buffer.get_rank_buffer(scaleup_rank_idx); + + // Expanding mode must not be backward + if constexpr (kUseExpandedLayout) + EP_DEVICE_ASSERT(topk_weights == nullptr); + + // Tail issuer + // `st.release.sys` is pretty slow, so do it by an interval + int update_counter = 0; + int stored_num_tokens_sent[kNumScaleupRanksPerLane] = {}; + int stored_old_num_tokens_sent[kNumScaleupRanksPerLane] = {}; + const auto tail_ptr = workspace_layout.get_channel_scaleup_tail_ptr( + channel_idx, scaleup_rank_idx); + const auto update_tails = [&](const bool& finish = false) { + ++update_counter; + if (finish or update_counter == kNumScaleupUpdateInterval) { + // Wait all TMA stores to finish + ptx::tma_store_wait(); + __syncwarp(); + +// Issue +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) { + if (const auto j = i * 32 + lane_idx; + i < (kNumScaleupRanksPerLane - 1) or + j < kNumScaleupRanks) { + // NOTES: save some traffic with + // `stored_old_num_tokens_sent` Also, we cannot rewrite + // a finished slot, if the peer is going to clean it + if (stored_num_tokens_sent[i] != + stored_old_num_tokens_sent[i]) + ptx::st_release_sys( + gin.get_sym_ptr( + tail_ptr, j), + stored_num_tokens_sent[i]); + stored_old_num_tokens_sent[i] = + stored_num_tokens_sent[i]; + } + } + update_counter = 0; + } + __syncwarp(); + }; + + // Shape of `channel_linked_list`: `[kNumChannels, + // kNumMaxTokensPerChannel + 1, kNumScaleupRanks]` Iterate until all + // scale-up peers finish + int dst_scaleup_rank_idx = channel_idx; + int stored_ll_idx[kNumScaleupRanksPerLane] = {}, + stored_token_idx[kNumScaleupRanksPerLane] = {}; +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) + stored_token_idx[i] = -1; + while (true) { +// Load token indices in the list +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) { + const auto j = i * 32 + lane_idx; + stored_token_idx[i] = + i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks + ? __ldg( + channel_linked_list + + channel_idx * + (kNumScaleoutRanks * kNumMaxTokensPerChannel + + 1) * + kNumScaleupRanks + + stored_ll_idx[i] * kNumScaleupRanks + j) + : -1; + } + __syncwarp(); + + // Check whether all ranks are finished + bool exited = true; +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) + exited &= ptx::all(stored_token_idx[i] < 0); + if (exited) break; + + // Process tokens for all ranks together using bitmask to skip + // inactive ranks + EP_STATIC_ASSERT(kNumScaleupRanks <= 64, + "Too many scale-up ranks for 64-bit mask"); + using mask_t = std::conditional_t<(kNumScaleupRanks <= 32), + uint32_t, uint64_t>; + mask_t wip_mask = 0; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + wip_mask |= + static_cast(ptx::gather(stored_token_idx[j] >= 0)) + << (j * 32); + while (wip_mask) { + // Find next active rank after `dst_scaleup_rank_idx` + // (round-robin) + const auto start = + (dst_scaleup_rank_idx + 1) % kNumScaleupRanks; + const auto hi_mask = (wip_mask >> start) << start; + dst_scaleup_rank_idx = + hi_mask ? ptx::ffs(hi_mask) : ptx::ffs(wip_mask); + wip_mask ^= static_cast(1) << dst_scaleup_rank_idx; + + // Exchange token index from the owning lane using static + // partition iteration + int token_idx = -1; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + const auto src_lane_idx = dst_scaleup_rank_idx - j * 32; + token_idx = src_lane_idx == lane_idx ? stored_token_idx[j] + : token_idx; + } + token_idx = ptx::exchange(token_idx, dst_scaleup_rank_idx % 32); + + // Get source metadata and decide the destination buffer + constexpr int kMetadataStride = 2 + kNumTopk; + const auto src_global_token_idx = + __ldg(src_metadata + token_idx * kMetadataStride + 0); + const auto src_token_idx = + src_global_token_idx % kNumMaxTokensPerRank; + const auto src_scaleout_rank_idx = + src_global_token_idx / + (kNumMaxTokensPerRank * kNumScaleupRanks); + auto token_buffer = [&]() { + if constexpr (kUseScaleupRankLayout) { + const auto src_slot_idx = + __ldg(src_metadata + token_idx * kMetadataStride + + 1) / + kNumTopk; + return scaleup_buffer.get_token_buffer(src_slot_idx); + } else { + const auto master_topk_idx = + __ldg(src_metadata + token_idx * kMetadataStride + + 1) % + kNumTopk; + return scaleup_buffer.get_rank_buffer(master_topk_idx) + .get_token_buffer(src_scaleout_rank_idx * + kNumMaxTokensPerRank + + src_token_idx); + } + }(); + token_buffer.set_base_ptr( + gin.get_sym_ptr( + token_buffer.get_base_ptr(), dst_scaleup_rank_idx)); + + // Some checks + EP_STATIC_ASSERT( + kHidden % (32 * sizeof(int4) / sizeof(nv_bfloat16)) == 0, + "Invalid hidden"); + + // Read source indices for expand mode + int stored_topk_slot_idx = -1; + if constexpr (kUseExpandedLayout) { + if (lane_idx < kNumTopk) + stored_topk_slot_idx = + __ldg(src_metadata + token_idx * kMetadataStride + + (2 + lane_idx)); + __syncwarp(); + } + + // 3 cases: + // - no-expand, expand + no-reduce + // - expand + reduce + // - expand + send all + auto reduce_valid_mask = ptx::gather(stored_topk_slot_idx >= 0); + auto no_local_reduce = + not kUseExpandedLayout or (kAllowMultipleReduction and + __popc(reduce_valid_mask) == 1); + if (no_local_reduce) { + int token_idx_in_tensor = token_idx; + if constexpr (kUseExpandedLayout) + token_idx_in_tensor = ptx::exchange( + stored_topk_slot_idx, + ptx::get_master_lane_idx(reduce_valid_mask)); + + // Directly load + if (ptx::elect_one_sync()) { + const auto load_ptr = math::advance_ptr( + x, static_cast(token_idx_in_tensor) * + kNumHiddenBytes); + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), load_ptr, + mbarrier_ptr, kNumHiddenBytes); + } + __syncwarp(); + } else if constexpr (kAllowMultipleReduction) { + // Do local reduction + // Sort valid top-k indices to front + int topk_slot_idx[kNumTopk]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, [=](const int& idx) { + return ptx::exchange(stored_topk_slot_idx, idx); + }); + + // Reduce into shared memory + constexpr int kUnrollFactor = + get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, + static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ + [=](const int& slot_idx) { + return math::advance_ptr( + x, slot_idx * + static_cast(kNumHiddenBytes)); + }, + /* Wait buffer release */ + [=]() { + ptx::tma_store_wait(); + __syncwarp(); + }); + ptx::tma_store_fence(); + __syncwarp(); + } else { +// No local reduction, send all data (expanded send) +#pragma unroll + for (int k = 0; k < kNumTopk; ++k) { + int topk_slot_idx = + ptx::exchange(stored_topk_slot_idx, k); + if (topk_slot_idx < 0) continue; + + if (ptx::elect_one_sync()) { + // Load + const auto load_ptr = math::advance_ptr( + x, static_cast(kDoExpandedSend + ? topk_slot_idx + : token_idx) * + kNumHiddenBytes); + ptx::tma_store_wait(); + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + load_ptr, mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, + phase); + // NOTES: We don't need to care about `topk_weights` + // since we are in expand mode + + // Store + const auto dst_token_buffer = + scaleup_buffer.get_rank_buffer(k) + .get_token_buffer(src_scaleout_rank_idx * + kNumMaxTokensPerRank + + src_token_idx); + ptx::tma_store_1d( + gin.get_sym_ptr( + dst_token_buffer.get_base_ptr(), + dst_scaleup_rank_idx), + tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + } + __syncwarp(); + } + } + + // Write top-k weights + if (not kUseExpandedLayout and topk_weights != nullptr and + lane_idx < kNumTopk) { + const float value = + __ldg(topk_weights + (token_idx * kNumTopk + lane_idx)); + tma_buffer.get_topk_weights_ptr()[lane_idx] = value; + ptx::tma_store_fence(); + } + __syncwarp(); + + // Issue TMA stores into remote scale-up buffer + // NOTES: `kDoExpandedSend` mode has already issued + if (not kDoExpandedSend and ptx::elect_one_sync()) { + // Wait TMA arrival (only for non-reduced cases) + if (no_local_reduce) { + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, + kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + } + + // Issue stores + ptx::tma_store_1d(token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + } +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + stored_num_tokens_sent[j] += + (j * 32 + lane_idx) == dst_scaleup_rank_idx; + __syncwarp(); + } + + // Update the tails together + // NOTES: TMA wait is inside + update_tails(); + +// Move linked list +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) + stored_ll_idx[i] += (stored_token_idx[i] >= 0); + } + + // Update for the unissued ones + update_tails(true); + } else { + const auto forward_warp_idx = warp_idx - kNumScaleupWarps; + const auto channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx; + + // Adjust registers + if constexpr (kAdjustRegisters) + ptx::warpgroup_reg_alloc(); + + // Shift into the right buffer + scaleout_send_buffer = scaleout_send_buffer.get_channel_buffer< + kNumScaleoutRanks * kNumMaxTokensPerChannel>(channel_idx); + + // Shape of `token_metadata_at_forward`: `[kNumChannels, + // kNumScaleoutRanks * kNumMaxTokensPerChannel + 1, + // kNumForwardMetadataDims]` + constexpr int kNumForwardMetadataDims = 2 + kNumTopk * 2; + token_metadata_at_forward += + channel_idx * ((kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * + kNumForwardMetadataDims); + + // Overlap TMA stores and reduction + int last_src_scaleout_rank_idx = -1; + int last_is_token_last_in_chunk = 0; + void* last_recv_token_buffer_ptr = nullptr; + void* last_send_token_buffer_ptr = nullptr; + const auto flush_last_tma_and_issue_rdma = [&]() { + if (last_src_scaleout_rank_idx >= 0 and ptx::elect_one_sync()) { + ptx::tma_store_wait(); + + // Issue only if not local rank + if (last_src_scaleout_rank_idx != scaleout_rank_idx) { + gin.put( + last_recv_token_buffer_ptr, last_send_token_buffer_ptr, + token_layout.get_num_bytes(), + last_src_scaleout_rank_idx, + last_is_token_last_in_chunk ? 0 : 0); + } + } + __syncwarp(); + }; + + // Replay the dispatch + int stored_num_tokens_recv[kNumScaleupRanksPerLane] = {}, + stored_cached_scaleup_tail[kNumScaleupRanksPerLane] = {}; + for (int i = 0;; ++i) { + const auto src_token_global_idx = + __ldg(token_metadata_at_forward + i * kNumForwardMetadataDims); + const auto is_token_last_in_chunk = __ldg( + token_metadata_at_forward + i * kNumForwardMetadataDims + 1); + const auto src_rank_idx = + src_token_global_idx / kNumMaxTokensPerRank; + const auto src_scaleout_rank_idx = src_rank_idx / kNumScaleupRanks; + const auto src_token_idx = + src_token_global_idx % kNumMaxTokensPerRank; + auto stored_src_scaleup_rank_idx = + lane_idx < kNumTopk + ? __ldg(token_metadata_at_forward + + i * kNumForwardMetadataDims + 2 + lane_idx) + : -1; + auto stored_src_slot_idx = lane_idx < kNumTopk + ? __ldg(token_metadata_at_forward + + i * kNumForwardMetadataDims + + 2 + kNumTopk + lane_idx) + : -1; + if (src_token_global_idx < 0) break; + + // Scaleup rank mask + EP_STATIC_ASSERT(kNumScaleupRanks <= 64, "Too many scale-up peers"); + using mask_t = std::conditional_t; + const auto scaleup_mask = + ptx::reduce_or(stored_src_scaleup_rank_idx >= 0 + ? (mask_t(1) << stored_src_scaleup_rank_idx) + : mask_t(0)); + bool stored_is_scaleup_rank_needed[kNumScaleupRanksPerLane]; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + stored_is_scaleup_rank_needed[j] = + (scaleup_mask >> (j * 32 + lane_idx)) & 1; + + // Wait all tails to arrive + comm::timeout_while([&](const bool& + is_last_check) { + bool arrived = true; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + arrived &= not stored_is_scaleup_rank_needed[j] or + stored_num_tokens_recv[j] < + stored_cached_scaleup_tail[j]; + if (ptx::all(arrived)) return true; + +// Reload cached +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + const auto k = j * 32 + lane_idx; + stored_cached_scaleup_tail[j] = + j < (kNumScaleupRanksPerLane - 1) or + k < kNumScaleupRanks + ? ptx::ld_acquire_sys( + workspace_layout.get_channel_scaleup_tail_ptr( + channel_idx, k)) + : -1; + } + + // Timeout + if (is_last_check) { +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + printf( + "DeepEP combine (scale-up wait) timeout, " + "scale-out: %d/%d, scale-up: %d/%d, " + "channel: %d, lane: %d, recv: %d, tail: %d " + "(wait=%d)\n", + scaleout_rank_idx, kNumScaleoutRanks, + scaleup_rank_idx, kNumScaleupRanks, channel_idx, + j * 32 + lane_idx, stored_num_tokens_recv[j], + stored_cached_scaleup_tail[j], + stored_is_scaleup_rank_needed[j]); + } + } + return false; + }); + +// Increase received count +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + stored_num_tokens_recv[j] += + static_cast(stored_is_scaleup_rank_needed[j]); + + if constexpr (not kAllowMultipleReduction) { + // Cases where multiple reduction is disabled. We need to + // forward all data from scaleup peers to scaleout peers + // TODO: Let scale-up warps directly put data into + // `send_buffer`? + const auto src_slot_idx = + src_scaleout_rank_idx * kNumMaxTokensPerRank + + src_token_idx; + auto topk_valid_mask = + kUseExpandedLayout + ? ptx::gather(stored_src_scaleup_rank_idx >= 0) + : ptx::gather( + ptx::deduplicate(stored_src_scaleup_rank_idx, + lane_idx) and + stored_src_scaleup_rank_idx >= + 0); // Deduplicate w.r.t. scaleup rank index + // if expanded mode is disabled + if (ptx::elect_one_sync()) { +#pragma unroll + for (int k = 0; k < kNumTopk; ++k) { + if ((topk_valid_mask & (1u << k)) == 0u) continue; + + // Issue TMA load, and wait + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + scaleup_buffer.get_rank_buffer(k) + .get_token_buffer(src_slot_idx) + .get_base_ptr(), + mbarrier_ptr, + token_layout.get_num_bytes()); + ptx::mbarrier_arrive_and_set_tx( + mbarrier_ptr, token_layout.get_num_bytes()); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + + // Issue TMA store, and wait + const auto recv_buffer_ptr = + scaleout_recv_buffer.get_rank_buffer(k) + .get_token_buffer(src_token_idx) + .get_base_ptr(); + const auto send_buffer_ptr = + src_scaleout_rank_idx == scaleout_rank_idx + ? recv_buffer_ptr + : scaleout_send_buffer.get_rank_buffer(k) + .get_token_buffer(i) + .get_base_ptr(); + ptx::tma_store_1d(send_buffer_ptr, + tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + ptx::tma_store_wait(); + + // Issue IBGDA + topk_valid_mask ^= 1u << k; + if (src_scaleout_rank_idx != scaleout_rank_idx) { + gin.put( + recv_buffer_ptr, send_buffer_ptr, + token_layout.get_num_bytes(), + src_scaleout_rank_idx, + topk_valid_mask == 0 and is_token_last_in_chunk + ? 0 + : 0); + } + } + } + __syncwarp(); + } else { + // NOTES: we must do deduplicate and only add once from one rank + auto reduce_valid_mask = ptx::gather( + ptx::deduplicate(stored_src_scaleup_rank_idx, lane_idx) and + stored_src_scaleup_rank_idx >= 0); + + // Calculate the source buffer index + int stored_src_buffer_idx = 0; + if constexpr (kUseScaleupRankLayout) { + stored_src_buffer_idx = + stored_src_scaleup_rank_idx * + scaleup_buffer.num_max_tokens_per_rank + + stored_src_slot_idx; + } else { + const auto src_slot_idx = + src_scaleout_rank_idx * kNumMaxTokensPerRank + + src_token_idx; + stored_src_buffer_idx = + stored_src_slot_idx == -1 + ? -1 + : lane_idx * + scaleup_buffer.num_max_tokens_per_rank + + src_slot_idx; + } + + // Preprocess top-k indices + int topk_slot_idx[kNumTokensInScaleupLayout]; + compute_topk_slots( + topk_slot_idx, reduce_valid_mask, [=](const int& idx) { + return ptx::exchange(stored_src_buffer_idx, idx); + }); + + // Do reduce + constexpr int kUnrollFactor = + get_max_unroll_factor(); + combine_reduce( + lane_idx, topk_slot_idx, + static_cast(tma_buffer.get_base_ptr()), + /* Get source base */ + [=](const int& slot_idx) { + return static_cast( + scaleup_buffer.get_token_buffer(slot_idx, true) + .get_base_ptr()); + }, + /* Wait buffer release */ + [=]() { flush_last_tma_and_issue_rdma(); }); + + // Merge topk weights + // NOTES: the slot indices must follow the master lane + stored_src_buffer_idx = ptx::exchange( + stored_src_buffer_idx, ptx::get_master_lane_idx(ptx::match( + stored_src_scaleup_rank_idx))); + if (not kUseExpandedLayout and + stored_src_scaleup_rank_idx >= 0) { + tma_buffer.get_topk_weights_ptr()[lane_idx] = + scaleup_buffer + .get_token_buffer(stored_src_buffer_idx, true) + .get_topk_weights_ptr()[lane_idx]; + } + ptx::tma_store_fence(); + __syncwarp(); // Necessary to let the leader lane see the + // writes + + // Assign send and receive buffers + // NOTES: as we only have 1 destination, we will use "send" as + // "recv" for local transfer + int scaleout_recv_buffer_rank_idx; + if constexpr (kUseScaleoutRankLayout) { + scaleout_recv_buffer_rank_idx = scaleout_rank_idx; + } else { + const int src_topk_idx = ptx::get_master_lane_idx( + ptx::gather(stored_src_scaleup_rank_idx >= 0)); + scaleout_recv_buffer_rank_idx = src_topk_idx; + } + const auto recv_token_buffer = + scaleout_recv_buffer + .get_rank_buffer(scaleout_recv_buffer_rank_idx) + .get_token_buffer(src_token_idx); + const auto send_token_buffer = + src_scaleout_rank_idx == scaleout_rank_idx + ? recv_token_buffer + : scaleout_send_buffer.get_token_buffer(i); + + // Write into scale-out send buffer or local rank recv buffer + // bypass + if (ptx::elect_one_sync()) { + ptx::tma_store_1d(send_token_buffer.get_base_ptr(), + tma_buffer.get_base_ptr(), + token_layout.get_num_bytes()); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Record RDMA info to issue later + last_src_scaleout_rank_idx = src_scaleout_rank_idx; + last_is_token_last_in_chunk = is_token_last_in_chunk; + last_recv_token_buffer_ptr = recv_token_buffer.get_base_ptr(); + last_send_token_buffer_ptr = send_token_buffer.get_base_ptr(); + } + } + + // Issue the last RDMA + if constexpr (kAllowMultipleReduction) flush_last_tma_and_issue_rdma(); + +// Clean scaleup tails +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + const auto k = j * 32 + lane_idx; + if (j < (kNumScaleupRanksPerLane - 1) or k < kNumScaleupRanks) + *workspace_layout.get_channel_scaleup_tail_ptr(channel_idx, k) = + 0; + } + __syncwarp(); + + // Update, wait and clean + EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Invalid ranks"); + if (lane_idx < kNumScaleoutRanks) { + // Update remote tails + const auto expected_signal = math::pack2(1, 0); + gin.red_add_rel( + workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, scaleout_rank_idx), + expected_signal, lane_idx); + + // Wait tail arrival + const auto wait_ptr = + workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, lane_idx); + comm::timeout_while([=](const bool& + is_last_check) { + const auto signal = ptx::ld_acquire_sys(wait_ptr); + if (signal == expected_signal) { + // Clean for next usages + *wait_ptr = 0; + return true; + } + + if (is_last_check) { + printf( + "DeepEP combine (scale-out wait all) timeout, " + "scale-out: %d/%d, scale-up: %d/%d, " + "channel: %d, lane: %d, signal: %lld, expected: %lld\n", + scaleout_rank_idx, kNumScaleoutRanks, scaleup_rank_idx, + kNumScaleupRanks, channel_idx, lane_idx, signal, + expected_signal); + } + return false; + }); + } + __syncwarp(); + } + + // No barrier at epilogue +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh new file mode 100644 index 0000000000..9211911d87 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_hybrid_dispatch_official.cuh @@ -0,0 +1,888 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +#include +#include +#include +#include +#include +#include + +namespace mooncake::elastic { + +template < + bool kDoCPUSync, bool kReuseSlotIndices, int kNumSMs, int kNumNotifyWarps, + int kNumScaleoutWarps, int kNumForwardWarps, int kNumScaleoutRanks, + int kNumScaleupRanks, int kNumHiddenBytes, int kNumSFPacks, + int kNumMaxTokensPerRank, int kNumExperts, int kNumTopk, + int kExpertAlignment, int kNumQPs, int64_t kNumTimeoutCycles, + int kNumScaleupRanksPerLane = math::constexpr_ceil_div(kNumScaleupRanks, + 32), + int kNumChannelsPerSM = kNumScaleoutWarps, + int kNumChannels = kNumScaleoutWarps * kNumSMs, + int kNumMaxTokensPerChannel = math::constexpr_ceil_div(kNumMaxTokensPerRank, + kNumChannels), + int kScaleoutUpdateInterval = 3, + int kNumSlotsPerForwardChunk = kScaleoutUpdateInterval, + int kNumRanks = kNumScaleoutRanks * kNumScaleupRanks, + int kNumNotifyThreads = kNumNotifyWarps * 32, + int kNumScaleoutSendThreads = kNumScaleoutWarps * 32, + int kNumForwardThreads = kNumForwardWarps * 32, + int kNumThreads = kNumNotifyThreads + kNumScaleoutSendThreads + + kNumForwardThreads> +__global__ void __launch_bounds__(kNumThreads, 1) + hybrid_dispatch_impl(void* x, sf_pack_t* sf, topk_idx_t* topk_idx, + float* topk_weights, topk_idx_t* copied_topk_idx, + int* cumulative_local_expert_recv_stats, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, + int* dst_buffer_slot_idx, + int* token_metadata_at_forward, const int num_tokens, + const int sf_token_stride, const int sf_hidden_stride, + // TODO(NCCL): so many params, plans to optimize? + const device::CommCtx comm_ctx, void* buffer, + void* workspace, void* mapped_host_workspace, + const int scaleout_rank_idx, + const int scaleup_rank_idx) { + constexpr int kNumExpertsPerRank = kNumExperts / kNumRanks; + constexpr int kNumExpertsPerScaleout = kNumExperts / kNumScaleoutRanks; + EP_STATIC_ASSERT(kNumExperts % kNumScaleupRanks == 0, + "Invalid number of experts or ranks"); + EP_STATIC_ASSERT(kNumNotifyWarps % 4 == 0, "Invalid warpgroup size"); + EP_STATIC_ASSERT(kNumScaleoutWarps == kNumForwardWarps, + "Invalid warp size"); + + // Utils + // NOTES: a warp is a channel (different channels may share QPs) + const auto sm_idx = static_cast(blockIdx.x), + thread_idx = static_cast(threadIdx.x); + const auto warp_idx = ptx::get_warp_idx(), lane_idx = ptx::get_lane_idx(); + const auto rank_idx = + scaleout_rank_idx * kNumScaleupRanks + scaleup_rank_idx; + + // Workspaces + const auto workspace_layout = layout::WorkspaceLayout( + workspace, kNumScaleoutRanks, kNumScaleupRanks, kNumExperts); + const auto host_workspace_layout = + layout::WorkspaceLayout(mapped_host_workspace, kNumScaleoutRanks, + kNumScaleupRanks, kNumExperts); + + // The kernel uses a fixed space of dynamic shared memory (no static shared + // memory) + extern __shared__ __align__(ptx::kNumTMAAlignBytes) int8_t smem[]; + constexpr int kNumSmemBytesForNotify = + kNumNotifyThreads > 0 ? math::constexpr_align(kNumRanks + kNumExperts, + kNumNotifyThreads) * + sizeof(int) + : 0; + EP_STATIC_ASSERT(kNumSmemBytesForNotify % ptx::kNumTMAAlignBytes == 0, + "Invalid TMA alignment"); + + // Named barrier indices + constexpr int kNotifyBarrierIndex = 1; + + // Mooncake Gin handle + // Each warp is a channel + const auto [qp_idx, sharing_mode] = + comm::get_qp_mode 0)>( + sm_idx, (warp_idx - kNumNotifyWarps) % kNumChannelsPerSM, + warp_idx < kNumNotifyWarps); + const auto gin = transport::MooncakeGin( + comm_ctx, qp_idx, sharing_mode, kNumQPs, scaleout_rank_idx, + scaleup_rank_idx, kNumScaleupRanks, kNumRanks); + + // Global parallel barriers for scale-out subteam and scale-up subteam + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, + thread_idx); + + // The golden layout during the whole process for both scale-out and forward + // warps + const auto token_layout = layout::TokenLayout( + kNumHiddenBytes, kNumSFPacks * sizeof(sf_pack_t), kNumTopk, true); + const auto tma_buffer = + layout::BufferLayout( + token_layout, kNumScaleoutWarps + kNumForwardWarps, 1, + math::advance_ptr(smem, kNumSmemBytesForNotify)) + .get_rank_buffer(warp_idx - kNumNotifyWarps) + .get_token_buffer(0); + + // All the buffers + auto scaleup_buffer = layout::BufferLayout( + token_layout, kNumScaleupRanks, + kNumScaleoutRanks * kNumMaxTokensPerRank, buffer); + auto scaleout_send_buffer = + layout::BufferLayout(token_layout, 1, kNumMaxTokensPerRank, + scaleup_buffer.get_buffer_end_ptr()); + auto scaleout_recv_buffer = layout::BufferLayout( + token_layout, kNumScaleoutRanks, kNumChannels * kNumMaxTokensPerChannel, + scaleout_send_buffer.get_buffer_end_ptr()); + + // Init TMA for scale-out and forward warps + ptx::arrival_phase phase = 0; + const auto mbarrier_ptr = tma_buffer.get_mbarrier_ptr(); + if (warp_idx >= kNumNotifyWarps and ptx::elect_one_sync()) + ptx::mbarrier_init_with_fence(mbarrier_ptr, 1); + __syncwarp(); + + // Different warp roles + if (warp_idx < kNumNotifyWarps) { + // Assign shared memory + constexpr int kNumAlignedElems = kNumSmemBytesForNotify / sizeof(int); + const auto rank_expert_count = math::advance_ptr(smem, 0); + + // Clean initial counts + // NOTES: if you want to change the order of different warp roles, + // please take care of the `thread_idx` + int *rank_count = rank_expert_count, + *expert_count = rank_expert_count + kNumRanks; +#pragma unroll + for (int i = 0; i < kNumAlignedElems / kNumNotifyThreads; ++i) + rank_expert_count[i * kNumNotifyThreads + thread_idx] = 0; + ptx::named_barrier(kNotifyBarrierIndex); + + // Atomic add on shared memory + EP_STATIC_ASSERT(kNumTopk <= 32, "Insufficient lanes"); + const auto global_warp_idx = sm_idx * kNumNotifyWarps + warp_idx; + for (int i = global_warp_idx; i < num_tokens; + i += kNumNotifyWarps * kNumSMs) { + // Expert choice can not be redundant + // NOTES: no assertions here as they are expensive + const auto dst_expert_idx = + lane_idx < kNumTopk ? static_cast(__ldg( + topk_idx + i * kNumTopk + lane_idx)) + : -1; + if (dst_expert_idx >= 0) + atomicAdd_block(expert_count + dst_expert_idx, 1); + + // Rank choice should do deduplication here + const auto dst_rank_idx = + dst_expert_idx >= 0 ? dst_expert_idx / kNumExpertsPerRank : -1; + if (ptx::deduplicate(dst_rank_idx, lane_idx) and dst_rank_idx >= 0) + atomicAdd_block(rank_count + dst_rank_idx, 1); + } + ptx::named_barrier(kNotifyBarrierIndex); + +// Do full-grid reduction +#pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + const int64_t counter = (1ll << 32ll) | rank_expert_count[i]; + ptx::red_add( + workspace_layout.get_notify_reduction_workspace_ptr() + i, + counter); + } + + // Do the remaining work by SM 0 + if (sm_idx == 0) { +// Reduce all SM's count +// Wait all SMs' arrival +#pragma unroll + for (int i = thread_idx; i < kNumRanks + kNumExperts; + i += kNumNotifyThreads) { + comm::timeout_while([=](const bool& + is_last_check) { + const auto status = ptx::ld_volatile( + workspace_layout.get_notify_reduction_workspace_ptr() + + i); + if ((status >> 32) == kNumSMs) { + // Encode and write into the send buffer + workspace_layout + .get_scaleout_rank_expert_count_ptr()[i] = + math::encode_decode_positive(status & + 0xffffffffll); + + // Clean for the next usage + workspace_layout + .get_notify_reduction_workspace_ptr()[i] = 0; + return true; + } + + if (is_last_check) { + printf( + "DeepEP hybrid notify (GPU reduction) timeout, " + "scale-out: %d/%d, scale-up: %d/%d, " + "thread: %d, status: %d | %d, expected: %d\n", + scaleout_rank_idx, kNumScaleoutRanks, + scaleup_rank_idx, kNumScaleupRanks, thread_idx, + static_cast(status >> 32), + static_cast(status & 0xffffffff), kNumSMs); + } + return false; + }); + } + ptx::named_barrier(kNotifyBarrierIndex); + + // Issue scaleout writes to peers + EP_STATIC_ASSERT( + kReuseSlotIndices or kNumScaleoutRanks <= kNumNotifyThreads, + "kNumScaleoutRanks must be less than kNumNotifyThreads"); + if (thread_idx < kNumScaleoutRanks) { + const auto dst_scaleout_rank_idx = thread_idx; + gin.put( + workspace_layout.get_scaleout_rank_count_ptr( + scaleout_rank_idx), + workspace_layout.get_scaleout_rank_count_ptr( + dst_scaleout_rank_idx), + kNumScaleupRanks * sizeof(int), dst_scaleout_rank_idx, 0); + gin.put( + workspace_layout.get_scaleout_expert_count_ptr( + scaleout_rank_idx), + workspace_layout.get_scaleout_expert_count_ptr( + dst_scaleout_rank_idx), + kNumExpertsPerScaleout * sizeof(int), + dst_scaleout_rank_idx); + } + __syncwarp(); + + // Util functions to get metadata from scale-out peers + // NOTES: this is correct as RDMA operations has a minimum write + // granularity of 1024 bytes (a whole integer write is atomic) + const auto recv_and_reduce = [=](const auto& get_ptr_func, + const bool& is_expert_reduction = + false) -> int { + int count = 0; +#pragma unroll + for (int j = 0; j < kNumScaleoutRanks; ++j) { + const auto ptr = get_ptr_func(j); + int decoded; + comm::timeout_while< + kNumTimeoutCycles>([&](const bool& is_last_check) { + decoded = math::encode_decode_positive( + ptx::ld_acquire_sys(ptr)); + if (math::is_decoded_positive_ready(decoded)) + return true; + + if (is_last_check) { + printf( + "DeepEP hybrid notify (scale-out %s reduction) " + "timeout, " + "scale-out: %d, scale-up: %d, " + "thread: %d, wait scale-out: %d, decoded: %d\n", + is_expert_reduction ? "expert" : "rank", + scaleout_rank_idx, scaleup_rank_idx, thread_idx, + j, decoded); + } + return false; + }); + + // Add and clean for next usages + count += decoded, *ptr = 0; + } + return count; + }; + +// Write into all scale-up peers' rank-level counters +#pragma unroll + for (int i = thread_idx; i < kNumScaleupRanks; + i += kNumNotifyThreads) { + // Wait scale-out arrival and reduce + const auto count = + recv_and_reduce([=](const int& scaleout_peer_idx) { + return workspace_layout + .get_scaleout_rank_count_ptr( + scaleout_peer_idx, i); + }); + + // Write into the remote scale-up peer + const int64_t counter = + (static_cast(kNumScaleupRanks) << 32ll) | count; + gin.put_value( + workspace_layout.get_scaleup_rank_count_ptr() + + scaleup_rank_idx, + counter, i); + } + __syncwarp(); + +// Atomic add into all scale-up peers' expert-level counters +#pragma unroll + for (int i = thread_idx; i < kNumExpertsPerScaleout; + i += kNumNotifyThreads) { + // Wait scale-out arrival and reduce + const auto count = recv_and_reduce( + [=](const int& scaleout_peer_idx) { + return workspace_layout + .get_scaleout_expert_count_ptr( + scaleout_peer_idx, i); + }, + true); + + // Write into the remote scale-up peer + const int64_t counter = (1ll << 32ll) | count; + const auto dst_scaleup_rank_idx = i / kNumExpertsPerRank; + const auto expert_idx_in_dst_rank = i % kNumExpertsPerRank; + gin.red_add_rel( + workspace_layout.get_scaleup_expert_count_ptr() + + expert_idx_in_dst_rank, + counter, dst_scaleup_rank_idx); + } + // There are shared memory reads above, a barrier is necessary + ptx::named_barrier(kNotifyBarrierIndex); + + // NOTES: from now on, the `rank` and `expert`s size change into the + // local size + expert_count = rank_expert_count + kNumScaleupRanks; + + // Wait local counters to be ready + // NOTES: here we only care the prefix sum by scale-up peers (used + // for later epilogue), not all ranks + EP_STATIC_ASSERT( + kNumNotifyWarps == 0 or kNumScaleupRanks + kNumExpertsPerRank <= + kNumNotifyWarps * 32, + "Insufficient notify threads"); + comm::timeout_while( + thread_idx < kNumScaleupRanks + kNumExpertsPerRank, + [&](const bool& is_last_check) { + const auto status = ptx::ld_volatile( + workspace_layout + .get_scaleup_rank_expert_count_ptr() + + thread_idx); + if ((status >> 32ull) == kNumScaleupRanks) { + // Clean GPU workspace and write into host workspace + const auto count = + static_cast(status & 0xffffffffll); + const auto aligned_count = math::align( + count, thread_idx < kNumScaleupRanks + ? 1 + : kExpertAlignment); + + workspace_layout.get_scaleup_rank_expert_count_ptr< + false>()[thread_idx] = 0; + if constexpr (kDoCPUSync) { + host_workspace_layout + .get_scaleup_rank_expert_count_ptr< + false>()[thread_idx] = + math::encode_decode_positive(aligned_count); + } + + // Update statistics counters + if (cumulative_local_expert_recv_stats != nullptr and + thread_idx >= kNumScaleupRanks) + atomicAdd(cumulative_local_expert_recv_stats + + (thread_idx - kNumScaleupRanks), + count); + + // Save for later prefix sum calculation + rank_expert_count[thread_idx] = aligned_count; + return true; + } + + if (is_last_check) { + printf( + "DeepEP hybrid notify (scale-up reduction) timeout," + "scale-out: %d/%d, scale-up: %d/%d, " + "thread: %d, status: %d | %d, expected: %d\n", + scaleout_rank_idx, kNumScaleoutRanks, + scaleup_rank_idx, kNumScaleupRanks, thread_idx, + static_cast(status >> 32), + static_cast(status & 0xffffffff), + kNumScaleupRanks); + } + return false; + }); + ptx::named_barrier(kNotifyBarrierIndex); + + // Do prefix sum by the warps of the first SM + // NOTES: we may have fast implementation with `cub::BlockScan`, but + // it is too heavy to use + const auto do_psum = [=](const int* count, int* out, const int n, + const int is_exclusive) { + int psum = 0; +#pragma unroll + for (int i = 0; i < math::ceil_div(n + is_exclusive, 32); ++i) { + const auto idx = i * 32 + lane_idx; + const auto mem_idx = idx - is_exclusive; + const auto value = + (0 <= mem_idx and mem_idx < n) ? count[mem_idx] : 0; + const auto sum = + psum + ptx::warp_inclusive_sum(value, lane_idx); + + // Store into global memory + if (idx < n + is_exclusive) out[idx] = sum; + + // Update `psum` by using the last lane's value + psum = ptx::exchange(sum, 31); + } + }; + if (warp_idx == 0) { + // Inclusive prefix sum + do_psum(rank_count, psum_num_recv_tokens_per_scaleup_rank, + kNumScaleupRanks, 0); + } else if (warp_idx == 1) { + // Exclusive prefix sum for later expanding + do_psum(expert_count, psum_num_recv_tokens_per_expert, + kNumExpertsPerRank, 1); + } + } + } else if (warp_idx < kNumNotifyWarps + kNumScaleoutWarps) { + const int scaleout_warp_idx = warp_idx - kNumNotifyWarps; + const int channel_idx = sm_idx * kNumChannelsPerSM + scaleout_warp_idx; + scaleout_recv_buffer = + scaleout_recv_buffer.get_rank_buffer(scaleout_rank_idx); + scaleout_recv_buffer = + scaleout_recv_buffer.get_channel_buffer( + channel_idx); + + // Channel metadata maintenance + EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, + "Invalid number of scale-out ranks"); + int stored_scaleout_tail = 0, stored_old_scaleout_tail = 0; + const auto update_scaleout_tail = [&](const bool& finish_flag = false) { + if (lane_idx < kNumScaleoutRanks and + (stored_scaleout_tail >= + stored_old_scaleout_tail + kScaleoutUpdateInterval or + finish_flag)) { + const auto signaled_tail = math::pack2( + finish_flag, stored_scaleout_tail); + const auto ptr = + workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, scaleout_rank_idx); + const auto old_signaled_tail = + math::pack2(0, stored_old_scaleout_tail); + + // NOTES: the "release" scope will be `sys` for the local rank + // (we may involve NVLink so not `gpu`) For RDMA requests, + // "release" is ensured by "atomic" + gin.red_add_rel( + ptr, signaled_tail - old_signaled_tail, lane_idx, + transport::kRedAddReleaseLowWordLast); + stored_old_scaleout_tail = stored_scaleout_tail; + } + __syncwarp(); + }; + + // Preload next token + const auto preload_next_token = [&](const int& token_idx) { + if (token_idx >= num_tokens) return; + + // Issue TMA load + const auto token_i64_idx = static_cast(token_idx); + if (ptx::elect_one_sync()) { + ptx::tma_load_1d( + tma_buffer.get_hidden_ptr(), + math::advance_ptr(x, token_i64_idx * kNumHiddenBytes), + mbarrier_ptr, kNumHiddenBytes); + } + __syncwarp(); + + // Issue SF `cp.async` + if constexpr (kNumSFPacks > 0) { + EP_STATIC_ASSERT(sizeof(sf_pack_t) % 4 == 0, + "Unaligned SF element type"); + const auto gmem_src_ptr = math::advance_ptr( + sf, token_i64_idx * sf_token_stride * sizeof(sf_pack_t)); + const auto smem_dst_ptr = tma_buffer.get_sf_ptr(); + + constexpr auto kNumFullIters = kNumSFPacks / 32; +#pragma unroll + for (int k = 0; k < kNumFullIters; ++k) { + ptx::cp_async_ca( + gmem_src_ptr + (k * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + k * 32 + lane_idx); + } + if (kNumFullIters * 32 + lane_idx < kNumSFPacks) { + ptx::cp_async_ca( + gmem_src_ptr + + (kNumFullIters * 32 + lane_idx) * sf_hidden_stride, + smem_dst_ptr + kNumFullIters * 32 + lane_idx); + } + ptx::cp_async_mbarrier_arrive(mbarrier_ptr); + __syncwarp(); + } + }; + + // Iterate all tokens + preload_next_token(channel_idx); + for (int token_idx = channel_idx; token_idx < num_tokens; + token_idx += kNumChannels) { + // Load top-k indices and weights + EP_STATIC_ASSERT(kNumTopk <= 32, + "Insufficient lanes for loading top-k indices"); + int stored_dst_scaleout_rank_idx = -1; + if (lane_idx < kNumTopk) { + const auto uncasted_dst_expert_idx = + __ldg(topk_idx + token_idx * kNumTopk + lane_idx); + const auto dst_expert_idx = + static_cast(uncasted_dst_expert_idx); + stored_dst_scaleout_rank_idx = + dst_expert_idx >= 0 + ? dst_expert_idx / kNumExpertsPerScaleout + : -1; + tma_buffer.get_topk_idx_ptr()[lane_idx] = dst_expert_idx; + if (topk_weights != nullptr) + tma_buffer.get_topk_weights_ptr()[lane_idx] = + __ldg(topk_weights + token_idx * kNumTopk + lane_idx); + if (copied_topk_idx != nullptr) + copied_topk_idx[token_idx * kNumTopk + lane_idx] = + uncasted_dst_expert_idx; + } + __syncwarp(); + + // Add source metadata (rank index and token index) + if (ptx::elect_one_sync()) + *tma_buffer.get_src_token_global_idx_ptr() = + rank_idx * kNumMaxTokensPerRank + token_idx; + ptx::tma_store_fence(); + __syncwarp(); + + // Deduplicate ranks and assign slots + int stored_dst_slot_idx = -1; + const auto stored_old_slot_idx = ptx::exchange( + stored_scaleout_tail, stored_dst_scaleout_rank_idx >= 0 + ? stored_dst_scaleout_rank_idx + : 0); + if (ptx::deduplicate(stored_dst_scaleout_rank_idx, lane_idx) and + stored_dst_scaleout_rank_idx >= 0) + stored_dst_slot_idx = stored_old_slot_idx; + + // Update scale-out tail + const auto scaleout_rank_mask = + ptx::reduce_or(stored_dst_scaleout_rank_idx >= 0 + ? (1u << stored_dst_scaleout_rank_idx) + : 0u); + stored_scaleout_tail += (scaleout_rank_mask >> lane_idx) & 1; + + // Wait TMA arrival and issue the TMA store into send buffer + if (ptx::elect_one_sync()) { + ptx::mbarrier_arrive_and_set_tx(mbarrier_ptr, kNumHiddenBytes); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + + // So if no ranks will go by RDMA, we skip the send buffer + // stores + if (scaleout_rank_mask ^ (1 << scaleout_rank_idx)) { + ptx::tma_store_1d( + scaleout_send_buffer.get_token_buffer(token_idx) + .get_base_ptr(), + tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + } + } + __syncwarp(); + + // Local rank can be bypassed + if (stored_dst_slot_idx >= 0 and + stored_dst_scaleout_rank_idx == scaleout_rank_idx) { + ptx::tma_store_1d( + scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + } + ptx::tma_store_commit(); + ptx::tma_store_wait(); + __syncwarp(); + + // Preload the next token (overlapping with the IBGDA issues) + preload_next_token(token_idx + kNumChannels); + + // Issue IBGDA requests + if (stored_dst_slot_idx >= 0 and + stored_dst_scaleout_rank_idx != scaleout_rank_idx) { + gin.put( + scaleout_recv_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + scaleout_send_buffer.get_token_buffer(token_idx) + .get_base_ptr(), + tma_buffer.get_num_bytes(), + stored_dst_scaleout_rank_idx, 0); + } + __syncwarp(); + + // Issue scale-out tail update + update_scaleout_tail(); + } + + // Flush unflushed tails + update_scaleout_tail(true); + } else { + const int forward_warp_idx = + warp_idx - (kNumNotifyWarps + kNumScaleoutWarps); + const int channel_idx = sm_idx * kNumChannelsPerSM + forward_warp_idx; + scaleout_recv_buffer = + scaleout_recv_buffer.get_channel_buffer( + channel_idx); + scaleup_buffer = scaleup_buffer.get_rank_buffer(scaleup_rank_idx); + + // Shape of `token_metadata_at_forward`: `[kNumChannels, + // kNumScaleoutRanks * kNumMaxTokensPerChannel + 1, + // kNumForwardMetadataDims]` + constexpr int kNumForwardMetadataDims = 2 + kNumTopk * 2; + token_metadata_at_forward += + channel_idx * ((kNumScaleoutRanks * kNumMaxTokensPerChannel + 1) * + kNumForwardMetadataDims); + + // Shape of `dst_buffer_slot_idx`: `[kNumChannels, kNumScaleoutRanks, + // kNumMaxTokensPerChannel, kNumTopk]` + dst_buffer_slot_idx += + channel_idx * + (kNumScaleoutRanks * kNumMaxTokensPerChannel * kNumTopk); + + // Transform linked list index + const auto transform_linked_list_idx = [=](const int& idx) { + constexpr int kNumTokensInLinkedList = + kNumMaxTokensPerChannel * kNumScaleoutRanks + 1; + return channel_idx * (kNumTokensInLinkedList * kNumScaleupRanks) + + idx * kNumScaleupRanks + scaleup_rank_idx; + }; + + // Forward tokens from scale-out ranks + EP_STATIC_ASSERT(kNumScaleoutRanks <= 32, "Too many scale-out ranks"); + int num_tokens_processed = 0; + int stored_scaleout_old_tail_idx = 0; + int stored_scaleup_send_counters[kNumScaleupRanksPerLane] = {}; + int stored_finish_flag = lane_idx >= kNumScaleoutRanks; + int stored_scaleout_tail_idx = 0; + int recv_scaleout_rank_idx = channel_idx % kNumScaleoutRanks; + uint32_t wip_mask; + while ((wip_mask = ptx::gather(stored_scaleout_tail_idx > + stored_scaleout_old_tail_idx or + stored_finish_flag == 0))) { + // Pick next rank in round-robin + const auto offset = + (recv_scaleout_rank_idx + 1) % kNumScaleoutRanks; + const auto hi_mask = (wip_mask >> offset) << offset; + recv_scaleout_rank_idx = + hi_mask ? ptx::ffs(hi_mask) : ptx::ffs(wip_mask); + + // Wait for this rank to have data (or finish) + comm::timeout_while([&](const bool& + is_last_check) { + const uint32_t arrived_or_finished = + stored_scaleout_tail_idx > stored_scaleout_old_tail_idx or + stored_finish_flag > 0; + if (ptx::exchange(arrived_or_finished, recv_scaleout_rank_idx)) + return true; + + // Timeout + if (is_last_check) { + if (lane_idx < kNumScaleoutRanks) { + printf( + "DeepEP hybrid dispatch (forwarding) timeout, " + "scale-out: %d, scale-up: %d, " + "channel: %d, lane: %d, old scale-out tail: %d, " + "scale-out tail: (%d, %d)\n", + scaleout_rank_idx, scaleup_rank_idx, channel_idx, + lane_idx, stored_scaleout_old_tail_idx, + stored_finish_flag, stored_scaleout_tail_idx); + } + return false; + } + + // Read new signaled tails + if (lane_idx < kNumScaleoutRanks) { + const auto signaled_tail = ptx::ld_acquire_sys( + workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, lane_idx)); + math::unpack2(signaled_tail, + stored_finish_flag, + stored_scaleout_tail_idx); + } + __syncwarp(); + return false; + }); + + // Process one chunk from the current rank + const auto start_slot_idx = ptx::exchange( + stored_scaleout_old_tail_idx, recv_scaleout_rank_idx); + const auto end_slot_idx = std::min( + ptx::exchange(stored_scaleout_tail_idx, recv_scaleout_rank_idx), + start_slot_idx + kNumSlotsPerForwardChunk); + if (lane_idx == recv_scaleout_rank_idx) + stored_scaleout_old_tail_idx = end_slot_idx; + + const auto recv_buffer = + scaleout_recv_buffer.get_rank_buffer(recv_scaleout_rank_idx); + for (int slot_idx = start_slot_idx; slot_idx < end_slot_idx; + ++slot_idx) { + const auto token_buffer = + recv_buffer.get_token_buffer(slot_idx); + + // Wait TMA arrival + ptx::tma_store_wait(); + __syncwarp(); + + // TMA load into shared memory + if (ptx::elect_one_sync()) { + ptx::tma_load_1d(tma_buffer.get_base_ptr(), + token_buffer.get_base_ptr(), mbarrier_ptr, + token_layout.get_num_bytes()); + ptx::mbarrier_arrive_and_set_tx( + mbarrier_ptr, token_layout.get_num_bytes()); + ptx::mbarrier_wait_and_flip_phase(mbarrier_ptr, phase); + } + __syncwarp(); + + // Read top-k indices + EP_STATIC_ASSERT(kNumTopk <= 32, "Too many top-k selections"); + int stored_dst_scaleup_rank_idx = -1; + auto dst_expert_idx = + lane_idx < kNumTopk + ? tma_buffer.get_topk_idx_ptr()[lane_idx] + : -1; + dst_expert_idx -= scaleout_rank_idx * kNumExpertsPerScaleout; + stored_dst_scaleup_rank_idx = + 0 <= dst_expert_idx and + dst_expert_idx < kNumExpertsPerScaleout + ? dst_expert_idx / kNumExpertsPerRank + : -1; + + // Write the per-scaleup channel index for this token + int linked_list_idx = -1; +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) { + const auto src_lane_idx = + stored_dst_scaleup_rank_idx - j * 32; + const bool valid = 0 <= src_lane_idx and src_lane_idx < 32; + const auto exchanged = + ptx::exchange(stored_scaleup_send_counters[j], + valid ? src_lane_idx : 0); + linked_list_idx = valid ? exchanged : linked_list_idx; + } + if (not kReuseSlotIndices and lane_idx < kNumTopk) { + tma_buffer.get_linked_list_idx_ptr()[lane_idx] = + transform_linked_list_idx(linked_list_idx); + ptx::tma_store_fence(); + } + __syncwarp(); + + // Deduplicate for scale-up ranks + int stored_dst_slot_idx = -1; + const auto dst_slot_idx_ptr = + dst_buffer_slot_idx + + recv_scaleout_rank_idx * + (kNumMaxTokensPerChannel * kNumTopk) + + slot_idx * kNumTopk; + if constexpr (kReuseSlotIndices) { + if (lane_idx < kNumTopk) + stored_dst_slot_idx = + __ldg(dst_slot_idx_ptr + lane_idx); + } else { + // Deduplicate for NVLink ranks + if (ptx::deduplicate(stored_dst_scaleup_rank_idx, + lane_idx) and + stored_dst_scaleup_rank_idx >= 0) + stored_dst_slot_idx = atomicAdd( + workspace_layout + .get_scaleup_atomic_sender_counter() + + stored_dst_scaleup_rank_idx, + 1); + } + __syncwarp(); + + // Issue TMAs + if (stored_dst_slot_idx >= 0) { + const auto dst_ptr = + gin.get_sym_ptr( + scaleup_buffer.get_token_buffer(stored_dst_slot_idx) + .get_base_ptr(), + stored_dst_scaleup_rank_idx); + ptx::tma_store_1d(dst_ptr, tma_buffer.get_base_ptr(), + tma_buffer.get_num_bytes()); + ptx::tma_store_commit(); + } + __syncwarp(); + + // Add per-scale-up counter + EP_STATIC_ASSERT(kNumScaleupRanks <= 64, + "Invalid number of scale-up peers"); + using mask_t = std::conditional_t; + const auto scaleup_send_mask = ptx::reduce_or( + stored_dst_scaleup_rank_idx >= 0 + ? (mask_t(1) << stored_dst_scaleup_rank_idx) + : mask_t(0)); +#pragma unroll + for (int j = 0; j < kNumScaleupRanksPerLane; ++j) + stored_scaleup_send_counters[j] += + (scaleup_send_mask >> (j * 32 + lane_idx)) & 1; + + // Record metadata at forward + if constexpr (not kReuseSlotIndices) { + EP_STATIC_ASSERT(kNumTopk <= 32, + "Invalid number of selections"); + const auto metadata_ptr = + token_metadata_at_forward + + num_tokens_processed * kNumForwardMetadataDims; + + // Source token index and last token index flag + if (ptx::elect_one_sync()) { + metadata_ptr[0] = + tma_buffer.get_src_token_global_idx_ptr()[0]; + metadata_ptr[1] = slot_idx == (end_slot_idx - 1); + } + + // Second, original top-k indices and destination slots + if (lane_idx < kNumTopk) { + metadata_ptr[2 + lane_idx] = + stored_dst_scaleup_rank_idx; + metadata_ptr[2 + kNumTopk + lane_idx] = + stored_dst_slot_idx; + dst_slot_idx_ptr[lane_idx] = stored_dst_slot_idx; + } + } + num_tokens_processed += 1; + __syncwarp(); + } + } + + // Assign the source token index part of the metadata into `-1` as an + // ending mark + if (not kReuseSlotIndices and ptx::elect_one_sync()) + token_metadata_at_forward[num_tokens_processed * + kNumForwardMetadataDims] = -1; + __syncwarp(); + + // Update linked list's ending position + if constexpr (not kReuseSlotIndices) { + const auto tail_ptr = workspace_layout.get_channel_scaleup_tail_ptr( + channel_idx, scaleup_rank_idx); +#pragma unroll + for (int i = 0; i < kNumScaleupRanksPerLane; ++i) { + if (const auto j = i * 32 + lane_idx; + i < (kNumScaleupRanksPerLane - 1) or j < kNumScaleupRanks) { + ptx::st_relaxed_sys( + gin.get_sym_ptr(tail_ptr, j), + transform_linked_list_idx( + stored_scaleup_send_counters[i])); + } + } + } + __syncwarp(); + + // Clean tails for next usages + if (lane_idx < kNumScaleoutRanks) + *workspace_layout.get_scaleout_channel_signaled_tail_ptr( + channel_idx, lane_idx) = 0; + __syncwarp(); + } + + // Scale-up barrier to ensure data arrival + // As scale-out tokens have already been consumed by forwarders, no need to + // do scale-out barrier again + comm::gpu_barrier( + gin, workspace_layout, scaleout_rank_idx, scaleup_rank_idx, sm_idx, + thread_idx, /* do not scale-out */ false, true); + + // Trigger the copy epilogue kernel +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + + // Clean scale-up counters + // All scale-out counters should be cleaned before + EP_STATIC_ASSERT(kNumScaleupRanks <= kNumThreads, "Insufficient threads"); + if (not kReuseSlotIndices and sm_idx == 0 and thread_idx < kNumScaleupRanks) + workspace_layout.get_scaleup_atomic_sender_counter()[thread_idx] = 0; +} + +} // namespace mooncake::elastic diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_launch.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_launch.cuh new file mode 100644 index 0000000000..e21c86a170 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_launch.cuh @@ -0,0 +1,78 @@ +#pragma once + +#include + +#include +#include + +namespace mooncake { + +struct ElasticLaunchContext { + void* gdr_buffer = nullptr; + const int32_t* nvlink_available = nullptr; + void* const* ipc_peer_ptrs = nullptr; + void* raddrs = nullptr; + void* rkeys = nullptr; + void* qp_devctxs = nullptr; + const void* rdma_send_signal_buffer = nullptr; + const void* rdma_recv_signal_buffer = nullptr; + void* buffer = nullptr; + void* workspace = nullptr; + void* mapped_host_workspace = nullptr; + int rank = 0; + int num_ranks = 1; + int scaleout_rank_idx = 0; + int scaleup_rank_idx = 0; + int num_scaleout_ranks = 1; + int num_scaleup_ranks = 1; + bool is_scaleup_nvlink = true; + int num_qps = 1; + int64_t timeout_cycles = -1; +}; + +void launch_elastic_dispatch_deterministic_prologue( + const int64_t* topk_idx, int* rank_count_buffer, int* dst_buffer_slot_idx, + int num_tokens, int num_max_tokens_per_rank, int num_experts, int num_topk, + int scaleup_rank_idx, int num_scaleup_ranks, int num_sms, + int num_smem_bytes, cudaStream_t stream); + +void launch_mooncake_elastic_dispatch( + void* x, void* sf, int64_t* topk_idx, float* topk_weights, + int64_t* copied_topk_idx, int* cumulative_local_expert_recv_stats, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, int* dst_buffer_slot_idx, + int* token_metadata_at_forward, int num_tokens, int num_max_tokens_per_rank, + int hidden, int elem_size, int num_sf_packs, int sf_token_stride, + int sf_hidden_stride, int num_experts, int num_topk, int expert_alignment, + int num_sms, int num_channels_per_sm, int num_smem_bytes, bool cached_mode, + bool deterministic, bool do_cpu_sync, const ElasticLaunchContext& ctx, + cudaStream_t stream); + +void launch_mooncake_elastic_dispatch_copy_epilogue( + void* recv_x, void* recv_sf, int64_t* recv_topk_idx, + float* recv_topk_weights, int* recv_src_metadata, int* channel_linked_list, + int num_recv_tokens, int num_max_tokens_per_rank, int hidden, int elem_size, + int num_sf_packs, int recv_sf_token_stride, int recv_sf_hidden_stride, + int num_experts, int num_topk, int num_sms, int num_smem_bytes, + int num_channels, bool do_expand, bool cached_mode, + const ElasticLaunchContext& ctx, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, cudaStream_t stream); + +void* launch_mooncake_elastic_combine( + void* x, float* topk_weights, int* src_metadata, + int* psum_num_recv_tokens_per_scaleup_rank, int* token_metadata_at_forward, + int* channel_linked_list, int num_reduced_tokens, + int num_max_tokens_per_rank, int hidden, int num_experts, int num_topk, + int num_sms, int num_smem_bytes, int num_channels, bool use_expanded_layout, + bool allow_multiple_reduction, const ElasticLaunchContext& ctx, + cudaStream_t stream); + +void launch_mooncake_elastic_combine_reduce_epilogue( + void* combined_x, float* combined_topk_weights, int64_t* combined_topk_idx, + int num_combined_tokens, int num_max_tokens_per_rank, int hidden, + int num_experts, int num_topk, void* reduce_buffer, void* bias_0, + void* bias_1, int num_sms, int num_smem_bytes, bool use_expanded_layout, + bool allow_multiple_reduction, const ElasticLaunchContext& ctx, + cudaStream_t stream); + +} // namespace mooncake diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_layout.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_layout.cuh new file mode 100644 index 0000000000..f03d5943bb --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_layout.cuh @@ -0,0 +1,374 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include +#include +#include + +namespace mooncake::elastic::layout { + +struct WorkspaceLayout { + void* workspace; + + int num_ranks; + int num_scaleout_ranks, num_scaleup_ranks; + int num_experts, num_experts_per_rank; + + // We want to fix the layout position for all settings, + // so that one buffer can be reused for all cases + static constexpr int kNumMaxRanks = 1024; + static constexpr int kNumMaxExperts = 2048; + static constexpr int kNumMaxExpertsPerRank = 256; + static constexpr int kNumMaxInflightAGRS = 32; + + // Mooncake Device API does not rely on NCCL GIN remote RED on a single + // symmetric signal word. Use per-source-rank signal slots for both phases: + // each sender atomically updates its own slot with release semantics and + // receivers poll the full slot vector. Keep an independent counter/slot + // vector for each logical barrier tag, as hybrid kernels mix world and + // scale-up-only barriers in the same workspace and therefore must not share + // phase/sign state across tags. + static constexpr int kNumBarrierTags = 16; + static constexpr int64_t kNumBarrierBytesPerTag = + sizeof(unsigned long long) + 2 * kNumMaxRanks * sizeof(int); + static constexpr int64_t kNumBarrierSignalBytes = + kNumBarrierTags * kNumBarrierBytesPerTag; + + __forceinline__ __device__ __host__ + WorkspaceLayout(void* workspace, const int& num_scaleout_ranks, + const int& num_scaleup_ranks, const int& num_experts) + : workspace(workspace), + num_ranks(num_scaleout_ranks * num_scaleup_ranks), + num_scaleout_ranks(num_scaleout_ranks), + num_scaleup_ranks(num_scaleup_ranks), + num_experts(num_experts) { + num_experts_per_rank = num_experts / num_ranks; + EP_UNIFIED_ASSERT(num_experts % num_ranks == 0); + EP_UNIFIED_ASSERT(num_ranks <= kNumMaxRanks); + EP_UNIFIED_ASSERT(num_experts <= kNumMaxExperts); + EP_UNIFIED_ASSERT(num_experts_per_rank <= kNumMaxExpertsPerRank); + } + + static int64_t get_num_bytes() { + // Pure NVLink scaleup barrier signals + int64_t num_bytes = 0; + num_bytes += kNumBarrierSignalBytes; + + // Notify reduction workspace + num_bytes += (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t); + + // Scaleup notify threads + // Rank send/recv count + num_bytes += kNumMaxRanks * sizeof(int64_t) * 2; + // Expert send/recv count + num_bytes += kNumMaxExperts * sizeof(int64_t) * 2; + + // Scaleup atomic sender count + num_bytes += kNumMaxRanks * sizeof(int); + + // Scaleout notify threads + // Rank send/recv count + num_bytes += kNumMaxRanks * sizeof(int) * 2; + // Expert send/recv count + num_bytes += kNumMaxExperts * sizeof(int) * 2; + + // Scaleout channel metadata (finish flag and tails) + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int64_t); + + // Channel aggregated into the scaleup domains + // Also reused for channel scaleup tail + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int); + + // Rank send/recv count, for PP prev/next ranks + num_bytes += 2 * 2 * sizeof(int64_t); + + // AGRS signals + num_bytes += (kNumMaxInflightAGRS + 1) * kNumMaxRanks * sizeof(int); + + // Ensure LDG.256 work + return math::align(num_bytes, 32); + } + + __forceinline__ __device__ __host__ unsigned long long* + get_nvl_barrier_counter_ptr(int tag = 0) const { + EP_UNIFIED_ASSERT(tag >= 0 && tag < kNumBarrierTags); + return math::advance_ptr( + workspace, tag * kNumBarrierBytesPerTag); + } + + __forceinline__ __device__ __host__ int* get_nvl_barrier_signal_ptr( + int tag, int phase) const { + EP_UNIFIED_ASSERT(tag >= 0 && tag < kNumBarrierTags); + EP_UNIFIED_ASSERT(phase >= 0 && phase < 2); + return math::advance_ptr(workspace, + tag * kNumBarrierBytesPerTag + + sizeof(unsigned long long) + + phase * kNumMaxRanks * sizeof(int)); + } + + __forceinline__ __device__ __host__ int64_t* + get_notify_reduction_workspace_ptr() const { + return math::advance_ptr(workspace, kNumBarrierSignalBytes); + } + + template + __forceinline__ __device__ __host__ int64_t* + get_scaleup_rank_expert_count_ptr() const { + const auto base_ptr = math::advance_ptr( + get_notify_reduction_workspace_ptr(), + (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t)); + return base_ptr + (kIsSendBuffer ? 0 : kNumMaxRanks + kNumMaxExperts); + } + + template + __forceinline__ __device__ __host__ int64_t* get_scaleup_rank_count_ptr() + const { + return get_scaleup_rank_expert_count_ptr(); + } + + template + __forceinline__ __device__ __host__ int64_t* get_scaleup_expert_count_ptr() + const { + return get_scaleup_rank_expert_count_ptr() + + num_scaleup_ranks; + } + + __forceinline__ __device__ __host__ int* get_scaleup_atomic_sender_counter() + const { + return math::advance_ptr( + get_scaleup_rank_expert_count_ptr(), + 2 * (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t)); + } + + template + __forceinline__ __device__ __host__ int* + get_scaleout_rank_expert_count_ptr() const { + const auto base_ptr = math::advance_ptr( + get_scaleup_atomic_sender_counter(), kNumMaxRanks * sizeof(int)); + return base_ptr + (kIsSendBuffer ? 0 : kNumMaxRanks + kNumMaxExperts); + } + + template + __forceinline__ __device__ __host__ int* get_scaleout_rank_count_ptr( + const int& scaleout_rank_idx = 0, + const int& scaleup_rank_idx = 0) const { + const auto base_ptr = + get_scaleout_rank_expert_count_ptr(); + return base_ptr + scaleout_rank_idx * num_scaleup_ranks + + scaleup_rank_idx; + } + + template + __forceinline__ __device__ __host__ int* get_scaleout_expert_count_ptr( + const int& scaleout_rank_idx = 0, const int& expert_idx = 0) const { + const auto base_ptr = + get_scaleout_rank_expert_count_ptr() + num_ranks; + return base_ptr + + scaleout_rank_idx * (num_scaleup_ranks * num_experts_per_rank) + + expert_idx; + } + + __forceinline__ __device__ __host__ int64_t* + get_scaleout_channel_signaled_tail_ptr(const int& channel_idx, + const int& scaleout_rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_scaleout_rank_expert_count_ptr(), + (kNumMaxRanks + kNumMaxExperts) * sizeof(int) * 2); + return base_ptr + + (channel_idx * num_scaleout_ranks + scaleout_rank_idx); + } + + __forceinline__ __device__ __host__ int* get_channel_scaleup_tail_ptr( + const int& channel_idx, const int& scaleup_rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_scaleout_channel_signaled_tail_ptr(0, 0), + kNumMaxRanks * kNumMaxChannels * sizeof(int64_t)); + return base_ptr + (channel_idx * num_scaleup_ranks + scaleup_rank_idx); + } + + __forceinline__ __device__ __host__ int64_t* get_pp_send_count_ptr( + const int& offset) const { + const auto base_ptr = math::advance_ptr( + get_channel_scaleup_tail_ptr(0, 0), + kNumMaxRanks * kNumMaxChannels * sizeof(int)); + return base_ptr + offset; + } + + __forceinline__ __device__ __host__ int64_t* get_pp_recv_count_ptr( + const int& offset) const { + const auto base_ptr = math::advance_ptr( + get_pp_send_count_ptr(0), 2 * sizeof(int64_t)); + return base_ptr + offset; + } + + __forceinline__ __device__ __host__ int* get_agrs_recv_signal_ptr( + const int& slot, const int& rank_idx) const { + const auto base_ptr = math::advance_ptr(get_pp_recv_count_ptr(0), + 2 * sizeof(int64_t)); + return base_ptr + slot * kNumMaxRanks + rank_idx; + } + + __forceinline__ __device__ __host__ int* get_agrs_session_signal_ptr( + const int& rank_idx) const { + const auto base_ptr = math::advance_ptr( + get_agrs_recv_signal_ptr(0, 0), + kNumMaxInflightAGRS * kNumMaxRanks * sizeof(int)); + return base_ptr + rank_idx; + } +}; + +struct TokenLayout { + int num_hidden_bytes, num_sf_bytes; + // NOTES: the top-k index is always 32-bit + bool with_metadata; + int num_topk, num_metadata_bytes; + void* base; + + __forceinline__ __device__ __host__ TokenLayout(const int& num_hidden_bytes, + const int& num_sf_bytes, + const int& num_topk, + const bool& with_metadata, + void* base = nullptr) + : num_hidden_bytes(num_hidden_bytes), + num_sf_bytes(num_sf_bytes), + // Metadata includes: top-k indices, weight and source rank/token + // index + with_metadata(with_metadata), + num_topk(num_topk), + num_metadata_bytes( + num_topk * (sizeof(int) + sizeof(float)) + + (with_metadata ? (1 + num_topk) * sizeof(int) : 0)), + base(base) { + EP_STATIC_ASSERT(sizeof(int) == sizeof(float), + "Invalid size assumption"); + EP_UNIFIED_ASSERT(num_hidden_bytes % ptx::kNumTMAAlignBytes == 0); + } + + template + __forceinline__ __device__ __host__ dtype_t get_num_bytes() const { + const auto num_bytes = + math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes) + + math::align(num_sf_bytes, ptx::kNumTMAAlignBytes) + + math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes) + + math::align(kWithMBarrier ? sizeof(ptx::mbarrier) : 0, + ptx::kNumTMAAlignBytes); + return static_cast(num_bytes); + } + + __forceinline__ __device__ __host__ void* get_base_ptr() const { + return base; + } + + __forceinline__ __device__ __host__ void set_base_ptr(void* ptr) { + base = ptr; + } + + __forceinline__ __device__ __host__ void* get_hidden_ptr() const { + return get_base_ptr(); + } + + __forceinline__ __device__ __host__ sf_pack_t* get_sf_ptr() const { + return math::advance_ptr( + base, math::align(num_hidden_bytes, ptx::kNumTMAAlignBytes)); + } + + __forceinline__ __device__ __host__ int* get_metadata_ptr() const { + return math::advance_ptr( + get_sf_ptr(), math::align(num_sf_bytes, ptx::kNumTMAAlignBytes)); + } + + __forceinline__ __device__ __host__ int* get_topk_idx_ptr() const { + return get_metadata_ptr(); + } + + __forceinline__ __device__ __host__ float* get_topk_weights_ptr() const { + return math::advance_ptr(get_metadata_ptr(), + num_topk * sizeof(int)); + } + + __forceinline__ __device__ __host__ int* get_src_token_global_idx_ptr() + const { + return math::advance_ptr(get_topk_weights_ptr(), + num_topk * sizeof(float)); + } + + __forceinline__ __device__ __host__ int* get_linked_list_idx_ptr() const { + return get_src_token_global_idx_ptr() + 1; + } + + __forceinline__ __device__ ptx::mbarrier* get_mbarrier_ptr() const { + return math::advance_ptr( + get_metadata_ptr(), + math::align(num_metadata_bytes, ptx::kNumTMAAlignBytes)); + } +}; + +template +struct BufferLayout { + TokenLayout token_layout; + int num_ranks; + int num_max_tokens_per_rank; + + void* base; + + __forceinline__ __device__ __host__ + BufferLayout(const TokenLayout& token_layout, const int& num_ranks, + const int& max_num_tokens_per_rank, void* base = nullptr) + : token_layout(token_layout), + num_ranks(num_ranks), + num_max_tokens_per_rank(max_num_tokens_per_rank), + base(base) {} + + __forceinline__ __device__ __host__ int64_t + get_num_bytes_per_token() const { + return token_layout.get_num_bytes(); + } + + __forceinline__ __device__ __host__ int64_t get_num_bytes_per_rank() const { + return num_max_tokens_per_rank * get_num_bytes_per_token(); + } + + __forceinline__ __device__ __host__ int64_t get_num_bytes() const { + return get_num_bytes_per_rank() * num_ranks; + } + + __forceinline__ __device__ __host__ void* get_buffer_end_ptr() const { + return math::advance_ptr(base, get_num_bytes()); + } + + __forceinline__ __device__ __host__ BufferLayout + get_rank_buffer(const int& rank_idx) const { + return BufferLayout( + token_layout, 1, num_max_tokens_per_rank, + static_cast(base) + get_num_bytes_per_rank() * rank_idx); + } + + template + __forceinline__ __device__ __host__ BufferLayout + get_channel_buffer(const int& channel_idx) const { + EP_UNIFIED_ASSERT(kNumTokensPerChannel > 0); + return BufferLayout( + token_layout, + // Do not use `num_max_tokens_per_rank / kNumTokensPerChannel` as + // the false stride + num_ranks, num_max_tokens_per_rank, + static_cast(base) + + get_num_bytes_per_token() * kNumTokensPerChannel * channel_idx); + } + + __forceinline__ __device__ __host__ TokenLayout + get_token_buffer(const int& token_idx, const bool& global = false) const { + EP_UNIFIED_ASSERT(num_ranks == 1 or global); + return TokenLayout( + token_layout.num_hidden_bytes, token_layout.num_sf_bytes, + token_layout.num_topk, token_layout.with_metadata, + static_cast(base) + + token_layout.get_num_bytes() * + token_idx); + } +}; + +} // namespace mooncake::elastic::layout diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_math.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_math.cuh new file mode 100644 index 0000000000..89dbcd0fcc --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_math.cuh @@ -0,0 +1,83 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include + +namespace mooncake::elastic::math { + +template +__forceinline__ __device__ __host__ T ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +template +__forceinline__ __device__ __host__ constexpr T constexpr_ceil_div(T a, T b) { + return (a + b - 1) / b; +} + +template +__forceinline__ __device__ __host__ T align(T a, T b) { + return (kDoCeilAlignment ? ceil_div(a, b) : (a / b)) * b; +} + +template +__forceinline__ __device__ __host__ constexpr T constexpr_align(T a, T b) { + return (kDoCeilAlignment ? constexpr_ceil_div(a, b) : (a / b)) * b; +} + +template +__forceinline__ __device__ __host__ bool is_decoded_positive_ready( + const dtype_t& value) { + return value >= 0; +} + +template +__forceinline__ __device__ __host__ dtype_t +encode_decode_positive(const dtype_t& value) { + return -value - static_cast(1); +} + +template +__forceinline__ __device__ __host__ dtype_t* advance_ptr( + void* ptr, const int64_t num_bytes) { + return reinterpret_cast(static_cast(ptr) + num_bytes); +} + +__forceinline__ __device__ __host__ ptrdiff_t ptr_diff(const void* ptr, + const void* base) { + return static_cast(ptr) - static_cast(base); +} + +template +__device__ __forceinline__ dtype_b_t pack2(const dtype_a_t& x, + const dtype_a_t& y) { + EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), + "Invalid dtypes"); + dtype_b_t packed; + auto unpacked_ptr = reinterpret_cast(&packed); + unpacked_ptr[0] = x, unpacked_ptr[1] = y; + return packed; +} + +template +__device__ __forceinline__ std::tuple unpack2( + const dtype_b_t& packed) { + EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), + "Invalid dtypes"); + auto unpacked_ptr = reinterpret_cast(&packed); + dtype_a_t x = unpacked_ptr[0], y = unpacked_ptr[1]; + return {x, y}; +} + +template +__device__ __forceinline__ void unpack2(const dtype_b_t& packed, dtype_a_t& x, + dtype_a_t& y) { + EP_STATIC_ASSERT(sizeof(dtype_a_t) * 2 == sizeof(dtype_b_t), + "Invalid dtypes"); + auto unpacked_ptr = reinterpret_cast(&packed); + x = unpacked_ptr[0], y = unpacked_ptr[1]; +} + +} // namespace mooncake::elastic::math diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_ptx.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_ptx.cuh new file mode 100644 index 0000000000..f1fa20650e --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_ptx.cuh @@ -0,0 +1,735 @@ +// Ported from DeepEP official elastic source. +// Mooncake changes: namespace switched to mooncake::elastic and NCCL GIN +// transport references are replaced with Mooncake Device API adapters. +#pragma once + +#include +#include + +#include +#include + +namespace mooncake::elastic::ptx { + +// Host-side placeholder with the same size/alignment as +// cuda::barrier (a single uint64_t atomic), so that +// sizeof(mbarrier) is consistent across host and device. +struct alignas(8) mbarrier { + uint64_t __placeholder; +}; +using arrival_phase = uint32_t; + +// More than TMA, `longlong4` requires 32 bytes aligned +static constexpr int kNumTMAAlignBytes = 32; + +#ifdef __CUDACC__ + +/// Exceptions +__forceinline__ __device__ void trap() { +#ifdef MOONCAKE_EP_USE_MUSA + return; +#else + asm volatile("trap;"); +#endif +} + +/// Thread layout +__forceinline__ __device__ int get_warp_idx() { + return __shfl_sync(0xffffffff, threadIdx.x / 32, 0); +} + +__forceinline__ __device__ int get_lane_idx() { +#ifdef MOONCAKE_EP_USE_MUSA + return static_cast(threadIdx.x) & 31; +#else + int lane_idx; + asm volatile("mov.s32 %0, %laneid;" : "=r"(lane_idx)); + return lane_idx; +#endif +} + +/// Election +__forceinline__ __device__ int elect_one_sync() { +#if !defined(MOONCAKE_EP_USE_MUSA) && !defined(DISABLE_SM90_FEATURES) && \ + defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + int pred = 0; + asm volatile( + "{\n" + ".reg .b32 %%rx;\n" + ".reg .pred %%px;\n" + " elect.sync %%rx|%%px, %1;\n" + "@%%px mov.s32 %0, 1;\n" + "}\n" + : "+r"(pred) + : "r"(0xffffffff)); + return pred; +#else + return get_lane_idx() == 0; +#endif +} + +/// TMA and `cp.async` +__forceinline__ __device__ void mbarrier_init_with_fence( + mbarrier* ptr, const int& arrive_count = 1) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("mbarrier.init.shared::cta.b64 [%1], %0;" ::"r"(arrive_count), + "r"(static_cast(__cvta_generic_to_shared(ptr)))); + asm volatile("fence.mbarrier_init.release.cluster;" ::); +#endif +} + +__forceinline__ __device__ void mbarrier_invalidate(mbarrier* ptr) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("mbarrier.inval.shared::cta.b64 [%0];" ::"r"( + static_cast(__cvta_generic_to_shared(ptr)))); +#endif +} + +__forceinline__ __device__ void mbarrier_arrive(mbarrier* ptr) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0]; \n\t" ::"r"( + static_cast(__cvta_generic_to_shared(ptr)))); +#endif +} + +__forceinline__ __device__ void mbarrier_arrive_and_set_tx( + mbarrier* ptr, const int& num_bytes) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0; \n\t" ::"r"( + num_bytes), + "r"(static_cast(__cvta_generic_to_shared(ptr)))); +#endif +} + +__forceinline__ __device__ void mbarrier_wait_and_flip_phase( + mbarrier* ptr, arrival_phase& phase) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile( + "{\n\t" + ".reg .pred P1; \n\t" + "LAB_WAIT: \n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1, %2; \n\t" + "@P1 bra DONE; \n\t" + "bra LAB_WAIT; \n\t" + "DONE: \n\t" + "}" ::"r"(static_cast(__cvta_generic_to_shared(ptr))), + "r"(phase), "r"(0x989680)); +#endif + phase ^= 1; +} + +__forceinline__ __device__ void tma_store_fence() { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("fence.proxy.async.shared::cta;"); +#endif +} + +template +__forceinline__ __device__ void tma_store_wait() { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("cp.async.bulk.wait_group %0;" ::"n"(kNumRemainingWaits) + : "memory"); +#endif +} + +enum TMACacheHint : int64_t { + kEvictFirst = 0x12f0000000000000ll, + kEvictNormal = 0x1000000000000000ll +}; + +__forceinline__ __device__ void tma_load_1d( + const void* dst_ptr, const void* src_ptr, mbarrier* ptr, + const int& num_bytes, + const TMACacheHint& hint = TMACacheHint::kEvictFirst) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + // NOTES: normally, the loaded part will be evicted soon + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::" + "cache_hint [%0], [%1], %2, [%3], %4;\n" ::"r"( + static_cast(__cvta_generic_to_shared(dst_ptr))), + "l"(src_ptr), "r"(num_bytes), + "r"(static_cast(__cvta_generic_to_shared(ptr))), "l"(hint) + : "memory"); +#else + const auto dst_addr = reinterpret_cast(dst_ptr); + const auto src_addr = reinterpret_cast(src_ptr); + if (((dst_addr | src_addr | static_cast(num_bytes)) & + (sizeof(int4) - 1)) == 0) { + auto* dst = reinterpret_cast(const_cast(dst_ptr)); + const auto* src = reinterpret_cast(src_ptr); + const int num_vecs = num_bytes / static_cast(sizeof(int4)); + for (int i = 0; i < num_vecs; ++i) dst[i] = src[i]; + } else { + auto* dst = static_cast(const_cast(dst_ptr)); + const auto* src = static_cast(src_ptr); + for (int i = 0; i < num_bytes; ++i) dst[i] = src[i]; + } +#endif +} + +__forceinline__ __device__ void tma_load_1d_warp( + const void* dst_ptr, const void* src_ptr, mbarrier* ptr, + const int& num_bytes, const int& lane_idx, + const TMACacheHint& hint = TMACacheHint::kEvictFirst) { +#ifdef MOONCAKE_EP_USE_MUSA + const auto dst_addr = reinterpret_cast(dst_ptr); + const auto src_addr = reinterpret_cast(src_ptr); + if (((dst_addr | src_addr | static_cast(num_bytes)) & + (sizeof(int4) - 1)) == 0) { + auto* dst = reinterpret_cast(const_cast(dst_ptr)); + const auto* src = reinterpret_cast(src_ptr); + const int num_vecs = num_bytes / static_cast(sizeof(int4)); + for (int i = lane_idx; i < num_vecs; i += 32) dst[i] = src[i]; + } else { + auto* dst = static_cast(const_cast(dst_ptr)); + const auto* src = static_cast(src_ptr); + for (int i = lane_idx; i < num_bytes; i += 32) dst[i] = src[i]; + } +#else + if (elect_one_sync()) tma_load_1d(dst_ptr, src_ptr, ptr, num_bytes, hint); +#endif +} + +__forceinline__ __device__ void tma_store_1d( + const void* dst_ptr, const void* src_ptr, const int& num_bytes, + const TMACacheHint& hint = TMACacheHint::kEvictNormal) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + // NOTES: normally, the stored part will be used soon + asm volatile( + "cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint [%0], " + "[%1], %2, %3;\n" ::"l"(dst_ptr), + "r"(static_cast(__cvta_generic_to_shared(src_ptr))), + "r"(num_bytes), "l"(hint) + : "memory"); +#else + const auto dst_addr = reinterpret_cast(dst_ptr); + const auto src_addr = reinterpret_cast(src_ptr); + if (((dst_addr | src_addr | static_cast(num_bytes)) & + (sizeof(int4) - 1)) == 0) { + auto* dst = reinterpret_cast(const_cast(dst_ptr)); + const auto* src = reinterpret_cast(src_ptr); + const int num_vecs = num_bytes / static_cast(sizeof(int4)); + for (int i = 0; i < num_vecs; ++i) { +#ifdef MOONCAKE_EP_USE_MUSA + const volatile int* src_words = + reinterpret_cast(src + i); + volatile int* dst_words = reinterpret_cast(dst + i); + dst_words[0] = src_words[0]; + dst_words[1] = src_words[1]; + dst_words[2] = src_words[2]; + dst_words[3] = src_words[3]; +#else + dst[i] = src[i]; +#endif + } + } else { +#ifdef MOONCAKE_EP_USE_MUSA + auto* dst = static_cast(const_cast(dst_ptr)); + const auto* src = static_cast(src_ptr); +#else + auto* dst = static_cast(const_cast(dst_ptr)); + const auto* src = static_cast(src_ptr); +#endif + for (int i = 0; i < num_bytes; ++i) dst[i] = src[i]; + } +#endif +} + +__forceinline__ __device__ void tma_store_commit() { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 900) + asm volatile("cp.async.bulk.commit_group;"); +#endif +} + +template +__forceinline__ __device__ void cp_async_ca(const dtype_t* gmem_src, + const dtype_t* smem_dst) { + EP_STATIC_ASSERT( + sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8 or sizeof(dtype_t) == 16, + "Invalid dtype bytes"); +#ifdef MOONCAKE_EP_USE_MUSA + *const_cast(smem_dst) = *gmem_src; +#else + asm volatile( + "cp.async.ca.shared::cta.global.L2::128B [%0], [%1], %2;\n" ::"r"( + static_cast(__cvta_generic_to_shared(smem_dst))), + "l"(gmem_src), "n"(sizeof(dtype_t))); +#endif +} + +__forceinline__ __device__ void cp_async_mbarrier_arrive(mbarrier* ptr) { +#ifdef MOONCAKE_EP_USE_MUSA + (void)ptr; +#else + asm volatile("cp.async.mbarrier.arrive.shared::cta.b64 [%0];\n" ::"r"( + static_cast(__cvta_generic_to_shared(ptr)))); +#endif +} + +/// Barriers +template +__forceinline__ __device__ void named_barrier(const int& idx) { +#ifdef MOONCAKE_EP_USE_MUSA + (void)idx; + __threadfence_block(); + __syncwarp(); +#else + // Equivalent to `barrier.sync.aligned`, which requires all threads run the + // same location of code + asm volatile("bar.sync %0, %1;" ::"r"(idx), "r"(kNumThreads)); +#endif +} + +/// LD/ST instructions +__forceinline__ __device__ int4 +ldg_with_gez_pred(const int4* ptr, const int& value, + const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + int4 ret = make_int4(0, 0, 0, 0); +#ifdef MOONCAKE_EP_USE_MUSA + (void)cache_hint; + if (value >= 0) { + const volatile int* words = reinterpret_cast(ptr); + ret.x = words[0]; + ret.y = words[1]; + ret.z = words[2]; + ret.w = words[3]; + } +#else + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s32 {%0, %1, %2, " + "%3}, [%4], %6;\n\t" + "}" + : "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); +#endif + return ret; +} + +__forceinline__ __device__ int4 +ldg_with_gtz_pred(const int4* ptr, const int& value, + const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + int4 ret = make_int4(0, 0, 0, 0); +#ifdef MOONCAKE_EP_USE_MUSA + (void)cache_hint; + if (value > 0) { + const volatile int* words = reinterpret_cast(ptr); + ret.x = words[0]; + ret.y = words[1]; + ret.z = words[2]; + ret.w = words[3]; + } +#else + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.gt.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s32 {%0, %1, %2, " + "%3}, [%4], %6;\n\t" + "}" + : "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); +#endif + return ret; +} + +__forceinline__ __device__ int4 +ld_with_gez_pred(const int4* ptr, const int& value, + const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + int4 ret = make_int4(0, 0, 0, 0); +#ifdef MOONCAKE_EP_USE_MUSA + (void)cache_hint; + if (value >= 0) { + const volatile int* words = reinterpret_cast(ptr); + ret.x = words[0]; + ret.y = words[1]; + ret.z = words[2]; + ret.w = words[3]; + } +#else + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.v4.s32 {%0, %1, %2, " + "%3}, [%4], %6;\n\t" + "}" + : "+r"(ret.x), "+r"(ret.y), "+r"(ret.z), "+r"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); +#endif + return ret; +} + +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 1000) +__forceinline__ __device__ longlong4_t +ldg_with_gez_pred(const longlong4_t* ptr, const int& value, + const TMACacheHint& cache_hint = TMACacheHint::kEvictFirst) { + longlong4_t ret = make_longlong4_t(0, 0, 0, 0); + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %5, 0;\n\t" + " @p ld.L1::no_allocate.L2::cache_hint.global.nc.v4.s64 {%0, %1, %2, " + "%3}, [%4], %6;\n\t" + "}" + : "+l"(ret.x), "+l"(ret.y), "+l"(ret.z), "+l"(ret.w) + : "l"(ptr), "r"(value), "l"(cache_hint) + : "memory"); + return ret; +} + +__forceinline__ __device__ longlong4_t ldg(const longlong4_t* ptr) { + longlong4_t ret; + asm volatile( + "ld.L1::no_allocate.global.nc.v4.s64 {%0, %1, %2, %3}, [%4];\n\t" + : "=l"(ret.x), "=l"(ret.y), "=l"(ret.z), "=l"(ret.w) + : "l"(ptr) + : "memory"); + return ret; +} +#endif + +__forceinline__ __device__ int4 ldg(const int4* ptr) { +#ifdef MOONCAKE_EP_USE_MUSA + const volatile int* words = reinterpret_cast(ptr); + int4 ret; + ret.x = words[0]; + ret.y = words[1]; + ret.z = words[2]; + ret.w = words[3]; + return ret; +#else + return __ldg(ptr); +#endif +} + +__forceinline__ __device__ void st_na(int4* ptr, const int4& value) { +#ifdef MOONCAKE_EP_USE_MUSA + volatile int* words = reinterpret_cast(ptr); + words[0] = value.x; + words[1] = value.y; + words[2] = value.z; + words[3] = value.w; +#else + *ptr = value; +#endif +} + +template +__forceinline__ __device__ void st_with_gez_pred(dtype_t* ptr, dtype_t value, + const int& condition) { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4, "Invalid data type"); + auto view = *reinterpret_cast(&value); +#ifdef MOONCAKE_EP_USE_MUSA + if (condition >= 0) *ptr = value; +#else + asm volatile( + "{\n\t" + " .reg .pred p;\n\t" + " setp.ge.s32 p, %2, 0;\n\t" + " @p st.global.s32 [%0], %1;\n\t" + "}" ::"l"(ptr), + "r"(view), "r"(condition) + : "memory"); +#endif +} + +template +__forceinline__ __device__ dtype_t ld_volatile(const void* ptr) { +#ifdef MOONCAKE_EP_USE_MUSA + const volatile dtype_t* typed = + reinterpret_cast(ptr); + return *typed; +#else + if constexpr (sizeof(dtype_t) == 4) { + uint32_t value; + asm volatile("ld.volatile.global.u32 %0, [%1];" + : "=r"(value) + : "l"(ptr)); + return reinterpret_cast(value); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t value; + asm volatile("ld.volatile.global.u64 %0, [%1];" + : "=l"(value) + : "l"(ptr)); + return reinterpret_cast(value); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, + "Invalid data type length"); + } +#endif +} + +__forceinline__ __device__ void red_add(const int64_t* ptr, + const int64_t& value) { +#ifdef MOONCAKE_EP_USE_MUSA + atomicAdd(const_cast( + reinterpret_cast(ptr)), + static_cast(value)); + __threadfence_system(); +#else + // TODO(NVCC): why don't NVCC support `s64`? + // Mooncake elastic consumers can poll these counters from a different CTA + // immediately after the RED producer issues the update. Keep the update as + // a single 64-bit RED so packed counters stay atomic, but use release scope + // instead of a relaxed GPU-scope RED to avoid occasional visibility stalls + // in the notify reduction path. + asm volatile("red.release.gpu.global.add.u64 [%0], %1;" ::"l"(ptr), + "l"(value) + : "memory"); +#endif +} + +__forceinline__ __device__ void red_add_rel_sys(const int* ptr, + const int& value) { +#ifdef MOONCAKE_EP_USE_MUSA + atomicAdd(const_cast(ptr), value); + __threadfence_system(); +#else + asm volatile("red.release.sys.global.add.s32 [%0], %1;" ::"l"(ptr), + "r"(value)); +#endif +} + +__forceinline__ __device__ void red_add_rel_sys(const int64_t* ptr, + const int64_t& value) { +#ifdef MOONCAKE_EP_USE_MUSA + atomicAdd(const_cast( + reinterpret_cast(ptr)), + static_cast(value)); + __threadfence_system(); +#else + asm volatile("red.release.sys.global.add.u64 [%0], %1;" ::"l"(ptr), + "l"(value)); +#endif +} + +template +__forceinline__ __device__ dtype_t ld_acquire_sys(const dtype_t* ptr) { +#ifdef MOONCAKE_EP_USE_MUSA + const volatile dtype_t* typed = + reinterpret_cast(ptr); + const dtype_t value = *typed; + __threadfence_system(); + return value; +#else + if constexpr (sizeof(dtype_t) == 4) { + uint32_t value; + asm volatile("ld.acquire.sys.L1::no_allocate.global.u32 %0, [%1];" + : "=r"(value) + : "l"(ptr)); + return reinterpret_cast(value); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t value; + asm volatile("ld.acquire.sys.L1::no_allocate.global.u64 %0, [%1];" + : "=l"(value) + : "l"(ptr)); + return reinterpret_cast(value); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, + "Invalid data type length"); + } +#endif +} + +template +__forceinline__ __device__ void st_relaxed_sys(void* ptr, dtype_t value) { +#ifdef MOONCAKE_EP_USE_MUSA + *static_cast(ptr) = value; +#else + if constexpr (sizeof(dtype_t) == 4) { + uint32_t int_value = reinterpret_cast(value); + asm volatile("st.relaxed.sys.global.u32 [%0], %1;" ::"l"(ptr), + "r"(int_value)); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t int_value = reinterpret_cast(value); + asm volatile("st.relaxed.sys.global.u64 [%0], %1;" ::"l"(ptr), + "l"(int_value)); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, + "Invalid data type length"); + } +#endif +} + +template +__forceinline__ __device__ void st_release_sys(void* ptr, dtype_t value) { +#ifdef MOONCAKE_EP_USE_MUSA + __threadfence_system(); + *static_cast(ptr) = value; + __threadfence_system(); +#else + if constexpr (sizeof(dtype_t) == 4) { + uint32_t int_value = reinterpret_cast(value); + asm volatile("st.release.sys.global.u32 [%0], %1;" ::"l"(ptr), + "r"(int_value)); + } else if constexpr (sizeof(dtype_t) == 8) { + uint64_t int_value = reinterpret_cast(value); + asm volatile("st.release.sys.global.u64 [%0], %1;" ::"l"(ptr), + "l"(int_value)); + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 4 or sizeof(dtype_t) == 8, + "Invalid data type length"); + } +#endif +} + +// Adjust registers +template +__device__ __forceinline__ void warpgroup_reg_alloc() { +#ifndef MOONCAKE_EP_USE_MUSA + asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" : : "n"(kNumRegs)); +#endif +} + +template +__device__ __forceinline__ void warpgroup_reg_dealloc() { +#ifndef MOONCAKE_EP_USE_MUSA + asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" : : "n"(kNumRegs)); +#endif +} + +/// General fences +__device__ __forceinline__ void fence_acq_rel_sys() { +#ifdef MOONCAKE_EP_USE_MUSA + __threadfence_system(); +#else + asm volatile("fence.acq_rel.sys;" ::: "memory"); +#endif +} + +/// Intrinsics +template +__device__ __forceinline__ dtype_t exchange(dtype_t ptr, + const int& src_lane_idx) { + EP_STATIC_ASSERT(sizeof(dtype_t) % sizeof(int) == 0, ""); + const auto send_int_values = reinterpret_cast(&ptr); + dtype_t recv_dtype; + auto recv_int_values = reinterpret_cast(&recv_dtype); +#pragma unroll + for (int i = 0; i < sizeof(dtype_t) / sizeof(int); ++i) + recv_int_values[i] = + __shfl_sync(0xffffffff, send_int_values[i], src_lane_idx); + return recv_dtype; +} + +__device__ __forceinline__ unsigned gather(const bool& value) { + return __ballot_sync(0xffffffff, value); +} + +__device__ __forceinline__ bool all(const bool& value) { + return __all_sync(0xffffffff, value); +} + +__device__ __forceinline__ bool any(const bool& value) { + return __any_sync(0xffffffff, value); +} + +__device__ __forceinline__ unsigned reduce_or(const unsigned& value) { + return __reduce_or_sync(0xffffffff, value); +} + +__device__ __forceinline__ unsigned long long reduce_or( + const unsigned long long& value) { + const auto low = __reduce_or_sync(0xffffffff, static_cast(value)); + const auto high = + __reduce_or_sync(0xffffffff, static_cast(value >> 32)); + return (static_cast(high) << 32) | low; +} + +__device__ __forceinline__ int reduce_add(const int& value) { + return __reduce_add_sync(0xffffffff, value); +} + +__device__ __forceinline__ unsigned match(const int& value) { + return __match_any_sync(0xffffffff, value); +} + +__device__ __forceinline__ int fns(const unsigned& value, const int& offset) { + return __fns(value, 0, offset); +} + +template +__device__ __forceinline__ auto ffs(const dtype_t& value) { + if constexpr (sizeof(dtype_t) == 4) { + return __ffs(static_cast(value)) - 1; + } else { + EP_STATIC_ASSERT(sizeof(dtype_t) == 8, "Invalid data type"); + return __ffsll(static_cast(value)) - 1; + } +} + +__device__ __forceinline__ int get_master_lane_idx(const unsigned& mask) { +#ifdef MOONCAKE_EP_USE_MUSA + return 31 - __clz(mask); +#else + // Equivalent to `31 - __clz(mask)` + int highest_idx; + asm volatile("bfind.u32 %0, %1;" : "=r"(highest_idx) : "r"(mask)); + return highest_idx; +#endif +} + +__device__ __forceinline__ bool deduplicate(const int& value, + const int& lane_idx) { + return get_master_lane_idx(match(value)) == lane_idx; +} + +__device__ __forceinline__ int warp_inclusive_sum(int value, + const int& lane_idx) { +#pragma unroll + for (int offset = 1; offset < 32; offset <<= 1) { + const auto synced = __shfl_up_sync(0xffffffff, value, offset); + if (lane_idx >= offset) value += synced; + } + return value; +} + +__device__ __forceinline__ float2 fadd2(const float2& a, const float2& b) { +#if !defined(MOONCAKE_EP_USE_MUSA) && defined(__CUDA_ARCH__) && \ + (__CUDA_ARCH__ >= 1000) + return __fadd2_rn(a, b); +#else + return {a.x + b.x, a.y + b.y}; +#endif +} + +__device__ __forceinline__ void accumulate(float2& a, nv_bfloat162 b) { +#ifdef MOONCAKE_EP_USE_MUSA + a.x += __low2float(b); + a.y += __high2float(b); +#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + // Use `add.rn.f32.bf16` instruction to perform fused (cast + add) operation + // on SM100 + asm("add.rn.f32.bf16 %0, %1, %0;\n" + : "+f"(a.x) + : "h"(*reinterpret_cast(&b.x))); + asm("add.rn.f32.bf16 %0, %1, %0;\n" + : "+f"(a.y) + : "h"(*reinterpret_cast(&b.y))); +#else + const auto [x, y] = __bfloat1622float2(b); + a.x += x, a.y += y; +#endif +} + +#endif + +} // namespace mooncake::elastic::ptx diff --git a/mooncake-ep/include/elastic/mooncake_ep_elastic_transport.cuh b/mooncake-ep/include/elastic/mooncake_ep_elastic_transport.cuh new file mode 100644 index 0000000000..7855ce3a63 --- /dev/null +++ b/mooncake-ep/include/elastic/mooncake_ep_elastic_transport.cuh @@ -0,0 +1,261 @@ +#pragma once + +#include +#include + +#include +#include +#include + +namespace mooncake::elastic::transport { + +struct WorldTeam {}; +struct ScaleupTeam {}; +struct ScaleoutTeam {}; + +constexpr int kRedAddReleaseHighWordLast = 0; +constexpr int kRedAddReleaseLowWordLast = 1 << 0; + +// Mooncake Device API adapter for DeepEP's NCCL GIN usage. +// +// DeepEP elastic kernels express all remote communication through a small GIN +// surface: symmetric-pointer translation, put, put_value, RED/add style signals +// and QP flushes. Mooncake maps that surface onto Device API semantics: +// +// get_sym_ptr -> mc_route_put, returning local/P2P peer VA or nullptr +// put -> local/P2P warp copy, otherwise mc_rdma_put +// put_value -> local/P2P release store, otherwise mc_rdma_put/mc_signal +// flush -> no-op; Device API operations are ordered by release/fence +// and +// explicit kernel barriers in the imported elastic kernels +// +// Team tags are kept as types so official DeepEP template code can remain close +// to the source while the actual routing is decided by Mooncake CommCtx. +struct MooncakeGin { + device::CommCtx ctx; + int qp_idx = 0; + int sharing_mode = 0; + int qps_per_rank = 1; + int scaleout_rank_idx = 0; + int scaleup_rank_idx = 0; + int num_scaleup_ranks = 0; + + __device__ __forceinline__ MooncakeGin( + const device::CommCtx& ctx, int qp_idx, int sharing_mode, int num_qps, + int scaleout_rank_idx = 0, int scaleup_rank_idx = 0, + int num_scaleup_ranks = 0, int num_ranks = 1) + : ctx(ctx), + qp_idx(qp_idx), + sharing_mode(sharing_mode), + qps_per_rank(max(1, num_qps / max(1, num_ranks))), + scaleout_rank_idx(scaleout_rank_idx), + scaleup_rank_idx(scaleup_rank_idx), + num_scaleup_ranks(num_scaleup_ranks) {} + + template + __device__ __forceinline__ int world_rank(int dst_rank) const { + if (num_scaleup_ranks <= 0) return dst_rank; + if constexpr (std::is_same_v) { + return scaleout_rank_idx * num_scaleup_ranks + dst_rank; + } else if constexpr (std::is_same_v) { + return dst_rank * num_scaleup_ranks + scaleup_rank_idx; + } else { + return dst_rank; + } + } + + template + __device__ __forceinline__ bool is_nvlink_accessible(int dst_rank) const { + dst_rank = world_rank(dst_rank); + return dst_rank == ctx.rank || + device::mc_comm_p2p_available(ctx, dst_rank); + } + + template + __device__ __forceinline__ void* get_sym_ptr(void* ptr, + int dst_rank) const { + dst_rank = world_rank(dst_rank); + return device::mc_route_put(ctx, dst_rank, ptr); + } + + template + __device__ __forceinline__ const void* get_sym_ptr(const void* ptr, + int dst_rank) const { + dst_rank = world_rank(dst_rank); + return device::mc_route_put(ctx, dst_rank, const_cast(ptr)); + } + + template + __device__ __forceinline__ void put(void* dst_ptr, const void* src_ptr, + int num_bytes, int dst_rank, + int /*flags*/ = 0) const { + dst_rank = world_rank(dst_rank); + auto routed = device::mc_route_put(ctx, dst_rank, dst_ptr); + if (routed != nullptr) { + const auto src_addr = reinterpret_cast(src_ptr); + const auto dst_addr = reinterpret_cast(routed); + if (((src_addr | dst_addr | static_cast(num_bytes)) & + (sizeof(int4) - 1)) == 0) { + const auto* src = reinterpret_cast(src_ptr); + auto* dst = reinterpret_cast(routed); + const int num_int4 = num_bytes / static_cast(sizeof(int4)); + for (int i = 0; i < num_int4; ++i) { +#ifdef MOONCAKE_EP_USE_MUSA + ptx::st_na(dst + i, device::mc_ld_nc(src + i)); +#else + dst[i] = device::mc_ld_nc(src + i); +#endif + } + } else { + auto* dst_bytes = reinterpret_cast(routed); + const auto* src_bytes = + reinterpret_cast(src_ptr); + for (int i = 0; i < num_bytes; ++i) { +#ifdef MOONCAKE_EP_USE_MUSA + reinterpret_cast(dst_bytes)[i] = + reinterpret_cast(src_bytes)[i]; +#else + dst_bytes[i] = src_bytes[i]; +#endif + } + } + // `put` is used both by full data-moving warps and by individual + // notify lanes. Do not place a full-warp barrier inside the + // transport primitive: divergent notify calls would deadlock. Each + // participating lane copies the complete payload for its request, + // so a system fence is sufficient to publish the writes. + __threadfence_system(); + } else { + device::mc_rdma_put(ctx, qp_idx, dst_rank, qps_per_rank, src_ptr, + dst_ptr, static_cast(num_bytes), 0); + } + } + + template + __device__ __forceinline__ void put_value(value_t* dst_ptr, value_t value, + int dst_rank, + int flags = 0) const { + dst_rank = world_rank(dst_rank); + auto* routed = + static_cast(device::mc_route_put(ctx, dst_rank, dst_ptr)); + if (routed != nullptr) { + if constexpr (sizeof(value_t) == sizeof(int32_t)) { + device::mc_st_release(reinterpret_cast(routed), + static_cast(value)); + } else { + *routed = value; + __threadfence_system(); + } + } else { + if constexpr (sizeof(value_t) == sizeof(int32_t)) { + device::mc_signal(ctx, dst_rank, qp_idx, qps_per_rank, + reinterpret_cast(dst_ptr), + static_cast(value)); + } else { + // Device RDMA WRITE sources must be registered GDR addresses; + // a by-value scalar lives in thread-local storage and is not a + // valid IBGDA source. Current elastic uses remote int64 + // put_value only for single-writer, zeroed notify slots, so a + // split 32-bit RED add is equivalent to writing the packed + // word. + auto* words = reinterpret_cast(dst_ptr); + const auto signed_value = static_cast(value); + const auto low = static_cast( + static_cast(signed_value) & 0xffffffffull); + const auto high = static_cast(signed_value >> 32); + if ((flags & kRedAddReleaseLowWordLast) == 0) { + if (low != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words, low); + } + if (high != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words + 1, high); + } + } else { + if (high != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words + 1, high); + } + if (low != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words, low); + } + } + } + } + } + + template + __device__ __forceinline__ void red_add_rel(value_t* dst_ptr, value_t value, + int dst_rank, + int flags = 0) const { + if constexpr (sizeof(value_t) == sizeof(int32_t)) { + dst_rank = world_rank(dst_rank); + auto* routed = + static_cast(device::mc_route_put(ctx, dst_rank, dst_ptr)); + if (routed != nullptr) { + device::mc_atomic_add_release(routed, static_cast(value)); + } else { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + reinterpret_cast(dst_ptr), + static_cast(value)); + } + } else if constexpr (sizeof(value_t) == sizeof(uint64_t) || + sizeof(value_t) == sizeof(int64_t)) { + dst_rank = world_rank(dst_rank); + auto* routed = static_cast( + device::mc_route_put(ctx, dst_rank, dst_ptr)); + if (routed != nullptr) { + // Some official elastic paths use the high 32 bits as the + // readiness word (notify counters), while others use the low 32 + // bits as the terminal flag (hybrid channel tails). Splitting + // a 64-bit RED into two 32-bit atomics can therefore publish + // the wrong half first for one of the protocols. Use one + // system- scope 64-bit RED on the routed local/P2P VA so the + // packed value is updated atomically with release ordering. + ptx::red_add_rel_sys(routed, static_cast(value)); + } else { + // Mooncake's current Device API only exposes 32-bit remote + // reduction. Do not emulate the 64-bit add with an RDMA WRITE + // from a thread-local scalar: IBGDA WQEs use the registered GDR + // buffer lkey, so a stack/local address is not a valid DMA + // source on true cross-node runs. Split the packed signal into + // two 32-bit remote reductions instead, publishing the + // readiness word last. Most notify counters use high word as + // the ready count; hybrid channel tails use low word as the + // finish flag. + auto* words = reinterpret_cast(dst_ptr); + const auto signed_value = static_cast(value); + const auto low = static_cast( + static_cast(signed_value) & 0xffffffffull); + const auto high = static_cast(signed_value >> 32); + if ((flags & kRedAddReleaseLowWordLast) == 0) { + if (low != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words, low); + } + if (high != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words + 1, high); + } + } else { + if (high != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words + 1, high); + } + if (low != 0) { + device::mc_red_add(ctx, dst_rank, qp_idx, qps_per_rank, + words, low); + } + } + } + } else { + put_value(dst_ptr, value, dst_rank, flags); + } + } + + __device__ __forceinline__ void flush() const { __threadfence_system(); } +}; + +} // namespace mooncake::elastic::transport diff --git a/mooncake-ep/include/mooncake_ep_api.cuh b/mooncake-ep/include/mooncake_ep_api.cuh index 2ab1c632a4..1a560d1b90 100644 --- a/mooncake-ep/include/mooncake_ep_api.cuh +++ b/mooncake-ep/include/mooncake_ep_api.cuh @@ -16,7 +16,7 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales, int hidden, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, bool use_fp8, void* workspace, cudaStream_t stream, int64_t timeout_ticks, - int phases); + int phases, int active_qps_per_rank); void mark_phase_ack(void* mxa_buffer, const int32_t* nvlink_available, void* const* ipc_peer_ptrs, int* ack_buffer, int rank, @@ -42,6 +42,6 @@ void combine(void* combined_x, int32_t* active_ranks, void* mxa_buffer, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, void* workspace, cudaStream_t stream, int64_t timeout_ticks, int phases, - bool zero_copy); + bool zero_copy, int active_qps_per_rank); } // namespace mooncake diff --git a/mooncake-ep/include/mooncake_ep_buffer.h b/mooncake-ep/include/mooncake_ep_buffer.h index 6088ce20af..9307936c80 100644 --- a/mooncake-ep/include/mooncake_ep_buffer.h +++ b/mooncake-ep/include/mooncake_ep_buffer.h @@ -16,6 +16,7 @@ namespace mooncake { class TransferEngine; +class MooncakeElasticBuffer; // MAX_QP_COUNT is defined in mooncake_ep_configs.cuh (shared with kernel code). @@ -64,6 +65,8 @@ struct BufferPair { struct MooncakeEpBuffer { private: + friend class MooncakeElasticBuffer; + // Device info and communication int device_id; int rank, num_ranks; @@ -89,6 +92,11 @@ struct MooncakeEpBuffer { bool ibgda_disabled_ = false; int USE_QP_COUNT = MAX_QP_COUNT; + // Cap on active RoCE QPs per peer: spreading small EP messages across too + // many QP/doorbell/progress streams hurts when GPUs share an HCA. Default + // 8; override at runtime with MOONCAKE_EP_ACTIVE_QPS_PER_RANK (>= per-rank + // QP count disables). + int active_qps_cap_ = 8; // Stream for communication at::cuda::CUDAStream comm_stream; diff --git a/mooncake-ep/include/mooncake_ep_configs.cuh b/mooncake-ep/include/mooncake_ep_configs.cuh index 103f6b8f09..1e7f0c2149 100644 --- a/mooncake-ep/include/mooncake_ep_configs.cuh +++ b/mooncake-ep/include/mooncake_ep_configs.cuh @@ -40,16 +40,22 @@ #endif #include +#ifndef MOONCAKE_EP_USE_MACA #include -#include #include +#endif +#include + +#if defined(MOONCAKE_EP_USE_MUSA) || defined(MOONCAKE_EP_USE_MACA) +#define MOONCAKE_EP_SPLIT_SEND_RECV 1 +#endif // torchada maps nv_bfloat16 → __mt_bfloat16 which is an incomplete type on // MUSA, so sizeof(__mt_bfloat16) fails. mt_bfloat16 (the complete typedef in // musa_bf16.hpp) requires the MUSA device compiler (mcc) and cannot be // included from host .cpp files. Use EP_BF16_SIZE: sizeof(nv_bfloat16) on // CUDA, hardcoded 2 on MUSA (both are 2 bytes). -#ifdef MOONCAKE_EP_USE_MUSA +#if defined(MOONCAKE_EP_USE_MUSA) || defined(MOONCAKE_EP_USE_MACA) #define EP_BF16_SIZE 2 #else #define EP_BF16_SIZE sizeof(nv_bfloat16) diff --git a/mooncake-ep/include/mooncake_ep_device.h b/mooncake-ep/include/mooncake_ep_device.h index 244751bec0..e88417ff7f 100644 --- a/mooncake-ep/include/mooncake_ep_device.h +++ b/mooncake-ep/include/mooncake_ep_device.h @@ -20,8 +20,13 @@ __device__ __forceinline__ ep_fp8x2_storage_t ep_cvt_float2_to_fp8x2(float2 x) { #endif // -- Device intrinsics (MUSA doesn't have __ldg / __activemask) -------------- -#ifndef __ldg -#define __ldg(ptr) (*(ptr)) +#if (defined(__CUDACC__) || defined(__MCC__)) && \ + !defined(MOONCAKE_EP_MUSA_LDG_DEFINED) +#define MOONCAKE_EP_MUSA_LDG_DEFINED +template +__device__ __forceinline__ dtype_t __ldg(const dtype_t* ptr) { + return *ptr; +} #endif #ifndef __activemask #define __activemask() (0xffffffff) @@ -49,7 +54,49 @@ __forceinline__ __device__ int get_lane_id() { return threadIdx.x % 32; } } \ } -#else // !MOONCAKE_EP_USE_MUSA +#elif defined(MOONCAKE_EP_USE_MACA) + +// -- FP8 types --------------------------------------------------------------- +// MetaX C500 does not support the FP8 EP path. The dispatch template is +// still compiled because the host selects the kernel through a runtime bool, so +// provide storage stubs and reject actual FP8 use in the host/Python wrappers. +#include +using ep_fp8_storage_t = uint8_t; +using ep_fp8x2_storage_t = uint16_t; +#if defined(__CUDACC__) || defined(__MCC__) +__device__ __forceinline__ ep_fp8x2_storage_t ep_cvt_float2_to_fp8x2(float2) { + return 0; +} +#endif + +// -- Device intrinsics ------------------------------------------------------- +#ifndef __activemask +#define __activemask() (0xffffffff) +#endif + +#if defined(__CUDACC__) || defined(__MCC__) +__forceinline__ __device__ int get_lane_id() { return threadIdx.x % 32; } +#endif + +// -- Kernel launch (MACA: no cooperative launch) ----------------------------- +#define EP_LAUNCH_BOUNDS(max_threads, min_blocks) + +#define SETUP_LAUNCH_CONFIG(num_sms, num_threads, stream) \ + dim3 _grid(num_sms); \ + dim3 _block(num_threads); \ + cudaStream_t _stream = stream + +#define LAUNCH_KERNEL(config, kernel, ...) \ + kernel<<<_grid, _block, 0, _stream>>>(__VA_ARGS__); \ + { \ + auto _err = cudaGetLastError(); \ + if (_err != cudaSuccess) { \ + fprintf(stderr, "[EP] kernel launch failed: %s\n", \ + cudaGetErrorString(_err)); \ + } \ + } + +#else // !MOONCAKE_EP_USE_MUSA && !MOONCAKE_EP_USE_MACA // -- FP8 types (CUDA native names) ------------------------------------------- #include @@ -86,7 +133,9 @@ __forceinline__ __device__ int get_lane_id() { #define LAUNCH_KERNEL(config, kernel, ...) \ CUDA_CHECK(cudaLaunchKernelEx(config, kernel, ##__VA_ARGS__)) -#endif // MOONCAKE_EP_USE_MUSA +#endif // MOONCAKE_EP_USE_MUSA / MOONCAKE_EP_USE_MACA // Both platforms need IB verbs +#ifndef MOONCAKE_EP_USE_MACA #include +#endif diff --git a/mooncake-ep/include/mooncake_ep_event.h b/mooncake-ep/include/mooncake_ep_event.h index 6dc6eb63fc..4809704983 100644 --- a/mooncake-ep/include/mooncake_ep_event.h +++ b/mooncake-ep/include/mooncake_ep_event.h @@ -25,6 +25,8 @@ struct EventHandle { void current_stream_wait() const { at::cuda::getCurrentCUDAStream().unwrap().wait(*event); } + + void synchronize() const { event->synchronize(); } }; inline torch::Event create_event(const at::cuda::CUDAStream& s) { diff --git a/mooncake-ep/include/mooncake_ep_exception.cuh b/mooncake-ep/include/mooncake_ep_exception.cuh index 060744594a..64061dcd9e 100644 --- a/mooncake-ep/include/mooncake_ep_exception.cuh +++ b/mooncake-ep/include/mooncake_ep_exception.cuh @@ -7,6 +7,8 @@ #define EP_STATIC_ASSERT(cond, reason) static_assert(cond, reason) #endif +#ifndef MOONCAKE_EP_EXCEPTION_CLASS_DEFINED +#define MOONCAKE_EP_EXCEPTION_CLASS_DEFINED class EPException : public std::exception { private: std::string message = {}; @@ -20,6 +22,7 @@ class EPException : public std::exception { const char* what() const noexcept override { return message.c_str(); } }; +#endif #ifndef CUDA_CHECK #define CUDA_CHECK(cmd) \ @@ -42,7 +45,7 @@ class EPException : public std::exception { #endif #ifndef EP_DEVICE_ASSERT -#ifdef MOONCAKE_EP_USE_MUSA +#if defined(MOONCAKE_EP_USE_MUSA) || defined(MOONCAKE_EP_USE_MACA) // MUSA SDK 4.3.x can turn kernels that merely contain a device-side __trap() // branch into illegal memory accesses, even when the assertion condition is // true. Keep these invariants as host/static checks on MUSA builds. diff --git a/mooncake-ep/include/mooncake_ep_utils.cuh b/mooncake-ep/include/mooncake_ep_utils.cuh index 90492f67ed..7b6b194c2b 100644 --- a/mooncake-ep/include/mooncake_ep_utils.cuh +++ b/mooncake-ep/include/mooncake_ep_utils.cuh @@ -48,7 +48,7 @@ struct VecInt<16> { }; // ---- TMA / mbarrier helpers (CUDA only) ---- -#ifndef MOONCAKE_EP_USE_MUSA +#if !defined(MOONCAKE_EP_USE_MUSA) && !defined(MOONCAKE_EP_USE_MACA) __device__ __forceinline__ void fence_view_async_shared() { asm volatile("fence.proxy.async.shared::cta; \n" ::); @@ -136,7 +136,7 @@ __device__ __forceinline__ void tma_store_wait() { asm volatile("cp.async.bulk.wait_group.read %0;" ::"n"(N) : "memory"); } -#endif // MOONCAKE_EP_USE_MUSA +#endif // !MOONCAKE_EP_USE_MUSA && !MOONCAKE_EP_USE_MACA template __host__ __device__ dtype_t cell_div(dtype_t a, dtype_t b) { diff --git a/mooncake-ep/setup.py b/mooncake-ep/setup.py index 0345c54865..2a4f388e1c 100644 --- a/mooncake-ep/setup.py +++ b/mooncake-ep/setup.py @@ -5,9 +5,15 @@ import torch use_musa = os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in {"1", "ON", "TRUE", "YES"} +use_maca = ( + os.getenv("MOONCAKE_EP_USE_MACA", "").upper() in {"1", "ON", "TRUE", "YES"} + or (hasattr(torch.version, "maca") and torch.version.maca is not None) +) if use_musa: try: - import torchada # noqa: F401 + import importlib + + importlib.import_module("torchada") except ImportError as e: raise ImportError( "torchada is required to build the MUSA EP extension. " @@ -28,16 +34,40 @@ abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI) current_dir = os.path.abspath(os.path.dirname(__file__)) +repo_dir = os.path.abspath(os.path.join(current_dir, os.pardir)) +sysroot_dir = os.path.join(repo_dir, ".deps", "sysroot", "usr") + + +def existing_dirs(*paths): + return [path for path in paths if os.path.isdir(path)] + + +sysroot_include_dirs = existing_dirs( + os.path.join(sysroot_dir, "include"), + os.path.join(sysroot_dir, "include", "jsoncpp"), + os.path.join(sysroot_dir, "include", "libnl3"), +) +sysroot_library_dirs = existing_dirs( + os.path.join(sysroot_dir, "lib", "x86_64-linux-gnu"), + os.path.join(sysroot_dir, "lib"), +) abi_define = f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}" cxx_args = [abi_define, "-std=c++20", "-O3", "-g0"] cuda_libraries = ["ibverbs", "mlx5"] cuda_library_dirs = [] +include_dirs = [ + os.path.join(current_dir, "include"), + os.path.join(current_dir, "../mooncake-transfer-engine/include"), +] if use_musa: cuda_libraries = [] - musa_defines = ["-DUSE_MUSA", "-DMOONCAKE_EP_USE_MUSA=1"] + musa_defines = [ + "-DUSE_MUSA", + "-DMOONCAKE_EP_USE_MUSA=1", + ] cxx_args += musa_defines # torchada maps the "nvcc" key to "mcc". device_args = [ @@ -48,6 +78,18 @@ "--cuda-gpu-arch=mp_31", "-O3", ] +elif use_maca: + cuda_libraries = [] + cuda_library_dirs = sysroot_library_dirs.copy() + include_dirs += sysroot_include_dirs + maca_defines = ["-DUSE_MACA", "-DMOONCAKE_EP_USE_MACA=1"] + cxx_args += maca_defines + device_args = [ + abi_define, + *maca_defines, + "-std=c++20", + "-O3", + ] else: cxx_args.append("-DUSE_CUDA") device_args = [ @@ -72,14 +114,13 @@ ext_modules=[ CUDAExtension( name=module_name, - include_dirs=[ - os.path.join(current_dir, "include"), - os.path.join(current_dir, "../mooncake-transfer-engine/include"), - ], + include_dirs=include_dirs, sources=[ "src/ep_py.cpp", "src/mooncake_ep_buffer.cpp", + "src/mooncake_ep_elastic_buffer.cpp", "src/mooncake_ep_kernel.cu", + "src/mooncake_ep_elastic_kernel.cu", ], extra_compile_args={"cxx": cxx_args, "nvcc": device_args}, libraries=cuda_libraries, diff --git a/mooncake-ep/src/CMakeLists.txt b/mooncake-ep/src/CMakeLists.txt index 574ab514c0..a102f5011e 100644 --- a/mooncake-ep/src/CMakeLists.txt +++ b/mooncake-ep/src/CMakeLists.txt @@ -1,4 +1,4 @@ -add_library(mooncake_ep ep_py.cpp mooncake_ep_buffer.cpp mooncake_ep_kernel.cu) +add_library(mooncake_ep ep_py.cpp mooncake_ep_buffer.cpp mooncake_ep_elastic_buffer.cpp mooncake_ep_kernel.cu mooncake_ep_elastic_kernel.cu) set_target_properties(mooncake_ep PROPERTIES POSITION_INDEPENDENT_CODE ON) target_link_libraries(mooncake_ep PUBLIC ${TORCH_LIBRARIES} transfer_engine ibverbs mlx5) diff --git a/mooncake-ep/src/ep_py.cpp b/mooncake-ep/src/ep_py.cpp index bd3cf53a99..02307c0caf 100644 --- a/mooncake-ep/src/ep_py.cpp +++ b/mooncake-ep/src/ep_py.cpp @@ -1,4 +1,5 @@ #include +#include #include #include #include @@ -13,10 +14,56 @@ namespace mooncake { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("get_ep_buffer_size_hint", &get_ep_buffer_size_hint); + m.def("calculate_elastic_buffer_size", + &MooncakeElasticBuffer::calculate_buffer_size); py::class_(m, "EventHandle") .def(py::init<>()) - .def("current_stream_wait", &EventHandle::current_stream_wait); + .def("current_stream_wait", &EventHandle::current_stream_wait) + .def("synchronize", &EventHandle::synchronize); + + py::class_(m, "ElasticNativeHandle") + .def(py::init<>()) + .def_readwrite("do_expand", &ElasticNativeHandle::do_expand) + .def_readwrite("num_experts", &ElasticNativeHandle::num_experts) + .def_readwrite("expert_alignment", + &ElasticNativeHandle::expert_alignment) + .def_readwrite("num_max_tokens_per_rank", + &ElasticNativeHandle::num_max_tokens_per_rank) + .def_readwrite("num_sms", &ElasticNativeHandle::num_sms) + .def_readwrite("topk_idx", &ElasticNativeHandle::topk_idx) + .def_readwrite( + "psum_num_recv_tokens_per_scaleup_rank", + &ElasticNativeHandle::psum_num_recv_tokens_per_scaleup_rank) + .def_readwrite("psum_num_recv_tokens_per_expert", + &ElasticNativeHandle::psum_num_recv_tokens_per_expert) + .def_readwrite("recv_src_metadata", + &ElasticNativeHandle::recv_src_metadata) + .def_readwrite("recv_layout_range", + &ElasticNativeHandle::recv_layout_range) + .def_readwrite("dst_buffer_slot_idx", + &ElasticNativeHandle::dst_buffer_slot_idx) + .def_readwrite("token_metadata_at_forward", + &ElasticNativeHandle::token_metadata_at_forward) + .def_readwrite("channel_linked_list", + &ElasticNativeHandle::channel_linked_list) + .def_readwrite("num_recv_tokens_per_expert_list", + &ElasticNativeHandle::num_recv_tokens_per_expert_list); + + py::class_(m, "ElasticDispatchOutput") + .def_readonly("recv_x", &ElasticDispatchOutput::recv_x) + .def_readonly("recv_x_scales", &ElasticDispatchOutput::recv_x_scales) + .def_readonly("recv_topk_idx", &ElasticDispatchOutput::recv_topk_idx) + .def_readonly("recv_topk_weights", + &ElasticDispatchOutput::recv_topk_weights) + .def_readonly("handle", &ElasticDispatchOutput::handle) + .def_readonly("event", &ElasticDispatchOutput::event); + + py::class_(m, "ElasticCombineOutput") + .def_readonly("combined_x", &ElasticCombineOutput::combined_x) + .def_readonly("combined_topk_weights", + &ElasticCombineOutput::combined_topk_weights) + .def_readonly("event", &ElasticCombineOutput::event); m.attr("MAX_QP_COUNT") = pybind11::int_(MAX_QP_COUNT); @@ -38,6 +85,50 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("combine", &MooncakeEpBuffer::combine) .def("get_next_combine_buffer", &MooncakeEpBuffer::get_next_combine_buffer); + + py::class_(m, "ElasticBuffer") + .def(py::init(), + py::arg("rank"), py::arg("num_ranks"), py::arg("num_buffer_bytes"), + py::arg("num_max_tokens_per_rank"), py::arg("hidden"), + py::arg("num_topk"), py::arg("use_fp8_dispatch"), + py::arg("deterministic"), py::arg("allow_hybrid_mode"), + py::arg("allow_multiple_reduction"), + py::arg("prefer_overlap_with_compute"), py::arg("sl_idx"), + py::arg("num_allocated_qps"), py::arg("num_cpu_timeout_secs"), + py::arg("num_gpu_timeout_secs")) + .def_static("calculate_buffer_size", + &MooncakeElasticBuffer::calculate_buffer_size) + .def("get_physical_domain_size", + &MooncakeElasticBuffer::get_physical_domain_size) + .def("get_logical_domain_size", + &MooncakeElasticBuffer::get_logical_domain_size) + .def("get_theoretical_num_sms", + &MooncakeElasticBuffer::get_theoretical_num_sms) + .def("ibgda_disabled", &MooncakeElasticBuffer::ibgda_disabled) + .def("use_fast_path", &MooncakeElasticBuffer::use_fast_path) + .def("update_local_qpns", &MooncakeElasticBuffer::update_local_qpns) + .def("is_roce", &MooncakeElasticBuffer::is_roce) + .def("sync_ibgda_peers", &MooncakeElasticBuffer::sync_ibgda_peers) + .def("get_mr_info", &MooncakeElasticBuffer::get_mr_info) + .def("get_gid", &MooncakeElasticBuffer::get_gid) + .def("get_local_qpns", &MooncakeElasticBuffer::get_local_qpns) + .def("get_local_lids", &MooncakeElasticBuffer::get_local_lids) + .def("get_ipc_handle", &MooncakeElasticBuffer::get_ipc_handle) + .def("sync_nvlink_ipc_handles", + &MooncakeElasticBuffer::sync_nvlink_ipc_handles) + .def("dispatch", &MooncakeElasticBuffer::dispatch, py::arg("x"), + py::arg("sf"), py::arg("topk_idx"), py::arg("topk_weights"), + py::arg("active_ranks"), py::arg("num_experts"), + py::arg("num_max_tokens_per_rank"), py::arg("expert_alignment"), + py::arg("num_sms"), py::arg("do_expand"), py::arg("do_cpu_sync"), + py::arg("async_with_compute_stream"), + py::arg("cached_handle") = std::nullopt) + .def("combine", &MooncakeElasticBuffer::combine, py::arg("x"), + py::arg("handle"), py::arg("topk_weights"), + py::arg("active_ranks"), py::arg("num_sms"), + py::arg("async_with_compute_stream"), + py::arg("out") = std::nullopt); } } // namespace mooncake diff --git a/mooncake-ep/src/mooncake_ep_buffer.cpp b/mooncake-ep/src/mooncake_ep_buffer.cpp index e8819f36c8..b587e29103 100644 --- a/mooncake-ep/src/mooncake_ep_buffer.cpp +++ b/mooncake-ep/src/mooncake_ep_buffer.cpp @@ -1,10 +1,21 @@ #include #include +#include +#include #include #include namespace mooncake { +namespace { + +int active_qps_per_rank_for_ep(int qps_per_rank, bool is_roce, int cap) { + if (!is_roce) return qps_per_rank; + return std::min(qps_per_rank, cap); +} + +} // namespace + // Initialize an RDMA transport: register memory, allocate control buffer, // create QPs. Returns true on success, false if IBGDA is unavailable. static bool initRdmaTransport(device::RdmaTransport* t, void* gdr_buffer, @@ -23,6 +34,14 @@ static bool initRdmaTransport(device::RdmaTransport* t, void* gdr_buffer, return ret == 0; } +static bool macaHostPhaseFenceCoversPeers() { +#ifdef MOONCAKE_EP_USE_MACA + return true; +#else + return false; +#endif +} + MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks, int64_t num_ep_buffer_bytes, TransferEngine* engine) @@ -31,6 +50,23 @@ MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks, num_ep_buffer_bytes(num_ep_buffer_bytes), comm_stream(at::cuda::getStreamFromPool(true)) { USE_QP_COUNT = MAX_QP_COUNT / num_ranks * num_ranks; + + // Optional runtime override for the RoCE active-QP cap (default 8). + // Set MOONCAKE_EP_ACTIVE_QPS_PER_RANK to a value >= the per-rank QP count + // (e.g. 256) to effectively disable the cap. + if (const char* env = std::getenv("MOONCAKE_EP_ACTIVE_QPS_PER_RANK")) { + char* end = nullptr; + long v = std::strtol(env, &end, 10); + if (end != env && *end == '\0' && v > 0) { + active_qps_cap_ = static_cast(v); + } else { + LOG(WARNING) << "[EP] ignoring invalid " + "MOONCAKE_EP_ACTIVE_QPS_PER_RANK='" + << env << "'"; + } + } + LOG(INFO) << "[EP] RoCE active QPs/rank cap = " << active_qps_cap_; + // Get ranks CUDA_CHECK(cudaGetDevice(&device_id)); CUDA_CHECK(cudaDeviceGetAttribute(&clock_rate_khz, cudaDevAttrClockRate, @@ -200,9 +236,12 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x, rdma_transport_ ? rdma_transport_->qpDevCtxsPtr() : nullptr; int32_t* nvlink_avail = p2p_transport_->availableTablePtr(); void** ipc_ptrs = p2p_transport_->peerPtrsTablePtr(); + int active_qps_per_rank = active_qps_per_rank_for_ep( + USE_QP_COUNT / num_ranks, rdma_transport_ && rdma_transport_->isRoce(), + active_qps_cap_); auto mark_send_done = [=]() { -#ifdef MOONCAKE_EP_USE_MUSA +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV mooncake::mark_phase_ack(gdr_buffer, nvlink_avail, ipc_ptrs, buffer.rdma_send_signal_buffer, rank, num_ranks, phase_epoch, launch_stream); @@ -210,7 +249,7 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x, }; auto wait_peer_send_done = [=]() { -#ifdef MOONCAKE_EP_USE_MUSA +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV mooncake::wait_phase_ack(buffer.rdma_send_signal_buffer, rank, num_ranks, phase_epoch, launch_stream, timeout_ticks); @@ -218,7 +257,7 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x, }; auto mark_and_wait_peer_send_done = [=]() { -#ifdef MOONCAKE_EP_USE_MUSA +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV mooncake::mark_and_wait_phase_ack( gdr_buffer, nvlink_avail, ipc_ptrs, buffer.rdma_send_signal_buffer, rank, num_ranks, phase_epoch, launch_stream, timeout_ticks); @@ -238,13 +277,13 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x, topk_idx.data_ptr(), next_buffer.rdma_recv_signal_buffer, num_tokens, hidden, num_max_dispatch_tokens_per_rank, num_topk, num_experts, rank, num_ranks, use_fp8, workspace, launch_stream, - timeout_ticks, phases); + timeout_ticks, phases, active_qps_per_rank); }; if (return_recv_hook) { launcher(LOW_LATENCY_SEND_PHASE); mark_send_done(); } else { -#ifdef MOONCAKE_EP_USE_MUSA +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV launcher(LOW_LATENCY_SEND_PHASE); mark_and_wait_peer_send_done(); launcher(LOW_LATENCY_RECV_PHASE); @@ -260,6 +299,8 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x, // before the stream-wait happens, so in Python API, we must wrap // all tensors into the event handle. event = EventHandle(launch_stream); + } else if (return_recv_hook && macaHostPhaseFenceCoversPeers()) { + event = EventHandle(launch_stream); } else if (not return_recv_hook) { stream_wait(compute_stream, launch_stream); } @@ -268,7 +309,7 @@ MooncakeEpBuffer::dispatch(const torch::Tensor& x, std::optional> recv_hook = std::nullopt; if (return_recv_hook) recv_hook = [=]() { - wait_peer_send_done(); + if (!macaHostPhaseFenceCoversPeers()) wait_peer_send_done(); launcher(LOW_LATENCY_RECV_PHASE); }; @@ -356,9 +397,12 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx, rdma_transport_ ? rdma_transport_->qpDevCtxsPtr() : nullptr; int32_t* nvlink_avail = p2p_transport_->availableTablePtr(); void** ipc_ptrs = p2p_transport_->peerPtrsTablePtr(); + int active_qps_per_rank = active_qps_per_rank_for_ep( + USE_QP_COUNT / num_ranks, rdma_transport_ && rdma_transport_->isRoce(), + active_qps_cap_); auto mark_send_done = [=]() { -#ifdef MOONCAKE_EP_USE_MUSA +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV mooncake::mark_phase_ack(gdr_buffer, nvlink_avail, ipc_ptrs, buffer.rdma_send_signal_buffer, rank, num_ranks, phase_epoch, launch_stream); @@ -366,7 +410,7 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx, }; auto wait_peer_send_done = [=]() { -#ifdef MOONCAKE_EP_USE_MUSA +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV mooncake::wait_phase_ack(buffer.rdma_send_signal_buffer, rank, num_ranks, phase_epoch, launch_stream, timeout_ticks); @@ -374,7 +418,7 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx, }; auto mark_and_wait_peer_send_done = [=]() { -#ifdef MOONCAKE_EP_USE_MUSA +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV mooncake::mark_and_wait_phase_ack( gdr_buffer, nvlink_avail, ipc_ptrs, buffer.rdma_send_signal_buffer, rank, num_ranks, phase_epoch, launch_stream, timeout_ticks); @@ -394,13 +438,13 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx, next_buffer.rdma_recv_signal_buffer, num_combined_tokens, hidden, num_max_dispatch_tokens_per_rank, num_topk, num_experts, rank, num_ranks, workspace, launch_stream, timeout_ticks, phases, - zero_copy); + zero_copy, active_qps_per_rank); }; if (return_recv_hook) { launcher(LOW_LATENCY_SEND_PHASE); mark_send_done(); } else { -#ifdef MOONCAKE_EP_USE_MUSA +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV launcher(LOW_LATENCY_SEND_PHASE); mark_and_wait_peer_send_done(); launcher(LOW_LATENCY_RECV_PHASE); @@ -416,6 +460,8 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx, // before the stream-wait happens, so in Python API, we must wrap // all tensors into the event handle. event = EventHandle(launch_stream); + } else if (return_recv_hook && macaHostPhaseFenceCoversPeers()) { + event = EventHandle(launch_stream); } else if (not return_recv_hook) { stream_wait(compute_stream, launch_stream); } @@ -424,7 +470,7 @@ MooncakeEpBuffer::combine(const torch::Tensor& x, const torch::Tensor& topk_idx, std::optional> recv_hook = std::nullopt; if (return_recv_hook) recv_hook = [=]() { - wait_peer_send_done(); + if (!macaHostPhaseFenceCoversPeers()) wait_peer_send_done(); launcher(LOW_LATENCY_RECV_PHASE); }; diff --git a/mooncake-ep/src/mooncake_ep_elastic_buffer.cpp b/mooncake-ep/src/mooncake_ep_elastic_buffer.cpp new file mode 100644 index 0000000000..7d3d0052cf --- /dev/null +++ b/mooncake-ep/src/mooncake_ep_elastic_buffer.cpp @@ -0,0 +1,632 @@ +#include +#include + +#include +#include +#include +#include + +#include + +namespace mooncake { +namespace { + +int64_t ceil_div_i64(int64_t x, int64_t y) { return (x + y - 1) / y; } + +constexpr int kElasticHybridChannelsPerSm = 4; + +int64_t align_i64(int64_t x, int64_t alignment) { + return ceil_div_i64(x, alignment) * alignment; +} + +int getenv_int(const char* name, int default_value) { + const char* value = std::getenv(name); + if (value == nullptr || value[0] == '\0') return default_value; + return std::max(1, std::atoi(value)); +} + +int hybrid_num_channels(int num_sms) { + return std::max(1, num_sms) * kElasticHybridChannelsPerSm; +} + +int hybrid_num_max_tokens_per_channel(int num_max_tokens_per_rank, + int num_sms) { + return static_cast( + ceil_div_i64(num_max_tokens_per_rank, hybrid_num_channels(num_sms))); +} + +int64_t elastic_workspace_num_bytes() { + constexpr int64_t kNumMaxRanks = 1024; + constexpr int64_t kNumMaxExperts = 2048; + constexpr int64_t kNumMaxChannels = 8 * 160; + constexpr int64_t kNumMaxInflightAGRS = 32; + constexpr int64_t kNumBarrierTags = 16; + + int64_t num_bytes = 0; + num_bytes += kNumBarrierTags * + (sizeof(unsigned long long) + 2 * kNumMaxRanks * sizeof(int)); + num_bytes += (kNumMaxRanks + kNumMaxExperts) * sizeof(int64_t); + num_bytes += kNumMaxRanks * sizeof(int64_t) * 2; + num_bytes += kNumMaxExperts * sizeof(int64_t) * 2; + num_bytes += kNumMaxRanks * sizeof(int); + num_bytes += kNumMaxRanks * sizeof(int) * 2; + num_bytes += kNumMaxExperts * sizeof(int) * 2; + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int64_t); + num_bytes += kNumMaxRanks * kNumMaxChannels * sizeof(int); + num_bytes += 2 * 2 * sizeof(int64_t); + num_bytes += (kNumMaxInflightAGRS + 1) * kNumMaxRanks * sizeof(int); + return align_i64(num_bytes, 32); +} + +int64_t elastic_atomic_scratch_num_bytes() { + return elastic_workspace_num_bytes(); +} + +int device_smem_bytes() { +#ifdef MOONCAKE_EP_USE_MUSA + return 0; +#else + int device = 0; + cudaGetDevice(&device); + int value = 0; + cudaDeviceGetAttribute(&value, cudaDevAttrMaxSharedMemoryPerBlockOptin, + device); + return value > 0 ? value : 98304; +#endif +} + +} // namespace + +ElasticLaunchContext MooncakeElasticBuffer::make_launch_context( + MooncakeEpBuffer& buffer, const ElasticTopology& topology, + void* mapped_host_workspace, int64_t timeout_cycles) { + ElasticLaunchContext ctx; + auto* rdma = buffer.rdma_transport_; + auto* gdr_base = static_cast(buffer.gdr_buffer); + // Mooncake P2P/RDMA Device API translates remote pointers as offsets from + // the registered GDR buffer base. DeepEP elastic writes both `buffer` and + // `workspace` pointers to peer ranks through GIN, so both regions must live + // inside the same peer-visible registered allocation. The elastic buffer + // size reserves `elastic_workspace_num_bytes()` first; use that prefix as + // the workspace. RDMA atomics also need a separate local response area: + // mlx5 atomics write the fetched old value to the WQE local address, so + // reusing the remote signal workspace as `local_atomic_base` can corrupt + // the barrier/signal slots. Reserve an equal-sized scratch prefix after + // the workspace, then place the communication buffer after both prefixes. + const auto workspace_bytes = elastic_workspace_num_bytes(); + const auto atomic_scratch_bytes = elastic_atomic_scratch_num_bytes(); + ctx.gdr_buffer = gdr_base; + ctx.nvlink_available = buffer.p2p_transport_->availableTablePtr(); + ctx.ipc_peer_ptrs = buffer.p2p_transport_->peerPtrsTablePtr(); + ctx.raddrs = rdma ? rdma->raddrsPtr() : nullptr; + ctx.rkeys = rdma ? rdma->rkeysPtr() : nullptr; + ctx.qp_devctxs = rdma ? rdma->qpDevCtxsPtr() : nullptr; + ctx.rdma_send_signal_buffer = gdr_base + workspace_bytes; + ctx.rdma_recv_signal_buffer = gdr_base; + ctx.workspace = gdr_base; + ctx.buffer = gdr_base + workspace_bytes + atomic_scratch_bytes; + ctx.mapped_host_workspace = mapped_host_workspace; + ctx.rank = topology.rank_idx; + ctx.num_ranks = topology.num_ranks; + ctx.scaleout_rank_idx = topology.scaleout_rank_idx; + ctx.scaleup_rank_idx = topology.scaleup_rank_idx; + ctx.num_scaleout_ranks = topology.num_scaleout_ranks; + ctx.num_scaleup_ranks = topology.num_scaleup_ranks; + ctx.is_scaleup_nvlink = true; + ctx.num_qps = buffer.USE_QP_COUNT; + ctx.timeout_cycles = timeout_cycles; + return ctx; +} + +MooncakeElasticBuffer::MooncakeElasticBuffer( + int rank, int num_ranks, int64_t num_buffer_bytes, + int64_t num_max_tokens_per_rank, int64_t hidden, int64_t num_topk, + bool use_fp8_dispatch, bool deterministic, bool allow_hybrid_mode, + bool allow_multiple_reduction, bool prefer_overlap_with_compute, int sl_idx, + int num_allocated_qps, int num_cpu_timeout_secs, int num_gpu_timeout_secs) { + config_.num_max_tokens_per_rank = num_max_tokens_per_rank; + config_.hidden = hidden; + config_.num_topk = num_topk; + config_.use_fp8_dispatch = use_fp8_dispatch; + config_.deterministic = deterministic; + config_.allow_hybrid_mode = allow_hybrid_mode; + config_.allow_multiple_reduction = allow_multiple_reduction; + config_.prefer_overlap_with_compute = prefer_overlap_with_compute; + config_.sl_idx = sl_idx; + config_.num_allocated_qps = num_allocated_qps; + config_.num_cpu_timeout_secs = num_cpu_timeout_secs; + config_.num_gpu_timeout_secs = num_gpu_timeout_secs; + + topology_ = discover_topology(rank, num_ranks, allow_hybrid_mode); + if (!allow_multiple_reduction) { + throw std::runtime_error( + "Mooncake ElasticBuffer currently supports only " + "allow_multiple_reduction=true"); + } + if (num_buffer_bytes == 0) { + num_buffer_bytes = calculate_buffer_size( + num_ranks, num_max_tokens_per_rank, hidden, num_topk, + use_fp8_dispatch, allow_hybrid_mode, allow_multiple_reduction); + } + native_buffer_ = + std::make_unique(rank, num_ranks, num_buffer_bytes); + host_workspace_bytes_ = elastic_workspace_num_bytes(); + CUDA_CHECK(cudaHostAlloc(&host_workspace_, host_workspace_bytes_, + cudaHostAllocMapped)); + CUDA_CHECK( + cudaHostGetDevicePointer(&mapped_host_workspace_, host_workspace_, 0)); + std::memset(host_workspace_, 0, host_workspace_bytes_); +} + +MooncakeElasticBuffer::~MooncakeElasticBuffer() { + if (host_workspace_ != nullptr) { + cudaFreeHost(host_workspace_); + host_workspace_ = nullptr; + mapped_host_workspace_ = nullptr; + } +} + +int64_t MooncakeElasticBuffer::calculate_buffer_size( + int num_ranks, int64_t num_max_tokens_per_rank, int64_t hidden, + int64_t num_topk, bool use_fp8_dispatch, bool allow_hybrid_mode, + bool allow_multiple_reduction) { + num_topk = std::max(1, num_topk); + const int64_t dtype_bytes = use_fp8_dispatch ? 1 : 2; + const int64_t scale_bytes = + use_fp8_dispatch ? ceil_div_i64(hidden, 128) * 4 : 0; + const int64_t token_bytes = + align_i64(hidden * dtype_bytes, 32) + align_i64(scale_bytes, 32); + const int64_t metadata_bytes = align_i64( + num_topk * (sizeof(int) + sizeof(float)) + (1 + num_topk) * sizeof(int), + 32); + const int64_t per_slot_bytes = token_bytes + metadata_bytes; + const int64_t dispatch_bytes = + num_ranks * num_max_tokens_per_rank * num_topk * per_slot_bytes * 2; + const int64_t combine_factor = allow_multiple_reduction ? 3 : 4; + const int64_t combine_bytes = dispatch_bytes * combine_factor; + const int64_t hybrid_factor = allow_hybrid_mode && num_ranks > 1 ? 2 : 1; + return elastic_workspace_num_bytes() + elastic_atomic_scratch_num_bytes() + + hybrid_factor * (dispatch_bytes + combine_bytes); +} + +std::tuple MooncakeElasticBuffer::get_physical_domain_size() const { + return {topology_.num_rdma_ranks, topology_.num_nvlink_ranks}; +} + +std::tuple MooncakeElasticBuffer::get_logical_domain_size() const { + return {topology_.num_scaleout_ranks, topology_.num_scaleup_ranks}; +} + +int MooncakeElasticBuffer::get_theoretical_num_sms(int num_experts, + int num_topk) const { + int device = 0; + cudaGetDevice(&device); + cudaDeviceProp prop{}; + cudaGetDeviceProperties(&prop, device); + if (config_.prefer_overlap_with_compute) { + return std::max(1, std::min(24, prop.multiProcessorCount / 4)); + } + return std::max(1, std::min({40, prop.multiProcessorCount / 2, + std::max(1, num_experts * num_topk)})); +} + +ElasticDispatchOutput MooncakeElasticBuffer::dispatch( + const torch::Tensor& x, const std::optional& sf, + const torch::Tensor& topk_idx, + const std::optional& topk_weights, + torch::Tensor& active_ranks, int num_experts, int num_max_tokens_per_rank, + int expert_alignment, int num_sms, bool do_expand, bool do_cpu_sync, + bool async_with_compute_stream, + const std::optional& cached_handle) { + EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous()); + const bool use_sf = sf.has_value(); + if (use_sf) { + EP_HOST_ASSERT(x.element_size() == 1); + EP_HOST_ASSERT(sf->dim() == 2 && sf->is_cuda()); + EP_HOST_ASSERT(sf->scalar_type() == torch::kFloat32 || + sf->scalar_type() == torch::kInt32); + EP_HOST_ASSERT(sf->size(0) == x.size(0)); + } else { + EP_HOST_ASSERT(!config_.use_fp8_dispatch); + EP_HOST_ASSERT(x.scalar_type() == torch::kBFloat16); + } + EP_HOST_ASSERT(topk_idx.dim() == 2 && topk_idx.is_contiguous()); + EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64); + EP_HOST_ASSERT(x.size(0) == topk_idx.size(0)); + EP_HOST_ASSERT(num_experts % topology_.num_ranks == 0); + + const int num_tokens = static_cast(x.size(0)); + const int hidden = static_cast(x.size(1)); + const int num_topk = static_cast(topk_idx.size(1)); + const int num_sf_packs = use_sf ? static_cast(sf->size(1)) : 0; + const int sf_token_stride = use_sf ? static_cast(sf->stride(0)) : 0; + const int sf_hidden_stride = use_sf ? static_cast(sf->stride(1)) : 0; + const int num_local_experts = num_experts / topology_.num_ranks; + // The copy epilogue uses `kNumMaxTokensPerRank * kNumRanks` as the + // no-CPU-sync sentinel and then reads the real local receive count from the + // GPU prefix-sum tensor. In hybrid mode each scale-up peer may receive + // tokens forwarded from every scale-out rank, so the conservative output + // capacity and sentinel must cover the full logical world, not just the + // intra-node scale-up domain. + const int num_recv_tokens = num_max_tokens_per_rank * topology_.num_ranks; + const int num_smem_bytes = device_smem_bytes(); + const int num_channels_per_sm = 1; + const int num_channels = num_sms * num_channels_per_sm; + const bool cached_mode = cached_handle.has_value(); + const bool use_hybrid = topology_.num_scaleout_ranks != 1; + const int hybrid_channels = use_hybrid ? hybrid_num_channels(num_sms) : 0; + const int hybrid_max_tokens_per_channel = + use_hybrid ? hybrid_num_max_tokens_per_channel(num_max_tokens_per_rank, + num_sms) + : 0; + if (cached_mode) { + const auto& handle = cached_handle.value(); + EP_HOST_ASSERT(!handle.do_expand && !do_expand); + EP_HOST_ASSERT(handle.num_experts == num_experts); + EP_HOST_ASSERT(handle.expert_alignment == expert_alignment); + EP_HOST_ASSERT(handle.num_max_tokens_per_rank == + num_max_tokens_per_rank); + EP_HOST_ASSERT(handle.num_sms == num_sms); + if (use_hybrid) { + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.dim() == 4); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(0) == + hybrid_channels); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(1) == + topology_.num_scaleout_ranks); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(2) == + hybrid_max_tokens_per_channel); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(3) == num_topk); + EP_HOST_ASSERT(handle.token_metadata_at_forward.has_value()); + EP_HOST_ASSERT(handle.channel_linked_list.has_value()); + } else { + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.dim() == 2); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(0) == num_tokens); + EP_HOST_ASSERT(handle.dst_buffer_slot_idx.size(1) == num_topk); + } + } + + auto compute_stream = at::cuda::getCurrentCUDAStream(); + auto launch_stream = native_buffer_->comm_stream; + stream_wait(launch_stream, compute_stream); + + const int64_t timeout_cycles = + config_.num_gpu_timeout_secs < 0 + ? -1 + : static_cast(native_buffer_->clock_rate_khz) * + static_cast(config_.num_gpu_timeout_secs) * 1000; + auto launch_ctx = make_launch_context( + *native_buffer_, topology_, mapped_host_workspace_, timeout_cycles); + + auto psum_num_recv_tokens_per_scaleup_rank = + cached_mode ? cached_handle->psum_num_recv_tokens_per_scaleup_rank + : torch::empty({topology_.num_scaleup_ranks}, + torch::TensorOptions() + .dtype(torch::kInt32) + .device(x.device())); + auto psum_num_recv_tokens_per_expert = + cached_mode + ? cached_handle->psum_num_recv_tokens_per_expert + : torch::empty({num_local_experts + 1}, torch::TensorOptions() + .dtype(torch::kInt32) + .device(x.device())); + auto dst_buffer_slot_idx = + cached_mode + ? cached_handle->dst_buffer_slot_idx + : (use_hybrid ? torch::empty( + {hybrid_channels, topology_.num_scaleout_ranks, + hybrid_max_tokens_per_channel, num_topk}, + torch::TensorOptions() + .dtype(torch::kInt32) + .device(x.device())) + : torch::empty({num_tokens, num_topk}, + torch::TensorOptions() + .dtype(torch::kInt32) + .device(x.device()))); + std::optional token_metadata_at_forward = std::nullopt; + std::optional channel_linked_list = std::nullopt; + if (use_hybrid) { + if (cached_mode) { + token_metadata_at_forward = + cached_handle->token_metadata_at_forward; + channel_linked_list = cached_handle->channel_linked_list; + } else { + const int forward_metadata_dims = 2 + num_topk * 2; + token_metadata_at_forward = torch::empty( + {hybrid_channels, + topology_.num_scaleout_ranks * hybrid_max_tokens_per_channel + + 1, + forward_metadata_dims}, + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + channel_linked_list = torch::empty( + {hybrid_channels, + topology_.num_scaleout_ranks * hybrid_max_tokens_per_channel + + 1, + topology_.num_scaleup_ranks}, + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + } + } + std::optional deterministic_rank_count_buffer = std::nullopt; +#ifdef MOONCAKE_EP_USE_MUSA + // MUSA non-hybrid dispatch always runs + // launch_musa_elastic_prepare_dispatch(), which assigns slots and publishes + // counts without cooperative grid sync. + const bool run_deterministic_prologue = false; +#else + const bool run_deterministic_prologue = + config_.deterministic && !cached_mode && !use_hybrid; +#endif + if (run_deterministic_prologue) { + deterministic_rank_count_buffer = torch::empty( + {num_sms, topology_.num_scaleup_ranks}, + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + launch_elastic_dispatch_deterministic_prologue( + topk_idx.data_ptr(), + deterministic_rank_count_buffer.value().data_ptr(), + dst_buffer_slot_idx.data_ptr(), num_tokens, + num_max_tokens_per_rank, num_experts, num_topk, + topology_.scaleup_rank_idx, topology_.num_scaleup_ranks, num_sms, + num_smem_bytes, launch_stream.stream()); + } + + launch_mooncake_elastic_dispatch( + x.data_ptr(), use_sf ? const_cast(sf->data_ptr()) : nullptr, + const_cast(topk_idx.data_ptr()), + topk_weights.has_value() + ? const_cast(topk_weights->data_ptr()) + : nullptr, + nullptr, nullptr, psum_num_recv_tokens_per_scaleup_rank.data_ptr(), + psum_num_recv_tokens_per_expert.data_ptr(), + dst_buffer_slot_idx.data_ptr(), + token_metadata_at_forward.has_value() + ? token_metadata_at_forward->data_ptr() + : nullptr, + num_tokens, num_max_tokens_per_rank, hidden, + static_cast(x.element_size()), num_sf_packs, sf_token_stride, + sf_hidden_stride, num_experts, num_topk, expert_alignment, num_sms, + use_hybrid ? kElasticHybridChannelsPerSm : num_channels_per_sm, + num_smem_bytes, cached_mode, config_.deterministic, false, launch_ctx, + launch_stream.stream()); + + const int num_recv_output_capacity = + do_expand ? num_recv_tokens * num_topk : num_recv_tokens; + auto recv_x = torch::empty({num_recv_output_capacity, hidden}, x.options()); + auto recv_x_scales = std::optional(); + void* recv_x_scales_ptr = nullptr; + int recv_sf_token_stride = 0; + int recv_sf_hidden_stride = 0; + if (use_sf) { + recv_x_scales = torch::empty({num_recv_output_capacity, num_sf_packs}, + sf->options()); + recv_x_scales_ptr = recv_x_scales->data_ptr(); + recv_sf_token_stride = static_cast(recv_x_scales->stride(0)); + recv_sf_hidden_stride = static_cast(recv_x_scales->stride(1)); + } + auto recv_topk_idx = + torch::empty({num_recv_tokens, num_topk}, topk_idx.options()); + auto recv_topk_weights = std::optional(); + float* recv_topk_weights_ptr = nullptr; + if (topk_weights.has_value()) { + recv_topk_weights = do_expand + ? torch::empty({num_recv_output_capacity}, + topk_weights->options()) + : torch::empty({num_recv_tokens, num_topk}, + topk_weights->options()); + recv_topk_weights_ptr = recv_topk_weights->data_ptr(); + } + auto recv_src_metadata = torch::empty( + {num_recv_tokens, num_topk + 2}, + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + auto handle_psum_num_recv_tokens_per_expert = + do_expand + ? psum_num_recv_tokens_per_expert.slice(0, 0, num_local_experts) + : psum_num_recv_tokens_per_expert.slice(0, 1, + num_local_experts + 1); + auto epilogue_psum_num_recv_tokens_per_expert = + do_expand ? psum_num_recv_tokens_per_expert + : handle_psum_num_recv_tokens_per_expert; + + launch_mooncake_elastic_dispatch_copy_epilogue( + recv_x.data_ptr(), recv_x_scales_ptr, recv_topk_idx.data_ptr(), + recv_topk_weights_ptr, recv_src_metadata.data_ptr(), + channel_linked_list.has_value() ? channel_linked_list->data_ptr() + : nullptr, + num_recv_tokens, num_max_tokens_per_rank, hidden, + static_cast(x.element_size()), num_sf_packs, recv_sf_token_stride, + recv_sf_hidden_stride, num_experts, num_topk, num_sms, num_smem_bytes, + use_hybrid ? hybrid_channels : num_channels, do_expand, cached_mode, + launch_ctx, psum_num_recv_tokens_per_scaleup_rank.data_ptr(), + epilogue_psum_num_recv_tokens_per_expert.data_ptr(), + launch_stream.stream()); + + if (do_cpu_sync || !async_with_compute_stream) { + stream_wait(compute_stream, launch_stream); + } + std::optional event = std::nullopt; + if (async_with_compute_stream) { + event = EventHandle(launch_stream); + } + + std::vector num_recv_tokens_per_expert_list; + int actual_num_recv_tokens = num_recv_tokens; + int actual_num_output_tokens = num_recv_tokens; + if (do_cpu_sync) { + auto scaleup_psum_cpu = psum_num_recv_tokens_per_scaleup_rank.cpu(); + auto expert_psum_cpu = psum_num_recv_tokens_per_expert.cpu(); + const auto* scaleup_psum = scaleup_psum_cpu.data_ptr(); + const auto* expert_psum = expert_psum_cpu.data_ptr(); + actual_num_recv_tokens = scaleup_psum[topology_.num_scaleup_ranks - 1]; + EP_HOST_ASSERT(actual_num_recv_tokens >= 0 && + actual_num_recv_tokens <= num_recv_tokens); + actual_num_output_tokens = actual_num_recv_tokens; + + num_recv_tokens_per_expert_list.reserve(num_local_experts); + const auto align_count = [expert_alignment](int value) { + return ((value + expert_alignment - 1) / expert_alignment) * + expert_alignment; + }; + if (do_expand) { + int previous_psum = 0; + for (int i = 0; i < num_local_experts; ++i) { + const int count = expert_psum[i] - align_count(previous_psum); + EP_HOST_ASSERT(count >= 0); + num_recv_tokens_per_expert_list.push_back(count); + previous_psum = expert_psum[i]; + } + actual_num_output_tokens = + num_local_experts == 0 ? 0 : expert_psum[num_local_experts - 1]; + } else { + for (int i = 0; i < num_local_experts; ++i) { + const int count = expert_psum[i + 1] - expert_psum[i]; + EP_HOST_ASSERT(count >= 0); + num_recv_tokens_per_expert_list.push_back(count); + } + } + EP_HOST_ASSERT(actual_num_output_tokens >= 0 && + actual_num_output_tokens <= recv_x.size(0)); + + recv_x = recv_x.slice(0, 0, actual_num_output_tokens); + if (recv_x_scales.has_value()) { + recv_x_scales = + recv_x_scales->slice(0, 0, actual_num_output_tokens); + } + recv_topk_idx = recv_topk_idx.slice(0, 0, actual_num_recv_tokens); + if (recv_topk_weights.has_value()) { + recv_topk_weights = + recv_topk_weights->slice(0, 0, actual_num_output_tokens); + } + recv_src_metadata = + recv_src_metadata.slice(0, 0, actual_num_recv_tokens); + } + + ElasticNativeHandle handle; + handle.do_expand = do_expand; + handle.num_experts = num_experts; + handle.expert_alignment = expert_alignment; + handle.num_max_tokens_per_rank = num_max_tokens_per_rank; + handle.num_sms = num_sms; + handle.topk_idx = cached_mode ? cached_handle->topk_idx : topk_idx.clone(); + handle.psum_num_recv_tokens_per_expert = + handle_psum_num_recv_tokens_per_expert; + handle.psum_num_recv_tokens_per_scaleup_rank = + psum_num_recv_tokens_per_scaleup_rank; + handle.recv_src_metadata = recv_src_metadata; + handle.recv_layout_range = torch::empty( + {0}, torch::TensorOptions().dtype(torch::kInt64).device(x.device())); + handle.dst_buffer_slot_idx = dst_buffer_slot_idx; + handle.token_metadata_at_forward = token_metadata_at_forward; + handle.channel_linked_list = channel_linked_list; + handle.num_recv_tokens_per_expert_list = num_recv_tokens_per_expert_list; + + ElasticDispatchOutput output; + output.recv_x = recv_x; + output.recv_x_scales = recv_x_scales; + output.recv_topk_idx = recv_topk_idx; + output.recv_topk_weights = recv_topk_weights; + output.handle = handle; + output.event = event; + return output; +} + +ElasticCombineOutput MooncakeElasticBuffer::combine( + const torch::Tensor& x, const ElasticNativeHandle& handle, + const std::optional& topk_weights, + torch::Tensor& active_ranks, int num_sms, bool async_with_compute_stream, + const std::optional& out) { + EP_HOST_ASSERT(x.dim() == 2 && x.is_contiguous()); + EP_HOST_ASSERT(x.scalar_type() == torch::kBFloat16); + torch::Tensor weights = topk_weights.value_or(torch::Tensor()); + if (!weights.defined()) { + weights = torch::ones( + handle.topk_idx.sizes(), + torch::TensorOptions().dtype(torch::kFloat32).device(x.device())); + } + const int hidden = static_cast(x.size(1)); + const int num_topk = static_cast(handle.topk_idx.size(1)); + const int num_combined_tokens = static_cast(handle.topk_idx.size(0)); + const int num_smem_bytes = device_smem_bytes(); + const int num_channels = std::max(1, num_sms); + const bool use_hybrid = topology_.num_scaleout_ranks != 1; + const int hybrid_channels = use_hybrid ? hybrid_num_channels(num_sms) : 0; + auto compute_stream = at::cuda::getCurrentCUDAStream(); + auto launch_stream = native_buffer_->comm_stream; + stream_wait(launch_stream, compute_stream); + const int64_t timeout_cycles = + config_.num_gpu_timeout_secs < 0 + ? -1 + : static_cast(native_buffer_->clock_rate_khz) * + static_cast(config_.num_gpu_timeout_secs) * 1000; + auto launch_ctx = make_launch_context( + *native_buffer_, topology_, mapped_host_workspace_, timeout_cycles); + auto psum_num_recv_tokens_per_scaleup_rank = + handle.psum_num_recv_tokens_per_scaleup_rank; + void* reduce_buffer = launch_mooncake_elastic_combine( + x.data_ptr(), weights.data_ptr(), + const_cast(handle.recv_src_metadata.data_ptr()), + psum_num_recv_tokens_per_scaleup_rank.data_ptr(), + handle.token_metadata_at_forward.has_value() + ? handle.token_metadata_at_forward->data_ptr() + : nullptr, + handle.channel_linked_list.has_value() + ? handle.channel_linked_list->data_ptr() + : nullptr, + static_cast(x.size(0)), handle.num_max_tokens_per_rank, hidden, + handle.num_experts, num_topk, num_sms, num_smem_bytes, + use_hybrid ? hybrid_channels : num_channels, handle.do_expand, + config_.allow_multiple_reduction, launch_ctx, launch_stream.stream()); + + torch::Tensor combined_x = + out.has_value() + ? out.value() + : torch::empty({num_combined_tokens, hidden}, x.options()); + launch_mooncake_elastic_combine_reduce_epilogue( + combined_x.data_ptr(), weights.data_ptr(), + const_cast(handle.topk_idx.data_ptr()), + num_combined_tokens, handle.num_max_tokens_per_rank, hidden, + handle.num_experts, num_topk, reduce_buffer, nullptr, nullptr, num_sms, + num_smem_bytes, handle.do_expand, config_.allow_multiple_reduction, + launch_ctx, launch_stream.stream()); + + if (!async_with_compute_stream) { + stream_wait(compute_stream, launch_stream); + } + std::optional event = std::nullopt; + if (async_with_compute_stream) event = EventHandle(launch_stream); + (void)active_ranks; + + ElasticCombineOutput output; + output.combined_x = combined_x; + output.combined_topk_weights = std::nullopt; + output.event = event; + return output; +} + +ElasticTopology MooncakeElasticBuffer::discover_topology( + int rank, int num_ranks, bool allow_hybrid_mode) { + int device_count = 1; + cudaGetDeviceCount(&device_count); + int num_local_ranks = + getenv_int("MOONCAKE_EP_NUM_LOCAL_RANKS", + std::max(1, std::min(num_ranks, device_count))); + num_local_ranks = std::max(1, std::min(num_local_ranks, num_ranks)); + + ElasticTopology topology; + topology.rank_idx = rank; + topology.num_ranks = num_ranks; + topology.num_rdma_ranks = + static_cast(ceil_div_i64(num_ranks, num_local_ranks)); + topology.num_nvlink_ranks = num_local_ranks; + if (allow_hybrid_mode && topology.num_rdma_ranks > 1) { + topology.num_scaleout_ranks = topology.num_rdma_ranks; + topology.num_scaleup_ranks = topology.num_nvlink_ranks; + topology.hybrid_enabled = true; + } else { + topology.num_scaleout_ranks = 1; + topology.num_scaleup_ranks = num_ranks; + topology.hybrid_enabled = false; + } + topology.scaleout_rank_idx = rank / topology.num_scaleup_ranks; + topology.scaleup_rank_idx = rank % topology.num_scaleup_ranks; + return topology; +} + +} // namespace mooncake diff --git a/mooncake-ep/src/mooncake_ep_elastic_kernel.cu b/mooncake-ep/src/mooncake_ep_elastic_kernel.cu new file mode 100644 index 0000000000..a5d3e383ff --- /dev/null +++ b/mooncake-ep/src/mooncake_ep_elastic_kernel.cu @@ -0,0 +1,973 @@ +// clang-format off + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace mooncake { +namespace { + +constexpr int kElasticNumNotifyWarps = 4; +#ifdef MOONCAKE_EP_USE_MUSA +constexpr int kElasticNumDispatchWarps = 4; +constexpr int kElasticNumEpilogueWarps = 4; +#else +constexpr int kElasticNumDispatchWarps = 8; +constexpr int kElasticNumEpilogueWarps = 8; +#endif +constexpr int kElasticNumHybridScaleoutWarps = 4; +constexpr int kElasticNumHybridForwardWarps = 4; +constexpr int kElasticNumHybridScaleupWarps = 4; +constexpr int kElasticNumQPs = MAX_QP_COUNT; +constexpr int64_t kElasticTimeoutCycles = NUM_TIMEOUT_CYCLES; + +inline int ceil_div(int x, int y) { return (x + y - 1) / y; } + +inline int hybrid_num_channels(int num_sms) { + return num_sms * kElasticNumHybridForwardWarps; +} + +inline void* hybrid_combine_reduce_buffer_ptr(void* buffer, int hidden, + int num_topk, + int num_max_tokens_per_rank, + int num_scaleout_ranks, + int num_scaleup_ranks, + bool allow_multiple_reduction) { + const int num_tokens_in_scaleup_layout = + allow_multiple_reduction && num_scaleup_ranks <= num_topk + ? num_scaleup_ranks + : num_topk; + const auto token_layout = elastic::layout::TokenLayout( + hidden * static_cast(sizeof(nv_bfloat16)), 0, num_topk, false); + const auto scaleup_buffer = elastic::layout::BufferLayout( + token_layout, num_tokens_in_scaleup_layout, + num_scaleout_ranks * num_max_tokens_per_rank, buffer); + return scaleup_buffer.get_buffer_end_ptr(); +} + +inline int dispatch_smem_bytes(int hidden, int elem_size, int num_sf_packs, + int num_topk, int num_ranks, + int num_experts, int num_notify_warps, + int num_dispatch_warps) { + const int notify_smem_bytes = num_notify_warps == 0 + ? 0 + : elastic::math::align(num_ranks + num_experts, num_notify_warps * 32) * + static_cast(sizeof(int)); + const auto token_layout = + elastic::layout::TokenLayout(hidden * elem_size, + num_sf_packs * sizeof(sf_pack_t), + num_topk, true); + return notify_smem_bytes + + num_dispatch_warps * static_cast(token_layout.get_num_bytes()); +} + +inline int dispatch_epilogue_smem_bytes(int hidden, int elem_size, + int num_sf_packs, int num_topk, + int num_warps) { + const auto token_layout = + elastic::layout::TokenLayout(hidden * elem_size, + num_sf_packs * sizeof(sf_pack_t), + num_topk, true); + return num_warps * static_cast(token_layout.get_num_bytes()); +} + +inline int combine_smem_bytes(int hidden, int num_topk, int num_warps) { + const auto token_layout = elastic::layout::TokenLayout( + hidden * static_cast(sizeof(nv_bfloat16)), 0, num_topk, false); + return num_warps * static_cast(token_layout.get_num_bytes()); +} + +inline int combine_epilogue_smem_bytes(int hidden, int num_warps) { + const auto token_layout = elastic::layout::TokenLayout( + hidden * static_cast(sizeof(nv_bfloat16)), 0, 0, false); + return num_warps * static_cast(token_layout.get_num_bytes()); +} + +inline device::CommCtx make_comm_ctx(const ElasticLaunchContext& ctx) { + device::CommCtx comm_ctx{}; + comm_ctx.rank = ctx.rank; + comm_ctx.p2p.available = ctx.nvlink_available; + comm_ctx.p2p.peer_ptrs = ctx.ipc_peer_ptrs; + comm_ctx.p2p.local_base = ctx.gdr_buffer; + comm_ctx.ibgda.qp_devctxs = + reinterpret_cast(ctx.qp_devctxs); + comm_ctx.ibgda.raddrs = reinterpret_cast(ctx.raddrs); + comm_ctx.ibgda.rkeys = reinterpret_cast(ctx.rkeys); + comm_ctx.ibgda.local_atomic_base = ctx.rdma_send_signal_buffer; + comm_ctx.ibgda.remote_atomic_base = ctx.rdma_recv_signal_buffer; + return comm_ctx; +} + +#ifdef MOONCAKE_EP_USE_MUSA + +// MUSA currently cannot rely on CUDA-style cooperative grid synchronization for +// the official elastic dispatch notify/prologue. These prepare kernels keep the +// dispatch algorithm semantics unchanged (slot assignment, count publication and +// prefix-sum generation), but split that prologue into ordinary launches with an +// explicit scale-up barrier so peer count writes cannot race with local clears. +// Payload movement still goes through the common dispatch kernel and Mooncake +// Device API transport primitives; this is only a no-cooperative-grid-sync +// metadata preparation fallback. + +__global__ void musa_elastic_prepare_init_kernel(void* workspace, + int num_scaleup_ranks, + int num_experts) { + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const int tid = blockIdx.x * blockDim.x + threadIdx.x; + const int stride = blockDim.x * gridDim.x; + for (int i = tid; i < num_scaleup_ranks + num_experts; i += stride) { + layout.get_scaleup_rank_expert_count_ptr()[i] = 0; + layout.get_scaleup_rank_expert_count_ptr()[i] = 0; + } + for (int i = tid; i < num_scaleup_ranks; i += stride) { + layout.get_scaleup_atomic_sender_counter()[i] = 0; + } +} + +__global__ void musa_elastic_prepare_clear_barrier_kernel( + device::CommCtx comm_ctx, void* workspace, int rank_idx, + int num_scaleup_ranks, int num_experts, int64_t timeout_cycles) { + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const auto gin = elastic::transport::MooncakeGin( + comm_ctx, 0, 0, 1, 0, rank_idx, num_scaleup_ranks, num_scaleup_ranks); + constexpr int kTag = elastic::comm::kDeviceBarrierTag; + const int status = + static_cast((*layout.get_nvl_barrier_counter_ptr(kTag)) & 3); + const int phase = status & 1; + const int sign = status >> 1; + const int* base_signal = layout.get_nvl_barrier_signal_ptr(kTag, phase); + + if (threadIdx.x < num_scaleup_ranks) { + auto* dst_ptr = const_cast(base_signal) + rank_idx; + gin.red_add_rel( + dst_ptr, sign ? -1 : 1, threadIdx.x); + } + __syncthreads(); + + if (threadIdx.x == 0) { + atomicAdd(layout.get_nvl_barrier_counter_ptr(kTag), 1ULL); + const int target = sign ? 0 : num_scaleup_ranks; + const auto start_clock = clock64(); + while (true) { + int sum = 0; + for (int i = 0; i < num_scaleup_ranks; ++i) { + sum += elastic::ptx::ld_acquire_sys( + const_cast(base_signal) + i); + } + if (sum == target) break; + if (timeout_cycles >= 0 && + clock64() - start_clock >= timeout_cycles) { + printf( + "MUSA prepare clear barrier timeout, rank=%d sum=%d " + "target=%d\n", + rank_idx, sum, target); + break; + } + } + } +} + +__global__ void musa_elastic_assign_slots_kernel( + const int64_t* topk_idx, int* dst_buffer_slot_idx, void* workspace, + int num_tokens, int num_max_tokens_per_rank, int num_experts, int num_topk, + int num_scaleup_ranks, int rank_idx) { + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const int num_experts_per_rank = num_experts / num_scaleup_ranks; + const int token_stride = blockDim.x * gridDim.x; + for (int token_idx = blockIdx.x * blockDim.x + threadIdx.x; + token_idx < num_tokens; token_idx += token_stride) { + int seen_ranks[32]; + int num_seen = 0; + for (int k = 0; k < num_topk; ++k) { + dst_buffer_slot_idx[token_idx * num_topk + k] = -1; + } + for (int k = 0; k < num_topk; ++k) { + const int expert_idx = + static_cast(topk_idx[token_idx * num_topk + k]); + if (expert_idx < 0) continue; + const int dst_rank = expert_idx / num_experts_per_rank; + bool duplicate_rank = false; + for (int i = 0; i < num_seen; ++i) + duplicate_rank |= (seen_ranks[i] == dst_rank); + if (!duplicate_rank) { + seen_ranks[num_seen++] = dst_rank; + const int slot = atomicAdd( + layout.get_scaleup_atomic_sender_counter() + dst_rank, 1); + dst_buffer_slot_idx[token_idx * num_topk + k] = + rank_idx * num_max_tokens_per_rank + slot; + atomicAdd(reinterpret_cast( + layout.get_scaleup_rank_count_ptr() + + dst_rank), + 1ULL); + } + atomicAdd(reinterpret_cast( + layout.get_scaleup_expert_count_ptr() + + expert_idx), + 1ULL); + } + } +} + +__global__ void musa_elastic_publish_counts_kernel( + device::CommCtx comm_ctx, void* workspace, int rank_idx, + int num_scaleup_ranks, int num_experts) { + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const int num_experts_per_rank = num_experts / num_scaleup_ranks; + const auto gin = elastic::transport::MooncakeGin( + comm_ctx, 0, 0, 1, 0, rank_idx, num_scaleup_ranks, num_scaleup_ranks); + const int tid = blockIdx.x * blockDim.x + threadIdx.x; + const int stride = blockDim.x * gridDim.x; + for (int dst_rank = tid; dst_rank < num_scaleup_ranks; dst_rank += stride) { + const auto count = static_cast( + layout.get_scaleup_rank_count_ptr()[dst_rank]); + auto* dst = reinterpret_cast( + layout.get_scaleup_rank_count_ptr() + rank_idx); + const auto encoded_count = + elastic::math::encode_decode_positive(count); + if (dst_rank == rank_idx) { + elastic::ptx::st_release_sys(dst, encoded_count); + } else { + gin.put_value( + dst, encoded_count, dst_rank, 0); + } + } + for (int expert_idx = tid; expert_idx < num_experts; expert_idx += stride) { + const int dst_rank = expert_idx / num_experts_per_rank; + const int local_expert_idx = expert_idx % num_experts_per_rank; + const auto count = static_cast( + layout.get_scaleup_expert_count_ptr()[expert_idx]); + auto* dst = reinterpret_cast( + layout.get_scaleup_expert_count_ptr() + + rank_idx * num_experts_per_rank + local_expert_idx); + const auto encoded_count = + elastic::math::encode_decode_positive(count); + if (dst_rank == rank_idx) { + elastic::ptx::st_release_sys(dst, encoded_count); + } else { + gin.put_value( + dst, encoded_count, dst_rank, 0); + } + } +} + +__global__ void musa_elastic_wait_prefix_counts_kernel( + void* workspace, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, int rank_idx, int num_scaleup_ranks, + int num_experts, int expert_alignment, int64_t timeout_cycles) { + (void)rank_idx; + const auto layout = elastic::layout::WorkspaceLayout( + workspace, 1, num_scaleup_ranks, num_experts); + const int num_experts_per_rank = num_experts / num_scaleup_ranks; + if (threadIdx.x == 0 && blockIdx.x == 0) { + int psum = 0; + for (int src_rank = 0; src_rank < num_scaleup_ranks; ++src_rank) { + auto* ptr = layout.get_scaleup_rank_count_ptr() + src_rank; + const auto start_clock = clock64(); + auto* word_ptr = reinterpret_cast(ptr); + int count = elastic::math::encode_decode_positive( + elastic::ptx::ld_acquire_sys(word_ptr)); + while (!elastic::math::is_decoded_positive_ready(count)) { + if (timeout_cycles >= 0 && + clock64() - start_clock >= timeout_cycles) { + printf("MUSA prepare rank-count timeout, self=%d src=%d decoded=%d\n", + rank_idx, src_rank, count); + count = 0; + break; + } + count = elastic::math::encode_decode_positive( + elastic::ptx::ld_acquire_sys(word_ptr)); + } + *ptr = 0; + psum += count; + psum_num_recv_tokens_per_scaleup_rank[src_rank] = psum; + } + psum = 0; + psum_num_recv_tokens_per_expert[0] = 0; + for (int expert_idx = 0; expert_idx < num_experts_per_rank; + ++expert_idx) { + int count = 0; + for (int src_rank = 0; src_rank < num_scaleup_ranks; ++src_rank) { + auto* ptr = layout.get_scaleup_expert_count_ptr() + + src_rank * num_experts_per_rank + expert_idx; + const auto start_clock = clock64(); + auto* word_ptr = reinterpret_cast(ptr); + int encoded_count = elastic::math::encode_decode_positive( + elastic::ptx::ld_acquire_sys(word_ptr)); + while (!elastic::math::is_decoded_positive_ready(encoded_count)) { + if (timeout_cycles >= 0 && + clock64() - start_clock >= timeout_cycles) { + printf("MUSA prepare expert-count timeout, self=%d src=%d expert=%d decoded=%d\n", + rank_idx, src_rank, expert_idx, encoded_count); + encoded_count = 0; + break; + } + encoded_count = elastic::math::encode_decode_positive( + elastic::ptx::ld_acquire_sys(word_ptr)); + } + count += encoded_count; + *ptr = 0; + } + psum += elastic::math::align(count, expert_alignment); + psum_num_recv_tokens_per_expert[expert_idx + 1] = psum; + } + } +} + +void launch_musa_elastic_prepare_dispatch( + const int64_t* topk_idx, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, int* dst_buffer_slot_idx, + int num_tokens, int num_max_tokens_per_rank, int num_experts, int num_topk, + int expert_alignment, const device::CommCtx& comm_ctx, + const ElasticLaunchContext& ctx, cudaStream_t stream) { + constexpr int kThreads = 256; + const int blocks = std::max(1, std::min(128, ceil_div(num_tokens, kThreads))); + musa_elastic_prepare_init_kernel<<>>( + ctx.workspace, ctx.num_scaleup_ranks, num_experts); + CUDA_RUNTIME_CHECK(cudaGetLastError()); + // Each rank clears its local receive-count slots in the init kernel above. + // Without a cross-rank phase boundary, a fast peer may publish counts into + // this rank while its init kernel is still clearing the same slots, losing + // the peer write. CUDA's cooperative prologue gets this ordering from grid + // synchronization; MUSA needs an explicit scale-up barrier before publish. + musa_elastic_prepare_clear_barrier_kernel<<<1, kThreads, 0, stream>>>( + comm_ctx, ctx.workspace, ctx.scaleup_rank_idx, ctx.num_scaleup_ranks, + num_experts, ctx.timeout_cycles); + CUDA_RUNTIME_CHECK(cudaGetLastError()); + musa_elastic_assign_slots_kernel<<>>( + topk_idx, dst_buffer_slot_idx, ctx.workspace, num_tokens, + num_max_tokens_per_rank, num_experts, num_topk, ctx.num_scaleup_ranks, + ctx.scaleup_rank_idx); + CUDA_RUNTIME_CHECK(cudaGetLastError()); + musa_elastic_publish_counts_kernel<<>>( + comm_ctx, ctx.workspace, ctx.scaleup_rank_idx, ctx.num_scaleup_ranks, + num_experts); + CUDA_RUNTIME_CHECK(cudaGetLastError()); + musa_elastic_wait_prefix_counts_kernel<<<1, 1, 0, stream>>>( + ctx.workspace, psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert, ctx.scaleup_rank_idx, + ctx.num_scaleup_ranks, num_experts, expert_alignment, + ctx.timeout_cycles); + CUDA_RUNTIME_CHECK(cudaGetLastError()); +} + +#endif + +template +void launch_cooperative(Kernel kernel, int num_sms, int num_threads, + int smem_bytes, cudaStream_t stream, Args... args) { +#ifndef MOONCAKE_EP_USE_MUSA + CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes)); +#endif +#ifdef MOONCAKE_EP_USE_MUSA + kernel<<>>(args...); + CUDA_RUNTIME_CHECK(cudaGetLastError()); +#else + cudaLaunchConfig_t cfg = {{num_sms, 1, 1}, {num_threads, 1, 1}, + static_cast(smem_bytes), stream, + nullptr, 0}; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeCooperative; + attr[0].val.cooperative = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + CUDA_RUNTIME_CHECK(cudaLaunchKernelEx(&cfg, kernel, args...)); +#endif +} + +[[noreturn]] void unsupported_elastic_config(const char* op, int hidden, + int num_experts, int num_topk, + int num_max_tokens_per_rank, + int num_sms, int ranks) { + throw std::runtime_error( + std::string("Unsupported Mooncake elastic ") + op + + " static-template config: hidden=" + std::to_string(hidden) + + ", experts=" + std::to_string(num_experts) + + ", topk=" + std::to_string(num_topk) + + ", max_tokens=" + std::to_string(num_max_tokens_per_rank) + + ", num_sms=" + std::to_string(num_sms) + + ", ranks=" + std::to_string(ranks)); +} + +} // namespace + +void launch_elastic_dispatch_deterministic_prologue( + const int64_t* topk_idx, int* rank_count_buffer, int* dst_buffer_slot_idx, + int num_tokens, + int num_max_tokens_per_rank, int num_experts, int num_topk, + int scaleup_rank_idx, int num_scaleup_ranks, int num_sms, + int num_smem_bytes, cudaStream_t stream) { + constexpr int kNumWarps = kElasticNumEpilogueWarps; + constexpr int kNumThreads = kNumWarps * 32; + const int smem_bytes = (1 + 2 * kNumWarps) * num_scaleup_ranks * sizeof(int); + (void)num_smem_bytes; + +#define LAUNCH_PROLOGUE(HIDDEN, EXPERTS, TOPK, MAXTOK, SMS, RANKS) \ + do { \ + auto kernel = elastic::dispatch_deterministic_prologue_impl< \ + SMS, kNumWarps, RANKS, MAXTOK, EXPERTS, TOPK>; \ + launch_cooperative(kernel, SMS, kNumThreads, smem_bytes, stream, \ + const_cast(topk_idx), rank_count_buffer, \ + dst_buffer_slot_idx, num_tokens, \ + scaleup_rank_idx); \ + } while (false) + +#define TRY_PROLOGUE(H, E, K, M, S, R) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + num_scaleup_ranks == R) { \ + LAUNCH_PROLOGUE(H, E, K, M, S, R); \ + return; \ + } + + const int hidden = 0; + (void)hidden; +#ifdef MOONCAKE_EP_USE_MUSA + // Keep the MUSA compile set intentionally small while validating the + // native elastic scale-up path; MUSA non-hybrid dispatch prepares slots in + // a separate kernel and does not call this CUDA cooperative prologue. + TRY_PROLOGUE(0, 256, 8, 128, 24, 2); + TRY_PROLOGUE(0, 256, 8, 128, 24, 8); +#else + // Common production MoE shapes; hidden is irrelevant for this prologue. + TRY_PROLOGUE(0, 256, 8, 128, 24, 8); + TRY_PROLOGUE(0, 256, 8, 128, 24, 2); +#endif + +#undef TRY_PROLOGUE +#undef LAUNCH_PROLOGUE + unsupported_elastic_config("deterministic_prologue", 0, num_experts, + num_topk, num_max_tokens_per_rank, num_sms, + num_scaleup_ranks); +} + +void launch_mooncake_elastic_dispatch( + void* x, void* sf, int64_t* topk_idx, float* topk_weights, + int64_t* copied_topk_idx, int* cumulative_local_expert_recv_stats, + int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, int* dst_buffer_slot_idx, + int* token_metadata_at_forward, int num_tokens, + int num_max_tokens_per_rank, int hidden, int elem_size, int num_sf_packs, + int sf_token_stride, int sf_hidden_stride, int num_experts, int num_topk, + int expert_alignment, int num_sms, int num_channels_per_sm, + int num_smem_bytes, bool cached_mode, bool deterministic, + bool do_cpu_sync, const ElasticLaunchContext& ctx, cudaStream_t stream) { +#ifdef MOONCAKE_EP_USE_MUSA + const bool musa_use_prepared_slots = !cached_mode && ctx.num_scaleout_ranks == 1; +#else + const bool musa_use_prepared_slots = false; +#endif + const bool effective_cached_mode = cached_mode || musa_use_prepared_slots; + const int num_notify_warps = effective_cached_mode ? 0 : kElasticNumNotifyWarps; + const int num_dispatch_warps = kElasticNumDispatchWarps; + const int num_threads = (num_notify_warps + num_dispatch_warps) * 32; + const int smem_bytes = std::max( + num_smem_bytes, + dispatch_smem_bytes(hidden, elem_size, num_sf_packs, num_topk, + ctx.num_scaleup_ranks, num_experts, + num_notify_warps, num_dispatch_warps)); + const bool reuse_slot_indices = effective_cached_mode || deterministic; + const auto comm_ctx = make_comm_ctx(ctx); + (void)num_channels_per_sm; + +#ifdef MOONCAKE_EP_USE_MUSA + if (musa_use_prepared_slots) { + launch_musa_elastic_prepare_dispatch( + topk_idx, psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert, dst_buffer_slot_idx, num_tokens, + num_max_tokens_per_rank, num_experts, num_topk, expert_alignment, + comm_ctx, ctx, stream); + } +#endif + +#ifndef MOONCAKE_EP_USE_MUSA + if (ctx.num_scaleout_ranks != 1) { + const bool hybrid_reuse_slot_indices = cached_mode; + const int hybrid_dispatch_warps = + kElasticNumHybridScaleoutWarps + kElasticNumHybridForwardWarps; + const int hybrid_threads = + (num_notify_warps + hybrid_dispatch_warps) * 32; + const int hybrid_smem_bytes = std::max( + num_smem_bytes, + dispatch_smem_bytes(hidden, elem_size, num_sf_packs, num_topk, + ctx.num_scaleout_ranks * ctx.num_scaleup_ranks, + num_experts, num_notify_warps, + hybrid_dispatch_warps)); + +#define LAUNCH_HYBRID_DISPATCH(HB, SFP, E, K, M, S, SO, SU) \ + do { \ + constexpr int kHiddenBytes = (HB); \ + constexpr int kNumSFPacks = (SFP); \ + if (cached_mode) { \ + auto kernel = elastic::hybrid_dispatch_impl< \ + false, true, S, 0, kElasticNumHybridScaleoutWarps, \ + kElasticNumHybridForwardWarps, SO, SU, kHiddenBytes, \ + kNumSFPacks, M, E, K, 1, kElasticNumQPs, \ + kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, hybrid_threads, \ + hybrid_smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, \ + token_metadata_at_forward, num_tokens, \ + sf_token_stride, sf_hidden_stride, \ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.mapped_host_workspace, \ + ctx.scaleout_rank_idx, \ + ctx.scaleup_rank_idx); \ + } else if (hybrid_reuse_slot_indices) { \ + auto kernel = elastic::hybrid_dispatch_impl< \ + false, true, S, kElasticNumNotifyWarps, \ + kElasticNumHybridScaleoutWarps, \ + kElasticNumHybridForwardWarps, SO, SU, kHiddenBytes, \ + kNumSFPacks, M, E, K, 1, kElasticNumQPs, \ + kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, hybrid_threads, \ + hybrid_smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, \ + token_metadata_at_forward, num_tokens, \ + sf_token_stride, sf_hidden_stride, \ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.mapped_host_workspace, \ + ctx.scaleout_rank_idx, \ + ctx.scaleup_rank_idx); \ + } else { \ + auto kernel = elastic::hybrid_dispatch_impl< \ + false, false, S, kElasticNumNotifyWarps, \ + kElasticNumHybridScaleoutWarps, \ + kElasticNumHybridForwardWarps, SO, SU, kHiddenBytes, \ + kNumSFPacks, M, E, K, 1, kElasticNumQPs, \ + kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, hybrid_threads, \ + hybrid_smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, \ + token_metadata_at_forward, num_tokens, \ + sf_token_stride, sf_hidden_stride, \ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.mapped_host_workspace, \ + ctx.scaleout_rank_idx, \ + ctx.scaleup_rank_idx); \ + } \ + } while (false) + +#define TRY_HYBRID_DISPATCH_TYPED(H, E, K, M, S, SO, SU, EL, SFP) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleout_ranks == SO && ctx.num_scaleup_ranks == SU && \ + elem_size == EL && num_sf_packs == SFP && expert_alignment == 1 && \ + !do_cpu_sync) { \ + LAUNCH_HYBRID_DISPATCH((H) * (EL), SFP, E, K, M, S, SO, SU); \ + return; \ + } + +#define TRY_HYBRID_DISPATCH(H, E, K, M, S, SO, SU) \ + TRY_HYBRID_DISPATCH_TYPED(H, E, K, M, S, SO, SU, \ + static_cast(sizeof(nv_bfloat16)), 0); \ + TRY_HYBRID_DISPATCH_TYPED(H, E, K, M, S, SO, SU, 1, (H) / 128) + +#define TRY_HYBRID_DISPATCH_SHAPE(H, E, K, M, S) \ + TRY_HYBRID_DISPATCH(H, E, K, M, S, 2, 4); \ + TRY_HYBRID_DISPATCH(H, E, K, M, S, 2, 8) + + TRY_HYBRID_DISPATCH_SHAPE(4096, 256, 8, 128, 24); + +#undef TRY_HYBRID_DISPATCH_SHAPE +#undef TRY_HYBRID_DISPATCH +#undef TRY_HYBRID_DISPATCH_TYPED +#undef LAUNCH_HYBRID_DISPATCH + } +#endif + +#define LAUNCH_DISPATCH(HB, SFP, E, K, M, S, R) \ + do { \ + constexpr int kHiddenBytes = (HB); \ + constexpr int kNumSFPacks = (SFP); \ + if (effective_cached_mode) { \ + auto kernel = elastic::dispatch_impl< \ + true, false, true, S, 0, kElasticNumDispatchWarps, R, \ + kHiddenBytes, kNumSFPacks, M, E, K, 1, kElasticNumQPs, \ + kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, num_tokens, sf_token_stride,\ + sf_hidden_stride, comm_ctx, ctx.buffer, \ + ctx.workspace, ctx.mapped_host_workspace, \ + ctx.scaleup_rank_idx); \ + } else if (reuse_slot_indices) { \ + auto kernel = elastic::dispatch_impl< \ + true, false, true, S, kElasticNumNotifyWarps, \ + kElasticNumDispatchWarps, R, kHiddenBytes, kNumSFPacks, M, E, K, 1, \ + kElasticNumQPs, kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, num_tokens, sf_token_stride,\ + sf_hidden_stride, comm_ctx, ctx.buffer, \ + ctx.workspace, ctx.mapped_host_workspace, \ + ctx.scaleup_rank_idx); \ + } else { \ + auto kernel = elastic::dispatch_impl< \ + true, false, false, S, kElasticNumNotifyWarps, \ + kElasticNumDispatchWarps, R, kHiddenBytes, kNumSFPacks, M, E, K, 1, \ + kElasticNumQPs, kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, x, \ + static_cast(sf), topk_idx, \ + topk_weights, copied_topk_idx, \ + cumulative_local_expert_recv_stats, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, \ + dst_buffer_slot_idx, num_tokens, sf_token_stride,\ + sf_hidden_stride, comm_ctx, ctx.buffer, \ + ctx.workspace, ctx.mapped_host_workspace, \ + ctx.scaleup_rank_idx); \ + } \ + } while (false) + +#define TRY_DISPATCH_TYPED(H, E, K, M, S, R, EL, SFP) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleup_ranks == R && elem_size == EL && \ + num_sf_packs == SFP && expert_alignment == 1 && !do_cpu_sync) { \ + LAUNCH_DISPATCH((H) * (EL), SFP, E, K, M, S, R); \ + return; \ + } + +#define TRY_DISPATCH(H, E, K, M, S, R) \ + TRY_DISPATCH_TYPED(H, E, K, M, S, R, static_cast(sizeof(nv_bfloat16)), 0); \ + TRY_DISPATCH_TYPED(H, E, K, M, S, R, 1, (H) / 128) + +#ifdef MOONCAKE_EP_USE_MUSA + TRY_DISPATCH_TYPED(4096, 256, 8, 128, 24, 2, + static_cast(sizeof(nv_bfloat16)), 0); + TRY_DISPATCH_TYPED(4096, 256, 8, 128, 24, 8, + static_cast(sizeof(nv_bfloat16)), 0); +#else + TRY_DISPATCH(4096, 256, 8, 128, 24, 8); + TRY_DISPATCH(4096, 256, 8, 128, 24, 2); +#endif + +#undef TRY_DISPATCH +#undef TRY_DISPATCH_TYPED +#undef LAUNCH_DISPATCH + unsupported_elastic_config("dispatch", hidden, num_experts, num_topk, + num_max_tokens_per_rank, num_sms, + ctx.num_scaleup_ranks); +} + +void launch_mooncake_elastic_dispatch_copy_epilogue( + void* recv_x, void* recv_sf, int64_t* recv_topk_idx, + float* recv_topk_weights, int* recv_src_metadata, + int* channel_linked_list, int num_recv_tokens, int num_max_tokens_per_rank, + int hidden, int elem_size, int num_sf_packs, int recv_sf_token_stride, + int recv_sf_hidden_stride, int num_experts, int num_topk, int num_sms, + int num_smem_bytes, int num_channels, bool do_expand, bool cached_mode, + const ElasticLaunchContext& ctx, int* psum_num_recv_tokens_per_scaleup_rank, + int* psum_num_recv_tokens_per_expert, cudaStream_t stream) { + const int num_threads = kElasticNumEpilogueWarps * 32; + const int smem_bytes = std::max( + num_smem_bytes, + dispatch_epilogue_smem_bytes(hidden, elem_size, num_sf_packs, num_topk, + kElasticNumEpilogueWarps)); + +#ifndef MOONCAKE_EP_USE_MUSA + if (ctx.num_scaleout_ranks != 1) { +#define LAUNCH_HYBRID_DISPATCH_EPILOGUE(HB, SFP, E, K, M, S, SO, SU, C) \ + do { \ + constexpr int kHiddenBytes = (HB); \ + constexpr int kNumSFPacks = (SFP); \ + auto kernel = do_expand ? \ + elastic::dispatch_copy_epilogue_impl< \ + true, false, S, C, kElasticNumEpilogueWarps, SO, SU, \ + kHiddenBytes, kNumSFPacks, M, E, K> : \ + (cached_mode ? \ + elastic::dispatch_copy_epilogue_impl< \ + false, true, S, C, kElasticNumEpilogueWarps, SO, SU, \ + kHiddenBytes, kNumSFPacks, M, E, K> : \ + elastic::dispatch_copy_epilogue_impl< \ + false, false, S, C, kElasticNumEpilogueWarps, SO, SU, \ + kHiddenBytes, kNumSFPacks, M, E, K>); \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, \ + ctx.buffer, ctx.workspace, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, recv_x, \ + static_cast(recv_sf), \ + recv_topk_idx, recv_topk_weights, \ + recv_src_metadata, channel_linked_list, \ + num_recv_tokens, recv_sf_token_stride, \ + recv_sf_hidden_stride, ctx.scaleout_rank_idx, \ + ctx.scaleup_rank_idx); \ + } while (false) + +#define TRY_HYBRID_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, SO, SU, EL, SFP) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleout_ranks == SO && ctx.num_scaleup_ranks == SU && \ + elem_size == EL && num_sf_packs == SFP && \ + num_channels == hybrid_num_channels(S)) { \ + LAUNCH_HYBRID_DISPATCH_EPILOGUE((H) * (EL), SFP, E, K, M, S, SO, SU, \ + (S) * kElasticNumHybridForwardWarps); \ + return; \ + } + +#define TRY_HYBRID_DISPATCH_EPILOGUE(H, E, K, M, S, SO, SU) \ + TRY_HYBRID_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, SO, SU, \ + static_cast(sizeof(nv_bfloat16)), 0); \ + TRY_HYBRID_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, SO, SU, 1, (H) / 128) + +#define TRY_HYBRID_DISPATCH_EPILOGUE_SHAPE(H, E, K, M, S) \ + TRY_HYBRID_DISPATCH_EPILOGUE(H, E, K, M, S, 2, 4); \ + TRY_HYBRID_DISPATCH_EPILOGUE(H, E, K, M, S, 2, 8) + + TRY_HYBRID_DISPATCH_EPILOGUE_SHAPE(4096, 256, 8, 128, 24); + +#undef TRY_HYBRID_DISPATCH_EPILOGUE_SHAPE +#undef TRY_HYBRID_DISPATCH_EPILOGUE +#undef TRY_HYBRID_DISPATCH_EPILOGUE_TYPED +#undef LAUNCH_HYBRID_DISPATCH_EPILOGUE + } +#endif + +#define LAUNCH_DISPATCH_EPILOGUE(HB, SFP, E, K, M, S, R) \ + do { \ + constexpr int kHiddenBytes = (HB); \ + constexpr int kNumSFPacks = (SFP); \ + auto kernel = do_expand ? \ + elastic::dispatch_copy_epilogue_impl< \ + true, false, S, 1, kElasticNumEpilogueWarps, 1, R, \ + kHiddenBytes, kNumSFPacks, M, E, K> : \ + (cached_mode ? \ + elastic::dispatch_copy_epilogue_impl< \ + false, true, S, 1, kElasticNumEpilogueWarps, 1, R, \ + kHiddenBytes, kNumSFPacks, M, E, K> : \ + elastic::dispatch_copy_epilogue_impl< \ + false, false, S, 1, kElasticNumEpilogueWarps, 1, R, \ + kHiddenBytes, kNumSFPacks, M, E, K>); \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, \ + ctx.buffer, ctx.workspace, \ + psum_num_recv_tokens_per_scaleup_rank, \ + psum_num_recv_tokens_per_expert, recv_x, \ + static_cast(recv_sf), recv_topk_idx, \ + recv_topk_weights, recv_src_metadata, \ + channel_linked_list, num_recv_tokens, \ + recv_sf_token_stride, recv_sf_hidden_stride, \ + ctx.scaleout_rank_idx, ctx.scaleup_rank_idx); \ + } while (false) + +#define TRY_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, R, EL, SFP) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleup_ranks == R && elem_size == EL && \ + num_sf_packs == SFP) { \ + LAUNCH_DISPATCH_EPILOGUE((H) * (EL), SFP, E, K, M, S, R); \ + return; \ + } + +#define TRY_DISPATCH_EPILOGUE(H, E, K, M, S, R) \ + TRY_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, R, static_cast(sizeof(nv_bfloat16)), 0); \ + TRY_DISPATCH_EPILOGUE_TYPED(H, E, K, M, S, R, 1, (H) / 128) + +#ifdef MOONCAKE_EP_USE_MUSA + TRY_DISPATCH_EPILOGUE_TYPED(4096, 256, 8, 128, 24, 2, + static_cast(sizeof(nv_bfloat16)), 0); + TRY_DISPATCH_EPILOGUE_TYPED(4096, 256, 8, 128, 24, 8, + static_cast(sizeof(nv_bfloat16)), 0); +#else + TRY_DISPATCH_EPILOGUE(4096, 256, 8, 128, 24, 8); + TRY_DISPATCH_EPILOGUE(4096, 256, 8, 128, 24, 2); +#endif + +#undef TRY_DISPATCH_EPILOGUE +#undef TRY_DISPATCH_EPILOGUE_TYPED +#undef LAUNCH_DISPATCH_EPILOGUE + unsupported_elastic_config("dispatch_copy_epilogue", hidden, num_experts, + num_topk, num_max_tokens_per_rank, num_sms, + ctx.num_scaleup_ranks); +} + +void* launch_mooncake_elastic_combine( + void* x, float* topk_weights, int* src_metadata, + int* psum_num_recv_tokens_per_scaleup_rank, + int* token_metadata_at_forward, int* channel_linked_list, + int num_reduced_tokens, int num_max_tokens_per_rank, int hidden, + int num_experts, int num_topk, int num_sms, int num_smem_bytes, + int num_channels, bool use_expanded_layout, bool allow_multiple_reduction, + const ElasticLaunchContext& ctx, cudaStream_t stream) { + const int num_threads = kElasticNumEpilogueWarps * 32; + const int smem_bytes = std::max( + num_smem_bytes, combine_smem_bytes(hidden, num_topk, kElasticNumEpilogueWarps)); + const auto comm_ctx = make_comm_ctx(ctx); + (void)token_metadata_at_forward; + (void)channel_linked_list; + +#ifndef MOONCAKE_EP_USE_MUSA + if (ctx.num_scaleout_ranks != 1) { + const int hybrid_combine_warps = + kElasticNumHybridScaleupWarps + kElasticNumHybridForwardWarps; + const int hybrid_threads = hybrid_combine_warps * 32; + const int hybrid_smem_bytes = std::max( + num_smem_bytes, + combine_smem_bytes(hidden, num_topk, hybrid_combine_warps)); + +#define LAUNCH_HYBRID_COMBINE(H, E, K, M, S, SO, SU) \ + do { \ + auto kernel = elastic::hybrid_combine_impl< \ + false, true, S, kElasticNumHybridScaleupWarps, \ + kElasticNumHybridForwardWarps, SO, SU, H, M, E, K, \ + kElasticNumQPs, kElasticTimeoutCycles>; \ + launch_cooperative(kernel, S, hybrid_threads, hybrid_smem_bytes, \ + stream, static_cast(x), \ + topk_weights, src_metadata, \ + psum_num_recv_tokens_per_scaleup_rank, \ + token_metadata_at_forward, channel_linked_list, \ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.scaleout_rank_idx, ctx.scaleup_rank_idx, \ + num_reduced_tokens); \ + } while (false) + +#define TRY_HYBRID_COMBINE(H, E, K, M, S, SO, SU) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleout_ranks == SO && ctx.num_scaleup_ranks == SU && \ + allow_multiple_reduction && !use_expanded_layout && \ + num_channels == hybrid_num_channels(S) && \ + token_metadata_at_forward != nullptr && channel_linked_list != nullptr) { \ + LAUNCH_HYBRID_COMBINE(H, E, K, M, S, SO, SU); \ + return hybrid_combine_reduce_buffer_ptr( \ + ctx.buffer, H, K, M, SO, SU, allow_multiple_reduction); \ + } + +#define TRY_HYBRID_COMBINE_SHAPE(H, E, K, M, S) \ + TRY_HYBRID_COMBINE(H, E, K, M, S, 2, 4); \ + TRY_HYBRID_COMBINE(H, E, K, M, S, 2, 8) + + TRY_HYBRID_COMBINE_SHAPE(4096, 256, 8, 128, 24); + +#undef TRY_HYBRID_COMBINE_SHAPE +#undef TRY_HYBRID_COMBINE +#undef LAUNCH_HYBRID_COMBINE + } +#endif + + (void)num_channels; + +#define LAUNCH_COMBINE(H, E, K, M, S, R) \ + do { \ + auto kernel = elastic::combine_impl; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, \ + static_cast(x), topk_weights, \ + src_metadata, psum_num_recv_tokens_per_scaleup_rank,\ + comm_ctx, ctx.buffer, ctx.workspace, \ + ctx.scaleup_rank_idx, num_reduced_tokens); \ + } while (false) + +#define TRY_COMBINE(H, E, K, M, S, R) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleup_ranks == R && !use_expanded_layout && \ + allow_multiple_reduction) { \ + LAUNCH_COMBINE(H, E, K, M, S, R); \ + return ctx.buffer; \ + } + +#ifdef MOONCAKE_EP_USE_MUSA + TRY_COMBINE(4096, 256, 8, 128, 24, 2); + TRY_COMBINE(4096, 256, 8, 128, 24, 8); +#else + TRY_COMBINE(4096, 256, 8, 128, 24, 8); + TRY_COMBINE(4096, 256, 8, 128, 24, 2); +#endif + +#undef TRY_COMBINE +#undef LAUNCH_COMBINE + unsupported_elastic_config("combine", hidden, num_experts, num_topk, + num_max_tokens_per_rank, num_sms, + ctx.num_scaleup_ranks); +} + +void launch_mooncake_elastic_combine_reduce_epilogue( + void* combined_x, float* combined_topk_weights, int64_t* combined_topk_idx, + int num_combined_tokens, int num_max_tokens_per_rank, int hidden, + int num_experts, int num_topk, void* reduce_buffer, void* bias_0, + void* bias_1, int num_sms, int num_smem_bytes, bool use_expanded_layout, + bool allow_multiple_reduction, const ElasticLaunchContext& ctx, + cudaStream_t stream) { + const int num_threads = kElasticNumEpilogueWarps * 32; + const int smem_bytes = std::max( + num_smem_bytes, combine_epilogue_smem_bytes(hidden, kElasticNumEpilogueWarps)); + +#define LAUNCH_COMBINE_EPILOGUE(H, E, K, M, S, SO, SU) \ + do { \ + auto kernel = elastic::combine_reduce_epilogue_impl< \ + false, true, S, kElasticNumEpilogueWarps, SO, SU, H, M, E, K>; \ + launch_cooperative(kernel, S, num_threads, smem_bytes, stream, \ + static_cast(combined_x), \ + combined_topk_weights, combined_topk_idx, \ + reduce_buffer, bias_0, bias_1, num_combined_tokens, \ + ctx.scaleout_rank_idx, ctx.scaleup_rank_idx); \ + } while (false) + +#define TRY_COMBINE_EPILOGUE(H, E, K, M, S, SO, SU) \ + if (hidden == H && num_experts == E && num_topk == K && \ + num_max_tokens_per_rank == M && num_sms == S && \ + ctx.num_scaleout_ranks == SO && ctx.num_scaleup_ranks == SU && \ + !use_expanded_layout && allow_multiple_reduction) { \ + LAUNCH_COMBINE_EPILOGUE(H, E, K, M, S, SO, SU); \ + return; \ + } + +#ifdef MOONCAKE_EP_USE_MUSA + TRY_COMBINE_EPILOGUE(4096, 256, 8, 128, 24, 1, 2); + TRY_COMBINE_EPILOGUE(4096, 256, 8, 128, 24, 1, 8); +#else + TRY_COMBINE_EPILOGUE(4096, 256, 8, 128, 24, 1, 8); + TRY_COMBINE_EPILOGUE(4096, 256, 8, 128, 24, 1, 2); + +#define TRY_HYBRID_COMBINE_EPILOGUE_SHAPE(H, E, K, M, S) \ + TRY_COMBINE_EPILOGUE(H, E, K, M, S, 2, 4); \ + TRY_COMBINE_EPILOGUE(H, E, K, M, S, 2, 8) + + TRY_HYBRID_COMBINE_EPILOGUE_SHAPE(4096, 256, 8, 128, 24); +#endif + +#undef TRY_HYBRID_COMBINE_EPILOGUE_SHAPE + +#undef TRY_COMBINE_EPILOGUE +#undef LAUNCH_COMBINE_EPILOGUE + unsupported_elastic_config("combine_reduce_epilogue", hidden, num_experts, + num_topk, num_max_tokens_per_rank, num_sms, + ctx.num_scaleup_ranks); +} + +} // namespace mooncake diff --git a/mooncake-ep/src/mooncake_ep_kernel.cu b/mooncake-ep/src/mooncake_ep_kernel.cu index 0ef1509784..bc96441181 100644 --- a/mooncake-ep/src/mooncake_ep_kernel.cu +++ b/mooncake-ep/src/mooncake_ep_kernel.cu @@ -28,6 +28,15 @@ using mooncake::device::mc_atomic_add_release; using mooncake::device::mc_fence; using mooncake::device::mc_fence_barrier_fence; +__device__ __forceinline__ int ep_qp_channel(int expert_local_idx, + int qps_per_rank, + int active_qps_per_rank) { + int active_qps = active_qps_per_rank; + if (active_qps <= 0 || active_qps > qps_per_rank) + active_qps = qps_per_rank; + return expert_local_idx % active_qps; +} + __global__ void mark_phase_ack_kernel(void* mxa_buffer, const int32_t* nvlink_available, void* const* ipc_peer_ptrs, @@ -143,7 +152,7 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales, int num_tokens, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, int64_t timeout_ticks, - int phases) { + int phases, int active_qps_per_rank) { const auto sm_id = static_cast(blockIdx.x); const auto thread_id = static_cast(threadIdx.x); const auto warp_id = thread_id / 32, lane_id = get_lane_id(); @@ -153,6 +162,22 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales, const auto warp_group_id = warp_id / kNumWarpsPerGroup; const auto sub_warp_id = warp_id % kNumWarpsPerGroup; const auto responsible_expert_idx = sm_id * kNumWarpGroups + warp_group_id; +#ifdef MOONCAKE_EP_USE_MACA + // C500 reports 64-thread hardware warps. Do not split the last hardware + // warp by assigning only the final 32-thread pseudo-warp to count work. + // Reserve one full warp group from the data path, but write counts from a + // single 32-thread lane group to avoid duplicate per-expert increments. + const bool is_count_warp = warp_group_id == kNumWarpGroups - 1; + const bool is_count_worker = is_count_warp && sub_warp_id == 0; + const bool is_data_warp = warp_group_id < kNumWarpGroups - 1; + const int num_send_threads = + (kNumWarpGroups - 1) * kNumWarpsPerGroup * 32; +#else + const bool is_count_warp = warp_id == num_warps - 1; + const bool is_count_worker = is_count_warp; + const bool is_data_warp = warp_id < num_warps - 1; + const int num_send_threads = (num_warps - 1) * 32; +#endif // FP8 staffs constexpr int kNumPerChannels = 128; @@ -183,14 +208,16 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales, // Expert counts __shared__ int shared_num_tokens_sent_per_expert[kNumWarpGroups]; - // There are 2 kinds of warps in this part: - // 1. The first-kind warps for FP8 cast and sending top-k tokens - // 2. The last warp for reading `topk_idx` and count for per-expert information - if (warp_id < num_warps - 1) { + // There are 2 kinds of execution lanes in this part: + // 1. Data lanes for FP8 cast and sending top-k tokens. + // 2. Count lanes for reading `topk_idx` and per-expert token counts. + // MACA reserves a full warp group for the count path; CUDA keeps the + // original final 32-thread warp behavior. + if (is_data_warp) { constexpr int kNumElemsPerRead = sizeof(int4) / EP_BF16_SIZE; EP_DEVICE_ASSERT(kHidden % kNumElemsPerRead == 0); EP_STATIC_ASSERT(kNumElemsPerRead * 32 % kNumPerChannels == 0, "Invalid vectorization"); - const auto num_threads = (num_warps - 1) * 32; + const auto num_threads = num_send_threads; const size_t hidden_bf16_int4 = kHidden / kNumElemsPerRead; for (int token_idx = sm_id; token_idx < num_tokens; token_idx += num_sms) { @@ -265,8 +292,12 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales, mc_fence(); } else { // IBGDA path — send directly from source buffer - mc_rdma_put(comm_ctx, dst_expert_local_idx % num_qp_per_rank, dst_rank, num_qp_per_rank, - src_ptr, dst_ptr, num_bytes_per_msg, lane_id); + mc_rdma_put(comm_ctx, + ep_qp_channel(dst_expert_local_idx, + num_qp_per_rank, + active_qps_per_rank), + dst_rank, num_qp_per_rank, src_ptr, dst_ptr, + num_bytes_per_msg, lane_id); } // Increase counter after finishing @@ -274,15 +305,17 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales, lane_id == 0 ? mc_atomic_add_release(atomic_finish_counter_per_expert + dst_expert_idx, 1) : 0; } } - } else if (warp_id == num_warps - 1) { -#ifdef MOONCAKE_EP_USE_MUSA + } else if (is_count_warp) { +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV // Participate in __syncthreads() barriers from data warps. // Each token iteration in the send loop above calls - // __syncthreads() once; the count warp must match. + // __syncthreads() once; the count path must match. for (int token_idx = sm_id; token_idx < num_tokens; token_idx += num_sms) { __syncthreads(); } #endif + } + if (is_count_worker) { EP_DEVICE_ASSERT(num_sms > 1); if (sm_id == 0) { // The first SM is also responsible for cleaning the next buffer @@ -332,8 +365,11 @@ dispatch(void* packed_recv_x, float* packed_recv_x_scales, while (mc_ld_acquire(atomic_finish_counter_per_expert + responsible_expert_idx) != FINISHED_SUM_TAG * 2); if (dst_rank != rank) { int* signal_ptr = rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank; - mc_red_add(comm_ctx, dst_rank, dst_expert_local_idx % num_qp_per_rank, num_qp_per_rank, - signal_ptr, static_cast(-num_tokens_sent - 1)); + mc_red_add(comm_ctx, dst_rank, + ep_qp_channel(dst_expert_local_idx, num_qp_per_rank, + active_qps_per_rank), + num_qp_per_rank, signal_ptr, + static_cast(-num_tokens_sent - 1)); } else { mc_st_release(rdma_recv_signal_buffer + dst_expert_local_idx * num_ranks + rank, -num_tokens_sent - 1); } @@ -445,7 +481,8 @@ void dispatch(void* packed_recv_x, float* packed_recv_x_scales, int* next_clean_buffer, int num_tokens, int hidden, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, bool use_fp8, - void* workspace, cudaStream_t stream, int64_t timeout_ticks, int phases) { + void* workspace, cudaStream_t stream, int64_t timeout_ticks, + int phases, int active_qps_per_rank) { constexpr int kNumMaxTopK = 11; constexpr int kNumWarpsPerGroup = 4; #ifdef MOONCAKE_EP_USE_MUSA @@ -482,7 +519,8 @@ LAUNCH_KERNEL(&cfg, dispatch_func, \ atomic_counter_per_expert, atomic_finish_counter_per_expert, \ next_clean_buffer, \ num_tokens, num_max_dispatch_tokens_per_rank, \ - num_topk, num_experts, rank, num_ranks, timeout_ticks, phases); } break + num_topk, num_experts, rank, num_ranks, timeout_ticks, phases, \ + active_qps_per_rank); } break SETUP_LAUNCH_CONFIG(num_sms, num_warps * 32, stream); SWITCH_HIDDEN(DISPATCH_LAUNCH_CASE); @@ -506,7 +544,7 @@ combine(void* combined_x, int32_t* active_ranks, int num_max_dispatch_tokens_per_rank, int num_experts, int rank, int num_ranks, int64_t timeout_ticks, - int phases, bool zero_copy) { + int phases, bool zero_copy, int active_qps_per_rank) { const auto sm_id = static_cast(blockIdx.x); const auto num_sms = static_cast(gridDim.x); const auto thread_id = static_cast(threadIdx.x); @@ -590,8 +628,11 @@ combine(void* combined_x, int32_t* active_ranks, if (not zero_copy) UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, buf_int4_ptr, x_int4, mc_ld_nc, mc_st_na); __syncwarp(); - mc_rdma_put(comm_ctx, local_expert_idx % num_qp_per_rank, dst_rank, num_qp_per_rank, - buf_ptr, dst_ptr, num_bytes_per_slot, lane_id); + mc_rdma_put(comm_ctx, + ep_qp_channel(local_expert_idx, num_qp_per_rank, + active_qps_per_rank), + dst_rank, num_qp_per_rank, buf_ptr, dst_ptr, + num_bytes_per_slot, lane_id); } } // Put finishing flag @@ -601,7 +642,10 @@ combine(void* combined_x, int32_t* active_ranks, while (mc_ld_acquire(atomic_clean_flag) == 0); if (dst_rank != rank) { int* signal_ptr = rdma_recv_signal_buffer + global_expert_idx; - mc_signal(comm_ctx, dst_rank, local_expert_idx % num_qp_per_rank, num_qp_per_rank, signal_ptr, 1); + mc_signal(comm_ctx, dst_rank, + ep_qp_channel(local_expert_idx, num_qp_per_rank, + active_qps_per_rank), + num_qp_per_rank, signal_ptr, 1); } else { mc_st_release(rdma_recv_signal_buffer + global_expert_idx, 1); } @@ -634,9 +678,9 @@ combine(void* combined_x, int32_t* active_ranks, } } } -#ifdef MOONCAKE_EP_USE_MUSA - // mc_grid_sync() is a no-op on MUSA; use a block-wide fence/barrier before - // reduction so threads see peer writes. +#ifdef MOONCAKE_EP_SPLIT_SEND_RECV + // mc_grid_sync() is a no-op on split-kernel platforms; use a block-wide + // fence/barrier before reduction so threads see peer writes. __syncthreads(); mc_fence(); __syncthreads(); @@ -662,6 +706,10 @@ combine(void* combined_x, int32_t* active_ranks, float combined_values[kNumElemsPerInt4] = {0.0f}; #pragma unroll for (int i = 0; i < num_topk; ++ i) if (reg_topk_idx[i] >= 0) { + // Skip experts on inactive ranks (timed out during combine recv) + int expert_src_rank = reg_topk_idx[i] / num_local_experts; + if (!active_ranks[expert_src_rank]) + continue; // Read from sources auto rdma_buffer_type = reinterpret_cast(reinterpret_cast(rdma_recv_data_buffer) + (reg_topk_idx[i] * num_max_dispatch_tokens_per_rank + token_idx) * num_bytes_per_slot); auto rdma_buffer_row = reinterpret_cast(rdma_buffer_type); @@ -698,7 +746,8 @@ void combine(void* combined_x, int32_t* active_ranks, int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank, int num_topk, int num_experts, int rank, int num_ranks, void* workspace, cudaStream_t stream, - int64_t timeout_ticks, int phases, bool zero_copy) { + int64_t timeout_ticks, int phases, bool zero_copy, + int active_qps_per_rank) { constexpr int kNumWarpsPerGroup = 4; constexpr int kNumWarpGroups = 8; constexpr int kNumMaxTopk = 11; @@ -727,7 +776,7 @@ LAUNCH_KERNEL(&cfg, combine_func, \ num_combined_tokens, hidden, num_topk, \ num_max_dispatch_tokens_per_rank, \ num_experts, rank, num_ranks, \ - timeout_ticks, phases, zero_copy); } break + timeout_ticks, phases, zero_copy, active_qps_per_rank); } break SETUP_LAUNCH_CONFIG(num_sms, num_warps * 32, stream); SWITCH_HIDDEN(COMBINE_LAUNCH_CASE); diff --git a/mooncake-ep/tests/test_elastic_buffer.py b/mooncake-ep/tests/test_elastic_buffer.py new file mode 100644 index 0000000000..53f049303c --- /dev/null +++ b/mooncake-ep/tests/test_elastic_buffer.py @@ -0,0 +1,412 @@ +#!/usr/bin/env python3 +"""Correctness smoke for the public Mooncake ElasticBuffer API. + +This test is intentionally self-contained: it exercises the PR's elastic +dispatch/combine path through the public ElasticBuffer wrapper and checks the +result against a deterministic PyTorch reference derived from the routing +metadata. + +Typical single-node usage: + + MOONCAKE_EP_NUM_LOCAL_RANKS=8 \ + torchrun --standalone --nproc_per_node=8 \ + mooncake-ep/tests/test_elastic_buffer.py --quick + +For multi-node hybrid validation, launch the same script with ``torchrun +--nnodes`` and set ``MOONCAKE_EP_NUM_LOCAL_RANKS`` to the number of GPUs per +node. +""" + +from __future__ import annotations + +import argparse +import os +from dataclasses import dataclass + +import torch +import torch.distributed as dist +import torch.testing as testing + +from mooncake.mooncake_elastic_buffer import ElasticBuffer + + +def using_musa_backend() -> bool: + return os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in { + "1", + "ON", + "TRUE", + "YES", + } + + +def import_torchada_if_needed() -> None: + if not using_musa_backend(): + return + import torchada # noqa: F401 — maps torch.cuda.* to torch.musa.* on MUSA + + +def distributed_barrier() -> None: + if using_musa_backend(): + dist.barrier(device_ids=[torch.cuda.current_device()]) + else: + dist.barrier() + + +@dataclass(frozen=True) +class RoutePlan: + topk_idx: torch.Tensor + expected_recv_tokens: int + expected_combine_factor: int + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Test Mooncake ElasticBuffer") + parser.add_argument("--num-tokens", type=int, default=64) + parser.add_argument("--max-tokens", type=int, default=0) + parser.add_argument("--hidden", type=int, default=4096) + parser.add_argument("--num-experts", type=int, default=256) + parser.add_argument("--num-topk", type=int, default=8) + parser.add_argument("--num-sms", type=int, default=24) + parser.add_argument( + "--route", + choices=("alltoall", "local", "cross"), + default="alltoall", + help="Expert routing pattern to generate.", + ) + parser.add_argument( + "--allow-hybrid-mode", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument("--seed", type=int, default=2026) + parser.add_argument( + "--quick", + action="store_true", + help="Use smaller defaults suitable for reviewer smoke tests.", + ) + return parser.parse_args() + + +def init_distributed(seed: int) -> tuple[int, int, int]: + import_torchada_if_needed() + if not dist.is_initialized(): + dist.init_process_group("nccl") + + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", rank % torch.cuda.device_count())) + torch.cuda.set_device(local_rank) + torch.set_default_device("cuda") + torch.set_default_dtype(torch.bfloat16) + torch.manual_seed(seed + rank) + return rank, local_rank, world_size + + +def make_route_plan( + *, + rank: int, + world_size: int, + buffer: ElasticBuffer, + num_tokens: int, + num_topk: int, + num_experts: int, + route: str, +) -> RoutePlan: + local_experts = num_experts // world_size + if local_experts <= 0: + raise ValueError("num_experts must be at least world_size") + + expert_offsets = torch.arange(num_topk, device="cuda", dtype=torch.long) % local_experts + + if route == "cross" and buffer.num_scaleout_ranks > 1: + dst_scaleout = (buffer.scaleout_rank_idx + 1) % buffer.num_scaleout_ranks + dst_rank = dst_scaleout * buffer.num_scaleup_ranks + buffer.scaleup_rank_idx + choices = dst_rank * local_experts + expert_offsets + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens, + 1, + ) + + if route == "local" or (route == "cross" and buffer.num_scaleout_ranks == 1): + choices = rank * local_experts + expert_offsets + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens, + 1, + ) + + dst_ranks = (rank + torch.arange(num_topk, device="cuda", dtype=torch.long)) % world_size + choices = dst_ranks * local_experts + expert_offsets + unique_dst_ranks = int(torch.unique(dst_ranks).numel()) + return RoutePlan( + choices.view(1, num_topk).repeat(num_tokens, 1).contiguous(), + num_tokens * unique_dst_ranks, + unique_dst_ranks, + ) + + +def make_input( + *, rank: int, num_tokens: int, hidden: int, multiplier: int, addend: int = 0 +) -> torch.Tensor: + base = torch.arange(num_tokens * hidden, device="cuda", dtype=torch.float32) + base = base.view(num_tokens, hidden) + return (base + rank * multiplier + addend).to(torch.bfloat16).contiguous() + + +def check_dispatch_payload( + *, + rank: int, + recv_x: torch.Tensor, + handle, + expected_recv_tokens: int, + max_tokens: int, + num_tokens: int, + hidden: int, + multiplier: int, + addend: int = 0, +) -> int: + actual = int(handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()) + if actual != expected_recv_tokens: + raise AssertionError( + f"rank={rank}: got {actual} received tokens, " + f"expected {expected_recv_tokens}" + ) + + src_global = handle.recv_src_metadata[:actual, 0].long() + src_rank = torch.div(src_global, max_tokens, rounding_mode="floor") + src_token = src_global % max_tokens + if not bool((src_token < num_tokens).all()): + raise AssertionError(f"rank={rank}: invalid source token index in metadata") + + base = torch.arange(num_tokens * hidden, device="cuda", dtype=torch.float32) + base = base.view(num_tokens, hidden) + expected = (base[src_token] + src_rank.view(-1, 1).float() * multiplier + addend) + expected = expected.to(torch.bfloat16) + if not torch.equal(recv_x[:actual], expected): + diff = (recv_x[:actual].float() - expected.float()).abs().max().item() + raise AssertionError(f"rank={rank}: dispatch payload mismatch, max_diff={diff}") + return actual + + +def check_combined( + *, rank: int, combined: torch.Tensor, expected: torch.Tensor, label: str +) -> None: + testing.assert_close( + combined, + expected, + rtol=5e-2, + atol=1e-3, + msg=lambda msg: f"rank={rank}: {label} combine mismatch: {msg}", + ) + + +def main() -> None: + args = parse_args() + if args.quick: + args.num_tokens = min(args.num_tokens, 32) + + rank, _local_rank, world_size = init_distributed(args.seed) + max_tokens = args.max_tokens or max(128, args.num_tokens) + num_experts = args.num_experts + + if num_experts % world_size != 0: + raise ValueError("num_experts must be divisible by world_size") + + buffer = ElasticBuffer( + dist.group.WORLD, + num_max_tokens_per_rank=max_tokens, + hidden=args.hidden, + num_topk=args.num_topk, + use_fp8_dispatch=False, + deterministic=False, + allow_hybrid_mode=args.allow_hybrid_mode, + allow_multiple_reduction=True, + num_gpu_timeout_secs=10, + ) + + route_plan = make_route_plan( + rank=rank, + world_size=world_size, + buffer=buffer, + num_tokens=args.num_tokens, + num_topk=args.num_topk, + num_experts=num_experts, + route=args.route, + ) + weights = torch.ones((args.num_tokens, args.num_topk), device="cuda", dtype=torch.float32) + + # CPU-sync dispatch: exact output extent and CPU-side expert counts. + x0 = make_input( + rank=rank, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=1_000_000, + ) + recv0, idx0, w0, handle0, _ = buffer.dispatch( + x0, + topk_idx=route_plan.topk_idx, + topk_weights=weights, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + expert_alignment=1, + do_cpu_sync=True, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + actual0 = check_dispatch_payload( + rank=rank, + recv_x=recv0, + handle=handle0, + expected_recv_tokens=route_plan.expected_recv_tokens, + max_tokens=max_tokens, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=1_000_000, + ) + if recv0.shape[0] != actual0: + raise AssertionError(f"rank={rank}: CPU-sync dispatch returned extra rows") + if len(handle0.num_recv_tokens_per_expert_list) != num_experts // world_size: + raise AssertionError(f"rank={rank}: invalid per-expert count length") + + combined0, _, _ = buffer.combine( + recv0.contiguous(), + handle0, + topk_weights=w0[:actual0].contiguous(), + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + expected0 = (x0.float() * route_plan.expected_combine_factor).to(torch.bfloat16) + check_combined(rank=rank, combined=combined0, expected=expected0, label="sync") + + # Cached-handle dispatch: DeepEP-style path without top-k / weight inputs. + x1 = make_input( + rank=rank, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=2_000_000, + addend=17, + ) + recv1, idx1, _w1, handle1, _ = buffer.dispatch( + x1, + handle=handle0, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + actual1 = check_dispatch_payload( + rank=rank, + recv_x=recv1, + handle=handle1, + expected_recv_tokens=route_plan.expected_recv_tokens, + max_tokens=max_tokens, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=2_000_000, + addend=17, + ) + if not torch.equal(idx1[:actual1], idx0[:actual1]): + raise AssertionError(f"rank={rank}: cached dispatch top-k metadata mismatch") + combined1, _, _ = buffer.combine( + recv1[:actual1].contiguous(), + handle1, + topk_weights=None, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + expected1 = (x1.float() * route_plan.expected_combine_factor).to(torch.bfloat16) + check_combined(rank=rank, combined=combined1, expected=expected1, label="cached") + + # Async no-CPU-sync dispatch: output keeps capacity, metadata provides exact count. + x2 = make_input( + rank=rank, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=3_000_000, + addend=31, + ) + recv2, _idx2, w2, handle2, event2 = buffer.dispatch( + x2, + topk_idx=route_plan.topk_idx, + topk_weights=weights, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + expert_alignment=1, + do_cpu_sync=False, + num_sms=args.num_sms, + async_with_compute_stream=True, + ) + event2.current_stream_wait() + torch.cuda.synchronize() + actual2 = check_dispatch_payload( + rank=rank, + recv_x=recv2, + handle=handle2, + expected_recv_tokens=route_plan.expected_recv_tokens, + max_tokens=max_tokens, + num_tokens=args.num_tokens, + hidden=args.hidden, + multiplier=3_000_000, + addend=31, + ) + if handle2.num_recv_tokens_per_expert_list: + raise AssertionError(f"rank={rank}: no-CPU-sync path should not return CPU counts") + combined2, _, event3 = buffer.combine( + recv2[:actual2].contiguous(), + handle2, + topk_weights=w2[:actual2].contiguous(), + num_sms=args.num_sms, + async_with_compute_stream=True, + ) + event3.current_stream_wait() + torch.cuda.synchronize() + expected2 = (x2.float() * route_plan.expected_combine_factor).to(torch.bfloat16) + check_combined(rank=rank, combined=combined2, expected=expected2, label="async") + + # Expanded dispatch layout: check output extent and metadata shape. + expanded_recv, expanded_idx, expanded_w, expanded_handle, _ = buffer.dispatch( + x0, + topk_idx=route_plan.topk_idx, + topk_weights=weights, + num_experts=num_experts, + num_max_tokens_per_rank=max_tokens, + expert_alignment=1, + do_expand=True, + do_cpu_sync=True, + num_sms=args.num_sms, + async_with_compute_stream=False, + ) + torch.cuda.synchronize() + expanded_actual = int(expanded_handle.psum_num_recv_tokens_per_scaleup_rank[-1].item()) + if expanded_actual != route_plan.expected_recv_tokens: + raise AssertionError(f"rank={rank}: expanded dispatch received-token mismatch") + expanded_output = int(expanded_handle.psum_num_recv_tokens_per_expert[-1].item()) + if expanded_recv.shape[0] != expanded_output: + raise AssertionError(f"rank={rank}: expanded output extent mismatch") + if expanded_w is not None and expanded_w.shape[0] != expanded_output: + raise AssertionError(f"rank={rank}: expanded weights extent mismatch") + if expanded_idx.shape[0] != expanded_actual: + raise AssertionError(f"rank={rank}: expanded metadata extent mismatch") + + distributed_barrier() + if rank == 0: + print( + "MOONCAKE_ELASTIC_TEST_OK", + f"world={world_size}", + f"route={args.route}", + f"recv={route_plan.expected_recv_tokens}", + f"expanded={expanded_output}", + f"scaleout={buffer.num_scaleout_ranks}", + f"scaleup={buffer.num_scaleup_ranks}", + flush=True, + ) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/mooncake-ep/tests/test_ep_grid.py b/mooncake-ep/tests/test_ep_grid.py index 2d5e42ddea..c48dff165d 100644 --- a/mooncake-ep/tests/test_ep_grid.py +++ b/mooncake-ep/tests/test_ep_grid.py @@ -22,6 +22,13 @@ def using_musa_backend() -> bool: } +def using_maca_backend() -> bool: + return ( + os.getenv("MOONCAKE_EP_USE_MACA", "").upper() in {"1", "ON", "TRUE", "YES"} + or bool(getattr(torch.version, "maca", None)) + ) + + def import_torchada_if_needed(): if not using_musa_backend(): return @@ -55,10 +62,11 @@ def run_test_iteration( return_recv_hook: bool, use_fallback: bool, fail_rank: int, + buf: Buffer = None, ): - assert not ( - async_finish and return_recv_hook - ), "Should be filtered out by generate_tests." + assert not (async_finish and return_recv_hook), ( + "Should be filtered out by generate_tests." + ) torch.manual_seed(2026 + rank) scale = 1.0 - 0.05 * (rank / num_ranks) @@ -102,7 +110,8 @@ def get_mock_factor(expert_id): num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint( max_tokens, hidden, num_ranks, num_experts ) - buf = Buffer(group, num_ep_buffer_bytes) + if buf is None: + buf = Buffer(group, num_ep_buffer_bytes) if use_fallback: buf._use_fallback = True @@ -230,7 +239,164 @@ def worker(rank, world_size, config_dict): num_ranks=world_size, **config_dict, ) - except Exception as e: + except Exception: + traceback.print_exc() + raise + + os._exit(0) + + +def run_stale_data_test( + group: dist.ProcessGroup, + rank: int, + num_ranks: int, + max_tokens: int, + hidden: int, + num_experts: int, + top_k: int, +): + """Verify combine does not read stale data from a previous round when + a peer rank times out. + + Two rounds share the same buffer: + Round 1 — normal dispatch + combine (fills rdma_recv_data_buffer). + Round 2 — both ranks dispatch; rank 1 exits before combine; rank 0 + combines with a short timeout. + + Without the fix the reduction loop reads stale Round-1 data for rank 1's + experts, producing silently wrong results. + """ + num_local_experts = num_experts // num_ranks + fail_rank = 1 + + def get_mock_factor(expert_id): + return expert_id * 0.1 + 1.0 + + num_ep_buffer_bytes = Buffer.get_ep_buffer_size_hint( + max_tokens, hidden, num_ranks, num_experts + ) + buf = Buffer(group, num_ep_buffer_bytes) + + # Round 1: normal round to fill buffer with data + run_test_iteration( + group=group, + rank=rank, + num_ranks=num_ranks, + max_tokens=max_tokens, + hidden=hidden, + num_experts=num_experts, + top_k=top_k, + use_fp8=False, + zero_copy=False, + async_finish=True, + return_recv_hook=False, + use_fallback=False, + fail_rank=-1, + buf=buf, + ) + + # Round 2: dispatch, then rank 1 exits before combine + torch.manual_seed(200 + rank) + scale = 1.0 - 0.05 * (rank / num_ranks) + num_tokens = int(max_tokens * scale) + x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16) + scores = torch.randn((num_tokens, num_experts), dtype=torch.float32) + topk_idx = torch.topk(scores, top_k, dim=-1)[1] + topk_weights = torch.softmax( + torch.rand(num_tokens, top_k, dtype=torch.float32), dim=-1 + ) + active_ranks = torch.ones((num_ranks,), dtype=torch.int32) + + # Expected output excluding fail_rank's experts + factors = get_mock_factor(topk_idx) + valid_mask = (topk_idx // num_local_experts) != fail_rank + factors = factors * valid_mask.to(factors.dtype) + expected = (x * (factors * topk_weights).sum(dim=1, keepdim=True)).to( + torch.bfloat16 + ) + + dist.barrier(group) + + recv_x, _, handle, event, _ = buf.dispatch( + x, + topk_idx, + active_ranks, + num_max_dispatch_tokens_per_rank=max_tokens, + num_experts=num_experts, + timeout_us=-1, + use_fp8=False, + async_finish=True, + ) + event.current_stream_wait() + torch.cuda.synchronize() + + dist.barrier(group) + + if rank == fail_rank: + os._exit(0) + + expert_out = torch.empty_like(recv_x) + for le in range(num_local_experts): + eid = rank * num_local_experts + le + expert_out[le] = recv_x[le] * get_mock_factor(eid) + expert_out = expert_out.to(torch.bfloat16) + + out = torch.zeros_like(x) + combined_x, event, _ = buf.combine( + expert_out, + topk_idx, + topk_weights, + active_ranks, + timeout_us=5_000, + handle=handle, + async_finish=True, + out=out, + ) + event.current_stream_wait() + torch.cuda.synchronize() + + assert active_ranks[fail_rank].item() == 0, ( + f"[Rank {rank}] active_ranks[{fail_rank}] should be 0 after timeout" + ) + assert not torch.isnan(combined_x).any().item() + testing.assert_close( + combined_x, + expected, + rtol=5e-2, + atol=1e-3, + msg=lambda m: f"[Rank {rank}] Stale data read in combine reduction. {m}", + ) + + +def stale_data_worker(rank, world_size): + import_torchada_if_needed() + + device_filter = [ + f + for f in os.getenv("DEVICE_FILTER", "mlx5_1,mlx5_2,mlx5_3,mlx5_4").split(",") + if f + ] + if device_filter: + pg.set_device_filter(device_filter) + + torch.cuda.set_device(rank) + torch.set_default_dtype(torch.bfloat16) + torch.set_default_device("cuda") + + dist.init_process_group(backend="mooncake", rank=rank, world_size=world_size) + group = dist.group.WORLD + + try: + run_stale_data_test( + group=group, + rank=rank, + num_ranks=world_size, + max_tokens=256, + hidden=2048, + num_experts=288, + top_k=8, + ) + except Exception: traceback.print_exc() raise @@ -253,6 +419,15 @@ def run_single_config(self, config_dict): daemon=False, ) + def test_combine_stale_data(self): + mp.spawn( + stale_data_worker, + args=(self.world_size,), + nprocs=self.world_size, + join=True, + daemon=False, + ) + def make_test_name(cfg): parts = ["test_ep"] @@ -289,8 +464,9 @@ def make_test_name(cfg): def generate_tests(): + fp8_options = [False] if using_maca_backend() else [False, True] test_grid = { - "use_fp8": [False, True], + "use_fp8": fp8_options, "zero_copy": [False, True], "async_finish": [False, True], "return_recv_hook": [False, True], diff --git a/mooncake-integration/CMakeLists.txt b/mooncake-integration/CMakeLists.txt index 7c84a3aff3..c7fe22818b 100644 --- a/mooncake-integration/CMakeLists.txt +++ b/mooncake-integration/CMakeLists.txt @@ -27,7 +27,7 @@ include_directories("../mooncake-transfer-engine/include") find_package( Python3 - COMPONENTS Interpreter Development + COMPONENTS Interpreter Development.Module REQUIRED) execute_process( @@ -52,6 +52,9 @@ if(WITH_TE) if(USE_EFA) target_compile_definitions(engine PRIVATE USE_EFA) endif() + if(USE_CXI) + target_compile_definitions(engine PRIVATE USE_CXI) + endif() # Propagate USE_TENT compile definition to engine target if(USE_TENT) @@ -60,7 +63,6 @@ if(WITH_TE) target_link_libraries(engine PRIVATE $) - target_link_libraries(engine PRIVATE ${Python3_LIBRARIES}) target_include_directories(engine PRIVATE ${Python3_INCLUDE_DIRS}) if(USE_ASCEND_DIRECT) target_link_libraries(engine PUBLIC ascendcl transfer_engine glog::glog @@ -102,9 +104,14 @@ if(USE_UBSHMEM) endif() if(WITH_STORE) - pybind11_add_module(store ${SOURCES} ${CACHE_ALLOCATOR_SOURCES} - store/store_py.cpp store/buffer_pool.cpp - store/engram_store_py.cpp integration_utils.h) + pybind11_add_module( + store + ${SOURCES} + ${CACHE_ALLOCATOR_SOURCES} + store/store_py.cpp + store/buffer_pool.cpp + store/engram_store_py.cpp + integration_utils.h) set_target_properties(store PROPERTIES INSTALL_RPATH "$ORIGIN") if(USE_ASCEND_DIRECT) target_link_libraries( @@ -207,6 +214,7 @@ if(WITH_EP) FILES "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/ep.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_ep_buffer.py" + "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_elastic_buffer.py" "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/pg.py" DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}) # ep.so / pg.so link against engine.so by that exact bare name. Create a diff --git a/mooncake-integration/store/store_py.cpp b/mooncake-integration/store/store_py.cpp index 4a55a2e11e..166d940c5f 100644 --- a/mooncake-integration/store/store_py.cpp +++ b/mooncake-integration/store/store_py.cpp @@ -3,7 +3,9 @@ #include #include +#include #include +#include #include #include @@ -13,6 +15,8 @@ #include "types.h" #include "memory_alloc.h" #include "ssd_register_client.h" +#include "device/accelerator_registry.h" +#include "device/cuda_ipc_buffer.h" #include // for atexit #include @@ -87,55 +91,51 @@ std::optional parse_tensor_metadata_from_buffer( std::pair calculate_shard_range(int64_t dim_size, int rank, int shard_count); -PyTensorInfo extract_tensor_info(const py::object &tensor, - const std::string &key_name = "") { - PyTensorInfo info = { - 0, - 0, - {}, - py::none(), - }; - +PyTensorInfo extract_tensor_info_impl(const py::object &tensor, + const std::string &key_name, + bool require_contiguous, + const char *role) { + PyTensorInfo info = {0, 0, {}, py::none()}; if (!(tensor.attr("__class__") .attr("__name__") .cast() .find("Tensor") != std::string::npos)) { - LOG(ERROR) << "Input " << (key_name.empty() ? "" : "for " + key_name) + LOG(ERROR) << role << " " << (key_name.empty() ? "" : "for " + key_name) << " is not a PyTorch tensor"; return info; } try { - py::object contiguous_tensor = tensor.attr("contiguous")(); - info.owner = contiguous_tensor; - info.data_ptr = contiguous_tensor.attr("data_ptr")().cast(); - size_t numel = contiguous_tensor.attr("numel")().cast(); - size_t element_size = - contiguous_tensor.attr("element_size")().cast(); + if (require_contiguous && + !tensor.attr("is_contiguous")().cast()) { + LOG(ERROR) << role << " " + << (key_name.empty() ? "" : "for " + key_name) + << " must be contiguous"; + return info; + } + + py::object inspected = + require_contiguous ? tensor : tensor.attr("contiguous")(); + info.owner = inspected; + info.data_ptr = inspected.attr("data_ptr")().cast(); + size_t numel = inspected.attr("numel")().cast(); + size_t element_size = inspected.attr("element_size")().cast(); info.tensor_size = numel * element_size; - pybind11::object shape_obj = tensor.attr("shape"); - pybind11::object dtype_obj = tensor.attr("dtype"); - - TensorDtype dtype_enum = get_tensor_dtype(dtype_obj); + TensorDtype dtype_enum = get_tensor_dtype(inspected.attr("dtype")); if (dtype_enum == TensorDtype::UNKNOWN) { LOG(ERROR) << "Unsupported tensor dtype" << (key_name.empty() ? "" : " for " + key_name); return {0, 0, {}, py::none()}; } - - pybind11::tuple shape_tuple = - pybind11::cast(shape_obj); - int32_t ndim = static_cast(shape_tuple.size()); - + int32_t ndim = static_cast(py::len(inspected.attr("shape"))); if (ndim > static_cast(kMaxTensorDims)) { LOG(ERROR) << "Tensor has more than " << kMaxTensorDims << " dimensions: " << ndim; return {0, 0, {}, py::none()}; } - info.metadata = - build_full_tensor_metadata(tensor, dtype_enum, info.tensor_size); + build_full_tensor_metadata(inspected, dtype_enum, info.tensor_size); } catch (const std::exception &e) { LOG(ERROR) << "Error extracting tensor info: " << e.what(); return {0, 0, {}, py::none()}; @@ -144,6 +144,17 @@ PyTensorInfo extract_tensor_info(const py::object &tensor, return info; } +PyTensorInfo extract_tensor_info(const py::object &tensor, + const std::string &key_name = "") { + return extract_tensor_info_impl(tensor, key_name, false, "Input"); +} + +PyTensorInfo extract_tensor_destination_info(const py::object &tensor, + const std::string &key_name = "") { + return extract_tensor_info_impl(tensor, key_name, true, + "Destination tensor"); +} + py::tuple tensor_shape_tuple(const TensorMetadata &metadata) { std::vector shape_vec; shape_vec.reserve(metadata.header.ndim); @@ -160,14 +171,15 @@ void append_tensor_payload_span(std::vector> &values, } pybind11::object buffer_to_tensor(BufferHandle *buffer_handle, char *usr_buffer, - int64_t data_length) { + int64_t data_length, + bool own_user_buffer = false) { if (!buffer_handle && !usr_buffer) return pybind11::none(); if (buffer_handle && usr_buffer) return pybind11::none(); - bool take_ownership = !!buffer_handle; + bool take_ownership = !!buffer_handle || own_user_buffer; size_t total_length; char *exported_data; - if (take_ownership) { + if (buffer_handle) { total_length = buffer_handle->size(); if (total_length < sizeof(TensorMetadata)) { LOG(ERROR) << "Invalid data format: insufficient data for metadata"; @@ -179,16 +191,29 @@ pybind11::object buffer_to_tensor(BufferHandle *buffer_handle, char *usr_buffer, exported_data = new char[total_length]; if (!exported_data) return pybind11::none(); - memcpy(exported_data, buffer_handle->ptr(), total_length); + if (!mooncake::device::GetAcceleratorRegistry() + .RuntimeAccelerators() + .CopyToHost(exported_data, buffer_handle->ptr(), + total_length)) { + LOG(ERROR) << "Failed to copy buffer to host memory"; + delete[] exported_data; + return pybind11::none(); + } } else { exported_data = usr_buffer; if (data_length < 0) { + if (take_ownership) { + delete[] exported_data; + } LOG(ERROR) << "Get tensor into failed with error code: " << data_length; return pybind11::none(); } total_length = static_cast(data_length); if (total_length < sizeof(TensorMetadata)) { + if (take_ownership) { + delete[] exported_data; + } LOG(ERROR) << "Invalid data format: insufficient data for metadata"; return pybind11::none(); } @@ -282,6 +307,77 @@ pybind11::object buffer_to_tensor(BufferHandle *buffer_handle, char *usr_buffer, } } +py::tuple serialize_tensor_metadata(py::object tensor) { + PyTensorInfo info = extract_tensor_info(tensor); + if (!info.valid()) { + throw py::value_error( + "unsupported tensor for Mooncake tensor serialization"); + } + + TensorMetadata metadata = info.metadata; + metadata.header.data_bytes = info.tensor_size; + return py::make_tuple(py::bytes(reinterpret_cast(&metadata), + sizeof(TensorMetadata)), + info.data_ptr, info.tensor_size, info.owner); +} + +size_t tensor_metadata_size() { return sizeof(TensorMetadata); } + +bool tensor_destination_matches_metadata(const PyTensorInfo &target, + const ParsedTensorMetadata &stored, + const std::string &key, + const std::string &context) { + if (target.metadata.header.dtype != stored.metadata.header.dtype || + target.metadata.header.ndim != stored.metadata.header.ndim || + target.tensor_size != stored.data_bytes) { + LOG(ERROR) << context << ": destination tensor mismatch for key " + << key; + return false; + } + return TensorShapeToVector(target.metadata.layout.local_shape, + target.metadata.header.ndim) == + TensorShapeToVector(stored.metadata.layout.local_shape, + stored.metadata.header.ndim); +} + +template +bool apply_indexed_results(const char *context, + const std::vector &op_results, + const std::vector &original_indices, + std::vector &results) { + if (op_results.size() != original_indices.size()) { + LOG(ERROR) << context << ": unexpected result count " + << op_results.size() << ", expected " + << original_indices.size(); + for (size_t index : original_indices) { + results[index] = + static_cast(toInt(ErrorCode::INTERNAL_ERROR)); + } + return false; + } + for (size_t i = 0; i < op_results.size(); ++i) { + results[original_indices[i]] = op_results[i]; + } + return true; +} + +py::object deserialize_tensor_from_bytes(py::bytes payload) { + std::string data = payload; + if (data.size() > + static_cast(std::numeric_limits::max())) { + throw py::value_error("serialized tensor payload is too large"); + } + + char *owned_data = new char[data.size()]; + memcpy(owned_data, data.data(), data.size()); + py::object tensor = buffer_to_tensor( + nullptr, owned_data, static_cast(data.size()), true); + if (tensor.is_none()) { + throw py::value_error("invalid serialized Mooncake tensor payload"); + } + return tensor; +} + std::vector> CastAddrs2Ptrs( const std::vector> &all_buffer_ptrs) { std::vector> all_buffers; @@ -437,6 +533,20 @@ class MooncakeStorePyWrapper { } py::gil_scoped_acquire acquire_gil; + auto runtime_accelerator = + mooncake::device::GetAcceleratorRegistry() + .RuntimeAccelerators(); + if (runtime_accelerator.FindDeviceForPointer( + buffer_handle->ptr())) { + std::string host_buf(buffer_handle->size(), '\0'); + if (!runtime_accelerator.CopyToHost(host_buf.data(), + buffer_handle->ptr(), + buffer_handle->size())) { + LOG(ERROR) << "Failed to copy buffer to host memory"; + return pybind11::none(); + } + return pybind11::bytes(host_buf); + } return pybind11::bytes((char *)buffer_handle->ptr(), buffer_handle->size()); } @@ -463,10 +573,28 @@ class MooncakeStorePyWrapper { std::vector results; results.reserve(batch_data.size()); + auto runtime_accelerator = + mooncake::device::GetAcceleratorRegistry() + .RuntimeAccelerators(); + for (const auto &data : batch_data) { - results.emplace_back( - data ? pybind11::bytes((char *)data->ptr(), data->size()) - : kNullString); + if (!data) { + results.emplace_back(kNullString); + continue; + } + if (runtime_accelerator.FindDeviceForPointer(data->ptr())) { + std::string host_buf(data->size(), '\0'); + if (!runtime_accelerator.CopyToHost( + host_buf.data(), data->ptr(), data->size())) { + LOG(ERROR) << "Failed to copy buffer to host memory"; + results.emplace_back(kNullString); + continue; + } + results.emplace_back(pybind11::bytes(host_buf)); + } else { + results.emplace_back( + pybind11::bytes((char *)data->ptr(), data->size())); + } } return results; } @@ -529,17 +657,17 @@ class MooncakeStorePyWrapper { pybind11::object get_tensor_into(const std::string &key, uintptr_t buffer_ptr, size_t size) { - char *buffer = reinterpret_cast(buffer_ptr); if (!is_client_initialized()) { LOG(ERROR) << "Client is not initialized"; return pybind11::none(); } - - if (use_dummy_client_) { - LOG(ERROR) << "get_tensor is not supported for dummy client now"; + if (buffer_ptr == 0) { + LOG(ERROR) << "Buffer pointer cannot be null"; return pybind11::none(); } + char *buffer = reinterpret_cast(buffer_ptr); + int64_t total_length; { py::gil_scoped_release release_gil; @@ -553,29 +681,34 @@ class MooncakeStorePyWrapper { const std::vector &keys, const std::vector &buffer_ptrs, const std::vector &sizes) { - std::vector buffers; - buffers.reserve(buffer_ptrs.size()); - for (uintptr_t ptr : buffer_ptrs) { - buffers.push_back(reinterpret_cast(ptr)); - } - - if (!is_client_initialized()) { - LOG(ERROR) << "Client is not initialized"; + auto invalid_params = [&keys]() { py::list empty_list; for (size_t i = 0; i < keys.size(); ++i) { empty_list.append(to_py_ret(ErrorCode::INVALID_PARAMS)); } return empty_list; - } + }; - if (use_dummy_client_) { - LOG(ERROR) << "batch_get_tensor is not supported for dummy client " - "now"; - py::list empty_list; - for (size_t i = 0; i < keys.size(); ++i) { - empty_list.append(to_py_ret(ErrorCode::INVALID_PARAMS)); + if (!is_client_initialized()) { + LOG(ERROR) << "Client is not initialized"; + return invalid_params(); + } + if (keys.size() != buffer_ptrs.size() || keys.size() != sizes.size()) { + LOG(ERROR) << "Size mismatch: keys, buffer_ptrs, and sizes must " + "have the same length"; + return invalid_params(); + } + for (uintptr_t ptr : buffer_ptrs) { + if (ptr == 0) { + LOG(ERROR) << "Buffer pointer cannot be null"; + return invalid_params(); } - return empty_list; + } + + std::vector buffers; + buffers.reserve(buffer_ptrs.size()); + for (uintptr_t ptr : buffer_ptrs) { + buffers.push_back(reinterpret_cast(ptr)); } // Phase 1: Batch Get Buffers (GIL Released) @@ -586,14 +719,10 @@ class MooncakeStorePyWrapper { total_lengths = store_->batch_get_into(keys, buffers, sizes); } - if (keys.size() != buffer_ptrs.size() || keys.size() != sizes.size()) { - LOG(ERROR) << "Size mismatch: keys, buffer_ptrs, and sizes must " - "have the same length"; - py::list empty_list; - for (size_t i = 0; i < keys.size(); ++i) { - empty_list.append(to_py_ret(ErrorCode::INVALID_PARAMS)); - } - return empty_list; + if (total_lengths.size() != keys.size()) { + LOG(ERROR) << "Unexpected batch_get_into result count " + << total_lengths.size() << ", expected " << keys.size(); + return invalid_params(); } py::list results_list; @@ -606,6 +735,136 @@ class MooncakeStorePyWrapper { return results_list; } + std::vector> + batch_get_tensor_metadata_prefixes(const std::vector &keys, + const std::string &context) { + std::vector> metadata(keys.size()); + if (keys.empty()) return metadata; + + const size_t scratch_size = keys.size() * sizeof(TensorMetadata); + auto scratch = store_->allocate_client_buffer(scratch_size); + if (!scratch) { + LOG(ERROR) << context << ": failed to allocate metadata buffer"; + return metadata; + } + + std::vector buffers{scratch->ptr()}; + std::vector> all_keys{keys}; + std::vector>> all_dst_offsets(1); + std::vector>> all_src_offsets( + 1, std::vector>(keys.size(), {0})); + std::vector>> all_sizes( + 1, std::vector>(keys.size(), + {sizeof(TensorMetadata)})); + all_dst_offsets[0].reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + all_dst_offsets[0].push_back({i * sizeof(TensorMetadata)}); + } + + std::vector>> results; + { + py::gil_scoped_release release_gil; + results = + store_->get_into_ranges(buffers, all_keys, all_dst_offsets, + all_src_offsets, all_sizes, nullptr); + } + if (results.size() != 1 || results[0].size() != keys.size()) { + LOG(ERROR) << context << ": metadata read result size mismatch"; + return metadata; + } + + const char *base = static_cast(scratch->ptr()); + for (size_t i = 0; i < keys.size(); ++i) { + if (results[0][i].size() != 1 || + results[0][i][0] != + static_cast(sizeof(TensorMetadata))) { + continue; + } + metadata[i] = parse_tensor_metadata_from_prefix( + base + i * sizeof(TensorMetadata), context, keys[i]); + } + return metadata; + } + + int64_t get_tensor_into_cuda(const std::string &key, + pybind11::object tensor) { + py::list tensors; + tensors.append(tensor); + auto results = batch_get_tensor_into_cuda({key}, tensors); + return results.empty() ? to_py_ret(ErrorCode::INVALID_PARAMS) + : results[0]; + } + + std::vector batch_get_tensor_into_cuda( + const std::vector &keys, pybind11::list tensors_list) { + std::vector results(keys.size(), + to_py_ret(ErrorCode::INVALID_PARAMS)); + const std::string context = "batch_get_tensor_into_cuda"; + if (!is_client_initialized()) { + LOG(ERROR) << "Client is not initialized"; + return results; + } + if (keys.size() != tensors_list.size()) { + LOG(ERROR) << context << ": keys and tensors size mismatch"; + return results; + } + + if (!use_dummy_client_) { + LOG(ERROR) << context + << ": CUDA IPC tensor read requires a dummy client"; + return results; + } + + auto metadata = batch_get_tensor_metadata_prefixes(keys, context); + std::vector read_requests; + std::vector original_indices; + read_requests.reserve(keys.size()); + original_indices.reserve(keys.size()); + + for (size_t i = 0; i < keys.size(); ++i) { + PyTensorInfo target = extract_tensor_destination_info( + tensors_list[i].cast(), keys[i]); + if (!target.valid() || !metadata[i].has_value() || + !tensor_destination_matches_metadata(target, *metadata[i], + keys[i], context)) { + continue; + } + if (target.tensor_size == 0) { + results[i] = 0; + continue; + } + auto dst_buffer = mooncake::device::ExportCudaIpcBuffer( + reinterpret_cast(target.data_ptr), target.tensor_size); + if (!dst_buffer) { + LOG(ERROR) << context + << ": destination tensor is not CUDA IPC exportable " + << "for key " << keys[i]; + continue; + } + read_requests.push_back(CudaIpcReadRequest{ + .key = keys[i], + .destination = *dst_buffer, + .source_offset = metadata[i]->data_offset, + .size = metadata[i]->data_bytes, + }); + original_indices.push_back(i); + } + + if (!read_requests.empty()) { + std::vector op_results; + { + py::gil_scoped_release release_gil; + auto dummy_client = + std::static_pointer_cast(store_); + op_results = + dummy_client->batch_get_into_cuda_ipc(read_requests); + } + apply_indexed_results(context.c_str(), op_results, original_indices, + results); + } + return results; + } + pybind11::object get_tensor_with_tp_into(const std::string &key, uintptr_t buffer_ptr, size_t size, int tp_rank = 0, int tp_size = 1, @@ -615,12 +874,6 @@ class MooncakeStorePyWrapper { return pybind11::none(); } - if (use_dummy_client_) { - LOG(ERROR) - << "get_tensor_into is not supported for dummy client now"; - return pybind11::none(); - } - if (tp_size <= 1) { return get_tensor_into(key, buffer_ptr, size); } @@ -645,16 +898,6 @@ class MooncakeStorePyWrapper { return empty_list; } - if (use_dummy_client_) { - LOG(ERROR) << "batch_get_tensor_with_tp_into is not supported for " - "dummy client"; - py::list empty_list; - for (size_t i = 0; i < base_keys.size(); ++i) { - empty_list.append(py::none()); - } - return empty_list; - } - // If tp_size is 1, it's just a normal batch_get_tensor_into if (tp_size <= 1) { return batch_get_tensor_into(base_keys, buffer_ptrs, sizes); @@ -676,6 +919,9 @@ class MooncakeStorePyWrapper { // Validation & Metadata extraction (GIL Held) auto info = extract_tensor_info(tensor, key); if (!info.valid()) return to_py_ret(ErrorCode::INVALID_PARAMS); + if (use_dummy_client_) { + return put_tensor_info_impl(key, info, config); + } // Prepare spans std::vector> values; @@ -693,9 +939,8 @@ class MooncakeStorePyWrapper { } int put_tensor(const std::string &key, pybind11::object tensor) { - if (!is_client_initialized() || use_dummy_client_) { - LOG(ERROR) << "Client not initialized or Dummy client not " - "supported for tensors"; + if (!is_client_initialized()) { + LOG(ERROR) << "Client not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } return put_tensor_impl(key, tensor, @@ -783,9 +1028,8 @@ class MooncakeStorePyWrapper { int put_tensor_with_tp(const std::string &key, pybind11::object tensor, int tp_rank = 0, int tp_size = 1, int split_dim = 0) { - if (!is_client_initialized() || use_dummy_client_) { - LOG(ERROR) << "Client not initialized or Dummy client not " - "supported for tensors"; + if (!is_client_initialized()) { + LOG(ERROR) << "Client not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } if (tp_size <= 1) return put_tensor(key, tensor); @@ -798,6 +1042,10 @@ class MooncakeStorePyWrapper { const std::vector &keys, const std::vector &infos, const ReplicateConfig &config = ReplicateConfig{}) { + if (auto cuda_ipc_results = + try_dummy_cuda_ipc_batch_put_tensor_impl(keys, infos, config)) { + return *cuda_ipc_results; + } return batch_write_tensor_impl( keys, infos, config, "put", [this](const std::vector &write_keys, @@ -822,7 +1070,7 @@ class MooncakeStorePyWrapper { std::vector batch_put_tensor(const std::vector &keys, const pybind11::list &tensors_list) { - if (!is_client_initialized() || use_dummy_client_) + if (!is_client_initialized()) return std::vector(keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); @@ -912,7 +1160,7 @@ class MooncakeStorePyWrapper { const pybind11::list &tensors_list, int tp_rank = 0, int tp_size = 1, int split_dim = 0) { if (tp_size <= 1) return batch_put_tensor(base_keys, tensors_list); - if (!is_client_initialized() || use_dummy_client_) + if (!is_client_initialized()) return std::vector(base_keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); @@ -939,10 +1187,6 @@ class MooncakeStorePyWrapper { LOG(ERROR) << "Client is not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } - if (use_dummy_client_) { - LOG(ERROR) << "put_tensor_from is not supported for dummy client"; - return to_py_ret(ErrorCode::INVALID_PARAMS); - } if (size < sizeof(TensorMetadata)) { LOG(ERROR) << "Buffer size too small for tensor metadata"; return to_py_ret(ErrorCode::INVALID_PARAMS); @@ -965,12 +1209,6 @@ class MooncakeStorePyWrapper { return std::vector(keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); } - if (use_dummy_client_) { - LOG(ERROR) - << "batch_put_tensor_from is not supported for dummy client"; - return std::vector(keys.size(), - to_py_ret(ErrorCode::INVALID_PARAMS)); - } if (keys.empty()) { return std::vector(); } @@ -1013,6 +1251,12 @@ class MooncakeStorePyWrapper { int put_tensor_info_impl(const std::string &key, const PyTensorInfo &info, const ReplicateConfig &config) { if (!info.valid()) return to_py_ret(ErrorCode::INVALID_PARAMS); + if (use_dummy_client_) { + auto results = batch_put_tensor_infos_impl( + {key}, std::vector{info}, config); + return results.empty() ? to_py_ret(ErrorCode::INTERNAL_ERROR) + : results[0]; + } std::vector> values; values.emplace_back(reinterpret_cast(&info.metadata), @@ -1054,11 +1298,6 @@ class MooncakeStorePyWrapper { LOG(ERROR) << "Client is not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } - if (use_dummy_client_) { - LOG(ERROR) - << "put_tensor_with_tp_from is not supported for dummy client"; - return to_py_ret(ErrorCode::INVALID_PARAMS); - } if (buffer_ptr == 0) { LOG(ERROR) << "Buffer pointer cannot be null"; return to_py_ret(ErrorCode::INVALID_PARAMS); @@ -1093,12 +1332,6 @@ class MooncakeStorePyWrapper { return std::vector(base_keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); } - if (use_dummy_client_) { - LOG(ERROR) << "batch_put_tensor_with_tp_from is not supported for " - "dummy client"; - return std::vector(base_keys.size(), - to_py_ret(ErrorCode::INVALID_PARAMS)); - } if (tp_size <= 1) { return batch_put_tensor_from(base_keys, buffer_ptrs, sizes); } @@ -1225,6 +1458,12 @@ class MooncakeStorePyWrapper { const PyTensorInfo &info, const ReplicateConfig &config) { if (!info.valid()) return to_py_ret(ErrorCode::INVALID_PARAMS); + if (use_dummy_client_) { + auto results = batch_upsert_tensor_infos_impl( + {key}, std::vector{info}, config); + return results.empty() ? to_py_ret(ErrorCode::INTERNAL_ERROR) + : results[0]; + } std::vector> values; values.emplace_back(reinterpret_cast(&info.metadata), @@ -1243,6 +1482,9 @@ class MooncakeStorePyWrapper { const ReplicateConfig &config) { auto info = extract_tensor_info(tensor, key); if (!info.valid()) return to_py_ret(ErrorCode::INVALID_PARAMS); + if (use_dummy_client_) { + return upsert_tensor_info_impl(key, info, config); + } std::vector> values; values.emplace_back(reinterpret_cast(&info.metadata), @@ -1258,9 +1500,8 @@ class MooncakeStorePyWrapper { } int upsert_tensor(const std::string &key, pybind11::object tensor) { - if (!is_client_initialized() || use_dummy_client_) { - LOG(ERROR) << "Client not initialized or Dummy client not " - "supported for tensors"; + if (!is_client_initialized()) { + LOG(ERROR) << "Client not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } return upsert_tensor_impl(key, tensor, ReplicateConfig{}); @@ -1300,12 +1541,6 @@ class MooncakeStorePyWrapper { return std::vector(keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); } - if (use_dummy_client_) { - LOG(ERROR) - << "batch_upsert_tensor_from is not supported for dummy client"; - return std::vector(keys.size(), - to_py_ret(ErrorCode::INVALID_PARAMS)); - } if (keys.empty()) { return std::vector(); } @@ -1346,6 +1581,21 @@ class MooncakeStorePyWrapper { ReplicateConfig{}); } + std::vector batch_upsert_tensor_infos_impl( + const std::vector &keys, + const std::vector &infos, + const ReplicateConfig &config = ReplicateConfig{}) { + return batch_write_tensor_impl( + keys, infos, config, "upsert", + [this](const std::vector &write_keys, + const std::vector &buffer_ptrs, + const std::vector &buffer_sizes, + const ReplicateConfig &write_config) { + return store_->batch_upsert_from(write_keys, buffer_ptrs, + buffer_sizes, write_config); + }); + } + std::vector batch_upsert_tensor_impl( const std::vector &keys, const pybind11::list &tensors_list, @@ -1357,76 +1607,15 @@ class MooncakeStorePyWrapper { } std::vector infos(keys.size()); - std::vector results(keys.size(), 0); - - // 1. Extract Metadata (GIL Held) for (size_t i = 0; i < keys.size(); ++i) { infos[i] = extract_tensor_info(tensors_list[i], keys[i]); - if (!infos[i].valid()) - results[i] = to_py_ret(ErrorCode::INVALID_PARAMS); } - - // 2. Prepare Buffers and Execute (GIL Released) - { - py::gil_scoped_release release_gil; - - std::vector valid_keys; - std::vector buffer_ptrs; - std::vector buffer_sizes; - std::vector original_indices; - - std::vector> temp_allocations; - - for (size_t i = 0; i < infos.size(); ++i) { - if (!infos[i].valid()) continue; - - size_t total_size = - sizeof(TensorMetadata) + infos[i].tensor_size; - auto alloc_result = - store_->client_buffer_allocator_->allocate(total_size); - - if (!alloc_result) { - LOG(ERROR) - << "Failed to allocate buffer for key: " << keys[i]; - results[i] = to_py_ret(ErrorCode::INVALID_PARAMS); - continue; - } - - // Copy Metadata & Data - char *dst = static_cast(alloc_result->ptr()); - memcpy(dst, &infos[i].metadata, sizeof(TensorMetadata)); - if (infos[i].tensor_size > 0) { - memcpy(dst + sizeof(TensorMetadata), - reinterpret_cast(infos[i].data_ptr), - infos[i].tensor_size); - } - - valid_keys.push_back(keys[i]); - buffer_ptrs.push_back(alloc_result->ptr()); - buffer_sizes.push_back(total_size); - original_indices.push_back(i); - - temp_allocations.push_back( - std::make_unique(std::move(*alloc_result))); - } - - if (!valid_keys.empty()) { - ReplicateConfig write_config = - MakeIndexedConfig(config, original_indices); - std::vector op_results = store_->batch_upsert_from( - valid_keys, buffer_ptrs, buffer_sizes, write_config); - for (size_t i = 0; i < op_results.size(); ++i) { - results[original_indices[i]] = op_results[i]; - } - } - } - - return results; + return batch_upsert_tensor_infos_impl(keys, infos, config); } std::vector batch_upsert_tensor(const std::vector &keys, const pybind11::list &tensors_list) { - if (!is_client_initialized() || use_dummy_client_) + if (!is_client_initialized()) return std::vector(keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); @@ -1450,11 +1639,6 @@ class MooncakeStorePyWrapper { LOG(ERROR) << "Client is not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } - if (use_dummy_client_) { - LOG(ERROR) - << "upsert_tensor_from is not supported for dummy client"; - return to_py_ret(ErrorCode::INVALID_PARAMS); - } if (size < sizeof(TensorMetadata)) { LOG(ERROR) << "Buffer size too small for tensor metadata"; return to_py_ret(ErrorCode::INVALID_PARAMS); @@ -1483,9 +1667,8 @@ class MooncakeStorePyWrapper { int upsert_pub_tensor(const std::string &key, pybind11::object tensor, const ReplicateConfig &config = ReplicateConfig{}) { - if (!is_client_initialized() || use_dummy_client_) { - LOG(ERROR) << "Client not initialized or Dummy client not " - "supported for tensors"; + if (!is_client_initialized()) { + LOG(ERROR) << "Client not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } @@ -1499,7 +1682,7 @@ class MooncakeStorePyWrapper { const std::vector &keys, const pybind11::list &tensors_list, const ReplicateConfig &config = ReplicateConfig{}) { - if (!is_client_initialized() || use_dummy_client_) + if (!is_client_initialized()) return std::vector(keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); @@ -1521,9 +1704,8 @@ class MooncakeStorePyWrapper { // --- End Upsert tensor methods --- int pub_tensor(const std::string &key, pybind11::object tensor, const ReplicateConfig &config = ReplicateConfig{}) { - if (!is_client_initialized() || use_dummy_client_) { - LOG(ERROR) << "Client not initialized or Dummy client not " - "supported for tensors"; + if (!is_client_initialized()) { + LOG(ERROR) << "Client not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } @@ -1537,9 +1719,8 @@ class MooncakeStorePyWrapper { const ReplicateConfig &config = ReplicateConfig{}, int tp_rank = 0, int tp_size = 1, int split_dim = 0) { - if (!is_client_initialized() || use_dummy_client_) { - LOG(ERROR) << "Client not initialized or Dummy client not " - "supported for tensors"; + if (!is_client_initialized()) { + LOG(ERROR) << "Client not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } @@ -1556,7 +1737,7 @@ class MooncakeStorePyWrapper { const std::vector &keys, const pybind11::list &tensors_list, const ReplicateConfig &config = ReplicateConfig{}) { - if (!is_client_initialized() || use_dummy_client_) + if (!is_client_initialized()) return std::vector(keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); @@ -1582,7 +1763,7 @@ class MooncakeStorePyWrapper { int tp_size = 1, int split_dim = 0) { if (tp_size <= 1) return batch_pub_tensor(base_keys, tensors_list, config); - if (!is_client_initialized() || use_dummy_client_) + if (!is_client_initialized()) return std::vector(base_keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); @@ -1609,12 +1790,6 @@ class MooncakeStorePyWrapper { return to_py_ret(ErrorCode::INVALID_PARAMS); } - if (use_dummy_client_) { - LOG(ERROR) << "save_tensor_to_safetensor is not supported for " - << "dummy client now"; - return to_py_ret(ErrorCode::INVALID_PARAMS); - } - pybind11::object tensor = get_tensor(key); if (tensor.is_none()) { LOG(ERROR) << "Failed to fetch tensor for key: " << key; @@ -1644,12 +1819,6 @@ class MooncakeStorePyWrapper { return pybind11::none(); } - if (use_dummy_client_) { - LOG(ERROR) << "load_tensor_from_safetensor is not supported for " - << "dummy client now"; - return pybind11::none(); - } - try { auto safetensors_torch = py::module_::import("safetensors.torch"); py::dict loaded = safetensors_torch.attr("load_file")(file_name); @@ -1723,6 +1892,14 @@ class MooncakeDistributedNoFRegisterPyWrapper { }; PYBIND11_MODULE(store, m) { + m.def("_serialize_tensor", &serialize_tensor_metadata, + "Inspect a torch tensor as Mooncake tensor metadata, data pointer, " + "size, and owner."); + m.def("_tensor_metadata_size", &tensor_metadata_size, + "Return the serialized Mooncake tensor metadata size."); + m.def("_deserialize_tensor", &deserialize_tensor_from_bytes, + "Deserialize Mooncake tensor metadata plus payload bytes."); + // Object data type classification py::enum_(m, "ObjectDataType") .value("UNKNOWN", ObjectDataType::UNKNOWN) @@ -2028,7 +2205,9 @@ PYBIND11_MODULE(store, m) { const py::object &engine = py::none(), bool enable_ssd_offload = false, const std::string &ssd_offload_path = "", - const std::string &tenant_id = "default") { + const std::string &tenant_id = "default", + bool enable_client_http_server = false, + int client_http_port = DEFAULT_CLIENT_HTTP_PORT) { auto real_client = self.init_real_client(); std::shared_ptr transfer_engine = nullptr; @@ -2040,14 +2219,17 @@ PYBIND11_MODULE(store, m) { local_hostname, metadata_server, global_segment_size, local_buffer_size, protocol, rdma_devices, master_server_addr, transfer_engine, "", enable_ssd_offload, - ssd_offload_path, tenant_id); + ssd_offload_path, tenant_id, enable_client_http_server, + client_http_port); }, py::arg("local_hostname"), py::arg("metadata_server"), py::arg("global_segment_size"), py::arg("local_buffer_size"), py::arg("protocol"), py::arg("rdma_devices"), py::arg("master_server_addr"), py::arg("engine") = py::none(), py::arg("enable_ssd_offload") = false, - py::arg("ssd_offload_path") = "", py::arg("tenant_id") = "default") + py::arg("ssd_offload_path") = "", py::arg("tenant_id") = "default", + py::arg("enable_client_http_server") = false, + py::arg("client_http_port") = DEFAULT_CLIENT_HTTP_PORT) .def( "setup", [](MooncakeStorePyWrapper &self, const py::dict &config_dict) { @@ -2078,7 +2260,11 @@ PYBIND11_MODULE(store, m) { " ipc_socket_path: IPC socket path.\n" " enable_ssd_offload: Enable SSD offload (default false).\n" " ssd_offload_path: SSD storage directory path (overrides env " - "var).") + "var).\n" + " tenant_id: Tenant identifier (default 'default').\n" + " enable_client_http_server: Enable client HTTP endpoints " + "(default false).\n" + " client_http_port: Client HTTP metrics port (default 9300).") .def( "setup_dummy", [](MooncakeStorePyWrapper &self, size_t mem_pool_size, @@ -2305,6 +2491,13 @@ PYBIND11_MODULE(store, m) { "Get tensors directly into pre-allocated buffers for " "multiple " "keys") + .def("get_tensor_into_cuda", + &MooncakeStorePyWrapper::get_tensor_into_cuda, py::arg("key"), + py::arg("tensor"), "Get tensor payload into a CUDA tensor") + .def("batch_get_tensor_into_cuda", + &MooncakeStorePyWrapper::batch_get_tensor_into_cuda, + py::arg("keys"), py::arg("tensors_list"), + "Get tensor payloads into pre-allocated CUDA tensors") .def( "get_tensor_with_tp_into", &MooncakeStorePyWrapper::get_tensor_with_tp_into, py::arg("key"), diff --git a/mooncake-integration/store/store_py_parallel_write.h b/mooncake-integration/store/store_py_parallel_write.h index 2bba450a8f..9f23e3d501 100644 --- a/mooncake-integration/store/store_py_parallel_write.h +++ b/mooncake-integration/store/store_py_parallel_write.h @@ -14,12 +14,82 @@ int write_manifest_impl(const std::string &key, return ret; } +std::optional> try_dummy_cuda_ipc_batch_put_tensor_impl( + const std::vector &keys, + const std::vector &infos, const ReplicateConfig &config) { + if (!use_dummy_client_ || keys.size() != infos.size()) return std::nullopt; + + std::vector results(keys.size(), 0); + py::gil_scoped_release release_gil; + + std::vector write_requests; + std::vector original_indices; + std::vector> metadata_allocations; + write_requests.reserve(infos.size()); + original_indices.reserve(infos.size()); + metadata_allocations.reserve(infos.size()); + + for (size_t i = 0; i < infos.size(); ++i) { + if (!infos[i].valid()) { + results[i] = to_py_ret(ErrorCode::INVALID_PARAMS); + continue; + } + if (infos[i].tensor_size == 0) return std::nullopt; + + auto payload = mooncake::device::ExportCudaIpcBuffer( + reinterpret_cast(infos[i].data_ptr), + infos[i].tensor_size); + if (!payload) return std::nullopt; + + size_t metadata_size = infos[i].metadata.header.data_offset; + auto metadata = store_->allocate_client_buffer(metadata_size); + if (!metadata) { + results[i] = to_py_ret(ErrorCode::NO_AVAILABLE_HANDLE); + continue; + } + + std::memcpy(metadata->ptr(), &infos[i].metadata, metadata_size); + write_requests.push_back(CudaIpcWriteRequest{ + .key = keys[i], + .metadata = + CudaIpcShmBufferRef{ + .ptr = reinterpret_cast(metadata->ptr()), + .size = static_cast(metadata_size), + }, + .payload = *payload, + }); + original_indices.push_back(i); + metadata_allocations.push_back( + std::make_unique(std::move(*metadata))); + } + + if (!write_requests.empty()) { + ReplicateConfig write_config = + MakeIndexedConfig(config, original_indices); + auto dummy_client = std::static_pointer_cast(store_); + std::vector op_results = + dummy_client->batch_put_from_cuda_ipc(write_requests, write_config); + if (!apply_indexed_results("put", op_results, original_indices, + results)) { + return results; + } + } + return results; +} + template std::vector batch_write_tensor_impl(const std::vector &keys, const std::vector &infos, const ReplicateConfig &config, const char *operation_name, BatchWriteFromFn &&batch_write_from) { + if (keys.size() != infos.size()) { + LOG(ERROR) << operation_name + << ": keys and tensor infos must have the same length"; + return std::vector(keys.size(), + to_py_ret(ErrorCode::INVALID_PARAMS)); + } + auto group_ids_error = ValidateGroupIdsForBatchConfig(config, keys.size(), operation_name); if (!group_ids_error.empty()) { @@ -30,6 +100,14 @@ std::vector batch_write_tensor_impl(const std::vector &keys, { py::gil_scoped_release release_gil; + if (!store_->client_buffer_allocator_ && !use_dummy_client_) { + LOG(ERROR) << operation_name + << ": client buffer allocator is not available"; + return std::vector(keys.size(), + to_py_ret(ErrorCode::INVALID_PARAMS)); + } + auto runtime_accelerator = + mooncake::device::GetAcceleratorRegistry().RuntimeAccelerators(); std::vector valid_keys; std::vector buffer_ptrs; @@ -45,8 +123,7 @@ std::vector batch_write_tensor_impl(const std::vector &keys, size_t total_size = infos[i].metadata.header.data_offset + infos[i].tensor_size; - auto alloc_result = - store_->client_buffer_allocator_->allocate(total_size); + auto alloc_result = store_->allocate_client_buffer(total_size); if (!alloc_result) { LOG(ERROR) << "Failed to allocate buffer for " << operation_name @@ -59,9 +136,15 @@ std::vector batch_write_tensor_impl(const std::vector &keys, std::memcpy(dst, &infos[i].metadata, infos[i].metadata.header.data_offset); if (infos[i].tensor_size > 0) { - std::memcpy(dst + infos[i].metadata.header.data_offset, - reinterpret_cast(infos[i].data_ptr), - infos[i].tensor_size); + if (!runtime_accelerator.CopyToHost( + dst + infos[i].metadata.header.data_offset, + reinterpret_cast(infos[i].data_ptr), + infos[i].tensor_size)) { + LOG(ERROR) << "Failed to copy tensor payload for " + << operation_name << " key: " << keys[i]; + results[i] = to_py_ret(ErrorCode::INVALID_PARAMS); + continue; + } } valid_keys.push_back(keys[i]); @@ -77,8 +160,9 @@ std::vector batch_write_tensor_impl(const std::vector &keys, MakeIndexedConfig(config, original_indices); std::vector op_results = batch_write_from( valid_keys, buffer_ptrs, buffer_sizes, write_config); - for (size_t i = 0; i < op_results.size(); ++i) { - results[original_indices[i]] = op_results[i]; + if (!apply_indexed_results(operation_name, op_results, + original_indices, results)) { + return results; } } } @@ -87,10 +171,8 @@ std::vector batch_write_tensor_impl(const std::vector &keys, } bool ensure_tensor_write_supported(const char *operation_name) const { - if (!is_client_initialized() || use_dummy_client_) { - LOG(ERROR) << operation_name - << ": client not initialized or dummy client not " - "supported for tensors"; + if (!is_client_initialized()) { + LOG(ERROR) << operation_name << ": client not initialized"; return false; } return true; @@ -927,7 +1009,7 @@ std::vector batch_put_tensor_with_parallelism( keys, tensors_list.size(), parallelisms, writer_partitions, "batch_put_tensor_with_parallelism", [this, &keys, &tensors_list, &config]() { - if (!is_client_initialized() || use_dummy_client_) { + if (!is_client_initialized()) { return std::vector(keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); } @@ -1060,7 +1142,7 @@ std::vector batch_put_tensor_with_parallelism_from( "batch_put_tensor_with_parallelism_from", [this, &keys, &buffer_ptrs, &sizes, &config]() { if (!is_default_replicate_config(config)) { - if (!is_client_initialized() || use_dummy_client_) { + if (!is_client_initialized()) { return std::vector( keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); } @@ -1300,11 +1382,6 @@ int execute_upsert_tensor_with_parallelism_from_route( LOG(ERROR) << "Client is not initialized"; return to_py_ret(ErrorCode::INVALID_PARAMS); } - if (use_dummy_client_) { - LOG(ERROR) << "upsert_tensor_with_parallelism_from is not " - "supported for dummy client"; - return to_py_ret(ErrorCode::INVALID_PARAMS); - } int validate_result = validate_replicate_config(write_config); if (validate_result) return validate_result; py::gil_scoped_release release_gil; @@ -1383,7 +1460,7 @@ std::vector batch_upsert_tensor_with_parallelism( keys, tensors_list.size(), parallelisms, writer_partitions, "batch_upsert_tensor_with_parallelism", [this, &keys, &tensors_list, &config]() { - if (!is_client_initialized() || use_dummy_client_) { + if (!is_client_initialized()) { return std::vector(keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); } @@ -1431,12 +1508,6 @@ std::vector batch_upsert_tensor_with_parallelism_from( return std::vector( keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); } - if (use_dummy_client_) { - LOG(ERROR) << "batch_upsert_tensor_with_parallelism_from " - "is not supported for dummy client"; - return std::vector( - keys.size(), to_py_ret(ErrorCode::INVALID_PARAMS)); - } int validate_result = validate_replicate_config(config); if (validate_result) { return std::vector(keys.size(), validate_result); diff --git a/mooncake-integration/transfer_engine/transfer_engine_py.cpp b/mooncake-integration/transfer_engine/transfer_engine_py.cpp index d0438751a9..39d31b75ba 100644 --- a/mooncake-integration/transfer_engine/transfer_engine_py.cpp +++ b/mooncake-integration/transfer_engine/transfer_engine_py.cpp @@ -22,7 +22,7 @@ #include "transport/rpc_communicator/rpc_interface.h" #ifdef USE_TENT -#include "tent/runtime/transport_selector.h" +#include "tent/common/types.h" #endif #ifdef USE_EFA @@ -207,6 +207,16 @@ int TransferEnginePy::initializeExt(const char* local_hostname, << engine_->getLocalTopology()->getHcaList().size() << " devices."; } +#elif defined(USE_CXI) + + bool use_cxi = (proto == "cxi"); + engine_ = std::make_unique(false, device_filter); + if (use_cxi) { + engine_->getLocalTopology()->discover(device_filter); + LOG(INFO) << "Topology discovery complete for CXI. Found " + << engine_->getLocalTopology()->getHcaList().size() + << " devices."; + } #else engine_ = std::make_unique(true, device_filter); #endif @@ -247,6 +257,29 @@ int TransferEnginePy::initializeExt(const char* local_hostname, } LOG(INFO) << "TCP transport installed successfully"; } +#elif defined(USE_CXI) + if (use_cxi) { + LOG(INFO) + << "Installing CXI transport as requested by protocol parameter"; + auto transport = engine_->installTransport("cxi", nullptr); + if (!transport) { + LOG(ERROR) << "Failed to install CXI transport"; + return -1; + } + LOG(INFO) << "CXI transport installed successfully"; + } else { + // For non-EFA protocols (e.g. TCP), manually install TCP transport + // since auto_discover is disabled to prevent RDMA installation + // (RDMA QP creation fails on EFA devices). + LOG(INFO) + << "Installing TCP transport (auto_discover disabled in CXI build)"; + auto transport = engine_->installTransport("tcp", nullptr); + if (!transport) { + LOG(ERROR) << "Failed to install TCP transport"; + return -1; + } + LOG(INFO) << "TCP transport installed successfully"; + } #endif free_list_.resize(kSlabSizeKBTabLen); @@ -328,7 +361,7 @@ int TransferEnginePy::freeManagedBuffer(uintptr_t buffer_addr, size_t length) { static int parseTransportHint(const std::string& name) { #ifdef USE_TENT if (name.empty()) return mooncake::tent::UNSPEC; - auto type = mooncake::tent::TransportSelector::parseTransportType(name); + auto type = mooncake::tent::parseTransportType(name); if (type == mooncake::tent::UNSPEC && name != "unspec") { throw std::invalid_argument( "Unknown transport_hint '" + name + diff --git a/mooncake-p2p-store/src/p2pstore/core.go b/mooncake-p2p-store/src/p2pstore/core.go index 33fee16ca1..1119e01a42 100644 --- a/mooncake-p2p-store/src/p2pstore/core.go +++ b/mooncake-p2p-store/src/p2pstore/core.go @@ -268,6 +268,7 @@ func (store *P2PStore) doGetReplica(ctx context.Context, payload *Payload, addrL if err != nil { return err } + offset = 0 for ; offset < size; offset += maxShardSize { source := addr + uintptr(offset) shard := payload.Shards[taskID] @@ -275,7 +276,7 @@ func (store *P2PStore) doGetReplica(ctx context.Context, payload *Payload, addrL wg.Add(1) go func() { defer wg.Done() - err = store.performTransfer(ctx, source, shard) + err := store.performTransfer(ctx, source, shard) if err != nil { select { case errChan <- err: diff --git a/mooncake-p2p-store/src/p2pstore/go.mod b/mooncake-p2p-store/src/p2pstore/go.mod index 8d29b5ab85..5929721056 100644 --- a/mooncake-p2p-store/src/p2pstore/go.mod +++ b/mooncake-p2p-store/src/p2pstore/go.mod @@ -1,8 +1,6 @@ module github.com/kvcache-ai/Mooncake/mooncake-p2p-store/src/p2pstore -go 1.23.0 - -toolchain go1.24.1 +go 1.25.0 require go.etcd.io/etcd/client/v3 v3.5.15 @@ -16,12 +14,11 @@ require ( go.uber.org/atomic v1.7.0 // indirect go.uber.org/multierr v1.6.0 // indirect go.uber.org/zap v1.17.0 // indirect - golang.org/x/net v0.38.0 // indirect - golang.org/x/sys v0.31.0 // indirect - golang.org/x/text v0.23.0 // indirect - google.golang.org/genproto v0.0.0-20230822172742-b8732ec3820d // indirect - google.golang.org/genproto/googleapis/api v0.0.0-20230822172742-b8732ec3820d // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect - google.golang.org/grpc v1.59.0 // indirect - google.golang.org/protobuf v1.33.0 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/text v0.37.0 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20251202230838-ff82c1b0f217 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect + google.golang.org/grpc v1.79.3 // indirect + google.golang.org/protobuf v1.36.10 // indirect ) diff --git a/mooncake-p2p-store/src/p2pstore/registered_memory.go b/mooncake-p2p-store/src/p2pstore/registered_memory.go index 9333d91a32..099e3d2e4b 100644 --- a/mooncake-p2p-store/src/p2pstore/registered_memory.go +++ b/mooncake-p2p-store/src/p2pstore/registered_memory.go @@ -82,7 +82,6 @@ func (memory *RegisteredMemory) Add(addr uintptr, length uint64, maxShardSize ui if err != nil { select { case errChan <- err: - close(errChan) return default: } @@ -104,6 +103,15 @@ func (memory *RegisteredMemory) Add(addr uintptr, length uint64, maxShardSize ui log.Println("cascading error:", unregisterErr) } } + memory.mu.Lock() + for idx, entry := range memory.bufferList { + if entry.addr == addr && entry.length == length { + memory.bufferList = append(memory.bufferList[:idx], + memory.bufferList[idx+1:]...) + break + } + } + memory.mu.Unlock() return err } @@ -117,21 +125,26 @@ func (memory *RegisteredMemory) Remove(addr uintptr, length uint64, maxShardSize memory.mu.Lock() found := false + lastRef := false for idx, entry := range memory.bufferList { if entry.addr == addr && entry.length == length { found = true - entry.refCount-- - if entry.refCount == 0 { + memory.bufferList[idx].refCount-- + if memory.bufferList[idx].refCount == 0 { + lastRef = true memory.bufferList = append(memory.bufferList[:idx], memory.bufferList[idx+1:]...) - break } + break } } memory.mu.Unlock() if !found { return ErrInvalidArgument } + if !lastRef { + return nil + } var wg sync.WaitGroup errChan := make(chan error, 1) diff --git a/mooncake-pg/CMakeLists.txt b/mooncake-pg/CMakeLists.txt index e77ce45133..8feaffb7c2 100644 --- a/mooncake-pg/CMakeLists.txt +++ b/mooncake-pg/CMakeLists.txt @@ -1,33 +1,5 @@ cmake_minimum_required(VERSION 3.16) project(mooncake-pg) -# Find PyTorch's CMake prefix path -execute_process( - COMMAND ${PYTHON_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)" - OUTPUT_VARIABLE PYTORCH_CMAKE_PATH - OUTPUT_STRIP_TRAILING_WHITESPACE -) -if(NOT PYTORCH_CMAKE_PATH) - message(WARNING "Could not find PyTorch CMake path! Please set Torch_DIR.") -else () - message(STATUS "Found PyTorch CMake path: ${PYTORCH_CMAKE_PATH}") - list(APPEND CMAKE_PREFIX_PATH "${PYTORCH_CMAKE_PATH}/Torch") -endif() - -set(TORCH_CUDA_ARCH_LIST "8.0;9.0") - -find_package(CUDAToolkit REQUIRED) -# https://discuss.pytorch.org/t/failed-to-find-nvtoolsext/179635/13 -if(NOT TARGET CUDA::nvToolsExt AND TARGET CUDA::nvtx3) - add_library(CUDA::nvToolsExt INTERFACE IMPORTED) - target_compile_definitions( - CUDA::nvToolsExt INTERFACE - TORCH_CUDA_USE_NVTX3 - ) - target_link_libraries(CUDA::nvToolsExt INTERFACE CUDA::nvtx3) -endif() -find_package(Torch REQUIRED) -include_directories(${TORCH_INCLUDE_DIRS}) - include_directories(include) add_subdirectory(src) diff --git a/mooncake-pg/include/comm_types.h b/mooncake-pg/include/comm_types.h new file mode 100644 index 0000000000..2f409ab856 --- /dev/null +++ b/mooncake-pg/include/comm_types.h @@ -0,0 +1,122 @@ +#ifndef MOONCAKE_PG_COMM_TYPES_H +#define MOONCAKE_PG_COMM_TYPES_H + +#include +#include +#include +#include +#include + +#include "error_types.h" + +namespace mooncake { + +enum class OpType : uint8_t { + Unknown = 0, + Broadcast, + AllReduce, + AllGather, + ReduceScatter, + AllToAll, + Barrier, + Reduce, + Gather, + Scatter, + Send, + Recv, +}; + +enum class DataType : uint8_t { + Int8 = 0, + Uint8, + Int16, + Uint16, + Int32, + Uint32, + Int64, + Uint64, + Float16, + Float32, + Float64, + Bfloat16, + Bool, + Float8e4m3fn, + Float8e5m2, + Float8e4m3fnuz, + Float8e5m2fnuz, + Float8e8m0fnu, +}; + +inline size_t elementSize(DataType dataType) { + switch (dataType) { + case DataType::Int8: + case DataType::Uint8: + case DataType::Bool: + case DataType::Float8e4m3fn: + case DataType::Float8e5m2: + case DataType::Float8e4m3fnuz: + case DataType::Float8e5m2fnuz: + case DataType::Float8e8m0fnu: + return 1; + case DataType::Int16: + case DataType::Uint16: + case DataType::Float16: + case DataType::Bfloat16: + return 2; + case DataType::Int32: + case DataType::Uint32: + case DataType::Float32: + return 4; + case DataType::Int64: + case DataType::Uint64: + case DataType::Float64: + return 8; + } + PG_ASSERT(false, + "unsupported Mooncake datatype: ", static_cast(dataType)); +} + +enum class ReduceOp : uint8_t { + Sum = 0, + Avg = 1, + Product = 2, + Min = 3, + Max = 4, +}; + +class WorkCompletion { + public: + explicit WorkCompletion(std::shared_future completion) + : completion_(std::move(completion)) {} + + bool isCompleted() const { + if (completion_.wait_for(std::chrono::microseconds(0)) != + std::future_status::ready) { + return false; + } + completion_.get(); + return true; + } + + bool wait(std::chrono::microseconds timeout) const { + if (timeout.count() < 0) { + completion_.wait(); + } else if (completion_.wait_for(timeout) != std::future_status::ready) { + return false; + } + completion_.get(); + return true; + } + + private: + std::shared_future completion_; +}; + +struct CudaTaskSubmissionToken { + size_t task_id; + uint64_t sequence; +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_COMM_TYPES_H diff --git a/mooncake-pg/include/connection_poller.h b/mooncake-pg/include/connection_poller.h deleted file mode 100644 index 42417b1efb..0000000000 --- a/mooncake-pg/include/connection_poller.h +++ /dev/null @@ -1,257 +0,0 @@ -#ifndef MOONCAKE_PG_CONNECTION_POLLER_H -#define MOONCAKE_PG_CONNECTION_POLLER_H - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include - -namespace mooncake { - -enum class PeerConnectionState { - WAITING_STORE, - WAITING_WARMUP_TRANSFER, - WAITING_PEER_WARMUP, - CONNECTED, - EXPIRING, -}; - -struct PeerConnection { - static constexpr size_t kCheckStoreInitialBackoffMs = 8; - static constexpr size_t kCheckStoreMaxBackoffMs = 1024; - - PeerConnectionState state{PeerConnectionState::WAITING_STORE}; - std::optional warmupBatchId{std::nullopt}; - std::optional segmentId{std::nullopt}; - - // Whether this peer has been counted in `totalConnectedPeers_`. - // Note that ConnectionPoller may establish connections for ranks beyond - // the current groupSize_ (when pollingLimit_ > groupSize_). In that case we - // delay counting until the rank officially enters the group. - bool countedInGroup{false}; - - // Back off to avoid frequently checking store. - std::chrono::steady_clock::time_point last_check_store; - size_t check_store_backoff_ms{kCheckStoreInitialBackoffMs}; - - void increaseCheckStoreBackoff() { - check_store_backoff_ms = - (std::min)(check_store_backoff_ms * 2, - PeerConnection::kCheckStoreMaxBackoffMs); - } - - void resetCheckStoreBackoff() { - check_store_backoff_ms = kCheckStoreInitialBackoffMs; - } -}; - -class ConnectionContext { - static constexpr size_t kDrainPollerTimeoutMs = 5000; // 5s - friend class ConnectionPoller; - - int backendIndex_; - int rank_; - - std::atomic groupSize_; - - // Upper bound of peer polling range. - // This can be larger than groupSize_ so that existing ranks can observe - // joiners without calling extendGroupSizeTo(). - std::atomic pollingLimit_; - - bool isDummy_; - - // A mark tracking the group size for which all ranks - // in [0, establishedGroupSize_) have been successfully - // connected at least once (they may disconnect afterwards). - // Mainly used in `waitUntilNewRanksConnected()`. - std::atomic establishedGroupSize_; - - uint64_t* local2global_rank_map_; - c10::intrusive_ptr<::c10d::Store> store_; - - std::shared_ptr meta_; - std::shared_ptr p2p_proxy_; - TransferEngine* engine_; - - std::atomic totalConnectedPeers_{0}; - std::atomic isShutdown_{false}; - - PeerConnection peerStates_[kMaxNumRanks]; - - // On MNNVL, warmup is skipped because CPU heap buffers aren't - // fabric-accessible for cross-node NVLink writes. - bool skip_warmup_; - - // warmup_send_region_ and warmup_recv_region_ are managed by - // ConnectionContext. nullptr when skip_warmup_ is true. - int32_t* warmup_send_region_; - int32_t* warmup_recv_region_; - - std::mutex backend_wakeup_mutex_; - std::condition_variable backend_wakeup_cv_; - - bool resource_abandoned_{false}; - - public: - ConnectionContext(int backendIndex, int rank, int size, bool isDummy, - uint64_t* local2global_rank_map, - c10::intrusive_ptr<::c10d::Store> store, - std::shared_ptr meta, - std::shared_ptr p2p_proxy, - TransferEngine* engine); - ~ConnectionContext(); - - int32_t* warmup_send_region() const { return warmup_send_region_; } - int32_t* warmup_recv_region() const { return warmup_recv_region_; } - - /** - * @brief Get the total number of actively connected peers. - * @return The count of peers currently in the CONNECTED state. - */ - int getTotalConnectedPeers() const; - - /** - * @brief Expands the group to a new size. - * - * @note This is a non-blocking operation. Callers must invoke - * `waitUntilNewRanksConnected()` prior to initiating any - * subsequent communications (e.g., send, recv, putTaskCpu, - * putTaskCuda) to ensure the new peers are ready. - * - * @param newGroupSize The target size for the extended group. - */ - void extendGroupSizeTo(int newGroupSize); - - // Allow polling ranks beyond groupSize_ without changing groupSize_. - void setPollingLimitTo(int pollingLimit); - - /** - * @brief Checks whether all peers within the group have - * established connections. - * - * @return True if all peers are fully connected. - */ - bool isAllPeerConnected() const; - - /** - * @brief Blocks until all peers in the group are connected. - * - * This method is primarily used during backend initialization. - * Upon completion, it set `establishedGroupSize_` to the - * current `groupSize_`. - */ - void waitUntilAllConnected(); - - void bootstrapLocalPeer(const std::string& localServerName, - const SegmentInfo& localRankInfo); - - /** - * @brief Blocks until all newly added ranks in the - * extended group are connected. - * - * Specifically, it waits for pending ranks in the range - * `[establishedGroupSize_, groupSize_)` to reach the connected state. - * This should be called before starting new communications if - * `extendGroupSizeTo()` has been invoked. - * Upon completion, it set `establishedGroupSize_` to the - * current `groupSize_`. - */ - void waitUntilNewRanksConnected(); - - void shutdown(); - - void setDummy(bool isDummy) { isDummy_ = isDummy; } - - /** - * @brief Waits for the poller to stop all peer connections gracefully. - * - * Blocks until all peer connections have transitioned to the EXPIRING state - * or the timeout expires. Used during shutdown to ensure no pending - * transfers are active before resource cleanup. - * - * @return True if all peers stopped within the timeout; false otherwise. - */ - bool drainPoller() const; - - /** - * @brief Abandons resources instead of releasing them properly. - * - * When a hung operation prevents clean shutdown, this method marks - * resources as abandoned to prevent crashes during cleanup. - */ - void abandonResources(); - - static std::string getServerNameStoreKey(int backendIndex, int rank) { - return "server_name_" + std::to_string(backendIndex) + "_" + - std::to_string(rank); - } - static std::string getBufferStoreKey(int backendIndex, int rank) { - return "buffer_" + std::to_string(backendIndex) + "_" + - std::to_string(rank); - } - static std::string getExtensionStateStoreKey(int backendIndex, int rank) { - return "extension_state_" + std::to_string(backendIndex) + "_" + - std::to_string(rank); - } - - private: - // For ConnectionManager - bool poll(); - bool tryStop(); - bool isStopped() const; - - // Internal helpers - bool pollPeer(int pollingRank); -}; - -class ConnectionPoller { - static constexpr size_t kConnectingIdleSleepMs = 50; - static constexpr size_t kAllConnectedIdleSleepMs = 200; - - public: - static ConnectionPoller& GetInstance() { - // leaky singleton to avoid destructor fiasco problem - static ConnectionPoller* instance = new ConnectionPoller; - return *instance; - } - - void registerContext(const std::shared_ptr& ctx); - void removeContext(const std::shared_ptr& ctx); - - void wakeup() { - std::lock_guard lock(wakeup_mutex_); - wakeup_cv_.notify_all(); - } - // the global ranks - bool global_peerConnected_[kMaxNumRanks]{}; - - private: - ConnectionPoller(); - void ensureThreadStarted(); - void pollerLoop(); - bool processContext(const std::shared_ptr& ctx); - bool processPeer(const std::shared_ptr& ctx, - int pollingRank); - - std::mutex wakeup_mutex_; - std::condition_variable wakeup_cv_; - std::thread pollerThread_; - std::atomic pollerThreadStarted_{false}; - - std::mutex contexts_mutex_; - std::atomic contexts_version_{0}; - std::vector> contexts_; -}; - -} // namespace mooncake - -#endif // MOONCAKE_PG_CONNECTION_POLLER_H diff --git a/mooncake-pg/include/control_plane/agent.h b/mooncake-pg/include/control_plane/agent.h new file mode 100644 index 0000000000..a82e717c33 --- /dev/null +++ b/mooncake-pg/include/control_plane/agent.h @@ -0,0 +1,98 @@ +#ifndef MOONCAKE_PG_AGENT_H +#define MOONCAKE_PG_AGENT_H + +#include +#include +#include +#include +#include + +#include "error_types.h" +#include "rpc.h" + +namespace mooncake { + +// AgentStateMachine - Pure state machine for the control-plane client. +class AgentStateMachine { + public: + AgentStateMachine(GlobalRank rank, int max_world_size); + + AgentApplyResult registerGroup(const GroupView& group); + void unregisterGroup(GroupId group_id); + + AgentApplyResult handlePeerJoined(const PeerJoinedPush& push); + AgentApplyResult handleRankStateUpdate(const RankStatePush& push); + PGResult applyGroupView(const GroupView& view); + PGResult handleViewUpdate(const ViewUpdatePush& push); + + HeartbeatRequest buildHeartbeat() const; + + AgentApplyResult applyRegisterAgentResponse( + const RegisterAgentResponse& resp); + AgentApplyResult reset(uint64_t new_session_id); + + AgentApplyResult pushLinkEvent(const LinkEvent& event); + std::optional getLinkEventReport() const; + void handleLinkEventReportAck(const LinkEventReportAck& ack); + + GroupView getGroupView(GroupId group_id) const; + + enum class CoordinatorConnection { + Connected, + AgentRegistering, + Disconnected + }; + CoordinatorConnection getCoordinatorConnection() const { + return coordinator_connection_; + } + void setCoordinatorConnection(CoordinatorConnection state) { + coordinator_connection_ = state; + } + + uint64_t getAgentSessionId() const { + return agent_session_id_.load(std::memory_order_acquire); + } + + uint64_t getRankEpoch() const { + return self_rank_epoch_.load(std::memory_order_acquire); + } + + private: + GlobalRank rank_; + int max_world_size_; + + std::atomic agent_session_id_{0}; + std::atomic self_rank_epoch_{0}; + + std::unordered_map groups_; + + std::vector global_rank_states_; + std::vector global_rank_epochs_; + std::vector global_rank_state_versions_; + std::vector> rank_connections_; + + std::vector observed_link_state_; + std::vector observed_target_rank_epochs_; + uint64_t link_state_version_ = 0; + uint64_t acked_link_state_version_ = 0; + + CoordinatorConnection coordinator_connection_ = + CoordinatorConnection::Disconnected; + + bool rankInRange(GlobalRank rank) const { + return 0 <= rank && rank < max_world_size_; + } + + void appendApplyViewEffect(const GroupView& view, + AgentApplyResult& effects) const; + void appendApplyViewEffectsForRank(GlobalRank rank, + AgentApplyResult& effects) const; + void resetRankForNewEpoch(GlobalRank rank, uint64_t rank_epoch, + AgentApplyResult& effects); + bool recordLinkEvent(GlobalRank peer, uint64_t target_rank_epoch, + LinkEvent::EventType type); +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_AGENT_H diff --git a/mooncake-pg/include/control_plane/agent_host.h b/mooncake-pg/include/control_plane/agent_host.h new file mode 100644 index 0000000000..2e5ec36249 --- /dev/null +++ b/mooncake-pg/include/control_plane/agent_host.h @@ -0,0 +1,251 @@ +#ifndef MOONCAKE_PG_AGENT_HOST_H +#define MOONCAKE_PG_AGENT_HOST_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "agent.h" +#include "rpc.h" +#include "serialized_executor.h" +#include "link_manager.h" + +#include "error_types.h" +namespace mooncake { + +class RpcServer; +class RpcClient; +class MooncakeCommunicator; + +// ========================================================================= +// Control Plane Architecture (Agent side) +// ========================================================================= +// +// Each rank runs one AgentHost. It owns the AgentStateMachine (pure state +// machine) and drives it via a SerializedExecutor. +// +// MooncakeCommunicator AgentHost +// +-----------------+ +---------------------------+ +// | proposeActivate |-> (sync) --->| call Coordinator RPC | +// | registerGroup |-> post() --->| agent_.registerGroup() | +// | pushLinkEvent |-> post() --->| agent_.pushLinkEvent() | +// +-----------------+ +---------------------------+ +// | +// SerializedExecutor (tick) +// | +// +--------------------------------+ +// | AgentStateMachine | +// | (pure state machine, no I/O) | +// +--------------------------------+ +// | +// returns Effect list +// | +// +--------------------------------+ +// | runEffects() | +// | EnablePeerProbe -> LinkManager| +// | SendLinkEventReport -> RPC | +// | ApplyViewToCommunicator -> comm| +// | ... | +// +--------------------------------+ + +// AgentInterface - control-plane service interface exposed to +// MooncakeCommunicator. +class AgentInterface { + public: + virtual ~AgentInterface() = default; + + virtual PGResult waitUntilRegistered( + std::chrono::milliseconds timeout) = 0; + + virtual PGResult waitUntilGroupReady( + GroupId group_id, std::chrono::milliseconds timeout) = 0; + + virtual PGResult waitUntilRankActive( + GroupId group_id, GlobalRank rank, + std::chrono::milliseconds timeout) = 0; + + // Returns an empty GroupId when the Coordinator rejects this group. The + // rejected group is not inserted into the process-scoped Agent state. + virtual PGResult registerGroup( + GroupBootstrapId group_bootstrap_id, int32_t max_group_size, + std::vector rank_order, + GroupBootstrapIdResolvePolicy resolve_policy, bool auto_deactivate, + MooncakeCommunicator* communicator) = 0; + + virtual void detachCommunicator(GroupId group_id) = 0; + + virtual PGResult unregisterGroup(GroupId group_id) = 0; + + virtual PGResult confirmReadyForActivation(GroupId group_id) = 0; + + virtual PGResult publishLocalEndpoint( + GroupEndpointPublication endpoint) = 0; + + virtual PGResult proposeActivate( + GroupId group_id, const std::vector& ranks) = 0; + + virtual PGResult proposeDeactivate( + GroupId group_id, const std::vector& ranks) = 0; + + virtual void pushLinkEvent(const LinkEvent& event) = 0; + + virtual PGResult syncAfterFailure( + GroupId group_id) = 0; +}; + +class AgentHost; + +// AgentRpcServiceImpl - thin RPC handler for Coordinator->Agent pushes. +class AgentRpcServiceImpl : public AgentRpcService { + public: + explicit AgentRpcServiceImpl(AgentHost& host) : host_(host) {} + + void onPeerJoined(PeerJoinedPush push) override; + void onRankStateUpdate(RankStatePush push) override; + void onViewUpdate(coro_rpc::context ctx, + ViewUpdatePush push) override; + + private: + AgentHost& host_; +}; + +// AgentHost - execution host for the agent state machine. +class AgentHost : public AgentInterface { + public: + // Throttle repeated registerAgent error logs. + static constexpr auto kAgentRegisterErrorLogInterval = + std::chrono::seconds(5); + static constexpr auto kHeartbeatInterval = std::chrono::seconds(1); + + AgentHost(std::string coordinator_addr, const std::string& host_ip, + GlobalRank rank, int max_world_size, LinkManager& link_manager, + int64_t fault_reconciliation_window_us); + + ~AgentHost() override; + + PGResult start(); + void shutdown(); + void setFaultReconciliationWindow(int64_t timeout_us); + + PGResult waitUntilRegistered( + std::chrono::milliseconds timeout) override; + PGResult waitUntilGroupReady( + GroupId group_id, std::chrono::milliseconds timeout) override; + PGResult waitUntilRankActive( + GroupId group_id, GlobalRank rank, + std::chrono::milliseconds timeout) override; + + PGResult registerGroup( + GroupBootstrapId group_bootstrap_id, int32_t max_group_size, + std::vector rank_order, + GroupBootstrapIdResolvePolicy resolve_policy, bool auto_deactivate, + MooncakeCommunicator* communicator) override; + void detachCommunicator(GroupId group_id) override; + PGResult unregisterGroup(GroupId group_id) override; + PGResult confirmReadyForActivation(GroupId group_id) override; + PGResult publishLocalEndpoint( + GroupEndpointPublication endpoint) override; + + PGResult proposeActivate( + GroupId group_id, const std::vector& ranks) override; + + PGResult proposeDeactivate( + GroupId group_id, const std::vector& ranks) override; + + void pushLinkEvent(const LinkEvent& event) override; + + PGResult syncAfterFailure( + GroupId group_id) override; + + void postPeerJoined(PeerJoinedPush push); + void postRankStateUpdate(RankStatePush push); + void postViewUpdate(coro_rpc::context ctx, + ViewUpdatePush push); + + private: + AgentStateMachine agent_; + SerializedExecutor executor_; + + LinkManager& link_manager_; + + std::string host_ip_; + GlobalRank rank_; + int max_world_size_; + + std::string coordinator_addr_; + std::atomic fault_reconciliation_window_us_; + uint64_t agent_session_id_ = 0; + bool agent_session_initialized_ = false; + std::atomic shutdown_requested_{false}; + std::chrono::steady_clock::time_point next_heartbeat_at_; + + // RPC infrastructure. + std::unique_ptr rpc_server_; + std::unique_ptr rpc_client_; + std::unique_ptr rpc_impl_; + + // Bootstrap synchronization: one-shot latch with executor-managed promises. + bool agent_registration_done_ = false; + std::vector>> + agent_registration_promises_; + + // Throttling state for registerAgent error logs + std::chrono::steady_clock::time_point last_agent_register_error_log_time_; + + // group_ready_promises_ is fulfilled when registerGroup returns and + // the GroupView is applied. + std::unordered_map>>> + group_ready_promises_; + + // rank_active_promises_[group_id][rank] is fulfilled when a ViewUpdate + // push activates `rank` in `group_id`. Used by extension/replacement + // ranks to block in MooncakeCommunicator::joinGroup() until activation. + std::unordered_map< + GroupId, + std::unordered_map>>>> + rank_active_promises_; + + // Communicator registry: for view application and link reset. + // Accessed only from the executor thread. + std::unordered_map communicators_; + + void startAgentRegistration(bool start_new_session = false); + bool shouldLogAgentRegistrationError(); + void unregisterAgent(); + void tick(); + + PGResult sendPublishEndpointRpc(GroupEndpointPublication endpoint); + + void sendLinkEventReport(LinkEventReport report); + + PGResult proposeViewUpdateInternal( + GroupId group_id, const std::vector& ranks, + bool is_activation); + + void runEffects(const AgentApplyResult& effects); + template + void forEachCommunicator(F&& func) { + for (auto& [group_id, communicator] : communicators_) { + func(communicator); + } + } + template + void withCommunicator(GroupId group_id, F&& func) { + auto it = communicators_.find(group_id); + if (it != communicators_.end()) { + func(it->second); + } + } +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_AGENT_HOST_H diff --git a/mooncake-pg/include/control_plane/control_types.h b/mooncake-pg/include/control_plane/control_types.h new file mode 100644 index 0000000000..645e8f24a0 --- /dev/null +++ b/mooncake-pg/include/control_plane/control_types.h @@ -0,0 +1,170 @@ +#ifndef MOONCAKE_PG_CONTROL_PLANE_CONTROL_TYPES_H +#define MOONCAKE_PG_CONTROL_PLANE_CONTROL_TYPES_H + +#include +#include +#include +#include + +namespace mooncake { + +// There are two rank namespaces that are easy to confuse: +// +// * GlobalRank - process-wide identifier, range 0 .. max_world_size-1. +// Used for process-level states +// * InGroupRank - group-local identifier, range 0 .. group_size-1. +// Used inside a single process group and mapped to a +// GlobalRank through GroupView::rank_order. +using GlobalRank = int32_t; +using InGroupRank = int32_t; + +// Bootstrap ID: Backend type ("cpu:" or "device:") + PyTorch-assigned group_id. +using GroupBootstrapId = std::string; +// Coordinator-assigned unique group id +using GroupId = std::string; + +constexpr GlobalRank kInvalidGlobalRank = -1; +constexpr int kMaxNumRanks = 64; + +// Resolves a registration only against runtime groups stored under the same +// GroupBootstrapId, i.e. the same device kind and PyTorch group id. +// An exact match requires both rank_order and max_group_size to be equal. +// +// CreateOrAttach: +// * Attach when exactly one existing group is an exact match. +// * Create a new group when there is no exact match. +// * Reject multiple exact matches and never modify an existing rank_order. +// +// AttachOrExtend: +// * Attach when exactly one existing group is an exact match. +// * Otherwise, extend only when exactly one existing rank_order is a proper +// prefix, max_group_size matches, and the registering rank is in the +// appended suffix. +// * Reject zero or multiple compatible groups and never create a new +// group. +enum class GroupBootstrapIdResolvePolicy : uint8_t { + CreateOrAttach = 0, + AttachOrExtend = 1, +}; + +// Process-level state for a rank. +// All transitions are driven by the Coordinator. +enum class RankState : uint8_t { + Offline = 0, + Synced = 1, + Healthy = 2, +}; + +// Group-level, per-(group_id, rank) buffer/sync/P2P addresses. +struct GroupEndpointInfo { + // Coordinator-assigned endpoint version. + // The Agent publishes with 0 (it does not know the epoch); + // the Coordinator fills it in before pushing the ViewUpdate. + uint64_t endpoint_epoch = 0; + + // collective + uint64_t send_buffer[2] = {}; + uint64_t recv_buffer[2] = {}; + uint64_t send_sync[2] = {}; + uint64_t recv_sync[2] = {}; + + // p2p + uint64_t p2p_credit_region = 0; + uint64_t p2p_ack_region = 0; + + bool operator==(const GroupEndpointInfo&) const = default; +}; + +// State of one rank inside a GroupView. +// +// Founding member: +// None ------------(initial group declaration)-----------> Active +// +// Joining member: +// None ----------------(registerGroup)-------------------> Inactive +// Inactive ---(joinGroup confirms local preparation)-----> AwaitingActivation +// AwaitingActivation --(Coordinator activates the rank)--> Active +// +// Active -----------------------(deactivate)-----------------------> Inactive +// AwaitingActivation ------(new agent session or offline)----------> Inactive +// Inactive/AwaitingActivation/Active ----(unregisterGroup)---------> Left +// Left --------------------------(registerGroup)-------------------> Inactive +// +// Only Active participates in collectives. +// AwaitingActivation may become activatable once its endpoint, health, and +// connectivity are ready; Inactive may not. +enum class GroupMemberState : uint8_t { + None = 0, // slot has not registered with this group + Inactive = 1, // registered, but not ready for activation + AwaitingActivation = 2, // local join preparation is complete + Active = 3, // committed collective participant + Left = 4, // explicitly left (called destroy_group) +}; + +// Rank state inside a single GroupView. +struct GroupMember { + GroupMemberState status = GroupMemberState::None; + std::optional endpoint; + + bool isNone() const { return status == GroupMemberState::None; } + bool isActive() const { return status == GroupMemberState::Active; } + bool isAwaitingActivation() const { + return status == GroupMemberState::AwaitingActivation; + } + bool isMember() const { + return status == GroupMemberState::Inactive || + status == GroupMemberState::AwaitingActivation || + status == GroupMemberState::Active; + } + bool hasLeft() const { return status == GroupMemberState::Left; } + bool hasEndpoint() const { return endpoint.has_value(); } + + bool operator==(const GroupMember&) const = default; +}; + +// Group lifecycle status. +// +// Bootstrapping - collecting endpoints and waiting for all active ranks +// to become Healthy with valid endpoints. +// BootstrapSyncing - Coordinator initiated 2PC barrier; waiting for all +// active ranks to ACK the initial ViewUpdate. +// If a peer dies here, waitUntilGroupReady() hangs +// until its timeout. +// Ready - barrier complete; all ranks ready for data-plane +// transfers. +enum class GroupStatus : uint8_t { + Bootstrapping = 0, + BootstrapSyncing = 1, + Ready = 2, +}; + +// Runtime state for a group. +struct GroupView { + GroupId group_id; + GroupStatus status = GroupStatus::Bootstrapping; + uint64_t epoch = 0; + bool auto_deactivate = true; + int32_t max_group_size = 0; // fixed in-group slot capacity + std::vector rank_order; // InGroupRank -> GlobalRank + std::vector members; // indexed by GlobalRank + + bool operator==(const GroupView&) const = default; +}; + +struct LinkEvent { + enum class EventType : uint8_t { + None = 0, + Success = 1, + Failure = 2, + }; + + std::vector events; + // The Coordinator-assigned epoch of the target rank observed by the + // event source. This is parallel to events and prevents a late event for + // an old process incarnation from being attributed to its replacement. + std::vector target_rank_epochs; +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_CONTROL_PLANE_CONTROL_TYPES_H diff --git a/mooncake-pg/include/control_plane/coordinator.h b/mooncake-pg/include/control_plane/coordinator.h new file mode 100644 index 0000000000..d40764dc05 --- /dev/null +++ b/mooncake-pg/include/control_plane/coordinator.h @@ -0,0 +1,322 @@ +#ifndef MOONCAKE_PG_COORDINATOR_H +#define MOONCAKE_PG_COORDINATOR_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rpc.h" + +namespace mooncake { + +// CoordinatorStateMachine - abstract interface for the control-plane server +// state machine. +class CoordinatorStateMachine { + public: + virtual ~CoordinatorStateMachine() = default; + + virtual CoordinatorApplyResult handleRegisterAgent( + const RegisterAgentRequest& req) = 0; + + virtual CoordinatorApplyResult handleHeartbeat( + const HeartbeatRequest& req) = 0; + + virtual CoordinatorApplyResult + handleUnregisterAgent(const UnregisterAgentRequest& req) = 0; + + virtual CoordinatorApplyResult handleRegisterGroup( + const RegisterGroupRequest& req) = 0; + + virtual CoordinatorApplyResult + handleUnregisterGroup(const UnregisterGroupRequest& req) = 0; + + virtual CoordinatorApplyResult + handleConfirmReadyForActivation( + const ConfirmReadyForActivationRequest& req) = 0; + + virtual CoordinatorApplyResult + handlePublishEndpoint(const PublishEndpointRequest& req) = 0; + + virtual CoordinatorApplyResult handleProposeViewUpdate( + uint64_t propose_id, const ProposeViewUpdateRequest& req) = 0; + + virtual CoordinatorApplyResult handleLinkEventReport( + const LinkEventReport& req) = 0; + + virtual CoordinatorApplyResult handleSyncAfterFailure( + uint64_t sync_id, const SyncAfterFailureRequest& req) = 0; + + virtual CoordinatorApplyResult handleViewUpdateAck(GroupId group_id, + GlobalRank rank, + uint64_t epoch, + bool applied) = 0; + + virtual CoordinatorApplyResult tick() = 0; + + virtual CoordinatorApplyResult requestShutdown() = 0; +}; + +// CentralizedCoordinatorStateMachine - single-node implementation of the +// Coordinator state machine. +class CentralizedCoordinatorStateMachine : public CoordinatorStateMachine { + public: + explicit CentralizedCoordinatorStateMachine( + int max_world_size, + std::chrono::microseconds fault_reconciliation_window = + std::chrono::microseconds(50000)); + + void setFaultReconciliationWindow( + std::chrono::microseconds fault_reconciliation_window); + + CoordinatorApplyResult handleRegisterAgent( + const RegisterAgentRequest& req) override; + + CoordinatorApplyResult handleHeartbeat( + const HeartbeatRequest& req) override; + + CoordinatorApplyResult handleUnregisterAgent( + const UnregisterAgentRequest& req) override; + + CoordinatorApplyResult handleRegisterGroup( + const RegisterGroupRequest& req) override; + + CoordinatorApplyResult handleUnregisterGroup( + const UnregisterGroupRequest& req) override; + + CoordinatorApplyResult + handleConfirmReadyForActivation( + const ConfirmReadyForActivationRequest& req) override; + + CoordinatorApplyResult handlePublishEndpoint( + const PublishEndpointRequest& req) override; + + CoordinatorApplyResult handleProposeViewUpdate( + uint64_t propose_id, const ProposeViewUpdateRequest& req) override; + + CoordinatorApplyResult handleLinkEventReport( + const LinkEventReport& req) override; + + CoordinatorApplyResult handleSyncAfterFailure( + uint64_t sync_id, const SyncAfterFailureRequest& req) override; + + CoordinatorApplyResult handleViewUpdateAck(GroupId group_id, + GlobalRank rank, + uint64_t epoch, + bool applied) override; + + CoordinatorApplyResult tick() override; + + CoordinatorApplyResult requestShutdown() override; + + RankState getRankState(GlobalRank rank) const { + if (!rankInRange(rank)) return RankState::Offline; + return ranks_[rank].state; + } + + const std::string& getAgentAddr(GlobalRank rank) const { + return ranks_[rank].agent_addr; + } + + private: + int max_world_size_; + + struct RankInfo { + RankState state = RankState::Offline; + std::string agent_addr; + std::string te_server_name; + // Agent-generated key for one logical registration. + uint64_t agent_session_id = 0; + // Coordinator-assigned, monotonically increasing incarnation of this + // GlobalRank. Zero means no Agent has ever been accepted for it. + uint64_t rank_epoch = 0; + // Monotonically increasing version of the authoritative rank state. + uint64_t rank_state_version = 0; + std::chrono::steady_clock::time_point last_heartbeat; + std::vector link_status; + uint64_t last_link_event_report_id = 0; + uint64_t warmup_recv_addr = 0; + }; + + // Per-GlobalRank coordinator state. + std::vector ranks_; + + std::unordered_map group_views_; + + // A bootstrap id may name multiple runtime groups. The resolve policy + // distinguishes creation from attach/append resolution within the bucket. + std::unordered_map> + group_ids_by_bootstrap_id_; + std::unordered_map group_bootstrap_ids_; + uint64_t next_group_id_ = 1; + + // Coordinator-assigned endpoint epoch counter per GlobalRank. + // Incremented on every successful publishEndpoint for that rank so the + // Agent can detect endpoint changes. + std::vector endpoint_epochs_; + + struct PendingViewUpdateBarrier { + GroupId group_id; + uint64_t epoch = 0; + std::unordered_set waiting_acks; + std::unordered_set dropped_ranks; + std::optional deadline; + + struct ProposalCommit { + uint64_t propose_id = 0; + }; + struct BootstrapCommit {}; + + std::variant commit = + BootstrapCommit{}; + }; + std::unordered_map> + pending_barriers_; + + struct PendingProposal { + uint64_t propose_id = 0; + ProposeViewUpdateRequest request; + std::chrono::steady_clock::time_point deadline; + }; + // Membership proposals are linearized per group. A queued proposal is + // admitted only after the preceding membership barrier has committed. + std::unordered_map> pending_proposals_; + + struct PendingSync { + uint64_t sync_id = 0; + uint64_t agent_session_id = 0; + std::optional link_event_report_ack; + }; + + using PendingSyncs = std::unordered_map< + GroupId, std::unordered_map>>; + + struct FaultReconciliationContext { + bool active = false; + std::chrono::steady_clock::time_point deadline; + PendingSyncs pending_syncs; + }; + FaultReconciliationContext reconciliation_ctx_; + std::chrono::microseconds fault_reconciliation_window_; + + // requestShutdown() freezes the ranks whose current Agent sessions must + // end before the state machine asks the Host to stop serving RPCs. New + // sessions cannot take ownership after this snapshot is created. + bool shutdown_requested_ = false; + bool shutdown_confirmed_ = false; + std::unordered_set shutdown_pending_ranks_; + + static constexpr auto kHeartbeatTimeout = std::chrono::seconds(30); + + bool invalidateAgentSession(GlobalRank rank); + + void handleTimedOutAgent(GlobalRank rank, const char* reason, + std::vector& effects); + void tryConfirmShutdown(std::vector& effects); + + // Recompute the authoritative healthy set (max clique) and update + // rank-state between Healthy and Synced. Emits rank-state effects. + void updateRankStates(std::vector& effects); + + // For every auto_deactivate + ready group, mark active ranks that are not + // in the current healthy set as inactive. Increments view epoch and emits + // a ViewUpdate when at least one rank is pruned. + void applyAutoDeactivate(std::vector& effects); + + // Opens a fault reconciliation window if it is not open. + // An existing window is not extended. + void tryOpenReconciliationWindow(); + void tryCloseReconciliationWindow(std::vector& effects); + std::optional processLinkEventReport( + const LinkEventReport& report, std::vector& effects); + + void populateRegisterAgentResponse(RegisterAgentResponse& response, + GlobalRank rank) const; + + SyncAfterFailureResponse makeSyncResponse(SyncAfterFailureStatus status, + GroupId group_id) const; + void resolvePendingSyncs(std::vector& effects); + + bool isMutuallyConnected(GlobalRank a, GlobalRank b) const; + + // Preserve existing healthy ranks that are still mutually connected, + // then extend with new candidates that have full connectivity to all + // current healthy members. + std::vector extendHealthySet() const; + + // Bootstrap state machine driver. Advances groups through: + // Bootstrapping -> BootstrapSyncing (when all active ranks are Healthy + // and have published endpoints) + // BootstrapSyncing -> Ready (when all active ranks have ACKed) + // + // Called after every state-changing operation. + void checkGroupTransitions(std::vector& effects); + + void processGroupRegistration(const RegisterGroupRequest& request, + const GroupId& group_id, + std::vector& effects); + + bool validateGroupRegistration(const RegisterGroupRequest& request, + RegisterGroupResponse& response) const; + std::optional resolveGroupId(const RegisterGroupRequest& request, + RegisterGroupResponse& response, + bool& new_group); + void bindGroupBootstrapId(GroupId group_id, + GroupBootstrapId group_bootstrap_id); + + bool canEraseGroup(const GroupView& view) const; + void eraseGroup(GroupId group_id, std::vector& effects); + + bool isActivatableSet(GroupId group_id, + const std::vector& new_ranks, + const GroupView& old_view) const; + + bool isRankActivatable(GroupId group_id, GlobalRank rank, + const std::vector& future_active) const; + + // Admit proposals in Coordinator arrival order. At most one membership + // proposal per group may wait on a ViewUpdate barrier at a time. + void tryAdmitPendingProposals(GroupId group_id, + std::vector& effects); + void rejectPendingProposals(GroupId group_id, GlobalRank rank, + const std::string& reason, + std::vector& effects); + void dropRankFromPendingBarriers(GroupId group_id, GlobalRank rank, + std::vector& effects); + + // Helpers for barrier (proposal, bootstrap, ...) lifecycle. + void commitBarrier(PendingViewUpdateBarrier barrier, + std::vector& effects); + + // Compute the ACK set for a ViewUpdate barrier (proposal, bootstrap, ...). + std::unordered_set computeBarrierAckSet( + const GroupView& old_view, const GroupView& new_view) const; + + // Reject pending syncs for `group_id` / `group_id, rank`. + // Emits ReplySync for each pending sync_id. + void rejectPendingSyncs(GroupId group_id, GlobalRank rank, + const std::string& reason, + std::vector& effects); + + // Request validation: rank must be in range, online, and matching session. + bool hasValidSession(GlobalRank rank, uint64_t session_id) const { + return rankInRange(rank) && ranks_[rank].state != RankState::Offline && + ranks_[rank].agent_session_id == session_id; + } + + bool rankInRange(GlobalRank rank) const { + return 0 <= rank && rank < max_world_size_; + } + + CoordinatorEffect makeRankStateEffect(GlobalRank rank); +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_COORDINATOR_H diff --git a/mooncake-pg/include/control_plane/coordinator_host.h b/mooncake-pg/include/control_plane/coordinator_host.h new file mode 100644 index 0000000000..cb93a482cc --- /dev/null +++ b/mooncake-pg/include/control_plane/coordinator_host.h @@ -0,0 +1,195 @@ +#ifndef MOONCAKE_PG_COORDINATOR_HOST_H +#define MOONCAKE_PG_COORDINATOR_HOST_H + +#include +#include +#include +#include +#include +#include + +#include "coordinator.h" +#include "rpc.h" +#include "rpc_runtime.h" +#include "serialized_executor.h" +#include "error_types.h" + +namespace mooncake { + +class RpcServer; +class RpcClient; + +// ========================================================================= +// Control Plane Architecture (Coordinator side) +// ========================================================================= +// +// The CentralizedCoordinator runs inside Rank 0's process. It is the +// authoritative source of truth for rank health (RankState) and group +// membership (GroupView). +// +// Agent (any rank) CoordinatorHost (Rank 0) +// +-----------------+ +---------------------------+ +// | registerAgent |--- RPC ----->| postRegisterAgent() | +// | heartbeat |--- RPC ----->| postHeartbeat() | +// | proposeViewUpd |--- RPC ----->| postProposeViewUpdate() | +// | publishEndpoint |--- RPC ----->| postPublishEndpoint() | +// | reportLinkEvent |--- RPC ----->| postLinkEventReport() | +// +-----------------+ +---------------------------+ +// | +// SerializedExecutor +// | +// +------------------------------------+ +// | CentralizedCoordinatorStateMachine | +// | (pure state machine, no I/O) | +// +------------------------------------+ +// | +// returns Effect list +// | +// +-----------------------------------+ +// | runEffects() | +// | BroadcastRankState -> broadcast | +// | PushViewUpdate -> callAsync | +// | ReplyProposal -> reply | +// +-----------------------------------+ +// + +class CoordinatorHost; + +// CoordinatorRpcServiceImpl - thin RPC handler that forwards all calls +// to CoordinatorHost::post*(). +class CoordinatorRpcServiceImpl : public CoordinatorRpcService { + public: + explicit CoordinatorRpcServiceImpl(CoordinatorHost& host) : host_(host) {} + + void registerAgent(coro_rpc::context ctx, + RegisterAgentRequest req) override; + + void heartbeat(coro_rpc::context ctx, + HeartbeatRequest req) override; + + void unregisterAgent(coro_rpc::context ctx, + UnregisterAgentRequest req) override; + + void registerGroup(coro_rpc::context ctx, + RegisterGroupRequest req) override; + + void unregisterGroup(coro_rpc::context ctx, + UnregisterGroupRequest req) override; + + void confirmReadyForActivation( + coro_rpc::context ctx, + ConfirmReadyForActivationRequest req) override; + + void publishEndpoint(coro_rpc::context ctx, + PublishEndpointRequest req) override; + + void proposeViewUpdate(coro_rpc::context ctx, + ProposeViewUpdateRequest req) override; + + void reportLinkEvent(coro_rpc::context ctx, + LinkEventReport req) override; + + void syncAfterFailure(coro_rpc::context ctx, + SyncAfterFailureRequest req) override; + + private: + CoordinatorHost& host_; +}; + +// CoordinatorHost - execution host for the Coordinator state machine. +class CoordinatorHost { + public: + CoordinatorHost(const std::string& host_ip, int max_world_size, + int64_t fault_reconciliation_window_us); + + ~CoordinatorHost(); + + PGResult start(); + void shutdown(); + PGResult setFaultReconciliationWindow(int64_t timeout_us); + + const std::string& getListenAddr() const { return listen_addr_; } + + void postRegisterAgent(coro_rpc::context ctx, + RegisterAgentRequest req); + + void postHeartbeat(coro_rpc::context ctx, + HeartbeatRequest req); + + void postUnregisterAgent(coro_rpc::context ctx, + UnregisterAgentRequest req); + + void postRegisterGroup(coro_rpc::context ctx, + RegisterGroupRequest req); + + void postUnregisterGroup(coro_rpc::context ctx, + UnregisterGroupRequest req); + + void postConfirmReadyForActivation( + coro_rpc::context ctx, + ConfirmReadyForActivationRequest req); + + void postProposeViewUpdate(coro_rpc::context ctx, + ProposeViewUpdateRequest req); + + void postPublishEndpoint(coro_rpc::context ctx, + PublishEndpointRequest req); + + void postLinkEventReport(coro_rpc::context ctx, + LinkEventReport req); + + void postSyncAfterFailure(coro_rpc::context ctx, + SyncAfterFailureRequest req); + + void postViewUpdateAck(GroupId group_id, GlobalRank rank, uint64_t epoch, + bool applied); + + private: + CentralizedCoordinatorStateMachine state_machine_; + SerializedExecutor executor_; + + std::string host_ip_; + std::string listen_addr_; + int max_world_size_; + + // RPC infrastructure. + std::unique_ptr rpc_server_; + std::unique_ptr rpc_client_; + std::unique_ptr rpc_impl_; + + // Host only maintains deferred response context mapping. + // Related states is inside CentralizedCoordinatorStateMachine; + + uint64_t next_propose_id_{1}; + std::unordered_map> + pending_proposal_resps_; + + uint64_t next_sync_id_{1}; + std::unordered_map> + pending_sync_resps_; + + static constexpr auto kShutdownDrainTimeout = std::chrono::seconds(30); + + // The Host requests shutdown from the state machine and stops the RPC + // server once all sessions in the shutdown snapshot have unregistered. + std::atomic shutdown_requested_{false}; + std::promise shutdown_confirmation_; + + void runEffects(const std::vector& effects); + void pushViewUpdate(const PushViewUpdate& effect); + + template + void pushToAgent(GlobalRank rank, const Push& msg) { + const auto& addr = state_machine_.getAgentAddr(rank); + if (addr.empty()) { + LOG(WARNING) << "[COORD] push target rank=" << rank + << " has no agent_addr; skipping"; + return; + } + rpc_client_->send(addr, msg); + } +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_COORDINATOR_HOST_H diff --git a/mooncake-pg/include/control_plane/link_manager.h b/mooncake-pg/include/control_plane/link_manager.h new file mode 100644 index 0000000000..e88d75082a --- /dev/null +++ b/mooncake-pg/include/control_plane/link_manager.h @@ -0,0 +1,161 @@ +#ifndef MOONCAKE_PG_LINK_MANAGER_H +#define MOONCAKE_PG_LINK_MANAGER_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "control_plane/control_types.h" +#include "error_types.h" + +namespace mooncake { + +struct TELinkUpEvent { + GlobalRank peer = kInvalidGlobalRank; + uint64_t target_rank_epoch = 0; +}; + +// Process-level manager for shared TE states. +class LinkManager { + public: + LinkManager() = default; + + PGResult init(GlobalRank rank, int max_world_size, + TransferEngine* engine); + void start(uint64_t self_rank_epoch); + void stop(); + + bool isInitialized() const { + return initialized_.load(std::memory_order_acquire); + } + + void shutdown(); + + std::string localServerName() const; + + uint64_t getWarmupRecvAddr() const; + + void enablePeerProbe(GlobalRank peer, uint64_t target_rank_epoch, + const std::string& server_name, + uint64_t warmup_recv_addr = 0); + + void disconnect(GlobalRank peer); + + void requestHealthCheck(GlobalRank peer); + + void stopReconnect(GlobalRank peer); + + bool isConnected(GlobalRank peer) const; + + using EventCallback = std::function; + void setEventCallback(EventCallback callback); + + std::optional resolvePeer( + GlobalRank peer) const; + + void refreshPeerSegment(GlobalRank peer); + + void publishLinkUp(GlobalRank peer, TransferMetadata::SegmentID target_id, + uint64_t target_rank_epoch); + void publishLinkDown(GlobalRank peer); + + ~LinkManager() { shutdown(); } + + LinkManager(const LinkManager&) = delete; + LinkManager& operator=(const LinkManager&) = delete; + + private: + bool rankInRange(GlobalRank peer) const { + return 0 <= peer && peer < max_world_size_; + } + + enum class PeerLinkState : uint8_t { + Idle = 0, + WaitingWarmupTransfer, + WaitingPeerWarmup, + Connected, + }; + + struct PeerLink { + PeerLinkState state = PeerLinkState::Idle; + uint64_t target_rank_epoch = 0; + std::string server_name; + bool is_candidate = false; + bool skip_warmup = false; + + std::optional target_id; + uint64_t warmup_recv_addr = 0; + + // Only meaningful while state == Connected; it does not prevent the + // application from resolving or using the existing target_id. + bool health_check_requested = false; + + std::optional probe_batch_id; + std::chrono::steady_clock::time_point next_probe_time; + static constexpr auto kProbeBackoffMin = + std::chrono::milliseconds(1000); + static constexpr auto kProbeBackoffMax = + std::chrono::milliseconds(10000); + std::chrono::milliseconds probe_backoff{kProbeBackoffMin}; + }; + + std::vector peers_; + mutable std::mutex peers_mutex_; + + struct PeerReadState { + std::atomic version{0}; + std::atomic link_connected{0}; + std::atomic target_id{}; + }; + + std::vector read_state_; + + GlobalRank rank_ = kInvalidGlobalRank; + int max_world_size_ = 0; + TransferEngine* engine_ = nullptr; + std::string local_server_name_; + + // Warmup region + std::unique_ptr warmup_send_region_; + std::unique_ptr warmup_recv_region_; + bool skip_warmup_ = false; + + EventCallback event_callback_; + mutable std::mutex event_callback_mutex_; + + // Poller thread + std::thread poller_thread_; + std::atomic poller_running_{false}; + std::mutex wakeup_mutex_; + std::condition_variable wakeup_cv_; + + std::atomic initialized_{false}; + std::atomic started_{false}; + std::atomic shutdown_{false}; + + void pollerLoop(); + bool advanceConnection(GlobalRank peer); + bool advanceHealthCheck(GlobalRank peer); + void tearDownPeerLink(GlobalRank peer); + void emit(TELinkUpEvent event); + void wakeup(); + + static bool supportFabricMem(); + + static constexpr size_t kPollerIdleSleepMs = 200; + static constexpr size_t kPollerActiveSleepMs = 10; +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_LINK_MANAGER_H diff --git a/mooncake-pg/include/control_plane/rpc.h b/mooncake-pg/include/control_plane/rpc.h new file mode 100644 index 0000000000..433aa4f564 --- /dev/null +++ b/mooncake-pg/include/control_plane/rpc.h @@ -0,0 +1,370 @@ +#ifndef MOONCAKE_PG_CONTROL_PLANE_RPC_H +#define MOONCAKE_PG_CONTROL_PLANE_RPC_H + +#include +#include +#include +#include +#include + +#include + +#include + +#include "control_plane/control_types.h" + +namespace mooncake { + +inline constexpr auto kProposalAdmissionTimeout = std::chrono::seconds(20); +inline constexpr auto kViewUpdateAckTimeout = std::chrono::seconds(20); + +// Agent -> Coordinator RPC messages + +struct RegisterAgentRequest { + GlobalRank rank = kInvalidGlobalRank; + std::string agent_addr; + std::string te_server_name; + uint64_t agent_session_id = 0; + uint64_t warmup_recv_addr = 0; +}; + +struct RankConnectionMetadata { + GlobalRank rank = kInvalidGlobalRank; + uint64_t rank_epoch = 0; + std::string agent_addr; + std::string te_server_name; + uint64_t warmup_recv_addr = 0; +}; + +struct RegisterAgentResponse { + bool success = false; + std::string reject_reason; + // The request reached the Coordinator, but this logical registration can + // no longer be accepted. + bool require_new_session = false; + // Coordinator-assigned epoch for the accepted incarnation of the + // registering rank. Zero means that no incarnation has been accepted. + uint64_t rank_epoch = 0; + std::vector all_rank_states; + std::vector all_rank_epochs; + std::vector all_rank_state_versions; + std::vector groups; + std::vector rank_connections; +}; + +struct LinkEventReport { + GlobalRank reporter_rank = kInvalidGlobalRank; + uint64_t agent_session_id = 0; + uint64_t reporter_rank_epoch = 0; + uint64_t report_id = 0; + std::vector events; + // Parallel to events. Each entry identifies the target + // incarnation against which the observation was made. + std::vector target_rank_epochs; +}; + +struct LinkEventReportAck { + GlobalRank reporter_rank = kInvalidGlobalRank; + uint64_t reporter_rank_epoch = 0; + uint64_t report_id = 0; +}; + +struct HeartbeatRequest { + GlobalRank rank = kInvalidGlobalRank; + uint64_t agent_session_id = 0; +}; + +struct HeartbeatResponse { + bool require_new_session = false; +}; + +struct UnregisterAgentRequest { + GlobalRank rank = kInvalidGlobalRank; + uint64_t agent_session_id = 0; +}; + +struct UnregisterAgentResponse { + bool success = false; + std::string reject_reason; +}; + +struct RegisterGroupRequest { + GlobalRank rank = kInvalidGlobalRank; + uint64_t agent_session_id = 0; + GroupBootstrapId group_bootstrap_id; + int32_t max_group_size = 0; + std::vector rank_order; + GroupBootstrapIdResolvePolicy resolve_policy = + GroupBootstrapIdResolvePolicy::CreateOrAttach; + bool auto_deactivate = true; +}; + +struct RegisterGroupResponse { + bool success = false; + std::string reject_reason; + GroupView view; +}; + +struct ConfirmReadyForActivationRequest { + GroupId group_id; + GlobalRank rank = kInvalidGlobalRank; + uint64_t agent_session_id = 0; +}; + +struct ConfirmReadyForActivationResponse { + bool success = false; + std::string reject_reason; +}; + +enum class ProposalStatus : uint8_t { + Rejected = 0, + Applied = 1, + AppliedWithDroppedRanks = 2, +}; + +struct ProposeViewUpdateRequest { + GroupId group_id; + GlobalRank source_rank = kInvalidGlobalRank; + uint64_t agent_session_id = 0; + std::vector requested_ranks; + bool is_activation = false; +}; + +struct ProposeViewUpdateResponse { + ProposalStatus status = ProposalStatus::Rejected; + uint64_t new_epoch = 0; + std::vector dropped_ranks; + std::string reject_reason; +}; + +struct GroupEndpointPublication { + GroupId group_id; + GroupEndpointInfo endpoint_info; +}; + +struct PublishEndpointRequest { + GlobalRank rank = kInvalidGlobalRank; + uint64_t agent_session_id = 0; + std::vector endpoints; +}; + +struct PublishEndpointResponse { + bool success = false; + std::string reject_reason; +}; + +struct UnregisterGroupRequest { + GroupId group_id; + GlobalRank rank = kInvalidGlobalRank; + uint64_t agent_session_id = 0; +}; + +struct UnregisterGroupResponse { + bool success = false; + std::string reject_reason; +}; + +struct SyncAfterFailureRequest { + GroupId group_id; + GlobalRank reporter_rank = kInvalidGlobalRank; + uint64_t agent_session_id = 0; + uint64_t current_epoch = 0; + // Piggybacked link event report. + std::optional link_event_report; +}; + +enum class SyncAfterFailureStatus : uint8_t { + Reconciled = 0, // A reconciliation window completed. + NoPending = 1, // No reconciliation was pending at request time. + Rejected = 2, // Invalid request (stale session, group not found). +}; + +struct SyncAfterFailureResponse { + SyncAfterFailureStatus status = SyncAfterFailureStatus::Rejected; + GroupView view; + // Piggybacked link event report ack. + std::optional link_event_report_ack; + std::string reject_reason; +}; + +// Coordinator -> Agent RPC messages + +struct PeerJoinedPush { + GlobalRank rank = kInvalidGlobalRank; + uint64_t rank_epoch = 0; + std::string te_server_name; + uint64_t warmup_recv_addr = 0; +}; + +struct RankStatePush { + GlobalRank rank = kInvalidGlobalRank; + uint64_t rank_epoch = 0; + uint64_t rank_state_version = 0; + RankState new_state = RankState::Offline; +}; + +struct ViewUpdatePush { + GroupView view; +}; + +struct ViewUpdateAck { + GlobalRank rank = kInvalidGlobalRank; + GroupId group_id; + uint64_t epoch = 0; + bool applied = false; + std::string error_msg; +}; + +// Coordinator effects + +struct BroadcastPeerJoined { + PeerJoinedPush push; +}; + +struct BroadcastRankState { + RankStatePush push; +}; + +struct PushViewUpdate { + GroupView view; +}; + +struct ReplyProposal { + uint64_t propose_id = 0; + ProposeViewUpdateResponse response; +}; + +struct ReplySync { + uint64_t sync_id = 0; + SyncAfterFailureResponse response; +}; + +struct ShutdownCoordinatorHost {}; + +using CoordinatorEffect = + std::variant; + +// Agent effects + +struct EnablePeerProbe { + GlobalRank rank = kInvalidGlobalRank; + uint64_t rank_epoch = 0; + std::string te_server_name; + uint64_t warmup_recv_addr = 0; +}; + +struct DisconnectLink { + GlobalRank peer = kInvalidGlobalRank; +}; + +struct RequestLinkHealthCheck { + GlobalRank peer = kInvalidGlobalRank; +}; + +struct SendLinkEventReport { + LinkEventReport report; +}; + +struct StopReconnect { + GlobalRank peer = kInvalidGlobalRank; +}; + +struct DisconnectAllLinks {}; + +struct ClearAllPeerMetadata {}; + +struct ApplyViewToCommunicator { + GroupView view; + std::vector rank_states; + std::vector rank_epochs; + std::vector activatable; +}; + +struct ResetPeerState { + GlobalRank peer = kInvalidGlobalRank; +}; + +struct RefreshPeerLink { + GlobalRank peer = kInvalidGlobalRank; +}; + +struct NotifyLinkRefreshed { + GlobalRank peer = kInvalidGlobalRank; +}; + +struct NotifyGroupReady { + GroupId group_id; +}; + +struct NotifyRanksActivated { + GroupId group_id; + std::vector ranks; +}; + +using AgentEffect = + std::variant; + +// Results produced by the Coordinator/Agent state machine + +template +struct CoordinatorApplyResult { + Response response; + std::vector effects; +}; + +template <> +struct CoordinatorApplyResult { + std::vector effects; +}; + +using AgentApplyResult = std::vector; + +// RPC Service interfaces + +class CoordinatorRpcService { + public: + virtual ~CoordinatorRpcService() = default; + + virtual void registerAgent(coro_rpc::context ctx, + RegisterAgentRequest req) = 0; + virtual void heartbeat(coro_rpc::context ctx, + HeartbeatRequest req) = 0; + virtual void unregisterAgent(coro_rpc::context ctx, + UnregisterAgentRequest req) = 0; + virtual void registerGroup(coro_rpc::context ctx, + RegisterGroupRequest req) = 0; + virtual void unregisterGroup(coro_rpc::context ctx, + UnregisterGroupRequest req) = 0; + virtual void confirmReadyForActivation( + coro_rpc::context ctx, + ConfirmReadyForActivationRequest req) = 0; + virtual void proposeViewUpdate( + coro_rpc::context ctx, + ProposeViewUpdateRequest req) = 0; + virtual void publishEndpoint(coro_rpc::context ctx, + PublishEndpointRequest req) = 0; + virtual void reportLinkEvent(coro_rpc::context ctx, + LinkEventReport req) = 0; + virtual void syncAfterFailure( + coro_rpc::context ctx, + SyncAfterFailureRequest req) = 0; +}; + +class AgentRpcService { + public: + virtual ~AgentRpcService() = default; + + virtual void onPeerJoined(PeerJoinedPush push) = 0; + virtual void onRankStateUpdate(RankStatePush push) = 0; + virtual void onViewUpdate(coro_rpc::context ctx, + ViewUpdatePush push) = 0; +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_CONTROL_PLANE_RPC_H diff --git a/mooncake-pg/include/control_plane/rpc_runtime.h b/mooncake-pg/include/control_plane/rpc_runtime.h new file mode 100644 index 0000000000..7bd8b16ee7 --- /dev/null +++ b/mooncake-pg/include/control_plane/rpc_runtime.h @@ -0,0 +1,242 @@ +#ifndef MOONCAKE_PG_RPC_RUNTIME_H +#define MOONCAKE_PG_RPC_RUNTIME_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include + +#include +#include + +#include "error_types.h" + +namespace mooncake { + +class RpcServer { + public: + explicit RpcServer(uint16_t port = 0, unsigned thread_num = 2); + + template + void registerHandler(util::class_type_t* impl) { + server_->register_handler(impl); + } + + bool start(); + std::string getListenAddr(const std::string& host_ip) const; + void shutdown(); + + private: + std::unique_ptr server_; + uint16_t port_; + unsigned thread_num_; +}; + +class RpcClient { + private: + template + using ResponseT = decltype(coro_rpc::get_return_type()); + + template + static PGResult> makeRpcFailure(std::string_view message) { + return makePGError(PGErrorCode::RpcError, + std::string(coro_rpc::get_func_name()) + + " RPC failed: " + std::string(message)); + } + + static std::string describeException(std::exception_ptr exception) { + try { + std::rethrow_exception(std::move(exception)); + } catch (const std::exception& e) { + return e.what(); + } catch (...) { + return "unknown transport exception"; + } + } + + public: + static constexpr auto kConnectTimeout = std::chrono::seconds(3); + static constexpr auto kDefaultRequestTimeout = + std::chrono::milliseconds(30000); + + explicit RpcClient( + std::chrono::milliseconds request_timeout = kDefaultRequestTimeout, + std::chrono::milliseconds connect_timeout = kConnectTimeout); + ~RpcClient() = default; + + template + PGResult> call( + const std::string& addr, Req req, + std::optional timeout = std::nullopt) { + auto client = createSyncClient(); + auto request_timeout = timeout.value_or(state_->request_timeout); + + auto ec = async_simple::coro::syncAwait(client->connect(addr)); + if (ec) { + return makeRpcFailure(ec.message()); + } + auto rpc_result = async_simple::coro::syncAwait( + client->call_for(request_timeout, std::move(req))); + if (!rpc_result) { + return makeRpcFailure(rpc_result.error().msg); + } + + if constexpr (std::is_void_v>) { + return {}; + } else { + return std::move(rpc_result).value(); + } + } + + // Async call with transport failures and the response delivered through + // one PGResult. The callback is not invoked after shutdown begins. + template + void callAsync(const std::string& addr, Req req, Callback cb) { + static_assert( + std::is_invocable_v>>, + "RpcClient::callAsync callback must accept PGResult"); + auto task = callAsyncCoroutine(state_, addr, std::move(req), + std::move(cb)); + spawn(std::move(task)); + } + + // Fire-and-forget. + template + void send(const std::string& addr, Req req) { + auto task = sendCoroutine(state_, addr, std::move(req)); + spawn(std::move(task)); + } + + bool isConnected(const std::string& addr) const; + bool tryReconnect(const std::string& addr); + + void shutdown() { state_->shutdown.store(true, std::memory_order_release); } + + private: + struct SharedState { + SharedState(std::chrono::milliseconds request_timeout, + std::chrono::milliseconds connect_timeout) + : request_timeout(request_timeout), + connect_timeout(connect_timeout) {} + + std::mutex mutex; + std::unordered_map> + clients; + std::atomic shutdown{false}; + const std::chrono::milliseconds request_timeout; + const std::chrono::milliseconds connect_timeout; + }; + + // Coroutine-based connect + cache lookup. + static async_simple::coro::Lazy> + getOrCreateClient(std::shared_ptr state, + const std::string& addr); + + // Spawn a coroutine on the global I/O executor. + static void spawn(async_simple::coro::Lazy task); + + // Create a coro_rpc_client with local io_context (for sync call()). + std::unique_ptr createSyncClient(); + + // Fire-and-forget coroutine: connect, send_request, discard result. + template + static async_simple::coro::Lazy sendCoroutine( + std::shared_ptr state, const std::string& addr, Req req) { + if (state->shutdown.load(std::memory_order_acquire)) co_return; + auto client = co_await getOrCreateClient(state, addr); + if (!client) co_return; + try { + coro_rpc::request_config_t config; + config.request_timeout_duration = state->request_timeout; + auto send_lazy = co_await client->template send_request( + std::move(config), std::move(req)); + co_await std::move(send_lazy); + } catch (const std::exception& e) { + if (!state->shutdown.load(std::memory_order_acquire)) { + VLOG(1) << "RpcClient: fire-and-forget RPC to " << addr + << " failed: " << e.what(); + } + } + } + + // Async call coroutine: connect, send_request, invoke callback. + template + static async_simple::coro::Lazy callAsyncCoroutine( + std::shared_ptr state, const std::string& addr, Req req, + Callback cb) { + if (state->shutdown.load(std::memory_order_acquire)) co_return; + + using Response = ResponseT; + auto complete = [&](PGResult result) { + if (!state->shutdown.load(std::memory_order_acquire)) { + cb(std::move(result)); + } + }; + + auto client_try = co_await getOrCreateClient(state, addr).coAwaitTry(); + if (client_try.hasError()) { + complete(makeRpcFailure( + describeException(client_try.getException()))); + co_return; + } + auto client = std::move(client_try).value(); + if (!client) { + complete(makeRpcFailure("failed to connect to " + addr)); + co_return; + } + if (state->shutdown.load(std::memory_order_acquire)) co_return; + + coro_rpc::request_config_t config; + config.request_timeout_duration = state->request_timeout; + auto send_operation = client->template send_request( + std::move(config), std::move(req)); + auto send_try = co_await std::move(send_operation).coAwaitTry(); + if (send_try.hasError()) { + complete(makeRpcFailure( + describeException(send_try.getException()))); + co_return; + } + + auto receive_operation = std::move(send_try).value(); + auto receive_try = co_await std::move(receive_operation).coAwaitTry(); + if (receive_try.hasError()) { + complete(makeRpcFailure( + describeException(receive_try.getException()))); + co_return; + } + + auto rpc_result = std::move(receive_try).value(); + if (!rpc_result) { + complete(makeRpcFailure(rpc_result.error().msg)); + } else if constexpr (std::is_void_v) { + complete(PGResult{}); + } else { + complete( + PGResult{std::move(rpc_result.value().result())}); + } + } + + std::shared_ptr state_; +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_RPC_RUNTIME_H diff --git a/mooncake-pg/include/control_plane/serialized_executor.h b/mooncake-pg/include/control_plane/serialized_executor.h new file mode 100644 index 0000000000..dfae5f2d1e --- /dev/null +++ b/mooncake-pg/include/control_plane/serialized_executor.h @@ -0,0 +1,214 @@ +#ifndef MOONCAKE_PG_SERIALIZED_EXECUTOR_H +#define MOONCAKE_PG_SERIALIZED_EXECUTOR_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "error_types.h" + +namespace mooncake { + +// Generic single-threaded serialized executor. +// +// All tasks posted to this executor are processed sequentially on a single +// dedicated thread. +// +// Thread safety: +// - post() may be called from any thread. +// - The tick callback and all tasks run on the executor thread. +class SerializedExecutor { + private: + // Flatten postAndWait task results: + // T -> PGResult + // PGResult -> PGResult + template + struct PostAndWaitResult { + using type = PGResult; + }; + + template + struct PostAndWaitResult> { + using type = PGResult; + }; + + template + using PostAndWaitResultT = typename PostAndWaitResult::type; + + public: + SerializedExecutor() = default; + explicit SerializedExecutor(std::string name) : name_(std::move(name)) {} + + ~SerializedExecutor() { shutdown(); } + + // Start the executor thread. + void start() { + if (running_.exchange(true, std::memory_order_acq_rel)) return; + thread_ = std::thread([this] { loop(); }); + } + + // Gracefully stop the executor thread. Blocks until the thread exits. + // + // Drains any tasks that were still in the queue when the loop exited so + // that RPC contexts (held inside lambda captures) receive a response + // rather than being leaked. + void shutdown() { + if (!running_.exchange(false, std::memory_order_acq_rel)) return; + cv_.notify_one(); + if (thread_.joinable()) thread_.join(); + // Execute any tasks that were posted between the last loop iteration + // and now (e.g. callbacks from the RPC background thread). These run + // on the caller's thread. + std::vector> leftover; + { + std::lock_guard lock(mutex_); + leftover.swap(queue_); + } + for (auto& task : leftover) { + try { + task(); + } catch (const std::exception& e) { + fprintf(stderr, + "SerializedExecutor(%s): unhandled exception in " + "leftover task during shutdown: %s\n", + name_.c_str(), e.what()); + } catch (...) { + fprintf(stderr, + "SerializedExecutor(%s): unhandled non-std exception " + "in leftover task during shutdown\n", + name_.c_str()); + } + } + } + + // Post a task and block the caller until it completes on the executor + // thread. The task's return value is wrapped in PGResult, while an existing + // PGResult is forwarded without nesting. Exceptions thrown inside the task + // are re-raised on the caller. + // + // Must NOT be called from the executor thread itself (deadlock). + template + auto postAndWait(F&& f) -> PostAndWaitResultT> { + using R = std::invoke_result_t; + using Result = PostAndWaitResultT; + std::promise promise; + auto future = promise.get_future(); + PG_TRY(post([&promise, fn = std::forward(f)]() mutable { + try { + if constexpr (std::is_void_v) { + fn(); + promise.set_value(); + } else { + promise.set_value(fn()); + } + } catch (...) { + promise.set_exception(std::current_exception()); + } + })); + if constexpr (std::is_void_v) { + future.get(); + return {}; + } else { + return Result{future.get()}; + } + } + + // Post a task to the executor. May be called from any thread. + // Tasks are processed in FIFO order on the executor thread. + // Returns InvalidState if the executor is not running. + PGResult post(std::function task) { + PG_VALIDATE_STATE(running_.load(std::memory_order_acquire), + "SerializedExecutor(" + name_ + ") is not running"); + { + std::lock_guard lock(mutex_); + // Double-check: running_ may have flipped between the check above + // and acquiring the mutex. If the executor has stopped, drop + // the task to avoid leaking it (shutdown already drained the + // queue). + PG_VALIDATE_STATE( + running_.load(std::memory_order_acquire), + "SerializedExecutor(" + name_ + ") is not running"); + queue_.push_back(std::move(task)); + } + cv_.notify_one(); + return {}; + } + + // Set a callback that fires after every batch of tasks, even empty + // batches (the wait_for timeout ensures it fires roughly every 50 ms). + // Only call this before start() or from the executor thread itself. + void setTickCallback(std::function cb) { + tick_callback_ = std::move(cb); + } + + private: + void loop() { + while (running_.load(std::memory_order_acquire)) { + std::vector> batch; + { + std::unique_lock lock(mutex_); + // Wait up to 50 ms for work. This also guarantees tick + // fires at least every ~50 ms even when idle. + cv_.wait_for(lock, std::chrono::milliseconds(50), [this] { + return !queue_.empty() || + !running_.load(std::memory_order_acquire); + }); + batch.swap(queue_); + } + for (auto& task : batch) { + try { + task(); + } catch (const std::exception& e) { + // Log and swallow: a single bad task must not kill the + // loop. + fprintf(stderr, + "SerializedExecutor(%s): unhandled exception in " + "task: %s\n", + name_.c_str(), e.what()); + } catch (...) { + fprintf( + stderr, + "SerializedExecutor(%s): unhandled non-std exception " + "in task\n", + name_.c_str()); + } + } + if (tick_callback_) { + try { + tick_callback_(); + } catch (const std::exception& e) { + fprintf(stderr, + "SerializedExecutor(%s): unhandled exception in " + "tick callback: %s\n", + name_.c_str(), e.what()); + } catch (...) { + fprintf( + stderr, + "SerializedExecutor(%s): unhandled non-std exception " + "in tick callback\n", + name_.c_str()); + } + } + } + } + + std::string name_; + std::thread thread_; + std::atomic running_{false}; + std::mutex mutex_; + std::condition_variable cv_; + std::vector> queue_; + std::function tick_callback_; +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_SERIALIZED_EXECUTOR_H diff --git a/mooncake-pg/include/error_types.h b/mooncake-pg/include/error_types.h new file mode 100644 index 0000000000..be4ff61940 --- /dev/null +++ b/mooncake-pg/include/error_types.h @@ -0,0 +1,145 @@ +#ifndef MOONCAKE_PG_ERROR_TYPES_H +#define MOONCAKE_PG_ERROR_TYPES_H + +#include +#include +#include +#include +#include + +#include + +namespace mooncake { + +class PGAssertionException : public std::runtime_error { + public: + using std::runtime_error::runtime_error; +}; + +namespace detail { + +template +[[noreturn]] inline void throwPGAssertFailure(Args&&... args) { + std::ostringstream message; + (message << ... << std::forward(args)); + throw PGAssertionException(message.str()); +} + +} // namespace detail + +// Keep the order and values synchronized with mooncakePgResult_t. +enum class PGErrorCode : uint8_t { + InvalidArgument = 1, + InvalidState = 2, + NotSupported = 3, + Timeout = 4, + ResourceBusy = 5, + TransferEngineError = 6, + RpcError = 7, + SystemError = 8, + InternalError = 9, +}; + +struct PGError { + PGErrorCode code; + std::string message; +}; + +template +using PGResult = ylt::expected; + +inline auto makePGError(PGError error) { + return ylt::unexpected{std::move(error)}; +} + +inline auto makePGError(PGErrorCode code, std::string message) { + return makePGError(PGError{code, std::move(message)}); +} + +} // namespace mooncake + +#define PG_ASSERT(condition, ...) \ + do { \ + if (!(condition)) { \ + ::mooncake::detail::throwPGAssertFailure(__VA_ARGS__); \ + } \ + } while (false) + +#define PG_ASSERT_CUDA(expression) \ + do { \ + const auto pg_cuda_error_internal = (expression); \ + PG_ASSERT(pg_cuda_error_internal == cudaSuccess, #expression, \ + " failed: ", cudaGetErrorString(pg_cuda_error_internal)); \ + } while (false) + +#define PG_DETAIL_TRY(expression) \ + do { \ + auto&& pg_result_internal = (expression); \ + if (!pg_result_internal.has_value()) { \ + return ::mooncake::makePGError( \ + std::move(pg_result_internal).error()); \ + } \ + } while (false) + +#define PG_DETAIL_CONCAT_IMPL(left, right) left##right +#define PG_DETAIL_CONCAT(left, right) PG_DETAIL_CONCAT_IMPL(left, right) + +#define PG_DETAIL_TRY_ASSIGN_IMPL(result_name, lhs, expression) \ + auto&& result_name = (expression); \ + if (!result_name.has_value()) { \ + return ::mooncake::makePGError(std::move(result_name).error()); \ + } \ + lhs = std::move(result_name).value() + +#define PG_DETAIL_TRY_ASSIGN(lhs, ...) \ + PG_DETAIL_TRY_ASSIGN_IMPL(PG_DETAIL_CONCAT(pg_result_internal_, __LINE__), \ + lhs, (__VA_ARGS__)) + +#define PG_DETAIL_FIRST(first, ...) first + +// Propagate an error with PG_TRY(expression). PG_TRY(lhs, expression) also +// assigns the successful value to lhs, which may be a declaration such as +// `auto value`. +#define PG_TRY(first, ...) \ + PG_DETAIL_FIRST(__VA_OPT__(PG_DETAIL_TRY_ASSIGN, ) \ + PG_DETAIL_TRY)(first __VA_OPT__(, ) __VA_ARGS__) + +#define PG_TRY_TE(expression) \ + do { \ + const int pg_te_error_internal = (expression); \ + if (pg_te_error_internal != 0) { \ + return ::mooncake::makePGError( \ + ::mooncake::PGErrorCode::TransferEngineError, \ + std::string(#expression) + \ + " failed, rc=" + std::to_string(pg_te_error_internal)); \ + } \ + } while (false) + +#define PG_TRY_CUDA(expression) \ + do { \ + const auto pg_cuda_error_internal = (expression); \ + if (pg_cuda_error_internal != cudaSuccess) { \ + return ::mooncake::makePGError( \ + ::mooncake::PGErrorCode::SystemError, \ + std::string(#expression) + \ + " failed: " + cudaGetErrorString(pg_cuda_error_internal)); \ + } \ + } while (false) + +#define PG_VALIDATE_ARG(condition, message) \ + do { \ + if (!(condition)) { \ + return ::mooncake::makePGError( \ + ::mooncake::PGErrorCode::InvalidArgument, (message)); \ + } \ + } while (false) + +#define PG_VALIDATE_STATE(condition, message) \ + do { \ + if (!(condition)) { \ + return ::mooncake::makePGError( \ + ::mooncake::PGErrorCode::InvalidState, (message)); \ + } \ + } while (false) + +#endif // MOONCAKE_PG_ERROR_TYPES_H diff --git a/mooncake-pg/include/gpu_runtime.h b/mooncake-pg/include/gpu_runtime.h new file mode 100644 index 0000000000..5e20ee2dcc --- /dev/null +++ b/mooncake-pg/include/gpu_runtime.h @@ -0,0 +1,81 @@ +#ifndef MOONCAKE_PG_GPU_RUNTIME_H +#define MOONCAKE_PG_GPU_RUNTIME_H + +#include + +namespace mooncake { + +class GpuDeviceGuard { + public: + explicit GpuDeviceGuard(int device); + ~GpuDeviceGuard() noexcept; + + GpuDeviceGuard(const GpuDeviceGuard&) = delete; + GpuDeviceGuard& operator=(const GpuDeviceGuard&) = delete; + + private: + int previous_device_ = -1; + bool restore_device_ = false; +}; + +class GpuEvent; + +class GpuStream { + public: + GpuStream() = delete; + ~GpuStream() noexcept; + + GpuStream(const GpuStream&) = delete; + GpuStream& operator=(const GpuStream&) = delete; + + GpuStream(GpuStream&& other) noexcept; + GpuStream& operator=(GpuStream&& other) noexcept; + + [[nodiscard]] static GpuStream createNonBlocking(int device); + [[nodiscard]] static GpuStream borrow(cudaStream_t stream, int device); + + [[nodiscard]] cudaStream_t get() const noexcept { return stream_; } + + [[nodiscard]] int deviceIndex() const noexcept { return device_index_; } + + [[nodiscard]] bool isCapturing() const; + + void waitEvent(const GpuEvent& event) const; + + private: + GpuStream(cudaStream_t stream, int device, bool owns_stream) noexcept; + + void reset() noexcept; + void moveFrom(GpuStream&& other) noexcept; + + cudaStream_t stream_ = nullptr; + int device_index_ = -1; + bool owns_stream_ = false; +}; + +class GpuEvent { + public: + explicit GpuEvent(int device, unsigned int flags = cudaEventDisableTiming); + ~GpuEvent() noexcept; + + GpuEvent(const GpuEvent&) = delete; + GpuEvent& operator=(const GpuEvent&) = delete; + + GpuEvent(GpuEvent&& other) noexcept; + GpuEvent& operator=(GpuEvent&& other) noexcept; + + void record(const GpuStream& stream); + + private: + friend class GpuStream; + + void reset() noexcept; + void moveFrom(GpuEvent&& other) noexcept; + + cudaEvent_t event_ = nullptr; + int device_index_ = -1; +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_GPU_RUNTIME_H diff --git a/mooncake-pg/include/mooncake_backend.h b/mooncake-pg/include/mooncake_backend.h deleted file mode 100644 index 81727175fa..0000000000 --- a/mooncake-pg/include/mooncake_backend.h +++ /dev/null @@ -1,268 +0,0 @@ -#ifndef MOONCAKE_BACKEND_H -#define MOONCAKE_BACKEND_H - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -namespace mooncake { - -// Forward declaration – MooncakeP2PShim holds a non-owning pointer to -// MooncakeBackend, which is defined below. -class MooncakeBackend; - -// Lightweight Backend shim that delegates P2P send/recv back to the owning -// MooncakeBackend. PyTorch's P2P dispatch (batch_isend_irecv, isend, irecv) -// requires getBackend() to return a registered c10d::Backend instance. -// Since MooncakeBackend inherits from ProcessGroup (not Backend), we register -// this shim in the ProcessGroup's deviceTypeToBackend_ map so that the P2P -// path can find it. The shim holds a non-owning pointer to its owner and -// delegates only the operations that the P2P dispatch path calls (send, recv, -// getBackendName, supportsCoalescing). -class MooncakeP2PShim final : public ::c10d::Backend { - public: - explicit MooncakeP2PShim(MooncakeBackend* owner); - - const std::string getBackendName() const override; - - bool supportsCoalescing() const override { return false; } - - c10::intrusive_ptr send(std::vector& tensors, - int dstRank, int tag) override; - - c10::intrusive_ptr recv(std::vector& tensors, - int srcRank, int tag) override; - - c10::intrusive_ptr recvAnysource( - std::vector& tensors, int tag) override; - - c10::intrusive_ptr barrier( - const c10d::BarrierOptions& opts) override; - - private: - // Non-owning: the shim is stored in ProcessGroup's backend maps which are - // cleared on destruction, and MooncakeBackend always outlives the shim. - MooncakeBackend* owner_; -}; - -class MooncakeBackend final : public ::c10d::ProcessGroup { - public: - struct MooncakeBackendOptions final : torch::CustomClassHolder { - explicit MooncakeBackendOptions(at::Tensor activeRanks) - : activeRanks_{activeRanks} {} - MooncakeBackendOptions(at::Tensor activeRanks, bool isExtension) - : activeRanks_{activeRanks}, isExtension_{isExtension} {} - MooncakeBackendOptions(at::Tensor activeRanks, bool isExtension, - int maxWorldSize) - : activeRanks_{activeRanks}, - isExtension_{isExtension}, - maxWorldSize_{maxWorldSize} {} - - ~MooncakeBackendOptions() override = default; - - at::Tensor activeRanks_; - bool isExtension_ = false; - // Optional upper bound for connection polling / reserved rank slots. - // When > 0, the backend may pre-size internal rank metadata to this - // value (while PyTorch's group_size() remains unchanged). - int maxWorldSize_ = -1; - }; - - /** - * @brief Construct a Mooncake process-group backend instance. - * - * `distBackendOpts` contains the PyTorch process-group information for this - * backend instance. `options` contains Mooncake-specific settings and may - * be null when callers omit `pg_options`. - * - * @param distBackendOpts Process-group information supplied by PyTorch. - * @param options *Optional* Mooncake-specific backend options. - * @param isCpu Whether to initialize the CPU backend variant. - */ - MooncakeBackend(c10d::DistributedBackendOptions distBackendOpts, - c10::intrusive_ptr options, - bool isCpu = false); - - ~MooncakeBackend() override; - - const std::string getBackendName() const override; - - int getSize() const override { return meta_ ? meta_->activeSize : size_; } - - // Point-to-point send/recv for torch.distributed P2POp/batch_isend_irecv. - // Only single-tensor ops are supported. - c10::intrusive_ptr send(std::vector& tensors, - int dstRank, int tag) override; - - c10::intrusive_ptr recv(std::vector& tensors, - int srcRank, int tag) override; - - c10::intrusive_ptr broadcast( - std::vector& tensors, - const c10d::BroadcastOptions& opts) override; - - c10::intrusive_ptr allreduce( - std::vector& tensors, - const c10d::AllreduceOptions& opts) override; - - c10::intrusive_ptr allgather( - std::vector>& outputTensors, - std::vector& inputTensors, - const c10d::AllgatherOptions& opts) override; - - c10::intrusive_ptr _allgather_base( - at::Tensor& outputBuffer, at::Tensor& inputBuffer, - const c10d::AllgatherOptions& opts) override; - - c10::intrusive_ptr _reduce_scatter_base( - at::Tensor& outputBuffer, at::Tensor& inputBuffer, - const c10d::ReduceScatterOptions& opts) override; - - c10::intrusive_ptr alltoall( - std::vector& outputTensors, - std::vector& inputTensors, - const c10d::AllToAllOptions& opts) override; - - c10::intrusive_ptr barrier( - const c10d::BarrierOptions& opts) override; - - c10::intrusive_ptr reduce( - std::vector& tensors, - const c10d::ReduceOptions& opts) override; - - c10::intrusive_ptr gather( - std::vector>& outputTensors, - std::vector& inputTensors, - const c10d::GatherOptions& opts) override; - - c10::intrusive_ptr scatter( - std::vector& outputTensors, - std::vector>& inputTensors, - const c10d::ScatterOptions& opts) override; - - void shutdown() override; - - static void setHostIp(const std::string& hostIp) { hostIp_ = hostIp; } - - static void setDeviceFilter(std::vector filters) { - engine_->setWhitelistFilters(std::move(filters)); - } - - /// Set an external TransferEngine to be used by MooncakeBackend - /// instead of creating its own. Must be called before - /// init_process_group(backend="mooncake"). The engine must already - /// be initialized. The caller is responsible for ensuring the engine - /// outlives all MooncakeBackend instances. Pass nullptr to reset. - static void setExternalEngine(TransferEngine* engine); - - std::string getPreferredHca(std::string location) { - static std::once_flag topo_once; - static std::shared_ptr topology; - static TopologyMatrix matrix; - std::call_once(topo_once, [this] { - // FIXME: getLocalTopology is deprecated in TENT - topology = engine_->getLocalTopology(); - if (topology) { - matrix = topology->getMatrix(); - } - if (!topology || matrix.empty()) { - topology = std::make_shared(); - topology->discover(); - matrix = topology->getMatrix(); - } - }); - - auto it = matrix.find(location); - if (it == matrix.end()) { - LOG(INFO) << "Topology is " << topology->toJson(); - LOG(ERROR) << "Topology entry not found for location: " << location; - return ""; - } - if (it->second.preferred_hca.empty()) { - LOG(INFO) << "Topology is " << topology->toJson(); - LOG(ERROR) << "Preferred HCA list is empty for location: " - << location; - return ""; - } - return it->second.preferred_hca[0]; - } - - at::Tensor getActiveRanksTensor() { return meta_->activeRanksTensor; } - - int getNumSyncedRanks(); - - void extendGroupSizeTo(int size); - - std::vector getPeerState(const std::vector& ranks); - - void recoverRanks(const std::vector& ranks); - - void joinGroup(); - - private: - void waitForExtensionState(); - void publishLocalPeerMetadata(); - void setLocalOnlyActiveRanks(); - void syncActiveRanksTensor(); - - static TransferEngine* engine_; - std::shared_ptr worker_; - static bool engineInitialized_; - static int backendIndex_; - // External engine injection: when set, MooncakeBackend uses this engine - // instead of the default self-created one. Non-owning pointer. - // The caller is responsible for ensuring the engine outlives all - // MooncakeBackend instances. - static TransferEngine* externalEngine_; - const c10::intrusive_ptr options_; - bool isCpu_{false}; - static std::string hostIp_; - void* send_buffer_[2]; - void* recv_buffer_[2]; - int32_t* cpu_sync_send_region_[2]; - int32_t* cpu_sync_recv_region_[2]; - SegmentInfo rank_info; - std::shared_ptr meta_; - bool isShutdown_{false}; - uint64_t local2global_rank_map_[kMaxNumRanks]; - std::string localServerName_; - - // P2P async infrastructure - // p2p_proxy_ is created in MooncakeBackend, but can live longer than - // MooncakeBackend. Because it is shared in P2PDeviceWorker, which must - // ensure P2PProxy's resources are not released until all transfers are - // completed. - std::shared_ptr p2p_proxy_; - // p2p_device_worker_ is created in P2PDeviceWorkerManager, - // and is shared between backends in the same device. - std::shared_ptr p2p_device_worker_; - - // Connection Poller Context - // Similar to p2p_proxy_, connection_ctx_ is created in MooncakeBackend, but - // can live longer than MooncakeBackend. - std::shared_ptr connection_ctx_; - bool connectionPollerRegistered_{false}; -}; - -struct ExtensionState { - std::vector activeRanks; - std::vector p2pEpochs; - int taskCount = -1; -}; -std::vector serialize(const ExtensionState& state); -ExtensionState deserialize(const std::vector& buffer); - -} // namespace mooncake - -#endif // MOONCAKE_BACKEND_H diff --git a/mooncake-pg/include/mooncake_communicator.h b/mooncake-pg/include/mooncake_communicator.h new file mode 100644 index 0000000000..8ec4545ff3 --- /dev/null +++ b/mooncake-pg/include/mooncake_communicator.h @@ -0,0 +1,315 @@ +#ifndef MOONCAKE_PG_COMMUNICATOR_H +#define MOONCAKE_PG_COMMUNICATOR_H + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "control_plane/agent_host.h" +#include "control_plane/coordinator_host.h" +#include "control_plane/link_manager.h" +#include "error_types.h" +#include "mooncake_pg.h" +#include "mooncake_worker.cuh" +#include "p2p_proxy.h" +#include "comm_types.h" + +namespace mooncake { + +static constexpr size_t kDefaultCollectiveTimeoutUs = 10000000; // 10 s +static constexpr int64_t kDefaultP2PTimeoutUs = 10000000; // 10 s + +// Must be greater than collective_timeout_us so that timeout-based +// failure reporters can contribute before the reconciliation window +// expires. (Some ranks report failures based on timeout, while others +// report based on failure status.) +static constexpr int64_t kDefaultFaultReconciliationWindowUs = + 3 * kDefaultCollectiveTimeoutUs; + +struct MooncakePGContext { + std::string host_ip = "127.0.0.1"; + size_t collective_timeout_us = kDefaultCollectiveTimeoutUs; + int64_t p2p_timeout_us = kDefaultP2PTimeoutUs; + int64_t fault_reconciliation_window_us = + kDefaultFaultReconciliationWindowUs; + + std::unique_ptr owned_engine = + std::make_unique(true); + TransferEngine* engine = owned_engine.get(); + bool engine_initialized = false; + int global_rank = -1; + int max_world_size = 0; + + LinkManager link_manager; + MooncakeWorkerManager worker_manager; + P2PDeviceWorkerManager p2p_device_worker_manager; + // Coordinator (rank 0 only). + // It must be started before the local AgentHost connects to it. + std::unique_ptr coordinator_host; + std::unique_ptr agent_host; + + MooncakePGContext() = default; + ~MooncakePGContext(); + + // Non-copyable: engine points to either owned_engine or an external engine + // whose lifetime is controlled by the caller. + MooncakePGContext(const MooncakePGContext&) = delete; + MooncakePGContext& operator=(const MooncakePGContext&) = delete; + + PGResult initialize(int rank, int world_size); + PGResult launchCoordinator(); + PGResult connectCoordinator(const std::string& coordinator_address); + PGResult setHostIp(std::string value); + PGResult setExternalEngine(TransferEngine* transfer_engine); + PGResult setDeviceFilter(std::vector filters); + PGResult setCollectiveTimeout(size_t timeout_us); + PGResult setP2PTimeout(int64_t timeout_us); + PGResult setFaultReconciliationWindow(int64_t timeout_us); + PGResult incrementCommUseCount(); + void decrementCommUseCount() noexcept; + PGResult shutdown(); + + private: + PGResult checkRunning() const; + + std::vector device_filters_; + std::mutex state_mutex_; + size_t comm_use_count_ = 0; + bool initialized_ = false; + bool shutdown_requested_ = false; +}; + +struct MooncakeCommunicatorConfig { + int rank = 0; + int size = 1; + int max_group_size = -1; + std::vector global_ranks; + GroupBootstrapId group_bootstrap_id; + bool is_cpu = false; + int device_index = -1; + GroupBootstrapIdResolvePolicy group_resolve_policy = + GroupBootstrapIdResolvePolicy::CreateOrAttach; + bool auto_deactivate_on_failure = true; + bool auto_sync_on_failure = true; + + // Optional caller-owned mirror of the communicator's active ranks. Its + // memory location is independent of the communicator's device. + int32_t* active_ranks_mirror = nullptr; + size_t active_ranks_mirror_count = 0; + bool active_ranks_mirror_is_device = false; + int active_ranks_mirror_device_index = -1; +}; + +class MooncakeCommunicator { + public: + static PGResult> create( + MooncakePGContext& context, MooncakeCommunicatorConfig config); + ~MooncakeCommunicator(); + + MooncakeCommunicator(const MooncakeCommunicator&) = delete; + MooncakeCommunicator& operator=(const MooncakeCommunicator&) = delete; + + int getRank() const { return rank_; } + int getSize() const; + int getMaxGroupSize() const { return max_group_size_; } + bool isCpu() const { return is_cpu_; } + + PGResult> sendCpu( + const void* buffer, size_t count, DataType datatype, int peer, + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count); + PGResult> sendGpu( + const void* buffer, size_t count, DataType datatype, int peer, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> recvCpu( + void* buffer, size_t count, DataType datatype, int peer, + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count); + PGResult> recvGpu( + void* buffer, size_t count, DataType datatype, int peer, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + + PGResult> broadcastCpu( + const void* send_buffer, void* recv_buffer, size_t count, + DataType datatype, int root, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult broadcastGpu(const void* send_buffer, void* recv_buffer, + size_t count, DataType datatype, int root, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> allReduceCpu( + const void* send_buffer, void* recv_buffer, size_t count, + DataType datatype, ReduceOp op, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult allReduceGpu(const void* send_buffer, void* recv_buffer, + size_t count, DataType datatype, ReduceOp op, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> allGatherCpu( + const void* send_buffer, void* recv_buffer, size_t count, + DataType datatype, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult allGatherGpu(const void* send_buffer, void* recv_buffer, + size_t count, DataType datatype, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> reduceScatterCpu( + const void* send_buffer, void* recv_buffer, size_t count, + DataType datatype, ReduceOp op, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult reduceScatterGpu(const void* send_buffer, void* recv_buffer, + size_t count, DataType datatype, + ReduceOp op, cudaStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> allToAllCpu( + const void* send_buffer, void* recv_buffer, size_t count, + DataType datatype, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult allToAllGpu(const void* send_buffer, void* recv_buffer, + size_t count, DataType datatype, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> barrierCpu( + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count); + PGResult barrierGpu(cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> reduceCpu( + const void* send_buffer, void* recv_buffer, size_t count, + DataType datatype, ReduceOp op, int root, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult reduceGpu(const void* send_buffer, void* recv_buffer, + size_t count, DataType datatype, ReduceOp op, + int root, cudaStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> gatherCpu( + const void* send_buffer, void* recv_buffer, size_t count, + DataType datatype, int root, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult gatherGpu(const void* send_buffer, void* recv_buffer, + size_t count, DataType datatype, int root, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> scatterCpu( + const void* send_buffer, void* recv_buffer, size_t count, + DataType datatype, int root, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult scatterGpu(const void* send_buffer, void* recv_buffer, + size_t count, DataType datatype, int root, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + + PGResult shutdown(); + std::vector getActiveRanks() const; + int getNumSyncedRanks() const; + PGResult> getPeerState( + const std::vector& ranks) const; + PGResult activateRanks( + const std::vector& ranks); + PGResult deactivateRanks( + const std::vector& ranks); + PGResult joinGroup(); + + // Returns the current GroupView epoch. + // Epoch starts at 0 (bootstrap) and increments on membership changes, + // auto-deactivation, and recovery. + uint64_t getCurrentEpoch() const; + + // Notify the Coordinator of a detected failure and block until a membership + // decision has been made and the Agent has ACKed the resulting ViewUpdate. + PGResult syncAfterFailure(); + + // Update the data-plane view. Called by AgentHost when a ViewUpdatePush is + // received or rank states change. rank_states and activatable are computed + // by the state machine. + void applyViewUpdate(const GroupView& view, + const std::vector& rank_states, + const std::vector& rank_epochs, + const std::vector& activatable); + // Called by AgentHost when a TE link to a peer comes back up. + void onPeerLinkReset(InGroupRank peer); + + // Called by NotifyLinkRefreshed effect: refresh the cached TE segment ID + // for `local` (InGroupRank) from the LinkManager. If the link is not up, + // segmentID is set to -1. + void refreshSegmentID(InGroupRank local); + GroupEndpointPublication buildEndpointMetadata() const; + + AgentInterface& getAgent() { return agent_; } + + private: + MooncakeCommunicator(MooncakePGContext& context, + const MooncakeCommunicatorConfig& config); + PGResult initialize(MooncakeCommunicatorConfig config); + + PGResult> enqueueSend( + const void* buffer, size_t count, DataType datatype, int peer, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + PGResult> enqueueRecv( + void* buffer, size_t count, DataType datatype, int peer, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count); + + // Guard: checks that the rank is not Offline (always) and, for collectives, + // that it is active in this group. Called at the top of every operation. + PGResult checkOpState(OpType op) const; + + // Validate and initialize the caller-owned failed-ranks output. + PGResult initializeFailedRanksHint( + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) const; + + // Reject operations if this communicator is invalid. + PGResult checkValidGroup(const char* operation) const; + + // A rejected registration has no Coordinator-assigned group id and is + // restricted to local-only collectives. + bool isValidGroup() const { return meta_ && !meta_->group_id.empty(); } + + // Sync the caller-provided host/device active-ranks mirror from the current + // GroupView. + void syncActiveRanksMirror() const; + + MooncakePGContext& context_; + AgentInterface& agent_; + int rank_ = 0; + int initial_size_ = 1; + int max_group_size_ = + 1; // per-group capacity (max active members for this group) + int device_index_ = -1; + bool is_cpu_ = false; + bool is_shutdown_ = false; + int32_t* active_ranks_mirror_ = nullptr; + bool active_ranks_mirror_is_device_ = false; + int active_ranks_mirror_device_index_ = -1; + std::optional active_ranks_mirror_stream_; + + std::shared_ptr worker_; + std::array send_buffer_{}; + std::array recv_buffer_{}; + std::array cpu_sync_send_region_{}; + std::array cpu_sync_recv_region_{}; + std::shared_ptr meta_; + + // P2P async infrastructure. p2p_proxy_ is created by this communicator but + // can live longer because P2PDeviceWorker retains it until all transfers + // complete. + std::shared_ptr p2p_proxy_; + + // Created by P2PDeviceWorkerManager and shared between communicators on the + // same device. + std::shared_ptr p2p_device_worker_; +}; + +} // namespace mooncake + +#endif // MOONCAKE_PG_COMMUNICATOR_H diff --git a/mooncake-pg/include/mooncake_pg.h b/mooncake-pg/include/mooncake_pg.h new file mode 100644 index 0000000000..597e60c93a --- /dev/null +++ b/mooncake-pg/include/mooncake_pg.h @@ -0,0 +1,329 @@ +#ifndef MOONCAKE_PG_H_ +#define MOONCAKE_PG_H_ + +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +#if defined(__GNUC__) +#define MOONCAKE_PG_EXPORT __attribute__((visibility("default"))) +#else +#define MOONCAKE_PG_EXPORT +#endif + +#define MOONCAKE_PG_MAX_RANKS 64 +#define MOONCAKE_PG_MAX_ERROR_STRING 256 +#define MOONCAKE_PG_CONFIG_UNDEF_INT INT_MIN + +#define MOONCAKE_PG_COMM_CONFIG_MAGIC 0x20B6DE6B +#define MOONCAKE_PG_COMM_CONFIG_VERSION 1u + +typedef struct mooncakePgContext* mooncakePgContext_t; +typedef struct mooncakePgComm* mooncakePgComm_t; +typedef struct mooncakePgCompletion* mooncakePgCompletion_t; +typedef void* mooncakePgStream_t; + +/* Keep non-success entries synchronized with C++ PGErrorCode. */ +typedef enum mooncakePgResult { + mooncakePgSuccess = 0, + mooncakePgInvalidArgument = 1, + mooncakePgInvalidState = 2, + mooncakePgNotSupported = 3, + mooncakePgTimeout = 4, + mooncakePgResourceBusy = 5, + mooncakePgTransferEngineError = 6, + mooncakePgRpcError = 7, + mooncakePgSystemError = 8, + mooncakePgInternalError = 9, +} mooncakePgResult_t; + +typedef enum mooncakePgDataType { + mooncakePgInt8 = 0, + mooncakePgUint8 = 1, + mooncakePgInt16 = 2, + mooncakePgUint16 = 3, + mooncakePgInt32 = 4, + mooncakePgUint32 = 5, + mooncakePgInt64 = 6, + mooncakePgUint64 = 7, + mooncakePgFloat16 = 8, + mooncakePgFloat32 = 9, + mooncakePgFloat64 = 10, + mooncakePgBfloat16 = 11, + mooncakePgBool = 12, + mooncakePgFloat8e4m3fn = 13, + mooncakePgFloat8e5m2 = 14, + mooncakePgFloat8e4m3fnuz = 15, + mooncakePgFloat8e5m2fnuz = 16, + mooncakePgFloat8e8m0fnu = 17, +} mooncakePgDataType_t; + +typedef enum mooncakePgReduceOp { + mooncakePgSum = 0, + mooncakePgAvg = 1, + mooncakePgProduct = 2, + mooncakePgMin = 3, + mooncakePgMax = 4, +} mooncakePgReduceOp_t; + +typedef enum mooncakePgDeviceType { + mooncakePgDeviceCpu = 0, + mooncakePgDeviceGpu = 1, +} mooncakePgDeviceType_t; + +typedef enum mooncakePgIdResolvePolicy { + mooncakePgIdResolveCreateOrAttach = 0, + mooncakePgIdResolveAttachOrExtend = 1, +} mooncakePgIdResolvePolicy_t; + +typedef struct mooncakePgCommConfig { + size_t structSize; + unsigned int magic; + unsigned int version; + const char* groupId; + int rank; + int size; + int maxGroupSize; + /* Required InGroupRank-to-GlobalRank mapping with exactly size entries. */ + const int32_t* globalRanks; + size_t globalRankCount; + mooncakePgDeviceType_t deviceType; + int deviceIndex; + mooncakePgIdResolvePolicy_t idResolvePolicy; + int autoDeactivateOnFailure; + int autoSyncOnFailure; + /* Optional caller-owned mirror of the communicator's active ranks. */ + int32_t* activeRanksMirror; + size_t activeRanksMirrorCount; + int activeRanksMirrorIsDevice; + int activeRanksMirrorDeviceIndex; +} mooncakePgCommConfig_t; + +#define MOONCAKE_PG_COMM_CONFIG_INITIALIZER \ + {sizeof(mooncakePgCommConfig_t), \ + MOONCAKE_PG_COMM_CONFIG_MAGIC, \ + MOONCAKE_PG_COMM_CONFIG_VERSION, \ + NULL, \ + MOONCAKE_PG_CONFIG_UNDEF_INT, \ + MOONCAKE_PG_CONFIG_UNDEF_INT, \ + MOONCAKE_PG_CONFIG_UNDEF_INT, \ + NULL, \ + 0, \ + mooncakePgDeviceGpu, \ + MOONCAKE_PG_CONFIG_UNDEF_INT, \ + mooncakePgIdResolveCreateOrAttach, \ + 1, \ + 1, \ + NULL, \ + 0, \ + 0, \ + MOONCAKE_PG_CONFIG_UNDEF_INT} + +typedef enum mooncakePgProposalStatus { + mooncakePgProposalRejected = 0, + mooncakePgProposalApplied = 1, + mooncakePgProposalAppliedWithDroppedRanks = 2, +} mooncakePgProposalStatus_t; + +typedef struct mooncakePgProposalResponse { + mooncakePgProposalStatus_t status; + uint64_t newEpoch; + size_t droppedRankCount; + int32_t droppedRanks[MOONCAKE_PG_MAX_RANKS]; + char rejectReason[MOONCAKE_PG_MAX_ERROR_STRING]; +} mooncakePgProposalResponse_t; + +typedef enum mooncakePgSyncAfterFailureStatus { + mooncakePgSyncReconciled = 0, + mooncakePgSyncNoPending = 1, + mooncakePgSyncRejected = 2, +} mooncakePgSyncAfterFailureStatus_t; + +typedef struct mooncakePgSyncAfterFailureResponse { + mooncakePgSyncAfterFailureStatus_t status; + char rejectReason[MOONCAKE_PG_MAX_ERROR_STRING]; +} mooncakePgSyncAfterFailureResponse_t; + +MOONCAKE_PG_EXPORT const char* mooncakePgGetErrorString( + mooncakePgResult_t result); +MOONCAKE_PG_EXPORT const char* mooncakePgGetLastError(void); + +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgContextCreate(mooncakePgContext_t* context); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgContextInitialize( + mooncakePgContext_t context, int globalRank, int maxWorldSize); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgContextLaunchCoordinator( + mooncakePgContext_t context, char* coordinatorAddressBuf, + size_t coordinatorAddressBufSize); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgContextConnectCoordinator( + mooncakePgContext_t context, const char* coordinatorAddress); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgContextSetHostIp(mooncakePgContext_t context, const char* hostIp); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgContextSetTransferEngine( + mooncakePgContext_t context, void* transferEngine); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgContextSetDeviceFilter( + mooncakePgContext_t context, const char* const* filters, + size_t filterCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgContextSetCollectiveTimeout( + mooncakePgContext_t context, size_t timeoutUs); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgContextSetP2PTimeout(mooncakePgContext_t context, int64_t timeoutUs); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgContextSetFaultReconciliationWindow(mooncakePgContext_t context, + int64_t timeoutUs); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgContextDestroy(mooncakePgContext_t context); + +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgCommCreate( + mooncakePgContext_t context, const mooncakePgCommConfig_t* config, + mooncakePgComm_t* comm); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgCommDestroy(mooncakePgComm_t comm); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgCommGetRank(mooncakePgComm_t comm, int* rank); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgCommGetSize(mooncakePgComm_t comm, int* size); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgCommGetMaxGroupSize(mooncakePgComm_t comm, int* maxGroupSize); + +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgBroadcastGpu(const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, int root, + mooncakePgComm_t comm, mooncakePgStream_t stream, + int32_t* failedRanksHint, size_t failedRanksHintCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgAllReduceGpu( + const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgReduceOp_t reduceOp, + mooncakePgComm_t comm, mooncakePgStream_t stream, int32_t* failedRanksHint, + size_t failedRanksHintCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgAllGatherGpu(const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgComm_t comm, + mooncakePgStream_t stream, int32_t* failedRanksHint, + size_t failedRanksHintCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgReduceScatterGpu( + const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgReduceOp_t reduceOp, + mooncakePgComm_t comm, mooncakePgStream_t stream, int32_t* failedRanksHint, + size_t failedRanksHintCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgAllToAllGpu(const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgComm_t comm, + mooncakePgStream_t stream, int32_t* failedRanksHint, + size_t failedRanksHintCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgReduceGpu( + const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgReduceOp_t reduceOp, int root, + mooncakePgComm_t comm, mooncakePgStream_t stream, int32_t* failedRanksHint, + size_t failedRanksHintCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgGatherGpu(const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, int root, + mooncakePgComm_t comm, mooncakePgStream_t stream, + int32_t* failedRanksHint, size_t failedRanksHintCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgScatterGpu(const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, int root, + mooncakePgComm_t comm, mooncakePgStream_t stream, + int32_t* failedRanksHint, size_t failedRanksHintCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgBarrierGpu(mooncakePgComm_t comm, mooncakePgStream_t stream, + int32_t* failedRanksHint, size_t failedRanksHintCount); + +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgBroadcastCpu( + const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, int root, mooncakePgComm_t comm, + int32_t* failedRanksHint, size_t failedRanksHintCount, + mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgAllReduceCpu( + const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgReduceOp_t reduceOp, + mooncakePgComm_t comm, int32_t* failedRanksHint, + size_t failedRanksHintCount, mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgAllGatherCpu(const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgComm_t comm, + int32_t* failedRanksHint, size_t failedRanksHintCount, + mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgReduceScatterCpu( + const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgReduceOp_t reduceOp, + mooncakePgComm_t comm, int32_t* failedRanksHint, + size_t failedRanksHintCount, mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgAllToAllCpu(const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgComm_t comm, + int32_t* failedRanksHint, size_t failedRanksHintCount, + mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgReduceCpu( + const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, mooncakePgReduceOp_t reduceOp, int root, + mooncakePgComm_t comm, int32_t* failedRanksHint, + size_t failedRanksHintCount, mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgGatherCpu( + const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, int root, mooncakePgComm_t comm, + int32_t* failedRanksHint, size_t failedRanksHintCount, + mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgScatterCpu( + const void* sendBuffer, void* recvBuffer, size_t count, + mooncakePgDataType_t dataType, int root, mooncakePgComm_t comm, + int32_t* failedRanksHint, size_t failedRanksHintCount, + mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgBarrierCpu( + mooncakePgComm_t comm, int32_t* failedRanksHint, + size_t failedRanksHintCount, mooncakePgCompletion_t* completion); + +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgSendGpu( + const void* sendBuffer, size_t count, mooncakePgDataType_t dataType, + int peer, mooncakePgComm_t comm, mooncakePgStream_t stream, + int32_t* failedRanksHint, size_t failedRanksHintCount, + mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgRecvGpu( + void* recvBuffer, size_t count, mooncakePgDataType_t dataType, int peer, + mooncakePgComm_t comm, mooncakePgStream_t stream, int32_t* failedRanksHint, + size_t failedRanksHintCount, mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgSendCpu( + const void* sendBuffer, size_t count, mooncakePgDataType_t dataType, + int peer, mooncakePgComm_t comm, int32_t* failedRanksHint, + size_t failedRanksHintCount, mooncakePgCompletion_t* completion); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgRecvCpu( + void* recvBuffer, size_t count, mooncakePgDataType_t dataType, int peer, + mooncakePgComm_t comm, int32_t* failedRanksHint, + size_t failedRanksHintCount, mooncakePgCompletion_t* completion); + +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgCompletionIsCompleted( + mooncakePgCompletion_t completion, int* completed); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgCompletionWait(mooncakePgCompletion_t completion, int64_t timeoutUs); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgCompletionDestroy(mooncakePgCompletion_t completion); + +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgCommGetActiveRanks( + mooncakePgComm_t comm, int32_t* activeRanks, size_t rankCount); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgCommGetPeerState(mooncakePgComm_t comm, const int32_t* ranks, + size_t rankCount, int32_t* peerStates); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgCommActivateRanks( + mooncakePgComm_t comm, const int32_t* ranks, size_t rankCount, + mooncakePgProposalResponse_t* response); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgCommDeactivateRanks( + mooncakePgComm_t comm, const int32_t* ranks, size_t rankCount, + mooncakePgProposalResponse_t* response); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgCommJoin(mooncakePgComm_t comm); +MOONCAKE_PG_EXPORT mooncakePgResult_t mooncakePgCommSyncAfterFailure( + mooncakePgComm_t comm, mooncakePgSyncAfterFailureResponse_t* response); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgCommGetEpoch(mooncakePgComm_t comm, uint64_t* epoch); +MOONCAKE_PG_EXPORT mooncakePgResult_t +mooncakePgCommGetNumSyncedRanks(mooncakePgComm_t comm, int* numSyncedRanks); + +#ifdef __cplusplus +} +#endif + +#endif // MOONCAKE_PG_H_ diff --git a/mooncake-pg/include/mooncake_worker.cuh b/mooncake-pg/include/mooncake_worker.cuh index d8a6c1a50c..6efdc58389 100644 --- a/mooncake-pg/include/mooncake_worker.cuh +++ b/mooncake-pg/include/mooncake_worker.cuh @@ -1,25 +1,20 @@ #ifndef MOONCAKE_WORKER_CUH #define MOONCAKE_WORKER_CUH -#if !defined(__MUSA__) -#include -#include -#include -#include -#include -#include -#else -// MUSA device compilation: minimal includes to avoid mcc compiler crash -#include -#include -#endif - -#include +#include +#include + +#include "control_plane/control_types.h" +#include "gpu_runtime.h" + #include +#include +#include "comm_types.h" #include -#include #include +#include +#include #include #include #include @@ -27,122 +22,118 @@ namespace mooncake { static constexpr size_t kBufferSize = 1u << 24; -static constexpr size_t kMaxNumRanks = 64; -struct SegmentInfo { - uint64_t send_buffer[2], recv_buffer[2], send_sync[2], recv_sync[2], - warmup_buffer[2]; - uint64_t p2p_credit_region; - uint64_t p2p_ack_region; +class MooncakeCommunicator; + +// Local collective extension state. Every communicator starts in Isolated. +// +// Founding member: +// Isolated --(Active view)--> Normal +// +// Joining member: +// Isolated --(joinGroup: drain preparation collectives)--> Quiescing +// -----------------(Active view)----------------> Normal +// +// Isolated admits local-only collectives with an {self} active ranks mask. +// Quiescing rejects new collectives while waiting for activation. Normal uses +// the Coordinator's committed membership as active ranks. These are local +// extension phases, not membership states: an auto-deactivated communicator +// remains Normal and fails its next collective through the inactive self bit. +enum class CollectiveExtensionState : uint8_t { + Isolated = 0, // Local-only collectives + Quiescing = 1, // awaiting activation; no collectives may be issued. + Normal = 2, // Collectives use the membership committed by the coordinator. }; struct TransferGroupMeta { - int rank; - int size; // capacity: number of slots allocated (incl. inactive) - int activeSize; // visible group size: number of ranks that participate + InGroupRank rank; + GlobalRank globalRank; + // rank_order maps InGroupRank (0 .. maxGroupSize-1) to GlobalRank. + GlobalRank rank_order[kMaxNumRanks]; + + int maxGroupSize; + // Highest active InGroupRank plus one. + std::atomic activeSize{0}; int taskCount; + + GroupId group_id; + std::atomic epoch{0}; + std::atomic extensionMode{ + CollectiveExtensionState::Isolated}; + bool* activeRanks; bool* activeRanksDevice; -#if !defined(__MUSA__) - at::Tensor activeRanksTensor; -#endif - bool peerConnected[kMaxNumRanks]{}; + bool* maybeActivatable; + RankState rankStates[kMaxNumRanks]; // per GlobalRank + uint64_t rankEpochs[kMaxNumRanks]; TransferEngine* engine; -#if !defined(__MUSA__) - c10::intrusive_ptr<::c10d::Store> store; -#endif - int bufferBaseIndex; - int backendIndex; TransferMetadata::SegmentID segmentIDs[kMaxNumRanks]; - SegmentInfo segmentInfos[kMaxNumRanks]; + GroupEndpointInfo segmentInfos[kMaxNumRanks]; + const size_t* collectiveTimeoutUs = nullptr; + MooncakeCommunicator* communicator = nullptr; + bool autoSyncOnFailure = true; }; -#if defined(__CUDACC__) || defined(__MUSA__) -__global__ -#endif - struct Task { - volatile bool active = false; - int opType = - 0; // c10d::OpType as int, for ABI compatibility with kernel code - size_t tensorSize; // In bytes - int64_t broadcastRoot; - int bufferOffset; - uint64_t submitSequence = 0; - BatchID batchID; - void* transferGroupMeta; -}; - -#if !defined(__MUSA__) -void launchReduceKernel(at::Tensor dst, size_t pos, size_t realSize, void* src, - size_t numRanks, c10d::ReduceOp op, bool* activeRanks, - cudaStream_t stream); - -void launchReduceCpu(at::Tensor dst, size_t pos, size_t realSize, void* src, - size_t numRanks, c10d::ReduceOp op, bool* activeRanks); -void preloadReduceKernels(); +void launchReduceKernel(void* dst, DataType datatype, size_t pos, + size_t realSize, void* src, size_t numRanks, + ReduceOp op, bool* activeRanks, cudaStream_t stream); -class ConnectionContext; - -struct CudaTaskSubmissionToken { - size_t task_id; - uint64_t sequence; -}; +void launchReduceCpu(void* dst, DataType datatype, size_t pos, size_t realSize, + void* src, size_t numRanks, ReduceOp op, + bool* activeRanks); class MooncakeWorker { public: explicit MooncakeWorker(int cuda_device_index = -1); ~MooncakeWorker(); - c10::intrusive_ptr putTaskCpu( - c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot, + std::unique_ptr putTaskCpu( + OpType opType, size_t dataSize, int64_t broadcastRoot, const std::shared_ptr& meta, - const std::shared_ptr& connection_ctx, + int32_t* failedRanksHint, const std::function& - tensorToBuffer, + copyToSendBuffer, const std::function& - bufferToTensor); + copyFromRecvBuffer); - c10::intrusive_ptr putTaskCuda( - c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot, + void putTaskCuda( + OpType opType, size_t dataSize, int64_t broadcastRoot, const std::shared_ptr& meta, - const std::shared_ptr& connection_ctx, - const at::cuda::CUDAStream& issue_stream, + cudaStream_t issueStream, int32_t* failedRanksHint, const std::function& tensorToBuffer, + cudaStream_t)>& copyToSendBuffer, const std::function& bufferToTensor); + cudaStream_t)>& copyFromRecvBuffer); void Start(); /** - * @brief Waits for all active collective tasks for the given backend to - * complete. + * @brief Waits for all active collective tasks for the given communicator + * to complete. * * Used during graceful shutdown to ensure no pending collective operations * are active before releasing resources. Blocks until all tasks complete * or the timeout expires. * - * @param meta The transfer group metadata identifying the backend. + * @param meta The transfer group metadata identifying the communicator. * @return True if all tasks completed within the timeout; false if timed * out. */ bool drainTasks(const TransferGroupMeta* meta) const; - bool waitUntilTasksSubmitted( - const std::vector& tasks, - std::chrono::milliseconds timeout) const; - private: void startWorker(); + void waitUntilTasksSubmitted( + const std::vector& tasks) const; static constexpr size_t kNumTasks_ = 4; - static constexpr size_t kPingTimeoutMicroseconds_ = 100; static constexpr size_t kDrainTasksTimeoutMs = 5000; // 5s std::atomic running_{false}; std::atomic started_{false}; int cuda_device_index_; + std::optional enqueue_stream_; Task *tasks_, *tasks_device_; bool hasCallback_[kNumTasks_]{}; @@ -158,11 +149,7 @@ class MooncakeWorker { class MooncakeWorkerManager { public: - static MooncakeWorkerManager& GetInstance() { - // leaky singleton to avoid destructor fiasco problem - static MooncakeWorkerManager* manager = new MooncakeWorkerManager; - return *manager; - } + MooncakeWorkerManager() = default; std::shared_ptr GetCPUWorker(); std::shared_ptr GetCUDAWorker(int cuda_device_index); @@ -175,7 +162,6 @@ class MooncakeWorkerManager { // detached threads must not outlive the MooncakeWorker object. std::unordered_map> workers_; }; -#endif // !defined(__MUSA__) } // namespace mooncake diff --git a/mooncake-pg/include/mooncake_worker_kernels.cuh b/mooncake-pg/include/mooncake_worker_kernels.cuh index 67b56b9e2c..35b6613f4e 100644 --- a/mooncake-pg/include/mooncake_worker_kernels.cuh +++ b/mooncake-pg/include/mooncake_worker_kernels.cuh @@ -1,24 +1,38 @@ #ifndef MOONCAKE_WORKER_KERNELS_CUH #define MOONCAKE_WORKER_KERNELS_CUH -// Include the main worker header for struct definitions (Task, SegmentInfo, -// TransferGroupMeta). When compiled by mcc (__MUSA__ defined), the torch- -// dependent parts are guarded out, making this safe for the MUSA compiler. -#include +#include +#include +#include namespace mooncake { +struct TransferGroupMeta; + +#if defined(__CUDACC__) || defined(__MUSA__) +__global__ +#endif + struct Task { + volatile bool active = false; + int opType = + 0; // c10d::OpType as int, for ABI compatibility with kernel code + size_t dataSize; // In bytes + int64_t broadcastRoot; + int bufferOffset; + uint64_t submitSequence = 0; + int32_t* failedRanksHint = nullptr; + bool resetFailedRanksHint = false; + void* transferGroupMeta; +}; + // Kernel function declarations — guarded so g++ doesn't see __global__ -// which it can't parse. Parameters use plain C++ types (int instead of -// c10d::OpType / c10d::ReduceOp::RedOpType) so that mcc can compile them -// without torch headers. #if defined(__CUDACC__) || defined(__MUSA__) -__global__ void enqueueTaskKernel(int opType, size_t tensorSize, +__global__ void enqueueTaskKernel(int opType, size_t dataSize, int64_t broadcastRoot, int bufferOffset, - uint64_t submitSequence, void* meta, - Task* tasks, int numRanks, - const bool* activeRanks, - int* activeRanksTensor, size_t taskId); + uint64_t submitSequence, + int32_t* failedRanksHint, + bool resetFailedRanksHint, void* meta, + Task* tasks, size_t taskId); template __global__ void reduceKernel(scalar_t* dst, const scalar_t* src, @@ -32,12 +46,11 @@ __global__ void reduceKernel(scalar_t* dst, const scalar_t* src, // Both CUDA and MUSA use cudaStream_t in the declaration: on MUSA, // cuda_alike.h typedefs cudaStream_t to musaStream_t. -void launchEnqueueTaskKernel(int opType, size_t tensorSize, - int64_t broadcastRoot, int bufferOffset, - uint64_t submitSequence, void* meta, Task* tasks, - int numRanks, const bool* activeRanks, - int* activeRanksTensor, size_t taskId, - cudaStream_t stream); +void launchEnqueueTaskKernel(int opType, size_t dataSize, int64_t broadcastRoot, + int bufferOffset, uint64_t submitSequence, + int32_t* failedRanksHint, + bool resetFailedRanksHint, void* meta, Task* tasks, + size_t taskId, cudaStream_t stream); void launchReduceKernel_uint8(uint8_t* dst, const uint8_t* src, size_t numElements, size_t numRanks, int op, diff --git a/mooncake-pg/include/p2p_proxy.h b/mooncake-pg/include/p2p_proxy.h index 69f508a725..6b17d126cf 100644 --- a/mooncake-pg/include/p2p_proxy.h +++ b/mooncake-pg/include/p2p_proxy.h @@ -2,13 +2,13 @@ #define MOONCAKE_P2P_PROXY_H #include -#include #include #include #include #include #include #include +#include #include #include #include @@ -63,7 +63,7 @@ namespace mooncake { // sender to write data into the designated RecvPool offset. // - The SENDER is passive. It polls its local CreditLane (written by the // receiver via RDMA). Only after receiving a credit does it allocate -// a staging buffer from SendPool, copy user tensor data into it, and +// a staging buffer from SendPool, copy user buffer data into it, and // perform the RDMA Write to the remote RecvPool. // - After the RDMA Write finishes, the sender writes an AckSlot back // to the receiver's AckLane to acknowledge the transfer. @@ -84,7 +84,7 @@ namespace mooncake { // "all chunks are reserved for receiving and none are left for sending" // can never happen. // -// Data flow example -- Rank 0 sends a 24 MiB tensor to Rank 1. +// Data flow example -- Rank 0 sends a 24 MiB buffer to Rank 1. // The five steps below are separated between Sender (Rank 0) and // Receiver (Rank 1) to make the protocol explicit. // @@ -98,11 +98,11 @@ namespace mooncake { // | chunk 1 : free | -----------> | slot 2 : {seq=2,off=1,len=8M} | // +-----------------------+ +-------------------------------+ // -// Step 2 -- Rank 0 stages user tensor into SendPool +// Step 2 -- Rank 0 stages the user buffer into SendPool // ------------------------------------------------------------------------ // Rank 0 (Sender) // +-----------------------+ -// | User tensor | +// | User buffer | // |[0..8M)[8..16M)[16..24)| // +-----------------------+ // | @@ -144,10 +144,10 @@ namespace mooncake { // | chunk 1 : [16..24M) | // +-----------------------+ // | -// | cudaMemcpyAsync (RecvPool -> user tensor) +// | cudaMemcpyAsync (RecvPool -> user buffer) // v // +-----------------------+ -// | User tensor (ready) | +// | User buffer (ready) | // |[0..8M)[8..16M)[16..24)| // +-----------------------+ // | @@ -274,29 +274,32 @@ class P2PProxy { public: friend class P2PDeviceWorker; - enum class OpStatus : uint8_t { kPending = 0, kSuccess = 1, kFailed = 2 }; + enum class IssueResult : uint8_t { kIssued, kNoCredit, kTimeout }; struct Options { bool is_cpu = false; int rank = 0; int size = 0; int cuda_device_index = -1; - std::chrono::milliseconds transfer_timeout_ms{30000}; + const int64_t* p2p_timeout_us = nullptr; }; struct SendOp { - at::Tensor tensor_; + const void* buffer_ = nullptr; + size_t size_ = 0; int peer_rank_ = -1; cudaStream_t cuda_stream_ = nullptr; - std::shared_ptr> status_; + std::shared_ptr> completion_; + int32_t* failed_ranks_hint_ = nullptr; }; struct RecvOp { - at::Tensor tensor_; - at::Tensor original_tensor_; + void* buffer_ = nullptr; + size_t size_ = 0; int peer_rank_ = -1; cudaStream_t cuda_stream_ = nullptr; - std::shared_ptr> status_; + std::shared_ptr> completion_; + int32_t* failed_ranks_hint_ = nullptr; }; P2PProxy(TransferEngine* engine, const Options& options); @@ -344,7 +347,7 @@ class P2PProxy { private: // Sender-side per-chunk state machine. enum class SendTaskState { - kCopyIn, // Copying tensor slice into local SendPool staging buffer. + kCopyIn, // Copying buffer slice into local SendPool staging buffer. // On CPU this is synchronous; on GPU we record a // cudaEvent on the op's stream and poll it. kWriteRemote, // Write from SendPool -> remote RecvPool. @@ -357,7 +360,7 @@ class P2PProxy { enum class RecvTaskState { kIssueCredit, // CreditSlot Write is in flight. kWaitAck, // Waiting for sender's AckSlot. - kCopyOut, // GPU: cudaMemcpyAsync RecvPool -> tensor in flight. + kCopyOut, // GPU: cudaMemcpyAsync RecvPool -> buffer in flight. kFinished, // Copy-out done, chunk returned to pool. kFailed, // Transport error or timeout; chunk returned to pool. }; @@ -365,17 +368,17 @@ class P2PProxy { struct SendOpContext; struct RecvOpContext; - // One chunk of a SendOp. The sender copies data from the user tensor + // One chunk of a SendOp. The sender copies data from the user buffer // into a SendPool staging buffer and then RDMA-writes it to the remote // RecvPool offset that the receiver issued in the CreditSlot. struct SendTransferTask { SendTransferTask() = default; - SendTransferTask(uint64_t tensor_offset_in, uint32_t chunk_len_in, + SendTransferTask(uint64_t buffer_offset_in, uint32_t chunk_len_in, void* staging_addr_in, uint64_t remote_addr_in, uint32_t sequence_in, uint32_t epoch_in); SendTaskState state_ = SendTaskState::kCopyIn; - uint64_t tensor_offset_ = 0; // Offset inside the user tensor. + uint64_t buffer_offset_ = 0; // Offset inside the user buffer. uint32_t chunk_len_ = 0; // Bytes in this chunk (<= kP2PChunkSize). void* staging_addr_ = nullptr; // Address inside SendPool. uint64_t remote_addr_ = 0; // Address inside REMOTE RecvPool. @@ -395,16 +398,17 @@ class P2PProxy { SendOpContext(SendOp&& op_in); std::deque tasks_; - std::shared_ptr> status_; + std::shared_ptr> completion_; - at::Tensor tensor_; + const void* buffer_ = nullptr; int peer_rank_ = -1; cudaStream_t cuda_stream_ = nullptr; uint64_t total_bytes_ = 0; - // Number of bytes already pulled from the user tensor into SendPool + // Number of bytes already pulled from the user buffer into SendPool // staging buffers. When bytes_staged_ == total_bytes_ every chunk // has at least entered the Copy-In stage. uint64_t bytes_staged_ = 0; + int32_t* failed_ranks_hint_ = nullptr; std::chrono::steady_clock::time_point last_update_time_; }; @@ -412,15 +416,15 @@ class P2PProxy { // One chunk of a RecvOp. The receiver allocates a RecvPool chunk, // issues it to the sender via a CreditSlot, waits for the sender // to RDMA-write data into it, and finally copies the data into the user - // tensor before returning the chunk to RecvPool. + // buffer before returning the chunk to RecvPool. struct RecvTransferTask { RecvTransferTask() = default; - RecvTransferTask(uint64_t tensor_offset_in, uint32_t chunk_len_in, + RecvTransferTask(uint64_t buffer_offset_in, uint32_t chunk_len_in, void* local_addr_in, uint32_t sequence_in, uint32_t epoch_in); RecvTaskState state_ = RecvTaskState::kIssueCredit; - uint64_t tensor_offset_ = 0; // Offset inside the user tensor. + uint64_t buffer_offset_ = 0; // Offset inside the user buffer. uint32_t chunk_len_ = 0; // Bytes in this chunk. void* local_addr_ = nullptr; // Address inside local RecvPool. uint32_t sequence_ = 0; // Sequence number in the control ring. @@ -440,16 +444,16 @@ class P2PProxy { RecvOpContext(RecvOp&& op_in); std::deque tasks_; - std::shared_ptr> status_; + std::shared_ptr> completion_; - at::Tensor tensor_; - at::Tensor original_tensor_; + void* buffer_ = nullptr; int peer_rank_ = -1; cudaStream_t cuda_stream_ = nullptr; uint64_t total_bytes_ = 0; + int32_t* failed_ranks_hint_ = nullptr; // Number of bytes for which a RecvPool chunk has been reserved and a // CreditSlot has been sent to the peer. When bytes_credited_ == - // total_bytes_ the entire tensor has been offered to the sender. + // total_bytes_ the entire buffer has been offered to the sender. uint64_t bytes_credited_ = 0; }; @@ -468,7 +472,7 @@ class P2PProxy { // to the peer to write data, and consumes AckSlots that the peer writes // back to acknowledge finished transfers. struct RecvPeerLane { - std::deque pending_recv_ops_; + std::deque pending_recv_ops_; std::optional active_recv_op_; // Sequence number of the next CreditSlot to issue to this peer. uint64_t credit_issue_seq_ = 0; @@ -498,7 +502,7 @@ class P2PProxy { bool isRecvOpCompleted(const RecvOpContext& op_ctx) const; void performRecvReset(int peer_rank); - bool tryIssueSendTask(SendOpContext& op_ctx, SendPeerLane& lane); + IssueResult tryIssueSendTask(SendOpContext& op_ctx, SendPeerLane& lane); bool stepSendTask(SendOpContext& op_ctx, SendTransferTask& task); bool stepSendCopyIn(SendTransferTask& task); bool stepSendWriteRemote(SendOpContext& op_ctx, SendTransferTask& task); @@ -506,8 +510,6 @@ class P2PProxy { bool isSendOpCompleted(const SendOpContext& op_ctx) const; void performSendReset(int peer_rank); - void reportBrokenPeer(int peer_rank); - // These helpers are used only during reset or shutdown. // In the normal execution path, task and lane resources are released // incrementally as work progresses. @@ -515,7 +517,16 @@ class P2PProxy { void releaseRecvTaskResources(RecvTransferTask& task) const; void resetSendLane(SendPeerLane& lane); void resetRecvLane(RecvPeerLane& lane); - void resetPeerControlLanes(int peer_rank); + void resetPeerCreditLanes(int peer_rank); + void resetPeerAckLanes(int peer_rank); + // Clean up the active op on a lane + void cleanupFailedSendOp(SendOpContext& op_ctx); + void cleanupFailedRecvOp(RecvOpContext& op_ctx); + // Reset P2P session and push link event to Agent. + void reportPeerFailure(int peer_rank); + // Clean up, mark kFailed, and report the failure. + void handleFailedSendOp(SendOpContext& op_ctx); + void handleFailedRecvOp(RecvOpContext& op_ctx); // Control lane addressing. // @@ -539,7 +550,7 @@ class P2PProxy { template bool isTimeout(const T& obj) const { return std::chrono::steady_clock::now() - obj.last_update_time_ > - transfer_timeout_ms_; + std::chrono::microseconds(*p2p_timeout_us_); } private: @@ -559,7 +570,7 @@ class P2PProxy { int rank_ = 0; int size_ = 0; int cuda_device_index_ = -1; - std::chrono::milliseconds transfer_timeout_ms_{5000}; // 5s + const int64_t* p2p_timeout_us_ = nullptr; P2PResources resources_; bool resource_abandoned_{false}; @@ -571,7 +582,7 @@ class P2PProxy { std::queue send_queue_; std::mutex send_queue_mutex_; - std::queue recv_queue_; + std::queue recv_queue_; std::mutex recv_queue_mutex_; std::array, kMaxNumRanks> reset_send_req_; @@ -589,9 +600,9 @@ class P2PProxy { std::array recv_peer_lanes_; }; -// P2PDeviceWorker instances are shared across multiple backends within the same -// process. Therefore, they must not be instantiated directly. Instead, obtain -// an instance through P2PDeviceWorkerManager. +// P2PDeviceWorker instances are shared across multiple communicators within the +// same process. Therefore, they must not be instantiated directly. Instead, +// obtain an instance through P2PDeviceWorkerManager. class P2PDeviceWorker { public: friend class P2PDeviceWorkerManager; @@ -636,7 +647,8 @@ class P2PDeviceWorker { // Memory footprint (with default env values): // - One direction (send or recv) : 128 MiB // - Total per device (send + recv) : 256 MiB - // This stays constant no matter how many backends or peer ranks are active. + // This stays constant no matter how many communicators or peer ranks are + // active. TransferEngine* engine_ = nullptr; P2PChunkPool send_pool_; P2PChunkPool recv_pool_; @@ -651,11 +663,7 @@ class P2PDeviceWorker { class P2PDeviceWorkerManager { public: - static P2PDeviceWorkerManager& getInstance() { - // leaky singleton to avoid destructor fiasco problem - static P2PDeviceWorkerManager* manager = new P2PDeviceWorkerManager; - return *manager; - } + P2PDeviceWorkerManager() = default; std::shared_ptr getCPUWorker(TransferEngine* engine); std::shared_ptr getCUDAWorker(int cuda_device_index, diff --git a/mooncake-pg/include/pg_utils.h b/mooncake-pg/include/pg_utils.h index 01304fb66e..3c037ab5db 100644 --- a/mooncake-pg/include/pg_utils.h +++ b/mooncake-pg/include/pg_utils.h @@ -3,10 +3,10 @@ #pragma once -#include -#include #include +#include #include +#include // For PAUSE macro #include @@ -181,6 +181,15 @@ class BackoffWaiter { uint32_t yield_count_{0}; std::chrono::microseconds current_sleep_; }; + +template +struct overloaded : Ts... { + using Ts::operator()...; +}; + +template +overloaded(Ts...) -> overloaded; + } // namespace mooncake #endif diff --git a/mooncake-pg/setup.py b/mooncake-pg/setup.py deleted file mode 100644 index e839339b5a..0000000000 --- a/mooncake-pg/setup.py +++ /dev/null @@ -1,104 +0,0 @@ -import os -import re - -from setuptools import setup -import torch - -use_musa = os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in {"1", "ON", "TRUE", "YES"} -if use_musa: - try: - import torchada # noqa: F401 - except ImportError as e: - raise ImportError( - "torchada is required to build the MUSA PG extension. " - "Please install it first using 'pip install torchada'." - ) from e - - -from torch.utils.cpp_extension import ( # noqa: E402 - BuildExtension, - CUDAExtension, - CUDA_HOME, -) - - -torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group() -version_suffix = "_" + torch_version.replace(".", "_") -module_name = "mooncake.pg" + version_suffix - -abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI) -current_dir = os.path.abspath(os.path.dirname(__file__)) - -abi_define = f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}" -cxx_args = [abi_define, "-std=c++20", "-O3", "-g0"] - -cuda_libraries = ["ibverbs", "mlx5"] -cuda_library_dirs = [] -use_maca = hasattr(torch.version, "maca") and torch.version.maca is not None - -if use_musa: - musa_defines = ["-DUSE_MUSA", "-DMOONCAKE_EP_USE_MUSA=1"] - cxx_args += musa_defines - # torchada maps the "nvcc" key to "mcc". - device_args = [ - abi_define, - *musa_defines, - "-std=c++20", - "--cuda-gpu-arch=mp_21", - "--cuda-gpu-arch=mp_31", - "-O3", - ] -else: - if use_maca: - cxx_args.append("-DUSE_MACA") - device_args = [ - abi_define, - "-std=c++20", - "-Xcompiler", - "-O3", - "-Xcompiler", - "-g0", - ] - if use_maca: - device_args.append("-DUSE_MACA") - # Link against the CUDA driver stub library if available. - # Same approach as mooncake-ep/setup.py. - if CUDA_HOME is not None: - cuda_stub_dir = os.path.join(CUDA_HOME, "lib64", "stubs") - cuda_stub_lib = os.path.join(cuda_stub_dir, "libcuda.so") - if os.path.exists(cuda_stub_lib): - cuda_libraries.insert(0, "cuda") - cuda_library_dirs.append(cuda_stub_dir) - -setup( - name=module_name, - ext_modules=[ - CUDAExtension( - name=module_name, - include_dirs=[ - os.path.join(current_dir, "include"), - os.path.join(current_dir, "../mooncake-transfer-engine/include"), - ], - sources=[ - "src/pg_py.cpp", - "src/mooncake_backend.cpp", - "src/p2p_proxy.cpp", - "src/mooncake_worker.cu", - "src/mooncake_worker_host.cpp", - "src/mooncake_worker_thread.cpp", - "src/connection_poller.cpp", - ], - extra_compile_args={"cxx": cxx_args, "nvcc": device_args}, - libraries=cuda_libraries, - library_dirs=cuda_library_dirs, - extra_link_args=[ - "-Wl,-rpath,$ORIGIN", - "-L" + os.path.join(current_dir, "../mooncake-wheel/mooncake"), - "-Wl,--push-state,--no-as-needed", - "-l:engine.so", - "-Wl,--pop-state", - ], - ) - ], - cmdclass={"build_ext": BuildExtension}, -) diff --git a/mooncake-pg/src/CMakeLists.txt b/mooncake-pg/src/CMakeLists.txt index 4f145a869f..3da240c604 100644 --- a/mooncake-pg/src/CMakeLists.txt +++ b/mooncake-pg/src/CMakeLists.txt @@ -1,4 +1,149 @@ -add_library(mooncake_pg connection_poller.cpp mooncake_backend.cpp mooncake_worker.cu mooncake_worker_thread.cpp p2p_proxy.cpp pg_py.cpp) +set(_pg_device_source "${CMAKE_CURRENT_SOURCE_DIR}/mooncake_worker.cu") +if(USE_CUDA) + # Preserve the existing sm_80/sm_90 coverage unless the caller selects + # architectures through CMake's standard variable or CUDAARCHS environment. + if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES AND + "$ENV{CUDAARCHS}" STREQUAL "") + set(CMAKE_CUDA_ARCHITECTURES 80-real 90-real) + endif() + enable_language(CUDA) + set(_pg_device_input "${_pg_device_source}") +elseif(USE_MUSA OR USE_MACA) + set(_pg_device_object + "${CMAKE_CURRENT_BINARY_DIR}/mooncake_worker_device.o") + set(_pg_device_depfile + "${CMAKE_CURRENT_BINARY_DIR}/mooncake_worker_device.d") + set(_pg_device_flags + -std=c++20 + -O3 + -MMD + -MT "${_pg_device_object}" + -MF "${_pg_device_depfile}" + "-I${CMAKE_CURRENT_SOURCE_DIR}/../include" + "-I${CMAKE_CURRENT_SOURCE_DIR}/../../mooncake-transfer-engine/include") + set(_pg_device_dependencies + "${_pg_device_source}" + "${CMAKE_CURRENT_SOURCE_DIR}/../include/mooncake_worker_kernels.cuh" + "${CMAKE_CURRENT_SOURCE_DIR}/../include/error_types.h" + "${CMAKE_CURRENT_SOURCE_DIR}/../../mooncake-transfer-engine/include/common.h" + "${CMAKE_CURRENT_SOURCE_DIR}/../../mooncake-transfer-engine/include/cuda_alike.h" + "${CMAKE_CURRENT_SOURCE_DIR}/../../mooncake-transfer-engine/include/transfer_engine.h") + set(_pg_depfile_argument) + if(CMAKE_GENERATOR MATCHES "^Ninja" OR + (CMAKE_GENERATOR MATCHES "Makefiles" AND + CMAKE_VERSION VERSION_GREATER_EQUAL 3.20)) + set(_pg_depfile_argument DEPFILE "${_pg_device_depfile}") + endif() -set_target_properties(mooncake_pg PROPERTIES POSITION_INDEPENDENT_CODE ON) -target_link_libraries(mooncake_pg PUBLIC ${TORCH_LIBRARIES} transfer_engine ibverbs mlx5) + if(USE_MUSA) + set(_pg_platform_name MUSA) + set(_pg_device_compiler_name mcc) + if(DEFINED ENV{MUSA_HOME} AND NOT "$ENV{MUSA_HOME}" STREQUAL "") + set(_pg_device_compiler_hint "$ENV{MUSA_HOME}/bin") + else() + set(_pg_device_compiler_hint /usr/local/musa/bin) + endif() + set(_pg_platform_device_flags + -x musa + -fPIC + --cuda-gpu-arch=mp_21 + --cuda-gpu-arch=mp_31) + elseif(USE_MACA) + set(_pg_platform_name MACA) + set(_pg_device_compiler_name mxcc) + set(_pg_device_compiler_hint "${MACA_ROOT}/bin") + set(_pg_platform_device_flags --compiler-options=-fPIC) + endif() + + find_program(_pg_device_compiler NAMES "${_pg_device_compiler_name}" + HINTS "${_pg_device_compiler_hint}") + if(NOT _pg_device_compiler) + message(FATAL_ERROR + "USE_${_pg_platform_name}=ON requires the " + "${_pg_platform_name} compiler (${_pg_device_compiler_name})") + endif() + list(APPEND _pg_device_flags + "-DUSE_${_pg_platform_name}" + ${_pg_platform_device_flags}) + + add_custom_command( + OUTPUT "${_pg_device_object}" + COMMAND "${_pg_device_compiler}" ${_pg_device_flags} + -c "${_pg_device_source}" -o "${_pg_device_object}" + DEPENDS ${_pg_device_dependencies} + ${_pg_depfile_argument} + COMMENT "Compiling the Mooncake PG ${_pg_platform_name} worker" + COMMAND_EXPAND_LISTS + VERBATIM) + set_source_files_properties("${_pg_device_object}" + PROPERTIES GENERATED TRUE EXTERNAL_OBJECT TRUE) + set(_pg_device_input "${_pg_device_object}") +else() + message(FATAL_ERROR + "Mooncake PG currently requires one of USE_CUDA, USE_MUSA, or USE_MACA") +endif() + +# Keep the GPU fatbin out of the host core so auditwheel can safely repair the +# core while the device library is added to the wheel afterwards. +add_library(mooncake_pg_device SHARED "${_pg_device_input}") +target_include_directories( + mooncake_pg_device + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include + ${CMAKE_CURRENT_SOURCE_DIR}/../../mooncake-transfer-engine/include) + +if(USE_CUDA) + # The standalone device worker only needs CUDA C++17. Host-side PG remains + # C++20 and must not leak its headers into the device translation unit. + set_target_properties( + mooncake_pg_device + PROPERTIES CUDA_STANDARD 17 + CUDA_STANDARD_REQUIRED ON + CUDA_EXTENSIONS OFF + CUDA_RUNTIME_LIBRARY Shared) + target_compile_options( + mooncake_pg_device PRIVATE $<$:-O3>) + target_link_libraries(mooncake_pg_device PRIVATE CUDA::cudart) +elseif(USE_MUSA) + set_target_properties(mooncake_pg_device PROPERTIES LINKER_LANGUAGE CXX) + target_link_libraries(mooncake_pg_device PRIVATE musa musart rt) +elseif(USE_MACA) + if(NOT DEFINED MACA_RUNTIME_LIBS) + set(MACA_RUNTIME_LIBS mcruntime mxc-runtime64 rt) + endif() + set_target_properties(mooncake_pg_device PROPERTIES LINKER_LANGUAGE CXX) + target_link_libraries(mooncake_pg_device PRIVATE ${MACA_RUNTIME_LIBS}) +endif() + +set_target_properties( + mooncake_pg_device + PROPERTIES POSITION_INDEPENDENT_CODE ON + BUILD_RPATH "$ORIGIN" + BUILD_RPATH_USE_ORIGIN YES) + +add_library(mooncake_pg SHARED + control_plane/link_manager.cpp + control_plane/coordinator.cpp + control_plane/coordinator_host.cpp + control_plane/agent.cpp + control_plane/agent_host.cpp + control_plane/rpc_runtime.cpp + gpu_runtime.cpp + mooncake_communicator.cpp + mooncake_pg.cpp + mooncake_worker_host.cpp + mooncake_worker_thread.cpp + p2p_proxy.cpp +) + +target_compile_features(mooncake_pg PRIVATE cxx_std_20) +target_include_directories( + mooncake_pg + PUBLIC $ + $) +target_link_libraries(mooncake_pg PRIVATE transfer_engine mooncake_pg_device) + +set_target_properties( + mooncake_pg + PROPERTIES POSITION_INDEPENDENT_CODE ON + BUILD_RPATH "$ORIGIN" + BUILD_RPATH_USE_ORIGIN YES) diff --git a/mooncake-pg/src/connection_poller.cpp b/mooncake-pg/src/connection_poller.cpp deleted file mode 100644 index e516cc4e1b..0000000000 --- a/mooncake-pg/src/connection_poller.cpp +++ /dev/null @@ -1,629 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include "memory_location.h" -#include "mooncake_worker.cuh" -#include "pg_utils.h" - -namespace mooncake { - -// Same check as nvlink_transport.cpp and mooncake_ep_buffer.cpp. -// On MNNVL clusters all GPUs support fabric mem handles, meaning -// NVLink transport can only access cuMemCreate(FABRIC) memory -// cross-node -- CPU heap buffers are invisible to remote peers. -#if !defined(MOONCAKE_EP_USE_MUSA) && !defined(USE_MACA) -static bool supportFabricMem() { - const char* nvlink_ipc = getenv("MC_USE_NVLINK_IPC"); - - bool fabric_enabled = nvlink_ipc && strcmp(nvlink_ipc, "0") == 0; - if (!fabric_enabled) return false; - - int num_devices = 0; - cudaError_t err = cudaGetDeviceCount(&num_devices); - if (err != cudaSuccess || num_devices == 0) return false; - - for (int dev = 0; dev < num_devices; ++dev) { - int supported = 0; - cuDeviceGetAttribute( - &supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, dev); - if (!supported) return false; - } - return true; -} -#else -// MUSA and MACA do not use NVIDIA NVLink fabric memory handles here. -static bool supportFabricMem() { return false; } -#endif -ConnectionContext::ConnectionContext(int backendIndex, int rank, int size, - bool isDummy, - uint64_t* local2global_rank_map, - c10::intrusive_ptr<::c10d::Store> store, - std::shared_ptr meta, - std::shared_ptr p2p_proxy, - TransferEngine* engine) - : backendIndex_(backendIndex), - rank_(rank), - groupSize_(size), - pollingLimit_(size), - isDummy_(isDummy), - establishedGroupSize_(0), - local2global_rank_map_(local2global_rank_map), - store_(std::move(store)), - meta_(std::move(meta)), - p2p_proxy_(std::move(p2p_proxy)), - engine_(engine), - skip_warmup_(supportFabricMem()) { - if (skip_warmup_) { - // On MNNVL clusters, CPU heap buffers aren't fabric-accessible so - // remote NVLink writes to them will fail. The fabric topology already - // guarantees connectivity, so we skip the warmup handshake entirely. - warmup_send_region_ = nullptr; - warmup_recv_region_ = nullptr; - return; - } - - warmup_send_region_ = new int32_t[kMaxNumRanks]{}; - warmup_send_region_[0] = 1; - int rc = engine_->registerLocalMemory( - warmup_send_region_, kMaxNumRanks * sizeof(int32_t), kWildcardLocation); - TORCH_CHECK(!rc, "Failed to register local memory for context."); - - warmup_recv_region_ = new int32_t[kMaxNumRanks]{}; - rc = engine_->registerLocalMemory( - warmup_recv_region_, kMaxNumRanks * sizeof(int32_t), kWildcardLocation); - TORCH_CHECK(!rc, "Failed to register local memory for context."); -} - -ConnectionContext::~ConnectionContext() { - if (resource_abandoned_) { - LOG(WARNING) << "Resource leak in ConnectionContext: cleanup skipped " - "due to hung operations."; - return; - } - - for (int i = 0; i < groupSize_; ++i) { - if (peerStates_[i].segmentId.has_value()) { - engine_->closeSegment(peerStates_[i].segmentId.value()); - } - } - - if (warmup_send_region_) { - engine_->unregisterLocalMemory(warmup_send_region_); - delete[] warmup_send_region_; - } - if (warmup_recv_region_) { - engine_->unregisterLocalMemory(warmup_recv_region_); - delete[] warmup_recv_region_; - } -} - -int ConnectionContext::getTotalConnectedPeers() const { - return totalConnectedPeers_.load(std::memory_order_acquire); -} - -void ConnectionContext::extendGroupSizeTo(int newGroupSize) { - const int oldGroupSize = groupSize_.load(std::memory_order_acquire); - if (newGroupSize == oldGroupSize) return; - - TORCH_CHECK( - newGroupSize >= 0 && static_cast(newGroupSize) < kMaxNumRanks, - "Size out of range"); - TORCH_CHECK(newGroupSize >= oldGroupSize, "newGroupSize < oldGroupSize"); - - // Reset local peer state for newly added ranks - for (int i = oldGroupSize; i < newGroupSize; ++i) { - meta_->peerConnected[i] = false; - } - - groupSize_.store(newGroupSize, std::memory_order_release); - // Keep polling range aligned with the largest known size. - pollingLimit_.store(newGroupSize, std::memory_order_release); -} - -void ConnectionContext::setPollingLimitTo(int pollingLimit) { - const int groupSize = groupSize_.load(std::memory_order_acquire); - TORCH_CHECK(pollingLimit >= groupSize, - "pollingLimit must be >= current groupSize"); - TORCH_CHECK( - pollingLimit >= 0 && static_cast(pollingLimit) < kMaxNumRanks, - "Size out of range"); - pollingLimit_.store(pollingLimit, std::memory_order_release); -} - -bool ConnectionContext::isAllPeerConnected() const { - return totalConnectedPeers_ == groupSize_; -} - -void ConnectionContext::waitUntilAllConnected() { - if (isAllPeerConnected()) return; - - std::unique_lock lock(backend_wakeup_mutex_); - backend_wakeup_cv_.wait(lock, [this]() { - return isAllPeerConnected() || - isShutdown_.load(std::memory_order_acquire); - }); - establishedGroupSize_.store(groupSize_, std::memory_order_release); -} - -void ConnectionContext::waitUntilNewRanksConnected() { - if (isDummy_) { - return; - } - const int targetGroupSize = groupSize_.load(std::memory_order_acquire); - const int established = - establishedGroupSize_.load(std::memory_order_acquire); - if (established >= targetGroupSize) { - return; - } - - std::unique_lock lock(backend_wakeup_mutex_); - backend_wakeup_cv_.wait(lock, [this, targetGroupSize, established]() { - if (isShutdown_.load(std::memory_order_acquire)) { - return true; - } - - for (int i = established; i < targetGroupSize; ++i) { - if (!meta_->peerConnected[i]) { - return false; - } - } - return true; - }); - - establishedGroupSize_.store(targetGroupSize, std::memory_order_release); -} - -void ConnectionContext::bootstrapLocalPeer(const std::string& localServerName, - const SegmentInfo& localRankInfo) { - auto& peerState = peerStates_[rank_]; - if (peerState.state == PeerConnectionState::CONNECTED) { - return; - } - - auto segment_id = engine_->openSegment(localServerName); - meta_->segmentIDs[rank_] = segment_id; - peerState.segmentId = segment_id; - memcpy(&meta_->segmentInfos[rank_], &localRankInfo, sizeof(SegmentInfo)); - - meta_->peerConnected[rank_] = true; - ConnectionPoller::GetInstance() - .global_peerConnected_[local2global_rank_map_[rank_]] = true; - peerState.state = PeerConnectionState::CONNECTED; - peerState.countedInGroup = true; - - { - std::lock_guard lock(backend_wakeup_mutex_); - totalConnectedPeers_.store(1, std::memory_order_release); - if (isAllPeerConnected()) { - backend_wakeup_cv_.notify_all(); - } - } -} - -void ConnectionContext::shutdown() { - // Notify backends that may be blocked in waitUntilAllConnected. - { - std::lock_guard backend_lock(backend_wakeup_mutex_); - isShutdown_.store(true, std::memory_order_release); - } - backend_wakeup_cv_.notify_all(); -} - -bool ConnectionContext::poll() { - if (isShutdown_.load(std::memory_order_acquire)) { - return false; - } - - bool did_work = false; - - // Poll all peers sequentially. - const int pollingLimit = pollingLimit_.load(std::memory_order_acquire); - for (int pollingRank = 0; pollingRank < pollingLimit; ++pollingRank) { - did_work |= pollPeer(pollingRank); - } - - return did_work; -} - -bool ConnectionContext::pollPeer(int pollingRank) { - auto globalPollingRank = local2global_rank_map_[pollingRank]; - auto& global_peerConnected_ = - ConnectionPoller::GetInstance().global_peerConnected_; - auto& peerState = peerStates_[pollingRank]; - bool state_changed = false; - - switch (peerState.state) { - case PeerConnectionState::WAITING_STORE: { - // See if we need backoff - auto now = std::chrono::steady_clock::now(); - auto elapsed = static_cast( - std::chrono::duration_cast( - now - peerState.last_check_store) - .count()); - if (elapsed < peerState.check_store_backoff_ms) { - return false; - } - - peerState.last_check_store = now; - - auto serverNameKey = - getServerNameStoreKey(backendIndex_, pollingRank); - auto bufferKey = getBufferStoreKey(backendIndex_, pollingRank); - - std::string peerServerName; - std::vector buffer_data; - try { - if (!store_->check({serverNameKey, bufferKey})) { - peerState.increaseCheckStoreBackoff(); - return false; - } - peerServerName = store_->get_to_str(serverNameKey); - buffer_data = store_->get(bufferKey); - - if (buffer_data.size() < sizeof(SegmentInfo)) { - LOG(WARNING) - << "Rank " << rank_ << " got invalid buffer data from " - << pollingRank << "."; - peerState.increaseCheckStoreBackoff(); - return false; - } - } catch (const std::exception& e) { - peerState.increaseCheckStoreBackoff(); - return false; - } - - auto segment_id = engine_->openSegment(peerServerName); - meta_->segmentIDs[pollingRank] = segment_id; - peerState.segmentId = segment_id; - - memcpy(&meta_->segmentInfos[pollingRank], buffer_data.data(), - sizeof(SegmentInfo)); - - if (skip_warmup_) { - // MNNVL: fabric guarantees connectivity, skip warmup write - // since CPU heap buffers aren't fabric-accessible anyway. - meta_->peerConnected[pollingRank] = true; - global_peerConnected_[globalPollingRank] = true; - peerState.state = PeerConnectionState::CONNECTED; - { - std::lock_guard lock(backend_wakeup_mutex_); - if (pollingRank < - groupSize_.load(std::memory_order_acquire)) { - totalConnectedPeers_.fetch_add( - 1, std::memory_order_release); - peerState.countedInGroup = true; - backend_wakeup_cv_.notify_all(); - } - } - } else if (pollingRank <= rank_) { - // Send a warmup request to establish connections - auto batchID = engine_->allocateBatchID(1); - engine_->submitTransfer( - batchID, - {TransferRequest{ - .opcode = TransferRequest::WRITE, - .source = warmup_send_region_, - .target_id = meta_->segmentIDs[pollingRank], - .target_offset = - meta_->segmentInfos[pollingRank].warmup_buffer[1] + - rank_ * sizeof(int32_t), - .length = sizeof(int32_t), - }}); - peerState.warmupBatchId = batchID; - peerState.state = PeerConnectionState::WAITING_WARMUP_TRANSFER; - } else { - // For pollingRank > rank_, wait for the peer's warmup write to - // arrive. - peerState.state = PeerConnectionState::WAITING_PEER_WARMUP; - } - peerState.resetCheckStoreBackoff(); - state_changed = true; - break; - } - - case PeerConnectionState::WAITING_WARMUP_TRANSFER: { - TransferStatus status; - engine_->getTransferStatus(peerState.warmupBatchId.value(), 0, - status); - if (status.s == TransferStatusEnum::COMPLETED) { - engine_->freeBatchID(peerState.warmupBatchId.value()); - peerState.warmupBatchId = std::nullopt; - meta_->peerConnected[pollingRank] = true; - global_peerConnected_[globalPollingRank] = true; - peerState.state = PeerConnectionState::CONNECTED; - - { - std::lock_guard lock(backend_wakeup_mutex_); - if (pollingRank < - groupSize_.load(std::memory_order_acquire)) { - totalConnectedPeers_.fetch_add( - 1, std::memory_order_release); - peerState.countedInGroup = true; - backend_wakeup_cv_.notify_all(); - } - } - state_changed = true; - } else if (status.s == TransferStatusEnum::FAILED) { - LOG(WARNING) << "Warmup request " << rank_ << " -> " - << pollingRank << " failed."; - // Free resources and retry - engine_->freeBatchID(peerState.warmupBatchId.value()); - engine_->closeSegment(peerState.segmentId.value()); - peerState.warmupBatchId = std::nullopt; - peerState.segmentId = std::nullopt; - peerState.state = PeerConnectionState::WAITING_STORE; - state_changed = true; - } - break; - } - - case PeerConnectionState::WAITING_PEER_WARMUP: { - if (*reinterpret_cast( - &warmup_recv_region_[pollingRank])) { - meta_->peerConnected[pollingRank] = true; - global_peerConnected_[globalPollingRank] = true; - peerState.state = PeerConnectionState::CONNECTED; - { - std::lock_guard lock(backend_wakeup_mutex_); - if (pollingRank < - groupSize_.load(std::memory_order_acquire)) { - totalConnectedPeers_.fetch_add( - 1, std::memory_order_release); - peerState.countedInGroup = true; - backend_wakeup_cv_.notify_all(); - } - } - state_changed = true; - } - break; - } - - case PeerConnectionState::CONNECTED: { - // A peer may be warmed up (CONNECTED) while still outside of the - // current groupSize_. When it later enters the group, we need to - // update totalConnectedPeers_ lazily here; otherwise - // waitUntilAllConnected()/isAllPeerConnected() may hang. - - // ATTENTION: Ensure consistency of local (meta_->peerConnected) - // and global (global_peerConnected_). - // - // Assuming there are two backends, and Backend A detects a failure. - // (by setting meta_->peerConnected[pollingRank] to false) - // - // For Backend A: - // Observes `Local=false, Global=true` and updates `Global=false`. - // - // For Backend B: - // Observes `Local=true, Global=false` and updates `Local=false`. - // - // Because ConnectionPoller runs as a single thread and processes - // contexts sequentially, no race conditions occur. Besides, the - // failure signal is guaranteed to propagate to all other contexts - // (e.g., Backend B) before the initiator (e.g., Backend A) is - // processed again. - // - // Thus, we ensure that if one backend disconnects, all backends - // disconnect. - - if (meta_->peerConnected[pollingRank] && - global_peerConnected_[globalPollingRank]) { - if (!peerState.countedInGroup && - pollingRank < groupSize_.load(std::memory_order_acquire)) { - std::lock_guard lock(backend_wakeup_mutex_); - totalConnectedPeers_.fetch_add(1, - std::memory_order_release); - peerState.countedInGroup = true; - backend_wakeup_cv_.notify_all(); - } - // happy path: both are connected. - break; - } - - // If we reach here, at least one peer connected (local or global) - // reports a failure. We must set both to false here. - global_peerConnected_[globalPollingRank] = false; - meta_->peerConnected[pollingRank] = false; - meta_->activeRanks[pollingRank] = false; - if (meta_->activeRanksTensor.device().is_cpu()) { - meta_->activeRanksTensor[pollingRank] = 0; - } - - // Reset store - try { - store_->deleteKey( - getServerNameStoreKey(backendIndex_, pollingRank)); - store_->deleteKey( - getBufferStoreKey(backendIndex_, pollingRank)); - store_->deleteKey( - getExtensionStateStoreKey(backendIndex_, pollingRank)); - } catch (const std::exception& e) { - LOG(WARNING) << "Rank " << rank_ - << " got an exception when deleteKey for peer " - << pollingRank << ": " << e.what(); - } - - // Reset warmup region - if (warmup_recv_region_) { - *reinterpret_cast( - &warmup_recv_region_[pollingRank]) = 0; - } - - // Reset P2PProxy states - p2p_proxy_->resetPeerState(pollingRank); - - // Back to WAITING_STORE to reconnect it. - peerState.state = PeerConnectionState::WAITING_STORE; - engine_->closeSegment(peerState.segmentId.value()); - peerState.segmentId = std::nullopt; - if (peerState.countedInGroup) { - totalConnectedPeers_.fetch_sub(1, std::memory_order_release); - peerState.countedInGroup = false; - } - state_changed = true; - break; - } - - case PeerConnectionState::EXPIRING: - TORCH_CHECK( - false, "Unexpected PeerConnectionState::EXPIRING in pollPeer."); - break; - } - - return state_changed; -} - -bool ConnectionContext::isStopped() const { - for (auto& peerState : peerStates_) { - if (peerState.state != PeerConnectionState::EXPIRING) { - return false; - } - } - return true; -} - -bool ConnectionContext::drainPoller() const { - BackoffWaiter waiter; - return waiter.wait_for(std::chrono::milliseconds(kDrainPollerTimeoutMs), - [this] { return isStopped(); }); -} - -void ConnectionContext::abandonResources() { resource_abandoned_ = true; } - -bool ConnectionContext::tryStop() { - bool stopped = true; - for (auto& peerState : peerStates_) { - if (peerState.state == PeerConnectionState::EXPIRING) { - continue; - } - - if (peerState.state != PeerConnectionState::WAITING_WARMUP_TRANSFER) { - peerState.state = PeerConnectionState::EXPIRING; - continue; - } - - // For WAITING_WARMUP_TRANSFER, wait for the existing transfer to - // complete so that we can safely release the registered memory. - TransferStatus status; - engine_->getTransferStatus(peerState.warmupBatchId.value(), 0, status); - - if (status.s == TransferStatusEnum::COMPLETED || - status.s == TransferStatusEnum::FAILED) { - engine_->freeBatchID(peerState.warmupBatchId.value()); - peerState.warmupBatchId = std::nullopt; - peerState.state = PeerConnectionState::EXPIRING; - } else { - stopped = false; - } - } - return stopped; -} - -ConnectionPoller::ConnectionPoller() = default; - -void ConnectionPoller::ensureThreadStarted() { - bool expected = false; - if (!pollerThreadStarted_.compare_exchange_strong( - expected, true, std::memory_order_acq_rel)) { - return; - } - pollerThread_ = std::thread([this] { pollerLoop(); }); -} - -void ConnectionPoller::registerContext( - const std::shared_ptr& ctx) { - ensureThreadStarted(); - { - std::lock_guard lock(contexts_mutex_); - contexts_.push_back(ctx); - contexts_version_.fetch_add(1, std::memory_order_release); - } - wakeup_cv_.notify_one(); -} - -void ConnectionPoller::removeContext( - const std::shared_ptr& ctx) { - TORCH_CHECK(ctx->isShutdown_, "connection context hasn't shutdown."); - - { - std::lock_guard lock(contexts_mutex_); - contexts_.erase(std::remove(contexts_.begin(), contexts_.end(), ctx), - contexts_.end()); - contexts_version_.fetch_add(1, std::memory_order_release); - } - wakeup_cv_.notify_one(); -} - -void ConnectionPoller::pollerLoop() { - // Similar to `local_proxies` in P2PDeviceWorker's worker loop, - // it is a thread-local cache to avoid locking too frequently. - std::vector> local_contexts; - uint64_t local_version = std::numeric_limits::max(); - - // Keep track of removed contexts and free them properly to ensure - // graceful backend shutdown. - std::vector> zombie_contexts; - - while (true) { - const auto current_version = - contexts_version_.load(std::memory_order_acquire); - if (local_version != current_version) { - std::vector> new_contexts; - { - std::lock_guard lock(contexts_mutex_); - new_contexts = contexts_; - // Reload to ensure version matches with new_contexts - local_version = - contexts_version_.load(std::memory_order_acquire); - } - - // Find zombie contexts - for (const auto& old_c : local_contexts) { - auto it = - std::find(new_contexts.begin(), new_contexts.end(), old_c); - if (it != new_contexts.end()) continue; - zombie_contexts.emplace_back(old_c); - } - - local_contexts = std::move(new_contexts); - } - - bool did_work = false; - bool all_connected = true; - - for (const auto& ctx : local_contexts) { - did_work |= ctx->poll(); - all_connected &= ctx->isAllPeerConnected(); - } - - for (auto it = zombie_contexts.begin(); it < zombie_contexts.end();) { - auto& zombie = *it; - if (zombie->tryStop()) - it = zombie_contexts.erase(it); - else { - ++it; - } - } - - if (did_work) continue; - - std::unique_lock lock(wakeup_mutex_); - auto sleep_ms = - all_connected ? kAllConnectedIdleSleepMs : kConnectingIdleSleepMs; - wakeup_cv_.wait_for(lock, std::chrono::milliseconds(sleep_ms), [&]() { - if (local_version != - contexts_version_.load(std::memory_order_acquire)) - return true; - return false; - }); - } -} -} // namespace mooncake diff --git a/mooncake-pg/src/control_plane/agent.cpp b/mooncake-pg/src/control_plane/agent.cpp new file mode 100644 index 0000000000..b826bff76d --- /dev/null +++ b/mooncake-pg/src/control_plane/agent.cpp @@ -0,0 +1,426 @@ +#include "control_plane/agent.h" + +#include +#include +#include + +#include "error_types.h" + +namespace mooncake { + +AgentStateMachine::AgentStateMachine(GlobalRank rank, int max_world_size) + : rank_(rank), max_world_size_(max_world_size) { + PG_ASSERT(max_world_size_ > 0 && max_world_size_ <= kMaxNumRanks, + "invalid max_world_size: ", max_world_size_); + global_rank_states_ = std::vector(max_world_size_); + global_rank_epochs_.assign(max_world_size_, 0); + global_rank_state_versions_.assign(max_world_size_, 0); + rank_connections_.resize(max_world_size_); + observed_link_state_.assign(max_world_size_, LinkEvent::EventType::None); + observed_target_rank_epochs_.assign(max_world_size_, 0); +} + +void AgentStateMachine::appendApplyViewEffect(const GroupView& view, + AgentApplyResult& effects) const { + std::vector activatable(view.rank_order.size()); + for (size_t i = 0; i < view.rank_order.size(); ++i) { + GlobalRank gr = view.rank_order[i]; + bool healthy = global_rank_states_[gr] == RankState::Healthy; + const auto& member = view.members[gr]; + activatable[i] = healthy && + (member.isActive() || member.isAwaitingActivation()) && + member.hasEndpoint(); + } + effects.push_back(ApplyViewToCommunicator{view, global_rank_states_, + global_rank_epochs_, + std::move(activatable)}); +} + +AgentApplyResult AgentStateMachine::registerGroup(const GroupView& group) { + AgentApplyResult effects; + groups_.insert_or_assign(group.group_id, group); + appendApplyViewEffect(group, effects); + return effects; +} + +void AgentStateMachine::unregisterGroup(GroupId group_id) { + groups_.erase(group_id); +} + +GroupView AgentStateMachine::getGroupView(GroupId group_id) const { + auto it = groups_.find(group_id); + if (it != groups_.end()) { + return it->second; + } + return GroupView{}; +} + +void AgentStateMachine::appendApplyViewEffectsForRank( + GlobalRank rank, AgentApplyResult& effects) const { + for (const auto& [group_id, view] : groups_) { + if (std::find(view.rank_order.begin(), view.rank_order.end(), rank) != + view.rank_order.end()) { + appendApplyViewEffect(view, effects); + } + } +} + +void AgentStateMachine::resetRankForNewEpoch(GlobalRank rank, + uint64_t rank_epoch, + AgentApplyResult& effects) { + global_rank_epochs_[rank] = rank_epoch; + global_rank_state_versions_[rank] = 0; + global_rank_states_[rank] = RankState::Offline; + rank_connections_[rank].reset(); + + if (rank != rank_) { + effects.push_back(DisconnectLink{rank}); + effects.push_back(StopReconnect{rank}); + effects.push_back(NotifyLinkRefreshed{rank}); + } +} + +AgentApplyResult AgentStateMachine::handlePeerJoined( + const PeerJoinedPush& push) { + AgentApplyResult effects; + if (!rankInRange(push.rank)) { + LOG(WARNING) << "AgentStateMachine: handlePeerJoined out-of-range rank " + << push.rank; + return effects; + } + if (push.rank == rank_) return effects; + + if (push.rank_epoch < global_rank_epochs_[push.rank]) return effects; + + const bool new_rank_epoch = + push.rank_epoch > global_rank_epochs_[push.rank]; + if (new_rank_epoch) { + resetRankForNewEpoch(push.rank, push.rank_epoch, effects); + appendApplyViewEffectsForRank(push.rank, effects); + } else if (global_rank_states_[push.rank] == RankState::Offline) { + // Offline is terminal within a rank epoch. A delayed PeerJoined for + // that epoch must not restart the old connection. + return effects; + } + + if (rank_connections_[push.rank].has_value()) return effects; + + rank_connections_[push.rank] = RankConnectionMetadata{ + .rank = push.rank, + .rank_epoch = push.rank_epoch, + .agent_addr = "", + .te_server_name = push.te_server_name, + .warmup_recv_addr = push.warmup_recv_addr, + }; + effects.push_back(EnablePeerProbe{push.rank, push.rank_epoch, + push.te_server_name, + push.warmup_recv_addr}); + return effects; +} + +AgentApplyResult AgentStateMachine::handleRankStateUpdate( + const RankStatePush& push) { + AgentApplyResult effects; + if (!rankInRange(push.rank)) { + LOG(WARNING) << "AgentStateMachine: handleRankStateUpdate out-of-range " + << push.rank; + return effects; + } + + if (push.rank_epoch < global_rank_epochs_[push.rank]) return effects; + + const bool new_rank_epoch = + push.rank_epoch > global_rank_epochs_[push.rank]; + if (new_rank_epoch) { + resetRankForNewEpoch(push.rank, push.rank_epoch, effects); + } + + if (push.rank_state_version <= global_rank_state_versions_[push.rank]) + return effects; + + global_rank_states_[push.rank] = push.new_state; + global_rank_state_versions_[push.rank] = push.rank_state_version; + appendApplyViewEffectsForRank(push.rank, effects); + + // Remote Offline: tear down TE link AND stop candidate probe. + if (push.rank != rank_ && push.new_state == RankState::Offline) { + rank_connections_[push.rank].reset(); + if (!new_rank_epoch) { + effects.push_back(DisconnectLink{push.rank}); + effects.push_back(StopReconnect{push.rank}); + effects.push_back(NotifyLinkRefreshed{push.rank}); + } + } + + return effects; +} + +PGResult AgentStateMachine::applyGroupView( + const GroupView& view) { + AgentApplyResult effects; + const auto& group_id = view.group_id; + auto it = groups_.find(group_id); + if (it == groups_.end()) { + LOG(WARNING) << "[AGENT] applyGroupView group=" << group_id + << " NOT FOUND in groups_ (epoch=" << view.epoch << ")"; + return makePGError( + PGErrorCode::InvalidState, + "group not found while applying GroupView: " + group_id); + } + + const auto& old_view = it->second; + + // View application is idempotent. A sync response may race with the + // regular ViewUpdate push for the same decision, or even arrive after a + // newer view has already been installed. + if (view.epoch < old_view.epoch) { + LOG(WARNING) << "[AGENT] Ignored stale GroupView for group=" << group_id + << " epoch=" << view.epoch; + return effects; + } else if (view.epoch == old_view.epoch) { + if (view == old_view) return effects; + LOG(ERROR) << "[AGENT] Ignored conflicting GroupView for group=" + << group_id << " epoch=" << view.epoch; + return makePGError(PGErrorCode::InternalError, + "conflicting GroupView for group " + group_id + + " at epoch " + std::to_string(view.epoch)); + } + + // Collect peers whose segment caches must be refreshed. + std::vector need_segment_refresh; + + // Detect endpoint updates. + for (size_t r = 0; r < view.members.size(); ++r) { + if (r == static_cast(rank_)) continue; + if (!view.members[r].isMember()) continue; + uint64_t old_epoch = 0; + uint64_t new_epoch = 0; + if (old_view.members[r].hasEndpoint()) { + old_epoch = old_view.members[r].endpoint->endpoint_epoch; + } + if (view.members[r].hasEndpoint()) { + new_epoch = view.members[r].endpoint->endpoint_epoch; + } + if (new_epoch != 0 && new_epoch != old_epoch) { + effects.push_back(RefreshPeerLink{static_cast(r)}); + need_segment_refresh.push_back(static_cast(r)); + } + } + + // Applying the view must happen-before waking group/rank waiters: callers + // may submit a collective as soon as waitUntilGroupReady()/joinGroup() + // returns, and that collective must observe the new data-plane metadata. + appendApplyViewEffect(view, effects); + + // Detect rank activation transitions. + if (!old_view.members.empty()) { + std::vector newly_activated; + for (size_t igr = 0; igr < view.rank_order.size(); ++igr) { + GlobalRank gr = view.rank_order[igr]; + if (!old_view.members[gr].isActive() && + view.members[gr].isActive()) { + newly_activated.push_back(gr); + } + } + if (!newly_activated.empty()) { + effects.push_back( + NotifyRanksActivated{group_id, std::move(newly_activated)}); + } + } + + // Detect Ready transition. + if (old_view.status != GroupStatus::Ready && + view.status == GroupStatus::Ready) { + effects.push_back(NotifyGroupReady{group_id}); + } + + it->second = view; + + // Must come AFTER ApplyViewToCommunicator: refreshSegmentID requires + // latest meta_->rank_order. + for (auto gr : need_segment_refresh) { + effects.push_back(NotifyLinkRefreshed{gr}); + } + + return effects; +} + +PGResult AgentStateMachine::handleViewUpdate( + const ViewUpdatePush& push) { + return applyGroupView(push.view); +} + +HeartbeatRequest AgentStateMachine::buildHeartbeat() const { + HeartbeatRequest req; + req.rank = rank_; + return req; +} + +AgentApplyResult AgentStateMachine::applyRegisterAgentResponse( + const RegisterAgentResponse& resp) { + AgentApplyResult effects; + + if (static_cast(resp.all_rank_states.size()) != max_world_size_ || + static_cast(resp.all_rank_epochs.size()) != max_world_size_ || + static_cast(resp.all_rank_state_versions.size()) != + max_world_size_) { + LOG(ERROR) << "AgentStateMachine: malformed RegisterAgentResponse"; + coordinator_connection_ = CoordinatorConnection::Disconnected; + return effects; + } + + self_rank_epoch_.store(resp.rank_epoch, std::memory_order_release); + + // PeerJoined and RankState pushes use independent RPCs and can arrive while + // the RegisterAgent response is in flight. Merge its snapshot monotonically + // so it cannot roll a rank back to an older epoch or state version. + for (int rank = 0; rank < max_world_size_; ++rank) { + const auto response_epoch = resp.all_rank_epochs[rank]; + const auto response_version = resp.all_rank_state_versions[rank]; + if (response_epoch < global_rank_epochs_[rank] || + (response_epoch == global_rank_epochs_[rank] && + response_version < global_rank_state_versions_[rank])) + continue; + + if (response_epoch > global_rank_epochs_[rank]) { + resetRankForNewEpoch(rank, response_epoch, effects); + } + global_rank_states_[rank] = resp.all_rank_states[rank]; + global_rank_epochs_[rank] = response_epoch; + global_rank_state_versions_[rank] = response_version; + } + + for (const auto& gv : resp.groups) { + groups_[gv.group_id] = gv; + appendApplyViewEffect(gv, effects); + } + + // The connection list belongs to the same snapshot. Do not install an + // entry invalidated by above. + for (const auto& connection : resp.rank_connections) { + if (connection.rank_epoch != global_rank_epochs_[connection.rank] || + global_rank_states_[connection.rank] == RankState::Offline) + continue; + + rank_connections_[connection.rank] = connection; + effects.push_back(EnablePeerProbe{ + .rank = connection.rank, + .rank_epoch = connection.rank_epoch, + .te_server_name = connection.te_server_name, + .warmup_recv_addr = connection.warmup_recv_addr, + }); + } + + coordinator_connection_ = CoordinatorConnection::Connected; + return effects; +} + +AgentApplyResult AgentStateMachine::reset(uint64_t new_session_id) { + AgentApplyResult effects; + + agent_session_id_.store(new_session_id, std::memory_order_release); + self_rank_epoch_.store(0, std::memory_order_release); + std::fill(observed_link_state_.begin(), observed_link_state_.end(), + LinkEvent::EventType::None); + std::fill(observed_target_rank_epochs_.begin(), + observed_target_rank_epochs_.end(), 0); + link_state_version_ = 0; + acked_link_state_version_ = 0; + groups_.clear(); + global_rank_states_ = std::vector(max_world_size_); + std::fill(global_rank_epochs_.begin(), global_rank_epochs_.end(), 0); + std::fill(global_rank_state_versions_.begin(), + global_rank_state_versions_.end(), 0); + for (auto& conn : rank_connections_) conn.reset(); + + effects.push_back(DisconnectAllLinks{}); + effects.push_back(ClearAllPeerMetadata{}); + + return effects; +} + +AgentApplyResult AgentStateMachine::pushLinkEvent(const LinkEvent& event) { + AgentApplyResult effects; + + if (event.events.size() != static_cast(max_world_size_) || + event.target_rank_epochs.size() != + static_cast(max_world_size_)) { + LOG(WARNING) << "AgentStateMachine: invalid LinkEvent size. " + << "Expected max_world_size=" << max_world_size_ + << "; dropping."; + return effects; + } + + bool changed = false; + for (int peer = 0; peer < max_world_size_; ++peer) { + auto type = event.events[peer]; + if (type == LinkEvent::EventType::None) continue; + const bool has_prior_link_observation = + observed_link_state_[peer] != LinkEvent::EventType::None; + if (!recordLinkEvent(peer, event.target_rank_epochs[peer], type)) + continue; + + changed = true; + if (type == LinkEvent::EventType::Failure) { + effects.push_back(RequestLinkHealthCheck{peer}); + } else if (peer != rank_) { + // Reset only when this success follows an earlier observation, + // i.e. a recovery or a new peer incarnation. + if (has_prior_link_observation) { + effects.push_back(ResetPeerState{peer}); + } + effects.push_back(NotifyLinkRefreshed{peer}); + } + } + + if (changed && + coordinator_connection_ == CoordinatorConnection::Connected) { + auto report = getLinkEventReport(); + if (report.has_value()) { + effects.push_back(SendLinkEventReport{std::move(*report)}); + } + } + + return effects; +} + +std::optional AgentStateMachine::getLinkEventReport() const { + if (link_state_version_ <= acked_link_state_version_) return std::nullopt; + + LinkEventReport report; + report.reporter_rank = rank_; + report.agent_session_id = getAgentSessionId(); + report.reporter_rank_epoch = getRankEpoch(); + report.report_id = link_state_version_; + report.events = observed_link_state_; + report.target_rank_epochs = observed_target_rank_epochs_; + return report; +} + +void AgentStateMachine::handleLinkEventReportAck( + const LinkEventReportAck& ack) { + if (ack.reporter_rank != rank_ || + ack.reporter_rank_epoch != getRankEpoch()) { + return; + } + acked_link_state_version_ = + std::max(acked_link_state_version_, ack.report_id); +} + +bool AgentStateMachine::recordLinkEvent(GlobalRank peer, + uint64_t target_rank_epoch, + LinkEvent::EventType type) { + if (target_rank_epoch != global_rank_epochs_[peer]) return false; + + if (observed_target_rank_epochs_[peer] == target_rank_epoch && + observed_link_state_[peer] == type) { + return false; + } + + observed_target_rank_epochs_[peer] = target_rank_epoch; + observed_link_state_[peer] = type; + ++link_state_version_; + return true; +} + +} // namespace mooncake diff --git a/mooncake-pg/src/control_plane/agent_host.cpp b/mooncake-pg/src/control_plane/agent_host.cpp new file mode 100644 index 0000000000..797ecb2981 --- /dev/null +++ b/mooncake-pg/src/control_plane/agent_host.cpp @@ -0,0 +1,716 @@ +#include "control_plane/agent_host.h" + +#include +#include +#include + +#include + +#include "mooncake_communicator.h" +#include "control_plane/link_manager.h" +#include "control_plane/rpc_runtime.h" +#include "pg_utils.h" + +namespace mooncake { + +namespace { + +// Generate a process-unique key for one logical registration. +uint64_t generateInitialAgentSessionId() { + auto now = std::chrono::steady_clock::now().time_since_epoch().count(); + uint64_t pid = static_cast(getpid()); + uint64_t base = (pid << 32) ^ static_cast(now); + return base == 0 ? 1 : base; +} + +template +PGResult callAndCheck(RpcClient& client, const std::string& addr, + Request request) { + PG_TRY(auto response, client.call(addr, std::move(request))); + if (!response.success) { + return makePGError(PGErrorCode::InvalidState, + std::string(coro_rpc::get_func_name()) + + " rejected: " + response.reject_reason); + } + return {}; +} + +} // namespace + +void AgentRpcServiceImpl::onPeerJoined(PeerJoinedPush push) { + host_.postPeerJoined(std::move(push)); +} + +void AgentRpcServiceImpl::onRankStateUpdate(RankStatePush push) { + host_.postRankStateUpdate(std::move(push)); +} + +void AgentRpcServiceImpl::onViewUpdate(coro_rpc::context ctx, + ViewUpdatePush push) { + host_.postViewUpdate(std::move(ctx), std::move(push)); +} + +AgentHost::AgentHost(std::string coordinator_addr, const std::string& host_ip, + GlobalRank rank, int max_world_size, + LinkManager& link_manager, + int64_t fault_reconciliation_window_us) + : agent_(rank, max_world_size), + executor_("AgentHost"), + link_manager_(link_manager), + host_ip_(host_ip), + rank_(rank), + max_world_size_(max_world_size), + coordinator_addr_(std::move(coordinator_addr)), + fault_reconciliation_window_us_(fault_reconciliation_window_us), + agent_session_id_(generateInitialAgentSessionId()), + rpc_client_(std::make_unique()) {} + +AgentHost::~AgentHost() { shutdown(); } + +void AgentHost::setFaultReconciliationWindow(int64_t timeout_us) { + fault_reconciliation_window_us_.store(timeout_us, + std::memory_order_relaxed); +} + +PGResult AgentHost::start() { + PG_VALIDATE_ARG(!coordinator_addr_.empty(), + "AgentHost coordinator address must not be empty"); + + PG_VALIDATE_STATE(!shutdown_requested_.load(std::memory_order_acquire), + "AgentHost cannot start after shutdown"); + + link_manager_.setEventCallback([this](TELinkUpEvent event) { + if (shutdown_requested_.load(std::memory_order_acquire)) return; + if (event.peer < 0 || event.peer >= max_world_size_) return; + LinkEvent link_event; + link_event.events.assign(max_world_size_, LinkEvent::EventType::None); + link_event.target_rank_epochs.assign(max_world_size_, 0); + link_event.events[event.peer] = LinkEvent::EventType::Success; + link_event.target_rank_epochs[event.peer] = event.target_rank_epoch; + pushLinkEvent(link_event); + }); + + rpc_server_ = std::make_unique(/*port=*/0, /*thread_num=*/2); + rpc_impl_ = std::make_unique(*this); + rpc_server_->registerHandler<&AgentRpcService::onPeerJoined, + &AgentRpcService::onRankStateUpdate, + &AgentRpcService::onViewUpdate>( + rpc_impl_.get()); + bool server_started = rpc_server_->start(); + if (!server_started) { + link_manager_.setEventCallback(nullptr); + rpc_impl_.reset(); + rpc_server_.reset(); + return makePGError(PGErrorCode::SystemError, + "AgentHost failed to start RPC server for rank " + + std::to_string(rank_)); + } + + executor_.setTickCallback([this]() { tick(); }); + executor_.start(); + + return executor_.post([this]() { startAgentRegistration(); }); +} + +void AgentHost::shutdown() { + if (shutdown_requested_.exchange(true, std::memory_order_acq_rel)) return; + + link_manager_.setEventCallback(nullptr); + if (rpc_server_) rpc_server_->shutdown(); + // Finish operations that callers already submitted, including explicit + // unregisterGroup calls, before releasing process-level rank ownership. + executor_.shutdown(); + link_manager_.stop(); + unregisterAgent(); + if (rpc_client_) { + rpc_client_->shutdown(); + rpc_client_.reset(); + } +} + +void AgentHost::unregisterAgent() { + if (!rpc_client_ || coordinator_addr_.empty()) return; + + const auto agent_session_id = agent_.getAgentSessionId(); + if (agent_session_id == 0) return; + + UnregisterAgentRequest req; + req.rank = rank_; + req.agent_session_id = agent_session_id; + auto result = rpc_client_->call<&CoordinatorRpcService::unregisterAgent>( + coordinator_addr_, std::move(req)); + if (!result.has_value()) { + LOG(WARNING) << "AgentHost: unregisterAgent RPC failed, rank=" << rank_ + << ": " << result.error().message; + return; + } + const auto& response = result.value(); + if (!response.success) { + LOG(WARNING) << "AgentHost: unregisterAgent rejected, rank=" << rank_ + << ": " << response.reject_reason; + } +} + +PGResult AgentHost::waitUntilRegistered( + std::chrono::milliseconds timeout) { + auto promise = std::make_shared>(); + auto future = promise->get_future(); + + PG_TRY(executor_.post([this, promise]() { + if (agent_registration_done_) { + promise->set_value(); + } else { + agent_registration_promises_.push_back(promise); + } + })); + + if (future.wait_for(timeout) != std::future_status::ready) { + // Timeout: remove the dangling promise on the executor thread. + executor_.post([this, promise]() { + std::erase(agent_registration_promises_, promise); + }); + return makePGError(PGErrorCode::Timeout, + "timed out waiting for agent registration"); + } + return {}; +} + +// Block until the group reaches Ready status (all active ranks ACKed the +// bootstrap ViewUpdate). Returns a timeout error if the deadline expires. +// +// Note: if a peer dies during the BootstrapSyncing phase, the Coordinator +// will not transition the group to Ready, and this call will hang until +// the timeout expires. The caller should handle this as a bootstrap failure. +PGResult AgentHost::waitUntilGroupReady( + GroupId group_id, std::chrono::milliseconds timeout) { + auto promise = std::make_shared>(); + auto future = promise->get_future(); + + PG_TRY(executor_.post([this, group_id, promise]() { + auto view = agent_.getGroupView(group_id); + if (view.status == GroupStatus::Ready) { + promise->set_value(view); + } else { + group_ready_promises_[group_id].push_back(promise); + } + })); + + if (future.wait_for(timeout) != std::future_status::ready) { + // Clean up the dangling promise before returning the timeout. + executor_.post([this, group_id, promise]() { + auto it = group_ready_promises_.find(group_id); + if (it != group_ready_promises_.end()) { + auto& vec = it->second; + vec.erase(std::remove(vec.begin(), vec.end(), promise), + vec.end()); + if (vec.empty()) group_ready_promises_.erase(it); + } + }); + return makePGError( + PGErrorCode::Timeout, + "waitUntilGroupReady timed out for group " + group_id); + } + return future.get(); +} + +PGResult AgentHost::waitUntilRankActive( + GroupId group_id, GlobalRank rank, std::chrono::milliseconds timeout) { + auto promise = std::make_shared>(); + auto future = promise->get_future(); + + PG_TRY(executor_.post([this, group_id, rank, promise]() { + auto view = agent_.getGroupView(group_id); + if (view.members[rank].isActive()) { + promise->set_value(); + } else { + rank_active_promises_[group_id][rank].push_back(promise); + } + })); + + if (future.wait_for(timeout) != std::future_status::ready) { + executor_.post([this, group_id, rank, promise]() { + auto it = rank_active_promises_.find(group_id); + if (it != rank_active_promises_.end()) { + auto rit = it->second.find(rank); + if (rit != it->second.end()) { + auto& vec = rit->second; + vec.erase(std::remove(vec.begin(), vec.end(), promise), + vec.end()); + if (vec.empty()) it->second.erase(rit); + } + if (it->second.empty()) rank_active_promises_.erase(it); + } + }); + return makePGError(PGErrorCode::Timeout, + "waitUntilRankActive timed out for rank " + + std::to_string(rank) + " in group " + group_id); + } + return {}; +} + +PGResult AgentHost::registerGroup( + GroupBootstrapId group_bootstrap_id, int32_t max_group_size, + std::vector rank_order, + GroupBootstrapIdResolvePolicy resolve_policy, bool auto_deactivate, + MooncakeCommunicator* communicator) { + return executor_.postAndWait( + [this, group_bootstrap_id = std::move(group_bootstrap_id), + max_group_size, rank_order = std::move(rank_order), resolve_policy, + auto_deactivate, communicator]() mutable -> PGResult { + RegisterGroupRequest req; + req.rank = rank_; + req.agent_session_id = agent_.getAgentSessionId(); + req.group_bootstrap_id = std::move(group_bootstrap_id); + req.max_group_size = max_group_size; + req.rank_order = std::move(rank_order); + req.resolve_policy = resolve_policy; + req.auto_deactivate = auto_deactivate; + + PG_TRY(auto resp, + rpc_client_->call<&CoordinatorRpcService::registerGroup>( + coordinator_addr_, std::move(req))); + + if (!resp.success) { + // A rejected group must not affect the process-scoped Agent. + // Return an empty id so this communicator can remain + // group-scoped and execute local-only collectives. + LOG(WARNING) + << "AgentHost: registerGroup rejected for rank=" << rank_ + << ": " << resp.reject_reason + << "; leaving this group out of Agent state and falling " + "back to local-only execution"; + return GroupId{}; + } + + const auto& group_id = resp.view.group_id; + communicators_.insert_or_assign(group_id, communicator); + runEffects(agent_.registerGroup(resp.view)); + return group_id; + }); +} + +void AgentHost::detachCommunicator(GroupId group_id) { + executor_.postAndWait( + [this, group_id]() { communicators_.erase(group_id); }); +} + +PGResult AgentHost::unregisterGroup(GroupId group_id) { + return executor_.postAndWait([this, group_id]() -> PGResult { + agent_.unregisterGroup(group_id); + + UnregisterGroupRequest req; + req.group_id = group_id; + req.rank = rank_; + req.agent_session_id = agent_.getAgentSessionId(); + return callAndCheck<&CoordinatorRpcService::unregisterGroup>( + *rpc_client_, coordinator_addr_, std::move(req)); + }); +} + +PGResult AgentHost::confirmReadyForActivation(GroupId group_id) { + ConfirmReadyForActivationRequest req; + req.group_id = std::move(group_id); + req.rank = rank_; + req.agent_session_id = agent_.getAgentSessionId(); + return callAndCheck<&CoordinatorRpcService::confirmReadyForActivation>( + *rpc_client_, coordinator_addr_, std::move(req)); +} + +PGResult AgentHost::sendPublishEndpointRpc( + GroupEndpointPublication endpoint) { + PublishEndpointRequest req; + req.rank = rank_; + req.agent_session_id = agent_.getAgentSessionId(); + req.endpoints.push_back(std::move(endpoint)); + return callAndCheck<&CoordinatorRpcService::publishEndpoint>( + *rpc_client_, coordinator_addr_, std::move(req)); +} + +PGResult AgentHost::publishLocalEndpoint( + GroupEndpointPublication endpoint) { + return executor_.postAndWait( + [this, endpoint = std::move(endpoint)]() mutable { + return sendPublishEndpointRpc(std::move(endpoint)); + }); +} + +void AgentHost::sendLinkEventReport(LinkEventReport report) { + if (!rpc_client_ || coordinator_addr_.empty()) return; + + const auto request_session = report.agent_session_id; + rpc_client_->callAsync<&CoordinatorRpcService::reportLinkEvent>( + coordinator_addr_, std::move(report), + [this, request_session](PGResult result) { + if (!result.has_value()) return; + auto ack = std::move(result).value(); + executor_.post([this, request_session, ack = std::move(ack)]() { + if (shutdown_requested_.load(std::memory_order_acquire)) return; + if (request_session != agent_.getAgentSessionId()) return; + agent_.handleLinkEventReportAck(ack); + }); + }); +} + +PGResult AgentHost::proposeViewUpdateInternal( + GroupId group_id, const std::vector& ranks, + bool is_activation) { + ProposeViewUpdateRequest req; + req.group_id = group_id; + req.source_rank = rank_; + req.agent_session_id = agent_.getAgentSessionId(); + req.requested_ranks = ranks; + req.is_activation = is_activation; + + const auto coordinator_timeout = + kProposalAdmissionTimeout + kViewUpdateAckTimeout; + const auto rpc_timeout = + std::max(RpcClient::kDefaultRequestTimeout, + std::chrono::duration_cast( + 2 * coordinator_timeout)); + return rpc_client_->call<&CoordinatorRpcService::proposeViewUpdate>( + coordinator_addr_, std::move(req), rpc_timeout); +} + +PGResult AgentHost::proposeActivate( + GroupId group_id, const std::vector& ranks) { + return proposeViewUpdateInternal(group_id, ranks, /*is_activation=*/true); +} + +PGResult AgentHost::proposeDeactivate( + GroupId group_id, const std::vector& ranks) { + return proposeViewUpdateInternal(group_id, ranks, /*is_activation=*/false); +} + +void AgentHost::pushLinkEvent(const LinkEvent& event) { + executor_.post( + [this, event]() { runEffects(agent_.pushLinkEvent(event)); }); +} + +PGResult AgentHost::syncAfterFailure( + GroupId group_id) { + SyncAfterFailureRequest req; + req.group_id = group_id; + + PG_TRY(executor_.postAndWait([this, &req]() { + req.reporter_rank = rank_; + req.agent_session_id = agent_.getAgentSessionId(); + req.link_event_report = agent_.getLinkEventReport(); + req.current_epoch = agent_.getGroupView(req.group_id).epoch; + })); + + // Synchronous RPC should be issued outside the executor. + // Blocking the serialized executor would stall all local state-machine + // tasks. + const auto reconciliation_window = std::chrono::microseconds( + fault_reconciliation_window_us_.load(std::memory_order_relaxed)); + const auto reconciliation_timeout = + std::chrono::ceil(reconciliation_window); + const auto rpc_timeout = + std::max(RpcClient::kDefaultRequestTimeout, 2 * reconciliation_timeout); + PG_TRY(auto response, + rpc_client_->call<&CoordinatorRpcService::syncAfterFailure>( + coordinator_addr_, req, rpc_timeout)); + + PG_TRY(executor_.postAndWait([this, request_session = req.agent_session_id, + &response]() -> PGResult { + PG_VALIDATE_STATE(request_session == agent_.getAgentSessionId(), + "agent session changed while syncing"); + + if (response.link_event_report_ack.has_value()) { + agent_.handleLinkEventReportAck(*response.link_event_report_ack); + } + + if (response.status != SyncAfterFailureStatus::Rejected) { + PG_TRY(auto effects, agent_.applyGroupView(response.view)); + runEffects(effects); + } + return {}; + })); + return response; +} + +void AgentHost::postPeerJoined(PeerJoinedPush push) { + executor_.post([this, push = std::move(push)]() { + runEffects(agent_.handlePeerJoined(push)); + }); +} + +void AgentHost::postRankStateUpdate(RankStatePush push) { + executor_.post([this, push = std::move(push)]() { + runEffects(agent_.handleRankStateUpdate(push)); + }); +} + +void AgentHost::postViewUpdate(coro_rpc::context ctx, + ViewUpdatePush push) { + auto group_id = push.view.group_id; + auto epoch = push.view.epoch; + + executor_.post([this, ctx = std::move(ctx), push = std::move(push), + group_id, epoch]() mutable { + auto apply_result = agent_.handleViewUpdate(push); + ViewUpdateAck ack{.rank = rank_, + .group_id = group_id, + .epoch = epoch, + .applied = false, + .error_msg = ""}; + if (apply_result.has_value()) { + runEffects(std::move(apply_result).value()); + ack.applied = true; + } else { + ack.error_msg = std::move(apply_result).error().message; + } + ctx.response_msg(std::move(ack)); + }); +} + +void AgentHost::startAgentRegistration(bool start_new_session) { + if (shutdown_requested_.load(std::memory_order_acquire)) return; + + // Avoid duplicate registration RPCs. This also covers the case where a + // heartbeat response callback asks for re-registration while another + // registration is already in flight. + if (agent_.getCoordinatorConnection() == + AgentStateMachine::CoordinatorConnection::AgentRegistering) { + return; + } + if (start_new_session) { + link_manager_.stop(); + ++agent_session_id_; + agent_session_initialized_ = false; + } + if (!agent_session_initialized_) { + runEffects(agent_.reset(agent_session_id_)); + agent_session_initialized_ = true; + } + + agent_.setCoordinatorConnection( + AgentStateMachine::CoordinatorConnection::AgentRegistering); + + RegisterAgentRequest req; + req.rank = rank_; + req.agent_addr = rpc_server_->getListenAddr(host_ip_); + req.te_server_name = link_manager_.localServerName(); + req.warmup_recv_addr = link_manager_.getWarmupRecvAddr(); + req.agent_session_id = agent_session_id_; + const uint64_t request_session_id = req.agent_session_id; + + rpc_client_->callAsync<&CoordinatorRpcService::registerAgent>( + coordinator_addr_, std::move(req), + [this, request_session_id](PGResult result) { + executor_.post([this, request_session_id, + result = std::move(result)]() mutable { + if (shutdown_requested_.load(std::memory_order_acquire)) return; + if (request_session_id != agent_.getAgentSessionId()) return; + + if (!result.has_value()) { + agent_.setCoordinatorConnection( + AgentStateMachine::CoordinatorConnection::Disconnected); + if (shouldLogAgentRegistrationError()) { + LOG(ERROR) << "AgentHost: registerAgent RPC failed: " + << result.error().message << "; will retry"; + } + return; + } + + auto resp = std::move(result).value(); + if (!resp.success) { + agent_.setCoordinatorConnection( + AgentStateMachine::CoordinatorConnection::Disconnected); + if (shouldLogAgentRegistrationError()) { + LOG(ERROR) << "AgentHost: registerAgent rejected: " + << resp.reject_reason << "; will retry"; + } + if (resp.require_new_session) { + startAgentRegistration(/*start_new_session=*/true); + } + return; + } + + auto effects = agent_.applyRegisterAgentResponse(resp); + runEffects(effects); + if (agent_.getCoordinatorConnection() != + AgentStateMachine::CoordinatorConnection::Connected) + return; + + link_manager_.start(agent_.getRankEpoch()); + + if (!agent_registration_done_) { + agent_registration_done_ = true; + for (auto& p : agent_registration_promises_) { + p->set_value(); + } + agent_registration_promises_.clear(); + } + + // Re-publish all local communicators' endpoints after (re-)reg. + // Old session endpoints were cleared by Coordinator. + forEachCommunicator([&](auto communicator) { + auto result = sendPublishEndpointRpc( + communicator->buildEndpointMetadata()); + if (!result.has_value()) { + LOG(ERROR) << "AgentHost: failed to re-publish " + "communicator endpoint: " + << result.error().message; + } + }); + }); + }); +} + +bool AgentHost::shouldLogAgentRegistrationError() { + const auto now = std::chrono::steady_clock::now(); + if (last_agent_register_error_log_time_.time_since_epoch() != + std::chrono::steady_clock::duration{} && + now - last_agent_register_error_log_time_ < + kAgentRegisterErrorLogInterval) { + return false; + } + last_agent_register_error_log_time_ = now; + return true; +} + +void AgentHost::tick() { + if (shutdown_requested_.load(std::memory_order_acquire)) return; + if (!rpc_client_) return; + + if (agent_.getCoordinatorConnection() == + AgentStateMachine::CoordinatorConnection::Disconnected) { + if (rpc_client_->tryReconnect(coordinator_addr_)) { + startAgentRegistration(); + } + return; + } + + if (agent_.getCoordinatorConnection() == + AgentStateMachine::CoordinatorConnection::AgentRegistering) { + return; + } + + auto now = std::chrono::steady_clock::now(); + if (now < next_heartbeat_at_) return; + next_heartbeat_at_ = now + kHeartbeatInterval; + + // Link reports are idempotent by report_id. Retry the latest unacknowledged + // snapshot with the heartbeat cadence when the request or its response is + // lost. + if (auto report = agent_.getLinkEventReport()) { + sendLinkEventReport(std::move(*report)); + } + + auto req = agent_.buildHeartbeat(); + req.agent_session_id = agent_.getAgentSessionId(); + auto request_session = req.agent_session_id; + + rpc_client_->callAsync<&CoordinatorRpcService::heartbeat>( + coordinator_addr_, std::move(req), + [this, request_session](PGResult result) { + if (!result.has_value()) return; + auto resp = std::move(result).value(); + executor_.post([this, request_session, resp]() { + if (shutdown_requested_.load(std::memory_order_acquire)) return; + if (request_session != agent_.getAgentSessionId()) return; + if (resp.require_new_session) { + // The current session is no longer valid. + startAgentRegistration(/*start_new_session=*/true); + } + }); + }); +} + +void AgentHost::runEffects(const AgentApplyResult& effects) { + for (const auto& effect : effects) { + std::visit( + overloaded{ + [this](const EnablePeerProbe& e) { + link_manager_.enablePeerProbe(e.rank, e.rank_epoch, + e.te_server_name, + e.warmup_recv_addr); + }, + [this](const DisconnectLink& e) { + link_manager_.disconnect(e.peer); + }, + [this](const RequestLinkHealthCheck& e) { + link_manager_.requestHealthCheck(e.peer); + }, + [this](const SendLinkEventReport& e) { + sendLinkEventReport(e.report); + }, + [this](const StopReconnect& e) { + link_manager_.stopReconnect(e.peer); + }, + [this](const RefreshPeerLink& e) { + link_manager_.refreshPeerSegment(e.peer); + }, + [this](const ResetPeerState& e) { + for (auto& [group_id, communicator] : communicators_) { + auto view = agent_.getGroupView(group_id); + for (int lr = 0; + lr < static_cast(view.rank_order.size()); + ++lr) { + if (view.rank_order[lr] == e.peer) { + communicator->onPeerLinkReset(lr); + break; + } + } + } + }, + [this](const NotifyLinkRefreshed& e) { + for (auto& [group_id, communicator] : communicators_) { + auto view = agent_.getGroupView(group_id); + for (int lr = 0; + lr < static_cast(view.rank_order.size()); + ++lr) { + if (view.rank_order[lr] == e.peer) { + communicator->refreshSegmentID(lr); + break; + } + } + } + }, + [this](const DisconnectAllLinks&) { + for (int i = 0; i < max_world_size_; ++i) { + if (i != rank_) { + link_manager_.disconnect(i); + } + } + }, + [this](const ClearAllPeerMetadata&) { + for (int i = 0; i < max_world_size_; ++i) { + if (i != rank_) { + link_manager_.publishLinkDown(i); + } + } + }, + [this](const ApplyViewToCommunicator& e) { + withCommunicator(e.view.group_id, [&](auto communicator) { + communicator->applyViewUpdate(e.view, e.rank_states, + e.rank_epochs, + e.activatable); + }); + }, + [this](const NotifyGroupReady& e) { + auto it = group_ready_promises_.find(e.group_id); + if (it == group_ready_promises_.end()) return; + auto view = agent_.getGroupView(e.group_id); + for (auto& p : it->second) p->set_value(view); + group_ready_promises_.erase(it); + }, + [this](const NotifyRanksActivated& e) { + auto it = rank_active_promises_.find(e.group_id); + if (it == rank_active_promises_.end()) return; + for (GlobalRank gr : e.ranks) { + auto rit = it->second.find(gr); + if (rit != it->second.end()) { + for (auto& p : rit->second) p->set_value(); + it->second.erase(rit); + } + } + if (it->second.empty()) rank_active_promises_.erase(it); + }, + }, + effect); + } +} + +} // namespace mooncake diff --git a/mooncake-pg/src/control_plane/coordinator.cpp b/mooncake-pg/src/control_plane/coordinator.cpp new file mode 100644 index 0000000000..da8a9dfbb0 --- /dev/null +++ b/mooncake-pg/src/control_plane/coordinator.cpp @@ -0,0 +1,1645 @@ +#include "control_plane/coordinator.h" + +#include +#include +#include +#include + +#include + +#include "error_types.h" +#include "pg_utils.h" + +namespace mooncake { + +CentralizedCoordinatorStateMachine::CentralizedCoordinatorStateMachine( + int max_world_size, std::chrono::microseconds fault_reconciliation_window) + : max_world_size_(max_world_size), + fault_reconciliation_window_(fault_reconciliation_window) { + PG_ASSERT(max_world_size_ > 0 && max_world_size_ <= kMaxNumRanks, + "invalid max_world_size: ", max_world_size_); + ranks_.resize(max_world_size_); + endpoint_epochs_.assign(max_world_size_, 0); + for (int r = 0; r < max_world_size_; ++r) { + ranks_[r].link_status.assign(max_world_size_, 0); + } +} + +void CentralizedCoordinatorStateMachine::setFaultReconciliationWindow( + std::chrono::microseconds fault_reconciliation_window) { + fault_reconciliation_window_ = fault_reconciliation_window; +} + +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleRegisterAgent( + const RegisterAgentRequest& req) { + CoordinatorApplyResult result; + if (!rankInRange(req.rank)) { + result.response.success = false; + result.response.reject_reason = "rank out of valid range"; + return result; + } + auto& info = ranks_[req.rank]; + + // agent_session_id is the idempotency key for a logical registration. + // Retrying an already-accepted registration must not invalidate link + // evidence, demote rank state, or rebroadcast lifecycle events. + const bool same_session = info.agent_session_id == req.agent_session_id; + if (same_session) { + if (info.state == RankState::Offline) { + result.response.success = false; + result.response.reject_reason = + "agent session is Offline; start a new registration session"; + result.response.require_new_session = true; + return result; + } + info.last_heartbeat = std::chrono::steady_clock::now(); + populateRegisterAgentResponse(result.response, req.rank); + return result; + } + + if (shutdown_requested_) { + result.response.success = false; + result.response.reject_reason = "coordinator is shutting down"; + return result; + } + + // A failed / auto-deactivated rank (Synced or Offline) may be replaced + // immediately. A different logical session may not take ownership from a + // Healthy rank. + if (info.state == RankState::Healthy) { + result.response.success = false; + result.response.reject_reason = + "rank already registered and is Healthy; replacement must wait " + "for the old process to leave the healthy set."; + return result; + } + + ++info.rank_epoch; + info.agent_addr = req.agent_addr; + info.te_server_name = req.te_server_name; + info.agent_session_id = req.agent_session_id; + info.warmup_recv_addr = req.warmup_recv_addr; + info.last_heartbeat = std::chrono::steady_clock::now(); + + // A new rank epoch invalidates both outgoing and incoming observations for + // the previous incarnation. No old edge is allowed to make the replacement + // Healthy before fresh, epoch-matched evidence arrives. + info.link_status.assign(max_world_size_, 0); + for (auto& peer : ranks_) { + peer.link_status[req.rank] = 0; + } + info.last_link_event_report_id = 0; + + for (auto& [group_id, view] : group_views_) { + auto& member = view.members[req.rank]; + // AwaitingActivation is an uncommitted promise made by the old Agent + // session, so a new rank epoch cancels it. Active membership is already + // committed and must only be changed by explicit or automatic + // deactivation paths, never by registration. + bool view_changed = false; + if (member.isAwaitingActivation()) { + member.status = GroupMemberState::Inactive; + view_changed = true; + } + + // Published endpoints belong to one rank epoch. Repeat the Offline + // reset here for replacements accepted before heartbeat timeout. + if (member.hasEndpoint()) { + member.endpoint = std::nullopt; + view_changed = true; + } + + if (view_changed) { + view.epoch++; + result.effects.push_back(PushViewUpdate{view}); + } + } + + info.state = RankState::Synced; + ++info.rank_state_version; + + result.effects.push_back(BroadcastPeerJoined{ + PeerJoinedPush{req.rank, info.rank_epoch, info.te_server_name, + info.warmup_recv_addr}}); + result.effects.push_back(makeRankStateEffect(req.rank)); + + populateRegisterAgentResponse(result.response, req.rank); + return result; +} + +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::requestShutdown() { + CoordinatorApplyResult result; + if (shutdown_requested_) return result; + + shutdown_requested_ = true; + for (GlobalRank rank = 0; rank < max_world_size_; ++rank) { + const auto& info = ranks_[rank]; + if (info.state != RankState::Offline) { + shutdown_pending_ranks_.insert(rank); + } + } + return result; +} + +void CentralizedCoordinatorStateMachine::populateRegisterAgentResponse( + RegisterAgentResponse& response, GlobalRank rank) const { + response.success = true; + response.rank_epoch = ranks_[rank].rank_epoch; + response.all_rank_states.resize(max_world_size_); + response.all_rank_epochs.resize(max_world_size_); + response.all_rank_state_versions.resize(max_world_size_); + for (int32_t i = 0; i < max_world_size_; ++i) { + response.all_rank_states[i] = ranks_[i].state; + response.all_rank_epochs[i] = ranks_[i].rank_epoch; + response.all_rank_state_versions[i] = ranks_[i].rank_state_version; + } + response.groups.reserve(group_views_.size()); + for (const auto& [group_id, view] : group_views_) { + response.groups.push_back(view); + } + response.rank_connections.reserve(max_world_size_); + for (int32_t i = 0; i < max_world_size_; ++i) { + if (i == rank || ranks_[i].state == RankState::Offline) continue; + RankConnectionMetadata connection; + connection.rank = i; + connection.rank_epoch = ranks_[i].rank_epoch; + connection.agent_addr = ranks_[i].agent_addr; + connection.te_server_name = ranks_[i].te_server_name; + connection.warmup_recv_addr = ranks_[i].warmup_recv_addr; + response.rank_connections.push_back(std::move(connection)); + } +} + +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleHeartbeat( + const HeartbeatRequest& req) { + CoordinatorApplyResult result; + if (!hasValidSession(req.rank, req.agent_session_id)) { + result.response.require_new_session = true; + return result; + } + auto& info = ranks_[req.rank]; + info.last_heartbeat = std::chrono::steady_clock::now(); + return result; +} + +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleUnregisterAgent( + const UnregisterAgentRequest& req) { + CoordinatorApplyResult result; + if (!rankInRange(req.rank)) { + result.response.reject_reason = "rank out of valid range"; + return result; + } + + auto& info = ranks_[req.rank]; + if (info.agent_session_id != req.agent_session_id) { + result.response.reject_reason = "stale agent_session_id"; + return result; + } + + // Agent lifetime is process-scoped and independent from every group. This + // RPC does not change GroupView; group lifecycle and fault handling remain + // separate operations. + if (invalidateAgentSession(req.rank)) { + result.effects.push_back(makeRankStateEffect(req.rank)); + updateRankStates(result.effects); + } + + result.response.success = true; + return result; +} + +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleRegisterGroup( + const RegisterGroupRequest& req) { + CoordinatorApplyResult result; + if (!hasValidSession(req.rank, req.agent_session_id)) { + result.response.success = false; + result.response.reject_reason = "rank out of range or stale session"; + return result; + } + + if (req.group_bootstrap_id.empty()) { + result.response.reject_reason = "group bootstrap id is empty"; + return result; + } + + if (!validateGroupRegistration(req, result.response)) { + return result; + } + + bool new_group = false; + auto group_id = resolveGroupId(req, result.response, new_group); + if (!group_id.has_value()) return result; + + processGroupRegistration(req, *group_id, result.effects); + + if (new_group) { + bindGroupBootstrapId(*group_id, req.group_bootstrap_id); + } + result.response.success = true; + result.response.view = group_views_.at(*group_id); + return result; +} + +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleConfirmReadyForActivation( + const ConfirmReadyForActivationRequest& req) { + CoordinatorApplyResult result; + if (!hasValidSession(req.rank, req.agent_session_id)) { + result.response.reject_reason = + "rank is out of range or has a stale session"; + return result; + } + + auto group_it = group_views_.find(req.group_id); + if (group_it == group_views_.end()) { + result.response.reject_reason = "group not found"; + return result; + } + + auto& view = group_it->second; + auto& member = view.members[req.rank]; + if (member.isAwaitingActivation()) { + result.response.success = true; + return result; + } + if (member.status != GroupMemberState::Inactive) { + result.response.reject_reason = "rank is not an inactive member"; + return result; + } + + member.status = GroupMemberState::AwaitingActivation; + view.epoch++; + result.effects.push_back(PushViewUpdate{view}); + result.response.success = true; + return result; +} + +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleUnregisterGroup( + const UnregisterGroupRequest& req) { + CoordinatorApplyResult result; + if (!hasValidSession(req.rank, req.agent_session_id)) { + result.response.reject_reason = + "rank is out of range, Offline, or has a stale session"; + return result; + } + + auto it = group_views_.find(req.group_id); + if (it == group_views_.end()) { + result.response.success = true; + return result; + } + + auto& view = it->second; + auto& member = view.members[req.rank]; + if (member.hasLeft()) { + result.response.success = true; + return result; + } + if (member.isNone()) { + result.response.reject_reason = "rank is not registered in the group"; + return result; + } + + member.status = GroupMemberState::Left; + member.endpoint = std::nullopt; + view.epoch++; + rejectPendingProposals(req.group_id, req.rank, "target rank left the group", + result.effects); + rejectPendingSyncs(req.group_id, req.rank, "rank left the group", + result.effects); + dropRankFromPendingBarriers(req.group_id, req.rank, result.effects); + + // Don't push a ViewUpdate when other members remain. The departing + // rank's unregister races with in-flight collectives on survivors: + // a ViewUpdate that changes activeRanks mid-collective may corrupt + // the result. + if (canEraseGroup(view)) { + eraseGroup(req.group_id, result.effects); + } + result.response.success = true; + return result; +} + +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handlePublishEndpoint( + const PublishEndpointRequest& req) { + CoordinatorApplyResult result; + if (!hasValidSession(req.rank, req.agent_session_id)) { + result.response.success = false; + result.response.reject_reason = "rank out of range or stale session"; + return result; + } + + for (const auto& ep : req.endpoints) { + auto it = group_views_.find(ep.group_id); + if (it == group_views_.end()) { + result.response.success = false; + result.response.reject_reason = "group not found"; + return result; + } + + auto& view = it->second; + auto& member = view.members[req.rank]; + member.endpoint = ep.endpoint_info; + member.endpoint->endpoint_epoch = ++endpoint_epochs_[req.rank]; + + if (member.isMember() && view.status == GroupStatus::Ready) { + view.epoch++; + result.effects.push_back(PushViewUpdate{view}); + } + } + + result.response.success = true; + checkGroupTransitions(result.effects); + return result; +} + +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleProposeViewUpdate( + uint64_t propose_id, const ProposeViewUpdateRequest& req) { + CoordinatorApplyResult result; + pending_proposals_[req.group_id].push_back(PendingProposal{ + propose_id, req, + std::chrono::steady_clock::now() + kProposalAdmissionTimeout}); + tryAdmitPendingProposals(req.group_id, result.effects); + + return result; +} + +// Update link_status from data-plane evidence. Negative transitions open the +// shared reconciliation window; the healthy-set and membership decision is +// deferred until that window closes. Positive-only transitions are applied +// immediately only when no reconciliation is already in progress. +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleLinkEventReport( + const LinkEventReport& req) { + CoordinatorApplyResult result; + if (auto ack = processLinkEventReport(req, result.effects)) { + result.response = *ack; + } + return result; +} + +void CentralizedCoordinatorStateMachine::tryOpenReconciliationWindow() { + if (!reconciliation_ctx_.active) { + reconciliation_ctx_.active = true; + reconciliation_ctx_.deadline = + std::chrono::steady_clock::now() + fault_reconciliation_window_; + reconciliation_ctx_.pending_syncs.clear(); + } +} + +void CentralizedCoordinatorStateMachine::tryCloseReconciliationWindow( + std::vector& effects) { + if (!reconciliation_ctx_.active) return; + if (std::chrono::steady_clock::now() < reconciliation_ctx_.deadline) { + return; + } + + LOG(INFO) << "[COORD] Reconciliation window expired."; + updateRankStates(effects); + applyAutoDeactivate(effects); + checkGroupTransitions(effects); + resolvePendingSyncs(effects); + + reconciliation_ctx_.active = false; +} + +std::optional +CentralizedCoordinatorStateMachine::processLinkEventReport( + const LinkEventReport& report, std::vector& effects) { + if (!hasValidSession(report.reporter_rank, report.agent_session_id)) { + return std::nullopt; + } + const auto& reporter_info = ranks_[report.reporter_rank]; + if (report.reporter_rank_epoch != reporter_info.rank_epoch) { + return std::nullopt; + } + + if (report.events.size() != static_cast(max_world_size_) || + report.target_rank_epochs.size() != + static_cast(max_world_size_)) { + LOG(WARNING) << "[COORD] invalid LinkEventReport vectors"; + return std::nullopt; + } + + LinkEventReportAck ack{report.reporter_rank, report.reporter_rank_epoch, + report.report_id}; + + auto& reporter = ranks_[report.reporter_rank]; + if (report.report_id <= reporter.last_link_event_report_id) return ack; + reporter.last_link_event_report_id = report.report_id; + + bool has_positive = false; + bool has_negative = false; + for (int32_t peer = 0; peer < max_world_size_; ++peer) { + auto type = report.events[peer]; + if (type == LinkEvent::EventType::None) continue; + + const auto& target = ranks_[peer]; + if (target.state == RankState::Offline || + report.target_rank_epochs[peer] != target.rank_epoch) { + continue; + } + + bool was_up = reporter.link_status[peer] != 0; + bool is_up = type == LinkEvent::EventType::Success; + if (was_up == is_up) continue; + + reporter.link_status[peer] = is_up ? 1 : 0; + if (is_up) { + has_positive = true; + } else { + has_negative = true; + } + } + + // Negative evidence opens a reconciliation window. Any positive changes + // in the same report are applied when the window closes. + if (has_negative) { + LOG(INFO) << "[COORD] LinkEventReport has negative -> try opening " + "reconciliation window"; + tryOpenReconciliationWindow(); + } else if (has_positive && !reconciliation_ctx_.active) { + // Positive-only changes do not need reconciliation. + updateRankStates(effects); + checkGroupTransitions(effects); + } + return ack; +} + +// handleSyncAfterFailure - sync-after-failure RPC handler. +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleSyncAfterFailure( + uint64_t sync_id, const SyncAfterFailureRequest& req) { + CoordinatorApplyResult result; + + if (!hasValidSession(req.reporter_rank, req.agent_session_id)) { + SyncAfterFailureResponse response; + response.status = SyncAfterFailureStatus::Rejected; + response.reject_reason = "rank out of range or stale session"; + result.effects.push_back(ReplySync{sync_id, response}); + return result; + } + auto view_it = group_views_.find(req.group_id); + if (view_it == group_views_.end()) { + SyncAfterFailureResponse response; + response.status = SyncAfterFailureStatus::Rejected; + response.reject_reason = "group not found"; + result.effects.push_back(ReplySync{sync_id, response}); + return result; + } + + std::optional link_event_report_ack; + + // Apply piggybacked link event report inline. + if (req.link_event_report.has_value() && + req.link_event_report->reporter_rank == req.reporter_rank && + req.link_event_report->agent_session_id == req.agent_session_id) { + link_event_report_ack = + processLinkEventReport(*req.link_event_report, result.effects); + } + + if (reconciliation_ctx_.active) { + reconciliation_ctx_.pending_syncs[req.group_id][req.reporter_rank] + .push_back(PendingSync{sync_id, req.agent_session_id, + std::move(link_event_report_ack)}); + return result; + } + + // The link state report was either already consumed by a completed window, + // or there is no pending decision. Return the current authoritative view + // and let AgentHost apply it synchronously before exposing the response. + auto response = + makeSyncResponse(SyncAfterFailureStatus::NoPending, req.group_id); + response.link_event_report_ack = std::move(link_event_report_ack); + result.effects.push_back(ReplySync{sync_id, std::move(response)}); + return result; +} + +// handleViewUpdateAck - unified ACK handler for all ViewUpdate pushes. +CoordinatorApplyResult +CentralizedCoordinatorStateMachine::handleViewUpdateAck(GroupId group_id, + GlobalRank rank, + uint64_t epoch, + bool applied) { + CoordinatorApplyResult result; + + if (!applied) return result; + + auto group_it = pending_barriers_.find(group_id); + if (group_it == pending_barriers_.end()) return result; + + auto epoch_it = group_it->second.find(epoch); + if (epoch_it == group_it->second.end()) return result; + + auto& barrier = epoch_it->second; + barrier.waiting_acks.erase(rank); + if (barrier.waiting_acks.empty()) { + auto completed = std::move(barrier); + group_it->second.erase(epoch_it); + if (group_it->second.empty()) pending_barriers_.erase(group_it); + commitBarrier(std::move(completed), result.effects); + } + + return result; +} + +CoordinatorApplyResult CentralizedCoordinatorStateMachine::tick() { + CoordinatorApplyResult result; + if (shutdown_confirmed_) return result; + + auto now = std::chrono::steady_clock::now(); + + // Heartbeat timeout + for (int rank = 0; rank < max_world_size_; ++rank) { + auto& info = ranks_[rank]; + if (info.state == RankState::Offline) continue; + if (now - info.last_heartbeat > kHeartbeatTimeout) { + handleTimedOutAgent(rank, "heartbeat timeout", result.effects); + } + } + + // Remove expired barriers first, checkGroupTransitions may create new + // bootstrap barriers and rehash this map. + std::vector expired_barriers; + for (auto group_it = pending_barriers_.begin(); + group_it != pending_barriers_.end();) { + auto& inner = group_it->second; + for (auto it = inner.begin(); it != inner.end();) { + auto& barrier = it->second; + if (!barrier.deadline.has_value() || now <= *barrier.deadline) { + ++it; + continue; + } + expired_barriers.push_back(std::move(barrier)); + it = inner.erase(it); + } + if (inner.empty()) { + group_it = pending_barriers_.erase(group_it); + } else { + ++group_it; + } + } + + for (auto& barrier : expired_barriers) { + std::vector timed_out(barrier.waiting_acks.begin(), + barrier.waiting_acks.end()); + barrier.dropped_ranks.insert(timed_out.begin(), timed_out.end()); + barrier.waiting_acks.clear(); + + // Only an ACK timeout invalidates the process-level Agent session. A + // graceful unregister is also reported as dropped by the barrier, but + // remains group-scoped (not a timed out agent). + for (GlobalRank rank : timed_out) { + handleTimedOutAgent(rank, "ViewUpdate barrier timeout", + result.effects); + } + commitBarrier(std::move(barrier), result.effects); + } + + tryCloseReconciliationWindow(result.effects); + + // A proposal may have been waiting for link readiness, rank state, or the + // preceding membership barrier. Iterate over a snapshot because admission + // removes empty queues. + std::vector pending_groups; + pending_groups.reserve(pending_proposals_.size()); + for (const auto& [group_id, _] : pending_proposals_) { + pending_groups.push_back(group_id); + } + for (const auto& group_id : pending_groups) { + tryAdmitPendingProposals(group_id, result.effects); + } + + // ShutdownCoordinatorHost must be the final state-machine effect. + // All deferred RPCs are resolved before the Host is allowed to stop + // serving requests. + tryConfirmShutdown(result.effects); + + return result; +} + +bool CentralizedCoordinatorStateMachine::invalidateAgentSession( + GlobalRank rank) { + if (ranks_[rank].state == RankState::Offline) return false; + + ranks_[rank].state = RankState::Offline; + ++ranks_[rank].rank_state_version; + ranks_[rank].link_status.assign(max_world_size_, 0); + + // Clear this rank's connectivity from all peers. + for (auto& peer : ranks_) { + if (static_cast(rank) < peer.link_status.size()) + peer.link_status[rank] = 0; + } + + if (shutdown_requested_) shutdown_pending_ranks_.erase(rank); + + return true; +} + +void CentralizedCoordinatorStateMachine::handleTimedOutAgent( + GlobalRank rank, const char* reason, + std::vector& effects) { + const auto previous_state = ranks_[rank].state; + if (!invalidateAgentSession(rank)) return; + + LOG(INFO) << "[COORD] handleTimedOutAgent rank=" << rank + << " state=" << static_cast(previous_state) + << " reason=" << reason; + + for (auto& [group_id, view] : group_views_) { + auto& member = view.members[rank]; + bool view_changed = false; + + // AwaitingActivation must be revoked in every group when the rank goes + // Offline, independently of that group's auto_deactivate policy. Active + // membership is handled below and is demoted only when auto_deactivate + // is enabled for the group. + if (member.isAwaitingActivation()) { + member.status = GroupMemberState::Inactive; + view_changed = true; + } + + // Endpoint validity is independent of collective membership. Once a + // rank is Offline, every group must discard its published endpoint and + // wait for AgentHost to publish it again after re-registration. + if (member.hasEndpoint()) { + member.endpoint = std::nullopt; + view_changed = true; + } + + if (view.auto_deactivate && member.isActive()) { + member.status = GroupMemberState::Inactive; + view_changed = true; + } + + if (view_changed) { + view.epoch++; + effects.push_back(PushViewUpdate{view}); + } + } + + effects.push_back(makeRankStateEffect(rank)); + updateRankStates(effects); + applyAutoDeactivate(effects); + checkGroupTransitions(effects); +} + +void CentralizedCoordinatorStateMachine::tryConfirmShutdown( + std::vector& effects) { + if (!shutdown_requested_ || shutdown_confirmed_ || + !shutdown_pending_ranks_.empty()) { + return; + } + + constexpr auto reason = "coordinator shutting down"; + + // A shutdown confirmation is terminal. Resolve every deferred response + // in the state machine first so ShutdownCoordinatorHost is the final + // effect ever emitted. + for (auto& group_barriers : pending_barriers_) { + for (auto& epoch_barrier : group_barriers.second) { + auto& barrier = epoch_barrier.second; + auto* commit = + std::get_if( + &barrier.commit); + if (commit == nullptr) continue; + effects.push_back(ReplyProposal{ + commit->propose_id, + {ProposalStatus::Rejected, barrier.epoch, {}, reason}}); + } + } + pending_barriers_.clear(); + + for (const auto& [group_id, proposals] : pending_proposals_) { + const auto view_it = group_views_.find(group_id); + const auto epoch = + view_it == group_views_.end() ? 0 : view_it->second.epoch; + for (const auto& proposal : proposals) { + effects.push_back( + ReplyProposal{proposal.propose_id, + {ProposalStatus::Rejected, epoch, {}, reason}}); + } + } + pending_proposals_.clear(); + + for (const auto& [group_id, ranks] : reconciliation_ctx_.pending_syncs) { + for (const auto& rank_syncs : ranks) { + const auto& pending_syncs = rank_syncs.second; + for (const auto& pending : pending_syncs) { + auto response = makeSyncResponse( + SyncAfterFailureStatus::Rejected, group_id); + response.link_event_report_ack = pending.link_event_report_ack; + response.reject_reason = reason; + effects.push_back( + ReplySync{pending.sync_id, std::move(response)}); + } + } + } + reconciliation_ctx_.pending_syncs.clear(); + reconciliation_ctx_.active = false; + + shutdown_confirmed_ = true; + effects.push_back(ShutdownCoordinatorHost{}); +} + +bool CentralizedCoordinatorStateMachine::isMutuallyConnected( + GlobalRank a, GlobalRank b) const { + PG_ASSERT(rankInRange(a) && rankInRange(b), + "isMutuallyConnected called with an out-of-range rank"); + if (ranks_[a].state == RankState::Offline || + ranks_[b].state == RankState::Offline) + return false; + return static_cast(b) < ranks_[a].link_status.size() && + static_cast(a) < ranks_[b].link_status.size() && + ranks_[a].link_status[b] != 0 && ranks_[b].link_status[a] != 0; +} + +std::vector CentralizedCoordinatorStateMachine::extendHealthySet() + const { + // Collect current Healthy ranks. + std::vector result; + for (int i = 0; i < max_world_size_; ++i) { + if (ranks_[i].state == RankState::Healthy && + isMutuallyConnected(i, i)) { + result.push_back(i); + } + } + + // Evict the least-connected rank until the set is a clique. + // (Focuses strictly on connection density; naturally terminates on + // singletons). + while (true) { + GlobalRank worst = kInvalidGlobalRank; + int worst_degree = std::numeric_limits::max(); + + for (GlobalRank r : result) { + int degree = 0; + for (GlobalRank other : result) { + if (r == other) continue; + if (isMutuallyConnected(r, other)) ++degree; + } + if (degree < worst_degree || + (degree == worst_degree && + (worst == kInvalidGlobalRank || r > worst))) { + worst_degree = degree; + worst = r; + } + } + + int expected = static_cast(result.size()) - 1; + if (worst_degree >= expected) break; + + result.erase(std::remove(result.begin(), result.end(), worst), + result.end()); + } + + // Evict isolated singletons + if (result.size() == 1) { + GlobalRank singleton = result[0]; + bool has_connections = false; + for (int other = 0; other < max_world_size_; ++other) { + if (other == singleton) continue; + if (ranks_[other].state == RankState::Offline) continue; + if (isMutuallyConnected(singleton, other)) { + has_connections = true; + break; + } + } + if (!has_connections) { + result.clear(); + } + } + + // Extend with new mutually-connected candidates. + for (int i = 0; i < max_world_size_; ++i) { + if (ranks_[i].state == RankState::Offline) continue; + // The diagonal is local data-plane readiness. A registered Agent + // remains Synced until its LinkManager reports the self-link up. + if (!isMutuallyConnected(i, i)) continue; + if (std::find(result.begin(), result.end(), i) != result.end()) + continue; + bool connected_to_all = true; + for (GlobalRank existing : result) { + if (!isMutuallyConnected(i, existing)) { + connected_to_all = false; + break; + } + } + if (connected_to_all) { + result.push_back(i); + } + } + + return result; +} + +void CentralizedCoordinatorStateMachine::updateRankStates( + std::vector& effects) { + auto healthy_set = extendHealthySet(); + + // Update per-rank Healthy / Synced state. + for (int i = 0; i < max_world_size_; ++i) { + if (ranks_[i].state == RankState::Offline) continue; + + bool in_healthy = std::find(healthy_set.begin(), healthy_set.end(), + i) != healthy_set.end(); + + if (in_healthy && ranks_[i].state != RankState::Healthy) { + ranks_[i].state = RankState::Healthy; + ++ranks_[i].rank_state_version; + effects.push_back(makeRankStateEffect(i)); + } else if (!in_healthy && ranks_[i].state == RankState::Healthy) { + ranks_[i].state = RankState::Synced; + ++ranks_[i].rank_state_version; + effects.push_back(makeRankStateEffect(i)); + } + } +} + +void CentralizedCoordinatorStateMachine::applyAutoDeactivate( + std::vector& effects) { + auto healthy_set = extendHealthySet(); + + // For auto_deactivate groups, remove unhealthy ranks from the active set. + // However, during bootstrap we do NOT do this: we wait for full mutual + // connectivity and let waitUntilGroupReady() time out if a peer is truly + // dead. + for (auto& [group_id, view] : group_views_) { + if (!view.auto_deactivate) continue; + if (view.status != GroupStatus::Ready) continue; + std::vector deactivated_ranks; + for (int i = 0; i < max_world_size_; ++i) { + if (!view.members[i].isActive()) continue; + bool in_healthy = std::find(healthy_set.begin(), healthy_set.end(), + i) != healthy_set.end(); + if (!in_healthy) { + view.members[i].status = GroupMemberState::Inactive; + view.members[i].endpoint = std::nullopt; + deactivated_ranks.push_back(i); + LOG(INFO) << "[COORD] auto_deactivate group=" << group_id + << " rank=" << i; + } + } + if (!deactivated_ranks.empty()) { + view.epoch++; + effects.push_back(PushViewUpdate{view}); + LOG(INFO) << "[COORD] auto_deactivate view update group=" + << group_id << " epoch=" << view.epoch; + } + } +} + +void CentralizedCoordinatorStateMachine::tryAdmitPendingProposals( + GroupId group_id, std::vector& effects) { + auto queue_it = pending_proposals_.find(group_id); + if (queue_it == pending_proposals_.end()) return; + + auto view_it = group_views_.find(group_id); + if (view_it == group_views_.end()) { + for (const auto& pending : queue_it->second) { + effects.push_back(ReplyProposal{ + pending.propose_id, + {ProposalStatus::Rejected, 0, {}, "group not found"}}); + } + pending_proposals_.erase(queue_it); + return; + } + + GroupView& view = view_it->second; + auto& queue = queue_it->second; + const auto now = std::chrono::steady_clock::now(); + auto reply_and_pop = [&](ProposalStatus status, const char* reason) { + effects.push_back(ReplyProposal{queue.front().propose_id, + {status, view.epoch, {}, reason}}); + queue.pop_front(); + }; + while (!queue.empty()) { + const auto& pending = queue.front(); + const auto& req = pending.request; + + if (!hasValidSession(req.source_rank, req.agent_session_id)) { + reply_and_pop(ProposalStatus::Rejected, + "source rank is Offline or has a stale session"); + continue; + } + if (view.status != GroupStatus::Ready) { + reply_and_pop(ProposalStatus::Rejected, "group is not ready"); + continue; + } + + bool has_invalid_target = false; + bool needs_membership_change = false; + + std::vector requested_global_ranks; + requested_global_ranks.reserve(req.requested_ranks.size()); + for (InGroupRank rank : req.requested_ranks) { + if (rank < 0 || + static_cast(rank) >= view.rank_order.size()) { + has_invalid_target = true; + break; + } + requested_global_ranks.push_back(view.rank_order[rank]); + } + + if (has_invalid_target) { + reply_and_pop(ProposalStatus::Rejected, + "target in-group rank is out of valid range"); + continue; + } + + for (GlobalRank rank : requested_global_ranks) { + switch (view.members[rank].status) { + case GroupMemberState::None: + case GroupMemberState::Left: + has_invalid_target = true; + break; + case GroupMemberState::Inactive: + case GroupMemberState::AwaitingActivation: + // Inactive is a valid activation target because joinGroup + // may advance it to AwaitingActivation while the proposal + // is queued. isActivatableSet still requires that + // transition before admitting the activation. + if (req.is_activation) needs_membership_change = true; + break; + case GroupMemberState::Active: + if (!req.is_activation) needs_membership_change = true; + break; + } + } + if (has_invalid_target) { + reply_and_pop(ProposalStatus::Rejected, + "target rank is not valid group member"); + continue; + } + + // A proposal updates group_views_ before its ViewUpdate barrier + // commits. While that barrier is pending, group_views_ contains the + // target membership, not a fully acknowledged membership. + auto barrier_it = pending_barriers_.find(group_id); + if (barrier_it != pending_barriers_.end() && + !barrier_it->second.empty()) { + return; + } + + if (!needs_membership_change) { + reply_and_pop(ProposalStatus::Applied, ""); + continue; + } + if (now > pending.deadline) { + reply_and_pop(ProposalStatus::Rejected, + "proposal admission timed out"); + continue; + } + + if (req.is_activation && + !isActivatableSet(group_id, requested_global_ranks, view)) { + // Keep the FIFO head pending until it times out. + return; + } + + GroupView old_view = view; + if (req.is_activation) { + for (GlobalRank rank : requested_global_ranks) { + view.members[rank].status = GroupMemberState::Active; + } + } else { + for (GlobalRank rank : requested_global_ranks) { + if (!view.members[rank].isActive()) continue; + view.members[rank].status = GroupMemberState::Inactive; + } + } + + const auto propose_id = pending.propose_id; + queue.pop_front(); + view.epoch++; + + auto required_acks = computeBarrierAckSet(old_view, view); + pending_barriers_[group_id][view.epoch] = PendingViewUpdateBarrier{ + group_id, + view.epoch, + std::move(required_acks), + {}, + now + kViewUpdateAckTimeout, + PendingViewUpdateBarrier::ProposalCommit{propose_id}}; + effects.push_back(PushViewUpdate{view}); + + if (queue.empty()) pending_proposals_.erase(queue_it); + return; + } + + pending_proposals_.erase(queue_it); +} + +void CentralizedCoordinatorStateMachine::rejectPendingProposals( + GroupId group_id, GlobalRank rank, const std::string& reason, + std::vector& effects) { + auto queue_it = pending_proposals_.find(group_id); + if (queue_it == pending_proposals_.end()) return; + + const auto view_it = group_views_.find(group_id); + if (view_it == group_views_.end()) return; + const auto& rank_order = view_it->second.rank_order; + const auto epoch = view_it->second.epoch; + + auto& queue = queue_it->second; + for (auto it = queue.begin(); it != queue.end();) { + const bool targets_rank = std::any_of( + it->request.requested_ranks.begin(), + it->request.requested_ranks.end(), [&](InGroupRank in_group_rank) { + return in_group_rank >= 0 && + static_cast(in_group_rank) < rank_order.size() && + rank_order[in_group_rank] == rank; + }); + if (!targets_rank) { + ++it; + continue; + } + effects.push_back(ReplyProposal{ + it->propose_id, {ProposalStatus::Rejected, epoch, {}, reason}}); + it = queue.erase(it); + } + + if (queue.empty()) pending_proposals_.erase(queue_it); +} + +void CentralizedCoordinatorStateMachine::dropRankFromPendingBarriers( + GroupId group_id, GlobalRank rank, + std::vector& effects) { + auto group_it = pending_barriers_.find(group_id); + if (group_it == pending_barriers_.end()) return; + + // Committing a barrier immediately retries the next queued proposal and + // may insert a new barrier for this group. So first collect completed + // barriers. + std::vector completed_barriers; + auto& barriers = group_it->second; + for (auto it = barriers.begin(); it != barriers.end();) { + auto& barrier = it->second; + if (barrier.waiting_acks.erase(rank) == 0) { + ++it; + continue; + } + + barrier.dropped_ranks.insert(rank); + if (!barrier.waiting_acks.empty()) { + ++it; + continue; + } + + completed_barriers.push_back(std::move(barrier)); + it = barriers.erase(it); + } + + if (barriers.empty()) pending_barriers_.erase(group_it); + + for (auto& barrier : completed_barriers) { + commitBarrier(std::move(barrier), effects); + } +} + +bool CentralizedCoordinatorStateMachine::isActivatableSet( + GroupId group_id, const std::vector& new_ranks, + const GroupView& old_view) const { + // Build the future active set: old active + new ranks. + std::vector future_active; + for (int i = 0; i < max_world_size_; ++i) { + if (old_view.members[i].isActive()) { + future_active.push_back(i); + } + } + for (GlobalRank r : new_ranks) { + if (!old_view.members[r].isActive()) { + future_active.push_back(r); + } + } + + // Every rank in the future set must be activatable with respect to the + // full future set. This guarantees all-to-all mutual connectivity: + // old <-> old, old <-> new, and new <-> new. + for (GlobalRank r : future_active) { + if (!isRankActivatable(group_id, r, future_active)) { + return false; + } + } + return true; +} + +bool CentralizedCoordinatorStateMachine::isRankActivatable( + GroupId group_id, GlobalRank rank, + const std::vector& future_active) const { + if (!rankInRange(rank)) { + return false; + } + if (ranks_[rank].state != RankState::Healthy) { + return false; + } + + for (GlobalRank other : future_active) { + if (other == rank) continue; + if (!isMutuallyConnected(rank, other)) { + return false; + } + } + + auto group = group_views_.find(group_id); + if (group == group_views_.end()) { + return false; + } + + const auto& member = group->second.members[rank]; + return (member.isActive() || member.isAwaitingActivation()) && + member.hasEndpoint(); +} + +void CentralizedCoordinatorStateMachine::checkGroupTransitions( + std::vector& effects) { + for (auto& [group_id, view] : group_views_) { + if (view.status == GroupStatus::Bootstrapping) { + // Collect all active ranks. + std::vector active; + bool has_any_active = false; + for (int i = 0; i < max_world_size_; ++i) { + if (!view.members[i].isActive()) continue; + has_any_active = true; + active.push_back(i); + } + + bool all_ready = true; + for (GlobalRank r : active) { + if (!isRankActivatable(group_id, r, active)) { + all_ready = false; + break; + } + } + + if (has_any_active && all_ready) { + // All active ranks have endpoints and are Healthy. + // Transition to BootstrapSyncing and initiate a barrier. + view.status = GroupStatus::BootstrapSyncing; + view.epoch++; + + auto required_acks = computeBarrierAckSet(view, view); + pending_barriers_[group_id][view.epoch] = + PendingViewUpdateBarrier{ + group_id, + view.epoch, + std::move(required_acks), + {}, + std::nullopt, + PendingViewUpdateBarrier::BootstrapCommit{}}; + + effects.push_back(PushViewUpdate{view}); + } + } + // BootstrapSyncing -> Ready is done in commitBarrier when all required + // ACKs arrive + } +} + +bool CentralizedCoordinatorStateMachine::validateGroupRegistration( + const RegisterGroupRequest& request, + RegisterGroupResponse& response) const { + if (request.max_group_size <= 0 || + request.max_group_size > max_world_size_) { + response.success = false; + response.reject_reason = "max_group_size is out of valid range"; + return false; + } + + if (request.rank_order.size() > + static_cast(request.max_group_size)) { + response.success = false; + response.reject_reason = "rank_order exceeds max_group_size"; + return false; + } + + // Validate joining_rank + if (!rankInRange(request.rank)) { + response.success = false; + response.reject_reason = "joining rank is out of valid range"; + return false; + } + + // Validate rank_order elements. + for (GlobalRank r : request.rank_order) { + if (!rankInRange(r)) { + response.success = false; + response.reject_reason = "rank_order contains invalid GlobalRank"; + return false; + } + } + + // Validate no duplicates in rank_order. + { + std::set seen(request.rank_order.begin(), + request.rank_order.end()); + if (seen.size() != request.rank_order.size()) { + response.success = false; + response.reject_reason = "rank_order contains duplicate ranks"; + return false; + } + } + + // The joining rank must be one of the ranks it declares in rank_order. + if (std::find(request.rank_order.begin(), request.rank_order.end(), + request.rank) == request.rank_order.end()) { + response.success = false; + response.reject_reason = "joining rank not in rank_order"; + return false; + } + + return true; +} + +static bool isRankOrderPrefix(const std::vector& prefix, + const std::vector& order) { + return prefix.size() <= order.size() && + std::equal(prefix.begin(), prefix.end(), order.begin()); +} + +std::optional CentralizedCoordinatorStateMachine::resolveGroupId( + const RegisterGroupRequest& request, RegisterGroupResponse& response, + bool& new_group) { + // GroupBootstrapId identifies a PyTorch group_id, not necessarily one + // runtime group. resolve_policy supplies the choice that cannot be inferred + // from rank-order relationships alone. Within one bootstrap-id bucket: + // + // * CreateOrAttach never modifies an existing rank order. Without an exact + // match, it creates a new runtime group even if rank orders overlap. + // * AttachOrExtend never creates a runtime group. Without an exact match, + // it must find one unique existing order that is a proper prefix of the + // request, has matching capacity, and appends the joining rank. + // + // More than one exact match is ambiguous. If there is no exact match, more + // than one append-compatible match is ambiguous. Callers that need to + // distinguish them must eventually provide distinct stable + // GroupBootstrapIds. processGroupRegistration() then applies the resolved + // registration. + static constexpr auto GROUP_ID_PREFIX = "mooncake_pg_"; + auto bucket = group_ids_by_bootstrap_id_.find(request.group_bootstrap_id); + + // Exact matching (Attach) is shared by both policies and always takes + // precedence. + std::vector exact_groups; + if (bucket != group_ids_by_bootstrap_id_.end()) { + for (const auto& group_id : bucket->second) { + const auto& view = group_views_.at(group_id); + if (view.max_group_size == request.max_group_size && + view.rank_order == request.rank_order) { + exact_groups.push_back(group_id); + } + } + } + + if (exact_groups.size() > 1) { + response.reject_reason = "ambiguous exact group matches"; + return std::nullopt; + } + + if (exact_groups.size() == 1) { + new_group = false; + return exact_groups.front(); + } + + switch (request.resolve_policy) { + case GroupBootstrapIdResolvePolicy::CreateOrAttach: + new_group = true; + return GROUP_ID_PREFIX + std::to_string(next_group_id_++); + + case GroupBootstrapIdResolvePolicy::AttachOrExtend: { + if (bucket == group_ids_by_bootstrap_id_.end()) { + response.reject_reason = "extension target not found"; + return std::nullopt; + } + + std::vector extension_groups; + for (const auto& group_id : bucket->second) { + const auto& view = group_views_.at(group_id); + const auto& existing_order = view.rank_order; + + if (view.max_group_size != request.max_group_size) { + continue; + } + if (existing_order.size() >= request.rank_order.size()) { + continue; + } + if (!isRankOrderPrefix(existing_order, request.rank_order)) { + continue; + } + + auto appended_begin = + request.rank_order.begin() + existing_order.size(); + if (std::find(appended_begin, request.rank_order.end(), + request.rank) == request.rank_order.end()) { + continue; + } + + extension_groups.push_back(group_id); + } + + if (extension_groups.size() > 1) { + response.reject_reason = "ambiguous extension target"; + return std::nullopt; + } + + if (extension_groups.empty()) { + response.reject_reason = + "no append-compatible extension target"; + return std::nullopt; + } + + new_group = false; + return extension_groups.front(); + } + + default: + response.reject_reason = + "invalid group bootstrap id resolve policy"; + return std::nullopt; + } +} + +void CentralizedCoordinatorStateMachine::bindGroupBootstrapId( + GroupId group_id, GroupBootstrapId group_bootstrap_id) { + group_ids_by_bootstrap_id_[group_bootstrap_id].push_back(group_id); + group_bootstrap_ids_.emplace(std::move(group_id), + std::move(group_bootstrap_id)); +} + +void CentralizedCoordinatorStateMachine::processGroupRegistration( + const RegisterGroupRequest& request, const GroupId& group_id, + std::vector& effects) { + auto it = group_views_.find(group_id); + if (it == group_views_.end()) { + // First declaration -> create group. + // Founding members are all entries in rank_order. + GroupView view; + view.group_id = group_id; + view.max_group_size = request.max_group_size; + view.rank_order = request.rank_order; + view.members.resize(max_world_size_); + for (GlobalRank r : request.rank_order) { + view.members[r].status = GroupMemberState::Active; + } + view.status = GroupStatus::Bootstrapping; + group_views_[group_id] = std::move(view); + group_views_[group_id].auto_deactivate = request.auto_deactivate; + return; + } + + auto& view = it->second; + + // If the request rank order is longer, the extra ranks are + // not activated here. They must be activated via a subsequent + // proposeViewUpdate (activate_rank / recover_ranks) from an existing + // active member. + // + // However, extend the existing rank_order with the new ranks now so that + // every member's ViewUpdate carries the correct rank_order (local->global + // mapping). + bool view_changed = false; + if (request.rank_order.size() > view.rank_order.size()) { + view.rank_order = request.rank_order; + view_changed = true; + } + + auto& joining_member = view.members[request.rank]; + if (joining_member.status == GroupMemberState::None || + joining_member.status == GroupMemberState::Left) { + joining_member.status = GroupMemberState::Inactive; + joining_member.endpoint = std::nullopt; + view_changed = true; + } + + // A Ready group that receives a registerGroup should push the authoritative + // Ready view to all members (including the newly-joined inactive rank) so + // that joining ranks can observe Ready and unblock waitUntilGroupReady(). + if (view.status == GroupStatus::Ready) { + // A changed payload must never be published under an epoch that agents + // may already have applied. Repeated registrations with no view change + // remain idempotent and reuse the current epoch. + if (view_changed) view.epoch++; + effects.push_back(PushViewUpdate{view}); + } +} + +// Private: helpers + +bool CentralizedCoordinatorStateMachine::canEraseGroup( + const GroupView& view) const { + return std::all_of(view.members.begin(), view.members.end(), + [](const GroupMember& m) { + return m.status == GroupMemberState::None || + m.status == GroupMemberState::Left; + }); +} + +void CentralizedCoordinatorStateMachine::eraseGroup( + GroupId group_id, std::vector& effects) { + // Erase any pending ViewUpdate barriers for this group so replies are not + // sent after the group is gone. + auto it = pending_barriers_.find(group_id); + if (it != pending_barriers_.end()) { + for (auto& [epoch, barrier] : it->second) { + if (auto* pc = + std::get_if( + &barrier.commit)) { + effects.push_back(ReplyProposal{ + pc->propose_id, + {ProposalStatus::Rejected, 0, {}, "group was destroyed"}}); + } + // Bootstrap barriers need no reply. + } + pending_barriers_.erase(it); + } + + auto proposal_it = pending_proposals_.find(group_id); + if (proposal_it != pending_proposals_.end()) { + for (const auto& pending : proposal_it->second) { + effects.push_back(ReplyProposal{ + pending.propose_id, + {ProposalStatus::Rejected, 0, {}, "group was destroyed"}}); + } + pending_proposals_.erase(proposal_it); + } + + auto sync_group_it = reconciliation_ctx_.pending_syncs.find(group_id); + if (sync_group_it != reconciliation_ctx_.pending_syncs.end()) { + for (const auto& rank_syncs : sync_group_it->second) { + for (const auto& pending : rank_syncs.second) { + auto response = makeSyncResponse( + SyncAfterFailureStatus::Rejected, group_id); + response.link_event_report_ack = pending.link_event_report_ack; + response.reject_reason = "group was destroyed"; + effects.push_back( + ReplySync{pending.sync_id, std::move(response)}); + } + } + reconciliation_ctx_.pending_syncs.erase(sync_group_it); + } + + auto group_bootstrap_id = group_bootstrap_ids_.at(group_id); + auto& group_ids = group_ids_by_bootstrap_id_.at(group_bootstrap_id); + group_ids.erase(std::remove(group_ids.begin(), group_ids.end(), group_id), + group_ids.end()); + if (group_ids.empty()) { + group_ids_by_bootstrap_id_.erase(group_bootstrap_id); + } + group_bootstrap_ids_.erase(group_id); + group_views_.erase(group_id); +} + +// Effect factories + +CoordinatorEffect CentralizedCoordinatorStateMachine::makeRankStateEffect( + GlobalRank rank) { + return BroadcastRankState{RankStatePush{rank, ranks_[rank].rank_epoch, + ranks_[rank].rank_state_version, + ranks_[rank].state}}; +} + +void CentralizedCoordinatorStateMachine::commitBarrier( + PendingViewUpdateBarrier barrier, std::vector& effects) { + std::visit( + overloaded{ + [&](const PendingViewUpdateBarrier::ProposalCommit& commit) { + ProposeViewUpdateResponse response{ + ProposalStatus::Applied, barrier.epoch, {}, ""}; + if (!barrier.dropped_ranks.empty()) { + response.status = ProposalStatus::AppliedWithDroppedRanks; + const auto group_it = group_views_.find(barrier.group_id); + if (group_it == group_views_.end()) { + LOG(ERROR) << "[COORD] cannot map dropped ranks for " + << "missing group " << barrier.group_id; + response.status = ProposalStatus::Rejected; + response.reject_reason = + "group disappeared while committing ViewUpdate"; + } else { + const auto& rank_order = group_it->second.rank_order; + response.dropped_ranks.reserve( + barrier.dropped_ranks.size()); + for (GlobalRank rank : barrier.dropped_ranks) { + const auto rank_it = std::find( + rank_order.begin(), rank_order.end(), rank); + if (rank_it == rank_order.end()) { + LOG(ERROR) + << "[COORD] dropped GlobalRank " << rank + << " is not in group " << barrier.group_id; + response.status = ProposalStatus::Rejected; + response.dropped_ranks.clear(); + response.reject_reason = + "dropped rank is not in the group"; + break; + } + response.dropped_ranks.push_back( + static_cast(std::distance( + rank_order.begin(), rank_it))); + } + std::sort(response.dropped_ranks.begin(), + response.dropped_ranks.end()); + } + } + effects.push_back(ReplyProposal{commit.propose_id, response}); + }, + [&](const PendingViewUpdateBarrier::BootstrapCommit&) { + auto it = group_views_.find(barrier.group_id); + if (it == group_views_.end()) return; + GroupView& view = it->second; + view.status = GroupStatus::Ready; + view.epoch++; + effects.push_back(PushViewUpdate{view}); + }, + }, + barrier.commit); + + // Barrier completion releases this group's proposal admission lane. + // Retry immediately on every completion path (ACK, timeout, or graceful + // unregister) instead of waiting for the next coordinator tick. + tryAdmitPendingProposals(barrier.group_id, effects); +} + +void CentralizedCoordinatorStateMachine::rejectPendingSyncs( + GroupId group_id, GlobalRank rank, const std::string& reason, + std::vector& effects) { + auto& pending_syncs = reconciliation_ctx_.pending_syncs; + auto group_it = pending_syncs.find(group_id); + if (group_it == pending_syncs.end()) return; + + auto rank_it = group_it->second.find(rank); + if (rank_it == group_it->second.end()) return; + + for (const PendingSync& pending : rank_it->second) { + auto resp = + makeSyncResponse(SyncAfterFailureStatus::Rejected, group_id); + resp.link_event_report_ack = pending.link_event_report_ack; + resp.reject_reason = reason; + effects.push_back(ReplySync{pending.sync_id, std::move(resp)}); + } + group_it->second.erase(rank_it); + if (group_it->second.empty()) { + pending_syncs.erase(group_it); + } +} + +SyncAfterFailureResponse CentralizedCoordinatorStateMachine::makeSyncResponse( + SyncAfterFailureStatus status, GroupId group_id) const { + SyncAfterFailureResponse response; + response.status = status; + + // piggybacked view update + if (status != SyncAfterFailureStatus::Rejected) { + auto view_it = group_views_.find(group_id); + if (view_it != group_views_.end()) { + response.view = view_it->second; + } + } + return response; +} + +void CentralizedCoordinatorStateMachine::resolvePendingSyncs( + std::vector& effects) { + for (auto& [group_id, ranks] : reconciliation_ctx_.pending_syncs) { + for (auto& [rank, pending_requests] : ranks) { + for (const PendingSync& pending : pending_requests) { + auto status = hasValidSession(rank, pending.agent_session_id) + ? SyncAfterFailureStatus::Reconciled + : SyncAfterFailureStatus::Rejected; + auto response = makeSyncResponse(status, group_id); + response.link_event_report_ack = pending.link_event_report_ack; + if (status == SyncAfterFailureStatus::Rejected) { + response.reject_reason = "stale agent session"; + } + effects.push_back( + ReplySync{pending.sync_id, std::move(response)}); + } + } + } + reconciliation_ctx_.pending_syncs.clear(); +} + +// computeBarrierAckSet -- ranks that must ACK before a proposal/bootstrap +// barrier can commit. Includes all online ranks active in either old or new +// view. +std::unordered_set +CentralizedCoordinatorStateMachine::computeBarrierAckSet( + const GroupView& old_view, const GroupView& new_view) const { + std::unordered_set acks; + for (int i = 0; i < max_world_size_; ++i) { + if (ranks_[i].state == RankState::Offline) continue; + if (old_view.members[i].isActive() || new_view.members[i].isActive()) { + acks.insert(i); + } + } + return acks; +} + +} // namespace mooncake diff --git a/mooncake-pg/src/control_plane/coordinator_host.cpp b/mooncake-pg/src/control_plane/coordinator_host.cpp new file mode 100644 index 0000000000..681bb5acf4 --- /dev/null +++ b/mooncake-pg/src/control_plane/coordinator_host.cpp @@ -0,0 +1,339 @@ +#include "control_plane/coordinator_host.h" + +#include +#include + +#include "control_plane/rpc.h" +#include "control_plane/rpc_runtime.h" +#include "pg_utils.h" + +namespace mooncake { + +void CoordinatorRpcServiceImpl::registerAgent( + coro_rpc::context ctx, RegisterAgentRequest req) { + host_.postRegisterAgent(std::move(ctx), std::move(req)); +} + +void CoordinatorRpcServiceImpl::heartbeat( + coro_rpc::context ctx, HeartbeatRequest req) { + host_.postHeartbeat(std::move(ctx), std::move(req)); +} + +void CoordinatorRpcServiceImpl::unregisterAgent( + coro_rpc::context ctx, + UnregisterAgentRequest req) { + host_.postUnregisterAgent(std::move(ctx), std::move(req)); +} + +void CoordinatorRpcServiceImpl::registerGroup( + coro_rpc::context ctx, RegisterGroupRequest req) { + host_.postRegisterGroup(std::move(ctx), std::move(req)); +} + +void CoordinatorRpcServiceImpl::unregisterGroup( + coro_rpc::context ctx, + UnregisterGroupRequest req) { + host_.postUnregisterGroup(std::move(ctx), std::move(req)); +} + +void CoordinatorRpcServiceImpl::confirmReadyForActivation( + coro_rpc::context ctx, + ConfirmReadyForActivationRequest req) { + host_.postConfirmReadyForActivation(std::move(ctx), std::move(req)); +} + +void CoordinatorRpcServiceImpl::proposeViewUpdate( + coro_rpc::context ctx, + ProposeViewUpdateRequest req) { + host_.postProposeViewUpdate(std::move(ctx), std::move(req)); +} +void CoordinatorRpcServiceImpl::publishEndpoint( + coro_rpc::context ctx, + PublishEndpointRequest req) { + host_.postPublishEndpoint(std::move(ctx), std::move(req)); +} + +void CoordinatorRpcServiceImpl::reportLinkEvent( + coro_rpc::context ctx, LinkEventReport req) { + host_.postLinkEventReport(std::move(ctx), std::move(req)); +} + +void CoordinatorRpcServiceImpl::syncAfterFailure( + coro_rpc::context ctx, + SyncAfterFailureRequest req) { + host_.postSyncAfterFailure(std::move(ctx), std::move(req)); +} + +CoordinatorHost::CoordinatorHost(const std::string& host_ip, int max_world_size, + int64_t fault_reconciliation_window_us) + : state_machine_(max_world_size, + std::chrono::microseconds(fault_reconciliation_window_us)), + executor_("CoordinatorHost"), + host_ip_(host_ip), + max_world_size_(max_world_size), + rpc_client_(std::make_unique()) {} + +CoordinatorHost::~CoordinatorHost() { shutdown(); } + +PGResult CoordinatorHost::setFaultReconciliationWindow( + int64_t timeout_us) { + return executor_.postAndWait([this, timeout_us] { + state_machine_.setFaultReconciliationWindow( + std::chrono::microseconds(timeout_us)); + }); +} + +PGResult CoordinatorHost::start() { + PG_VALIDATE_STATE(!shutdown_requested_.load(std::memory_order_acquire), + "CoordinatorHost cannot start after shutdown"); + + rpc_server_ = std::make_unique(/*port=*/0, /*thread_num=*/2); + rpc_impl_ = std::make_unique(*this); + rpc_server_ + ->registerHandler<&CoordinatorRpcService::registerAgent, + &CoordinatorRpcService::heartbeat, + &CoordinatorRpcService::unregisterAgent, + &CoordinatorRpcService::registerGroup, + &CoordinatorRpcService::unregisterGroup, + &CoordinatorRpcService::confirmReadyForActivation, + &CoordinatorRpcService::proposeViewUpdate, + &CoordinatorRpcService::publishEndpoint, + &CoordinatorRpcService::reportLinkEvent, + &CoordinatorRpcService::syncAfterFailure>( + rpc_impl_.get()); + + bool server_started = rpc_server_->start(); + if (!server_started) { + rpc_impl_.reset(); + rpc_server_.reset(); + return makePGError(PGErrorCode::SystemError, + "CoordinatorHost failed to start RPC server"); + } + + listen_addr_ = rpc_server_->getListenAddr(host_ip_); + + executor_.setTickCallback([this]() { + auto result = state_machine_.tick(); + runEffects(result.effects); + }); + + executor_.start(); + return {}; +} + +void CoordinatorHost::shutdown() { + if (shutdown_requested_.exchange(true, std::memory_order_acq_rel)) return; + + if (rpc_server_) { + auto shutdown_confirmation = shutdown_confirmation_.get_future(); + auto post_result = executor_.postAndWait([this]() { + auto result = state_machine_.requestShutdown(); + runEffects(result.effects); + }); + + if (!post_result.has_value()) { + LOG(WARNING) << "[COORD] failed to request shutdown: " + << post_result.error().message; + } else if (shutdown_confirmation.wait_for(kShutdownDrainTimeout) != + std::future_status::ready) { + LOG(WARNING) << "[COORD] shutdown drain timed out"; + } + rpc_server_->shutdown(); + } + + // Keep the executor alive while outbound callbacks finish; callbacks may + // still post their final state-machine work during client draining. + if (rpc_client_) rpc_client_->shutdown(); + executor_.shutdown(); +} + +void CoordinatorHost::postRegisterAgent( + coro_rpc::context ctx, RegisterAgentRequest req) { + executor_.post( + [this, ctx = std::move(ctx), req = std::move(req)]() mutable { + auto r = state_machine_.handleRegisterAgent(req); + runEffects(r.effects); + ctx.response_msg(std::move(r.response)); + }); +} + +void CoordinatorHost::postHeartbeat(coro_rpc::context ctx, + HeartbeatRequest req) { + executor_.post( + [this, ctx = std::move(ctx), req = std::move(req)]() mutable { + auto result = state_machine_.handleHeartbeat(req); + runEffects(result.effects); + ctx.response_msg(std::move(result.response)); + }); +} + +void CoordinatorHost::postUnregisterAgent( + coro_rpc::context ctx, + UnregisterAgentRequest req) { + executor_.post( + [this, ctx = std::move(ctx), req = std::move(req)]() mutable { + auto result = state_machine_.handleUnregisterAgent(req); + runEffects(result.effects); + ctx.response_msg(std::move(result.response)); + }); +} + +void CoordinatorHost::postRegisterGroup( + coro_rpc::context ctx, RegisterGroupRequest req) { + executor_.post( + [this, ctx = std::move(ctx), req = std::move(req)]() mutable { + auto result = state_machine_.handleRegisterGroup(req); + runEffects(result.effects); + ctx.response_msg(std::move(result.response)); + }); +} + +void CoordinatorHost::postUnregisterGroup( + coro_rpc::context ctx, + UnregisterGroupRequest req) { + executor_.post( + [this, ctx = std::move(ctx), req = std::move(req)]() mutable { + auto result = state_machine_.handleUnregisterGroup(req); + runEffects(result.effects); + ctx.response_msg(std::move(result.response)); + }); +} + +void CoordinatorHost::postConfirmReadyForActivation( + coro_rpc::context ctx, + ConfirmReadyForActivationRequest req) { + executor_.post( + [this, ctx = std::move(ctx), req = std::move(req)]() mutable { + auto result = state_machine_.handleConfirmReadyForActivation(req); + runEffects(result.effects); + ctx.response_msg(std::move(result.response)); + }); +} + +void CoordinatorHost::postProposeViewUpdate( + coro_rpc::context ctx, + ProposeViewUpdateRequest req) { + executor_.post([this, ctx = std::move(ctx), + req = std::move(req)]() mutable { + uint64_t propose_id = next_propose_id_++; + pending_proposal_resps_.emplace(propose_id, std::move(ctx)); + auto result = state_machine_.handleProposeViewUpdate(propose_id, req); + runEffects(result.effects); + }); +} + +void CoordinatorHost::postPublishEndpoint( + coro_rpc::context ctx, + PublishEndpointRequest req) { + executor_.post( + [this, ctx = std::move(ctx), req = std::move(req)]() mutable { + auto result = state_machine_.handlePublishEndpoint(req); + runEffects(result.effects); + ctx.response_msg(std::move(result.response)); + }); +} + +void CoordinatorHost::postLinkEventReport( + coro_rpc::context ctx, LinkEventReport req) { + executor_.post( + [this, ctx = std::move(ctx), req = std::move(req)]() mutable { + auto result = state_machine_.handleLinkEventReport(req); + runEffects(result.effects); + ctx.response_msg(std::move(result.response)); + }); +} + +void CoordinatorHost::postSyncAfterFailure( + coro_rpc::context ctx, + SyncAfterFailureRequest req) { + executor_.post( + [this, ctx = std::move(ctx), req = std::move(req)]() mutable { + uint64_t sync_id = next_sync_id_++; + pending_sync_resps_.emplace(sync_id, std::move(ctx)); + auto result = state_machine_.handleSyncAfterFailure(sync_id, req); + runEffects(result.effects); + }); +} + +void CoordinatorHost::postViewUpdateAck(GroupId group_id, GlobalRank rank, + uint64_t epoch, bool applied) { + executor_.post([this, group_id, rank, epoch, applied]() { + auto result = + state_machine_.handleViewUpdateAck(group_id, rank, epoch, applied); + runEffects(result.effects); + }); +} + +void CoordinatorHost::runEffects( + const std::vector& effects) { + for (const auto& effect : effects) { + std::visit( + overloaded{ + [this](const BroadcastRankState& e) { + for (int i = 0; i < max_world_size_; ++i) { + if (state_machine_.getRankState(i) != + RankState::Offline) { + pushToAgent<&AgentRpcService::onRankStateUpdate>( + i, e.push); + } + } + }, + [this](const PushViewUpdate& e) { pushViewUpdate(e); }, + [this](const ReplyProposal& e) { + auto it = pending_proposal_resps_.find(e.propose_id); + if (it != pending_proposal_resps_.end()) { + it->second.response_msg(e.response); + pending_proposal_resps_.erase(it); + } + }, + [this](const ReplySync& e) { + auto it = pending_sync_resps_.find(e.sync_id); + if (it != pending_sync_resps_.end()) { + it->second.response_msg(e.response); + pending_sync_resps_.erase(it); + } + }, + [this](const BroadcastPeerJoined& e) { + for (int i = 0; i < max_world_size_; ++i) { + if (i != e.push.rank && state_machine_.getRankState( + i) != RankState::Offline) { + pushToAgent<&AgentRpcService::onPeerJoined>(i, + e.push); + } + } + }, + [this](const ShutdownCoordinatorHost&) { + shutdown_confirmation_.set_value(); + }, + }, + effect); + } +} + +void CoordinatorHost::pushViewUpdate(const PushViewUpdate& effect) { + ViewUpdatePush push{effect.view}; + auto group_id = effect.view.group_id; + + for (int32_t i = 0; i < max_world_size_; ++i) { + const auto& member = effect.view.members[i]; + if (member.status == GroupMemberState::None || + member.status == GroupMemberState::Left) { + continue; + } + + const auto& addr = state_machine_.getAgentAddr(i); + if (state_machine_.getRankState(i) == RankState::Offline || + addr.empty()) + continue; + + rpc_client_->callAsync<&AgentRpcService::onViewUpdate>( + addr, push, + [this, group_id, rank = i](PGResult result) { + if (!result.has_value()) return; + auto ack = std::move(result).value(); + postViewUpdateAck(group_id, rank, ack.epoch, ack.applied); + }); + } +} + +} // namespace mooncake diff --git a/mooncake-pg/src/control_plane/link_manager.cpp b/mooncake-pg/src/control_plane/link_manager.cpp new file mode 100644 index 0000000000..2c332bfa2a --- /dev/null +++ b/mooncake-pg/src/control_plane/link_manager.cpp @@ -0,0 +1,555 @@ +#include "control_plane/link_manager.h" + +#include + +#ifndef USE_MUSA +#include +#include +#endif + +#include +#include + +#include "memory_location.h" +#include "pg_utils.h" + +namespace mooncake { + +#ifndef USE_MUSA +static bool checkSupportFabricMem() { + const char* nvlink_ipc = getenv("MC_USE_NVLINK_IPC"); + bool fabric_enabled = nvlink_ipc && strcmp(nvlink_ipc, "0") == 0; + if (!fabric_enabled) return false; + + int num_devices = 0; + cudaError_t err = cudaGetDeviceCount(&num_devices); + if (err != cudaSuccess || num_devices == 0) return false; + + for (int dev = 0; dev < num_devices; ++dev) { + int supported = 0; + cuDeviceGetAttribute( + &supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, dev); + if (!supported) return false; + } + return true; +} +#else +static bool checkSupportFabricMem() { return false; } +#endif + +bool LinkManager::supportFabricMem() { + static bool cached = checkSupportFabricMem(); + return cached; +} + +PGResult LinkManager::init(GlobalRank rank, int max_world_size, + TransferEngine* engine) { + PG_VALIDATE_STATE(!shutdown_.load(std::memory_order_acquire), + "LinkManager cannot be initialized after shutdown"); + if (initialized_.load(std::memory_order_acquire)) return {}; + PG_VALIDATE_ARG(max_world_size > 0 && max_world_size <= kMaxNumRanks, + "LinkManager world size is outside the supported range"); + + rank_ = rank; + max_world_size_ = max_world_size; + engine_ = engine; + local_server_name_ = engine_->getLocalIpAndPort(); + skip_warmup_ = supportFabricMem(); + peers_.resize(max_world_size_); + read_state_ = std::vector(max_world_size_); + + if (!skip_warmup_) { + auto warmup_send_region = std::make_unique(1); + auto warmup_recv_region = std::make_unique(max_world_size_); + std::memset(warmup_recv_region.get(), 0, + max_world_size_ * sizeof(int32_t)); + + PG_TRY_TE(engine_->registerLocalMemory( + warmup_send_region.get(), sizeof(int32_t), kWildcardLocation)); + + const int rc = engine_->registerLocalMemory( + warmup_recv_region.get(), max_world_size_ * sizeof(int32_t), + kWildcardLocation); + if (rc != 0) engine_->unregisterLocalMemory(warmup_send_region.get()); + PG_TRY_TE(rc); + + warmup_send_region_ = std::move(warmup_send_region); + warmup_recv_region_ = std::move(warmup_recv_region); + } + + initialized_.store(true, std::memory_order_release); + return {}; +} + +void LinkManager::start(uint64_t self_rank_epoch) { + if (!initialized_.load(std::memory_order_acquire)) { + LOG(ERROR) << "LinkManager: start() called before init()"; + return; + } + if (shutdown_.load(std::memory_order_acquire)) return; + if (started_.exchange(true, std::memory_order_acq_rel)) return; + + TransferMetadata::SegmentID self_target_id{}; + { + std::lock_guard lock(peers_mutex_); + auto& link = peers_[rank_]; + link.state = PeerLinkState::Connected; + link.is_candidate = false; + link.target_rank_epoch = self_rank_epoch; + link.target_id = engine_->openSegment(local_server_name_); + self_target_id = link.target_id.value(); + } + + publishLinkUp(rank_, self_target_id, self_rank_epoch); + + poller_running_.store(true, std::memory_order_release); + poller_thread_ = std::thread([this] { pollerLoop(); }); +} + +void LinkManager::stop() { + if (!initialized_.load(std::memory_order_acquire)) return; + + started_.store(false, std::memory_order_release); + poller_running_.store(false, std::memory_order_release); + wakeup(); + if (poller_thread_.joinable()) { + poller_thread_.join(); + } + + std::lock_guard lock(peers_mutex_); + for (int peer = 0; peer < max_world_size_; ++peer) { + if (peer == rank_) continue; + auto& link = peers_[peer]; + if (link.target_id.has_value()) { + engine_->closeSegment(link.target_id.value()); + link.target_id = std::nullopt; + } + if (link.probe_batch_id.has_value()) { + engine_->freeBatchID(link.probe_batch_id.value()); + link.probe_batch_id = std::nullopt; + } + // Segment metadata belongs to the TransferEngine. removeLocalSegment + // is only valid while refreshing a live peer, not during shutdown. + link.health_check_requested = false; + link.state = PeerLinkState::Idle; + link.is_candidate = false; + } + if (warmup_recv_region_) { + std::memset(warmup_recv_region_.get(), 0, + max_world_size_ * sizeof(int32_t)); + } +} + +void LinkManager::shutdown() { + if (shutdown_.exchange(true, std::memory_order_acq_rel)) return; + + stop(); + + // Release warmup regions. + if (warmup_send_region_) { + engine_->unregisterLocalMemory(warmup_send_region_.get()); + warmup_send_region_.reset(); + } + if (warmup_recv_region_) { + engine_->unregisterLocalMemory(warmup_recv_region_.get()); + warmup_recv_region_.reset(); + } +} + +std::string LinkManager::localServerName() const { return local_server_name_; } + +uint64_t LinkManager::getWarmupRecvAddr() const { + if (skip_warmup_ || !warmup_recv_region_) return 0; + return reinterpret_cast(warmup_recv_region_.get()); +} + +void LinkManager::enablePeerProbe(GlobalRank peer, uint64_t target_rank_epoch, + const std::string& server_name, + uint64_t warmup_recv_addr) { + if (peer == rank_) return; + if (!rankInRange(peer)) return; + + std::lock_guard lock(peers_mutex_); + auto& link = peers_[peer]; + + // Peer lifecycle pushes can arrive out of order. Never let metadata for + // an older rank incarnation replace the current probe target. + if (target_rank_epoch < link.target_rank_epoch) return; + if (target_rank_epoch > link.target_rank_epoch) { + tearDownPeerLink(peer); + link.target_rank_epoch = target_rank_epoch; + } + link.server_name = server_name; + link.warmup_recv_addr = warmup_recv_addr; + link.is_candidate = true; + link.skip_warmup = skip_warmup_; + + // Reset probe backoff so connection attempt starts promptly. + link.probe_backoff = PeerLink::kProbeBackoffMin; + link.next_probe_time = std::chrono::steady_clock::now(); + + wakeup(); +} + +void LinkManager::disconnect(GlobalRank peer) { + if (peer == rank_) return; + if (!rankInRange(peer)) return; + + std::lock_guard lock(peers_mutex_); + tearDownPeerLink(peer); +} + +void LinkManager::requestHealthCheck(GlobalRank peer) { + if (peer == rank_) return; + if (!rankInRange(peer)) return; + + std::lock_guard lock(peers_mutex_); + auto& link = peers_[peer]; + if (!link.is_candidate || link.state != PeerLinkState::Connected || + !link.target_id.has_value() || link.skip_warmup || + !warmup_send_region_ || link.warmup_recv_addr == 0) { + return; + } + if (link.health_check_requested) return; + + link.health_check_requested = true; + link.probe_backoff = PeerLink::kProbeBackoffMin; + link.next_probe_time = std::chrono::steady_clock::now(); + wakeup(); +} + +void LinkManager::stopReconnect(GlobalRank peer) { + if (peer == rank_) return; + if (!rankInRange(peer)) return; + + std::lock_guard lock(peers_mutex_); + auto& link = peers_[peer]; + link.is_candidate = false; + link.health_check_requested = false; + if (link.probe_batch_id.has_value()) { + engine_->freeBatchID(link.probe_batch_id.value()); + link.probe_batch_id = std::nullopt; + } +} + +bool LinkManager::isConnected(GlobalRank peer) const { + if (!rankInRange(peer)) return false; + return read_state_[peer].link_connected.load(std::memory_order_acquire) != + 0; +} + +void LinkManager::setEventCallback(EventCallback callback) { + std::lock_guard lock(event_callback_mutex_); + event_callback_ = std::move(callback); +} + +std::optional LinkManager::resolvePeer( + GlobalRank peer) const { + if (!rankInRange(peer)) return std::nullopt; + + auto& rs = read_state_[peer]; + + uint64_t v1 = rs.version.load(std::memory_order_acquire); + if (rs.link_connected.load(std::memory_order_acquire) == 0) + return std::nullopt; + auto target_id = rs.target_id.load(std::memory_order_acquire); + uint64_t v2 = rs.version.load(std::memory_order_acquire); + + if (v1 != v2) return std::nullopt; + + return target_id; +} + +void LinkManager::refreshPeerSegment(GlobalRank peer) { + if (peer == rank_) return; + if (!rankInRange(peer)) return; + + std::lock_guard lock(peers_mutex_); + auto& link = peers_[peer]; + + // Only meaningful for connected peers. + if (link.state != PeerLinkState::Connected) return; + if (!link.target_id.has_value()) return; + + engine_->closeSegment(link.target_id.value()); + engine_->removeLocalSegment(link.server_name); + link.target_id = engine_->openSegment(link.server_name); + + read_state_[peer].target_id.store(link.target_id.value(), + std::memory_order_relaxed); + read_state_[peer].version.fetch_add(1, std::memory_order_release); +} + +void LinkManager::publishLinkUp(GlobalRank peer, + TransferMetadata::SegmentID target_id, + uint64_t target_rank_epoch) { + if (!rankInRange(peer)) return; + read_state_[peer].target_id.store(target_id, std::memory_order_relaxed); + read_state_[peer].link_connected.store(1, std::memory_order_release); + read_state_[peer].version.fetch_add(1, std::memory_order_release); + emit(TELinkUpEvent{ + .peer = peer, + .target_rank_epoch = target_rank_epoch, + }); +} + +void LinkManager::publishLinkDown(GlobalRank peer) { + if (!rankInRange(peer)) return; + read_state_[peer].link_connected.store(0, std::memory_order_release); + read_state_[peer].version.fetch_add(1, std::memory_order_release); +} + +void LinkManager::tearDownPeerLink(GlobalRank peer) { + auto& link = peers_[peer]; + + if (link.target_id.has_value()) { + engine_->closeSegment(link.target_id.value()); + link.target_id = std::nullopt; + } + + if (!link.server_name.empty()) { + engine_->removeLocalSegment(link.server_name); + } + + if (link.probe_batch_id.has_value()) { + engine_->freeBatchID(link.probe_batch_id.value()); + link.probe_batch_id = std::nullopt; + } + link.health_check_requested = false; + + link.state = PeerLinkState::Idle; + + publishLinkDown(peer); +} + +void LinkManager::emit(TELinkUpEvent event) { + std::lock_guard lock(event_callback_mutex_); + if (event_callback_) { + event_callback_(std::move(event)); + } +} + +void LinkManager::wakeup() { wakeup_cv_.notify_one(); } + +void LinkManager::pollerLoop() { + while (poller_running_.load(std::memory_order_acquire)) { + bool did_work = false; + + // Find the next eligible peer to probe under a single lock. + GlobalRank connection_peer = kInvalidGlobalRank; + GlobalRank health_check_peer = kInvalidGlobalRank; + { + std::lock_guard lock(peers_mutex_); + auto now = std::chrono::steady_clock::now(); + + for (int peer = 0; peer < max_world_size_; ++peer) { + if (peer == rank_) continue; + + auto& link = peers_[peer]; + if (!link.is_candidate) continue; + if (link.state == PeerLinkState::Connected) { + if (link.health_check_requested && + now >= link.next_probe_time) { + health_check_peer = peer; + break; + } + continue; + } + if (now < link.next_probe_time) continue; + + // Schedule next probe attempt with exponential backoff. + link.next_probe_time = now + link.probe_backoff; + connection_peer = peer; + break; // probe one per iteration + } + } + + if (health_check_peer != kInvalidGlobalRank) { + did_work = advanceHealthCheck(health_check_peer); + } else if (connection_peer != kInvalidGlobalRank) { + did_work = advanceConnection(connection_peer); + } + + // Sleep a bit if idle, or a tiny bit if active (to avoid busy-loop). + { + std::unique_lock lock(wakeup_mutex_); + auto sleep_ms = + did_work ? kPollerActiveSleepMs : kPollerIdleSleepMs; + wakeup_cv_.wait_for(lock, std::chrono::milliseconds(sleep_ms)); + } + } +} + +bool LinkManager::advanceHealthCheck(GlobalRank peer) { + std::lock_guard lock(peers_mutex_); + auto& link = peers_[peer]; + if (!link.is_candidate || link.state != PeerLinkState::Connected || + !link.health_check_requested || !link.target_id.has_value()) { + return false; + } + + auto now = std::chrono::steady_clock::now(); + if (!link.probe_batch_id.has_value()) { + auto batch_id = engine_->allocateBatchID(1); + uint64_t target_offset = + link.warmup_recv_addr + rank_ * sizeof(int32_t); + engine_->submitTransfer(batch_id, + {TransferRequest{ + .opcode = TransferRequest::WRITE, + .source = warmup_send_region_.get(), + .target_id = *link.target_id, + .target_offset = target_offset, + .length = sizeof(int32_t), + }}); + link.probe_batch_id = batch_id; + link.next_probe_time = + now + std::chrono::milliseconds(kPollerActiveSleepMs); + return true; + } + + TransferStatus status; + engine_->getTransferStatus(*link.probe_batch_id, 0, status); + if (status.s == TransferStatusEnum::COMPLETED) { + engine_->freeBatchID(*link.probe_batch_id); + link.probe_batch_id = std::nullopt; + link.health_check_requested = false; + link.probe_backoff = PeerLink::kProbeBackoffMin; + emit(TELinkUpEvent{.peer = peer, + .target_rank_epoch = link.target_rank_epoch}); + return true; + } + + if (status.s == TransferStatusEnum::FAILED) { + LOG(WARNING) << "LinkManager: health check rank " << rank_ << " -> " + << peer << " FAILED"; + engine_->freeBatchID(*link.probe_batch_id); + link.probe_batch_id = std::nullopt; + link.next_probe_time = now + link.probe_backoff; + link.probe_backoff = + std::min(link.probe_backoff * 2, PeerLink::kProbeBackoffMax); + return false; + } + + link.next_probe_time = + now + std::chrono::milliseconds(kPollerActiveSleepMs); + return true; +} + +bool LinkManager::advanceConnection(GlobalRank peer) { + std::lock_guard lock(peers_mutex_); + auto& link = peers_[peer]; + + switch (link.state) { + case PeerLinkState::Idle: { + if (link.server_name.empty()) return false; + + auto segment_id = engine_->openSegment(link.server_name); + if (segment_id == static_cast(-1)) { + LOG(WARNING) + << "[LINK] openSegment failed rank=" << rank_ + << " peer=" << peer << " server_name=" << link.server_name; + link.probe_backoff = std::min(link.probe_backoff * 2, + PeerLink::kProbeBackoffMax); + return false; + } + + link.target_id = segment_id; + + if (link.skip_warmup) { + link.state = PeerLinkState::Connected; + link.probe_backoff = PeerLink::kProbeBackoffMin; + publishLinkUp(peer, segment_id, link.target_rank_epoch); + return true; + } + + // Warmup handshake based on total ordering to avoid both sides + // initiating warmup simultaneously. + if (peer <= rank_) { + // We initiate: write to peer's warmup recv region. + auto batch_id = engine_->allocateBatchID(1); + uint64_t target_offset = + link.warmup_recv_addr + rank_ * sizeof(int32_t); + + engine_->submitTransfer(batch_id, + {TransferRequest{ + .opcode = TransferRequest::WRITE, + .source = warmup_send_region_.get(), + .target_id = segment_id, + .target_offset = target_offset, + .length = sizeof(int32_t), + }}); + link.probe_batch_id = batch_id; + link.state = PeerLinkState::WaitingWarmupTransfer; + } else { + // Wait for the peer to warmup us. + link.state = PeerLinkState::WaitingPeerWarmup; + } + return true; + } + + case PeerLinkState::WaitingWarmupTransfer: { + if (!link.probe_batch_id.has_value()) { + link.state = PeerLinkState::Idle; + return false; + } + + TransferStatus status; + engine_->getTransferStatus(link.probe_batch_id.value(), 0, status); + + if (status.s == TransferStatusEnum::COMPLETED) { + engine_->freeBatchID(link.probe_batch_id.value()); + link.probe_batch_id = std::nullopt; + link.state = PeerLinkState::Connected; + link.probe_backoff = PeerLink::kProbeBackoffMin; + publishLinkUp(peer, link.target_id.value(), + link.target_rank_epoch); + return true; + } + + if (status.s == TransferStatusEnum::FAILED) { + LOG(WARNING) + << "LinkManager: warmup rank " << rank_ << " -> " << peer + << " FAILED" + << " warmup_recv_addr=" << (void*)link.warmup_recv_addr + << " target_offset=" + << (void*)(link.warmup_recv_addr + rank_ * sizeof(int32_t)); + engine_->freeBatchID(link.probe_batch_id.value()); + link.probe_batch_id = std::nullopt; + engine_->closeSegment(link.target_id.value()); + link.target_id = std::nullopt; + link.state = PeerLinkState::Idle; + link.probe_backoff = std::min(link.probe_backoff * 2, + PeerLink::kProbeBackoffMax); + return false; + } + return false; + } + + case PeerLinkState::WaitingPeerWarmup: { + if (!warmup_recv_region_) { + link.state = PeerLinkState::Idle; + return false; + } + + auto* warmup_flag = + reinterpret_cast(&warmup_recv_region_[peer]); + if (*warmup_flag) { + // Consume the one-shot signal here. + *warmup_flag = 0; + link.state = PeerLinkState::Connected; + link.probe_backoff = PeerLink::kProbeBackoffMin; + publishLinkUp(peer, link.target_id.value(), + link.target_rank_epoch); + return true; + } + return false; + } + + case PeerLinkState::Connected: + break; + } + + return false; +} + +} // namespace mooncake diff --git a/mooncake-pg/src/control_plane/rpc_runtime.cpp b/mooncake-pg/src/control_plane/rpc_runtime.cpp new file mode 100644 index 0000000000..1eef9768a1 --- /dev/null +++ b/mooncake-pg/src/control_plane/rpc_runtime.cpp @@ -0,0 +1,125 @@ +#include "control_plane/rpc_runtime.h" + +#include + +namespace mooncake { + +RpcClient::RpcClient(std::chrono::milliseconds request_timeout, + std::chrono::milliseconds connect_timeout) + : state_(std::make_shared(request_timeout, connect_timeout)) {} + +RpcServer::RpcServer(uint16_t port, unsigned thread_num) + : port_(port), thread_num_(thread_num) { + server_ = std::make_unique(thread_num_, port_); +} + +bool RpcServer::start() { + if (!server_) return false; + // Use async_start() instead of start() to avoid blocking the calling + // thread. + auto fut = server_->async_start(); + if (fut.hasResult()) { + // Listen failed - the future resolved immediately. + auto ec = std::move(fut).get(); + LOG(ERROR) << "RpcServer: failed to start: " << ec.message(); + return false; + } + return true; +} + +std::string RpcServer::getListenAddr(const std::string& host_ip) const { + if (!server_) return ""; + return host_ip + ":" + std::to_string(server_->port()); +} + +void RpcServer::shutdown() { + if (server_) { + server_->stop(); + server_.reset(); + } +} + +async_simple::coro::Lazy> +RpcClient::getOrCreateClient(std::shared_ptr state, + const std::string& addr) { + // Fast path: lookup under lock. + { + std::lock_guard lock(state->mutex); + auto it = state->clients.find(addr); + if (it != state->clients.end()) co_return it->second; + } + + // Slow path: create + connect outside the lock so that a slow TCP + // handshake suspends only this coroutine, not other peers'. + coro_rpc::coro_rpc_client::config config; + config.connect_timeout_duration = state->connect_timeout; + config.request_timeout_duration = state->request_timeout; + + auto client = std::make_shared( + coro_io::get_global_executor(), config); + + auto ec = co_await client->connect(addr); + if (ec) { + LOG(ERROR) << "RpcClient: connect failed to " << addr << ": " + << ec.message(); + co_return nullptr; + } + + // Double-check under lock. + { + std::lock_guard lock(state->mutex); + auto it = state->clients.find(addr); + if (it != state->clients.end()) co_return it->second; + + state->clients[addr] = client; + co_return client; + } +} + +void RpcClient::spawn(async_simple::coro::Lazy task) { + auto executor = coro_io::get_global_executor(); + std::move(task).via(executor).start([](auto&&) {}); +} + +std::unique_ptr RpcClient::createSyncClient() { + coro_rpc::coro_rpc_client::config config; + config.connect_timeout_duration = state_->connect_timeout; + config.request_timeout_duration = state_->request_timeout; + return std::make_unique( + coro_io::get_global_executor(), config); +} + +bool RpcClient::isConnected(const std::string& addr) const { + std::lock_guard lock(state_->mutex); + return state_->clients.find(addr) != state_->clients.end(); +} + +bool RpcClient::tryReconnect(const std::string& addr) { + // Evict the old entry under the lock. In-flight coroutines may still + // hold shared_ptr copies of the old client, so it stays alive until + // they complete - no use-after-free. + { + std::lock_guard lock(state_->mutex); + state_->clients.erase(addr); + } + + coro_rpc::coro_rpc_client::config config; + config.connect_timeout_duration = state_->connect_timeout; + config.request_timeout_duration = state_->request_timeout; + + auto client = std::make_shared( + coro_io::get_global_executor(), config); + + auto ec = async_simple::coro::syncAwait(client->connect(addr)); + if (ec) { + LOG(ERROR) << "RpcClient: reconnect failed to " << addr << ": " + << ec.message(); + return false; + } + + std::lock_guard lock(state_->mutex); + state_->clients[addr] = client; + return true; +} + +} // namespace mooncake diff --git a/mooncake-pg/src/gpu_runtime.cpp b/mooncake-pg/src/gpu_runtime.cpp new file mode 100644 index 0000000000..22619b3337 --- /dev/null +++ b/mooncake-pg/src/gpu_runtime.cpp @@ -0,0 +1,158 @@ +#include "gpu_runtime.h" + +#include +#include + +#include + +#include "error_types.h" + +namespace mooncake { +namespace { + +void warnCleanupFailure(const char* operation, const char* error) noexcept { + LOG(WARNING) << "Mooncake PG CUDA cleanup failed while " << operation + << ": " << (error ? error : "unknown error"); +} + +} // namespace + +GpuDeviceGuard::GpuDeviceGuard(int device) { + PG_ASSERT(device >= 0, "invalid CUDA device index"); + + PG_ASSERT_CUDA(cudaGetDevice(&previous_device_)); + if (previous_device_ == device) return; + + PG_ASSERT_CUDA(cudaSetDevice(device)); + restore_device_ = true; +} + +GpuDeviceGuard::~GpuDeviceGuard() noexcept { + if (!restore_device_) return; + const auto error = cudaSetDevice(previous_device_); + if (error != cudaSuccess) { + warnCleanupFailure("restore collective CUDA device", + cudaGetErrorString(error)); + } +} + +GpuStream::~GpuStream() noexcept { reset(); } + +GpuStream::GpuStream(GpuStream&& other) noexcept { moveFrom(std::move(other)); } + +GpuStream& GpuStream::operator=(GpuStream&& other) noexcept { + if (this != &other) { + reset(); + moveFrom(std::move(other)); + } + return *this; +} + +GpuStream GpuStream::createNonBlocking(int device) { + const GpuDeviceGuard device_guard(device); + cudaStream_t stream = nullptr; + PG_ASSERT_CUDA(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking)); + return GpuStream(stream, device, true); +} + +GpuStream GpuStream::borrow(cudaStream_t stream, int device) { + PG_ASSERT(device >= 0, "invalid CUDA device index"); + return GpuStream(stream, device, false); +} + +bool GpuStream::isCapturing() const { + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + PG_ASSERT_CUDA(cudaStreamIsCapturing(stream_, &capture_status)); + return capture_status != cudaStreamCaptureStatusNone; +} + +void GpuStream::waitEvent(const GpuEvent& event) const { + const GpuDeviceGuard device_guard(device_index_); + PG_ASSERT_CUDA(cudaStreamWaitEvent(stream_, event.event_, 0)); +} + +GpuStream::GpuStream(cudaStream_t stream, int device, bool owns_stream) noexcept + : stream_(stream), device_index_(device), owns_stream_(owns_stream) {} + +void GpuStream::reset() noexcept { + if (owns_stream_ && stream_) { + try { + const GpuDeviceGuard device_guard(device_index_); + const auto error = cudaStreamDestroy(stream_); + if (error != cudaSuccess) { + warnCleanupFailure("destroy stream", cudaGetErrorString(error)); + } + } catch (const std::exception& error) { + warnCleanupFailure("destroy stream", error.what()); + } catch (...) { + warnCleanupFailure("destroy stream", "unknown error"); + } + } + stream_ = nullptr; + device_index_ = -1; + owns_stream_ = false; +} + +void GpuStream::moveFrom(GpuStream&& other) noexcept { + stream_ = other.stream_; + device_index_ = other.device_index_; + owns_stream_ = other.owns_stream_; + other.stream_ = nullptr; + other.device_index_ = -1; + other.owns_stream_ = false; +} + +GpuEvent::GpuEvent(int device, unsigned int flags) : device_index_(device) { + PG_ASSERT(device >= 0, "invalid CUDA device index"); + + const GpuDeviceGuard device_guard(device); + PG_ASSERT_CUDA(cudaEventCreateWithFlags(&event_, flags)); +} + +GpuEvent::~GpuEvent() noexcept { reset(); } + +GpuEvent::GpuEvent(GpuEvent&& other) noexcept { moveFrom(std::move(other)); } + +GpuEvent& GpuEvent::operator=(GpuEvent&& other) noexcept { + if (this != &other) { + reset(); + moveFrom(std::move(other)); + } + return *this; +} + +void GpuEvent::record(const GpuStream& stream) { + PG_ASSERT(device_index_ == stream.deviceIndex(), + "CUDA event device does not match recording stream device"); + + const GpuDeviceGuard device_guard(device_index_); + PG_ASSERT_CUDA(cudaEventRecord(event_, stream.get())); +} + +void GpuEvent::reset() noexcept { + if (event_) { + try { + const GpuDeviceGuard device_guard(device_index_); + const auto error = cudaEventDestroy(event_); + if (error != cudaSuccess) { + warnCleanupFailure("destroy CUDA event", + cudaGetErrorString(error)); + } + } catch (const std::exception& error) { + warnCleanupFailure("destroy CUDA event", error.what()); + } catch (...) { + warnCleanupFailure("destroy CUDA event", "unknown error"); + } + } + event_ = nullptr; + device_index_ = -1; +} + +void GpuEvent::moveFrom(GpuEvent&& other) noexcept { + event_ = other.event_; + device_index_ = other.device_index_; + other.event_ = nullptr; + other.device_index_ = -1; +} + +} // namespace mooncake diff --git a/mooncake-pg/src/mooncake_backend.cpp b/mooncake-pg/src/mooncake_backend.cpp deleted file mode 100644 index e4c630a49a..0000000000 --- a/mooncake-pg/src/mooncake_backend.cpp +++ /dev/null @@ -1,1312 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include "connection_poller.h" -#include "memory_location.h" -#include "mooncake_worker.cuh" -#include "pg_utils.h" - -namespace mooncake { -namespace { - -#ifdef USE_MACA -static void requireMacaHostTransport() { - TORCH_CHECK(std::getenv("MC_MACA_HOST_TRANSPORT") != nullptr, - "MACA PG requires MC_MACA_HOST_TRANSPORT=1 so the transfer " - "engine uses a host transport."); -} - -#endif - -} // namespace - -constexpr const char* REGISTER_BUFFER_ERROR_MSG = - "Failed to register local memory."; -constexpr const char* MULTI_DEVICE_ERROR_MSG = - "Expecting one tensor only but got multiple."; -constexpr const char* SYNC_OP_ERROR_MSG = "Expecting async op but got sync op."; -constexpr const char* REDUCE_OP_ERROR_MSG = "Only support SUM."; -constexpr const char* SPARSE_ERROR_MSG = "Sparse op not supported."; -constexpr const char* REDUCE_DTYPE_ERROR_MSG = "Unsupported reduce dtype: "; -constexpr int kBarrierDummyTensorSize = 1; - -std::string MooncakeBackend::hostIp_ = "127.0.0.1"; -// leaky singleton to avoid destructor fiasco problem -TransferEngine* MooncakeBackend::engine_ = new TransferEngine(true); -// worker_ is now owned per backend instance via MooncakeWorkerManager. -bool MooncakeBackend::engineInitialized_ = false; -int MooncakeBackend::backendIndex_ = 0; -TransferEngine* MooncakeBackend::externalEngine_ = nullptr; - -std::vector serialize(const ExtensionState& state) { - uint32_t rankCount = static_cast(state.activeRanks.size()); - - // Calculate bytes needed for the bitmap: 1 bit per rank, rounded up to - // nearest byte - size_t bitmapSize = (rankCount + 7) / 8; - - // Total size = count field + bitmap + p2pEpochs[] + taskCount - size_t totalSize = sizeof(uint32_t) + bitmapSize + - sizeof(uint32_t) * rankCount + sizeof(int32_t); - - std::vector buffer(totalSize, 0); - uint8_t* ptr = buffer.data(); - - // 1. Store the number of ranks - std::memcpy(ptr, &rankCount, sizeof(uint32_t)); - ptr += sizeof(uint32_t); - - // 2. Store activeRanks as a bitset - for (size_t i = 0; i < rankCount; ++i) { - if (state.activeRanks[i]) { - // Set the i-th bit to 1 if the rank is active - ptr[i / 8] |= (1 << (i % 8)); - } - } - ptr += bitmapSize; - - // 3. Store per-peer p2pEpochs - for (size_t i = 0; i < rankCount; ++i) { - std::memcpy(ptr + i * sizeof(uint32_t), &state.p2pEpochs[i], - sizeof(uint32_t)); - } - ptr += sizeof(uint32_t) * rankCount; - - // 4. Store taskCount - int32_t taskCount = static_cast(state.taskCount); - std::memcpy(ptr, &taskCount, sizeof(int32_t)); - - return buffer; -} - -ExtensionState deserialize(const std::vector& buffer) { - ExtensionState state; - if (buffer.size() < sizeof(uint32_t)) return state; - - const uint8_t* ptr = buffer.data(); - - // 1. Read the number of ranks - uint32_t rankCount = 0; - std::memcpy(&rankCount, ptr, sizeof(uint32_t)); - ptr += sizeof(uint32_t); - - // Calculate expected total size and verify buffer is sufficient before - // proceeding with further reads. - size_t bitmapSize = (rankCount + 7) / 8; - size_t expectedSize = sizeof(uint32_t) + bitmapSize + - sizeof(uint32_t) * rankCount + sizeof(int32_t); - if (buffer.size() < expectedSize) return state; - - // 2. Read the bitmap and reconstruct the activeRanks vector - state.activeRanks.resize(rankCount); - for (size_t i = 0; i < rankCount; ++i) { - // Check if the i-th bit is set - bool isActive = ptr[i / 8] & (1 << (i % 8)); - state.activeRanks[i] = isActive; - } - ptr += bitmapSize; - - // 3. Read per-peer p2pEpochs - state.p2pEpochs.resize(rankCount); - for (size_t i = 0; i < rankCount; ++i) { - std::memcpy(&state.p2pEpochs[i], ptr, sizeof(uint32_t)); - ptr += sizeof(uint32_t); - } - - // 4. Read taskCount - int32_t taskCount = 0; - std::memcpy(&taskCount, ptr, sizeof(int32_t)); - state.taskCount = static_cast(taskCount); - - return state; -} - -// Async Work implementation for P2P operations processed by worker threads. -class MooncakeP2PWork : public ::c10d::Work { - public: - explicit MooncakeP2PWork( - std::shared_ptr> status) - : Work(-1, c10d::OpType::UNKNOWN), status_(status) {} - - bool isCompleted() override { - return status_->load(std::memory_order_acquire) != - P2PProxy::OpStatus::kPending; - } - - bool isSuccess() const override { - return status_->load(std::memory_order_acquire) == - P2PProxy::OpStatus::kSuccess; - } - - bool wait(std::chrono::milliseconds timeout) override { - BackoffWaiterConfig cfg{}; - cfg.max_sleep = std::chrono::microseconds(10); - BackoffWaiter waiter(cfg); - - bool done = false; - if (timeout.count() > 0) { - done = waiter.wait_for(timeout, [this] { - return status_->load(std::memory_order_acquire) != - P2PProxy::OpStatus::kPending; - }); - } else { - waiter.wait([this] { - return status_->load(std::memory_order_acquire) != - P2PProxy::OpStatus::kPending; - }); - done = true; - } - - if (!done) { - return false; - } - - if (status_->load(std::memory_order_acquire) == - P2PProxy::OpStatus::kFailed) { - TORCH_CHECK(false, "Mooncake P2P operation failed."); - } - return true; - } - - private: - std::shared_ptr> status_; -}; - -/** - * @brief Initialize Mooncake backend state from the PyTorch process-group - * information and optional Mooncake-specific options. - */ -MooncakeBackend::MooncakeBackend( - c10d::DistributedBackendOptions distBackendOpts, - c10::intrusive_ptr options, bool isCpu) - : ProcessGroup(distBackendOpts.store, distBackendOpts.group_rank, - distBackendOpts.group_size), - options_(std::move(options)), - isCpu_(isCpu) { - auto store = std::move(distBackendOpts.store); - const int rank = distBackendOpts.group_rank; - const int size = distBackendOpts.group_size; - const int max_size = (options_ && options_->maxWorldSize_ > 0) - ? options_->maxWorldSize_ - : size; - - TORCH_CHECK(max_size >= 0 && static_cast(max_size) <= kMaxNumRanks, - "max_world_size out of range"); - TORCH_CHECK(max_size >= size, - "max_world_size must be >= process group size"); - const auto& globalRanks = distBackendOpts.global_ranks_in_group; - - // Memory location for device specific buffers - // always kWildcardLocation for cpu backend - std::string location = kWildcardLocation; - if (!isCpu) { - int deviceCount = 0; - cudaError_t err = cudaGetDeviceCount(&deviceCount); - if (err == cudaSuccess && deviceCount != 0) { - int deviceId_; - err = cudaGetDevice(&deviceId_); - TORCH_CHECK(!err, c10::str("Failed to get device id")); - location = GPU_PREFIX + std::to_string(deviceId_); - } - } - - // Initialize transfer engine - if (externalEngine_) { - // Use externally-provided engine (already initialized), skip init. - engine_ = externalEngine_; - engineInitialized_ = true; - } else if (!engineInitialized_) { -#ifdef USE_MACA - requireMacaHostTransport(); -#endif - engine_->init(P2PHANDSHAKE, hostIp_); - engineInitialized_ = true; - } - localServerName_ = engine_->getLocalIpAndPort(); - // construct local to global rank map - if (globalRanks.size() == static_cast(size)) { - for (int i = 0; i < size; ++i) { - local2global_rank_map_[i] = static_cast(globalRanks[i]); - } - } else { - for (int i = 0; i < size; ++i) { - local2global_rank_map_[i] = i; - } - } - - // Fill the remaining slots for polling / future joiners. - for (int i = size; i < max_size; ++i) { - local2global_rank_map_[i] = i; - } - - // Register buffers - if (isCpu) { - for (size_t i = 0; i < 2; i++) { - send_buffer_[i] = malloc(kBufferSize); - TORCH_CHECK(send_buffer_[i], - c10::str("Failed to allocate CPU send buffer")); - - int rc = engine_->registerLocalMemory(send_buffer_[i], kBufferSize, - location); - TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG); - } - - for (size_t i = 0; i < 2; i++) { - recv_buffer_[i] = malloc(kBufferSize); - TORCH_CHECK(recv_buffer_[i], - c10::str("Failed to allocate CPU recv buffer")); - - int rc = engine_->registerLocalMemory(recv_buffer_[i], kBufferSize, - location); - TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG); - } - -#ifdef USE_MACA - } else { - for (size_t i = 0; i < 2; i++) { - cudaError_t err = cudaMalloc(&send_buffer_[i], kBufferSize); - TORCH_CHECK(!err, - c10::str("Failed to allocate MACA GPU send buffer")); - - int rc = engine_->registerLocalMemory(send_buffer_[i], kBufferSize, - location); - TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG); - } - - for (size_t i = 0; i < 2; i++) { - cudaError_t err = cudaMalloc(&recv_buffer_[i], kBufferSize); - TORCH_CHECK(!err, - c10::str("Failed to allocate MACA GPU recv buffer")); - - int rc = engine_->registerLocalMemory(recv_buffer_[i], kBufferSize, - location); - TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG); - } -#else - } else { - for (size_t i = 0; i < 2; i++) { - cudaError_t err = cudaMalloc(&send_buffer_[i], kBufferSize); - TORCH_CHECK(!err, c10::str("Failed to allocate CUDA send buffer")); - - int rc = engine_->registerLocalMemory(send_buffer_[i], kBufferSize, - location); - TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG); - } - - for (size_t i = 0; i < 2; i++) { - cudaError_t err = cudaMalloc(&recv_buffer_[i], kBufferSize); - TORCH_CHECK(!err, c10::str("Failed to allocate CUDA recv buffer")); - - int rc = engine_->registerLocalMemory(recv_buffer_[i], kBufferSize, - location); - TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG); - } -#endif - } - - // Register CPU sync regions - TORCH_CHECK(static_cast(size) <= kMaxNumRanks, - "The number of ranks exceeds the limit."); - for (size_t i = 0; i < 2; i++) { - cpu_sync_send_region_[i] = new int32_t[kMaxNumRanks]{}; - int rc = engine_->registerLocalMemory(cpu_sync_send_region_[i], - kMaxNumRanks * sizeof(int32_t), - kWildcardLocation); - TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG); - } - - for (size_t i = 0; i < 2; i++) { - cpu_sync_recv_region_[i] = new int32_t[kMaxNumRanks]{}; - int rc = engine_->registerLocalMemory(cpu_sync_recv_region_[i], - kMaxNumRanks * sizeof(int32_t), - kWildcardLocation); - TORCH_CHECK(!rc, REGISTER_BUFFER_ERROR_MSG); - } - - auto& dev_worker_mgr = P2PDeviceWorkerManager::getInstance(); - int cuda_device_index = isCpu_ ? -1 : at::cuda::current_device(); - - if (isCpu_) - p2p_device_worker_ = dev_worker_mgr.getCPUWorker(engine_); - else - p2p_device_worker_ = - dev_worker_mgr.getCUDAWorker(cuda_device_index, engine_); - - auto& worker_mgr = MooncakeWorkerManager::GetInstance(); - if (isCpu_) - worker_ = worker_mgr.GetCPUWorker(); - else - worker_ = worker_mgr.GetCUDAWorker(cuda_device_index); - if (!isCpu_) { - preloadReduceKernels(); - } - worker_->Start(); - - p2p_proxy_ = std::make_shared( - engine_, P2PProxy::Options{ - .is_cpu = isCpu_, - .rank = rank_, - .size = size_, - .cuda_device_index = cuda_device_index, - }); - p2p_device_worker_->registerProxy(p2p_proxy_); - - meta_ = std::make_shared(); - connection_ctx_ = std::make_shared( - backendIndex_, rank, size, options_ && options_->isExtension_, - local2global_rank_map_, store, meta_, p2p_proxy_, engine_); - - if (max_size != size) { - connection_ctx_->setPollingLimitTo(max_size); - } - - rank_info.send_buffer[0] = (uint64_t)send_buffer_[0]; - rank_info.send_buffer[1] = (uint64_t)send_buffer_[1]; - rank_info.recv_buffer[0] = (uint64_t)recv_buffer_[0]; - rank_info.recv_buffer[1] = (uint64_t)recv_buffer_[1]; - rank_info.send_sync[0] = (uint64_t)cpu_sync_send_region_[0]; - rank_info.send_sync[1] = (uint64_t)cpu_sync_send_region_[1]; - rank_info.recv_sync[0] = (uint64_t)cpu_sync_recv_region_[0]; - rank_info.recv_sync[1] = (uint64_t)cpu_sync_recv_region_[1]; - rank_info.warmup_buffer[0] = - (uint64_t)connection_ctx_->warmup_send_region(); - rank_info.warmup_buffer[1] = - (uint64_t)connection_ctx_->warmup_recv_region(); - rank_info.p2p_credit_region = (uint64_t)p2p_proxy_->credit_region(); - rank_info.p2p_ack_region = (uint64_t)p2p_proxy_->ack_region(); - - // Sync metadata - std::vector rank_info_bytes(sizeof(SegmentInfo)); - memcpy(rank_info_bytes.data(), &rank_info, sizeof(SegmentInfo)); - meta_->rank = rank; - // NOTE: meta_->size is intentionally initialized to max_world_size (when - // provided) so that healthy ranks can activate joiners via recoverRanks() - // without calling extendGroupSizeTo(). Inactive slots are masked by - // meta_->activeRanks / meta_->activeRanksTensor. - meta_->size = max_size; - // activeSize tracks the visible group size (returned by getSize() / - // dist.get_world_size()). It starts at the actual member count and grows - // when extendGroupSizeTo() or recoverRanks() expands the group. - // For extension ranks, activeSize equals world_size (= max_world_size); - // the local-only behavior before joinGroup() is ensured by activeRanks - // masking, not by a smaller activeSize. - meta_->activeSize = size; - meta_->taskCount = 0; - if (isCpu) { - meta_->activeRanks = new bool[kMaxNumRanks]; - } else { - cudaHostAlloc(&meta_->activeRanks, kMaxNumRanks * sizeof(bool), - cudaHostAllocMapped); - cudaHostGetDevicePointer(&meta_->activeRanksDevice, meta_->activeRanks, - 0); - } - for (size_t i = 0; i < kMaxNumRanks; ++i) { - meta_->activeRanks[i] = true; - } - - // Reserve extra slots as inactive so collectives won't wait on them. - for (int i = size; i < max_size; ++i) { - meta_->activeRanks[i] = false; - } - if (options_ && options_->activeRanks_.defined()) { - TORCH_CHECK(options_->activeRanks_.dtype() == at::kInt, - "activeRanks must be int."); - if (isCpu) { - TORCH_CHECK(options_->activeRanks_.device().is_cpu(), - "activeRanks must be on CPU."); - } else { - TORCH_CHECK( - options_->activeRanks_.device().type() == c10::DeviceType::CUDA, - "activeRanks must be on GPU."); - } - if (max_size != size) { - TORCH_CHECK(options_->activeRanks_.numel() == max_size, - "activeRanks must be sized to max_world_size when " - "max_world_size is set"); - } - meta_->activeRanksTensor = options_->activeRanks_; - } else { - meta_->activeRanksTensor = at::ones( - {max_size}, torch::dtype(torch::kInt32) - .device(isCpu ? torch::kCPU : torch::kCUDA)); - if (max_size != size) { - meta_->activeRanksTensor.slice(0, size, max_size).fill_(0); - } - } - meta_->engine = engine_; - meta_->store = store; - meta_->backendIndex = backendIndex_; - meta_->bufferBaseIndex = backendIndex_ * 10; - p2p_proxy_->bindMeta(meta_); - - connection_ctx_->bootstrapLocalPeer(localServerName_, rank_info); - if (options_ && options_->isExtension_) { - setLocalOnlyActiveRanks(); - } else { - publishLocalPeerMetadata(); - ConnectionPoller::GetInstance().registerContext(connection_ctx_); - connectionPollerRegistered_ = true; - connection_ctx_->waitUntilAllConnected(); - } - - // Register a lightweight Backend shim so that PyTorch's P2P dispatch path - // (batch_isend_irecv → _get_backend → getBackend) can find a registered - // Backend for this ProcessGroup. The shim delegates send/recv back to us. - auto deviceType = isCpu ? c10::DeviceType::CPU : c10::DeviceType::CUDA; - auto shim = c10::make_intrusive(this); - setBackend(deviceType, BackendType::CUSTOM, shim); -#ifndef MOONCAKE_EP_USE_MUSA - setDefaultBackend(BackendType::CUSTOM); -#endif - - // Increment backend index - ++backendIndex_; -} - -MooncakeBackend::~MooncakeBackend() { shutdown(); } - -const std::string MooncakeBackend::getBackendName() const { return "mooncake"; } - -// ---- MooncakeP2PShim implementation ---- - -MooncakeP2PShim::MooncakeP2PShim(MooncakeBackend* owner) - : Backend(owner->getRank(), owner->getSize()), owner_(owner) {} - -const std::string MooncakeP2PShim::getBackendName() const { return "mooncake"; } - -c10::intrusive_ptr MooncakeP2PShim::send( - std::vector& tensors, int dstRank, int tag) { - return owner_->send(tensors, dstRank, tag); -} - -c10::intrusive_ptr MooncakeP2PShim::recv( - std::vector& tensors, int srcRank, int tag) { - return owner_->recv(tensors, srcRank, tag); -} - -c10::intrusive_ptr MooncakeP2PShim::recvAnysource( - std::vector& tensors, int tag) { - // MooncakeBackend doesn't implement recvAnysource; fall back to - // the base class which will raise a clear error. - return ::c10d::Backend::recvAnysource(tensors, tag); -} - -c10::intrusive_ptr MooncakeP2PShim::barrier( - const c10d::BarrierOptions& opts) { - return owner_->barrier(opts); -} - -c10::intrusive_ptr MooncakeBackend::send( - std::vector& tensors, int dstRank, int tag) { - connection_ctx_->waitUntilNewRanksConnected(); - - (void)tag; - TORCH_CHECK(tensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - auto tensor = tensors.back(); - - TORCH_CHECK(meta_->store, "P2P send requires a valid Store."); - TORCH_CHECK(dstRank >= 0 && dstRank < meta_->size, - "P2P send: dstRank out of range."); - - auto contiguous = tensor.contiguous(); - auto status = std::make_shared>( - P2PProxy::OpStatus::kPending); - cudaStream_t stream = nullptr; - if (!isCpu_) { - auto current_stream = - at::cuda::getCurrentCUDAStream(contiguous.device().index()); - stream = current_stream.stream(); - } - - TORCH_CHECK(p2p_proxy_, "P2P send proxy is not initialized."); - p2p_proxy_->enqueueSend(P2PProxy::SendOp{ - .tensor_ = std::move(contiguous), - .peer_rank_ = dstRank, - .cuda_stream_ = stream, - .status_ = status, - }); - - return c10::make_intrusive(status); -} - -c10::intrusive_ptr MooncakeBackend::recv( - std::vector& tensors, int srcRank, int tag) { - connection_ctx_->waitUntilNewRanksConnected(); - - (void)tag; - TORCH_CHECK(tensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - auto tensor = tensors.back(); - - TORCH_CHECK(meta_->store, "P2P recv requires a valid Store."); - TORCH_CHECK(srcRank >= 0 && srcRank < meta_->size, - "P2P recv: srcRank out of range."); - - auto target = tensor.is_contiguous() ? tensor : tensor.contiguous(); - auto status = std::make_shared>( - P2PProxy::OpStatus::kPending); - cudaStream_t stream = nullptr; - if (!isCpu_) { - auto current_stream = - at::cuda::getCurrentCUDAStream(target.device().index()); - stream = current_stream.stream(); - } - - TORCH_CHECK(p2p_proxy_, "P2P recv proxy is not initialized."); - p2p_proxy_->enqueueRecv(P2PProxy::RecvOp{ - .tensor_ = target, - .original_tensor_ = tensor, - .peer_rank_ = srcRank, - .cuda_stream_ = stream, - .status_ = status, - }); - - return c10::make_intrusive(status); -} - -c10::intrusive_ptr MooncakeBackend::broadcast( - std::vector& tensors, const c10d::BroadcastOptions& opts) { - TORCH_CHECK(tensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - auto tensor = tensors.back(); - size_t tensorSize = tensor.numel() * tensor.element_size(); - int64_t root = opts.rootRank + opts.rootTensor; - bool isRoot = (root == rank_); - if (isCpu_) { - return worker_->putTaskCpu( - c10d::OpType::BROADCAST, tensorSize, root, meta_, connection_ctx_, - [=](void* dst, size_t pos, size_t realSize) { - if (isRoot) { - memcpy(dst, (char*)tensor.data_ptr() + pos, realSize); - } - }, - [=](void* src, size_t pos, size_t realSize) { - memcpy((char*)tensor.data_ptr() + pos, src, realSize); - }); - } else { - at::cuda::CUDAStream stream = - at::cuda::getCurrentCUDAStream(tensor.device().index()); - return worker_->putTaskCuda( - c10d::OpType::BROADCAST, tensorSize, root, meta_, connection_ctx_, - stream, - [=](void* dst, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - if (isRoot) { - cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos, - realSize, cudaMemcpyDeviceToDevice, - enq_stream); - } - }, - [=](void* src, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - cudaMemcpyAsync((char*)tensor.data_ptr() + pos, src, realSize, - cudaMemcpyDeviceToDevice, enq_stream); - }); - } -} - -c10::intrusive_ptr MooncakeBackend::allreduce( - std::vector& tensors, const c10d::AllreduceOptions& opts) { - TORCH_CHECK(tensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - TORCH_CHECK(opts.sparseIndices == std::nullopt, SPARSE_ERROR_MSG); - auto tensor = tensors.back(); - size_t tensorSize = tensor.numel() * tensor.element_size(); - if (isCpu_) { - auto numRanks = meta_->size; - return worker_->putTaskCpu( - c10d::OpType::ALLREDUCE, tensorSize, 0, meta_, connection_ctx_, - [=](void* dst, size_t pos, size_t realSize) { - memcpy(dst, (char*)tensor.data_ptr() + pos, realSize); - }, - [=, this](void* src, size_t pos, size_t realSize) { - memset((char*)tensor.data_ptr() + pos, 0, realSize); - launchReduceCpu(tensor, pos, realSize, src, numRanks, - opts.reduceOp, meta_->activeRanks); - }); - } else { - auto stream = at::cuda::getCurrentCUDAStream(tensor.device().index()); - return worker_->putTaskCuda( - c10d::OpType::ALLREDUCE, tensorSize, 0, meta_, connection_ctx_, - stream, - [=](void* dst, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos, realSize, - cudaMemcpyDeviceToDevice, enq_stream); - }, - [=, this](void* src, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - cudaMemsetAsync((char*)tensor.data_ptr() + pos, 0, realSize, - enq_stream); - launchReduceKernel(tensor, pos, realSize, src, meta_->size, - opts.reduceOp, meta_->activeRanksDevice, - enq_stream); - }); - } -} - -c10::intrusive_ptr MooncakeBackend::allgather( - std::vector>& outputTensors, - std::vector& inputTensors, const c10d::AllgatherOptions& opts) { - TORCH_CHECK(inputTensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - TORCH_CHECK(outputTensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - auto inputTensor = inputTensors.back(); - auto outputTensors_ = outputTensors.back(); - size_t tensorSize = inputTensor.numel() * inputTensor.element_size(); - if (isCpu_) { - return worker_->putTaskCpu( - c10d::OpType::ALLGATHER, tensorSize, 0, meta_, connection_ctx_, - [=](void* dst, size_t pos, size_t realSize) { - memcpy(dst, (char*)inputTensor.data_ptr() + pos, realSize); - }, - [=](void* src, size_t pos, size_t realSize) { - for (const auto j : c10::irange(outputTensors_.size())) { - memcpy((char*)outputTensors_[j].data_ptr() + pos, - (char*)src + j * realSize, realSize); - } - }); - } else { - auto stream = - at::cuda::getCurrentCUDAStream(inputTensor.device().index()); - return worker_->putTaskCuda( - c10d::OpType::ALLGATHER, tensorSize, 0, meta_, connection_ctx_, - stream, - [=](void* dst, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - cudaMemcpyAsync(dst, (char*)inputTensor.data_ptr() + pos, - realSize, cudaMemcpyDeviceToDevice, enq_stream); - }, - [=](void* src, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - for (const auto j : c10::irange(outputTensors_.size())) { - cudaMemcpyAsync((char*)outputTensors_[j].data_ptr() + pos, - (char*)src + j * realSize, realSize, - cudaMemcpyDeviceToDevice, enq_stream); - } - }); - } -} - -c10::intrusive_ptr MooncakeBackend::_allgather_base( - at::Tensor& outputBuffer, at::Tensor& inputBuffer, - const c10d::AllgatherOptions& opts) { - size_t tensorSize = inputBuffer.numel() * inputBuffer.element_size(); - if (isCpu_) { - return worker_->putTaskCpu( - c10d::OpType::_ALLGATHER_BASE, tensorSize, 0, meta_, - connection_ctx_, - [=](void* dst, size_t pos, size_t realSize) { - memcpy(dst, (char*)inputBuffer.data_ptr() + pos, realSize); - }, - [=, this](void* src, size_t pos, size_t realSize) { - for (int j = 0; j < meta_->size; ++j) { - if (!meta_->activeRanks[j]) continue; - memcpy( - (char*)outputBuffer.data_ptr() + j * tensorSize + pos, - (char*)src + j * realSize, realSize); - } - }); - } else { - auto stream = - at::cuda::getCurrentCUDAStream(inputBuffer.device().index()); - return worker_->putTaskCuda( - c10d::OpType::_ALLGATHER_BASE, tensorSize, 0, meta_, - connection_ctx_, stream, - [=](void* dst, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - cudaMemcpyAsync(dst, (char*)inputBuffer.data_ptr() + pos, - realSize, cudaMemcpyDeviceToDevice, enq_stream); - }, - [=, this](void* src, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - for (int j = 0; j < meta_->size; ++j) { - if (!meta_->activeRanks[j]) continue; - cudaMemcpyAsync( - (char*)outputBuffer.data_ptr() + j * tensorSize + pos, - (char*)src + j * realSize, realSize, - cudaMemcpyDeviceToDevice, enq_stream); - } - }); - } -} - -c10::intrusive_ptr MooncakeBackend::_reduce_scatter_base( - at::Tensor& outputBuffer, at::Tensor& inputBuffer, - const c10d::ReduceScatterOptions& opts) { - size_t tensorSize = outputBuffer.numel() * outputBuffer.element_size(); - if (isCpu_) { - auto numRanks = meta_->size; - return worker_->putTaskCpu( - c10d::OpType::_REDUCE_SCATTER_BASE, tensorSize, 0, meta_, - connection_ctx_, - [=, this](void* dst, size_t pos, size_t realSize) { - for (int j = 0; j < meta_->size; ++j) { - if (!meta_->activeRanks[j]) continue; - memcpy((char*)dst + j * realSize, - (char*)inputBuffer.data_ptr() + j * tensorSize + pos, - realSize); - } - }, - [=, this](void* src, size_t pos, size_t realSize) { - memset((char*)outputBuffer.data_ptr() + pos, 0, realSize); - launchReduceCpu(outputBuffer, pos, realSize, src, numRanks, - opts.reduceOp, meta_->activeRanks); - }); - } else { - auto stream = - at::cuda::getCurrentCUDAStream(inputBuffer.device().index()); - return worker_->putTaskCuda( - c10d::OpType::_REDUCE_SCATTER_BASE, tensorSize, 0, meta_, - connection_ctx_, stream, - [=, this](void* dst, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - for (int j = 0; j < meta_->size; ++j) { - if (!meta_->activeRanks[j]) continue; - cudaMemcpyAsync( - (char*)dst + j * realSize, - (char*)inputBuffer.data_ptr() + j * tensorSize + pos, - realSize, cudaMemcpyDeviceToDevice, enq_stream); - } - }, - [=, this](void* src, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - cudaMemsetAsync((char*)outputBuffer.data_ptr() + pos, 0, - realSize, enq_stream); - launchReduceKernel(outputBuffer, pos, realSize, src, - meta_->size, opts.reduceOp, - meta_->activeRanksDevice, enq_stream); - }); - } -} - -c10::intrusive_ptr MooncakeBackend::alltoall( - std::vector& outputTensors, - std::vector& inputTensors, const c10d::AllToAllOptions& opts) { - size_t tensorSize = - inputTensors[0].numel() * inputTensors[0].element_size(); - if (isCpu_) { - return worker_->putTaskCpu( - c10d::OpType::ALLTOALL, tensorSize, 0, meta_, connection_ctx_, - [=](void* dst, size_t pos, size_t realSize) { - for (const auto j : c10::irange(inputTensors.size())) { - memcpy((char*)dst + j * realSize, - (char*)inputTensors[j].data_ptr() + pos, realSize); - } - }, - [=](void* src, size_t pos, size_t realSize) { - for (const auto j : c10::irange(outputTensors.size())) { - memcpy((char*)outputTensors[j].data_ptr() + pos, - (char*)src + j * realSize, realSize); - } - }); - } else { - auto stream = - at::cuda::getCurrentCUDAStream(inputTensors[0].device().index()); - return worker_->putTaskCuda( - c10d::OpType::ALLTOALL, tensorSize, 0, meta_, connection_ctx_, - stream, - [=](void* dst, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - for (const auto j : c10::irange(inputTensors.size())) { - cudaMemcpyAsync((char*)dst + j * realSize, - (char*)inputTensors[j].data_ptr() + pos, - realSize, cudaMemcpyDeviceToDevice, - enq_stream); - } - }, - [=](void* src, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - for (const auto j : c10::irange(outputTensors.size())) { - cudaMemcpyAsync((char*)outputTensors[j].data_ptr() + pos, - (char*)src + j * realSize, realSize, - cudaMemcpyDeviceToDevice, enq_stream); - } - }); - } -} -c10::intrusive_ptr MooncakeBackend::barrier( - const c10d::BarrierOptions& opts) { - if (isCpu_) { - return worker_->putTaskCpu( - // a non-zero tensorSize is required to ensure the worker task for - // the barrier is created - c10d::OpType::BARRIER, kBarrierDummyTensorSize, 0, meta_, - connection_ctx_, [=](void*, size_t, size_t) {}, - [=](void*, size_t, size_t) {}); - } else { - auto device_index = at::cuda::current_device(); - auto stream = at::cuda::getCurrentCUDAStream(device_index); - return worker_->putTaskCuda( - c10d::OpType::BARRIER, kBarrierDummyTensorSize, 0, meta_, - connection_ctx_, stream, - [=](void*, size_t, size_t, const at::cuda::CUDAStream&) {}, - [=](void*, size_t, size_t, const at::cuda::CUDAStream&) {}); - } -} - -c10::intrusive_ptr MooncakeBackend::reduce( - std::vector& tensors, const c10d::ReduceOptions& opts) { - TORCH_CHECK(tensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - auto tensor = tensors.back(); - size_t tensorSize = tensor.numel() * tensor.element_size(); - int64_t root = opts.rootRank + opts.rootTensor; - bool isRoot = (root == rank_); - if (isCpu_) { - auto numRanks = meta_->size; - return worker_->putTaskCpu( - c10d::OpType::REDUCE, tensorSize, root, meta_, connection_ctx_, - [=](void* dst, size_t pos, size_t realSize) { - memcpy(dst, (char*)tensor.data_ptr() + pos, realSize); - }, - [=, this](void* src, size_t pos, size_t realSize) { - if (isRoot) { - memset((char*)tensor.data_ptr() + pos, 0, realSize); - launchReduceCpu(tensor, pos, realSize, src, numRanks, - opts.reduceOp, meta_->activeRanks); - } - }); - } else { - auto stream = at::cuda::getCurrentCUDAStream(tensor.device().index()); - return worker_->putTaskCuda( - c10d::OpType::REDUCE, tensorSize, root, meta_, connection_ctx_, - stream, - [=](void* dst, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - cudaMemcpyAsync(dst, (char*)tensor.data_ptr() + pos, realSize, - cudaMemcpyDeviceToDevice, enq_stream); - }, - [=, this](void* src, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - if (isRoot) { - cudaMemsetAsync((char*)tensor.data_ptr() + pos, 0, realSize, - enq_stream); - launchReduceKernel(tensor, pos, realSize, src, meta_->size, - opts.reduceOp, meta_->activeRanksDevice, - enq_stream); - } - }); - } -} - -c10::intrusive_ptr MooncakeBackend::gather( - std::vector>& outputTensors, - std::vector& inputTensors, const c10d::GatherOptions& opts) { - int64_t root = opts.rootRank; - bool isRoot = (root == rank_); - TORCH_CHECK(inputTensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - if (isRoot) { - TORCH_CHECK(outputTensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - } - auto inputTensor = inputTensors.back(); - size_t tensorSize = inputTensor.numel() * inputTensor.element_size(); - if (isCpu_) { - return worker_->putTaskCpu( - c10d::OpType::GATHER, tensorSize, root, meta_, connection_ctx_, - [=](void* dst, size_t pos, size_t realSize) { - memcpy(dst, (char*)inputTensor.data_ptr() + pos, realSize); - }, - [=](void* src, size_t pos, size_t realSize) { - if (isRoot) { - auto outputTensors_ = outputTensors.back(); - for (const auto j : c10::irange(outputTensors_.size())) { - memcpy((char*)outputTensors_[j].data_ptr() + pos, - (char*)src + j * realSize, realSize); - } - } - }); - } else { - auto stream = - at::cuda::getCurrentCUDAStream(inputTensor.device().index()); - return worker_->putTaskCuda( - c10d::OpType::GATHER, tensorSize, root, meta_, connection_ctx_, - stream, - [=](void* dst, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - cudaMemcpyAsync(dst, (char*)inputTensor.data_ptr() + pos, - realSize, cudaMemcpyDeviceToDevice, enq_stream); - }, - [=](void* src, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - if (isRoot) { - auto outputTensors_ = outputTensors.back(); - for (const auto j : c10::irange(outputTensors_.size())) { - cudaMemcpyAsync( - (char*)outputTensors_[j].data_ptr() + pos, - (char*)src + j * realSize, realSize, - cudaMemcpyDeviceToDevice, enq_stream); - } - } - }); - } -} - -c10::intrusive_ptr MooncakeBackend::scatter( - std::vector& outputTensors, - std::vector>& inputTensors, - const c10d::ScatterOptions& opts) { - int64_t root = opts.rootRank; - bool isRoot = (root == rank_); - if (isRoot) { - TORCH_CHECK(inputTensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - } - TORCH_CHECK(outputTensors.size() == 1, MULTI_DEVICE_ERROR_MSG); - auto outputTensor = outputTensors.back(); - size_t tensorSize = outputTensor.numel() * outputTensor.element_size(); - if (isCpu_) { - return worker_->putTaskCpu( - c10d::OpType::SCATTER, tensorSize, root, meta_, connection_ctx_, - [=](void* dst, size_t pos, size_t realSize) { - if (isRoot) { - auto inputTensors_ = inputTensors.back(); - for (const auto j : c10::irange(inputTensors_.size())) { - memcpy((char*)dst + j * realSize, - (char*)inputTensors_[j].data_ptr() + pos, - realSize); - } - } - }, - [=](void* src, size_t pos, size_t realSize) { - memcpy((char*)outputTensor.data_ptr() + pos, src, realSize); - }); - } else { - auto stream = - at::cuda::getCurrentCUDAStream(outputTensor.device().index()); - return worker_->putTaskCuda( - c10d::OpType::SCATTER, tensorSize, root, meta_, connection_ctx_, - stream, - [=](void* dst, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - if (isRoot) { - auto inputTensors_ = inputTensors.back(); - for (const auto j : c10::irange(inputTensors_.size())) { - cudaMemcpyAsync( - (char*)dst + j * realSize, - (char*)inputTensors_[j].data_ptr() + pos, realSize, - cudaMemcpyDeviceToDevice, enq_stream); - } - } - }, - [=](void* src, size_t pos, size_t realSize, - const at::cuda::CUDAStream& enq_stream) { - cudaMemcpyAsync((char*)outputTensor.data_ptr() + pos, src, - realSize, cudaMemcpyDeviceToDevice, enq_stream); - }); - } -} - -void MooncakeBackend::shutdown() { - if (isShutdown_) { - return; - } - isShutdown_ = true; - - // If we encounter any hung operations, don't release resources - // to avoid potential crash. Instead, we allow those resources to leak - // and rely on the OS to reclaim them later. - bool has_hung_operation = false; - - // Phase 1: Drain P2P tasks - p2p_device_worker_->removeProxy(p2p_proxy_); - has_hung_operation |= !p2p_proxy_->drainTasks(); - - // Phase 2: Drain collective tasks for this backend - has_hung_operation |= !worker_->drainTasks(meta_.get()); - - // Phase 3: Drain warm-up transfers for connection poller - connection_ctx_->shutdown(); - if (connectionPollerRegistered_) { - ConnectionPoller::GetInstance().removeContext(connection_ctx_); - has_hung_operation |= !connection_ctx_->drainPoller(); - connectionPollerRegistered_ = false; - } - - // Phase 4: CUDA synchronization - if (!isCpu_ && !has_hung_operation) { - cudaDeviceSynchronize(); - } - - // Phase 5: Release resources if no hung operations - if (has_hung_operation) { - p2p_proxy_->abandonResources(); - connection_ctx_->abandonResources(); - } - - if (!has_hung_operation) { - for (size_t i = 0; i < 2; i++) { - engine_->unregisterLocalMemory(cpu_sync_send_region_[i]); - engine_->unregisterLocalMemory(cpu_sync_recv_region_[i]); - engine_->unregisterLocalMemory(send_buffer_[i]); - engine_->unregisterLocalMemory(recv_buffer_[i]); - delete[] cpu_sync_send_region_[i]; - delete[] cpu_sync_recv_region_[i]; - if (isCpu_) { - free(send_buffer_[i]); - free(recv_buffer_[i]); - } else { - cudaFree(send_buffer_[i]); - cudaFree(recv_buffer_[i]); - } - } - if (isCpu_) { - delete[] meta_->activeRanks; - } else { - cudaFreeHost(meta_->activeRanks); - } - meta_->activeRanks = nullptr; - meta_->activeRanksDevice = nullptr; - } -} - -void MooncakeBackend::syncActiveRanksTensor() { - std::vector active_ranks(meta_->size); - for (int i = 0; i < meta_->size; ++i) { - active_ranks[i] = meta_->activeRanks[i] ? 1 : 0; - } - - auto cpu_tensor = torch::tensor(active_ranks, torch::dtype(torch::kInt32)); - if (!meta_->activeRanksTensor.defined() || - meta_->activeRanksTensor.size(0) != meta_->size) { - meta_->activeRanksTensor = - cpu_tensor.to(isCpu_ ? torch::kCPU : torch::kCUDA); - return; - } - - if (meta_->activeRanksTensor.device().is_cpu()) { - meta_->activeRanksTensor.copy_(cpu_tensor); - } else { - meta_->activeRanksTensor.copy_( - cpu_tensor.to(meta_->activeRanksTensor.device())); - } -} - -void MooncakeBackend::publishLocalPeerMetadata() { - TORCH_CHECK(meta_->store, - "Publishing local peer metadata requires a valid Store."); - - std::vector rank_info_bytes(sizeof(SegmentInfo)); - memcpy(rank_info_bytes.data(), &rank_info, sizeof(SegmentInfo)); - - auto bufferKey = - ConnectionContext::getBufferStoreKey(meta_->backendIndex, rank_); - meta_->store->set(bufferKey, rank_info_bytes); - - auto serverNameKey = - ConnectionContext::getServerNameStoreKey(meta_->backendIndex, rank_); - meta_->store->set(serverNameKey, localServerName_); -} - -void MooncakeBackend::setLocalOnlyActiveRanks() { - for (int i = 0; i < meta_->size; ++i) { - meta_->activeRanks[i] = (i == meta_->rank); - } - syncActiveRanksTensor(); -} - -void MooncakeBackend::waitForExtensionState() { - TORCH_CHECK(meta_->store, "Recovery join requires a valid Store."); - - auto state_key = ConnectionContext::getExtensionStateStoreKey( - meta_->backendIndex, rank_); - - BackoffWaiter waiter( - BackoffWaiterConfig::constantSleep(std::chrono::milliseconds(50))); - - waiter.wait([&] { return meta_->store->check({state_key}); }); - - auto state_data = meta_->store->get(state_key); - auto state = deserialize(state_data); - - // taskCount - meta_->taskCount = state.taskCount; - - // p2pEpochs - TORCH_CHECK(static_cast(meta_->size) == state.p2pEpochs.size(), - "Invalid p2pEpochs size"); - for (int i = 0; i < meta_->size; ++i) { - p2p_proxy_->setEpoch(i, state.p2pEpochs[i]); - } - - // activeRanks - TORCH_CHECK(static_cast(meta_->size) == state.activeRanks.size(), - "Invalid activeRanks"); - for (int i = 0; i < meta_->size; ++i) { - meta_->activeRanks[i] = state.activeRanks[i]; - } - syncActiveRanksTensor(); - - // activeSize: count the number of active ranks (contiguous from 0) - int newActiveSize = 0; - for (int i = 0; i < meta_->size; ++i) { - if (meta_->activeRanks[i]) { - newActiveSize = i + 1; - } - } - meta_->activeSize = newActiveSize; -} - -int MooncakeBackend::getNumSyncedRanks() { - std::vector tensors; - tensors.emplace_back(torch::tensor( - connection_ctx_->getTotalConnectedPeers(), - torch::dtype(torch::kInt).device(isCpu_ ? torch::kCPU : torch::kCUDA))); - c10d::AllreduceOptions opts{ - .reduceOp = c10d::ReduceOp::MIN, - }; - auto work = allreduce(tensors, opts); - work->wait(); - if (!isCpu_) { - auto stream = - at::cuda::getCurrentCUDAStream(tensors[0].device().index()); - cudaStreamSynchronize(stream); - } - return tensors[0].cpu().item(); -} - -void MooncakeBackend::extendGroupSizeTo(int newSize) { - const int oldSize = meta_->size; - const int oldActiveSize = meta_->activeSize; - if (newSize == oldSize) return; - - TORCH_CHECK(newSize >= 0 && static_cast(newSize) < kMaxNumRanks, - "Size out of range"); - TORCH_CHECK(newSize >= oldSize, "newSize < oldSize"); - - LOG(INFO) << "Backend " << backendIndex_ << " rank " << rank_ - << ": Group size extend to " << newSize; - - meta_->size = newSize; - meta_->activeSize = newSize; - meta_->taskCount = 0; - - // Initialize new rank's metadata - for (int i = oldSize; i < newSize; ++i) { - local2global_rank_map_[i] = i; - // IMPORTANT: Newly-extended ranks must start as inactive. - // They will only participate in collectives after healthy ranks - // explicitly activate them via recoverRanks(). This enables a - // two-phase scale-up protocol (extend capacity -> poll readiness - // -> recover/activate) and avoids collectives including ranks that - // haven't joined yet. - meta_->activeRanks[i] = false; - } - - auto& tensor = meta_->activeRanksTensor; - if (newSize > tensor.numel()) { - tensor.resize_({newSize}); - } - tensor.slice(0, oldSize, newSize).fill_(0); - - connection_ctx_->extendGroupSizeTo(newSize); - p2p_proxy_->extendGroupSizeTo(newSize); - // After extendGroupSizeTo, we don't `waitUntilNewRanksConnected` here - // but do it in the first task. This enables client code to overlap - // execution between `extendGroupSizeTo` and the first communication call. -} - -std::vector MooncakeBackend::getPeerState(const std::vector& ranks) { - bool activeRanksBackup[kMaxNumRanks]; - while (true) { - std::vector input; - for (const int rank : ranks) { - TORCH_CHECK(rank >= 0 && static_cast(rank) < kMaxNumRanks, - "Rank out of range"); - input.push_back(meta_->peerConnected[rank]); - } - for (int i = 0; i < meta_->size; i++) { - activeRanksBackup[i] = meta_->activeRanks[i]; - } - - std::vector tensors; - tensors.emplace_back(torch::tensor( - input, torch::dtype(torch::kInt) - .device(isCpu_ ? torch::kCPU : torch::kCUDA))); - c10d::AllreduceOptions opts{ - .reduceOp = c10d::ReduceOp::MIN, - }; - auto work = allreduce(tensors, opts); - work->wait(); - if (!isCpu_) { - auto stream = - at::cuda::getCurrentCUDAStream(tensors[0].device().index()); - cudaStreamSynchronize(stream); - } - bool activeRanksChanged = false; - for (int i = 0; i < meta_->size; i++) { - if (activeRanksBackup[i] != meta_->activeRanks[i]) { - activeRanksChanged = true; - break; - } - } - - if (!activeRanksChanged) { - std::vector output; - for (int i = 0; i < tensors[0].size(0); ++i) { - output.push_back(tensors[0].cpu()[i].item() != 0); - } - return output; - } - } -} - -void MooncakeBackend::recoverRanks(const std::vector& ranks) { - TORCH_CHECK(meta_->store, "Rank recovery requires a valid Store."); - - for (const int rank : ranks) { - TORCH_CHECK(rank >= 0 && static_cast(rank) < kMaxNumRanks, - "Rank out of range"); - TORCH_CHECK(meta_->peerConnected[rank]); - meta_->activeRanks[rank] = true; - } - - // Expand activeSize if any recovered rank is beyond the current boundary. - if (!ranks.empty()) { - const int max_rank = *std::max_element(ranks.begin(), ranks.end()); - if (max_rank >= meta_->activeSize) { - meta_->activeSize = max_rank + 1; - } - } - - syncActiveRanksTensor(); - std::vector epochs(meta_->size); - for (int i = 0; i < meta_->size; ++i) { - epochs[i] = p2p_proxy_->getEpoch(i); - } - ExtensionState state{ - .activeRanks = - std::vector(meta_->activeRanks, meta_->activeRanks + meta_->size), - .p2pEpochs = std::move(epochs), - .taskCount = meta_->taskCount}; - auto state_data = serialize(state); - for (const int rank : ranks) { - auto key = ConnectionContext::getExtensionStateStoreKey( - meta_->backendIndex, rank); - meta_->store->set(key, state_data); - } -} - -void MooncakeBackend::joinGroup() { - TORCH_CHECK(options_ && options_->isExtension_, - "joinGroup is only valid for extension backends."); - connection_ctx_->setDummy(false); - publishLocalPeerMetadata(); - if (!connectionPollerRegistered_) { - ConnectionPoller::GetInstance().registerContext(connection_ctx_); - connectionPollerRegistered_ = true; - } - connection_ctx_->waitUntilAllConnected(); - waitForExtensionState(); -} - -void MooncakeBackend::setExternalEngine(TransferEngine* engine) { - externalEngine_ = engine; - if (engine) { - LOG(INFO) << "MooncakeBackend: external TransferEngine set (ptr=" - << engine << ")"; - } -} -} // namespace mooncake diff --git a/mooncake-pg/src/mooncake_communicator.cpp b/mooncake-pg/src/mooncake_communicator.cpp new file mode 100644 index 0000000000..647b7361d0 --- /dev/null +++ b/mooncake-pg/src/mooncake_communicator.cpp @@ -0,0 +1,1627 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "common.h" +#include "error_types.h" +#include "gpu_runtime.h" +#include "memory_location.h" + +namespace mooncake { +namespace { + +// A non-zero operation size is required to ensure that the worker creates a +// task for the barrier. +constexpr size_t kBarrierDummySize = 1; + +void copyDeviceToDevice(void* dst, const void* src, size_t bytes, + cudaStream_t stream) { + PG_ASSERT_CUDA( + cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, stream)); +} + +PGResult checkBuffer(const void* buffer, size_t bytes, const char* name) { + PG_VALIDATE_ARG(buffer || bytes == 0, std::string(name) + " is null"); + return {}; +} + +PGResult checkRoot(int root, int max_group_size, const char* operation) { + PG_VALIDATE_ARG(root >= 0 && root < max_group_size, + std::string(operation) + " root is out of range"); + return {}; +} + +PGResult checkP2PPeer(const TransferGroupMeta& meta, int peer, + int max_group_size, const char* operation) { + PG_VALIDATE_ARG(peer >= 0 && peer < max_group_size, + std::string(operation) + " peer is out of range"); + // P2P may target inactive members, but reserved extension slots are not + // valid targets until they have an assigned global-rank mapping. + PG_VALIDATE_ARG( + meta.rank_order[peer] != kInvalidGlobalRank, + std::string(operation) + " peer is not assigned in this group"); + return {}; +} + +PGResult getByteCount(size_t count, DataType datatype) { + switch (datatype) { + case DataType::Int8: + case DataType::Uint8: + case DataType::Int16: + case DataType::Uint16: + case DataType::Int32: + case DataType::Uint32: + case DataType::Int64: + case DataType::Uint64: + case DataType::Float16: + case DataType::Float32: + case DataType::Float64: + case DataType::Bfloat16: + case DataType::Bool: + case DataType::Float8e4m3fn: + case DataType::Float8e5m2: + case DataType::Float8e4m3fnuz: + case DataType::Float8e5m2fnuz: + case DataType::Float8e8m0fnu: + break; + default: + return makePGError(PGErrorCode::InvalidArgument, + "unsupported Mooncake PG datatype"); + } + + const size_t element_size = elementSize(datatype); + PG_VALIDATE_ARG(count <= std::numeric_limits::max() / element_size, + "element count overflows size_t"); + return count * element_size; +} + +PGResult checkReduction(DataType datatype, ReduceOp op, bool is_cpu) { + switch (op) { + case ReduceOp::Sum: + case ReduceOp::Product: + case ReduceOp::Min: + case ReduceOp::Max: + break; + case ReduceOp::Avg: + return makePGError(PGErrorCode::NotSupported, + "average reduction is not supported"); + default: + return makePGError(PGErrorCode::NotSupported, + "reduction operation is not supported"); + } + + switch (datatype) { + case DataType::Uint8: + case DataType::Int8: + case DataType::Int16: + case DataType::Int32: + case DataType::Int64: + case DataType::Float32: + case DataType::Float64: + case DataType::Bool: + break; + case DataType::Bfloat16: + if (!is_cpu) break; + [[fallthrough]]; + default: + return makePGError(PGErrorCode::NotSupported, + "reduction datatype is not supported"); + } + + return {}; +} + +} // namespace + +MooncakePGContext::~MooncakePGContext() { + try { + auto result = shutdown(); + if (!result.has_value()) { + LOG(ERROR) + << "Mooncake PG context shutdown failed during destruction: " + << result.error().message; + } + } catch (const std::exception& error) { + LOG(ERROR) << "Mooncake PG context shutdown failed during destruction: " + << error.what(); + } catch (...) { + LOG(ERROR) << "Mooncake PG context shutdown failed during destruction"; + } +} + +PGResult MooncakePGContext::checkRunning() const { + PG_VALIDATE_STATE(!shutdown_requested_, + "Mooncake PG context is shutting down"); + return {}; +} + +PGResult MooncakePGContext::initialize(int rank, int world_size) { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + PG_VALIDATE_ARG(world_size > 0 && world_size <= kMaxNumRanks, + "max_world_size is outside the supported range"); + PG_VALIDATE_ARG(rank >= 0 && rank < world_size, + "global rank is outside the process world"); + if (initialized_) { + PG_VALIDATE_STATE( + global_rank == rank && max_world_size == world_size, + "Mooncake process context was initialized with a different rank " + "or world size"); + return {}; + } + + // Ordering constraint: AgentHost::start() sends registerAgent immediately, + // which includes LinkManager's localServerName() and getWarmupRecvAddr(). + // These must be non-empty, so the engine and LinkManager must be + // initialized before connectCoordinator starts the AgentHost. + if (!engine_initialized) { +#ifdef USE_MACA + PG_VALIDATE_ARG(std::getenv("MC_MACA_HOST_TRANSPORT") != nullptr, + "MACA PG requires MC_MACA_HOST_TRANSPORT=1"); +#endif + PG_TRY_TE(engine->init(P2PHANDSHAKE, host_ip)); + engine_initialized = true; + } + if (!link_manager.isInitialized()) { + PG_TRY(link_manager.init(rank, world_size, engine)); + } + + global_rank = rank; + max_world_size = world_size; + initialized_ = true; + return {}; +} + +PGResult MooncakePGContext::launchCoordinator() { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + PG_VALIDATE_STATE( + initialized_, + "Mooncake PG context must be initialized before launching the " + "coordinator"); + PG_VALIDATE_STATE(global_rank == 0, + "only global rank 0 may start the coordinator"); + if (!coordinator_host) { + auto candidate = std::make_unique( + host_ip, max_world_size, fault_reconciliation_window_us); + PG_TRY(candidate->start()); + coordinator_host = std::move(candidate); + } + const auto& address = coordinator_host->getListenAddr(); + if (address.empty()) { + return makePGError(PGErrorCode::SystemError, + "coordinator returned an empty address"); + } + return address; +} + +PGResult MooncakePGContext::connectCoordinator( + const std::string& coordinator_address) { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + PG_VALIDATE_STATE( + initialized_, + "Mooncake PG context must be initialized before connecting to the " + "coordinator"); + PG_VALIDATE_ARG(!coordinator_address.empty(), + "coordinator address must not be empty"); + if (!agent_host) { + auto candidate = std::make_unique( + coordinator_address, host_ip, global_rank, max_world_size, + link_manager, fault_reconciliation_window_us); + PG_TRY(candidate->start()); + agent_host = std::move(candidate); + } + return {}; +} + +PGResult MooncakePGContext::setHostIp(std::string value) { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + PG_VALIDATE_ARG(!value.empty(), "host IP must not be empty"); + PG_VALIDATE_STATE(!initialized_ || host_ip == value, + "host IP cannot be changed after context initialization"); + if (!initialized_) host_ip = std::move(value); + return {}; +} + +PGResult MooncakePGContext::setExternalEngine( + TransferEngine* transfer_engine) { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + auto* requested_engine = + transfer_engine ? transfer_engine : owned_engine.get(); + PG_VALIDATE_STATE( + !initialized_ || engine == requested_engine, + "transfer engine cannot be changed after context initialization"); + if (!initialized_) { + engine = requested_engine; + engine_initialized = transfer_engine != nullptr; + if (transfer_engine) { + const auto endpoint = engine->getLocalIpAndPort(); + const auto derived_host = getHostNameWithoutPort(endpoint); + PG_VALIDATE_STATE( + !derived_host.empty(), + "set_transfer_engine requires an initialized TransferEngine " + "with a local endpoint"); + host_ip = derived_host; + } + } + return {}; +} + +PGResult MooncakePGContext::setDeviceFilter( + std::vector filters) { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + std::sort(filters.begin(), filters.end()); + filters.erase(std::unique(filters.begin(), filters.end()), filters.end()); + PG_VALIDATE_STATE( + !initialized_ || device_filters_ == filters, + "device filters cannot be changed after context initialization"); + if (!initialized_) { + device_filters_ = filters; + engine->setWhitelistFilters(std::move(filters)); + } + return {}; +} + +PGResult MooncakePGContext::setCollectiveTimeout(size_t timeout_us) { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + collective_timeout_us = timeout_us; + return {}; +} + +PGResult MooncakePGContext::setP2PTimeout(int64_t timeout_us) { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + PG_VALIDATE_ARG(timeout_us >= 0, "P2P timeout must not be negative"); + p2p_timeout_us = timeout_us; + return {}; +} + +PGResult MooncakePGContext::setFaultReconciliationWindow( + int64_t timeout_us) { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + PG_VALIDATE_ARG(timeout_us >= 0, + "fault reconciliation window must not be negative"); + fault_reconciliation_window_us = timeout_us; + if (agent_host) { + agent_host->setFaultReconciliationWindow(timeout_us); + } + if (coordinator_host) { + PG_TRY(coordinator_host->setFaultReconciliationWindow(timeout_us)); + } + return {}; +} + +PGResult MooncakePGContext::incrementCommUseCount() { + std::lock_guard lock(state_mutex_); + PG_TRY(checkRunning()); + PG_VALIDATE_STATE( + initialized_, + "Mooncake PG context must be initialized before creating a " + "communicator"); + PG_VALIDATE_STATE( + agent_host, + "Mooncake PG context must connect to a coordinator before creating " + "a communicator"); + ++comm_use_count_; + return {}; +} + +void MooncakePGContext::decrementCommUseCount() noexcept { + std::lock_guard lock(state_mutex_); + if (comm_use_count_ == 0) { + LOG(ERROR) << "Mooncake PG communicator use count underflow"; + return; + } + --comm_use_count_; +} + +PGResult MooncakePGContext::shutdown() { + { + std::lock_guard lock(state_mutex_); + if (shutdown_requested_) return {}; + if (comm_use_count_ != 0) { + return makePGError( + PGErrorCode::ResourceBusy, + "Mooncake PG context still has active communicators"); + } + shutdown_requested_ = true; + } + + if (agent_host) agent_host->shutdown(); + link_manager.shutdown(); + agent_host.reset(); + if (coordinator_host) coordinator_host->shutdown(); + coordinator_host.reset(); + engine = nullptr; + return {}; +} + +/** + * @brief Initialize Mooncake communicator state + */ +PGResult> MooncakeCommunicator::create( + MooncakePGContext& context, MooncakeCommunicatorConfig config) { + PG_VALIDATE_STATE(context.engine, + "Mooncake PG context has no transfer engine"); + PG_VALIDATE_STATE( + context.agent_host, + "Mooncake PG context must connect to a coordinator before creating " + "a communicator"); + + auto communicator = std::unique_ptr( + new MooncakeCommunicator(context, config)); + PG_TRY(communicator->initialize(std::move(config))); + return communicator; +} + +MooncakeCommunicator::MooncakeCommunicator( + MooncakePGContext& context, const MooncakeCommunicatorConfig& config) + : context_(context), + agent_(*context_.agent_host), + rank_(config.rank), + initial_size_(config.size), + max_group_size_(config.max_group_size > 0 ? config.max_group_size + : config.size), + device_index_(config.device_index), + is_cpu_(config.is_cpu), + active_ranks_mirror_(config.active_ranks_mirror), + active_ranks_mirror_is_device_(config.active_ranks_mirror_is_device), + active_ranks_mirror_device_index_( + config.active_ranks_mirror_device_index) {} + +PGResult MooncakeCommunicator::initialize( + MooncakeCommunicatorConfig config) { + PG_VALIDATE_ARG(initial_size_ > 0 && initial_size_ <= max_group_size_, + "group size exceeds max_group_size"); + PG_VALIDATE_ARG(max_group_size_ > 0 && max_group_size_ <= kMaxNumRanks, + "max_group_size is outside the supported range"); + PG_VALIDATE_ARG(rank_ >= 0 && rank_ < initial_size_, + "rank is outside the initial group"); + PG_VALIDATE_ARG(!config.group_bootstrap_id.empty(), + "group bootstrap id must not be empty"); + PG_VALIDATE_ARG( + !config.auto_sync_on_failure || config.auto_deactivate_on_failure, + "auto_sync_on_failure requires auto_deactivate_on_failure"); + PG_VALIDATE_ARG( + !active_ranks_mirror_ || config.active_ranks_mirror_count >= + static_cast(max_group_size_), + "active-ranks mirror is too small"); + PG_VALIDATE_ARG(!active_ranks_mirror_ || !active_ranks_mirror_is_device_ || + active_ranks_mirror_device_index_ >= 0, + "device active-ranks mirror requires a valid device index"); + + PG_VALIDATE_ARG( + config.global_ranks.size() == static_cast(initial_size_), + "global rank count must equal communicator size"); + std::array seen_global_ranks{}; + for (const auto global_rank : config.global_ranks) { + PG_VALIDATE_ARG( + global_rank >= 0 && global_rank < context_.max_world_size, + "global rank is outside the process world"); + PG_VALIDATE_ARG(!seen_global_ranks[global_rank], + "global ranks contains duplicates"); + seen_global_ranks[global_rank] = true; + } + PG_VALIDATE_ARG(config.global_ranks[rank_] == context_.global_rank, + "communicator rank does not map to the process global " + "rank"); + auto initial_rank_order = std::move(config.global_ranks); + + // Memory location for device-specific buffers. Always kWildcardLocation for + // a CPU communicator. + std::unique_ptr device_guard; + std::string location = kWildcardLocation; + if (!is_cpu_) { + if (device_index_ < 0) { + PG_TRY_CUDA(cudaGetDevice(&device_index_)); + } + device_guard = std::make_unique(device_index_); + location = GPU_PREFIX + std::to_string(device_index_); + } + if (active_ranks_mirror_ && active_ranks_mirror_is_device_) { + active_ranks_mirror_stream_ = + GpuStream::createNonBlocking(active_ranks_mirror_device_index_); + } + + // Register collective buffers. + for (size_t index = 0; index < 2; ++index) { + if (is_cpu_) { + send_buffer_[index] = std::malloc(kBufferSize); + recv_buffer_[index] = std::malloc(kBufferSize); + PG_ASSERT(send_buffer_[index] && recv_buffer_[index], + "failed to allocate CPU collective buffers"); + } else { + PG_TRY_CUDA(cudaMalloc(&send_buffer_[index], kBufferSize)); + PG_TRY_CUDA(cudaMalloc(&recv_buffer_[index], kBufferSize)); + } + PG_TRY_TE(context_.engine->registerLocalMemory(send_buffer_[index], + kBufferSize, location)); + PG_TRY_TE(context_.engine->registerLocalMemory(recv_buffer_[index], + kBufferSize, location)); + + // Register CPU synchronization regions. + cpu_sync_send_region_[index] = new int32_t[kMaxNumRanks]{}; + cpu_sync_recv_region_[index] = new int32_t[kMaxNumRanks]{}; + PG_TRY_TE(context_.engine->registerLocalMemory( + cpu_sync_send_region_[index], kMaxNumRanks * sizeof(int32_t), + kWildcardLocation)); + PG_TRY_TE(context_.engine->registerLocalMemory( + cpu_sync_recv_region_[index], kMaxNumRanks * sizeof(int32_t), + kWildcardLocation)); + } + + if (is_cpu_) { + p2p_device_worker_ = + context_.p2p_device_worker_manager.getCPUWorker(context_.engine); + worker_ = context_.worker_manager.GetCPUWorker(); + } else { + p2p_device_worker_ = context_.p2p_device_worker_manager.getCUDAWorker( + device_index_, context_.engine); + worker_ = context_.worker_manager.GetCUDAWorker(device_index_); + preloadReduceKernels(); + } + worker_->Start(); + + p2p_proxy_ = std::make_shared( + context_.engine, + P2PProxy::Options{.is_cpu = is_cpu_, + .rank = rank_, + .size = max_group_size_, + .cuda_device_index = device_index_, + .p2p_timeout_us = &context_.p2p_timeout_us}); + p2p_device_worker_->registerProxy(p2p_proxy_); + + meta_ = std::make_shared(); + for (int index = 0; index < kMaxNumRanks; ++index) { + meta_->segmentIDs[index] = static_cast(-1); + meta_->rankEpochs[index] = 0; + meta_->rankStates[index] = RankState::Offline; + meta_->rank_order[index] = kInvalidGlobalRank; + } + meta_->rank = rank_; + meta_->globalRank = initial_rank_order[rank_]; + for (int index = 0; index < initial_size_; ++index) { + meta_->rank_order[index] = initial_rank_order[index]; + } + meta_->maxGroupSize = max_group_size_; // slot capacity + meta_->activeSize.store(initial_size_, std::memory_order_relaxed); + meta_->taskCount = 0; + meta_->collectiveTimeoutUs = &context_.collective_timeout_us; + meta_->engine = context_.engine; + meta_->communicator = this; + meta_->autoSyncOnFailure = config.auto_sync_on_failure; + p2p_proxy_->bindMeta(meta_); + + // Active ranks will be filled by applyViewUpdate, so only allocate their + // storage here. + meta_->maybeActivatable = new bool[max_group_size_]{}; + if (is_cpu_) { + meta_->activeRanks = new bool[max_group_size_]{}; + meta_->activeRanksDevice = meta_->activeRanks; + } else { + PG_TRY_CUDA(cudaHostAlloc(&meta_->activeRanks, + max_group_size_ * sizeof(bool), + cudaHostAllocMapped)); + PG_TRY_CUDA(cudaHostGetDevicePointer(&meta_->activeRanksDevice, + meta_->activeRanks, 0)); + std::fill_n(meta_->activeRanks, max_group_size_, false); + } + + // Initial local endpoint info. + meta_->segmentInfos[rank_] = GroupEndpointInfo{ + .send_buffer = {reinterpret_cast(send_buffer_[0]), + reinterpret_cast(send_buffer_[1])}, + .recv_buffer = {reinterpret_cast(recv_buffer_[0]), + reinterpret_cast(recv_buffer_[1])}, + .send_sync = {reinterpret_cast(cpu_sync_send_region_[0]), + reinterpret_cast(cpu_sync_send_region_[1])}, + .recv_sync = {reinterpret_cast(cpu_sync_recv_region_[0]), + reinterpret_cast(cpu_sync_recv_region_[1])}, + .p2p_credit_region = + reinterpret_cast(p2p_proxy_->credit_region()), + .p2p_ack_region = reinterpret_cast(p2p_proxy_->ack_region()), + }; + + // Control Plane Initialization + + // Wait for Agent registration. + PG_TRY(agent_.waitUntilRegistered(std::chrono::seconds(30))); + + // The PyTorch-provided group id is only a bootstrap id. The Coordinator + // resolves it together with rank order into a process-lifetime GroupId. CPU + // and device communicators use independent namespaces. + auto bootstrap_id = std::string(is_cpu_ ? "cpu:" : "device:") + + std::move(config.group_bootstrap_id); + + // Register this group with the Agent, publish the local endpoint, and block + // until the Coordinator says it is ready. Group registration is + // synchronous. + auto group_result = agent_.registerGroup( + std::move(bootstrap_id), max_group_size_, std::move(initial_rank_order), + config.group_resolve_policy, config.auto_deactivate_on_failure, this); + if (!group_result.has_value()) { + return makePGError(std::move(group_result).error()); + } + meta_->group_id = std::move(group_result.value()); + + if (!isValidGroup()) { + // Registration rejection is scoped to this communicator. Keep the Agent + // and every other group untouched, and use the pre-join local-only + // collective behavior with an effective {self} membership. + std::fill_n(meta_->activeRanks, max_group_size_, false); + meta_->activeRanks[rank_] = true; + meta_->autoSyncOnFailure = false; + syncActiveRanksMirror(); + refreshSegmentID(rank_); + LOG(WARNING) << "Mooncake communicator rank=" << meta_->globalRank + << " is using local-only execution because group " + "registration was rejected"; + return {}; + } + + PG_TRY(agent_.publishLocalEndpoint(buildEndpointMetadata())); + PG_TRY( + agent_.waitUntilGroupReady(meta_->group_id, std::chrono::seconds(300))); + + // Initialize all peer segment IDs from the LinkManager. Subsequent updates + // (endpoint changes, disconnects) are handled by NotifyLinkRefreshed. + for (int local = 0; local < max_group_size_; ++local) { + refreshSegmentID(local); + } + return {}; +} + +MooncakeCommunicator::~MooncakeCommunicator() { + try { + auto result = shutdown(); + if (!result.has_value()) { + LOG(ERROR) << "Mooncake communicator shutdown failed: " + << result.error().message; + } + } catch (const std::exception& error) { + LOG(ERROR) << "Mooncake communicator shutdown failed: " << error.what(); + } +} + +int MooncakeCommunicator::getSize() const { + if (!meta_ || meta_->extensionMode.load(std::memory_order_acquire) != + CollectiveExtensionState::Normal) { + return initial_size_; + } + return meta_->activeSize.load(std::memory_order_acquire); +} + +PGResult MooncakeCommunicator::checkOpState(OpType op) const { + PG_VALIDATE_STATE(!is_shutdown_, "communicator is shut down"); + PG_ASSERT(meta_, "initialized communicator has no group metadata"); + const auto mode = meta_->extensionMode.load(std::memory_order_acquire); + if (isValidGroup()) { + PG_VALIDATE_STATE( + meta_->rankStates[meta_->globalRank] != RankState::Offline, + "rank " + std::to_string(meta_->globalRank) + + " is offline and cannot perform operations"); + } + // P2P operations don't require the rank to be active in the group. + const bool is_p2p = op == OpType::Send || op == OpType::Recv; + if (!isValidGroup() && is_p2p) { + return makePGError(PGErrorCode::NotSupported, + "P2P is unavailable for an invalid Mooncake group"); + } + if (!is_p2p) { + PG_VALIDATE_STATE(mode != CollectiveExtensionState::Quiescing, + "rank is quiescing and cannot issue collectives"); + PG_VALIDATE_STATE(mode == CollectiveExtensionState::Isolated || + meta_->activeRanks[rank_], + "rank is not active in this group"); + } + return {}; +} + +PGResult MooncakeCommunicator::initializeFailedRanksHint( + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) const { + PG_VALIDATE_ARG(failed_ranks_hint, "failed-ranks hint is null"); + PG_VALIDATE_ARG( + failed_ranks_hint_count >= static_cast(max_group_size_), + "failed-ranks hint buffer is too small"); + std::fill_n(failed_ranks_hint, max_group_size_, int32_t{0}); + return {}; +} + +PGResult> MooncakeCommunicator::sendCpu( + const void* buffer, size_t count, DataType datatype, int peer, + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "sendCpu requires a CPU communicator"); + return enqueueSend(buffer, count, datatype, peer, nullptr, + failed_ranks_hint, failed_ranks_hint_count); +} + +PGResult> MooncakeCommunicator::sendGpu( + const void* buffer, size_t count, DataType datatype, int peer, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "sendGpu requires a GPU communicator"); + return enqueueSend(buffer, count, datatype, peer, stream, failed_ranks_hint, + failed_ranks_hint_count); +} + +PGResult> MooncakeCommunicator::enqueueSend( + const void* buffer, size_t count, DataType datatype, int peer, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_TRY(checkOpState(OpType::Send)); + PG_TRY(auto bytes, getByteCount(count, datatype)); + PG_VALIDATE_ARG(buffer || bytes == 0, "send buffer is null"); + PG_TRY(checkP2PPeer(*meta_, peer, max_group_size_, "P2P send")); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + auto completion = std::make_shared>(); + auto future = completion->get_future().share(); + auto result = std::make_unique(std::move(future)); + p2p_proxy_->enqueueSend(P2PProxy::SendOp{ + .buffer_ = buffer, + .size_ = bytes, + .peer_rank_ = peer, + .cuda_stream_ = stream, + .completion_ = completion, + .failed_ranks_hint_ = failed_ranks_hint, + }); + return result; +} + +PGResult> MooncakeCommunicator::recvCpu( + void* buffer, size_t count, DataType datatype, int peer, + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "recvCpu requires a CPU communicator"); + return enqueueRecv(buffer, count, datatype, peer, nullptr, + failed_ranks_hint, failed_ranks_hint_count); +} + +PGResult> MooncakeCommunicator::recvGpu( + void* buffer, size_t count, DataType datatype, int peer, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "recvGpu requires a GPU communicator"); + return enqueueRecv(buffer, count, datatype, peer, stream, failed_ranks_hint, + failed_ranks_hint_count); +} + +PGResult> MooncakeCommunicator::enqueueRecv( + void* buffer, size_t count, DataType datatype, int peer, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_TRY(checkOpState(OpType::Recv)); + PG_TRY(auto bytes, getByteCount(count, datatype)); + PG_VALIDATE_ARG(buffer || bytes == 0, "recv buffer is null"); + PG_TRY(checkP2PPeer(*meta_, peer, max_group_size_, "P2P recv")); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + auto completion = std::make_shared>(); + auto future = completion->get_future().share(); + auto result = std::make_unique(std::move(future)); + p2p_proxy_->enqueueRecv(P2PProxy::RecvOp{ + .buffer_ = buffer, + .size_ = bytes, + .peer_rank_ = peer, + .cuda_stream_ = stream, + .completion_ = completion, + .failed_ranks_hint_ = failed_ranks_hint, + }); + return result; +} + +PGResult> MooncakeCommunicator::broadcastCpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + int root, int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "broadcastCpu requires a CPU communicator"); + PG_TRY(checkOpState(OpType::Broadcast)); + PG_TRY(auto bytes, getByteCount(count, datatype)); + PG_TRY(checkRoot(root, max_group_size_, "broadcast")); + const bool is_root = root == rank_; + if (is_root) { + PG_TRY(checkBuffer(send_buffer, bytes, "send buffer")); + } + PG_TRY(checkBuffer(recv_buffer, bytes, "receive buffer")); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + return worker_->putTaskCpu( + OpType::Broadcast, bytes, root, meta_, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size) { + if (is_root) { + std::memcpy(dst, static_cast(send_buffer) + pos, + size); + } + }, + [=](void* src, size_t pos, size_t size) { + std::memcpy(static_cast(recv_buffer) + pos, src, size); + }); +} + +PGResult MooncakeCommunicator::broadcastGpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + int root, cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "broadcastGpu requires a GPU communicator"); + PG_TRY(checkOpState(OpType::Broadcast)); + PG_TRY(auto bytes, getByteCount(count, datatype)); + PG_TRY(checkRoot(root, max_group_size_, "broadcast")); + const bool is_root = root == rank_; + if (is_root) { + PG_TRY(checkBuffer(send_buffer, bytes, "send buffer")); + } + PG_TRY(checkBuffer(recv_buffer, bytes, "receive buffer")); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + worker_->putTaskCuda( + OpType::Broadcast, bytes, root, meta_, stream, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size, cudaStream_t enqueue_stream) { + if (is_root) { + copyDeviceToDevice(dst, + static_cast(send_buffer) + pos, + size, enqueue_stream); + } + }, + [=](void* src, size_t pos, size_t size, cudaStream_t enqueue_stream) { + copyDeviceToDevice(static_cast(recv_buffer) + pos, src, size, + enqueue_stream); + }); + return {}; +} + +PGResult> MooncakeCommunicator::allReduceCpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + ReduceOp op, int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "allReduceCpu requires a CPU communicator"); + PG_TRY(checkOpState(OpType::AllReduce)); + PG_TRY(auto bytes, getByteCount(count, datatype)); + PG_TRY(checkBuffer(send_buffer, bytes, "send buffer")); + PG_TRY(checkBuffer(recv_buffer, bytes, "receive buffer")); + PG_TRY(checkReduction(datatype, op, true)); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + return worker_->putTaskCpu( + OpType::AllReduce, bytes, 0, meta_, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size) { + std::memcpy(dst, static_cast(send_buffer) + pos, size); + }, + [=, this](void* src, size_t pos, size_t size) { + std::memset(static_cast(recv_buffer) + pos, 0, size); + launchReduceCpu(recv_buffer, datatype, pos, size, src, active_size, + op, meta_->activeRanks); + }); +} + +PGResult MooncakeCommunicator::allReduceGpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + ReduceOp op, cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "allReduceGpu requires a GPU communicator"); + PG_TRY(checkOpState(OpType::AllReduce)); + PG_TRY(auto bytes, getByteCount(count, datatype)); + PG_TRY(checkBuffer(send_buffer, bytes, "send buffer")); + PG_TRY(checkBuffer(recv_buffer, bytes, "receive buffer")); + PG_TRY(checkReduction(datatype, op, false)); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + worker_->putTaskCuda( + OpType::AllReduce, bytes, 0, meta_, stream, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size, cudaStream_t enqueue_stream) { + copyDeviceToDevice(dst, static_cast(send_buffer) + pos, + size, enqueue_stream); + }, + [=, this](void* src, size_t pos, size_t size, + cudaStream_t enqueue_stream) { + PG_ASSERT_CUDA( + cudaMemsetAsync(static_cast(recv_buffer) + pos, 0, size, + enqueue_stream)); + launchReduceKernel(recv_buffer, datatype, pos, size, src, + active_size, op, meta_->activeRanksDevice, + enqueue_stream); + }); + return {}; +} + +PGResult> MooncakeCommunicator::allGatherCpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "allGatherCpu requires a CPU communicator"); + PG_TRY(checkOpState(OpType::AllGather)); + PG_TRY(auto send_bytes, getByteCount(count, datatype)); + PG_TRY(checkBuffer(send_buffer, send_bytes, "send buffer")); + PG_TRY(checkBuffer(recv_buffer, send_bytes, "receive buffer")); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + return worker_->putTaskCpu( + OpType::AllGather, send_bytes, 0, meta_, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size) { + std::memcpy(dst, static_cast(send_buffer) + pos, size); + }, + [=, this](void* src, size_t pos, size_t size) { + for (int peer = 0; peer < active_size; ++peer) { + if (!meta_->activeRanks[peer]) continue; + std::memcpy( + static_cast(recv_buffer) + peer * send_bytes + pos, + static_cast(src) + peer * size, size); + } + }); +} + +PGResult MooncakeCommunicator::allGatherGpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "allGatherGpu requires a GPU communicator"); + PG_TRY(checkOpState(OpType::AllGather)); + PG_TRY(auto send_bytes, getByteCount(count, datatype)); + PG_TRY(checkBuffer(send_buffer, send_bytes, "send buffer")); + PG_TRY(checkBuffer(recv_buffer, send_bytes, "receive buffer")); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + worker_->putTaskCuda( + OpType::AllGather, send_bytes, 0, meta_, stream, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size, cudaStream_t enqueue_stream) { + copyDeviceToDevice(dst, static_cast(send_buffer) + pos, + size, enqueue_stream); + }, + [=, this](void* src, size_t pos, size_t size, + cudaStream_t enqueue_stream) { + for (int peer = 0; peer < active_size; ++peer) { + if (!meta_->activeRanks[peer]) continue; + copyDeviceToDevice( + static_cast(recv_buffer) + peer * send_bytes + pos, + static_cast(src) + peer * size, size, + enqueue_stream); + } + }); + return {}; +} + +PGResult> +MooncakeCommunicator::reduceScatterCpu(const void* send_buffer, + void* recv_buffer, size_t count, + DataType datatype, ReduceOp op, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "reduceScatterCpu requires a CPU communicator"); + PG_TRY(checkOpState(OpType::ReduceScatter)); + PG_TRY(auto recv_bytes, getByteCount(count, datatype)); + PG_TRY(checkBuffer(send_buffer, recv_bytes, "send buffer")); + PG_TRY(checkBuffer(recv_buffer, recv_bytes, "receive buffer")); + PG_TRY(checkReduction(datatype, op, true)); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + return worker_->putTaskCpu( + OpType::ReduceScatter, recv_bytes, 0, meta_, failed_ranks_hint, + [=, this](void* dst, size_t pos, size_t size) { + for (int peer = 0; peer < active_size; ++peer) { + if (!meta_->activeRanks[peer]) continue; + std::memcpy(static_cast(dst) + peer * size, + static_cast(send_buffer) + + peer * recv_bytes + pos, + size); + } + }, + [=, this](void* src, size_t pos, size_t size) { + std::memset(static_cast(recv_buffer) + pos, 0, size); + launchReduceCpu(recv_buffer, datatype, pos, size, src, active_size, + op, meta_->activeRanks); + }); +} + +PGResult MooncakeCommunicator::reduceScatterGpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + ReduceOp op, cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "reduceScatterGpu requires a GPU communicator"); + PG_TRY(checkOpState(OpType::ReduceScatter)); + PG_TRY(auto recv_bytes, getByteCount(count, datatype)); + PG_TRY(checkBuffer(send_buffer, recv_bytes, "send buffer")); + PG_TRY(checkBuffer(recv_buffer, recv_bytes, "receive buffer")); + PG_TRY(checkReduction(datatype, op, false)); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + worker_->putTaskCuda( + OpType::ReduceScatter, recv_bytes, 0, meta_, stream, failed_ranks_hint, + [=, this](void* dst, size_t pos, size_t size, + cudaStream_t enqueue_stream) { + for (int peer = 0; peer < active_size; ++peer) { + if (!meta_->activeRanks[peer]) continue; + copyDeviceToDevice(static_cast(dst) + peer * size, + static_cast(send_buffer) + + peer * recv_bytes + pos, + size, enqueue_stream); + } + }, + [=, this](void* src, size_t pos, size_t size, + cudaStream_t enqueue_stream) { + PG_ASSERT_CUDA( + cudaMemsetAsync(static_cast(recv_buffer) + pos, 0, size, + enqueue_stream)); + launchReduceKernel(recv_buffer, datatype, pos, size, src, + active_size, op, meta_->activeRanksDevice, + enqueue_stream); + }); + return {}; +} + +PGResult> MooncakeCommunicator::allToAllCpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "allToAllCpu requires a CPU communicator"); + PG_TRY(checkOpState(OpType::AllToAll)); + PG_TRY(auto peer_bytes, getByteCount(count, datatype)); + PG_TRY(checkBuffer(send_buffer, peer_bytes, "send buffer")); + PG_TRY(checkBuffer(recv_buffer, peer_bytes, "receive buffer")); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + return worker_->putTaskCpu( + OpType::AllToAll, peer_bytes, 0, meta_, failed_ranks_hint, + [=, this](void* dst, size_t pos, size_t size) { + for (int peer = 0; peer < active_size; ++peer) { + std::memcpy(static_cast(dst) + peer * size, + static_cast(send_buffer) + + peer * peer_bytes + pos, + size); + } + }, + [=, this](void* src, size_t pos, size_t size) { + for (int peer = 0; peer < active_size; ++peer) { + std::memcpy( + static_cast(recv_buffer) + peer * peer_bytes + pos, + static_cast(src) + peer * size, size); + } + }); +} + +PGResult MooncakeCommunicator::allToAllGpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "allToAllGpu requires a GPU communicator"); + PG_TRY(checkOpState(OpType::AllToAll)); + PG_TRY(auto peer_bytes, getByteCount(count, datatype)); + PG_TRY(checkBuffer(send_buffer, peer_bytes, "send buffer")); + PG_TRY(checkBuffer(recv_buffer, peer_bytes, "receive buffer")); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + worker_->putTaskCuda( + OpType::AllToAll, peer_bytes, 0, meta_, stream, failed_ranks_hint, + [=, this](void* dst, size_t pos, size_t size, + cudaStream_t enqueue_stream) { + for (int peer = 0; peer < active_size; ++peer) { + copyDeviceToDevice(static_cast(dst) + peer * size, + static_cast(send_buffer) + + peer * peer_bytes + pos, + size, enqueue_stream); + } + }, + [=, this](void* src, size_t pos, size_t size, + cudaStream_t enqueue_stream) { + for (int peer = 0; peer < active_size; ++peer) { + copyDeviceToDevice( + static_cast(recv_buffer) + peer * peer_bytes + pos, + static_cast(src) + peer * size, size, + enqueue_stream); + } + }); + return {}; +} + +PGResult> MooncakeCommunicator::barrierCpu( + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "barrierCpu requires a CPU communicator"); + PG_TRY(checkOpState(OpType::Barrier)); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + return worker_->putTaskCpu( + OpType::Barrier, kBarrierDummySize, 0, meta_, failed_ranks_hint, + [](void*, size_t, size_t) {}, [](void*, size_t, size_t) {}); +} + +PGResult MooncakeCommunicator::barrierGpu( + cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "barrierGpu requires a GPU communicator"); + PG_TRY(checkOpState(OpType::Barrier)); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + worker_->putTaskCuda( + OpType::Barrier, kBarrierDummySize, 0, meta_, stream, failed_ranks_hint, + [](void*, size_t, size_t, cudaStream_t) {}, + [](void*, size_t, size_t, cudaStream_t) {}); + return {}; +} + +PGResult> MooncakeCommunicator::reduceCpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + ReduceOp op, int root, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "reduceCpu requires a CPU communicator"); + PG_TRY(checkOpState(OpType::Reduce)); + PG_TRY(auto bytes, getByteCount(count, datatype)); + PG_TRY(checkRoot(root, max_group_size_, "reduce")); + PG_TRY(checkBuffer(send_buffer, bytes, "send buffer")); + const bool is_root = root == rank_; + if (is_root) { + PG_TRY(checkBuffer(recv_buffer, bytes, "receive buffer")); + } + PG_TRY(checkReduction(datatype, op, true)); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + return worker_->putTaskCpu( + OpType::Reduce, bytes, root, meta_, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size) { + std::memcpy(dst, static_cast(send_buffer) + pos, size); + }, + [=, this](void* src, size_t pos, size_t size) { + if (!is_root) return; + std::memset(static_cast(recv_buffer) + pos, 0, size); + launchReduceCpu(recv_buffer, datatype, pos, size, src, active_size, + op, meta_->activeRanks); + }); +} + +PGResult MooncakeCommunicator::reduceGpu(const void* send_buffer, + void* recv_buffer, size_t count, + DataType datatype, ReduceOp op, + int root, cudaStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "reduceGpu requires a GPU communicator"); + PG_TRY(checkOpState(OpType::Reduce)); + PG_TRY(auto bytes, getByteCount(count, datatype)); + PG_TRY(checkRoot(root, max_group_size_, "reduce")); + PG_TRY(checkBuffer(send_buffer, bytes, "send buffer")); + const bool is_root = root == rank_; + if (is_root) { + PG_TRY(checkBuffer(recv_buffer, bytes, "receive buffer")); + } + PG_TRY(checkReduction(datatype, op, false)); + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + worker_->putTaskCuda( + OpType::Reduce, bytes, root, meta_, stream, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size, cudaStream_t enqueue_stream) { + copyDeviceToDevice(dst, static_cast(send_buffer) + pos, + size, enqueue_stream); + }, + [=, this](void* src, size_t pos, size_t size, + cudaStream_t enqueue_stream) { + if (!is_root) return; + PG_ASSERT_CUDA( + cudaMemsetAsync(static_cast(recv_buffer) + pos, 0, size, + enqueue_stream)); + launchReduceKernel(recv_buffer, datatype, pos, size, src, + active_size, op, meta_->activeRanksDevice, + enqueue_stream); + }); + return {}; +} + +PGResult> MooncakeCommunicator::gatherCpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + int root, int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "gatherCpu requires a CPU communicator"); + PG_TRY(checkOpState(OpType::Gather)); + PG_TRY(auto send_bytes, getByteCount(count, datatype)); + PG_TRY(checkRoot(root, max_group_size_, "gather")); + PG_TRY(checkBuffer(send_buffer, send_bytes, "send buffer")); + const bool is_root = root == rank_; + if (is_root) { + PG_TRY(checkBuffer(recv_buffer, send_bytes, "receive buffer")); + } + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + return worker_->putTaskCpu( + OpType::Gather, send_bytes, root, meta_, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size) { + std::memcpy(dst, static_cast(send_buffer) + pos, size); + }, + [=, this](void* src, size_t pos, size_t size) { + if (!is_root) return; + for (int peer = 0; peer < active_size; ++peer) { + std::memcpy( + static_cast(recv_buffer) + peer * send_bytes + pos, + static_cast(src) + peer * size, size); + } + }); +} + +PGResult MooncakeCommunicator::gatherGpu(const void* send_buffer, + void* recv_buffer, size_t count, + DataType datatype, int root, + cudaStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "gatherGpu requires a GPU communicator"); + PG_TRY(checkOpState(OpType::Gather)); + PG_TRY(auto send_bytes, getByteCount(count, datatype)); + PG_TRY(checkRoot(root, max_group_size_, "gather")); + PG_TRY(checkBuffer(send_buffer, send_bytes, "send buffer")); + const bool is_root = root == rank_; + if (is_root) { + PG_TRY(checkBuffer(recv_buffer, send_bytes, "receive buffer")); + } + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + worker_->putTaskCuda( + OpType::Gather, send_bytes, root, meta_, stream, failed_ranks_hint, + [=](void* dst, size_t pos, size_t size, cudaStream_t enqueue_stream) { + copyDeviceToDevice(dst, static_cast(send_buffer) + pos, + size, enqueue_stream); + }, + [=, this](void* src, size_t pos, size_t size, + cudaStream_t enqueue_stream) { + if (!is_root) return; + for (int peer = 0; peer < active_size; ++peer) { + copyDeviceToDevice( + static_cast(recv_buffer) + peer * send_bytes + pos, + static_cast(src) + peer * size, size, + enqueue_stream); + } + }); + return {}; +} + +PGResult> MooncakeCommunicator::scatterCpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + int root, int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(is_cpu_, "scatterCpu requires a CPU communicator"); + PG_TRY(checkOpState(OpType::Scatter)); + PG_TRY(auto recv_bytes, getByteCount(count, datatype)); + PG_TRY(checkRoot(root, max_group_size_, "scatter")); + PG_TRY(checkBuffer(recv_buffer, recv_bytes, "receive buffer")); + const bool is_root = root == rank_; + if (is_root) { + PG_TRY(checkBuffer(send_buffer, recv_bytes, "send buffer")); + } + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + return worker_->putTaskCpu( + OpType::Scatter, recv_bytes, root, meta_, failed_ranks_hint, + [=, this](void* dst, size_t pos, size_t size) { + if (!is_root) return; + for (int peer = 0; peer < active_size; ++peer) { + std::memcpy(static_cast(dst) + peer * size, + static_cast(send_buffer) + + peer * recv_bytes + pos, + size); + } + }, + [=](void* src, size_t pos, size_t size) { + std::memcpy(static_cast(recv_buffer) + pos, src, size); + }); +} + +PGResult MooncakeCommunicator::scatterGpu( + const void* send_buffer, void* recv_buffer, size_t count, DataType datatype, + int root, cudaStream_t stream, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + PG_VALIDATE_STATE(!is_cpu_, "scatterGpu requires a GPU communicator"); + PG_TRY(checkOpState(OpType::Scatter)); + PG_TRY(auto recv_bytes, getByteCount(count, datatype)); + PG_TRY(checkRoot(root, max_group_size_, "scatter")); + PG_TRY(checkBuffer(recv_buffer, recv_bytes, "receive buffer")); + const bool is_root = root == rank_; + if (is_root) { + PG_TRY(checkBuffer(send_buffer, recv_bytes, "send buffer")); + } + PG_TRY( + initializeFailedRanksHint(failed_ranks_hint, failed_ranks_hint_count)); + const int active_size = getSize(); + worker_->putTaskCuda( + OpType::Scatter, recv_bytes, root, meta_, stream, failed_ranks_hint, + [=, this](void* dst, size_t pos, size_t size, + cudaStream_t enqueue_stream) { + if (!is_root) return; + for (int peer = 0; peer < active_size; ++peer) { + copyDeviceToDevice(static_cast(dst) + peer * size, + static_cast(send_buffer) + + peer * recv_bytes + pos, + size, enqueue_stream); + } + }, + [=](void* src, size_t pos, size_t size, cudaStream_t enqueue_stream) { + copyDeviceToDevice(static_cast(recv_buffer) + pos, src, size, + enqueue_stream); + }); + return {}; +} + +PGResult MooncakeCommunicator::shutdown() { + if (is_shutdown_) return {}; + std::unique_ptr device_guard; + const bool has_device_state = + !is_cpu_ && + (active_ranks_mirror_stream_.has_value() || send_buffer_[0] || + recv_buffer_[0] || worker_ || p2p_proxy_ || meta_); + if (has_device_state) { + device_guard = std::make_unique(device_index_); + } + is_shutdown_ = true; + // Remove this communicator from AgentHost's callback lookup before teardown + // so a concurrent ViewUpdate cannot call into it. Keep the group registered + // locally and at the Coordinator while worker tasks are draining because + // their failure path may still call syncAfterFailure(). + if (isValidGroup()) agent_.detachCommunicator(meta_->group_id); + + // If we encounter any hung operations, don't release resources to avoid a + // potential crash. Instead, allow those resources to leak and rely on the + // OS to reclaim them later. + bool has_hung_operation = false; + + // Phase 1: Drain P2P tasks. + if (p2p_device_worker_ && p2p_proxy_) { + p2p_device_worker_->removeProxy(p2p_proxy_); + has_hung_operation |= !p2p_proxy_->drainTasks(); + } + // Phase 2: Drain collective tasks for this communicator. + if (worker_ && meta_) { + has_hung_operation |= !worker_->drainTasks(meta_.get()); + } + // Phase 3: Device synchronization. + if (has_device_state && !has_hung_operation) cudaDeviceSynchronize(); + + // Phase 4: Release resources. + if (has_hung_operation && p2p_proxy_) p2p_proxy_->abandonResources(); + + if (!has_hung_operation && meta_) { + for (size_t index = 0; index < 2; ++index) { + context_.engine->unregisterLocalMemory( + cpu_sync_send_region_[index]); + context_.engine->unregisterLocalMemory( + cpu_sync_recv_region_[index]); + context_.engine->unregisterLocalMemory(send_buffer_[index]); + context_.engine->unregisterLocalMemory(recv_buffer_[index]); + delete[] cpu_sync_send_region_[index]; + delete[] cpu_sync_recv_region_[index]; + if (is_cpu_) { + std::free(send_buffer_[index]); + std::free(recv_buffer_[index]); + } else { + cudaFree(send_buffer_[index]); + cudaFree(recv_buffer_[index]); + } + } + delete[] meta_->maybeActivatable; + if (is_cpu_) { + delete[] meta_->activeRanks; + } else { + cudaFreeHost(meta_->activeRanks); + } + meta_->activeRanks = nullptr; + meta_->activeRanksDevice = nullptr; + meta_->maybeActivatable = nullptr; + } + // Prevent zombie P2PProxy workers from dereferencing this communicator + // after destruction. Must happen after drainTasks so in-flight failures can + // still be reported during shutdown. + if (meta_) meta_->communicator = nullptr; + + // The data-plane teardown has finished. Remove the group from the local + // Agent and notify the Coordinator that this rank has left it. + if (isValidGroup()) { + PG_TRY(agent_.unregisterGroup(meta_->group_id)); + } + return {}; +} + +std::vector MooncakeCommunicator::getActiveRanks() const { + std::vector result(max_group_size_, 0); + if (!meta_ || !meta_->activeRanks) return result; + for (int index = 0; index < max_group_size_; ++index) { + result[index] = meta_->activeRanks[index] ? 1 : 0; + } + return result; +} + +void MooncakeCommunicator::syncActiveRanksMirror() const { + if (!active_ranks_mirror_) return; + // The mirror is InGroupRank-indexed, in the same order as the caller-owned + // storage. + auto active_ranks = getActiveRanks(); + const size_t bytes = max_group_size_ * sizeof(int32_t); + if (active_ranks_mirror_is_device_) { + const GpuDeviceGuard device_guard(active_ranks_mirror_device_index_); + PG_ASSERT_CUDA(cudaMemcpyAsync( + active_ranks_mirror_, active_ranks.data(), bytes, + cudaMemcpyHostToDevice, active_ranks_mirror_stream_.value().get())); + } else { + std::memcpy(active_ranks_mirror_, active_ranks.data(), bytes); + } +} + +int MooncakeCommunicator::getNumSyncedRanks() const { + if (!meta_ || !meta_->maybeActivatable) return 0; + int count = 0; + for (int index = 0; index < max_group_size_; ++index) { + if (meta_->maybeActivatable[index]) ++count; + } + return count; +} + +PGResult MooncakeCommunicator::checkValidGroup( + const char* operation) const { + PG_VALIDATE_STATE(!is_shutdown_, "communicator is shut down"); + if (!isValidGroup()) { + return makePGError( + PGErrorCode::NotSupported, + std::string(operation) + + " is unavailable because this communicator is invalid"); + } + return {}; +} + +PGResult> MooncakeCommunicator::getPeerState( + const std::vector& ranks) const { + PG_TRY(checkValidGroup("getPeerState")); + std::vector result; + result.reserve(ranks.size()); + for (const int rank : ranks) { + PG_VALIDATE_ARG(rank >= 0 && rank < max_group_size_, + "peer rank is out of range"); + result.push_back(meta_->maybeActivatable[rank]); + } + return result; +} + +PGResult MooncakeCommunicator::activateRanks( + const std::vector& ranks) { + PG_TRY(checkValidGroup("activateRanks")); + for (const int rank : ranks) { + PG_VALIDATE_ARG(rank >= 0 && rank < max_group_size_, + "rank to activate is out of range"); + } + std::vector local_ranks(ranks.begin(), ranks.end()); + auto result = agent_.proposeActivate(meta_->group_id, local_ranks); + if (!result.has_value()) { + return makePGError(std::move(result).error()); + } + if (result.value().status == ProposalStatus::Rejected) { + LOG(WARNING) << "MooncakeCommunicator: activateRanks rejected: " + << result.value().reject_reason; + } + return result; +} + +PGResult MooncakeCommunicator::deactivateRanks( + const std::vector& ranks) { + PG_TRY(checkValidGroup("deactivateRanks")); + for (const int rank : ranks) { + PG_VALIDATE_ARG(rank >= 0 && rank < max_group_size_, + "rank to deactivate is out of range"); + } + std::vector local_ranks(ranks.begin(), ranks.end()); + auto result = agent_.proposeDeactivate(meta_->group_id, local_ranks); + if (!result.has_value()) { + return makePGError(std::move(result).error()); + } + if (result.value().status == ProposalStatus::Rejected) { + LOG(WARNING) << "MooncakeCommunicator: deactivateRanks rejected: " + << result.value().reject_reason; + } + return result; +} + +PGResult MooncakeCommunicator::joinGroup() { + PG_TRY(checkValidGroup("joinGroup")); + auto mode = meta_->extensionMode.load(std::memory_order_acquire); + PG_VALIDATE_STATE( + mode == CollectiveExtensionState::Isolated, + "joinGroup may only be called once on an isolated joining " + "communicator"); + // Stop admitting isolated collectives before advertising readiness. + meta_->extensionMode.store(CollectiveExtensionState::Quiescing, + std::memory_order_release); + if (!worker_->drainTasks(meta_.get())) { + return makePGError( + PGErrorCode::Timeout, + "timed out draining join preparation collectives for rank " + + std::to_string(meta_->globalRank)); + } + PG_TRY(agent_.confirmReadyForActivation(meta_->group_id)); + // Block until the Coordinator activates this rank in the group. + PG_TRY(agent_.waitUntilRankActive(meta_->group_id, meta_->globalRank, + std::chrono::seconds(300))); + const bool normal_and_active = + meta_->extensionMode.load(std::memory_order_acquire) == + CollectiveExtensionState::Normal && + meta_->activeRanks[rank_]; + PG_ASSERT(normal_and_active, "Bad waitUntilRankActive"); + LOG(INFO) << "joinGroup rank=" << meta_->globalRank + << " group=" << meta_->group_id << " activated"; + return {}; +} + +uint64_t MooncakeCommunicator::getCurrentEpoch() const { + return meta_ ? meta_->epoch.load(std::memory_order_acquire) : 0; +} + +PGResult MooncakeCommunicator::syncAfterFailure() { + PG_TRY(checkValidGroup("syncAfterFailure")); + return agent_.syncAfterFailure(meta_->group_id); +} + +void MooncakeCommunicator::applyViewUpdate( + const GroupView& view, const std::vector& rank_states, + const std::vector& rank_epochs, + const std::vector& activatable) { + if (!meta_) return; + + // Ignore stale views that arrive out of order + auto current_epoch = meta_->epoch.load(std::memory_order_acquire); + if (view.epoch < current_epoch) { + return; + } + + bool epoch_changed = current_epoch != view.epoch; + + // An authoritative view in which self is Active is the common commit point + // for enabling normal collective execution: + // + // founding ranks: Isolated -> Normal + // joining ranks: Quiescing -> Normal + // + // A non-Active view deliberately does not determine the local mode. A new + // joiner must remain Isolated until joinGroup is called; a joiner awaiting + // activation must remain Quiescing; and an auto-deactivated communicator + // must remain Normal so its inactive self bit makes the next collective + // fail fast. + auto mode = meta_->extensionMode.load(std::memory_order_acquire); + auto next_mode = mode; + if (view.members[meta_->globalRank].isActive()) { + next_mode = CollectiveExtensionState::Normal; + } + + PG_ASSERT( + static_cast(view.rank_order.size()) <= meta_->maxGroupSize, + "Bad group view"); + + // Preserve stable in-group rank slots: activeSize is the upper bound of the + // active rank space, not the number of set bits. For example, an active + // mask of [true, false, true] has activeSize == 3. + int active_size = 0; + for (size_t local_rank = 0; local_rank < view.rank_order.size(); + ++local_rank) { + const auto global_rank = view.rank_order[local_rank]; + if (view.members[global_rank].isActive()) { + active_size = static_cast(local_rank) + 1; + } + } + + std::vector previous_active_ranks(meta_->maxGroupSize); + for (int local_rank = 0; local_rank < meta_->maxGroupSize; ++local_rank) { + previous_active_ranks[local_rank] = meta_->activeRanks[local_rank]; + } + + // The execution mode determines the effective active ranks consumed by + // kernels. Isolated and Quiescing use a local-only mask; Normal follows the + // Coordinator's committed membership view. + switch (next_mode) { + case CollectiveExtensionState::Isolated: + case CollectiveExtensionState::Quiescing: + for (int local_rank = 0; local_rank < meta_->maxGroupSize; + ++local_rank) { + meta_->activeRanks[local_rank] = local_rank == rank_; + } + break; + case CollectiveExtensionState::Normal: + for (int local_rank = 0; local_rank < meta_->maxGroupSize; + ++local_rank) { + meta_->activeRanks[local_rank] = false; + } + for (size_t local_rank = 0; local_rank < view.rank_order.size(); + ++local_rank) { + const auto global_rank = view.rank_order[local_rank]; + meta_->activeRanks[local_rank] = + view.members[global_rank].isActive(); + } + break; + } + + // Only a change in execution mode or effective participants starts a new + // collective taskCount. Endpoint, AwaitingActivation updates, ... keep the + // current taskCount even though they advance the view epoch. + bool reset_task_count = next_mode != mode; + for (int local_rank = 0; local_rank < meta_->maxGroupSize; ++local_rank) { + reset_task_count |= + previous_active_ranks[local_rank] != meta_->activeRanks[local_rank]; + } + if (reset_task_count) meta_->taskCount = 0; + + // Rank order and endpoint metadata. + for (size_t local_rank = 0; local_rank < view.rank_order.size(); + ++local_rank) { + // rank order + meta_->rank_order[local_rank] = view.rank_order[local_rank]; + const auto global_rank = view.rank_order[local_rank]; + + const auto& member = view.members[global_rank]; + if (member.endpoint.has_value()) { + meta_->segmentInfos[local_rank] = *member.endpoint; + } + } + + // Rank states + for (size_t i = 0; i < rank_states.size(); ++i) { + meta_->rankStates[i] = rank_states[i]; + } + for (size_t i = 0; i < rank_epochs.size(); ++i) { + meta_->rankEpochs[i] = rank_epochs[i]; + } + + // Best-effort Activatable + for (size_t i = 0; i < activatable.size(); ++i) { + meta_->maybeActivatable[i] = activatable[i]; + } + + // Keep the caller-visible active-ranks mirror in sync with the view. + // FIXME: potential deadlock? + syncActiveRanksMirror(); + + // Publish the rank-space extent after the corresponding data-plane state. + // getSize() reads this from the application thread. + meta_->activeSize.store(active_size, std::memory_order_release); + + // Publish epoch AFTER all data-plane state (activeRanks, segmentInfos, + // etc.) is updated. This ensures that a thread observing the new epoch via + // getCurrentEpoch() (acquire) sees the complete membership state. + if (epoch_changed) { + meta_->epoch.store(view.epoch, std::memory_order_release); + } + + if (next_mode != mode) { + meta_->extensionMode.store(next_mode, std::memory_order_release); + } +} + +void MooncakeCommunicator::onPeerLinkReset(InGroupRank peer) { + if (is_shutdown_) return; + if (p2p_proxy_) p2p_proxy_->resetPeerState(peer); + if (peer >= 0 && peer < max_group_size_) { + meta_->segmentIDs[peer] = static_cast(-1); + } +} + +void MooncakeCommunicator::refreshSegmentID(InGroupRank local) { + if (local < 0 || local >= max_group_size_) return; + const auto handle = + context_.link_manager.resolvePeer(meta_->rank_order[local]); + meta_->segmentIDs[local] = + handle ? *handle : static_cast(-1); +} + +GroupEndpointPublication MooncakeCommunicator::buildEndpointMetadata() const { + return GroupEndpointPublication{ + .group_id = meta_->group_id, + .endpoint_info = meta_->segmentInfos[meta_->rank]}; +} + +} // namespace mooncake diff --git a/mooncake-pg/src/mooncake_pg.cpp b/mooncake-pg/src/mooncake_pg.cpp new file mode 100644 index 0000000000..b457284aa3 --- /dev/null +++ b/mooncake-pg/src/mooncake_pg.cpp @@ -0,0 +1,971 @@ +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "comm_types.h" +#include "error_types.h" + +using namespace mooncake; + +struct mooncakePgContext { + std::unique_ptr impl; +}; + +struct mooncakePgComm { + explicit mooncakePgComm(MooncakePGContext& context) : context(&context) {} + + ~mooncakePgComm() { + impl.reset(); + if (context_use_counted) context->decrementCommUseCount(); + } + + MooncakePGContext* context; + bool context_use_counted = false; + std::unique_ptr impl; +}; + +struct mooncakePgCompletion { + std::unique_ptr impl; +}; + +namespace { + +struct LastError { + mooncakePgResult_t result = mooncakePgSuccess; + // Keep catch paths allocation-free: assigning to a std::string here could + // throw and let an exception escape the C API boundary. + std::array message{}; +}; + +thread_local LastError g_last_error; + +mooncakePgResult_t setLastError(mooncakePgResult_t result, + const char* message) { + g_last_error.result = result; + std::snprintf(g_last_error.message.data(), g_last_error.message.size(), + "%s", message ? message : ""); + return result; +} + +mooncakePgResult_t setLastError(const PGError& error) { + const auto set = [&](mooncakePgResult_t result) { + return setLastError(result, error.message.c_str()); + }; + switch (error.code) { + case PGErrorCode::InvalidArgument: + return set(mooncakePgInvalidArgument); + case PGErrorCode::InvalidState: + return set(mooncakePgInvalidState); + case PGErrorCode::NotSupported: + return set(mooncakePgNotSupported); + case PGErrorCode::Timeout: + return set(mooncakePgTimeout); + case PGErrorCode::ResourceBusy: + return set(mooncakePgResourceBusy); + case PGErrorCode::TransferEngineError: + return set(mooncakePgTransferEngineError); + case PGErrorCode::RpcError: + return set(mooncakePgRpcError); + case PGErrorCode::SystemError: + return set(mooncakePgSystemError); + case PGErrorCode::InternalError: + return set(mooncakePgInternalError); + } + PG_ASSERT(false, "unknown PGErrorCode: ", static_cast(error.code)); +} + +template +mooncakePgResult_t asCApiResult(Function&& function) { + static_assert( + std::is_same_v, PGResult>); + try { + const auto result = function(); + return result.has_value() ? mooncakePgSuccess + : setLastError(result.error()); + } catch (const PGAssertionException& error) { + return setLastError(mooncakePgInternalError, error.what()); + } catch (const std::bad_alloc& error) { + return setLastError(mooncakePgSystemError, error.what()); + } catch (const std::exception& error) { + return setLastError(mooncakePgInternalError, error.what()); + } catch (...) { + return setLastError(mooncakePgInternalError, + "unknown Mooncake PG error"); + } +} + +PGResult parseCommConfig( + const mooncakePgCommConfig_t* config) { + PG_VALIDATE_ARG(config, "communicator config is null"); + PG_VALIDATE_ARG(config->structSize >= sizeof(*config) && + config->magic == MOONCAKE_PG_COMM_CONFIG_MAGIC && + config->version == MOONCAKE_PG_COMM_CONFIG_VERSION, + "invalid communicator config header"); + PG_VALIDATE_ARG(config->groupId, "communicator group ID is null"); + PG_VALIDATE_ARG(config->globalRanks, "communicator global ranks are null"); + PG_VALIDATE_ARG(config->globalRankCount != 0 && + config->globalRankCount <= MOONCAKE_PG_MAX_RANKS, + "invalid communicator global rank count"); + PG_VALIDATE_ARG( + config->activeRanksMirror || config->activeRanksMirrorCount == 0, + "active-ranks mirror is null"); + + MooncakeCommunicatorConfig internal; + internal.rank = config->rank; + internal.size = config->size; + internal.max_group_size = + config->maxGroupSize == MOONCAKE_PG_CONFIG_UNDEF_INT + ? config->size + : config->maxGroupSize; + internal.global_ranks.assign(config->globalRanks, + config->globalRanks + config->globalRankCount); + internal.group_bootstrap_id = config->groupId; + switch (config->deviceType) { + case mooncakePgDeviceCpu: + internal.is_cpu = true; + break; + case mooncakePgDeviceGpu: + internal.is_cpu = false; + break; + default: + return makePGError(PGErrorCode::InvalidArgument, + "invalid communicator device type"); + } + internal.device_index = config->deviceIndex == MOONCAKE_PG_CONFIG_UNDEF_INT + ? -1 + : config->deviceIndex; + switch (config->idResolvePolicy) { + case mooncakePgIdResolveCreateOrAttach: + internal.group_resolve_policy = + GroupBootstrapIdResolvePolicy::CreateOrAttach; + break; + case mooncakePgIdResolveAttachOrExtend: + internal.group_resolve_policy = + GroupBootstrapIdResolvePolicy::AttachOrExtend; + break; + default: + return makePGError(PGErrorCode::InvalidArgument, + "invalid communicator group resolve policy"); + } + internal.auto_deactivate_on_failure = config->autoDeactivateOnFailure != 0; + internal.auto_sync_on_failure = config->autoSyncOnFailure != 0; + internal.active_ranks_mirror = config->activeRanksMirror; + internal.active_ranks_mirror_count = config->activeRanksMirrorCount; + internal.active_ranks_mirror_is_device = + config->activeRanksMirrorIsDevice != 0; + if (internal.active_ranks_mirror_is_device) { + internal.active_ranks_mirror_device_index = + config->activeRanksMirrorDeviceIndex == MOONCAKE_PG_CONFIG_UNDEF_INT + ? -1 + : config->activeRanksMirrorDeviceIndex; + } + return internal; +} + +PGResult convertDataType(mooncakePgDataType_t data_type) { + switch (data_type) { + case mooncakePgInt8: + return DataType::Int8; + case mooncakePgUint8: + return DataType::Uint8; + case mooncakePgInt16: + return DataType::Int16; + case mooncakePgUint16: + return DataType::Uint16; + case mooncakePgInt32: + return DataType::Int32; + case mooncakePgUint32: + return DataType::Uint32; + case mooncakePgInt64: + return DataType::Int64; + case mooncakePgUint64: + return DataType::Uint64; + case mooncakePgFloat16: + return DataType::Float16; + case mooncakePgFloat32: + return DataType::Float32; + case mooncakePgFloat64: + return DataType::Float64; + case mooncakePgBfloat16: + return DataType::Bfloat16; + case mooncakePgBool: + return DataType::Bool; + case mooncakePgFloat8e4m3fn: + return DataType::Float8e4m3fn; + case mooncakePgFloat8e5m2: + return DataType::Float8e5m2; + case mooncakePgFloat8e4m3fnuz: + return DataType::Float8e4m3fnuz; + case mooncakePgFloat8e5m2fnuz: + return DataType::Float8e5m2fnuz; + case mooncakePgFloat8e8m0fnu: + return DataType::Float8e8m0fnu; + default: + return makePGError(PGErrorCode::InvalidArgument, + "unsupported Mooncake PG datatype"); + } +} + +PGResult convertReduceOp(mooncakePgReduceOp_t reduce_op) { + switch (reduce_op) { + case mooncakePgSum: + return ReduceOp::Sum; + case mooncakePgAvg: + return ReduceOp::Avg; + case mooncakePgProduct: + return ReduceOp::Product; + case mooncakePgMin: + return ReduceOp::Min; + case mooncakePgMax: + return ReduceOp::Max; + default: + return makePGError(PGErrorCode::InvalidArgument, + "unsupported Mooncake PG reduction operation"); + } +} + +cudaStream_t convertStream(mooncakePgStream_t stream) { + return reinterpret_cast(stream); +} + +using CompletionResult = PGResult>; + +template +mooncakePgResult_t invokeCommOpWithCompletion( + mooncakePgComm_t comm, mooncakePgCompletion_t* output_completion, + Launch&& launch) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(output_completion, "completion output is null"); + *output_completion = nullptr; + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + + auto output = std::make_unique(); + PG_TRY(auto completion, launch(*comm->impl)); + PG_ASSERT(completion, "operation returned no completion"); + + output->impl = std::move(completion); + *output_completion = output.release(); + return {}; + }); +} + +template +mooncakePgResult_t invokeCommOp(mooncakePgComm_t comm, Launch&& launch) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + return launch(*comm->impl); + }); +} + +PGResult> parseRanks(const int32_t* ranks, size_t rank_count) { + PG_VALIDATE_ARG(rank_count <= MOONCAKE_PG_MAX_RANKS, + "rank count exceeds maximum"); + PG_VALIDATE_ARG(rank_count == 0 || ranks, "ranks are null"); + std::vector output; + if (rank_count != 0) { + output.assign(ranks, ranks + rank_count); + } + return output; +} + +mooncakePgProposalResponse_t convertProposalResponse( + const ProposeViewUpdateResponse& response) { + mooncakePgProposalResponse_t converted{}; + converted.status = static_cast(response.status); + converted.newEpoch = response.new_epoch; + converted.droppedRankCount = + std::min(response.dropped_ranks.size(), + static_cast(MOONCAKE_PG_MAX_RANKS)); + std::copy_n(response.dropped_ranks.begin(), converted.droppedRankCount, + converted.droppedRanks); + std::snprintf(converted.rejectReason, sizeof(converted.rejectReason), "%s", + response.reject_reason.c_str()); + return converted; +} + +} // namespace + +const char* mooncakePgGetErrorString(mooncakePgResult_t result) { + switch (result) { + case mooncakePgSuccess: + return "success"; + case mooncakePgInvalidArgument: + return "invalid argument"; + case mooncakePgInvalidState: + return "invalid state"; + case mooncakePgNotSupported: + return "operation not supported"; + case mooncakePgTimeout: + return "operation timed out"; + case mooncakePgResourceBusy: + return "resource busy"; + case mooncakePgTransferEngineError: + return "transfer engine error"; + case mooncakePgRpcError: + return "RPC error"; + case mooncakePgSystemError: + return "system error"; + case mooncakePgInternalError: + return "internal error"; + } + return "unknown result"; +} + +const char* mooncakePgGetLastError(void) { return g_last_error.message.data(); } + +mooncakePgResult_t mooncakePgContextCreate(mooncakePgContext_t* context) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context, "context output is null"); + *context = nullptr; + + auto output = std::make_unique(); + output->impl = std::make_unique(); + *context = output.release(); + return {}; + }); +} + +mooncakePgResult_t mooncakePgContextInitialize(mooncakePgContext_t context, + int global_rank, + int max_world_size) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + return context->impl->initialize(global_rank, max_world_size); + }); +} + +mooncakePgResult_t mooncakePgContextLaunchCoordinator( + mooncakePgContext_t context, char* coordinator_address_buf, + size_t coordinator_address_buf_size) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + PG_VALIDATE_ARG(coordinator_address_buf, + "coordinator-address output is null"); + PG_TRY(auto value, context->impl->launchCoordinator()); + PG_VALIDATE_ARG(value.size() < coordinator_address_buf_size, + "coordinator-address output is too small"); + std::memcpy(coordinator_address_buf, value.c_str(), value.size() + 1); + return {}; + }); +} + +mooncakePgResult_t mooncakePgContextConnectCoordinator( + mooncakePgContext_t context, const char* coordinator_address) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + PG_VALIDATE_ARG(coordinator_address && coordinator_address[0] != '\0', + "coordinator address is null or empty"); + return context->impl->connectCoordinator(coordinator_address); + }); +} + +mooncakePgResult_t mooncakePgContextSetHostIp(mooncakePgContext_t context, + const char* host_ip) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + PG_VALIDATE_ARG(host_ip, "host IP is null"); + return context->impl->setHostIp(host_ip); + }); +} + +mooncakePgResult_t mooncakePgContextSetTransferEngine( + mooncakePgContext_t context, void* transfer_engine) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + return context->impl->setExternalEngine( + static_cast(transfer_engine)); + }); +} + +mooncakePgResult_t mooncakePgContextSetDeviceFilter(mooncakePgContext_t context, + const char* const* filters, + size_t filter_count) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + PG_VALIDATE_ARG(filter_count == 0 || filters, + "device filters are null"); + std::vector values; + values.reserve(filter_count); + for (size_t index = 0; index < filter_count; ++index) { + PG_VALIDATE_ARG(filters[index], "device filter is null"); + values.emplace_back(filters[index]); + } + return context->impl->setDeviceFilter(std::move(values)); + }); +} + +mooncakePgResult_t mooncakePgContextSetCollectiveTimeout( + mooncakePgContext_t context, size_t timeout_us) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + return context->impl->setCollectiveTimeout(timeout_us); + }); +} + +mooncakePgResult_t mooncakePgContextSetP2PTimeout(mooncakePgContext_t context, + int64_t timeout_us) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + return context->impl->setP2PTimeout(timeout_us); + }); +} + +mooncakePgResult_t mooncakePgContextSetFaultReconciliationWindow( + mooncakePgContext_t context, int64_t timeout_us) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + return context->impl->setFaultReconciliationWindow(timeout_us); + }); +} + +mooncakePgResult_t mooncakePgContextDestroy(mooncakePgContext_t context) { + return asCApiResult([&]() -> PGResult { + if (!context) return {}; + PG_VALIDATE_ARG(context->impl, "invalid context"); + auto result = context->impl->shutdown(); + if (result.has_value()) delete context; + return result; + }); +} + +mooncakePgResult_t mooncakePgCommCreate(mooncakePgContext_t context, + const mooncakePgCommConfig_t* config, + mooncakePgComm_t* comm) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm, "communicator output is null"); + *comm = nullptr; + PG_VALIDATE_ARG(context && context->impl, "invalid context"); + PG_TRY(auto internal, parseCommConfig(config)); + + auto output = std::make_unique(*context->impl); + PG_TRY(context->impl->incrementCommUseCount()); + output->context_use_counted = true; + PG_TRY(output->impl, MooncakeCommunicator::create(*context->impl, + std::move(internal))); + *comm = output.release(); + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommDestroy(mooncakePgComm_t comm) { + return asCApiResult([&]() -> PGResult { + std::unique_ptr holder(comm); + if (holder && holder->impl) { + return holder->impl->shutdown(); + } + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommGetRank(mooncakePgComm_t comm, int* rank) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + PG_VALIDATE_ARG(rank, "rank output is null"); + *rank = comm->impl->getRank(); + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommGetSize(mooncakePgComm_t comm, int* size) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + PG_VALIDATE_ARG(size, "size output is null"); + *size = comm->impl->getSize(); + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommGetMaxGroupSize(mooncakePgComm_t comm, + int* max_group_size) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + PG_VALIDATE_ARG(max_group_size, "maximum group-size output is null"); + *max_group_size = comm->impl->getMaxGroupSize(); + return {}; + }); +} + +mooncakePgResult_t mooncakePgBroadcastGpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, + int root, mooncakePgComm_t comm, + mooncakePgStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + return invokeCommOp( + comm, [&](MooncakeCommunicator& impl) -> PGResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.broadcastGpu(send_buffer, recv_buffer, count, + converted_data_type, root, + convertStream(stream), failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgAllReduceGpu( + const void* send_buffer, void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, mooncakePgReduceOp_t reduce_op, + mooncakePgComm_t comm, mooncakePgStream_t stream, + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + return invokeCommOp( + comm, [&](MooncakeCommunicator& impl) -> PGResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + PG_TRY(auto converted_reduce_op, convertReduceOp(reduce_op)); + return impl.allReduceGpu(send_buffer, recv_buffer, count, + converted_data_type, converted_reduce_op, + convertStream(stream), failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgAllGatherGpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, + mooncakePgComm_t comm, + mooncakePgStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + return invokeCommOp( + comm, [&](MooncakeCommunicator& impl) -> PGResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.allGatherGpu(send_buffer, recv_buffer, count, + converted_data_type, convertStream(stream), + failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgReduceScatterGpu( + const void* send_buffer, void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, mooncakePgReduceOp_t reduce_op, + mooncakePgComm_t comm, mooncakePgStream_t stream, + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + return invokeCommOp( + comm, [&](MooncakeCommunicator& impl) -> PGResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + PG_TRY(auto converted_reduce_op, convertReduceOp(reduce_op)); + return impl.reduceScatterGpu( + send_buffer, recv_buffer, count, converted_data_type, + converted_reduce_op, convertStream(stream), failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgAllToAllGpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, + mooncakePgComm_t comm, + mooncakePgStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + return invokeCommOp( + comm, [&](MooncakeCommunicator& impl) -> PGResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.allToAllGpu(send_buffer, recv_buffer, count, + converted_data_type, convertStream(stream), + failed_ranks_hint, failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgReduceGpu( + const void* send_buffer, void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, mooncakePgReduceOp_t reduce_op, int root, + mooncakePgComm_t comm, mooncakePgStream_t stream, + int32_t* failed_ranks_hint, size_t failed_ranks_hint_count) { + return invokeCommOp( + comm, [&](MooncakeCommunicator& impl) -> PGResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + PG_TRY(auto converted_reduce_op, convertReduceOp(reduce_op)); + return impl.reduceGpu(send_buffer, recv_buffer, count, + converted_data_type, converted_reduce_op, + root, convertStream(stream), + failed_ranks_hint, failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgGatherGpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, int root, + mooncakePgComm_t comm, + mooncakePgStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + return invokeCommOp( + comm, [&](MooncakeCommunicator& impl) -> PGResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.gatherGpu(send_buffer, recv_buffer, count, + converted_data_type, root, + convertStream(stream), failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgScatterGpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, + int root, mooncakePgComm_t comm, + mooncakePgStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + return invokeCommOp( + comm, [&](MooncakeCommunicator& impl) -> PGResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.scatterGpu(send_buffer, recv_buffer, count, + converted_data_type, root, + convertStream(stream), failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgBarrierGpu(mooncakePgComm_t comm, + mooncakePgStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count) { + return invokeCommOp(comm, [&](MooncakeCommunicator& impl) { + return impl.barrierGpu(convertStream(stream), failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgBroadcastCpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, + int root, mooncakePgComm_t comm, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.broadcastCpu( + send_buffer, recv_buffer, count, converted_data_type, root, + failed_ranks_hint, failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgAllReduceCpu( + const void* send_buffer, void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, mooncakePgReduceOp_t reduce_op, + mooncakePgComm_t comm, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + PG_TRY(auto converted_reduce_op, convertReduceOp(reduce_op)); + return impl.allReduceCpu(send_buffer, recv_buffer, count, + converted_data_type, converted_reduce_op, + failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgAllGatherCpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, + mooncakePgComm_t comm, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.allGatherCpu(send_buffer, recv_buffer, count, + converted_data_type, failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgReduceScatterCpu( + const void* send_buffer, void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, mooncakePgReduceOp_t reduce_op, + mooncakePgComm_t comm, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + PG_TRY(auto converted_reduce_op, convertReduceOp(reduce_op)); + return impl.reduceScatterCpu(send_buffer, recv_buffer, count, + converted_data_type, + converted_reduce_op, failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgAllToAllCpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, + mooncakePgComm_t comm, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.allToAllCpu(send_buffer, recv_buffer, count, + converted_data_type, failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgReduceCpu( + const void* send_buffer, void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, mooncakePgReduceOp_t reduce_op, int root, + mooncakePgComm_t comm, int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + PG_TRY(auto converted_reduce_op, convertReduceOp(reduce_op)); + return impl.reduceCpu(send_buffer, recv_buffer, count, + converted_data_type, converted_reduce_op, + root, failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgGatherCpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, int root, + mooncakePgComm_t comm, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.gatherCpu(send_buffer, recv_buffer, count, + converted_data_type, root, failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgScatterCpu(const void* send_buffer, + void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, + int root, mooncakePgComm_t comm, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.scatterCpu(send_buffer, recv_buffer, count, + converted_data_type, root, failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgBarrierCpu(mooncakePgComm_t comm, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) { + return impl.barrierCpu(failed_ranks_hint, failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgSendGpu(const void* send_buffer, size_t count, + mooncakePgDataType_t data_type, int peer, + mooncakePgComm_t comm, + mooncakePgStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.sendGpu(send_buffer, count, converted_data_type, peer, + convertStream(stream), failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgRecvGpu(void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, int peer, + mooncakePgComm_t comm, + mooncakePgStream_t stream, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.recvGpu(recv_buffer, count, converted_data_type, peer, + convertStream(stream), failed_ranks_hint, + failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgSendCpu(const void* send_buffer, size_t count, + mooncakePgDataType_t data_type, int peer, + mooncakePgComm_t comm, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.sendCpu(send_buffer, count, converted_data_type, peer, + failed_ranks_hint, failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgRecvCpu(void* recv_buffer, size_t count, + mooncakePgDataType_t data_type, int peer, + mooncakePgComm_t comm, + int32_t* failed_ranks_hint, + size_t failed_ranks_hint_count, + mooncakePgCompletion_t* completion) { + return invokeCommOpWithCompletion( + comm, completion, [&](MooncakeCommunicator& impl) -> CompletionResult { + PG_TRY(auto converted_data_type, convertDataType(data_type)); + return impl.recvCpu(recv_buffer, count, converted_data_type, peer, + failed_ranks_hint, failed_ranks_hint_count); + }); +} + +mooncakePgResult_t mooncakePgCompletionIsCompleted( + mooncakePgCompletion_t completion, int* completed) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(completion && completion->impl, "invalid completion"); + PG_VALIDATE_ARG(completed, "completed output is null"); + *completed = completion->impl->isCompleted() ? 1 : 0; + return {}; + }); +} + +mooncakePgResult_t mooncakePgCompletionWait(mooncakePgCompletion_t completion, + int64_t timeout_us) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(completion && completion->impl, "invalid completion"); + if (!completion->impl->wait(std::chrono::microseconds(timeout_us))) { + return makePGError(PGErrorCode::Timeout, + "completion wait timed out"); + } + return {}; + }); +} + +mooncakePgResult_t mooncakePgCompletionDestroy( + mooncakePgCompletion_t completion) { + return asCApiResult([&]() -> PGResult { + delete completion; + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommGetActiveRanks(mooncakePgComm_t comm, + int32_t* active_ranks, + size_t rank_count) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + const auto ranks = comm->impl->getActiveRanks(); + PG_VALIDATE_ARG(rank_count >= ranks.size(), + "active-ranks output is too small"); + PG_VALIDATE_ARG(ranks.empty() || active_ranks, + "active-ranks output is null"); + if (!ranks.empty()) { + std::copy(ranks.begin(), ranks.end(), active_ranks); + } + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommGetPeerState(mooncakePgComm_t comm, + const int32_t* ranks, + size_t rank_count, + int32_t* peer_states) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + PG_VALIDATE_ARG(rank_count == 0 || peer_states, + "peer-states output is null"); + PG_TRY(auto parsed_ranks, parseRanks(ranks, rank_count)); + PG_TRY(auto states, comm->impl->getPeerState(parsed_ranks)); + for (size_t index = 0; index < states.size(); ++index) { + peer_states[index] = states[index] ? 1 : 0; + } + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommActivateRanks( + mooncakePgComm_t comm, const int32_t* ranks, size_t rank_count, + mooncakePgProposalResponse_t* response) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + PG_VALIDATE_ARG(response, "proposal response is null"); + PG_TRY(auto parsed_ranks, parseRanks(ranks, rank_count)); + PG_TRY(auto proposal_response, comm->impl->activateRanks(parsed_ranks)); + *response = convertProposalResponse(proposal_response); + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommDeactivateRanks( + mooncakePgComm_t comm, const int32_t* ranks, size_t rank_count, + mooncakePgProposalResponse_t* response) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + PG_VALIDATE_ARG(response, "proposal response is null"); + PG_TRY(auto parsed_ranks, parseRanks(ranks, rank_count)); + PG_TRY(auto proposal_response, + comm->impl->deactivateRanks(parsed_ranks)); + *response = convertProposalResponse(proposal_response); + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommJoin(mooncakePgComm_t comm) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + return comm->impl->joinGroup(); + }); +} + +mooncakePgResult_t mooncakePgCommSyncAfterFailure( + mooncakePgComm_t comm, mooncakePgSyncAfterFailureResponse_t* response) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + PG_VALIDATE_ARG(response, "sync response is null"); + PG_TRY(auto sync_response, comm->impl->syncAfterFailure()); + std::memset(response, 0, sizeof(*response)); + response->status = static_cast( + sync_response.status); + std::snprintf(response->rejectReason, sizeof(response->rejectReason), + "%s", sync_response.reject_reason.c_str()); + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommGetEpoch(mooncakePgComm_t comm, + uint64_t* epoch) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + PG_VALIDATE_ARG(epoch, "epoch output is null"); + *epoch = comm->impl->getCurrentEpoch(); + return {}; + }); +} + +mooncakePgResult_t mooncakePgCommGetNumSyncedRanks(mooncakePgComm_t comm, + int* num_synced_ranks) { + return asCApiResult([&]() -> PGResult { + PG_VALIDATE_ARG(comm && comm->impl, "invalid communicator"); + PG_VALIDATE_ARG(num_synced_ranks, "num-synced-ranks output is null"); + *num_synced_ranks = comm->impl->getNumSyncedRanks(); + return {}; + }); +} diff --git a/mooncake-pg/src/mooncake_worker.cu b/mooncake-pg/src/mooncake_worker.cu index 127ecbcb45..46d18577e9 100644 --- a/mooncake-pg/src/mooncake_worker.cu +++ b/mooncake-pg/src/mooncake_worker.cu @@ -1,32 +1,34 @@ // mooncake_worker.cu — GPU kernel functions and launch wrappers. -// Compiled by nvcc (CUDA) and mcc (MUSA, via mooncake_worker.mu symlink). -// The __MUSA__ branch avoids torch headers to stay compatible with mcc. #include +#include +#include + #ifdef __MUSA__ #include +#else +#include #endif -namespace mooncake { +#include "error_types.h" -// ── Kernel functions ────────────────────────────────────────────── -// Both CUDA and MUSA share the same kernel bodies. Parameters use plain -// C++ types (int instead of c10d::OpType / c10d::ReduceOp::RedOpType) -// so that mcc can compile them without torch headers. +namespace mooncake { -__global__ void enqueueTaskKernel(int opType, size_t tensorSize, +__global__ void enqueueTaskKernel(int opType, size_t dataSize, int64_t broadcastRoot, int bufferOffset, - uint64_t submitSequence, void* meta, - Task* tasks, int numRanks, - const bool* activeRanks, - int* activeRanksTensor, size_t taskId) { + uint64_t submitSequence, + int32_t* failedRanksHint, + bool resetFailedRanksHint, void* meta, + Task* tasks, size_t taskId) { // Copy task into slot tasks[taskId].opType = opType; - tasks[taskId].tensorSize = tensorSize; + tasks[taskId].dataSize = dataSize; tasks[taskId].broadcastRoot = broadcastRoot; tasks[taskId].bufferOffset = bufferOffset; tasks[taskId].submitSequence = submitSequence; + tasks[taskId].failedRanksHint = failedRanksHint; + tasks[taskId].resetFailedRanksHint = resetFailedRanksHint; tasks[taskId].transferGroupMeta = meta; // Publish task metadata before notifying the host worker thread. @@ -37,9 +39,6 @@ __global__ void enqueueTaskKernel(int opType, size_t tensorSize, while (tasks[taskId].active) { __threadfence_system(); } - for (int i = 0; i < numRanks; ++i) { - activeRanksTensor[i] = activeRanks[i] ? 1 : 0; - } } template @@ -93,20 +92,19 @@ __global__ void reduceKernel(scalar_t* dst, const scalar_t* src, break; } } else { + const scalar_t val = src[rank * numElements + elem_idx]; switch (op) { case 0: // SUM - acc += src[rank * numElements + elem_idx]; + acc += val; break; case 2: // PRODUCT - acc *= src[rank * numElements + elem_idx]; + acc *= val; break; case 3: // MIN - acc = std::min( - src[rank * numElements + elem_idx], acc); + acc = acc < val ? acc : val; break; case 4: // MAX - acc = std::max( - src[rank * numElements + elem_idx], acc); + acc = val < acc ? acc : val; break; default: break; @@ -136,15 +134,14 @@ __global__ void reduceKernel(scalar_t* dst, const scalar_t* src, namespace mooncake { -void launchEnqueueTaskKernel(int opType, size_t tensorSize, - int64_t broadcastRoot, int bufferOffset, - uint64_t submitSequence, void* meta, Task* tasks, - int numRanks, const bool* activeRanks, - int* activeRanksTensor, size_t taskId, - cudaStream_t stream) { +void launchEnqueueTaskKernel(int opType, size_t dataSize, int64_t broadcastRoot, + int bufferOffset, uint64_t submitSequence, + int32_t* failedRanksHint, + bool resetFailedRanksHint, void* meta, Task* tasks, + size_t taskId, cudaStream_t stream) { enqueueTaskKernel<<<1, 1, 0, stream>>>( - opType, tensorSize, broadcastRoot, bufferOffset, submitSequence, meta, - tasks, numRanks, activeRanks, activeRanksTensor, taskId); + opType, dataSize, broadcastRoot, bufferOffset, submitSequence, + failedRanksHint, resetFailedRanksHint, meta, tasks, taskId); } #define DEF_LAUNCH_REDUCE(scalar_t, suffix) \ @@ -174,8 +171,8 @@ void launchReduceKernel_bf16(void* dst, const void* src, size_t numElements, (const mt_bfloat16*)src, numElements, numRanks, op, activeRanks); #else - reduceKernel<<<64, 256, 0, stream>>>((at::BFloat16*)dst, - (const at::BFloat16*)src, numElements, + reduceKernel<<<64, 256, 0, stream>>>((__nv_bfloat16*)dst, + (const __nv_bfloat16*)src, numElements, numRanks, op, activeRanks); #endif } @@ -189,8 +186,8 @@ void preloadReduceKernels() { auto preload = [](const char* name, auto kernel_ptr) { cudaFuncAttributes attr{}; auto err = cudaFuncGetAttributes(&attr, kernel_ptr); - TORCH_CHECK(err == cudaSuccess, "Failed to preload kernel ", name, ": ", - cudaGetErrorString(err)); + PG_ASSERT(err == cudaSuccess, "Failed to preload kernel ", name, ": ", + cudaGetErrorString(err)); }; preload("reduceKernel", reinterpret_cast(reduceKernel)); @@ -209,7 +206,7 @@ void preloadReduceKernels() { preload("reduceKernel", reinterpret_cast(reduceKernel)); preload("reduceKernel", - reinterpret_cast(reduceKernel)); + reinterpret_cast(reduceKernel<__nv_bfloat16>)); #endif } diff --git a/mooncake-pg/src/mooncake_worker_host.cpp b/mooncake-pg/src/mooncake_worker_host.cpp index 98c6d99426..7ae890157f 100644 --- a/mooncake-pg/src/mooncake_worker_host.cpp +++ b/mooncake-pg/src/mooncake_worker_host.cpp @@ -2,295 +2,140 @@ // Compiled by g++ for both CUDA and MUSA builds. Uses kernel launch wrappers // from mooncake_worker_kernels.cuh instead of <<<>>> syntax. -#include -#include #include #include #include #include -#include -#include "pg_utils.h" +#include "error_types.h" namespace mooncake { -class MooncakeWorkCpu : public ::c10d::Work { - public: - MooncakeWorkCpu(c10d::OpType opType, - c10::intrusive_ptr future, - std::shared_ptr meta) - : Work(-1, opType), - future_(std::move(future)), - meta_(std::move(meta)) {} - - bool isCompleted() override { return future_->completed(); } - - bool wait(std::chrono::milliseconds timeout) override { - future_->wait(); - return future_->completed() && !future_->hasError(); - } - - private: - c10::intrusive_ptr future_; - std::shared_ptr meta_; -}; - -class MooncakeWorkCuda : public ::c10d::Work { - public: - MooncakeWorkCuda(c10d::OpType opType, std::shared_ptr event, - std::shared_ptr meta, - const MooncakeWorker* worker, - std::vector submitted_tasks) - : Work(-1, opType), - event_(std::move(event)), - meta_(std::move(meta)), - worker_(worker), - submitted_tasks_(std::move(submitted_tasks)) {} - - bool isCompleted() override { return event_->query(); } - - bool wait(std::chrono::milliseconds timeout) override { - // Wait until the task has been submitted to TransferEngine: - // This tries to ensure that the CUDA kernels required for the transfer - // have been launched by the time `waitUntilTasksSubmitted` returns. - // - // Why is this needed? PyTorch documentation implies that collective - // operations should be enqueued when `wait()` returns. In practice, we - // found that violating this causes hangs. - // - // Our current hypothesis for the hang is: PyTorch assumes the kernels - // needed for the transfer are already launched when `wait` returns - // true. It may then launch subsequent operations after the collective - // (e.g., `.cpu()`). Such operations may acquire a process-wide lock in - // the CUDA runtime. Also, they may rely on the data produced by the - // collective, thus causing a synchronization on enq_stream. However, - // holding that runtime lock prevents cudaMemcpy(Async) in TE/TENT from - // launching. This means the transfer can't finish, and enq_stream won't - // complete. Thus, a deadlock occurs. - // (In practice, we found that replacing all cudaMemcpyAsync in TENT - // with cuMemcpyAsync actually alleviates this, which further suggests a - // deadlock in the CUDA runtime. However, that change is too invasive - // for TE/TENT, so we do not adopt it here.) - // - // Strictly speaking, the wait is needed for another reason: The current - // stream will be blocked on the event below. Any subsequent work on - // `current_stream` will wait on that event, which effectively waits for - // the task to be done. Therefore, we must ensure all kernels needed for - // the transfer task are launched BEFORE blocking the current stream, in - // case TE/TENT use `current_stream` to launch those kernels (though it - // is rare). - // - // Please note that this logic relies on the assumption that TE/TENT - // will launch all CUDA operations in `submitTransfer`. - // Unfortunately, TcpTransport in TE and TENT currently violates this - // assumption (cudaMemcpy(Async) may be called later from a callback), - // which can cause hangs in PG when a CUDA operation such as - // `x.cpu().item()` follows the collective. For TE's TcpTransport, the - // use of cudaMemcpy on the default stream may also contribute to the - // hang. - // - // Besides, for CPU-only transports (like RdmaTransport), - // waitUntilTasksSubmitted is totally unnecessary, but we keep it for - // uniform behavior to avoid invasive changes to TE/TENT. - bool submitted = true; - if (at::cuda::currentStreamCaptureStatus() == - c10::cuda::CaptureStatus::None) { - // Normal execution: block until tasks are submitted. - submitted = - worker_->waitUntilTasksSubmitted(submitted_tasks_, timeout); - } else { - // During CUDA graph capture, kernels are recorded but not actually - // executed. The enqueueTaskKernel would never run, so - // waitUntilTasksSubmitted would hang because the CPU worker thread - // never sees task.active == true. - // - // Note that this also means NvlinkTransport (and TcpTransport too, - // of course) won't work with CUDA Graphs: Kernels launched inside - // TE/TENT can't be captured by the graph, and during replay they - // are not ordered with the graph execution. This may trigger the - // same deadlock described above. - } - if (!submitted) return false; - - // Once all tasks have been submitted, use the event to synchronize - // the current stream and the enqueue stream, but do not wait on this - // event. - // - // See PyTorch docs for more details: - // https://docs.pytorch.org/docs/stable/distributed.html#synchronous-and-asynchronous-collective-operations - // "wait() - in the case of CPU collectives, will block the process - // until the operation is completed. In the case of CUDA collectives, - // will block the currently active CUDA stream until the operation - // is completed (but will not block the CPU)." - auto current_stream = at::cuda::getCurrentCUDAStream(); - event_->block(current_stream); - return true; - } - - protected: - std::shared_ptr event_; - std::shared_ptr meta_; - const MooncakeWorker* worker_; - std::vector submitted_tasks_; -}; - -class MooncakeBarrierWorkCuda : public MooncakeWorkCuda { - public: - using MooncakeWorkCuda::MooncakeWorkCuda; - - bool wait(std::chrono::milliseconds timeout) override { - // Skip host-side synchronization during CUDA graph capture. - // cudaEventSynchronize is not permitted while a stream is capturing. - if (at::cuda::currentStreamCaptureStatus() != - c10::cuda::CaptureStatus::None) { - // We still need stream-level synchronization so that subsequent - // operations on the capture stream are ordered after the barrier - // task on the enqueue stream. - auto current_stream = at::cuda::getCurrentCUDAStream(); - event_->block(current_stream); - return true; - } - - if (timeout == kNoTimeout) { - event_->synchronize(); - return true; - } - - BackoffWaiter waiter( - BackoffWaiterConfig::constantSleep(std::chrono::microseconds(10))); - return waiter.wait_for(timeout, [this] { return event_->query(); }); - } -}; - -void launchReduceKernel(at::Tensor dst, size_t pos, size_t realSize, void* src, - size_t numRanks, c10d::ReduceOp op, bool* activeRanks, - cudaStream_t stream) { - TORCH_CHECK(op == c10d::ReduceOp::SUM || op == c10d::ReduceOp::MIN || - op == c10d::ReduceOp::MAX || op == c10d::ReduceOp::PRODUCT, - "Only support SUM/MIN/MAX/PRODUCT for reduction."); - auto ptr = (char*)dst.data_ptr() + pos; - size_t num = realSize / dst.element_size(); - - switch (dst.scalar_type()) { - case c10::kByte: +void launchReduceKernel(void* dst, DataType dataType, size_t pos, + size_t realSize, void* src, size_t numRanks, + ReduceOp op, bool* activeRanks, cudaStream_t stream) { + PG_ASSERT(op == ReduceOp::Sum || op == ReduceOp::Min || + op == ReduceOp::Max || op == ReduceOp::Product, + "Only support SUM/MIN/MAX/PRODUCT for reduction."); + auto ptr = (char*)dst + pos; + size_t num = realSize / elementSize(dataType); + + switch (dataType) { + case DataType::Uint8: launchReduceKernel_uint8((uint8_t*)ptr, (uint8_t*)src, num, numRanks, (int)op, activeRanks, stream); break; - case c10::kChar: + case DataType::Int8: launchReduceKernel_int8((int8_t*)ptr, (int8_t*)src, num, numRanks, (int)op, activeRanks, stream); break; - case c10::kShort: + case DataType::Int16: launchReduceKernel_int16((int16_t*)ptr, (int16_t*)src, num, numRanks, (int)op, activeRanks, stream); break; - case c10::kInt: + case DataType::Int32: launchReduceKernel_int32((int*)ptr, (int*)src, num, numRanks, (int)op, activeRanks, stream); break; - case c10::kLong: + case DataType::Int64: launchReduceKernel_int64((int64_t*)ptr, (int64_t*)src, num, numRanks, (int)op, activeRanks, stream); break; - case c10::kFloat: + case DataType::Float32: launchReduceKernel_float((float*)ptr, (float*)src, num, numRanks, (int)op, activeRanks, stream); break; - case c10::kDouble: + case DataType::Float64: launchReduceKernel_double((double*)ptr, (double*)src, num, numRanks, (int)op, activeRanks, stream); break; - case c10::kBool: + case DataType::Bool: launchReduceKernel_bool((bool*)ptr, (bool*)src, num, numRanks, (int)op, activeRanks, stream); break; - case c10::kBFloat16: + case DataType::Bfloat16: launchReduceKernel_bf16(ptr, src, num, numRanks, (int)op, activeRanks, stream); break; default: - TORCH_CHECK(false, c10::str("Unsupported reduce dtype: ", - dst.scalar_type())); + PG_ASSERT(false, "Unsupported reduce dtype: ", (int)dataType); } } template -T applyReduceOp(const T& a, const T& b, c10d::ReduceOp op) { +T applyReduceOp(const T& a, const T& b, ReduceOp op) { switch (op) { - case c10d::ReduceOp::SUM: + case ReduceOp::Sum: return a + b; - case c10d::ReduceOp::PRODUCT: + case ReduceOp::Product: return a * b; - case c10d::ReduceOp::MIN: + case ReduceOp::Min: return std::min(a, b); - case c10d::ReduceOp::MAX: + case ReduceOp::Max: return std::max(a, b); default: - TORCH_CHECK(false, c10::str("Unsupported reduce op: ", op)); + PG_ASSERT(false, "Unsupported reduce op: ", (int)op); } } template void reduceCpu(T* dst, const T* src, size_t numElements, size_t numRanks, - c10d::ReduceOp op, bool* activeRanks) { - at::parallel_for(0, numElements, 1024, [&](int64_t begin, int64_t end) { - for (int64_t i = begin; i < end; ++i) { - bool valid = false; - T acc{}; - for (int64_t rank = 0; rank < numRanks; ++rank) { - if (activeRanks[rank]) { - if (!valid) { - acc = src[i + rank * numElements]; - valid = true; - } else { - acc = - applyReduceOp(acc, src[i + rank * numElements], op); - } + ReduceOp op, bool* activeRanks) { + for (size_t i = 0; i < numElements; ++i) { + bool valid = false; + T acc{}; + for (size_t rank = 0; rank < numRanks; ++rank) { + if (activeRanks[rank]) { + if (!valid) { + acc = src[i + rank * numElements]; + valid = true; + } else { + acc = applyReduceOp(acc, src[i + rank * numElements], op); } } - dst[i] = acc; } - }); + dst[i] = acc; + } } -void launchReduceCpu(at::Tensor dst, size_t pos, size_t realSize, void* src, - size_t numRanks, c10d::ReduceOp op, bool* activeRanks) { - auto ptr = (char*)dst.data_ptr() + pos; - size_t num = realSize / dst.element_size(); +void launchReduceCpu(void* dst, DataType dataType, size_t pos, size_t realSize, + void* src, size_t numRanks, ReduceOp op, + bool* activeRanks) { + auto ptr = (char*)dst + pos; + size_t num = realSize / elementSize(dataType); - switch (dst.scalar_type()) { - case c10::kByte: + switch (dataType) { + case DataType::Uint8: reduceCpu((uint8_t*)ptr, (uint8_t*)src, num, numRanks, op, activeRanks); break; - case c10::kChar: + case DataType::Int8: reduceCpu((int8_t*)ptr, (int8_t*)src, num, numRanks, op, activeRanks); break; - case c10::kShort: + case DataType::Int16: reduceCpu((int16_t*)ptr, (int16_t*)src, num, numRanks, op, activeRanks); break; - case c10::kInt: + case DataType::Int32: reduceCpu((int*)ptr, (int*)src, num, numRanks, op, activeRanks); break; - case c10::kLong: + case DataType::Int64: reduceCpu((int64_t*)ptr, (int64_t*)src, num, numRanks, op, activeRanks); break; - case c10::kFloat: + case DataType::Float32: reduceCpu((float*)ptr, (float*)src, num, numRanks, op, activeRanks); break; - case c10::kDouble: + case DataType::Float64: reduceCpu((double*)ptr, (double*)src, num, numRanks, op, activeRanks); break; - case c10::kBool: + case DataType::Bool: reduceCpu((bool*)ptr, (bool*)src, num, numRanks, op, activeRanks); break; default: - TORCH_CHECK(false, c10::str("Unsupported reduce dtype: ", - dst.scalar_type())); + PG_ASSERT(false, "Unsupported reduce dtype: ", (int)dataType); } } @@ -302,13 +147,18 @@ MooncakeWorker::MooncakeWorker(int cuda_device_index) cudaHostAlloc(&tasks_, kNumTasks_ * sizeof(Task), cudaHostAllocMapped); cudaHostGetDevicePointer(&tasks_device_, tasks_, 0); } else { - LOG(WARNING) << "No GPU device found. Only the `mooncake-cpu` backend " - "can be used."; tasks_ = new Task[kNumTasks_]; } + + if (cuda_device_index_ >= 0) { + enqueue_stream_ = GpuStream::createNonBlocking(cuda_device_index_); + } + for (size_t i = 0; i < kNumTasks_; ++i) { tasks_[i].active = false; tasks_[i].submitSequence = 0; + tasks_[i].failedRanksHint = nullptr; + tasks_[i].resetFailedRanksHint = false; submitted_task_sequence_[i].store(0, std::memory_order_relaxed); } } @@ -320,19 +170,18 @@ MooncakeWorker::~MooncakeWorker() { } } -c10::intrusive_ptr MooncakeWorker::putTaskCpu( - c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot, - const std::shared_ptr& meta, - const std::shared_ptr& connection_ctx, +std::unique_ptr MooncakeWorker::putTaskCpu( + OpType opType, size_t tensorSize, int64_t broadcastRoot, + const std::shared_ptr& meta, int32_t* failed_ranks_hint, const std::function& - tensorToBuffer, + copyToSendBuffer, const std::function& - bufferToTensor) { - connection_ctx->waitUntilNewRanksConnected(); - - size_t chunkSize = ((kBufferSize - 1) / meta->size) & ~(size_t)7; - auto future = c10::make_intrusive( - c10::ListType::create(c10::TensorType::get())); + copyFromRecvBuffer) { + PG_ASSERT(failed_ranks_hint, "failed-ranks hint is null"); + size_t chunkSize = ((kBufferSize - 1) / meta->maxGroupSize) & ~(size_t)7; + auto completion = std::make_shared>(); + auto future = completion->get_future().share(); + auto result = std::make_unique(std::move(future)); struct IterState { size_t currentPos = 0; @@ -344,41 +193,37 @@ c10::intrusive_ptr MooncakeWorker::putTaskCpu( processNextChunk; *processNextChunk = [this, weakProcessNextChunk, state, opType, tensorSize, - chunkSize, broadcastRoot, meta, tensorToBuffer, - bufferToTensor, future]() { + chunkSize, broadcastRoot, meta, copyToSendBuffer, + copyFromRecvBuffer, completion, failed_ranks_hint]() { auto processNextChunk = weakProcessNextChunk.lock(); if (state->currentPos >= tensorSize) { - future->markCompleted(c10::IValue()); + completion->set_value(); return; } int taskId = cpuTaskCount % 2; - TORCH_CHECK(!tasks_[taskId].active); - + PG_ASSERT(!tasks_[taskId].active, + "collective CPU task slot is still active"); size_t realSize = std::min(chunkSize, tensorSize - state->currentPos); int bufferOffset = meta->taskCount % 2; - tasks_[taskId].opType = (int)opType; - tasks_[taskId].tensorSize = realSize; + tasks_[taskId].dataSize = realSize; tasks_[taskId].broadcastRoot = broadcastRoot; tasks_[taskId].bufferOffset = bufferOffset; + tasks_[taskId].submitSequence = 0; + tasks_[taskId].failedRanksHint = failed_ranks_hint; + tasks_[taskId].resetFailedRanksHint = state->currentPos == 0; tasks_[taskId].transferGroupMeta = meta.get(); - tensorToBuffer( + copyToSendBuffer( (void*)meta->segmentInfos[meta->rank].send_buffer[bufferOffset], state->currentPos, realSize); hasCallback_[taskId] = true; - callbacks_[taskId] = [this, processNextChunk, state, meta, - bufferToTensor, bufferOffset, realSize, - future]() { - if (meta->activeRanksTensor.device().is_cpu()) { - for (int i = 0; i < meta->size; ++i) { - meta->activeRanksTensor[i] = meta->activeRanks[i] ? 1 : 0; - } - } - bufferToTensor( + callbacks_[taskId] = [processNextChunk, state, meta, copyFromRecvBuffer, + bufferOffset, realSize, completion]() { + copyFromRecvBuffer( (void*)meta->segmentInfos[meta->rank].recv_buffer[bufferOffset], state->currentPos, realSize); @@ -394,28 +239,27 @@ c10::intrusive_ptr MooncakeWorker::putTaskCpu( (*processNextChunk)(); - return c10::make_intrusive(opType, future, meta); + return result; } -c10::intrusive_ptr MooncakeWorker::putTaskCuda( - c10d::OpType opType, size_t tensorSize, int64_t broadcastRoot, - const std::shared_ptr& meta, - const std::shared_ptr& connection_ctx, - const at::cuda::CUDAStream& issue_stream, +void MooncakeWorker::putTaskCuda( + OpType opType, size_t tensorSize, int64_t broadcastRoot, + const std::shared_ptr& meta, cudaStream_t issueStream, + int32_t* failed_ranks_hint, const std::function& tensorToBuffer, + cudaStream_t)>& copyToSendBuffer, const std::function& bufferToTensor) { - connection_ctx->waitUntilNewRanksConnected(); - - size_t chunkSize = ((kBufferSize - 1) / meta->size) & ~(size_t)7; + cudaStream_t)>& copyFromRecvBuffer) { + size_t chunkSize = ((kBufferSize - 1) / meta->maxGroupSize) & ~(size_t)7; - at::cuda::CUDAStream enq_stream = - at::cuda::getStreamFromPool(false, issue_stream.device_index()); + const GpuDeviceGuard guard(cuda_device_index_); + const auto issue_stream = + GpuStream::borrow(issueStream, cuda_device_index_); + const auto& enq_stream = enqueue_stream_.value(); - auto event_start = std::make_shared(torch::kCUDA); - event_start->record(issue_stream); - event_start->block(enq_stream); + GpuEvent event_start(issue_stream.deviceIndex()); + event_start.record(issue_stream); + enq_stream.waitEvent(event_start); std::vector submitted_tasks; submitted_tasks.reserve((tensorSize + chunkSize - 1) / chunkSize); @@ -423,37 +267,38 @@ c10::intrusive_ptr MooncakeWorker::putTaskCuda( size_t realSize = std::min(tensorSize, pos + chunkSize) - pos; int taskId = cudaTaskCount % 2 + 2; int bufferOffset = meta->taskCount % 2; + const uint64_t taskSequence = next_cuda_task_sequence_.fetch_add(1, std::memory_order_relaxed); submitted_tasks.push_back( {.task_id = static_cast(taskId), .sequence = taskSequence}); - tensorToBuffer( + copyToSendBuffer( (void*)meta->segmentInfos[meta->rank].send_buffer[bufferOffset], - pos, realSize, enq_stream); + pos, realSize, enq_stream.get()); hasCallback_[taskId] = false; - launchEnqueueTaskKernel( - (int)opType, realSize, broadcastRoot, bufferOffset, taskSequence, - meta.get(), tasks_device_, meta->size, meta->activeRanksDevice, - meta->activeRanksTensor.data_ptr(), taskId, - enq_stream.stream()); - bufferToTensor( + + launchEnqueueTaskKernel((int)opType, realSize, broadcastRoot, + bufferOffset, taskSequence, failed_ranks_hint, + pos == 0, meta.get(), tasks_device_, taskId, + enq_stream.get()); + copyFromRecvBuffer( (void*)meta->segmentInfos[meta->rank].recv_buffer[bufferOffset], - pos, realSize, enq_stream); + pos, realSize, enq_stream.get()); ++cudaTaskCount; ++meta->taskCount; } - auto event_end = std::make_shared(torch::kCUDA); - event_end->record(enq_stream); - - if (opType == c10d::OpType::BARRIER) { - return c10::make_intrusive( - opType, event_end, meta, this, std::move(submitted_tasks)); + // During CUDA graph capture the kernels are recorded but not executed, so + // waiting for the worker thread to observe them would hang. + if (!issue_stream.isCapturing()) { + waitUntilTasksSubmitted(submitted_tasks); } - return c10::make_intrusive(opType, event_end, meta, this, - std::move(submitted_tasks)); + + GpuEvent event_end(enq_stream.deviceIndex()); + event_end.record(enq_stream); + issue_stream.waitEvent(event_end); } } // namespace mooncake diff --git a/mooncake-pg/src/mooncake_worker_thread.cpp b/mooncake-pg/src/mooncake_worker_thread.cpp index 56e0c5c069..7e21a13195 100644 --- a/mooncake-pg/src/mooncake_worker_thread.cpp +++ b/mooncake-pg/src/mooncake_worker_thread.cpp @@ -1,9 +1,13 @@ #include +#include #include #include +#include #include #include +#include "control_plane/rpc.h" #include "pg_utils.h" +#include "control_plane/agent_host.h" namespace mooncake { @@ -16,13 +20,6 @@ enum WorkerTaskStatus { static constexpr size_t kInvalidTaskId = static_cast(-1); -static void setActiveRanksTensorValue(TransferGroupMeta* group, int rank, - int value) { - if (group->activeRanksTensor.device().is_cpu()) { - group->activeRanksTensor[rank] = value; - } -} - void MooncakeWorker::Start() { bool expected = false; if (started_.compare_exchange_strong(expected, true)) { @@ -42,11 +39,10 @@ bool MooncakeWorker::drainTasks(const TransferGroupMeta* meta) const { }); } -bool MooncakeWorker::waitUntilTasksSubmitted( - const std::vector& tasks, - std::chrono::milliseconds timeout) const { +void MooncakeWorker::waitUntilTasksSubmitted( + const std::vector& tasks) const { if (tasks.empty()) { - return true; + return; } auto submitted = [this, &tasks] { @@ -65,11 +61,7 @@ bool MooncakeWorker::waitUntilTasksSubmitted( BackoffWaiter waiter( BackoffWaiterConfig::constantSleep(std::chrono::microseconds(10))); - if (timeout == kNoTimeout) { - waiter.wait(submitted); - return true; - } - return waiter.wait_for(timeout, submitted); + waiter.wait(submitted); } void MooncakeWorker::startWorker() { @@ -80,8 +72,29 @@ void MooncakeWorker::startWorker() { } std::atomic task_status[kNumTasks_]; using clock = std::chrono::high_resolution_clock; + + // Per-slot state owned exclusively by this worker thread. clock::time_point activeTime[kNumTasks_]; + BatchID batchIDs[kNumTasks_]{}; size_t rankToTaskId[kNumTasks_][kMaxNumRanks]; + int32_t failedRanks[kNumTasks_][kMaxNumRanks]{}; + int32_t attemptedRanks[kNumTasks_][kMaxNumRanks]{}; + int32_t* active_failed_ranks_hint[kNumTasks_]{}; + bool task_detected_failure[kNumTasks_]{}; + + auto hasFailed = [&failedRanks](size_t task_id, int rank) { + return failedRanks[task_id][rank] != 0; + }; + auto markFailed = [&failedRanks](size_t task_id, int rank) { + failedRanks[task_id][rank] = 1; + }; + auto hasAttempted = [&attemptedRanks](size_t task_id, int rank) { + return attemptedRanks[task_id][rank] != 0; + }; + auto markAttempted = [&attemptedRanks](size_t task_id, int rank) { + attemptedRanks[task_id][rank] = 1; + }; + while (running_) { PAUSE(); for (size_t i = 0; i < kNumTasks_; ++i) { @@ -92,14 +105,32 @@ void MooncakeWorker::startWorker() { } auto group = (TransferGroupMeta*)task.transferGroupMeta; - bool skipTransfer = - ((c10d::OpType)task.opType == c10d::OpType::BROADCAST && - group->rank != task.broadcastRoot) || - ((c10d::OpType)task.opType == c10d::OpType::SCATTER && - group->rank != task.broadcastRoot) || - (c10d::OpType)task.opType == c10d::OpType::BARRIER; + const auto op_type = static_cast(task.opType); + bool skipTransfer = (op_type == OpType::Broadcast && + group->rank != task.broadcastRoot) || + (op_type == OpType::Scatter && + group->rank != task.broadcastRoot) || + op_type == OpType::Barrier; + if (task_status[i].load(std::memory_order_acquire) == IDLE) { const auto submit_sequence = task.submitSequence; + // A slot is reused by unrelated tasks. Start with fresh + // task-local observations so stale failures/attempts from + // the previous occupant cannot affect this task. + for (size_t j = 0; j < kMaxNumRanks; ++j) { + failedRanks[i][j] = 0; + attemptedRanks[i][j] = 0; + } + task_detected_failure[i] = false; + active_failed_ranks_hint[i] = task.failedRanksHint; + + if (task.resetFailedRanksHint && + active_failed_ranks_hint[i]) { + // A graph replay carries the captured first-chunk flag, + // so it starts a fresh hint in the same tensor storage. + std::fill_n(active_failed_ranks_hint[i], + group->maxGroupSize, int32_t{0}); + } if (skipTransfer) { submitted_task_sequence_[i].store( submit_sequence, std::memory_order_release); @@ -111,33 +142,31 @@ void MooncakeWorker::startWorker() { rankToTaskId[i][j] = kInvalidTaskId; } std::vector entries; - for (int j = 0; j < group->size; ++j) { - if (!group->activeRanks[j]) { + entries.reserve(group->maxGroupSize); + for (int j = 0; j < group->maxGroupSize; ++j) { + if (!group->activeRanks[j] || hasFailed(i, j)) { continue; } - if (((c10d::OpType)task.opType == - c10d::OpType::GATHER || - (c10d::OpType)task.opType == - c10d::OpType::REDUCE) && + if ((op_type == OpType::Gather || + op_type == OpType::Reduce) && j != task.broadcastRoot) { continue; } + uint64_t source = group->segmentInfos[group->rank] .send_buffer[task.bufferOffset]; - switch ((c10d::OpType)task.opType) { - case c10d::OpType::BROADCAST: - case c10d::OpType::ALLREDUCE: - case c10d::OpType::ALLGATHER: - case c10d::OpType::_ALLGATHER_BASE: - case c10d::OpType::REDUCE: - case c10d::OpType::GATHER: + switch (op_type) { + case OpType::Broadcast: + case OpType::AllReduce: + case OpType::AllGather: + case OpType::Reduce: + case OpType::Gather: break; - case c10d::OpType::ALLTOALL_BASE: - case c10d::OpType::ALLTOALL: - case c10d::OpType::_REDUCE_SCATTER_BASE: - case c10d::OpType::SCATTER: - source += j * task.tensorSize; + case OpType::AllToAll: + case OpType::ReduceScatter: + case OpType::Scatter: + source += j * task.dataSize; break; default: break; @@ -146,19 +175,17 @@ void MooncakeWorker::startWorker() { group->segmentInfos[j] .recv_buffer[task.bufferOffset]; - switch ((c10d::OpType)task.opType) { - case c10d::OpType::BROADCAST: - case c10d::OpType::SCATTER: + switch (op_type) { + case OpType::Broadcast: + case OpType::Scatter: break; - case c10d::OpType::ALLREDUCE: - case c10d::OpType::ALLGATHER: - case c10d::OpType::_ALLGATHER_BASE: - case c10d::OpType::ALLTOALL_BASE: - case c10d::OpType::ALLTOALL: - case c10d::OpType::_REDUCE_SCATTER_BASE: - case c10d::OpType::REDUCE: - case c10d::OpType::GATHER: - target_offset += group->rank * task.tensorSize; + case OpType::AllReduce: + case OpType::AllGather: + case OpType::AllToAll: + case OpType::ReduceScatter: + case OpType::Reduce: + case OpType::Gather: + target_offset += group->rank * task.dataSize; break; default: @@ -170,12 +197,15 @@ void MooncakeWorker::startWorker() { .source = (void*)source, .target_id = group->segmentIDs[j], .target_offset = target_offset, - .length = task.tensorSize, + .length = task.dataSize, }); + + // Attempted to transfer to this peer + markAttempted(i, j); } - task.batchID = + batchIDs[i] = group->engine->allocateBatchID(entries.size()); - group->engine->submitTransfer(task.batchID, entries); + group->engine->submitTransfer(batchIDs[i], entries); submitted_task_sequence_[i].store( submit_sequence, std::memory_order_release); activeTime[i] = clock::now(); @@ -190,33 +220,38 @@ void MooncakeWorker::startWorker() { auto now = clock::now(); auto diff = std::chrono::duration_cast< std::chrono::microseconds>(now - activeTime[i]); - for (int j = 0; j < group->size; ++j) { - if (!group->activeRanks[j]) { + for (int j = 0; j < group->maxGroupSize; ++j) { + if (!group->activeRanks[j] || hasFailed(i, j)) { continue; } if (rankToTaskId[i][j] == kInvalidTaskId) { continue; } group->engine->getTransferStatus( - task.batchID, rankToTaskId[i][j], status); + batchIDs[i], rankToTaskId[i][j], status); if (status.s != TransferStatusEnum::COMPLETED) { - if (status.s == TransferStatusEnum::FAILED || - (j != group->rank && - diff.count() > kPingTimeoutMicroseconds_ && - group->engine->probePeerAliveByID( - group->segmentIDs[j]) != - PeerLiveness::Alive)) { + bool peer_dead = false; + if (status.s == TransferStatusEnum::FAILED) { + peer_dead = true; + } else if (j != group->rank && + diff.count() > + *group->collectiveTimeoutUs) { + peer_dead = + group->engine->probePeerAliveByID( + group->segmentIDs[j]) != + PeerLiveness::Alive; + } + if (peer_dead) { + markFailed(i, j); + task_detected_failure[i] = true; LOG(ERROR) - << "Rank " << group->rank - << " marking peer " << j - << " as broken during transferring op " - << (int)task.opType; - - // Set peerConnected to notify the - // connection poller to reconnect it. - group->peerConnected[j] = false; - group->activeRanks[j] = false; - setActiveRanksTensorValue(group, j, 0); + << "Rank " << group->globalRank + << " [DATA] transfer to peer " << j + << " (global=" << group->rank_order[j] + << ") failed for op " + << (int)task.opType + << " status=" << (int)status.s + << " diff_us=" << diff.count(); } else { batch_done = false; break; @@ -230,7 +265,7 @@ void MooncakeWorker::startWorker() { } if (!skipTransfer) { - auto s = group->engine->freeBatchID(task.batchID); + auto s = group->engine->freeBatchID(batchIDs[i]); if (!s.ok()) { LOG(WARNING) << "BatchID leaked due to freeBatchID " @@ -246,8 +281,9 @@ void MooncakeWorker::startWorker() { rankToTaskId[i][j] = kInvalidTaskId; } std::vector entries; - for (int j = 0; j < group->size; ++j) { - if (!group->activeRanks[j]) { + entries.reserve(group->maxGroupSize); + for (int j = 0; j < group->maxGroupSize; ++j) { + if (!group->activeRanks[j] || hasFailed(i, j)) { continue; } *source_ptr = 1; @@ -261,10 +297,11 @@ void MooncakeWorker::startWorker() { group->rank * sizeof(int32_t), .length = sizeof(int32_t), }); + markAttempted(i, j); } - task.batchID = + batchIDs[i] = group->engine->allocateBatchID(entries.size()); - group->engine->submitTransfer(task.batchID, entries); + group->engine->submitTransfer(batchIDs[i], entries); activeTime[i] = clock::now(); task_status[i].store(SIGNALED_1, std::memory_order_release); } else if (task_status[i].load(std::memory_order_acquire) == @@ -279,47 +316,112 @@ void MooncakeWorker::startWorker() { now - activeTime[i]); TransferStatus status; - for (int j = 0; j < group->size; ++j) { - if (!group->activeRanks[j]) { + for (int j = 0; j < group->maxGroupSize; ++j) { + if (!group->activeRanks[j] || hasFailed(i, j)) { continue; } if (rankToTaskId[i][j] == kInvalidTaskId) { continue; } group->engine->getTransferStatus( - task.batchID, rankToTaskId[i][j], status); + batchIDs[i], rankToTaskId[i][j], status); if (signal_ptr[j] != 1 || status.s != TransferStatusEnum::COMPLETED) { - if (status.s == TransferStatusEnum::FAILED || - (j != group->rank && - diff.count() > kPingTimeoutMicroseconds_ && - group->engine->probePeerAliveByID( - group->segmentIDs[j]) != - PeerLiveness::Alive)) { - LOG(ERROR) << "Rank " << group->rank - << " marking peer " << j - << " as broken during syncing op " - << (int)task.opType; - - // Set peerConnected to notify the - // connection poller to reconnect it. - group->peerConnected[j] = false; - group->activeRanks[j] = false; - setActiveRanksTensorValue(group, j, 0); + bool peer_dead = false; + if (status.s == TransferStatusEnum::FAILED) { + peer_dead = true; + } else if (j != group->rank && + diff.count() > + *group->collectiveTimeoutUs) { + peer_dead = group->engine->probePeerAliveByID( + group->segmentIDs[j]) != + PeerLiveness::Alive; + } + if (peer_dead) { + markFailed(i, j); + task_detected_failure[i] = true; + LOG(ERROR) + << "Rank " << group->globalRank + << " [SYNC] sync to peer " << j + << " (global=" << group->rank_order[j] + << ") failed for op " << (int)task.opType + << " status=" << (int)status.s + << " signal_ptr=" << (int)signal_ptr[j] + << " diff_us=" << diff.count(); } else { task_done = false; break; } } } - if (diff.count() > kPingTimeoutMicroseconds_) { + if (diff.count() > *group->collectiveTimeoutUs) { // reset timer activeTime[i] = clock::now(); } if (task_done) { - for (int j = 0; j < group->size; ++j) { + for (int j = 0; j < group->maxGroupSize; ++j) { signal_ptr[j] = 0; } + + if (active_failed_ranks_hint[i]) { + for (int rank = 0; rank < group->maxGroupSize; + ++rank) { + active_failed_ranks_hint[i][rank] |= + failedRanks[i][rank]; + } + active_failed_ranks_hint[i] = nullptr; + } + + // Push link event via communicator's Agent. + if (group->communicator) { + bool has_any_attempted = false; + for (int j = 0; j < group->maxGroupSize; ++j) { + if (hasAttempted(i, j)) { + has_any_attempted = true; + break; + } + } + if (has_any_attempted) { + LinkEvent event; + event.events.assign(kMaxNumRanks, + LinkEvent::EventType::None); + event.target_rank_epochs.assign(kMaxNumRanks, + 0); + for (int j = 0; j < group->maxGroupSize; ++j) { + const auto peer_global = + group->rank_order[j]; + if (!hasAttempted(i, j)) continue; + event.events[peer_global] = + hasFailed(i, j) + ? LinkEvent::EventType::Failure + : LinkEvent::EventType::Success; + event.target_rank_epochs[peer_global] = + group->rankEpochs[peer_global]; + } + group->communicator->getAgent().pushLinkEvent( + event); + } + } + + auto s = group->engine->freeBatchID(batchIDs[i]); + if (!s.ok()) { + LOG(WARNING) + << "BatchID leaked due to freeBatchID " + "failure (likely caused by a timeout): " + << s.message(); + } + + if (task_detected_failure[i] && + group->autoSyncOnFailure) { + auto result = + group->communicator->syncAfterFailure(); + PG_ASSERT(result.has_value() && + result.value().status != + SyncAfterFailureStatus::Rejected, + "syncAfterFailure failed for rank ", + group->globalRank); + } + task_status[i].store(DONE, std::memory_order_release); task.active = false; if (hasCallback_[i]) { @@ -329,13 +431,6 @@ void MooncakeWorker::startWorker() { hasCallback_[i] = false; callback(); } - auto s = group->engine->freeBatchID(task.batchID); - if (!s.ok()) { - LOG(WARNING) - << "BatchID leaked due to freeBatchID " - "failure (likely caused by a timeout): " - << s.message(); - } } } } diff --git a/mooncake-pg/src/p2p_proxy.cpp b/mooncake-pg/src/p2p_proxy.cpp index e56dcdf0c0..074eb692fe 100644 --- a/mooncake-pg/src/p2p_proxy.cpp +++ b/mooncake-pg/src/p2p_proxy.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -10,6 +11,7 @@ #include #include "memory_location.h" #include "pg_utils.h" +#include "control_plane/agent_host.h" namespace mooncake { @@ -25,11 +27,10 @@ void setCudaDeviceIfNeeded(bool is_cpu, int cuda_device_index, if (is_cpu) { return; } - TORCH_CHECK(cuda_device_index >= 0, context, - ": invalid CUDA device index."); + PG_ASSERT(cuda_device_index >= 0, context, ": invalid CUDA device index."); const cudaError_t set_device_error = cudaSetDevice(cuda_device_index); - TORCH_CHECK(set_device_error == cudaSuccess, context, ": ", - cudaGetErrorString(set_device_error)); + PG_ASSERT(set_device_error == cudaSuccess, context, ": ", + cudaGetErrorString(set_device_error)); } } // namespace @@ -162,13 +163,13 @@ P2PProxy::P2PProxy(TransferEngine* engine, const Options& options) rank_(options.rank), size_(options.size), cuda_device_index_(options.cuda_device_index), - transfer_timeout_ms_(options.transfer_timeout_ms) { + p2p_timeout_us_(options.p2p_timeout_us) { if (!is_cpu_ && cuda_device_index_ < 0) { int current_device = -1; const cudaError_t get_device_error = cudaGetDevice(¤t_device); - TORCH_CHECK(get_device_error == cudaSuccess, - "P2PProxy cudaGetDevice failed: ", - cudaGetErrorString(get_device_error)); + PG_ASSERT(get_device_error == cudaSuccess, + "P2PProxy cudaGetDevice failed: ", + cudaGetErrorString(get_device_error)); cuda_device_index_ = current_device; } allocateResources(); @@ -189,7 +190,7 @@ void P2PProxy::bindMeta(const std::shared_ptr& meta) { } void P2PProxy::extendGroupSizeTo(int new_size) { - TORCH_CHECK(new_size >= size_, "extendGroupSizeTo: new_size < size_"); + PG_ASSERT(new_size >= size_, "extendGroupSizeTo: new_size < size_"); size_ = new_size; } @@ -197,15 +198,15 @@ void P2PProxy::extendGroupSizeTo(int new_size) { // Send/Recv chunk pools are owned by P2PDeviceWorker and shared across // multiple P2PProxy instances on the same device. void P2PProxy::allocateResources() { - TORCH_CHECK(engine_, "P2PProxy engine is null."); + PG_ASSERT(engine_, "P2PProxy engine is null."); if (resources_.credit_region_ != nullptr || resources_.ack_region_ != nullptr) { return; } - TORCH_CHECK(size_ > 0, "P2PProxy invalid group size: ", size_); - TORCH_CHECK(static_cast(size_) <= kMaxNumRanks, - "P2PProxy group size exceeds kMaxNumRanks: ", size_); + PG_ASSERT(size_ > 0, "P2PProxy invalid group size: ", size_); + PG_ASSERT(static_cast(size_) <= kMaxNumRanks, + "P2PProxy group size exceeds kMaxNumRanks: ", size_); const size_t ctrl_slots = kMaxNumRanks * static_cast(kP2PControlRingSize); @@ -240,18 +241,18 @@ void P2PProxy::allocateResources() { send_peer_lanes_[peer_rank].copy_ready_events_) { const cudaError_t create_error = cudaEventCreateWithFlags( ©_ready_event, cudaEventDisableTiming); - TORCH_CHECK(create_error == cudaSuccess, - "Failed to create pooled send copy-ready event: ", - cudaGetErrorString(create_error)); + PG_ASSERT(create_error == cudaSuccess, + "Failed to create pooled send copy-ready event: ", + cudaGetErrorString(create_error)); } for (auto& copy_ready_event : recv_peer_lanes_[peer_rank].copy_ready_events_) { const cudaError_t create_error = cudaEventCreateWithFlags( ©_ready_event, cudaEventDisableTiming); - TORCH_CHECK(create_error == cudaSuccess, - "Failed to create pooled recv copy-ready event: ", - cudaGetErrorString(create_error)); + PG_ASSERT(create_error == cudaSuccess, + "Failed to create pooled recv copy-ready event: ", + cudaGetErrorString(create_error)); } } } @@ -259,12 +260,12 @@ void P2PProxy::allocateResources() { int rc = engine_->registerLocalMemory(resources_.credit_region_, ctrl_slots * sizeof(CreditSlot), kWildcardLocation); - TORCH_CHECK(rc == 0, "Failed to register P2P credit region"); + PG_ASSERT(rc == 0, "Failed to register P2P credit region"); rc = engine_->registerLocalMemory(resources_.ack_region_, ctrl_slots * sizeof(AckSlot), kWildcardLocation); - TORCH_CHECK(rc == 0, "Failed to register P2P ack region"); + PG_ASSERT(rc == 0, "Failed to register P2P ack region"); // Staging buffers for control messages. RDMA requires every // transfer source to live in a registered MR. @@ -272,19 +273,19 @@ void P2PProxy::allocateResources() { rc = engine_->registerLocalMemory(resources_.credit_staging_buf_, ctrl_slots * sizeof(CreditSlot), kWildcardLocation); - TORCH_CHECK(rc == 0, "Failed to register P2P credit staging region"); + PG_ASSERT(rc == 0, "Failed to register P2P credit staging region"); resources_.ack_staging_buf_ = new AckSlot[ctrl_slots]{}; rc = engine_->registerLocalMemory(resources_.ack_staging_buf_, ctrl_slots * sizeof(AckSlot), kWildcardLocation); - TORCH_CHECK(rc == 0, "Failed to register P2P ack staging region"); + PG_ASSERT(rc == 0, "Failed to register P2P ack staging region"); } void P2PProxy::resetPeerState(int peer_rank) { - TORCH_CHECK(peer_rank >= 0 && peer_rank < size_, - "ResetPeerState: peer_rank out of range: ", peer_rank, - " size: ", size_); + PG_ASSERT(peer_rank >= 0 && peer_rank < size_, + "ResetPeerState: peer_rank out of range: ", peer_rank, + " size: ", size_); // Epoch update: // We bump epoch_ so that any slots still in flight from the @@ -306,25 +307,73 @@ void P2PProxy::resetPeerState(int peer_rank) { // Reset sender state for peer_rank. void P2PProxy::performSendReset(int peer_rank) { resetSendLane(send_peer_lanes_[peer_rank]); - resetPeerControlLanes(peer_rank); + resetPeerAckLanes(peer_rank); } // Reset receiver state for peer_rank. void P2PProxy::performRecvReset(int peer_rank) { resetRecvLane(recv_peer_lanes_[peer_rank]); - resetPeerControlLanes(peer_rank); -} - -void P2PProxy::reportBrokenPeer(int peer_rank) { - resetPeerState(peer_rank); - // Set peerConnected to notify the connection poller to reconnect it. - meta_->peerConnected[peer_rank] = false; - meta_->activeRanks[peer_rank] = false; - if (meta_->activeRanksTensor.device().is_cpu()) { - meta_->activeRanksTensor[peer_rank] = 0; - } - LOG(ERROR) << "Rank " << meta_->rank << " marking peer " << peer_rank - << " as broken during P2P transfer."; + resetPeerCreditLanes(peer_rank); +} + +void P2PProxy::cleanupFailedSendOp(SendOpContext& op_ctx) { + for (auto& task : op_ctx.tasks_) { + releaseSendTaskResources(task); + } + op_ctx.tasks_.clear(); + active_send_tasks_.fetch_sub(1, std::memory_order_release); +} + +void P2PProxy::cleanupFailedRecvOp(RecvOpContext& op_ctx) { + for (auto& task : op_ctx.tasks_) { + releaseRecvTaskResources(task); + } + op_ctx.tasks_.clear(); + active_recv_tasks_.fetch_sub(1, std::memory_order_release); +} + +void P2PProxy::handleFailedSendOp(SendOpContext& op_ctx) { + cleanupFailedSendOp(op_ctx); + op_ctx.failed_ranks_hint_[op_ctx.peer_rank_] = 1; + // Reset P2P session state (epoch, lanes). + resetPeerState(op_ctx.peer_rank_); + // link event vector is indexed by GlobalRank. + const auto peer_global = meta_->rank_order[op_ctx.peer_rank_]; + const auto target_rank_epoch = meta_->rankEpochs[peer_global]; + if (meta_->communicator) { + LinkEvent event; + event.events.assign(kMaxNumRanks, LinkEvent::EventType::None); + event.target_rank_epochs.assign(kMaxNumRanks, 0); + event.events[peer_global] = LinkEvent::EventType::Failure; + event.target_rank_epochs[peer_global] = target_rank_epoch; + meta_->communicator->getAgent().pushLinkEvent(event); + } + op_ctx.completion_->set_value(); + LOG(ERROR) << "Rank " << meta_->rank << ": P2P SendOp to peer " + << op_ctx.peer_rank_ << " (global=" << peer_global + << ") failed."; +} + +void P2PProxy::handleFailedRecvOp(RecvOpContext& op_ctx) { + cleanupFailedRecvOp(op_ctx); + op_ctx.failed_ranks_hint_[op_ctx.peer_rank_] = 1; + // Reset P2P session state (epoch, lanes). + resetPeerState(op_ctx.peer_rank_); + // link event vector is indexed by GlobalRank. + const auto peer_global = meta_->rank_order[op_ctx.peer_rank_]; + const auto target_rank_epoch = meta_->rankEpochs[peer_global]; + if (meta_->communicator) { + LinkEvent event; + event.events.assign(kMaxNumRanks, LinkEvent::EventType::None); + event.target_rank_epochs.assign(kMaxNumRanks, 0); + event.events[peer_global] = LinkEvent::EventType::Failure; + event.target_rank_epochs[peer_global] = target_rank_epoch; + meta_->communicator->getAgent().pushLinkEvent(event); + } + op_ctx.completion_->set_value(); + LOG(ERROR) << "Rank " << meta_->rank << ": P2P RecvOp from peer " + << op_ctx.peer_rank_ << " (global=" << peer_global + << ") failed."; } void P2PProxy::releaseSendTaskResources(SendTransferTask& task) const { @@ -355,19 +404,12 @@ void P2PProxy::releaseRecvTaskResources(RecvTransferTask& task) const { void P2PProxy::resetSendLane(SendPeerLane& lane) { for (auto& pending : lane.pending_send_ops_) { - pending.status_->store(OpStatus::kFailed, std::memory_order_release); - active_send_tasks_.fetch_sub(1, std::memory_order_release); + handleFailedSendOp(pending); } lane.pending_send_ops_.clear(); if (lane.active_send_op_.has_value()) { - auto& op_ctx = lane.active_send_op_.value(); - for (auto& task : op_ctx.tasks_) { - releaseSendTaskResources(task); - } - op_ctx.tasks_.clear(); - op_ctx.status_->store(OpStatus::kFailed, std::memory_order_release); - active_send_tasks_.fetch_sub(1, std::memory_order_release); + handleFailedSendOp(*lane.active_send_op_); lane.active_send_op_.reset(); } lane.credit_consume_seq_ = 0; @@ -375,37 +417,34 @@ void P2PProxy::resetSendLane(SendPeerLane& lane) { void P2PProxy::resetRecvLane(RecvPeerLane& lane) { for (auto& pending : lane.pending_recv_ops_) { - pending.status_->store(OpStatus::kFailed, std::memory_order_release); - active_recv_tasks_.fetch_sub(1, std::memory_order_release); + handleFailedRecvOp(pending); } lane.pending_recv_ops_.clear(); if (lane.active_recv_op_.has_value()) { - auto& op_ctx = lane.active_recv_op_.value(); - for (auto& task : op_ctx.tasks_) { - releaseRecvTaskResources(task); - } - op_ctx.tasks_.clear(); - op_ctx.status_->store(OpStatus::kFailed, std::memory_order_release); - active_recv_tasks_.fetch_sub(1, std::memory_order_release); + handleFailedRecvOp(*lane.active_recv_op_); lane.active_recv_op_.reset(); } lane.credit_issue_seq_ = 0; lane.ack_consume_seq_ = 0; } -void P2PProxy::resetPeerControlLanes(int peer_rank) { +void P2PProxy::resetPeerCreditLanes(int peer_rank) { auto* credit_lane = getLocalCreditLane(peer_rank); - auto* ack_lane = getLocalAckLane(peer_rank); for (uint32_t i = 0; i < kP2PControlRingSize; ++i) { credit_lane[i].reset(); + } +} + +void P2PProxy::resetPeerAckLanes(int peer_rank) { + auto* ack_lane = getLocalAckLane(peer_rank); + for (uint32_t i = 0; i < kP2PControlRingSize; ++i) { ack_lane[i].reset(); } } void P2PProxy::releaseResources() { - TORCH_CHECK(!resource_abandoned_, - "Should not release abandoned resources."); + PG_ASSERT(!resource_abandoned_, "Should not release abandoned resources."); setCudaDeviceIfNeeded(is_cpu_, cuda_device_index_, "P2PProxy ReleaseResources cudaSetDevice failed"); @@ -422,9 +461,9 @@ void P2PProxy::releaseResources() { for (auto& ev : events) { if (ev == nullptr) continue; const cudaError_t err = cudaEventDestroy(ev); - TORCH_CHECK(err == cudaSuccess, - "Failed to destroy pooled copy-ready event: ", - cudaGetErrorString(err)); + PG_ASSERT(err == cudaSuccess, + "Failed to destroy pooled copy-ready event: ", + cudaGetErrorString(err)); ev = nullptr; } }; @@ -462,9 +501,6 @@ void P2PProxy::releaseResources() { void P2PProxy::abandonResources() { resource_abandoned_ = true; } void P2PProxy::enqueueSend(SendOp op) { - op.tensor_ = - op.tensor_.is_contiguous() ? op.tensor_ : op.tensor_.contiguous(); - { std::lock_guard lock(send_queue_mutex_); send_queue_.emplace(std::move(op)); @@ -476,16 +512,16 @@ void P2PProxy::enqueueSend(SendOp op) { void P2PProxy::enqueueRecv(RecvOp op) { { std::lock_guard lock(recv_queue_mutex_); - recv_queue_.push(std::move(op)); + recv_queue_.emplace(std::move(op)); } active_recv_tasks_.fetch_add(1, std::memory_order_release); if (device_worker_) device_worker_->wakeUpRecv(); } P2PProxy::SendTransferTask::SendTransferTask( - uint64_t tensor_offset_in, uint32_t chunk_len_in, void* staging_addr_in, + uint64_t buffer_offset_in, uint32_t chunk_len_in, void* staging_addr_in, uint64_t remote_addr_in, uint32_t sequence_in, uint32_t epoch_in) - : tensor_offset_(tensor_offset_in), + : buffer_offset_(buffer_offset_in), chunk_len_(chunk_len_in), staging_addr_(staging_addr_in), remote_addr_(remote_addr_in), @@ -495,21 +531,21 @@ P2PProxy::SendTransferTask::SendTransferTask( } P2PProxy::SendOpContext::SendOpContext(SendOp&& op_in) - : status_(std::move(op_in.status_)), - tensor_(std::move(op_in.tensor_)), + : completion_(std::move(op_in.completion_)), + buffer_(op_in.buffer_), peer_rank_(op_in.peer_rank_), - cuda_stream_(op_in.cuda_stream_) { - total_bytes_ = - tensor_.numel() * static_cast(tensor_.element_size()); + cuda_stream_(op_in.cuda_stream_), + failed_ranks_hint_(op_in.failed_ranks_hint_) { + total_bytes_ = op_in.size_; last_update_time_ = std::chrono::steady_clock::now(); } -P2PProxy::RecvTransferTask::RecvTransferTask(uint64_t tensor_offset_in, +P2PProxy::RecvTransferTask::RecvTransferTask(uint64_t buffer_offset_in, uint32_t chunk_len_in, void* local_addr_in, uint32_t sequence_in, uint32_t epoch_in) - : tensor_offset_(tensor_offset_in), + : buffer_offset_(buffer_offset_in), chunk_len_(chunk_len_in), local_addr_(local_addr_in), sequence_(sequence_in), @@ -518,13 +554,12 @@ P2PProxy::RecvTransferTask::RecvTransferTask(uint64_t tensor_offset_in, } P2PProxy::RecvOpContext::RecvOpContext(RecvOp&& op_in) - : status_(std::move(op_in.status_)), - tensor_(std::move(op_in.tensor_)), - original_tensor_(std::move(op_in.original_tensor_)), + : completion_(std::move(op_in.completion_)), + buffer_(op_in.buffer_), peer_rank_(op_in.peer_rank_), - cuda_stream_(op_in.cuda_stream_) { - total_bytes_ = - tensor_.numel() * static_cast(tensor_.element_size()); + cuda_stream_(op_in.cuda_stream_), + failed_ranks_hint_(op_in.failed_ranks_hint_) { + total_bytes_ = op_in.size_; } CreditSlot* P2PProxy::getLocalCreditLane(int peer_rank) const { @@ -646,8 +681,8 @@ bool P2PProxy::stepRecvTask(RecvTransferTask& task) { } bool P2PProxy::stepRecvIssueCredit(RecvTransferTask& task) { - TORCH_CHECK(task.credit_batch_id_.has_value(), - "Expected a credit_batch_id in tryIssueRecvTask"); + PG_ASSERT(task.credit_batch_id_.has_value(), + "Expected a credit_batch_id in tryIssueRecvTask"); TransferStatus credit_status; engine_->getTransferStatus(task.credit_batch_id_.value(), 0, credit_status); @@ -674,11 +709,10 @@ bool P2PProxy::stepRecvIssueCredit(RecvTransferTask& task) { // Poll the GPU Copy-Out event. // // After the AckSlot arrives we initiate cudaMemcpyAsync from the -// RecvPool chunk to the user tensor and record an event. This function +// RecvPool chunk to the user buffer and record an event. This function // waits for that event and returns the chunk to RecvPool. bool P2PProxy::stepRecvCopyOut(RecvTransferTask& task) { - TORCH_CHECK(task.copy_ready_event_ != nullptr, - "Expected a copy_ready_event"); + PG_ASSERT(task.copy_ready_event_ != nullptr, "Expected a copy_ready_event"); cudaError_t query_error = cudaEventQuery(task.copy_ready_event_); if (query_error == cudaSuccess) { task.copy_ready_event_ = nullptr; @@ -692,8 +726,8 @@ bool P2PProxy::stepRecvCopyOut(RecvTransferTask& task) { } task.copy_ready_event_ = nullptr; - TORCH_CHECK(false, "P2P recv cudaEventQuery failed: ", - cudaGetErrorString(query_error)); + PG_ASSERT(false, "P2P recv cudaEventQuery failed: ", + cudaGetErrorString(query_error)); return false; } @@ -701,23 +735,23 @@ bool P2PProxy::stepRecvCopyOut(RecvTransferTask& task) { // // Poll the local CreditLane for the next expected sequence. If the slot // matches our consume cursor we accept the credit: allocate a staging -// chunk from SendPool, copy the corresponding slice of the user tensor into +// chunk from SendPool, copy the corresponding slice of the user buffer into // it, and advance the cursor. If the pool is full or no credit has // arrived we return false immediately (non-blocking). -bool P2PProxy::tryIssueSendTask(SendOpContext& op_ctx, SendPeerLane& lane) { +P2PProxy::IssueResult P2PProxy::tryIssueSendTask(SendOpContext& op_ctx, + SendPeerLane& lane) { // Check for timeout while waiting for the peer's CreditSlot. if (isTimeout(op_ctx)) { LOG(ERROR) << "P2P wait-for-credit timeout, peer=" << op_ctx.peer_rank_; - op_ctx.status_->store(OpStatus::kFailed, std::memory_order_release); - return false; + return IssueResult::kTimeout; } if (op_ctx.bytes_staged_ >= op_ctx.total_bytes_) { - return false; + return IssueResult::kNoCredit; } if (op_ctx.tasks_.size() >= kP2PControlRingSize) { - return false; + return IssueResult::kNoCredit; } CreditSlot* local_credit = getLocalCreditLane(op_ctx.peer_rank_); @@ -731,7 +765,7 @@ bool P2PProxy::tryIssueSendTask(SendOpContext& op_ctx, SendPeerLane& lane) { uint32_t slot_seq = 0; if (!slot.tryLoad(recv_addr, chunk_len, slot_epoch, slot_seq)) { // Slot is either empty or torn (partial RDMA write). Retry next poll. - return false; + return IssueResult::kNoCredit; } // Step 2 -- Stale packet: the slot carries data from a previous epoch @@ -744,24 +778,24 @@ bool P2PProxy::tryIssueSendTask(SendOpContext& op_ctx, SendPeerLane& lane) { << " slot.epoch=" << slot_epoch << " curr_epoch=" << curr_epoch; slot.reset(); - return true; + return IssueResult::kIssued; } // Step 3 -- Sequence check: make sure this is exactly the slot we expect. if (slot_seq != static_cast(seq)) { - return false; + return IssueResult::kNoCredit; } void* staging_addr = send_pool_->acquire(); if (staging_addr == nullptr) { - return false; + return IssueResult::kNoCredit; } const uint32_t expected_chunk_len = static_cast(std::min( chunk_size_, op_ctx.total_bytes_ - op_ctx.bytes_staged_)); - TORCH_CHECK(chunk_len <= expected_chunk_len, - "P2P send got invalid chunk_len in credit slot."); + PG_ASSERT(chunk_len <= expected_chunk_len, + "P2P send got invalid chunk_len in credit slot."); op_ctx.tasks_.emplace_back(op_ctx.bytes_staged_, chunk_len, staging_addr, recv_addr, seq, slot_epoch); @@ -769,38 +803,37 @@ bool P2PProxy::tryIssueSendTask(SendOpContext& op_ctx, SendPeerLane& lane) { slot.reset(); - const auto* tensor_ptr = - static_cast(op_ctx.tensor_.data_ptr()); + const auto* buffer_ptr = static_cast(op_ctx.buffer_); if (is_cpu_) { - std::memcpy(staging_addr, tensor_ptr + task.tensor_offset_, + std::memcpy(staging_addr, buffer_ptr + task.buffer_offset_, task.chunk_len_); task.state_ = SendTaskState::kWriteRemote; } else { cudaError_t copy_error = cudaMemcpyAsync( - staging_addr, tensor_ptr + task.tensor_offset_, task.chunk_len_, + staging_addr, buffer_ptr + task.buffer_offset_, task.chunk_len_, cudaMemcpyDeviceToDevice, op_ctx.cuda_stream_); - TORCH_CHECK(!copy_error, "P2P send cudaMemcpyAsync failed: ", - cudaGetErrorString(copy_error)); + PG_ASSERT(!copy_error, "P2P send cudaMemcpyAsync failed: ", + cudaGetErrorString(copy_error)); const cudaEvent_t pooled_copy_ready_event = lane.copy_ready_events_[static_cast(seq % kP2PControlRingSize)]; - TORCH_CHECK(pooled_copy_ready_event != nullptr, - "P2P send pooled copy-ready event is not initialized."); + PG_ASSERT(pooled_copy_ready_event != nullptr, + "P2P send pooled copy-ready event is not initialized."); task.copy_ready_event_ = pooled_copy_ready_event; copy_error = cudaEventRecord(task.copy_ready_event_, op_ctx.cuda_stream_); if (copy_error != cudaSuccess) { task.copy_ready_event_ = nullptr; - TORCH_CHECK(false, "P2P send cudaEventRecord failed: ", - cudaGetErrorString(copy_error)); + PG_ASSERT(false, "P2P send cudaEventRecord failed: ", + cudaGetErrorString(copy_error)); } } ++lane.credit_consume_seq_; op_ctx.bytes_staged_ += task.chunk_len_; task.last_update_time_ = std::chrono::steady_clock::now(); - return true; + return IssueResult::kIssued; } // Drive a single sender chunk through its state machine. @@ -849,8 +882,8 @@ bool P2PProxy::stepSendCopyIn(SendTransferTask& task) { return true; } - TORCH_CHECK(false, "P2P send cudaEventQuery failed: ", - cudaGetErrorString(query_error)); + PG_ASSERT(false, "P2P send cudaEventQuery failed: ", + cudaGetErrorString(query_error)); return false; } @@ -959,7 +992,7 @@ bool P2PProxy::isRecvOpCompleted(const RecvOpContext& op_ctx) const { // 1. Drain the shared send_queue_ into the peer's pending_send_ops_. // 2. Promote the first pending op to active_send_op_. // 3. While we have credits (CreditSlots) and staging buffers, -// pull more tensor slices into SendPool (TryIssueSendTask). +// pull more buffer slices into SendPool (TryIssueSendTask). // 4. Advance every active chunk through its state machine // (Copy-In -> Write Remote -> Ack write). // 5. Erase fully-acknowledged chunks. When the op is empty and all @@ -998,8 +1031,7 @@ bool P2PProxy::stepSend() { SendOpContext op_ctx = std::move(lane.pending_send_ops_.front()); lane.pending_send_ops_.pop_front(); if (op_ctx.total_bytes_ == 0) { - op_ctx.status_->store(OpStatus::kSuccess, - std::memory_order_release); + op_ctx.completion_->set_value(); active_send_tasks_.fetch_sub(1, std::memory_order_release); did_work = true; continue; @@ -1017,40 +1049,46 @@ bool P2PProxy::stepSend() { auto& op_ctx = lane.active_send_op_.value(); // Pull as many credits as we have free chunks - while (tryIssueSendTask(op_ctx, lane)) { + IssueResult issue; + do { + issue = tryIssueSendTask(op_ctx, lane); + did_work |= issue == IssueResult::kIssued; + } while (issue == IssueResult::kIssued); + + if (issue == IssueResult::kTimeout) { + handleFailedSendOp(op_ctx); + lane.active_send_op_.reset(); did_work = true; + continue; } // Advance every chunk through the sender state machine. - // send op may fail due to credit-wait timeout in tryIssueSendTask. - auto op_status = op_ctx.status_->load(std::memory_order_acquire); - bool op_failed = op_status == OpStatus::kFailed; - if (!op_failed) { - for (auto it = op_ctx.tasks_.begin(); it != op_ctx.tasks_.end();) { - if (stepSendTask(op_ctx, *it)) { - did_work = true; - } - if (it->state_ == SendTaskState::kFinished) { - it = op_ctx.tasks_.erase(it); - did_work = true; - } else if (it->state_ == SendTaskState::kFailed) { - op_failed = true; - break; - } else { - ++it; - } + bool op_failed = false; + for (auto it = op_ctx.tasks_.begin(); it != op_ctx.tasks_.end();) { + if (stepSendTask(op_ctx, *it)) { + did_work = true; + } + if (it->state_ == SendTaskState::kFinished) { + it = op_ctx.tasks_.erase(it); + did_work = true; + } else if (it->state_ == SendTaskState::kFailed) { + op_failed = true; + break; + } else { + ++it; } } if (op_failed) { - reportBrokenPeer(peer_rank); + handleFailedSendOp(op_ctx); + lane.active_send_op_.reset(); + did_work = true; continue; } if (isSendOpCompleted(op_ctx)) { - op_ctx.status_->store(OpStatus::kSuccess, - std::memory_order_release); + op_ctx.completion_->set_value(); lane.active_send_op_.reset(); active_send_tasks_.fetch_sub(1, std::memory_order_release); did_work = true; @@ -1107,8 +1145,8 @@ bool P2PProxy::pollRecvAckSlot(RecvOpContext& op_ctx, RecvPeerLane& lane, } void* src_ptr = head_task.local_addr_; - auto* tensor_ptr = static_cast(op_ctx.tensor_.data_ptr()); - void* dst_ptr = tensor_ptr + head_task.tensor_offset_; + auto* buffer_ptr = static_cast(op_ctx.buffer_); + void* dst_ptr = buffer_ptr + head_task.buffer_offset_; if (is_cpu_) { std::memcpy(dst_ptr, src_ptr, head_task.chunk_len_); @@ -1119,20 +1157,20 @@ bool P2PProxy::pollRecvAckSlot(RecvOpContext& op_ctx, RecvPeerLane& lane, cudaError_t copy_error = cudaMemcpyAsync(dst_ptr, src_ptr, head_task.chunk_len_, cudaMemcpyDeviceToDevice, op_ctx.cuda_stream_); - TORCH_CHECK(!copy_error, "P2P recv cudaMemcpyAsync failed: ", - cudaGetErrorString(copy_error)); + PG_ASSERT(!copy_error, "P2P recv cudaMemcpyAsync failed: ", + cudaGetErrorString(copy_error)); const cudaEvent_t pooled_copy_ready_event = lane.copy_ready_events_[static_cast(head_task.sequence_ % kP2PControlRingSize)]; - TORCH_CHECK(pooled_copy_ready_event != nullptr, - "P2P recv pooled copy-ready event is not initialized."); + PG_ASSERT(pooled_copy_ready_event != nullptr, + "P2P recv pooled copy-ready event is not initialized."); head_task.copy_ready_event_ = pooled_copy_ready_event; copy_error = cudaEventRecord(head_task.copy_ready_event_, op_ctx.cuda_stream_); if (copy_error != cudaSuccess) { head_task.copy_ready_event_ = nullptr; - TORCH_CHECK(false, "P2P recv cudaEventRecord failed: ", - cudaGetErrorString(copy_error)); + PG_ASSERT(false, "P2P recv cudaEventRecord failed: ", + cudaGetErrorString(copy_error)); } head_task.state_ = RecvTaskState::kCopyOut; } @@ -1151,7 +1189,7 @@ bool P2PProxy::pollRecvAckSlot(RecvOpContext& op_ctx, RecvPeerLane& lane, // 3. While we have free RecvPool chunks, issue CreditSlots to the sender // (TryIssueRecvTask). // 4. Poll the local AckLane in order. When an AckSlot -// arrives, initiate Copy-Out from RecvPool to the user tensor. +// arrives, initiate Copy-Out from RecvPool to the user buffer. // 5. Advance every active chunk through its state machine // (IssueCredit -> Copy-Out -> erase). // 6. When all chunks are copied out, mark the op complete. @@ -1171,10 +1209,10 @@ bool P2PProxy::stepRecv() { { std::lock_guard lock(recv_queue_mutex_); while (!recv_queue_.empty()) { - RecvOp op = std::move(recv_queue_.front()); + RecvOpContext op_ctx = std::move(recv_queue_.front()); recv_queue_.pop(); - recv_peer_lanes_[op.peer_rank_].pending_recv_ops_.push_back( - std::move(op)); + recv_peer_lanes_[op_ctx.peer_rank_].pending_recv_ops_.push_back( + std::move(op_ctx)); did_work = true; } } @@ -1186,9 +1224,8 @@ bool P2PProxy::stepRecv() { lane.pending_recv_ops_.empty()) { continue; } - RecvOp recv_op = std::move(lane.pending_recv_ops_.front()); + RecvOpContext op_ctx = std::move(lane.pending_recv_ops_.front()); lane.pending_recv_ops_.pop_front(); - RecvOpContext op_ctx(std::move(recv_op)); lane.active_recv_op_ = std::move(op_ctx); did_work = true; } @@ -1202,7 +1239,7 @@ bool P2PProxy::stepRecv() { auto& op_ctx = lane.active_recv_op_.value(); - // Offer as many RecvPool chunks as we have free buffers and tensor + // Offer as many RecvPool chunks as we have free buffers and user-buffer // bytes remaining. while (tryIssueRecvTask(op_ctx, lane)) { did_work = true; @@ -1236,24 +1273,14 @@ bool P2PProxy::stepRecv() { } if (op_failed) { - reportBrokenPeer(peer_rank); + handleFailedRecvOp(op_ctx); + lane.active_recv_op_.reset(); did_work = true; continue; } if (isRecvOpCompleted(op_ctx)) { - if (!op_ctx.original_tensor_.is_contiguous()) { - (void)op_ctx.original_tensor_.copy_(op_ctx.tensor_); - if (!is_cpu_) { - const cudaError_t sync_error = cudaDeviceSynchronize(); - TORCH_CHECK(sync_error == cudaSuccess, - "P2P recv final copy cudaDeviceSynchronize " - "failed: ", - cudaGetErrorString(sync_error)); - } - } - op_ctx.status_->store(OpStatus::kSuccess, - std::memory_order_release); + op_ctx.completion_->set_value(); lane.active_recv_op_.reset(); active_recv_tasks_.fetch_sub(1, std::memory_order_release); did_work = true; @@ -1379,50 +1406,47 @@ void P2PDeviceWorker::initPools(TransferEngine* engine, chunk_size_ = getEnv_size_t("MOONCAKE_P2P_CHUNK_SIZE", kDefaultChunkSize); // chunk_len is uint32_t fields in control slots - TORCH_CHECK( + PG_ASSERT( chunk_size_ > 0 && chunk_size_ <= std::numeric_limits::max(), "Invalid MOONCAKE_P2P_CHUNK_SIZE: must be > 0 and <= 4GB"); - TORCH_CHECK( - pool_bytes_ > 0 && pool_bytes_ % chunk_size_ == 0, - "Invalid pool size and chunk size (must hold 'pool_bytes_ > 0 && " - "pool_bytes_ % chunk_size_ == 0')"); + PG_ASSERT(pool_bytes_ > 0 && pool_bytes_ % chunk_size_ == 0, + "Invalid pool size and chunk size (must hold 'pool_bytes_ > 0 && " + "pool_bytes_ % chunk_size_ == 0')"); num_chunks_ = static_cast(pool_bytes_ / chunk_size_); - TORCH_CHECK(num_chunks_ > 0, "P2PDeviceWorker: num_chunks_ must be > 0"); + PG_ASSERT(num_chunks_ > 0, "P2PDeviceWorker: num_chunks_ must be > 0"); if (is_cpu_) { send_pool_base_ = std::malloc(pool_bytes_); - TORCH_CHECK(send_pool_base_ != nullptr, - "Failed to allocate CPU P2P send pool"); + PG_ASSERT(send_pool_base_ != nullptr, + "Failed to allocate CPU P2P send pool"); int rc = engine->registerLocalMemory(send_pool_base_, pool_bytes_, location); - TORCH_CHECK(rc == 0, "Failed to register CPU P2P send pool"); + PG_ASSERT(rc == 0, "Failed to register CPU P2P send pool"); recv_pool_base_ = std::malloc(pool_bytes_); - TORCH_CHECK(recv_pool_base_ != nullptr, - "Failed to allocate CPU P2P recv pool"); + PG_ASSERT(recv_pool_base_ != nullptr, + "Failed to allocate CPU P2P recv pool"); rc = engine->registerLocalMemory(recv_pool_base_, pool_bytes_, location); - TORCH_CHECK(rc == 0, "Failed to register CPU P2P recv pool"); + PG_ASSERT(rc == 0, "Failed to register CPU P2P recv pool"); } else { setCudaDeviceIfNeeded(is_cpu_, cuda_device_index_, "P2PDeviceWorker initPools cudaSetDevice failed"); cudaError_t err = cudaMalloc(&send_pool_base_, pool_bytes_); - TORCH_CHECK( - err == cudaSuccess, - "Failed to allocate CUDA P2P send pool: ", cudaGetErrorString(err)); + PG_ASSERT(err == cudaSuccess, "Failed to allocate CUDA P2P send pool: ", + cudaGetErrorString(err)); int rc = engine->registerLocalMemory(send_pool_base_, pool_bytes_, location); - TORCH_CHECK(rc == 0, "Failed to register CUDA P2P send pool"); + PG_ASSERT(rc == 0, "Failed to register CUDA P2P send pool"); err = cudaMalloc(&recv_pool_base_, pool_bytes_); - TORCH_CHECK( - err == cudaSuccess, - "Failed to allocate CUDA P2P recv pool: ", cudaGetErrorString(err)); + PG_ASSERT(err == cudaSuccess, "Failed to allocate CUDA P2P recv pool: ", + cudaGetErrorString(err)); rc = engine->registerLocalMemory(recv_pool_base_, pool_bytes_, location); - TORCH_CHECK(rc == 0, "Failed to register CUDA P2P recv pool"); + PG_ASSERT(rc == 0, "Failed to register CUDA P2P recv pool"); } send_pool_.init(send_pool_base_, chunk_size_, num_chunks_); @@ -1468,11 +1492,11 @@ void workerThreadLoop(bool is_cpu, int cuda_device_index, // A thread-local cache for proxies to avoid locking too frequently. // This is efficient because proxy registration/removal is rare after - // backend initialization. + // communicator initialization. std::vector> local_proxies; uint64_t local_version = std::numeric_limits::max(); - // Some proxies might be removed (caused by backend shutdown) before + // Some proxies might be removed (caused by communicator shutdown) before // all its transfers completed. We must keep them explicitly to avoid its // resources released, and step their transfers till completion. std::vector> zombie_proxies; diff --git a/mooncake-pg/src/pg_py.cpp b/mooncake-pg/src/pg_py.cpp deleted file mode 100644 index e39daffeec..0000000000 --- a/mooncake-pg/src/pg_py.cpp +++ /dev/null @@ -1,139 +0,0 @@ -#include -#include -#include -#include -#include -#include - -namespace py = pybind11; - -namespace mooncake { - -c10::intrusive_ptr createMooncakeBackend( - c10d::DistributedBackendOptions distBackendOpts, - c10::intrusive_ptr - backendOptions) { - return c10::make_intrusive(std::move(distBackendOpts), - std::move(backendOptions)); -} - -c10::intrusive_ptr createMooncakeCpuBackend( - c10d::DistributedBackendOptions distBackendOpts, - c10::intrusive_ptr - backendOptions) { - return c10::make_intrusive( - std::move(distBackendOpts), std::move(backendOptions), true); -} - -__attribute__((constructor)) static void MooncakeBackendConstructor() { - py::object module = py::module::import("torch.distributed"); - py::object register_backend = - module.attr("Backend").attr("register_backend"); - py::dict kwargsCpu; - kwargsCpu["devices"] = py::make_tuple("cpu"); - register_backend("mooncake-cpu", py::cpp_function(createMooncakeCpuBackend), - /* extended_api */ true, **kwargsCpu); -#ifndef MOONCAKE_EP_USE_MUSA - py::dict kwargsCuda; - kwargsCuda["devices"] = py::make_tuple("cuda"); - register_backend("mooncake", py::cpp_function(createMooncakeBackend), - /* extended_api */ true, **kwargsCuda); -#else - py::dict kwargsMusa; - kwargsMusa["devices"] = py::make_tuple("musa"); - register_backend("mooncake", py::cpp_function(createMooncakeBackend), - /* extended_api */ true, **kwargsMusa); -#endif -} - -std::string getPreferredHca(c10::intrusive_ptr backend, - std::string location) { - auto mooncakeBackend = - c10::static_intrusive_pointer_cast(backend); - return mooncakeBackend->getPreferredHca(location); -} - -at::Tensor getActiveRanks(c10::intrusive_ptr backend) { - auto mooncakeBackend = - c10::static_intrusive_pointer_cast(backend); - return mooncakeBackend->getActiveRanksTensor(); -} - -int getNumSyncedRanks(c10::intrusive_ptr backend) { - auto mooncakeBackend = - c10::static_intrusive_pointer_cast(backend); - return mooncakeBackend->getNumSyncedRanks(); -} - -void extendGroupSizeTo(c10::intrusive_ptr backend, - int size) { - auto mooncakeBackend = - c10::static_intrusive_pointer_cast(backend); - mooncakeBackend->extendGroupSizeTo(size); -} - -std::vector getPeerState(c10::intrusive_ptr backend, - const std::vector& ranks) { - auto mooncakeBackend = - c10::static_intrusive_pointer_cast(backend); - return mooncakeBackend->getPeerState(ranks); -} - -void recoverRanks(c10::intrusive_ptr backend, - const std::vector& ranks) { - auto mooncakeBackend = - c10::static_intrusive_pointer_cast(backend); - mooncakeBackend->recoverRanks(ranks); -} - -void joinGroup(c10::intrusive_ptr backend) { - auto mooncakeBackend = - c10::static_intrusive_pointer_cast(backend); - mooncakeBackend->joinGroup(); -} - -/// Python-facing wrapper that extracts the raw TransferEngine* from a -/// mooncake.engine.TransferEngine Python object and passes it to -/// MooncakeBackend::setExternalEngine(). The caller must ensure the -/// TransferEnginePy object outlives all MooncakeBackend instances. -void setTransferEnginePy(pybind11::object engine_obj) { - if (engine_obj.is_none()) { - MooncakeBackend::setExternalEngine(nullptr); - return; - } - auto get_engine_ptr = engine_obj.attr("get_engine_ptr"); - uintptr_t ptr = get_engine_ptr().cast(); - auto* engine = reinterpret_cast(ptr); - MooncakeBackend::setExternalEngine(engine); -} - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("createMooncakeBackend", &createMooncakeBackend); - m.def("createMooncakeCpuBackend", &createMooncakeCpuBackend); - m.def("set_host_ip", &MooncakeBackend::setHostIp); - m.def("set_device_filter", &MooncakeBackend::setDeviceFilter); - m.def("set_transfer_engine", &setTransferEnginePy, py::arg("engine"), - "Set an external TransferEngine to be used by MooncakeBackend. " - "Must be called before init_process_group(). The engine must already " - "be initialized. Pass None to reset to default behavior. " - "The caller must ensure the TransferEngine object outlives all " - "MooncakeBackend instances."); - m.def("get_preferred_hca", &getPreferredHca); - m.def("get_active_ranks", &getActiveRanks); - m.def("get_num_synced_ranks", &getNumSyncedRanks); - m.def("extend_group_size_to", &extendGroupSizeTo); - m.def("get_peer_state", &getPeerState); - m.def("recover_ranks", &recoverRanks); - m.def("join_group", &joinGroup); - - py::class_>( - m, "MooncakeBackendOptions") - .def(py::init(), py::arg("active_ranks")) - .def(py::init(), py::arg("active_ranks"), - py::arg("is_extension")) - .def(py::init(), py::arg("active_ranks"), - py::arg("is_extension"), py::arg("max_world_size")); -} - -} // namespace mooncake diff --git a/mooncake-pg/tests/pg_test_utils.py b/mooncake-pg/tests/pg_test_utils.py index 87950c80bc..176859a5e7 100644 --- a/mooncake-pg/tests/pg_test_utils.py +++ b/mooncake-pg/tests/pg_test_utils.py @@ -116,38 +116,17 @@ def require_test_device(rank: int, device_type: str) -> torch.device: def mooncake_backend_options( - world_size: int, - device_type: str, + max_group_size: int, *, - active_value: int = 0, is_extension: bool = False, - max_world_size: int | None = None, + auto_deactivate_on_failure: bool = True, + auto_sync_on_failure: bool = True, ) -> pg.MooncakeBackendOptions: - device = torch.device(device_type) - tensor_size = world_size if max_world_size is None else int(max_world_size) - active_ranks = torch.full( - (tensor_size,), - int(active_value), - dtype=torch.int32, - device=device, - ) - if max_world_size is None: - if is_extension: - return pg.MooncakeBackendOptions(active_ranks, True) - return pg.MooncakeBackendOptions(active_ranks) - return pg.MooncakeBackendOptions(active_ranks, bool(is_extension), tensor_size) - - -def mooncake_cpu_options(world_size: int) -> pg.MooncakeBackendOptions: - return mooncake_backend_options(world_size, "cpu", active_value=0) - - -def mooncake_extension_cpu_options(world_size: int) -> pg.MooncakeBackendOptions: - return mooncake_backend_options( - world_size, - "cpu", - active_value=1, - is_extension=True, + return pg.MooncakeBackendOptions( + int(max_group_size), + bool(is_extension), + bool(auto_deactivate_on_failure), + bool(auto_sync_on_failure), ) @@ -159,9 +138,10 @@ def init_mooncake_group( device_type: str, device_filters: Sequence[str] | None = None, use_pg_options: bool = True, + max_group_size: int | None = None, is_extension: bool = False, - active_value: int | None = None, - max_world_size: int | None = None, + auto_deactivate_on_failure: bool = True, + auto_sync_on_failure: bool = True, ) -> torch.device: device = require_test_device(rank, device_type) configure_mooncake_device_filter(device_filters) @@ -171,37 +151,16 @@ def init_mooncake_group( "world_size": world_size, } if use_pg_options: - resolved_active_value = ( - 1 if is_extension else 0 if active_value is None else active_value - ) kwargs["pg_options"] = mooncake_backend_options( - world_size, - device_type, - active_value=resolved_active_value, + max_group_size if max_group_size is not None else world_size, is_extension=is_extension, - max_world_size=max_world_size, + auto_deactivate_on_failure=auto_deactivate_on_failure, + auto_sync_on_failure=auto_sync_on_failure, ) dist.init_process_group(**kwargs) return device -def init_mooncake_cpu_group( - rank: int, - world_size: int, - *, - device_filters: Sequence[str] | None = None, - use_pg_options: bool = True, -) -> None: - init_mooncake_group( - rank, - world_size, - backend_name="mooncake-cpu", - device_type="cpu", - device_filters=device_filters, - use_pg_options=use_pg_options, - ) - - def get_mooncake_backend(group=None, device_type: str = "cpu"): if group is None: group = dist.group.WORLD @@ -235,9 +194,10 @@ def init_group( world_size: int | None = None, device_filters: Sequence[str] | None = None, use_pg_options: bool = True, + max_group_size: int | None = None, is_extension: bool = False, - active_value: int | None = None, - max_world_size: int | None = None, + auto_deactivate_on_failure: bool = True, + auto_sync_on_failure: bool = True, ) -> torch.device: self._device = init_mooncake_group( self.proc_rank if rank is None else rank, @@ -248,9 +208,10 @@ def init_group( if device_filters is None else device_filters, use_pg_options=use_pg_options, + max_group_size=max_group_size, is_extension=is_extension, - active_value=active_value, - max_world_size=max_world_size, + auto_deactivate_on_failure=auto_deactivate_on_failure, + auto_sync_on_failure=auto_sync_on_failure, ) return self._device @@ -383,14 +344,75 @@ def wait_until( raise TimeoutError(f"timed out waiting for {description}") -def wait_for_spawn_context(ctx, timeout_s: float) -> None: - """Wait for spawn context with timeout; force kill if hung.""" +def _describe_signal(signum: int) -> str: + """Return a human-readable description for a signal number.""" + names = { + signal.SIGSEGV: "SIGSEGV (segmentation fault)", + signal.SIGABRT: "SIGABRT (abort)", + signal.SIGBUS: "SIGBUS (bus error)", + signal.SIGFPE: "SIGFPE (floating-point exception)", + signal.SIGILL: "SIGILL (illegal instruction)", + signal.SIGTERM: "SIGTERM (terminated)", + signal.SIGKILL: "SIGKILL (killed)", + signal.SIGQUIT: "SIGQUIT (quit)", + } + return names.get(signum, f"signal {signum}") + + +def _capture_signal_deaths(ctx, result_map) -> None: + """Inspect process exit codes and record any signal-killed processes. + + A negative exit code N means the process was killed by signal -N. + We only record a signal death when the worker did NOT already report + a result — otherwise we would overwrite a more specific error (e.g., an + AssertionError from a survivor) that happened before the crash. + """ + for rank_idx, process in enumerate(ctx.processes): + exitcode = process.exitcode + if exitcode is not None and exitcode < 0: + # Only fill in missing results; don't overwrite existing ones + if rank_idx in result_map: + continue + signum = -exitcode + description = _describe_signal(signum) + record_rank_error( + result_map, + rank_idx, + RuntimeError( + f"Rank {rank_idx} was killed by {description}. " + f"This usually indicates a native crash (segfault, abort, etc.). " + f"Check core dumps (ulimit -c) or run under a debugger." + ), + ) + + +def wait_for_spawn_context( + ctx, + timeout_s: float, + result_map=None, + process_exit_events: dict[int, object] | None = None, +) -> None: + """Wait for spawn context with timeout; force kill if hung. + + When *result_map* is provided, process exit codes are inspected and + any signal-killed ranks are recorded in the map before returning. + """ deadline = time.monotonic() + timeout_s + def publish_process_exits() -> None: + if process_exit_events is None: + return + for rank, event in process_exit_events.items(): + if not ctx.processes[rank].is_alive(): + event.set() + # Phase 1: Normal wait while time.monotonic() < deadline: + publish_process_exits() if not any(p.is_alive() for p in ctx.processes): # All processes exited (success or failure) + if result_map is not None: + _capture_signal_deaths(ctx, result_map) return time.sleep(0.1) @@ -424,6 +446,9 @@ def wait_for_spawn_context(ctx, timeout_s: float) -> None: break time.sleep(0.1) + if result_map is not None: + _capture_signal_deaths(ctx, result_map) + raise AssertionError(f"Spawn timed out after {timeout_s} seconds") @@ -472,7 +497,7 @@ def spawn_and_collect( nprocs=actual_nprocs, join=False, ) - wait_for_spawn_context(ctx, timeout_s) + wait_for_spawn_context(ctx, timeout_s, result_map=result_map) return collect_rank_results(result_map, actual_nprocs) @@ -503,12 +528,39 @@ def spawn_and_collect( ) def assert_all_ok(self, rows: list[dict]) -> None: + failures: list[str] = [] for row in rows: - if not row.get("ok", False): - self.fail( - f"rank {row.get('rank', '?')} failed with " - f"{row.get('error_type', 'UnknownError')}: {row.get('error', '')}" - ) + if row.get("ok", False): + continue + rank = row.get("rank", "?") + error_type = row.get("error_type", "UnknownError") + error_msg = row.get("error", "") + failures.append( + f" rank {rank}: {error_type}" + + (f" -- {error_msg}" if error_msg else "") + ) + + if not failures: + return + + succeeded_ranks = sorted( + r.get("rank", "?") for r in rows if r.get("ok", False) + ) + failed_ranks = sorted( + r.get("rank", "?") for r in rows if not r.get("ok", False) + ) + + report = [ + f"\n{'='*60}", + f"RANK ERROR REPORT ({len(failures)} failure(s))", + f"{'='*60}", + f"Succeeded ranks: {succeeded_ranks or '(none)'}", + f"Failed ranks: {failed_ranks or '(none)'}", + f"{'-'*60}", + "Failures:", + ] + failures + [f"{'='*60}"] + + self.fail("\n".join(report)) class BackendMultiProcessTestCase(MultiProcessTestCase): @@ -549,6 +601,7 @@ def spawn_backend_and_collect( nprocs: int | None = None, timeout_s: float | None = None, world_size: int | None = None, + process_exit_events: dict[int, object] | None = None, ) -> list[dict]: if self.backend_name is None or self.device_type is None: raise RuntimeError( @@ -594,7 +647,12 @@ def spawn_backend_and_collect( nprocs=actual_nprocs, join=False, ) - wait_for_spawn_context(ctx, resolved_timeout) + wait_for_spawn_context( + ctx, + resolved_timeout, + result_map=result_map, + process_exit_events=process_exit_events, + ) return collect_rank_results(result_map, actual_nprocs) diff --git a/mooncake-pg/tests/test_pg_collectives.py b/mooncake-pg/tests/test_pg_collectives.py index 52fbee7ad0..f0804e912c 100644 --- a/mooncake-pg/tests/test_pg_collectives.py +++ b/mooncake-pg/tests/test_pg_collectives.py @@ -8,7 +8,6 @@ MooncakePGCUDABackendTestCase, MooncakePGMUSABackendTestCase, MooncakePGWorkerContext, - wait_until, ) diff --git a/mooncake-pg/tests/test_pg_elastic.py b/mooncake-pg/tests/test_pg_elastic.py index 23f71008b8..114bdefc94 100644 --- a/mooncake-pg/tests/test_pg_elastic.py +++ b/mooncake-pg/tests/test_pg_elastic.py @@ -1,5 +1,4 @@ import os -import time import unittest import torch @@ -21,33 +20,9 @@ BROKEN_RANK = 1 -def _dynamic_world_size_worker( - ctx: MooncakePGWorkerContext, -) -> None: - """Worker for testing that dist.get_world_size() reflects dynamic size after extend.""" - initial_world_size = ctx.world_size - ctx.init_group() - backend = ctx.get_backend() - - initial_ws = dist.get_world_size() - assert initial_ws == initial_world_size, ( - f"rank {ctx.rank}: initial world_size={initial_ws}, expected {initial_world_size}" - ) - - pg.extend_group_size_to(backend, initial_world_size + 1) - - new_ws = dist.get_world_size() - assert new_ws == initial_world_size + 1, ( - f"rank {ctx.rank}: after extend world_size={new_ws}, expected {initial_world_size + 1}" - ) - - ctx.record_result({"initial_ws": initial_ws, "new_ws": new_ws}) - - def _extension_worker( ctx: MooncakePGWorkerContext, extend_event: mp.Event, - init_done_event: mp.Event, ) -> None: """Worker for testing extension mode - new ranks join existing group.""" initial_world_size = ctx.world_size - 1 @@ -59,12 +34,12 @@ def _extension_worker( # Original ranks device = ctx.init_group( world_size=initial_world_size, - max_world_size=ctx.world_size, + max_group_size=ctx.world_size, ) backend = ctx.get_backend() - # group_size should equal initial_world_size immediately after init - # (max_world_size only pre-allocates capacity, does not change visible size) + # max_group_size only reserves capacity; the visible world size starts + # at the number of active ranks. actual_ws = dist.get_world_size() assert actual_ws == initial_world_size, ( f"rank {ctx.proc_rank}: initial world_size={actual_ws}, " @@ -83,20 +58,22 @@ def _extension_worker( # Two-phase extension protocol: # 1) joiner publishes metadata + establishes transport readiness # 2) healthy ranks recover/activate it via recover_ranks() - # Note: get_peer_state() is collective among *healthy ranks*. wait_until( lambda: all(pg.get_peer_state(backend, join_ranks)), - timeout_s=30.0, - poll_interval_s=0.05, + timeout_s=10.0, + poll_interval_s=0.01, description=f"rank {ctx.proc_rank} waiting for joiner ready", ) - pg.recover_ranks(backend, join_ranks) + resp = pg.recover_ranks(backend, join_ranks) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank {ctx.proc_rank}: recover_ranks should apply, got {resp.status}" - # After recover_ranks, world_size should now reflect the expanded group + # The Coordinator has committed the joiner as active, so the visible + # rank-space extent now covers the full group. actual_ws_after = dist.get_world_size() assert actual_ws_after == ctx.world_size, ( f"rank {ctx.proc_rank}: world_size after recover={actual_ws_after}, " - f"expected max_world_size={ctx.world_size}" + f"expected max_group_size={ctx.world_size}" ) # Final collective @@ -116,39 +93,39 @@ def _extension_worker( device = ctx.init_group( rank=extension_rank, world_size=ctx.world_size, + max_group_size=ctx.world_size, is_extension=True, - max_world_size=ctx.world_size, ) backend = ctx.get_backend() - # group_size for extension rank equals world_size passed at init. - # Note: this is world_size (= max_world_size for joiners), not 1, - # because the joiner's activeSize is initialized to world_size. - # The local-only behavior is ensured by activeRanks masking, not - # by a smaller activeSize. + # PyTorch validates the declared full size before join_group() can run. + # Local-only behavior is provided by the active-rank mask, not by + # reporting a smaller world size from the isolated joiner. actual_ws = dist.get_world_size() assert actual_ws == ctx.world_size, ( f"extension rank: initial world_size={actual_ws}, " f"expected {ctx.world_size}" ) - # In extension mode, joiner starts in local-only collectives. - local_tensor = torch.tensor([extension_rank + 1], dtype=torch.int32, device=device) + # Before join_group, the joining backend executes collectives with an + # effective {self} mask and does not involve the existing ranks. + local_tensor = torch.tensor( + [extension_rank + 1], dtype=torch.int32, device=device + ) dist.all_reduce(local_tensor, op=dist.ReduceOp.SUM) - if int(local_tensor.cpu().item()) != extension_rank + 1: + local_value = int(local_tensor.cpu().item()) + if local_value != extension_rank + 1: raise AssertionError( - f"extension rank expected local-only sum {extension_rank + 1}, got {int(local_tensor.cpu().item())}" + f"extension rank expected local-only sum " + f"{extension_rank + 1}, got {local_value}" ) - # join_group publishes metadata and then blocks until recover_ranks() - # publishes the extension state. pg.join_group(backend) - # After joinGroup, world_size should reflect the full group actual_ws_after = dist.get_world_size() assert actual_ws_after == ctx.world_size, ( - f"extension rank: world_size after joinGroup={actual_ws_after}, " + f"extension rank: world_size after join_group={actual_ws_after}, " f"expected {ctx.world_size}" ) @@ -162,6 +139,119 @@ def _extension_worker( }) +def _extension_p2p_worker( + ctx: MooncakePGWorkerContext, + extend_event: mp.Event, + direction: str, +) -> None: + """Worker for testing P2P after max_group_size/recover_ranks scale-up.""" + initial_world_size = ctx.world_size - 1 + extension_rank = ctx.world_size - 1 + primary_peer = initial_world_size - 1 + join_ranks = [extension_rank] + + if initial_world_size < 1: + raise AssertionError("elastic P2P test expects at least one primary rank") + + if ctx.proc_rank < initial_world_size: + device = ctx.init_group( + world_size=initial_world_size, + max_group_size=ctx.world_size, + ) + backend = ctx.get_backend() + + if ctx.proc_rank == 0: + extend_event.set() + + wait_until( + lambda: all(pg.get_peer_state(backend, join_ranks)), + timeout_s=30.0, + poll_interval_s=0.05, + description=f"rank {ctx.proc_rank} waiting for joiner ready", + ) + resp = pg.recover_ranks(backend, join_ranks) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank {ctx.proc_rank}: recover_ranks should apply, got {resp.status}" + else: + if not extend_event.wait(timeout=30.0): + raise TimeoutError("timed out waiting for extend_event") + + device = ctx.init_group( + rank=extension_rank, + world_size=ctx.world_size, + max_group_size=ctx.world_size, + is_extension=True, + ) + backend = ctx.get_backend() + pg.join_group(backend) + + # Prove elastic collectives are functional before isolating P2P. + collective = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) + dist.all_reduce(collective, op=dist.ReduceOp.SUM) + expected_sum = ctx.world_size * (ctx.world_size + 1) // 2 + if int(collective.cpu().item()) != expected_sum: + raise AssertionError( + f"rank {ctx.proc_rank}: post-recovery all_reduce expected " + f"{expected_sum}, got {int(collective.cpu().item())}" + ) + + dist.barrier() + + if direction == "joiner_to_primary": + src_rank = extension_rank + dst_rank = primary_peer + elif direction == "primary_to_joiner": + src_rank = primary_peer + dst_rank = extension_rank + else: + raise AssertionError(f"unknown P2P direction: {direction}") + + numel = 1024 + if ctx.proc_rank == src_rank: + send_tensor = torch.full( + (numel,), src_rank, dtype=torch.int32, device=device + ) + works = dist.batch_isend_irecv( + [dist.P2POp(op=dist.isend, tensor=send_tensor, peer=dst_rank)] + ) + for work in works: + work.wait() + ctx.synchronize() + value = "sent" + elif ctx.proc_rank == dst_rank: + recv_tensor = torch.empty((numel,), dtype=torch.int32, device=device) + works = dist.batch_isend_irecv( + [dist.P2POp(op=dist.irecv, tensor=recv_tensor, peer=src_rank)] + ) + for work in works: + work.wait() + ctx.synchronize() + expected = torch.full_like(recv_tensor, src_rank) + if not torch.equal(recv_tensor.cpu(), expected.cpu()): + raise AssertionError( + f"rank {ctx.proc_rank}: received unexpected P2P payload " + f"from rank {src_rank}" + ) + value = "received" + else: + value = "idle" + + dist.barrier() + ctx.record_result( + { + "direction": direction, + "role": ( + "src" + if ctx.proc_rank == src_rank + else "dst" + if ctx.proc_rank == dst_rank + else "idle" + ), + "value": value, + } + ) + + def _extension_worker_with_subgroups( ctx: MooncakePGWorkerContext, extend_event: mp.Event, @@ -169,19 +259,9 @@ def _extension_worker_with_subgroups( """Multi-subgroup elastic extension test using split-ranks pattern. Layout (world_size=4, primary=[0,1], joiners=[2,3]): - group_a: primary ranks=[0], joiner ranks=[0,2], max_world_size=2 - group_b: primary ranks=[1], joiner ranks=[1,3], max_world_size=2 - group_c: primary ranks=[0,1], joiner ranks=[0,1,2,3], max_world_size=4 - - Primary ranks must initialize WORLD with world_size=initial_world_size and - create subgroups using only their current membership; joiners wait for the - extend signal, then init WORLD with the full world_size and create subgroups - using the full eventual membership. PyTorch's new_group uses a monotonic - call counter for the store prefix, so primary and joiner side land on the - same prefix as long as the call order matches. backendIndex_ is also - process-local and increments only when the rank is an actual member of the - new group, so it stays aligned across primaries and joiners that all call - new_group in the same order. + group_a: primary ranks=[0], extended ranks=[0,2], max_group_size=2 + group_b: primary ranks=[1], extended ranks=[1,3], max_group_size=2 + group_c: primary ranks=[0,1], extended ranks=[0,1,2,3], max_group_size=4 """ configure_mooncake_device_filter(ctx.device_filters) device = require_test_device(ctx.proc_rank, ctx.device_type) @@ -191,40 +271,34 @@ def _extension_worker_with_subgroups( join_ranks = [2, 3] is_joiner = ctx.proc_rank >= initial_world_size - a_active = torch.tensor([1, 0], dtype=torch.int32, device=device) - b_active = torch.tensor([1, 0], dtype=torch.int32, device=device) - c_active = torch.tensor([1, 1, 0, 0], dtype=torch.int32, device=device) - if not is_joiner: - # Primary ranks: init WORLD with world_size=2, max_world_size=4 - world_active = torch.tensor([1, 1, 0, 0], dtype=torch.int32, device=device) + # Primary ranks: init WORLD with world_size=2, max_group_size=4 dist_kwargs = { "backend": ctx.backend_name, "rank": ctx.proc_rank, "world_size": initial_world_size, - "pg_options": pg.MooncakeBackendOptions(world_active, False, ctx.world_size), + "pg_options": pg.MooncakeBackendOptions(ctx.world_size), } if ctx.device_type == "cuda": dist_kwargs["device_id"] = device dist.init_process_group(**dist_kwargs) world_backend = get_mooncake_backend(device_type=ctx.device_type) - # Subgroups with split-ranks pattern. All ranks in WORLD must call - # new_group in the same order even for groups they are not members of. + # Subgroups with split-ranks pattern group_a = dist.new_group( ranks=[0], backend=ctx.backend_name, - pg_options=pg.MooncakeBackendOptions(a_active, False, 2), + pg_options=pg.MooncakeBackendOptions(2), ) group_b = dist.new_group( ranks=[1], backend=ctx.backend_name, - pg_options=pg.MooncakeBackendOptions(b_active, False, 2), + pg_options=pg.MooncakeBackendOptions(2), ) group_c = dist.new_group( ranks=[0, 1], backend=ctx.backend_name, - pg_options=pg.MooncakeBackendOptions(c_active, False, 4), + pg_options=pg.MooncakeBackendOptions(4), ) a_backend = get_mooncake_backend(group_a, device_type=ctx.device_type) if ctx.proc_rank == 0 else None b_backend = get_mooncake_backend(group_b, device_type=ctx.device_type) if ctx.proc_rank == 1 else None @@ -245,43 +319,51 @@ def _extension_worker_with_subgroups( if ctx.proc_rank == 0: extend_event.set() - # WORLD: wait for joiners then recover + # WORLD: wait for joiners then activate wait_until( lambda: all(pg.get_peer_state(world_backend, join_ranks)), - timeout_s=60.0, + timeout_s=30.0, poll_interval_s=0.05, - description=f"rank {ctx.proc_rank} waiting for WORLD joiners", + description=f"rank {ctx.proc_rank} waiting for joiners", ) - pg.recover_ranks(world_backend, join_ranks) + resp = pg.recover_ranks(world_backend, join_ranks) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank {ctx.proc_rank}: recover_ranks(world) should apply, got {resp.status}" # group_a: rank 0 waits for joiner (local rank 1 = global rank 2) if ctx.proc_rank == 0: wait_until( lambda: pg.get_peer_state(a_backend, [1])[0], - timeout_s=60.0, + timeout_s=30.0, poll_interval_s=0.05, description="rank 0 waiting for group_a joiner", ) - pg.recover_ranks(a_backend, [1]) + resp = pg.recover_ranks(a_backend, [1]) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank 0: recover_ranks(group_a) should apply, got {resp.status}" # group_b: rank 1 waits for joiner (local rank 1 = global rank 3) if ctx.proc_rank == 1: wait_until( lambda: pg.get_peer_state(b_backend, [1])[0], - timeout_s=60.0, + timeout_s=30.0, poll_interval_s=0.05, description="rank 1 waiting for group_b joiner", ) - pg.recover_ranks(b_backend, [1]) + resp = pg.recover_ranks(b_backend, [1]) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank 1: recover_ranks(group_b) should apply, got {resp.status}" # group_c: both primaries wait for both joiners (local ranks 2,3) wait_until( lambda: all(pg.get_peer_state(c_backend, [2, 3])), - timeout_s=60.0, + timeout_s=30.0, poll_interval_s=0.05, description=f"rank {ctx.proc_rank} waiting for group_c joiners", ) - pg.recover_ranks(c_backend, [2, 3]) + resp = pg.recover_ranks(c_backend, [2, 3]) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank {ctx.proc_rank}: recover_ranks(group_c) should apply, got {resp.status}" # Post-activation: WORLD all 4 ranks → 1+2+3+4=10 t = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) @@ -315,12 +397,11 @@ def _extension_worker_with_subgroups( if not extend_event.wait(timeout=60.0): raise TimeoutError("timed out waiting for extend_event") - world_active = torch.tensor([1, 1, 0, 0], dtype=torch.int32, device=device) dist_kwargs = { "backend": ctx.backend_name, "rank": ctx.proc_rank, "world_size": ctx.world_size, - "pg_options": pg.MooncakeBackendOptions(world_active, True, ctx.world_size), + "pg_options": pg.MooncakeBackendOptions(ctx.world_size, True), } if ctx.device_type == "cuda": dist_kwargs["device_id"] = device @@ -331,28 +412,31 @@ def _extension_worker_with_subgroups( group_a = dist.new_group( ranks=[0, 2], backend=ctx.backend_name, - pg_options=pg.MooncakeBackendOptions(a_active, True, 2), + pg_options=pg.MooncakeBackendOptions(2, True), ) group_b = dist.new_group( ranks=[1, 3], backend=ctx.backend_name, - pg_options=pg.MooncakeBackendOptions(b_active, True, 2), + pg_options=pg.MooncakeBackendOptions(2, True), ) group_c = dist.new_group( ranks=[0, 1, 2, 3], backend=ctx.backend_name, - pg_options=pg.MooncakeBackendOptions(c_active, True, 4), + pg_options=pg.MooncakeBackendOptions(4, True), ) a_backend = get_mooncake_backend(group_a, device_type=ctx.device_type) if ctx.proc_rank == 2 else None b_backend = get_mooncake_backend(group_b, device_type=ctx.device_type) if ctx.proc_rank == 3 else None c_backend = get_mooncake_backend(group_c, device_type=ctx.device_type) - # Joiners are local-only until join_group is called + # Before any join_group call, WORLD collectives are local-only on each + # joining rank and do not involve the primary ranks. t = torch.tensor([ctx.proc_rank + 1], dtype=torch.int32, device=device) dist.all_reduce(t, op=dist.ReduceOp.SUM) - if int(t.cpu().item()) != ctx.proc_rank + 1: + local_value = int(t.cpu().item()) + if local_value != ctx.proc_rank + 1: raise AssertionError( - f"WORLD local-only: expected {ctx.proc_rank + 1}, got {int(t.cpu().item())}" + f"WORLD local-only: expected {ctx.proc_rank + 1}, " + f"got {local_value}" ) # Join groups in same order primaries created them @@ -409,6 +493,7 @@ def _run_allgather_reduce_scatter( input_t = torch.tensor([rank + 1], dtype=torch.int32, device=device) output_t = torch.zeros(active_world_size, dtype=torch.int32, device=device) dist.all_gather_into_tensor(output_t, input_t) + for j in range(active_world_size): expected = j + 1 got = int(output_t[j].item()) @@ -441,8 +526,8 @@ def _allgather_reduce_scatter_extension_worker( ) -> None: """Test _allgather_base and _reduce_scatter_base across elastic extension. - Layout: world_size=4, initial=3, extension_rank=3, max_world_size=4. - Pre-activation: 3 active ranks, max_world_size=4 → exercises the overflow path. + Layout: world_size=4, initial=3, extension_rank=3, max_group_size=4. + Pre-activation: 3 active ranks, max_group_size=4 → exercises the overflow path. Post-activation: 4 active ranks → exercises correctness after extension. """ configure_mooncake_device_filter(ctx.device_filters) @@ -455,19 +540,18 @@ def _allgather_reduce_scatter_extension_worker( is_joiner = ctx.proc_rank == extension_rank if not is_joiner: - active = torch.tensor([1, 1, 1, 0], dtype=torch.int32, device=device) dist_kwargs = { "backend": ctx.backend_name, "rank": ctx.proc_rank, "world_size": initial_world_size, - "pg_options": pg.MooncakeBackendOptions(active, False, ctx.world_size), + "pg_options": pg.MooncakeBackendOptions(ctx.world_size), } if ctx.device_type == "cuda": dist_kwargs["device_id"] = device dist.init_process_group(**dist_kwargs) backend = get_mooncake_backend(device_type=ctx.device_type) - # Pre-activation: 3 active ranks, max_world_size=4. + # Pre-activation: 3 active ranks, max_group_size=4. # This is the overflow path: buggy code would iterate 4 times into a # buffer sized for 3. _run_allgather_reduce_scatter(device, initial_world_size, ctx.proc_rank) @@ -481,7 +565,9 @@ def _allgather_reduce_scatter_extension_worker( poll_interval_s=0.05, description=f"rank {ctx.proc_rank} waiting for joiner", ) - pg.recover_ranks(backend, join_ranks) + resp = pg.recover_ranks(backend, join_ranks) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank {ctx.proc_rank}: recover_ranks should apply, got {resp.status}" # Post-activation: all 4 ranks active. _run_allgather_reduce_scatter(device, ctx.world_size, ctx.proc_rank) @@ -491,12 +577,11 @@ def _allgather_reduce_scatter_extension_worker( if not extend_event.wait(timeout=30.0): raise TimeoutError("timed out waiting for extend_event") - active = torch.tensor([1, 1, 1, 0], dtype=torch.int32, device=device) dist_kwargs = { "backend": ctx.backend_name, "rank": extension_rank, "world_size": ctx.world_size, - "pg_options": pg.MooncakeBackendOptions(active, True, ctx.world_size), + "pg_options": pg.MooncakeBackendOptions(ctx.world_size, True), } if ctx.device_type == "cuda": dist_kwargs["device_id"] = device @@ -514,14 +599,13 @@ def _allgather_reduce_scatter_extension_worker( def _allgather_reduce_scatter_recovery_worker( ctx: MooncakePGWorkerContext, broken_exited: mp.Event, - replacement_ready: mp.Event, start_recovery: mp.Event, ) -> None: """Test _allgather_base and _reduce_scatter_base across rank recovery. Layout: world_size=4, broken_rank=3, replacement takes rank 3. Pre-failure: 4 active ranks. - Post-failure (3 survivors): 3 active ranks, max_world_size=4 → overflow path. + Post-failure (3 survivors): 3 active ranks, max_group_size=4 → overflow path. Post-recovery: 4 active ranks again. """ broken_rank = ctx.world_size - 1 @@ -539,8 +623,9 @@ def _allgather_reduce_scatter_recovery_worker( broken_exited.set() os._exit(0) - # Survivors: 3 active ranks, max_world_size=4 → overflow path. + # Survivors: 3 active ranks, max_group_size=4 → overflow path. broken_exited.wait() + _run_allgather_reduce_scatter(device, ctx.world_size - 1, logical_rank) if logical_rank == 0: @@ -552,8 +637,9 @@ def _allgather_reduce_scatter_recovery_worker( poll_interval_s=2.0, description=f"rank {logical_rank} waiting for replacement", ) - replacement_ready.wait() - pg.recover_ranks(backend, [broken_rank]) + resp = pg.recover_ranks(backend, [broken_rank]) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" # Post-recovery: all 4 ranks active again. _run_allgather_reduce_scatter(device, ctx.world_size, logical_rank) @@ -563,7 +649,6 @@ def _allgather_reduce_scatter_recovery_worker( start_recovery.wait() device = ctx.init_group(rank=logical_rank, is_extension=True) backend = ctx.get_backend() - replacement_ready.set() pg.join_group(backend) # Post-recovery: all 4 ranks active. @@ -603,22 +688,40 @@ def _fault_detection_worker( def _replacement_recovery_worker( ctx: MooncakePGWorkerContext, broken_exited: mp.Event, - replacement_ready: mp.Event, start_recovery: mp.Event, + graceful_group_destroy: bool = False, ) -> None: - """Worker for testing replacement recovery.""" + """Worker for testing replacement after failure or graceful teardown.""" logical_rank = ctx.rank if ctx.proc_rank < ctx.world_size else BROKEN_RANK if ctx.proc_rank < ctx.world_size: # Original rank (0, 1, 2, or 3) device = ctx.init_group(rank=logical_rank) - # First collective with all ranks + # Round 1: all healthy tensor = torch.tensor([logical_rank], dtype=torch.int32, device=device) dist.all_reduce(tensor, op=dist.ReduceOp.SUM) if logical_rank == BROKEN_RANK: - # Broken rank exits + if graceful_group_destroy: + backend = ctx.get_backend() + # CUDA Work::wait() orders the current stream but does not + # block the host until the collective finishes. Complete + # Round 1 before deactivate_ranks changes the active-rank + # view used by that collective. + ctx.synchronize() + resp = pg.deactivate_ranks(backend, [logical_rank]) + assert resp.status == pg.ProposalStatus.Applied, \ + "graceful self-deactivation should apply, " \ + f"got {resp.status}: {resp.reject_reason}" + + dist.destroy_process_group() + ctx.record_result({"role": "gracefully_removed"}) + # Do not set broken_exited here. The parent sets the event only + # after it observes that this process has exited. + return + + # Broken rank exits without running any teardown. ctx.record_result({"role": "broken"}) broken_exited.set() os._exit(0) @@ -627,7 +730,7 @@ def _replacement_recovery_worker( broken_exited.wait() backend = ctx.get_backend() - # Run collective without broken rank + # Round 2: run collective without the departed rank. tensor = torch.tensor([logical_rank], dtype=torch.int32, device=device) dist.all_reduce(tensor, op=dist.ReduceOp.SUM) @@ -636,7 +739,6 @@ def _replacement_recovery_worker( start_recovery.set() # Wait for replacement to be connected (metadata published) - # Use longer poll interval to avoid overloading the connection poller wait_until( lambda: pg.get_peer_state(backend, [BROKEN_RANK])[0], timeout_s=30.0, @@ -644,11 +746,10 @@ def _replacement_recovery_worker( description=f"rank {logical_rank} waiting for replacement to connect", ) - # Wait for replacement to be ready for join_group - replacement_ready.wait() - # All ranks call recover_ranks to include replacement - pg.recover_ranks(backend, [BROKEN_RANK]) + resp = pg.recover_ranks(backend, [BROKEN_RANK]) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" # Final collective with all 4 ranks tensor = torch.tensor([logical_rank], dtype=torch.int32, device=device) @@ -660,13 +761,10 @@ def _replacement_recovery_worker( # Wait for signal to start start_recovery.wait() - # Replacement initializes with is_extension (local-only mode) + # Replacement initializes with is_extension device = ctx.init_group(rank=logical_rank, is_extension=True) backend = ctx.get_backend() - # Signal that we're initialized and ready for join_group - replacement_ready.set() - # join_group completes the connection and switches to global mode pg.join_group(backend) @@ -677,20 +775,147 @@ def _replacement_recovery_worker( ctx.record_result({"role": "replacement"}) -class _ElasticMixin: - world_size = 4 - spawn_timeout_s = 30.0 +def _manual_deactivate_recovery_worker( + ctx: MooncakePGWorkerContext, + broken_exited: mp.Event, + start_recovery: mp.Event, +) -> None: + """Manual deactivation and recovery with auto_deactivate disabled: + Round 1 (all healthy): failedRanks = all 0s, activeRanks = all 1s. + Round 2 (rank died, not yet deactivated): activeRanks unchanged. + Round 3: survivors sync, manually deactivate, and run as a reduced group. + Round 4: a replacement joins, survivors recover it, and the full group + runs again. + """ + logical_rank = ctx.rank if ctx.proc_rank < ctx.world_size else BROKEN_RANK - def test_dynamic_world_size(self) -> None: - """Test that dist.get_world_size() returns updated value after extend_group_size_to.""" - rows = self.spawn_backend_and_collect( - _dynamic_world_size_worker, + if ctx.proc_rank < ctx.world_size: + device = ctx.init_group( + rank=logical_rank, + auto_deactivate_on_failure=False, + auto_sync_on_failure=False, + ) + backend = ctx.get_backend() + + # Round 1: all healthy + expected_all = ctx.world_size * (ctx.world_size + 1) // 2 + tensor = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device) + work = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True) + work.wait() + assert int(tensor.cpu().item()) == expected_all + assert pg.get_local_success(work), \ + f"rank {ctx.rank}: round 1 should succeed locally" + + failed_ranks_hint = pg.get_failed_ranks_hint(work) + assert failed_ranks_hint.tolist() == [0] * ctx.world_size + + active_ranks = pg.get_active_ranks(backend) + assert active_ranks.cpu().tolist() == [1] * ctx.world_size + + if logical_rank == BROKEN_RANK: + ctx.record_result({"role": "broken"}) + broken_exited.set() + os._exit(0) + + broken_exited.wait() + + # Round 2: rank died, auto_deactivate=False ==> activeRanks unchanged. + # local_success=False because the dead rank is still in the group. + expected_reduced = expected_all - (BROKEN_RANK + 1) + tensor = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device) + work = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True) + work.wait() + assert int(tensor.cpu().item()) == expected_reduced + assert not pg.get_local_success(work), \ + f"rank {ctx.rank}: round 2 should detect broken rank" + + failed_ranks_hint = pg.get_failed_ranks_hint(work) + expected_failed_ranks_hint = [0] * ctx.world_size + expected_failed_ranks_hint[BROKEN_RANK] = 1 + assert failed_ranks_hint.tolist() == expected_failed_ranks_hint + + active_ranks = pg.get_active_ranks(backend) + assert active_ranks.cpu().tolist() == [1] * ctx.world_size + + # sync_after_failure waits for the shared reconciliation decision even + # when membership is managed manually. Its response applies the + # authoritative group view before returning, and the local failed-link + # observation must already be visible through get_peer_state(). + sync_resp = pg.sync_after_failure(backend) + assert sync_resp.status in ( + pg.SyncAfterFailureStatus.Reconciled, + pg.SyncAfterFailureStatus.NoPending, + ), f"rank {logical_rank}: sync_after_failure failed: {sync_resp.reject_reason}" + assert not pg.get_peer_state(backend, [BROKEN_RANK])[0], \ + f"rank {logical_rank}: failed rank remained locally activatable after sync" + + # Survivors deactivate the dead rank before issuing new collectives. + resp = pg.deactivate_ranks(backend, [BROKEN_RANK]) + assert resp.status in ( + pg.ProposalStatus.Applied, + pg.ProposalStatus.AppliedWithDroppedRanks, + ), f"rank {ctx.rank}: deactivate_ranks should apply, got {resp.status}" + # If ranks were dropped, they should only be the BROKEN_RANK + assert resp.status != pg.ProposalStatus.AppliedWithDroppedRanks or set( + resp.dropped_ranks + ) == { + BROKEN_RANK + }, f"rank {ctx.rank}: unexpected dropped ranks {resp.dropped_ranks}" + + # Round 3: after deactivate, collective with reduced group succeeds. + expected_reduced = expected_all - (BROKEN_RANK + 1) + tensor = torch.tensor([ctx.rank + 1], dtype=torch.int32, device=device) + work = dist.all_reduce(tensor, op=dist.ReduceOp.SUM, async_op=True) + work.wait() + assert int(tensor.cpu().item()) == expected_reduced + assert pg.get_local_success(work), \ + f"rank {ctx.rank}: round 3 should succeed after manual deactivate" + + if logical_rank == 0: + start_recovery.set() + + wait_until( + lambda: pg.get_peer_state(backend, [BROKEN_RANK])[0], timeout_s=30.0, + poll_interval_s=0.05, + description=f"rank {logical_rank} waiting for replacement", ) - self.assert_all_ok(rows) - for row in rows: - self.assertEqual(row["new_ws"], self.world_size + 1) + resp = pg.recover_ranks(backend, [BROKEN_RANK]) + assert resp.status == pg.ProposalStatus.Applied, \ + f"rank {logical_rank}: recover_ranks should apply, got {resp.status}" + + tensor = torch.tensor( + [logical_rank + 1], dtype=torch.int32, device=device + ) + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + + assert int(tensor.cpu().item()) == expected_all + + ctx.record_result({"role": "survivor"}) + + else: + if not start_recovery.wait(timeout=120.0): + raise TimeoutError("timed out waiting to start manual recovery") + device = ctx.init_group( + rank=logical_rank, + is_extension=True, + auto_deactivate_on_failure=False, + auto_sync_on_failure=False, + ) + backend = ctx.get_backend() + pg.join_group(backend) + + tensor = torch.tensor( + [logical_rank + 1], dtype=torch.int32, device=device + ) + dist.all_reduce(tensor, op=dist.ReduceOp.SUM) + ctx.record_result({"role": "replacement"}) + + +class _ElasticMixin: + world_size = 4 + spawn_timeout_s = 30.0 def test_failed_rank(self) -> None: """Test that survivors can continue collective after a rank fails.""" @@ -700,9 +925,11 @@ def test_failed_rank(self) -> None: rows = self.spawn_backend_and_collect( _fault_detection_worker, broken_exited, - timeout_s=30.0, + timeout_s=120.0, ) + self.assert_all_ok(rows) + # All survivors should complete survivor_rows = [r for r in rows if r.get("role") == "survivor"] self.assertEqual(len(survivor_rows), self.world_size - 1) @@ -715,18 +942,18 @@ def test_recovery(self) -> None: """Test that replacement can join and restore full collective.""" spawn_ctx = mp.get_context("spawn") broken_exited = spawn_ctx.Event() - replacement_ready = spawn_ctx.Event() start_recovery = spawn_ctx.Event() rows = self.spawn_backend_and_collect( _replacement_recovery_worker, broken_exited, - replacement_ready, start_recovery, nprocs=self.world_size + 1, - timeout_s=30.0, + timeout_s=120.0, ) + self.assert_all_ok(rows) + # Verify all participants completed survivor_rows = [r for r in rows if r.get("role") == "survivor"] replacement_rows = [r for r in rows if r.get("role") == "replacement"] @@ -736,26 +963,56 @@ def test_recovery(self) -> None: self.assertEqual(len(replacement_rows), 1) self.assertGreaterEqual(len(broken_rows), 1) + def test_recovery_after_graceful_group_destroy(self) -> None: + """A rank can destroy its group normally and later be replaced. """ + spawn_ctx = mp.get_context("spawn") + removed_process_exited = spawn_ctx.Event() + start_recovery = spawn_ctx.Event() + + rows = self.spawn_backend_and_collect( + _replacement_recovery_worker, + removed_process_exited, + start_recovery, + True, # graceful_group_destroy + nprocs=self.world_size + 1, + timeout_s=120.0, + # The graceful child cannot signal after its own destructors have + # run. Have the parent publish broken_exited only after observing + # that the old process, including its Agent teardown, has exited. + process_exit_events={BROKEN_RANK: removed_process_exited}, + ) + + self.assert_all_ok(rows) + + survivor_rows = [r for r in rows if r.get("role") == "survivor"] + replacement_rows = [r for r in rows if r.get("role") == "replacement"] + removed_rows = [ + r for r in rows if r.get("role") == "gracefully_removed" + ] + + self.assertEqual(len(survivor_rows), self.world_size - 1) + self.assertEqual(len(replacement_rows), 1) + self.assertEqual(len(removed_rows), 1) + def test_extension(self) -> None: """Test extension mode allows new ranks to join existing group.""" spawn_ctx = mp.get_context("spawn") extend_event = spawn_ctx.Event() - init_done_event = spawn_ctx.Event() # Spawn world_size processes: (world_size - 1) original + 1 extension rows = self.spawn_backend_and_collect( _extension_worker, extend_event, - init_done_event, nprocs=self.world_size, - timeout_s=30.0, + timeout_s=120.0, ) + self.assert_all_ok(rows) + # Verify all participants completed original_rows = [r for r in rows if r.get("role") == "original"] extension_rows = [r for r in rows if r.get("role") == "extension"] - # Original: world_size - 1 ranks, Extension: 1 rank self.assertEqual(len(original_rows), self.world_size - 1) self.assertEqual(len(extension_rows), 1) @@ -764,6 +1021,40 @@ def test_extension(self) -> None: for row in original_rows: self.assertEqual(row.get("baseline"), expected_baseline) + def test_extension_p2p_joiner_to_primary(self) -> None: + """Test P2P from an extension rank to an original rank after recovery.""" + spawn_ctx = mp.get_context("spawn") + extend_event = spawn_ctx.Event() + + rows = self.spawn_backend_and_collect( + _extension_p2p_worker, + extend_event, + "joiner_to_primary", + nprocs=self.world_size, + timeout_s=120.0, + ) + + self.assert_all_ok(rows) + self.assertEqual(len([r for r in rows if r.get("role") == "src"]), 1) + self.assertEqual(len([r for r in rows if r.get("role") == "dst"]), 1) + + def test_extension_p2p_primary_to_joiner(self) -> None: + """Test P2P from an original rank to an extension rank after recovery.""" + spawn_ctx = mp.get_context("spawn") + extend_event = spawn_ctx.Event() + + rows = self.spawn_backend_and_collect( + _extension_p2p_worker, + extend_event, + "primary_to_joiner", + nprocs=self.world_size, + timeout_s=120.0, + ) + + self.assert_all_ok(rows) + self.assertEqual(len([r for r in rows if r.get("role") == "src"]), 1) + self.assertEqual(len([r for r in rows if r.get("role") == "dst"]), 1) + def test_extension_with_subgroups(self) -> None: """Test extension with multiple disjoint subgroups using split-ranks pattern.""" spawn_ctx = mp.get_context("spawn") @@ -773,16 +1064,18 @@ def test_extension_with_subgroups(self) -> None: _extension_worker_with_subgroups, extend_event, nprocs=self.world_size, - timeout_s=60.0, + timeout_s=120.0, ) + self.assert_all_ok(rows) + result_rows = [r for r in rows if r.get("role") == "extension_subgroups"] self.assertEqual(len(result_rows), self.world_size) def test_allgather_reduce_scatter_extension(self) -> None: """Test _allgather_base/_reduce_scatter_base correctness across elastic extension. - Exercises the overflow path: pre-activation uses max_world_size=4 with only + Exercises the overflow path: pre-activation uses max_group_size=4 with only 3 active ranks, so the buggy code would access slot 3 of a size-3 buffer. """ spawn_ctx = mp.get_context("spawn") @@ -792,9 +1085,11 @@ def test_allgather_reduce_scatter_extension(self) -> None: _allgather_reduce_scatter_extension_worker, extend_event, nprocs=self.world_size, - timeout_s=60.0, + timeout_s=120.0, ) + self.assert_all_ok(rows) + primary_rows = [r for r in rows if r.get("role") == "primary"] joiner_rows = [r for r in rows if r.get("role") == "joiner"] self.assertEqual(len(primary_rows), self.world_size - 1) @@ -804,22 +1099,22 @@ def test_allgather_reduce_scatter_recovery(self) -> None: """Test _allgather_base/_reduce_scatter_base correctness across rank recovery. Exercises the overflow path: post-failure survivors run with 3 active ranks - and max_world_size=4, so the buggy code would access slot 3 of a size-3 buffer. + and max_group_size=4, so the buggy code would access slot 3 of a size-3 buffer. """ spawn_ctx = mp.get_context("spawn") broken_exited = spawn_ctx.Event() - replacement_ready = spawn_ctx.Event() start_recovery = spawn_ctx.Event() rows = self.spawn_backend_and_collect( _allgather_reduce_scatter_recovery_worker, broken_exited, - replacement_ready, start_recovery, nprocs=self.world_size + 1, - timeout_s=60.0, + timeout_s=120.0, ) + self.assert_all_ok(rows) + survivor_rows = [r for r in rows if r.get("role") == "survivor"] replacement_rows = [r for r in rows if r.get("role") == "replacement"] broken_rows = [r for r in rows if r.get("role") == "broken"] @@ -827,6 +1122,39 @@ def test_allgather_reduce_scatter_recovery(self) -> None: self.assertEqual(len(replacement_rows), 1) self.assertGreaterEqual(len(broken_rows), 1) + def test_manual_evict_recovery(self) -> None: + """Test manual evict and recovery with auto_deactivate disabled. + + Round 1: all healthy, failedRanks = all 0s. + Round 2: rank dies; sync reconciles but leaves membership unchanged. + Round 3: manually deactivate and run with the reduced group. + Round 4: activate a replacement and run with the full group. + """ + spawn_ctx = mp.get_context("spawn") + broken_exited = spawn_ctx.Event() + start_recovery = spawn_ctx.Event() + + rows = self.spawn_backend_and_collect( + _manual_deactivate_recovery_worker, + broken_exited, + start_recovery, + nprocs=self.world_size + 1, + timeout_s=120.0, + ) + + self.assert_all_ok(rows) + + # All survivors should complete + survivor_rows = [r for r in rows if r.get("role") == "survivor"] + self.assertEqual(len(survivor_rows), self.world_size - 1) + + replacement_rows = [r for r in rows if r.get("role") == "replacement"] + self.assertEqual(len(replacement_rows), 1) + + # Broken rank should have exited (may not have result) + broken_rows = [r for r in rows if r.get("role") == "broken"] + self.assertGreaterEqual(len(broken_rows), 1) + class TestMooncakePGElasticCPU( _ElasticMixin, MooncakePGCPUBackendTestCase diff --git a/mooncake-pg/tests/test_pg_init_functional.py b/mooncake-pg/tests/test_pg_init_functional.py index 66cf1acfc4..bd096ce589 100644 --- a/mooncake-pg/tests/test_pg_init_functional.py +++ b/mooncake-pg/tests/test_pg_init_functional.py @@ -117,6 +117,13 @@ def test_null_init(self) -> None: for row in rows: self.assertEqual(row["sum"], expected_sum) + def test_world_size_one(self) -> None: + """Test a process group with world size 1 initializes and works.""" + rows = self.spawn_backend_and_collect(_basic_init_worker, world_size=1) + self.assert_all_ok(rows) + for row in rows: + self.assertEqual(row["sum"], 1) + def test_subgroup_create_destroy(self) -> None: """Test subgroup creation and destruction.""" if self.world_size < 4: diff --git a/mooncake-pg/tests/test_pg_p2p.py b/mooncake-pg/tests/test_pg_p2p.py index 41be71fc0e..4fac899dd2 100644 --- a/mooncake-pg/tests/test_pg_p2p.py +++ b/mooncake-pg/tests/test_pg_p2p.py @@ -1,7 +1,10 @@ +import os import unittest import torch import torch.distributed as dist +import torch.multiprocessing as mp +from mooncake import pg from pg_test_utils import ( MooncakePGCPUBackendTestCase, @@ -123,6 +126,74 @@ def _multiple_senders_worker( ctx.record_result({"value": value}) +BROKEN_RANK = 1 + + +def _p2p_fault_detection_worker( + ctx: MooncakePGWorkerContext, + broken_exited, +) -> None: + """P2P fault tolerance: 2 rounds of all-to-all. + Round 1: all healthy -> failedRanks = 0s. + Round 2: rank 1 exited -> peer's failedRanks[1] = 1 + """ + + def p2p_all_to_all(ctx, device): + """All-to-all P2P send/recv. Returns work handles.""" + ops = [] + for p in range(ctx.world_size): + if p == ctx.rank: + continue + s = torch.tensor([ctx.rank], dtype=torch.int64, device=device) + r = torch.empty_like(s) + ops.append(dist.P2POp(op=dist.isend, tensor=s, peer=p)) + ops.append(dist.P2POp(op=dist.irecv, tensor=r, peer=p)) + return dist.batch_isend_irecv(ops) + + device = ctx.init_group() + + # Round 1: all healthy + works = p2p_all_to_all(ctx, device) + for w in works: + w.wait() + ctx.synchronize() + for w in works: + assert pg.get_local_success(w), \ + f"rank {ctx.rank} round 1: all P2P ops should succeed locally" + failed_ranks_hint = pg.get_failed_ranks_hint(w) + assert ( + failed_ranks_hint.cpu().tolist() == [0] * ctx.world_size + ), f"rank {ctx.rank} round 1: failed_ranks_hint={failed_ranks_hint.cpu().tolist()}" + + if ctx.rank == BROKEN_RANK: + ctx.record_result({"role": "broken"}) + broken_exited.set() + os._exit(0) + + broken_exited.wait() + + # Round 2: BROKEN_RANK is dead + works = p2p_all_to_all(ctx, device) + for w in works: + w.wait() + + normal_failed_ranks_hint = [0] * ctx.world_size + broken_peer_failed_ranks_hint = [0] * ctx.world_size + broken_peer_failed_ranks_hint[BROKEN_RANK] = 1 + peers = [p for p in range(ctx.world_size) if p != ctx.rank] + for w, peer in zip(works, [p for p in peers for _ in range(2)]): + failed_ranks_hint = pg.get_failed_ranks_hint(w) + expected = ( + broken_peer_failed_ranks_hint if peer == BROKEN_RANK else normal_failed_ranks_hint + ) + assert failed_ranks_hint.cpu().tolist() == expected + if peer == BROKEN_RANK: + assert not pg.get_local_success(w), \ + f"rank {ctx.rank} round 2: P2P with broken peer should fail locally" + + ctx.record_result({"role": "survivor"}) + + class _P2PMixin: world_size = 4 @@ -150,6 +221,26 @@ def test_ordering_between_two_ranks(self) -> None: rank1 = next(row for row in rows if row["rank"] == 1) self.assertEqual(rank1["value"], list(range(4))) + def test_p2p_fault_detection(self) -> None: + """Test P2P fault detection: failedRanks and activeRanks on P2P failure.""" + spawn_ctx = mp.get_context("spawn") + broken_exited = spawn_ctx.Event() + + rows = self.spawn_backend_and_collect( + _p2p_fault_detection_worker, + broken_exited, + world_size=3, + nprocs=3, + timeout_s=75.0, + ) + + self.assert_all_ok(rows) + + survivor_rows = [r for r in rows if r.get("role") == "survivor"] + broken_rows = [r for r in rows if r.get("role") == "broken"] + self.assertEqual(len(survivor_rows), 2) + self.assertGreaterEqual(len(broken_rows), 1) + def test_multiple_senders_to_same_receiver(self) -> None: if self.world_size < 3: self.skipTest("multiple-sender P2P coverage requires at least 3 ranks") diff --git a/mooncake-pg/BuildPgExt.cmake b/mooncake-pg/torch/BuildPgExt.cmake similarity index 69% rename from mooncake-pg/BuildPgExt.cmake rename to mooncake-pg/torch/BuildPgExt.cmake index ead67c3e8b..3332cf36e4 100644 --- a/mooncake-pg/BuildPgExt.cmake +++ b/mooncake-pg/torch/BuildPgExt.cmake @@ -3,28 +3,26 @@ # Invoked at build time via cmake -P from the root CMakeLists.txt when # WITH_EP=ON. Variables are passed with -D from the custom target: # -# SOURCE_DIR - mooncake-pg source directory +# SOURCE_DIR - mooncake-pg/torch source directory # EP_CUDA_MAJOR - CUDA major version (integer) # EP_CUDA_MINOR - CUDA minor version (integer) # EP_TORCH_VERSIONS - pipe-separated (|) PyTorch versions to build for # (empty = use the currently-installed torch) -# TORCH_CUDA_ARCH_LIST - pipe-separated CUDA arch list forwarded to torch # STAGING_DIR - destination directory for the built .so files -# ENGINE_SO_PATH - absolute path to the built engine.cpython-XYZ.so +# PG_CORE_SO_PATH - absolute path to the built libmooncake_pg.so +# PG_DEVICE_SO_PATH - absolute path to libmooncake_pg_device.so # EP_USE_MUSA - set to "1" when building for MUSA (MTLink path) +# EP_USE_MACA - set to "1" when building for MACA (MTLink path) cmake_minimum_required(VERSION 3.16) # Include common build utilities. -include("${SOURCE_DIR}/../mooncake-common/SetupPyTorchEnv.cmake") +include("${SOURCE_DIR}/../../mooncake-common/SetupPyTorchEnv.cmake") # Restore pipe-separated strings back to CMake semicolon-separated lists. if(EP_TORCH_VERSIONS) string(REPLACE "|" ";" EP_TORCH_VERSIONS "${EP_TORCH_VERSIONS}") endif() -if(TORCH_CUDA_ARCH_LIST) - string(REPLACE "|" ";" TORCH_CUDA_ARCH_LIST "${TORCH_CUDA_ARCH_LIST}") -endif() # --------------------------------------------------------------------------- # 1. Set up the build environment. @@ -35,30 +33,33 @@ endif() # file descriptors". set(ENV{MAKEFLAGS} "") set(ENV{MFLAGS} "") -set(ENV{TORCH_CUDA_ARCH_LIST} "${TORCH_CUDA_ARCH_LIST}") +if(NOT PG_CORE_SO_PATH OR NOT EXISTS "${PG_CORE_SO_PATH}") + message(FATAL_ERROR + "[PG] PG_CORE_SO_PATH is missing or does not exist: ${PG_CORE_SO_PATH}") +endif() +if(NOT PG_DEVICE_SO_PATH OR NOT EXISTS "${PG_DEVICE_SO_PATH}") + message(FATAL_ERROR + "[PG] PG_DEVICE_SO_PATH is missing or does not exist: ${PG_DEVICE_SO_PATH}") +endif() +set(ENV{MOONCAKE_PG_CORE_SO_PATH} "${PG_CORE_SO_PATH}") if(EP_USE_MUSA) set(ENV{MOONCAKE_EP_USE_MUSA} "1") else() unset(ENV{MOONCAKE_EP_USE_MUSA}) endif() - -# --------------------------------------------------------------------------- -# 2. Ensure engine.so exists in mooncake-wheel/mooncake/ for setup.py linking. -# --------------------------------------------------------------------------- -# setup.py links against -l:engine.so in ../mooncake-wheel/mooncake/. -# During the make phase only the versioned engine.cpython-XYZ.so exists in -# the build tree; create a bare engine.so symlink so the linker can find it. -set(_wheel_mooncake_dir "${SOURCE_DIR}/../mooncake-wheel/mooncake") -set(_engine_symlink "${_wheel_mooncake_dir}/engine.so") -if(ENGINE_SO_PATH AND NOT EXISTS "${_engine_symlink}") - message(STATUS "[PG] Creating engine.so symlink -> ${ENGINE_SO_PATH}") - execute_process( - COMMAND ${CMAKE_COMMAND} -E create_symlink "${ENGINE_SO_PATH}" "${_engine_symlink}" - ) +if(EP_USE_MACA) + set(ENV{MOONCAKE_EP_USE_MACA} "1") + if(DEFINED ENV{MACA_PATH}) + set(ENV{MACA_HOME} "$ENV{MACA_PATH}") + elseif(DEFINED ENV{MACA_HOME}) + set(ENV{MACA_PATH} "$ENV{MACA_HOME}") + endif() +else() + unset(ENV{MOONCAKE_EP_USE_MACA}) endif() # --------------------------------------------------------------------------- -# 3. Build the PG Python extension. +# 2. Build the PG Python extension. # --------------------------------------------------------------------------- if("${EP_TORCH_VERSIONS}" STREQUAL "") message(STATUS "[PG] Building with currently-installed PyTorch") @@ -87,10 +88,12 @@ else() endif() # --------------------------------------------------------------------------- -# 4. Copy the built .so files to the staging directory. +# 3. Stage only fatbin-bearing device and extension .so files. The host-only +# core is packaged before auditwheel so its dependencies are repaired. # --------------------------------------------------------------------------- file(MAKE_DIRECTORY "${STAGING_DIR}") file(GLOB _so_files "${SOURCE_DIR}/mooncake/*.so") +list(APPEND _so_files "${PG_DEVICE_SO_PATH}") foreach(_so IN LISTS _so_files) get_filename_component(_fname "${_so}" NAME) message(STATUS "[PG] Staging ${_fname} -> ${STAGING_DIR}") diff --git a/mooncake-pg/torch/include/mooncake_backend.h b/mooncake-pg/torch/include/mooncake_backend.h new file mode 100644 index 0000000000..62be7bc692 --- /dev/null +++ b/mooncake-pg/torch/include/mooncake_backend.h @@ -0,0 +1,185 @@ +#ifndef MOONCAKE_BACKEND_H +#define MOONCAKE_BACKEND_H + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include + +namespace mooncake { + +class MooncakeBackend final : public ::c10d::ProcessGroup { + public: + struct MooncakeBackendOptions final : torch::CustomClassHolder { + explicit MooncakeBackendOptions(int maxGroupSize) + : maxGroupSize_{maxGroupSize > 0 ? maxGroupSize : -1} {} + + // isExtension=false maps to CreateOrAttach; true maps to + // AttachOrExtend. + MooncakeBackendOptions(int maxGroupSize, bool isExtension) + : isExtension_{isExtension}, + maxGroupSize_{maxGroupSize > 0 ? maxGroupSize : -1} {} + MooncakeBackendOptions(int maxGroupSize, bool isExtension, + bool autoDeactivateOnFailure, + bool autoSyncOnFailure) + : isExtension_{isExtension}, + maxGroupSize_{maxGroupSize > 0 ? maxGroupSize : -1}, + autoDeactivateOnFailure_{autoDeactivateOnFailure}, + autoSyncOnFailure_{autoSyncOnFailure} {} + + // If activeRanks is provided, only its storage is used -- the contents + // are populated by the Coordinator. + explicit MooncakeBackendOptions(at::Tensor activeRanks) + : activeRanks_{std::move(activeRanks)} {} + + // Main-compatible tensor overloads use the same isExtension mapping. + MooncakeBackendOptions(at::Tensor activeRanks, bool isExtension) + : activeRanks_{std::move(activeRanks)}, isExtension_{isExtension} {} + MooncakeBackendOptions(at::Tensor activeRanks, bool isExtension, + int maxGroupSize) + : activeRanks_{std::move(activeRanks)}, + isExtension_{isExtension}, + maxGroupSize_{maxGroupSize > 0 ? maxGroupSize : -1} {} + + ~MooncakeBackendOptions() override = default; + + at::Tensor activeRanks_; + bool isExtension_ = false; + int maxGroupSize_ = -1; + + // Automatically deactivate failed ranks on timeout / operation failure. + // + // When true (default), failed ranks are removed from the active set + // automatically. When false, failures are only reported through + // per-operation failedRanks hints, so the caller can decide how to + // handle the failure. + // + // Default: MOONCAKE_PG_AUTO_DEACTIVATE_ON_FAILURE (1) + bool autoDeactivateOnFailure_ = true; + + // Fence a failed collective on Coordinator reconciliation. + // + // When true (default), the worker reports a locally detected transfer + // failure and waits for the authoritative membership view to be + // applied before completing the task. Consequently, CPU work and + // CUDA stream execution cannot pass the failed collective before the + // view is updated. CUDA Work::wait() itself remains asynchronous and + // does not imply that the host can immediately observe the new view. + // + // Requires autoDeactivateOnFailure_ == true. + // + // Default: MOONCAKE_PG_AUTO_SYNC_ON_FAILURE (1) + bool autoSyncOnFailure_ = true; + }; + + /** + * @brief Construct a Mooncake process-group backend instance. + * + * `distBackendOpts` contains the PyTorch process-group information for this + * backend instance. `options` contains Mooncake-specific settings and may + * be null when callers omit `pg_options`. + * + * @param distBackendOpts Process-group information supplied by PyTorch. + * @param options Optional Mooncake-specific backend options. + * @param context Process-wide Mooncake PG core context. + * @param isCpu Whether to initialize the CPU backend variant. + */ + MooncakeBackend(c10d::DistributedBackendOptions distBackendOpts, + c10::intrusive_ptr options, + mooncakePgContext_t context, bool isCpu = false); + ~MooncakeBackend() override; + + const std::string getBackendName() const override; + + // In Normal mode, return the active rank-space extent (highest active + // InGroupRank plus one). During bootstrap, PyTorch still needs the + // group_size declared at construction to validate future ranks passed to + // new_group() before joinGroup() can run: + // https://github.com/pytorch/pytorch/blob/release/2.13/torch/distributed/distributed_c10d.py#L6012 + int getSize() const override; + + // Point-to-point send/recv for torch.distributed P2POp/batch_isend_irecv. + // Only single-tensor ops are supported. + c10::intrusive_ptr send(std::vector& tensors, + int dstRank, int tag) override; + c10::intrusive_ptr recv(std::vector& tensors, + int srcRank, int tag) override; + c10::intrusive_ptr broadcast( + std::vector& tensors, + const c10d::BroadcastOptions& opts) override; + c10::intrusive_ptr allreduce( + std::vector& tensors, + const c10d::AllreduceOptions& opts) override; + c10::intrusive_ptr allgather( + std::vector>& outputTensors, + std::vector& inputTensors, + const c10d::AllgatherOptions& opts) override; + c10::intrusive_ptr _allgather_base( + at::Tensor& outputBuffer, at::Tensor& inputBuffer, + const c10d::AllgatherOptions& opts) override; + c10::intrusive_ptr _reduce_scatter_base( + at::Tensor& outputBuffer, at::Tensor& inputBuffer, + const c10d::ReduceScatterOptions& opts) override; + c10::intrusive_ptr alltoall( + std::vector& outputTensors, + std::vector& inputTensors, + const c10d::AllToAllOptions& opts) override; + c10::intrusive_ptr barrier( + const c10d::BarrierOptions& opts) override; + c10::intrusive_ptr reduce( + std::vector& tensors, + const c10d::ReduceOptions& opts) override; + c10::intrusive_ptr gather( + std::vector>& outputTensors, + std::vector& inputTensors, + const c10d::GatherOptions& opts) override; + c10::intrusive_ptr scatter( + std::vector& outputTensors, + std::vector>& inputTensors, + const c10d::ScatterOptions& opts) override; + + void shutdown() override; + + at::Tensor getActiveRanksTensor() { return activeRanks_; } + int getNumSyncedRanks(); + void extendGroupSizeTo(int size); + std::vector getPeerState(const std::vector& ranks); + mooncakePgProposalResponse_t activateRanks(const std::vector& ranks); + mooncakePgProposalResponse_t deactivateRanks(const std::vector& ranks); + void joinGroup(); + uint64_t getCurrentEpoch() const; + mooncakePgSyncAfterFailureResponse_t syncAfterFailure(); + + private: + template + c10::intrusive_ptr launchCollective( + c10d::OpType opType, const char* operation, + const at::Tensor& streamTensor, std::vector keepAlive, + std::function postCompletion, Args... args); + + template + c10::intrusive_ptr launchP2P( + c10d::OpType opType, const char* operation, + const at::Tensor& streamTensor, std::vector keepAlive, + std::function postCompletion, Args... args); + + const c10::intrusive_ptr options_; + mooncakePgComm_t comm_ = nullptr; + at::Tensor activeRanks_; + bool isCpu_ = false; + bool isShutdown_ = false; + int max_group_size_ = + 0; // per-group capacity (max active members for this group) + std::shared_ptr work_tracker_; +}; + +} // namespace mooncake + +#endif // MOONCAKE_BACKEND_H diff --git a/mooncake-pg/torch/include/torch_utils.h b/mooncake-pg/torch/include/torch_utils.h new file mode 100644 index 0000000000..9ed88724d7 --- /dev/null +++ b/mooncake-pg/torch/include/torch_utils.h @@ -0,0 +1,18 @@ +#ifndef MOONCAKE_PG_TORCH_UTILS_H +#define MOONCAKE_PG_TORCH_UTILS_H + +#include + +#include + +namespace mooncake { + +inline void checkResult(mooncakePgResult_t result, const char* operation) { + TORCH_CHECK(result == mooncakePgSuccess, operation, + " failed: ", mooncakePgGetErrorString(result), ": ", + mooncakePgGetLastError()); +} + +} // namespace mooncake + +#endif // MOONCAKE_PG_TORCH_UTILS_H diff --git a/mooncake-pg/torch/include/work_handles.h b/mooncake-pg/torch/include/work_handles.h new file mode 100644 index 0000000000..8881ac5263 --- /dev/null +++ b/mooncake-pg/torch/include/work_handles.h @@ -0,0 +1,152 @@ +#ifndef MOONCAKE_WORK_HANDLES_H +#define MOONCAKE_WORK_HANDLES_H + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace mooncake { + +// Per-operation failedRanksHint buffer +struct FailedRanksHint { + at::Tensor tensor; + + FailedRanksHint() = default; + explicit FailedRanksHint(at::Tensor tensor_in) + : tensor(std::move(tensor_in)) {} + + int32_t* data() { return tensor.data_ptr(); } + const int32_t* data() const { return tensor.data_ptr(); } + + bool isLocalSuccess() const; + + static FailedRanksHint allocate(int size); +}; + +// Owns resources for operations whose c10d::Work was released before the +// operation completed. Eviction is deliberately opportunistic so Tensor +// destruction and post-completion callbacks run on a user thread. Captured +// CUDA operations remain retained until shutdown because their graph may +// replay. +class MooncakeWorkTracker final { + public: + MooncakeWorkTracker(); + ~MooncakeWorkTracker(); + + void evictCompleted() noexcept; + void shutdown() noexcept; + + private: + friend class MooncakeWorkCpu; + friend class MooncakeWorkCuda; + friend class MooncakeP2PWork; + + struct RetiredResources { + // The tracker only owns this payload; it never inspects its type. + std::any resources; + std::function ready_to_release; + }; + + void retire(std::any resources, + std::function ready_to_release) noexcept; + void retainUntilShutdown(std::any resources) noexcept; + + std::mutex mutex_; + std::vector retired_; + std::vector retained_until_shutdown_; + bool is_shutdown_ = false; +}; + +// Collective Work handles +class MooncakeWorkCpu : public ::c10d::Work { + public: + MooncakeWorkCpu(c10d::OpType opType, mooncakePgCompletion_t completion, + FailedRanksHint failedRanksHint, + std::shared_ptr tracker, + std::vector keepAlive = {}, + std::function postCompletion = {}); + ~MooncakeWorkCpu() override; + + bool isCompleted() override; + bool wait(std::chrono::milliseconds timeout) override; + + at::Tensor getFailedRanksHint() const; + bool getLocalSuccess() const; + + private: + std::shared_ptr<::mooncakePgCompletion> completion_; + FailedRanksHint failed_ranks_hint_; + std::shared_ptr tracker_; + std::vector keep_alive_; + // Idempotent wrapper around the optional callback. + std::function post_completion_; +}; + +class MooncakeWorkCuda : public ::c10d::Work { + public: + MooncakeWorkCuda(c10d::OpType opType, std::shared_ptr event, + FailedRanksHint failedRanksHint, + std::shared_ptr tracker, + std::vector keepAlive = {}); + ~MooncakeWorkCuda() override; + + bool isCompleted() override { return event_->query(); } + bool wait(std::chrono::milliseconds timeout) override; + + at::Tensor getFailedRanksHint() const; + bool getLocalSuccess() const; + + protected: + std::shared_ptr event_; + + private: + bool is_captured_ = false; + FailedRanksHint failed_ranks_hint_; + std::shared_ptr tracker_; + std::vector keep_alive_; +}; + +class MooncakeBarrierWorkCuda : public MooncakeWorkCuda { + public: + using MooncakeWorkCuda::MooncakeWorkCuda; + bool wait(std::chrono::milliseconds timeout) override; +}; + +// P2P Work handle +class MooncakeP2PWork : public ::c10d::Work { + public: + MooncakeP2PWork(c10d::OpType opType, mooncakePgCompletion_t completion, + FailedRanksHint failedRanksHint, + std::shared_ptr tracker, + std::vector keepAlive = {}, + std::function postCompletion = {}); + ~MooncakeP2PWork() override; + + bool isCompleted() override; + bool isSuccess() const override; + bool wait(std::chrono::milliseconds timeout) override; + at::Tensor getFailedRanksHint() const; + bool getLocalSuccess() const; + + private: + std::shared_ptr<::mooncakePgCompletion> completion_; + FailedRanksHint failed_ranks_hint_; + std::shared_ptr tracker_; + std::vector keep_alive_; + // Idempotent wrapper around the optional callback. + std::function post_completion_; +}; + +} // namespace mooncake + +#endif // MOONCAKE_WORK_HANDLES_H diff --git a/mooncake-pg/torch/setup.py b/mooncake-pg/torch/setup.py new file mode 100644 index 0000000000..b957d52ad4 --- /dev/null +++ b/mooncake-pg/torch/setup.py @@ -0,0 +1,85 @@ +import os +import re + +from setuptools import setup +import torch + +use_musa = os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in {"1", "ON", "TRUE", "YES"} +if use_musa: + try: + import importlib + + importlib.import_module("torchada") + except ImportError as e: + raise ImportError( + "torchada is required to build the MUSA PG extension. " + "Please install it first using 'pip install torchada'." + ) from e + + +from torch.utils.cpp_extension import ( # noqa: E402 + BuildExtension, + CUDAExtension, +) + + +torch_version = re.match(r"\d+(?:\.\d+)*", torch.__version__).group() +version_suffix = "_" + torch_version.replace(".", "_") +module_name = "mooncake.pg" + version_suffix + +abi_flag = int(torch._C._GLIBCXX_USE_CXX11_ABI) +current_dir = os.path.abspath(os.path.dirname(__file__)) + +abi_define = f"-D_GLIBCXX_USE_CXX11_ABI={abi_flag}" + +pg_core_so_path = os.getenv("MOONCAKE_PG_CORE_SO_PATH", "") +if not os.path.isfile(pg_core_so_path): + raise RuntimeError( + "MOONCAKE_PG_CORE_SO_PATH is unset or does not name " + "libmooncake_pg.so" + ) + +cxx_args = [ + abi_define, + "-std=c++20", + "-O3", + "-g0", +] + +include_dirs = [ + os.path.join(current_dir, "include"), + os.path.join(current_dir, "../include"), + os.path.join(current_dir, "../../mooncake-transfer-engine/include"), +] +use_maca = ( + os.getenv("MOONCAKE_EP_USE_MACA", "").upper() in {"1", "ON", "TRUE", "YES"} + or (hasattr(torch.version, "maca") and torch.version.maca is not None) +) + +if use_musa: + musa_defines = ["-DUSE_MUSA", "-DMOONCAKE_EP_USE_MUSA=1"] + cxx_args += musa_defines +else: + if use_maca: + cxx_args += ["-DUSE_MACA", "-DMOONCAKE_EP_USE_MACA=1"] + +setup( + name=module_name, + ext_modules=[ + CUDAExtension( + name=module_name, + include_dirs=include_dirs, + sources=[ + "src/pg_py.cpp", + "src/mooncake_backend.cpp", + "src/work_handles.cpp", + ], + extra_compile_args={"cxx": cxx_args}, + extra_objects=[pg_core_so_path], + extra_link_args=[ + "-Wl,-rpath,$ORIGIN", + ], + ) + ], + cmdclass={"build_ext": BuildExtension}, +) diff --git a/mooncake-pg/torch/src/mooncake_backend.cpp b/mooncake-pg/torch/src/mooncake_backend.cpp new file mode 100644 index 0000000000..a5fbf6ca0f --- /dev/null +++ b/mooncake-pg/torch/src/mooncake_backend.cpp @@ -0,0 +1,782 @@ +#include +#include + +#include +#ifdef MOONCAKE_EP_USE_MUSA +#include +#else +#include +#endif +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace mooncake { +namespace { + +constexpr const char* kSingleTensorError = + "Expecting one tensor only but got multiple."; +constexpr const char* kSparseError = "Sparse op not supported."; +mooncakePgDataType_t tensorType(const at::Tensor& tensor) { + switch (tensor.scalar_type()) { + case at::kChar: + return mooncakePgInt8; + case at::kByte: + return mooncakePgUint8; + case at::kShort: + return mooncakePgInt16; + case at::kUInt16: + return mooncakePgUint16; + case at::kInt: + return mooncakePgInt32; + case at::kUInt32: + return mooncakePgUint32; + case at::kLong: + return mooncakePgInt64; + case at::kUInt64: + return mooncakePgUint64; + case at::kHalf: + return mooncakePgFloat16; + case at::kFloat: + return mooncakePgFloat32; + case at::kDouble: + return mooncakePgFloat64; + case at::kBool: + return mooncakePgBool; + case at::kBFloat16: + return mooncakePgBfloat16; + case at::kFloat8_e4m3fn: + return mooncakePgFloat8e4m3fn; + case at::kFloat8_e5m2: + return mooncakePgFloat8e5m2; + case at::kFloat8_e4m3fnuz: + return mooncakePgFloat8e4m3fnuz; + case at::kFloat8_e5m2fnuz: + return mooncakePgFloat8e5m2fnuz; + case at::kFloat8_e8m0fnu: + return mooncakePgFloat8e8m0fnu; + default: + TORCH_CHECK(false, "Unsupported Mooncake PG datatype: ", + tensor.scalar_type()); + } +} + +mooncakePgReduceOp_t convertReduceOp(const c10d::ReduceOp& reduce_op) { + switch (reduce_op) { + case c10d::ReduceOp::SUM: + return mooncakePgSum; + case c10d::ReduceOp::AVG: + return mooncakePgAvg; + case c10d::ReduceOp::PRODUCT: + return mooncakePgProduct; + case c10d::ReduceOp::MIN: + return mooncakePgMin; + case c10d::ReduceOp::MAX: + return mooncakePgMax; + default: + TORCH_CHECK(false, "Unsupported Mooncake PG op: ", reduce_op); + } +} + +size_t tensorCount(const at::Tensor& tensor) { + TORCH_CHECK(tensor.numel() >= 0, "invalid Tensor element count"); + return static_cast(tensor.numel()); +} + +c10::cuda::CUDAStream currentCudaStream(const at::Tensor& tensor) { + return c10::cuda::getCurrentCUDAStream(tensor.device().index()); +} + +mooncakePgStream_t convertStream(const c10::cuda::CUDAStream& stream) { + return reinterpret_cast(stream.stream()); +} + +void validateEqualPeerTensors(const std::vector& tensors, + const at::Tensor& reference, int active_size) { + TORCH_CHECK(tensors.size() == static_cast(active_size), + "Tensor list size must match active group size"); + for (const auto& tensor : tensors) { + TORCH_CHECK(tensor.scalar_type() == reference.scalar_type(), + "All peer tensors must have the same dtype"); + TORCH_CHECK(tensor.device() == reference.device(), + "All peer tensors must be on the same device"); + TORCH_CHECK(tensor.numel() == reference.numel(), + "All peer tensors must have the same number of elements"); + } +} + +void validateSingleBufferTensors(const at::Tensor& output, + const at::Tensor& input, + c10::DeviceType expected_device) { + TORCH_CHECK(input.device().type() == expected_device, + "Input tensor device does not match the backend device"); + TORCH_CHECK(output.device() == input.device(), + "Input and output tensors must be on the same device"); + TORCH_CHECK(output.scalar_type() == input.scalar_type(), + "Input and output tensors must have the same dtype"); + TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous"); + TORCH_CHECK(output.is_contiguous(), "Output tensor must be contiguous"); +} + +at::Tensor packPeerTensors(const std::vector& tensors, + const at::Tensor& reference, int active_size) { + validateEqualPeerTensors(tensors, reference, active_size); + const int64_t elements_per_peer = reference.numel(); + auto packed = + at::empty({elements_per_peer * static_cast(tensors.size())}, + reference.options()); + for (size_t index = 0; index < tensors.size(); ++index) { + packed + .narrow(0, static_cast(index) * elements_per_peer, + elements_per_peer) + .copy_(tensors[index].reshape({elements_per_peer})); + } + return packed; +} + +std::function makeCopyBackToPeerTensors( + at::Tensor packed, std::vector outputs) { + return [packed = std::move(packed), + outputs = std::move(outputs)]() mutable { + if (outputs.empty()) return; + const int64_t elements_per_peer = outputs.front().numel(); + for (size_t index = 0; index < outputs.size(); ++index) { + outputs[index].copy_( + packed + .narrow(0, static_cast(index) * elements_per_peer, + elements_per_peer) + .view(outputs[index].sizes())); + } + }; +} + +std::vector convertRanks(const std::vector& ranks) { + return std::vector(ranks.begin(), ranks.end()); +} + +// Lightweight Backend shim that delegates operations back to the owning +// MooncakeBackend. PyTorch's P2P dispatch (batch_isend_irecv, isend, irecv) +// requires getBackend() to return a registered c10d::Backend instance. +// Since MooncakeBackend inherits from ProcessGroup (not Backend), we register +// this shim in the ProcessGroup's deviceTypeToBackend_ map. The shim holds a +// non-owning pointer to its owner. +// +// PyTorch 2.13 added ProcessGroup::all_gather_single and +// ProcessGroup::reduce_scatter_single, and the deprecated single-buffer +// aliases now forward to those methods. They dispatch through c10d/Ops.cpp +// and ProcessGroup::getBackend(dev), so calls land on this registered shim +// instead of MooncakeBackend's _allgather_base and _reduce_scatter_base +// overrides. Delegate every collective MooncakeBackend implements so the +// shim exposes the same capabilities as its owner. +class MooncakeBackendShim final : public ::c10d::Backend { + public: + MooncakeBackendShim(MooncakeBackend* owner, int maxGroupSize) + : Backend(owner->getRank(), maxGroupSize), owner_(owner) {} + + const std::string getBackendName() const override { return "mooncake"; } + bool supportsCoalescing() const override { return false; } + + c10::intrusive_ptr send(std::vector& tensors, + int dstRank, int tag) override { + return owner_->send(tensors, dstRank, tag); + } + + c10::intrusive_ptr recv(std::vector& tensors, + int srcRank, int tag) override { + return owner_->recv(tensors, srcRank, tag); + } + + c10::intrusive_ptr recvAnysource( + std::vector& tensors, int tag) override { + // MooncakeBackend doesn't implement recvAnysource; fall back to the + // base class which will raise a clear error. + return ::c10d::Backend::recvAnysource(tensors, tag); + } + + c10::intrusive_ptr barrier( + const c10d::BarrierOptions& opts) override { + return owner_->barrier(opts); + } + + // Signatures mirror MooncakeBackend's overrides so the shim re-exposes the + // same c10d virtuals. + c10::intrusive_ptr broadcast( + std::vector& tensors, + const c10d::BroadcastOptions& opts) override { + return owner_->broadcast(tensors, opts); + } + + c10::intrusive_ptr allreduce( + std::vector& tensors, + const c10d::AllreduceOptions& opts) override { + return owner_->allreduce(tensors, opts); + } + + c10::intrusive_ptr allgather( + std::vector>& outputTensors, + std::vector& inputTensors, + const c10d::AllgatherOptions& opts) override { + return owner_->allgather(outputTensors, inputTensors, opts); + } + + c10::intrusive_ptr _allgather_base( + at::Tensor& outputBuffer, at::Tensor& inputBuffer, + const c10d::AllgatherOptions& opts) override { + return owner_->_allgather_base(outputBuffer, inputBuffer, opts); + } + + c10::intrusive_ptr _reduce_scatter_base( + at::Tensor& outputBuffer, at::Tensor& inputBuffer, + const c10d::ReduceScatterOptions& opts) override { + return owner_->_reduce_scatter_base(outputBuffer, inputBuffer, opts); + } + + c10::intrusive_ptr alltoall( + std::vector& outputTensors, + std::vector& inputTensors, + const c10d::AllToAllOptions& opts) override { + return owner_->alltoall(outputTensors, inputTensors, opts); + } + + c10::intrusive_ptr reduce( + std::vector& tensors, + const c10d::ReduceOptions& opts) override { + return owner_->reduce(tensors, opts); + } + + c10::intrusive_ptr gather( + std::vector>& outputTensors, + std::vector& inputTensors, + const c10d::GatherOptions& opts) override { + return owner_->gather(outputTensors, inputTensors, opts); + } + + c10::intrusive_ptr scatter( + std::vector& outputTensors, + std::vector>& inputTensors, + const c10d::ScatterOptions& opts) override { + return owner_->scatter(outputTensors, inputTensors, opts); + } + + private: + // Non-owning: the shim is stored in ProcessGroup's backend maps which are + // cleared on destruction, and MooncakeBackend always outlives the shim. + MooncakeBackend* owner_; +}; + +} // namespace + +/** + * @brief Initialize Mooncake backend state from the PyTorch process-group + * information and optional Mooncake-specific options. + */ +MooncakeBackend::MooncakeBackend( + c10d::DistributedBackendOptions distBackendOpts, + c10::intrusive_ptr options, + mooncakePgContext_t context, bool isCpu) + : ProcessGroup(distBackendOpts.store, distBackendOpts.group_rank, + distBackendOpts.group_size), + options_(std::move(options)), + isCpu_(isCpu), + work_tracker_(std::make_shared()) { + TORCH_CHECK(context, "Mooncake PG core context is null"); + + const int rank = distBackendOpts.group_rank; + const int size = distBackendOpts.group_size; + max_group_size_ = options_ && options_->maxGroupSize_ > 0 + ? options_->maxGroupSize_ + : size; + TORCH_CHECK(max_group_size_ >= size && max_group_size_ > 0 && + max_group_size_ <= MOONCAKE_PG_MAX_RANKS, + "max_group_size must be in [group_size, ", + MOONCAKE_PG_MAX_RANKS, "]"); + TORCH_CHECK(rank >= 0 && rank < size, "rank out of valid range"); + TORCH_CHECK(!distBackendOpts.group_id.empty(), + "MooncakeBackend: group_id must not be empty"); + + // Use user-provided tensor memory if available. Only its storage is used; + // the Coordinator populates its contents through the core communicator. + if (options_ && options_->activeRanks_.defined()) { + activeRanks_ = options_->activeRanks_; + TORCH_CHECK(activeRanks_.scalar_type() == at::kInt, + "active_ranks must have dtype int32"); + TORCH_CHECK(activeRanks_.is_contiguous(), + "active_ranks must be contiguous"); + TORCH_CHECK(activeRanks_.numel() >= max_group_size_, + "active_ranks is smaller than max_group_size"); + TORCH_CHECK(activeRanks_.is_cpu() || activeRanks_.is_cuda(), + "active_ranks must be on a CPU or supported GPU device"); + } else { + activeRanks_ = + at::empty({max_group_size_}, + torch::dtype(torch::kInt32) + .device(isCpu_ ? torch::kCPU : torch::kCUDA)); + } + // The mirror follows its tensor storage, independently of the communicator + // device used for collectives. + const bool active_ranks_mirror_is_device = !activeRanks_.is_cpu(); + + std::vector global_ranks; + if (distBackendOpts.global_ranks_in_group.empty()) { + global_ranks.resize(size); + std::iota(global_ranks.begin(), global_ranks.end(), 0); + } else { + TORCH_CHECK(distBackendOpts.global_ranks_in_group.size() == + static_cast(size), + "global_ranks_in_group must contain group_size entries"); + global_ranks.reserve(size); + for (const auto global_rank : distBackendOpts.global_ranks_in_group) { + TORCH_CHECK(global_rank >= 0 && global_rank < MOONCAKE_PG_MAX_RANKS, + "global rank is outside the supported range"); + global_ranks.push_back(static_cast(global_rank)); + } + } + + mooncakePgCommConfig_t config = MOONCAKE_PG_COMM_CONFIG_INITIALIZER; + + // PyTorch's group_id is only a bootstrap id. The Coordinator resolves it + // together with rank order into a process-lifetime GroupId. CPU and device + // backends use independent namespaces. + config.groupId = distBackendOpts.group_id.c_str(); + config.rank = rank; + config.size = size; + config.maxGroupSize = max_group_size_; + config.globalRanks = global_ranks.data(); + config.globalRankCount = global_ranks.size(); + config.deviceIndex = isCpu_ ? -1 : at::cuda::current_device(); + config.deviceType = isCpu_ ? mooncakePgDeviceCpu : mooncakePgDeviceGpu; + config.idResolvePolicy = options_ && options_->isExtension_ + ? mooncakePgIdResolveAttachOrExtend + : mooncakePgIdResolveCreateOrAttach; + config.autoDeactivateOnFailure = + options_ + ? options_->autoDeactivateOnFailure_ + : c10::utils::check_env("MOONCAKE_PG_AUTO_DEACTIVATE_ON_FAILURE") + .value_or(true); + config.autoSyncOnFailure = + options_ ? options_->autoSyncOnFailure_ + : c10::utils::check_env("MOONCAKE_PG_AUTO_SYNC_ON_FAILURE") + .value_or(true); + // auto_sync_on_failure requires auto_deactivate_on_failure. + TORCH_CHECK(!config.autoSyncOnFailure || config.autoDeactivateOnFailure, + "auto_sync_on_failure requires " + "auto_deactivate_on_failure=true"); + config.activeRanksMirror = activeRanks_.data_ptr(); + config.activeRanksMirrorCount = static_cast(activeRanks_.numel()); + config.activeRanksMirrorIsDevice = active_ranks_mirror_is_device ? 1 : 0; + config.activeRanksMirrorDeviceIndex = + active_ranks_mirror_is_device ? activeRanks_.get_device() : -1; + + checkResult(mooncakePgCommCreate(context, &config, &comm_), + "mooncakePgCommCreate"); + + // Register a lightweight Backend shim so PyTorch dispatch can find a + // registered Backend for this ProcessGroup. The shim delegates supported + // P2P and collective operations back to this backend. + const auto device_type = + isCpu_ ? c10::DeviceType::CPU : c10::DeviceType::CUDA; + auto shim = c10::make_intrusive(this, max_group_size_); + setBackend(device_type, BackendType::CUSTOM, shim); +#ifndef MOONCAKE_EP_USE_MUSA + setDefaultBackend(BackendType::CUSTOM); +#endif +} + +MooncakeBackend::~MooncakeBackend() { + try { + shutdown(); + } catch (const std::exception& error) { + TORCH_WARN("MooncakeBackend: shutdown failed during destruction: ", + error.what()); + } catch (...) { + TORCH_WARN("MooncakeBackend: shutdown failed during destruction"); + } +} + +const std::string MooncakeBackend::getBackendName() const { return "mooncake"; } + +int MooncakeBackend::getSize() const { + int size = 0; + checkResult(mooncakePgCommGetSize(comm_, &size), "mooncakePgCommGetSize"); + return size; +} + +template +c10::intrusive_ptr MooncakeBackend::launchCollective( + c10d::OpType opType, const char* operation, const at::Tensor& streamTensor, + std::vector keepAlive, std::function postCompletion, + Args... args) { + auto failed_ranks_hint = FailedRanksHint::allocate(max_group_size_); + const auto failed_ranks_hint_count = static_cast(max_group_size_); + if (isCpu_) { + mooncakePgCompletion_t completion = nullptr; + checkResult(CpuFn(args..., comm_, failed_ranks_hint.data(), + failed_ranks_hint_count, &completion), + operation); + work_tracker_->evictCompleted(); + return c10::make_intrusive( + opType, completion, std::move(failed_ranks_hint), work_tracker_, + std::move(keepAlive), std::move(postCompletion)); + } + + const auto stream = currentCudaStream(streamTensor); + checkResult(GpuFn(args..., comm_, convertStream(stream), + failed_ranks_hint.data(), failed_ranks_hint_count), + operation); + if (postCompletion) postCompletion(); + auto event = std::make_shared(c10::DeviceType::CUDA); + event->record(stream); + if (at::cuda::currentStreamCaptureStatus() == + c10::cuda::CaptureStatus::None) { + work_tracker_->evictCompleted(); + } + return c10::make_intrusive( + opType, std::move(event), std::move(failed_ranks_hint), work_tracker_, + std::move(keepAlive)); +} + +template +c10::intrusive_ptr MooncakeBackend::launchP2P( + c10d::OpType opType, const char* operation, const at::Tensor& streamTensor, + std::vector keepAlive, std::function postCompletion, + Args... args) { + auto failed_ranks_hint = FailedRanksHint::allocate(max_group_size_); + const auto failed_ranks_hint_count = static_cast(max_group_size_); + mooncakePgCompletion_t completion = nullptr; + if (isCpu_) { + checkResult(CpuFn(args..., comm_, failed_ranks_hint.data(), + failed_ranks_hint_count, &completion), + operation); + } else { + const auto stream = currentCudaStream(streamTensor); + checkResult(GpuFn(args..., comm_, convertStream(stream), + failed_ranks_hint.data(), failed_ranks_hint_count, + &completion), + operation); + } + + if (isCpu_ || at::cuda::currentStreamCaptureStatus() == + c10::cuda::CaptureStatus::None) { + work_tracker_->evictCompleted(); + } + return c10::make_intrusive( + opType, completion, std::move(failed_ranks_hint), work_tracker_, + std::move(keepAlive), std::move(postCompletion)); +} + +c10::intrusive_ptr MooncakeBackend::send( + std::vector& tensors, int dstRank, int tag) { + (void)tag; + TORCH_CHECK(tensors.size() == 1, kSingleTensorError); + auto tensor = tensors.back().contiguous(); + return launchP2P( + c10d::OpType::SEND, "mooncakePgSend", tensor, {tensor}, {}, + tensor.data_ptr(), tensorCount(tensor), tensorType(tensor), dstRank); +} + +c10::intrusive_ptr MooncakeBackend::recv( + std::vector& tensors, int srcRank, int tag) { + (void)tag; + TORCH_CHECK(tensors.size() == 1, kSingleTensorError); + auto output = tensors.back(); + const bool copy_back = !output.is_contiguous(); + auto target = copy_back ? output.contiguous() : output; + std::function post_completion; + if (copy_back) { + post_completion = [output, target, is_cpu = isCpu_]() mutable { + output.copy_(target); + if (!is_cpu) currentCudaStream(output).synchronize(); + }; + } + return launchP2P( + c10d::OpType::RECV, "mooncakePgRecv", target, {output, target}, + std::move(post_completion), target.data_ptr(), tensorCount(target), + tensorType(target), srcRank); +} + +c10::intrusive_ptr MooncakeBackend::broadcast( + std::vector& tensors, const c10d::BroadcastOptions& opts) { + TORCH_CHECK(tensors.size() == 1, kSingleTensorError); + auto tensor = tensors.back(); + const int root = opts.rootRank + opts.rootTensor; + return launchCollective( + c10d::OpType::BROADCAST, "mooncakePgBroadcast", tensor, {tensor}, {}, + tensor.data_ptr(), tensor.data_ptr(), tensorCount(tensor), + tensorType(tensor), root); +} + +c10::intrusive_ptr MooncakeBackend::allreduce( + std::vector& tensors, const c10d::AllreduceOptions& opts) { + TORCH_CHECK(tensors.size() == 1, kSingleTensorError); + TORCH_CHECK(opts.sparseIndices == std::nullopt, kSparseError); + auto tensor = tensors.back(); + return launchCollective( + c10d::OpType::ALLREDUCE, "mooncakePgAllReduce", tensor, {tensor}, {}, + tensor.data_ptr(), tensor.data_ptr(), tensorCount(tensor), + tensorType(tensor), convertReduceOp(opts.reduceOp)); +} + +c10::intrusive_ptr MooncakeBackend::allgather( + std::vector>& outputTensors, + std::vector& inputTensors, const c10d::AllgatherOptions&) { + TORCH_CHECK(inputTensors.size() == 1, kSingleTensorError); + TORCH_CHECK(outputTensors.size() == 1, kSingleTensorError); + auto input = inputTensors.back(); + auto outputs = outputTensors.back(); + const int active_size = getSize(); + validateEqualPeerTensors(outputs, input, active_size); + auto packed_output = + at::empty({input.numel() * static_cast(outputs.size())}, + input.options()); + + std::vector keep_alive{input, packed_output}; + keep_alive.insert(keep_alive.end(), outputs.begin(), outputs.end()); + auto post_completion = + makeCopyBackToPeerTensors(packed_output, std::move(outputs)); + return launchCollective( + c10d::OpType::ALLGATHER, "mooncakePgAllGather", input, + std::move(keep_alive), std::move(post_completion), input.data_ptr(), + packed_output.data_ptr(), tensorCount(input), tensorType(input)); +} + +c10::intrusive_ptr MooncakeBackend::_allgather_base( + at::Tensor& outputBuffer, at::Tensor& inputBuffer, + const c10d::AllgatherOptions&) { + validateSingleBufferTensors( + outputBuffer, inputBuffer, + isCpu_ ? c10::DeviceType::CPU : c10::DeviceType::CUDA); + + return launchCollective( + c10d::OpType::_ALLGATHER_BASE, "mooncakePgAllGather", inputBuffer, + {inputBuffer, outputBuffer}, {}, inputBuffer.data_ptr(), + outputBuffer.data_ptr(), tensorCount(inputBuffer), + tensorType(inputBuffer)); +} + +c10::intrusive_ptr MooncakeBackend::_reduce_scatter_base( + at::Tensor& outputBuffer, at::Tensor& inputBuffer, + const c10d::ReduceScatterOptions& opts) { + validateSingleBufferTensors( + outputBuffer, inputBuffer, + isCpu_ ? c10::DeviceType::CPU : c10::DeviceType::CUDA); + + return launchCollective( + c10d::OpType::_REDUCE_SCATTER_BASE, "mooncakePgReduceScatter", + outputBuffer, {inputBuffer, outputBuffer}, {}, inputBuffer.data_ptr(), + outputBuffer.data_ptr(), tensorCount(outputBuffer), + tensorType(outputBuffer), convertReduceOp(opts.reduceOp)); +} + +c10::intrusive_ptr MooncakeBackend::alltoall( + std::vector& outputTensors, + std::vector& inputTensors, const c10d::AllToAllOptions&) { + TORCH_CHECK(!inputTensors.empty() && !outputTensors.empty(), + "alltoall requires non-empty Tensor lists"); + const auto reference = inputTensors.front(); + const int active_size = getSize(); + validateEqualPeerTensors(outputTensors, reference, active_size); + auto packed_input = packPeerTensors(inputTensors, reference, active_size); + auto packed_output = at::empty( + {reference.numel() * static_cast(outputTensors.size())}, + reference.options()); + + std::vector keep_alive{packed_input, packed_output}; + keep_alive.insert(keep_alive.end(), inputTensors.begin(), + inputTensors.end()); + keep_alive.insert(keep_alive.end(), outputTensors.begin(), + outputTensors.end()); + auto post_completion = + makeCopyBackToPeerTensors(packed_output, outputTensors); + return launchCollective( + c10d::OpType::ALLTOALL, "mooncakePgAllToAll", reference, + std::move(keep_alive), std::move(post_completion), + packed_input.data_ptr(), packed_output.data_ptr(), + tensorCount(reference), tensorType(reference)); +} + +c10::intrusive_ptr MooncakeBackend::barrier( + const c10d::BarrierOptions&) { + auto failed_ranks_hint = FailedRanksHint::allocate(max_group_size_); + const auto failed_ranks_hint_count = static_cast(max_group_size_); + if (isCpu_) { + mooncakePgCompletion_t completion = nullptr; + checkResult(mooncakePgBarrierCpu(comm_, failed_ranks_hint.data(), + failed_ranks_hint_count, &completion), + "mooncakePgBarrierCpu"); + work_tracker_->evictCompleted(); + return c10::make_intrusive( + c10d::OpType::BARRIER, completion, std::move(failed_ranks_hint), + work_tracker_); + } + + const auto stream = c10::cuda::getCurrentCUDAStream(); + checkResult( + mooncakePgBarrierGpu(comm_, convertStream(stream), + failed_ranks_hint.data(), failed_ranks_hint_count), + "mooncakePgBarrierGpu"); + auto event = std::make_shared(c10::DeviceType::CUDA); + event->record(stream); + if (at::cuda::currentStreamCaptureStatus() == + c10::cuda::CaptureStatus::None) { + work_tracker_->evictCompleted(); + } + return c10::make_intrusive( + c10d::OpType::BARRIER, std::move(event), std::move(failed_ranks_hint), + work_tracker_); +} + +c10::intrusive_ptr MooncakeBackend::reduce( + std::vector& tensors, const c10d::ReduceOptions& opts) { + TORCH_CHECK(tensors.size() == 1, kSingleTensorError); + auto tensor = tensors.back(); + const int root = opts.rootRank + opts.rootTensor; + return launchCollective( + c10d::OpType::REDUCE, "mooncakePgReduce", tensor, {tensor}, {}, + tensor.data_ptr(), tensor.data_ptr(), tensorCount(tensor), + tensorType(tensor), convertReduceOp(opts.reduceOp), root); +} + +c10::intrusive_ptr MooncakeBackend::gather( + std::vector>& outputTensors, + std::vector& inputTensors, const c10d::GatherOptions& opts) { + TORCH_CHECK(inputTensors.size() == 1, kSingleTensorError); + const int root = opts.rootRank; + const bool is_root = root == rank_; + if (is_root) { + TORCH_CHECK(outputTensors.size() == 1, kSingleTensorError); + } + auto input = inputTensors.back(); + std::vector outputs; + at::Tensor packed_output; + if (is_root) { + outputs = outputTensors.back(); + const int active_size = getSize(); + validateEqualPeerTensors(outputs, input, active_size); + packed_output = + at::empty({input.numel() * static_cast(outputs.size())}, + input.options()); + } + + std::vector keep_alive{input, packed_output}; + keep_alive.insert(keep_alive.end(), outputs.begin(), outputs.end()); + auto post_completion = packed_output.defined() ? makeCopyBackToPeerTensors( + packed_output, outputs) + : std::function{}; + return launchCollective( + c10d::OpType::GATHER, "mooncakePgGather", input, std::move(keep_alive), + std::move(post_completion), input.data_ptr(), + packed_output.defined() ? packed_output.data_ptr() : nullptr, + tensorCount(input), tensorType(input), root); +} + +c10::intrusive_ptr MooncakeBackend::scatter( + std::vector& outputTensors, + std::vector>& inputTensors, + const c10d::ScatterOptions& opts) { + TORCH_CHECK(outputTensors.size() == 1, kSingleTensorError); + const int root = opts.rootRank; + const bool is_root = root == rank_; + if (is_root) { + TORCH_CHECK(inputTensors.size() == 1, kSingleTensorError); + } + auto output = outputTensors.back(); + at::Tensor packed_input; + std::vector inputs; + if (is_root) { + inputs = inputTensors.back(); + packed_input = packPeerTensors(inputs, output, getSize()); + } + + std::vector keep_alive{output, packed_input}; + keep_alive.insert(keep_alive.end(), inputs.begin(), inputs.end()); + return launchCollective( + c10d::OpType::SCATTER, "mooncakePgScatter", output, + std::move(keep_alive), {}, + packed_input.defined() ? packed_input.data_ptr() : nullptr, + output.data_ptr(), tensorCount(output), tensorType(output), root); +} + +void MooncakeBackend::shutdown() { + if (isShutdown_) return; + isShutdown_ = true; + auto comm = std::exchange(comm_, nullptr); + const auto result = comm ? mooncakePgCommDestroy(comm) : mooncakePgSuccess; + work_tracker_->shutdown(); + checkResult(result, "mooncakePgCommDestroy"); +} + +void MooncakeBackend::extendGroupSizeTo(int) { + // Deprecated: in the Coordinator-based path, group size is determined by + // GroupView.rank_order. This is a no-op stub kept for compatibility. + TORCH_WARN( + "MooncakeBackend::extendGroupSizeTo is deprecated; group size " + "is controlled by the Coordinator's GroupView."); +} + +int MooncakeBackend::getNumSyncedRanks() { + int num_synced_ranks = 0; + checkResult(mooncakePgCommGetNumSyncedRanks(comm_, &num_synced_ranks), + "mooncakePgCommGetNumSyncedRanks"); + return num_synced_ranks; +} + +std::vector MooncakeBackend::getPeerState(const std::vector& ranks) { + const auto core_ranks = convertRanks(ranks); + std::vector core_states(ranks.size(), 0); + checkResult( + mooncakePgCommGetPeerState(comm_, core_ranks.data(), core_ranks.size(), + core_states.data()), + "mooncakePgCommGetPeerState"); + std::vector result; + result.reserve(core_states.size()); + for (const int32_t state : core_states) result.push_back(state != 0); + return result; +} + +mooncakePgProposalResponse_t MooncakeBackend::activateRanks( + const std::vector& ranks) { + const auto core_ranks = convertRanks(ranks); + mooncakePgProposalResponse_t response{}; + checkResult(mooncakePgCommActivateRanks(comm_, core_ranks.data(), + core_ranks.size(), &response), + "mooncakePgCommActivateRanks"); + return response; +} + +mooncakePgProposalResponse_t MooncakeBackend::deactivateRanks( + const std::vector& ranks) { + const auto core_ranks = convertRanks(ranks); + mooncakePgProposalResponse_t response{}; + checkResult(mooncakePgCommDeactivateRanks(comm_, core_ranks.data(), + core_ranks.size(), &response), + "mooncakePgCommDeactivateRanks"); + return response; +} + +void MooncakeBackend::joinGroup() { + checkResult(mooncakePgCommJoin(comm_), "mooncakePgCommJoin"); +} + +uint64_t MooncakeBackend::getCurrentEpoch() const { + uint64_t epoch = 0; + checkResult(mooncakePgCommGetEpoch(comm_, &epoch), + "mooncakePgCommGetEpoch"); + return epoch; +} + +mooncakePgSyncAfterFailureResponse_t MooncakeBackend::syncAfterFailure() { + mooncakePgSyncAfterFailureResponse_t response{}; + checkResult(mooncakePgCommSyncAfterFailure(comm_, &response), + "mooncakePgCommSyncAfterFailure"); + return response; +} + +} // namespace mooncake diff --git a/mooncake-pg/torch/src/pg_py.cpp b/mooncake-pg/torch/src/pg_py.cpp new file mode 100644 index 0000000000..db1e441d7b --- /dev/null +++ b/mooncake-pg/torch/src/pg_py.cpp @@ -0,0 +1,400 @@ +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace py = pybind11; + +namespace mooncake { + +constexpr const char* kCoordinatorStoreKey = "coordinator_addr"; +constexpr size_t kCoordinatorAddressBufSize = 256; + +struct MooncakeProcessContext { + mooncakePgContext_t handle = nullptr; + c10::intrusive_ptr bootstrap_store; + + ~MooncakeProcessContext() { + if (handle) (void)mooncakePgContextDestroy(handle); + } +}; + +static MooncakeProcessContext g_ctx; +static std::once_flag g_create_context_once; +static std::once_flag g_init_control_plane_once; + +mooncakePgContext_t getContext() { + std::call_once(g_create_context_once, [] { + mooncakePgContext_t context = nullptr; + checkResult(mooncakePgContextCreate(&context), + "mooncakePgContextCreate"); + g_ctx.handle = context; + }); + return g_ctx.handle; +} + +static mooncakePgContext_t initControlPlane( + const c10::intrusive_ptr& store, int rank, + int max_world_size) { + auto context = getContext(); + std::call_once(g_init_control_plane_once, [&] { + // Ordering constraint: AgentHost::start() sends registerAgent + // immediately, which includes LinkManager's localServerName() and + // getWarmupRecvAddr(). These must be non-empty, so the engine and + // LinkManager must be initialized BEFORE AgentHost starts. + checkResult(mooncakePgContextInitialize(context, rank, max_world_size), + "mooncakePgContextInitialize"); + + // Rank 0 hosts the Coordinator in-process. + if (rank == 0) { + std::array + coordinator_address_buf{}; + checkResult(mooncakePgContextLaunchCoordinator( + context, coordinator_address_buf.data(), + coordinator_address_buf.size()), + "mooncakePgContextLaunchCoordinator"); + store->set(kCoordinatorStoreKey, + std::string(coordinator_address_buf.data())); + } + + store->wait({kCoordinatorStoreKey}); + const std::string value = store->get_to_str(kCoordinatorStoreKey); + TORCH_CHECK(!value.empty(), + "invalid Mooncake coordinator address in Store"); + checkResult(mooncakePgContextConnectCoordinator(context, value.c_str()), + "mooncakePgContextConnectCoordinator"); + + // Keep the first rendezvous Store alive for the process-wide control + // plane. In particular, this keeps rank 0's TCPStore server alive while + // the default ProcessGroup is destroyed and re-created. + g_ctx.bootstrap_store = store; + }); + return context; +} + +c10::intrusive_ptr createMooncakeBackend( + c10d::DistributedBackendOptions distBackendOpts, + c10::intrusive_ptr + backendOptions) { + int rank = distBackendOpts.group_rank; + auto context = + initControlPlane(distBackendOpts.store, rank, MOONCAKE_PG_MAX_RANKS); + auto backend = c10::make_intrusive( + std::move(distBackendOpts), std::move(backendOptions), context); + return backend; +} + +c10::intrusive_ptr createMooncakeCpuBackend( + c10d::DistributedBackendOptions distBackendOpts, + c10::intrusive_ptr + backendOptions) { + int rank = distBackendOpts.group_rank; + auto context = + initControlPlane(distBackendOpts.store, rank, MOONCAKE_PG_MAX_RANKS); + auto backend = c10::make_intrusive( + std::move(distBackendOpts), std::move(backendOptions), context, true); + return backend; +} + +__attribute__((constructor)) static void MooncakeBackendConstructor() { + py::object module = py::module::import("torch.distributed"); + py::object register_backend = + module.attr("Backend").attr("register_backend"); + py::dict kwargsCpu; + kwargsCpu["devices"] = py::make_tuple("cpu"); + register_backend("mooncake-cpu", py::cpp_function(createMooncakeCpuBackend), + /* extended_api */ true, **kwargsCpu); +#ifndef MOONCAKE_EP_USE_MUSA + py::dict kwargsCuda; + kwargsCuda["devices"] = py::make_tuple("cuda"); + register_backend("mooncake", py::cpp_function(createMooncakeBackend), + /* extended_api */ true, **kwargsCuda); +#else + py::dict kwargsMusa; + kwargsMusa["devices"] = py::make_tuple("musa"); + register_backend("mooncake", py::cpp_function(createMooncakeBackend), + /* extended_api */ true, **kwargsMusa); +#endif +} + +at::Tensor getActiveRanks(c10::intrusive_ptr backend) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + return mooncakeBackend->getActiveRanksTensor(); +} + +int getNumSyncedRanks(c10::intrusive_ptr backend) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + return mooncakeBackend->getNumSyncedRanks(); +} + +void extendGroupSizeTo(c10::intrusive_ptr backend, + int size) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + mooncakeBackend->extendGroupSizeTo(size); +} + +std::vector getPeerState(c10::intrusive_ptr backend, + const std::vector& ranks) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + return mooncakeBackend->getPeerState(ranks); +} + +mooncakePgProposalResponse_t recoverRanks( + c10::intrusive_ptr backend, + const std::vector& ranks) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + return mooncakeBackend->activateRanks(ranks); +} + +mooncakePgProposalResponse_t deactivateRanks( + c10::intrusive_ptr backend, + const std::vector& ranks) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + return mooncakeBackend->deactivateRanks(ranks); +} + +mooncakePgProposalResponse_t activateRanks( + c10::intrusive_ptr backend, + const std::vector& ranks) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + return mooncakeBackend->activateRanks(ranks); +} + +void joinGroup(c10::intrusive_ptr backend) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + mooncakeBackend->joinGroup(); +} + +at::Tensor getFailedRanksHint(c10::intrusive_ptr work) { + if (auto* w = dynamic_cast(work.get())) { + return w->getFailedRanksHint(); + } + if (auto* w = dynamic_cast(work.get())) { + return w->getFailedRanksHint(); + } + if (auto* w = dynamic_cast(work.get())) { + return w->getFailedRanksHint(); + } + return at::Tensor(); +} + +bool getLocalSuccess(c10::intrusive_ptr work) { + if (auto* w = dynamic_cast(work.get())) { + return w->getLocalSuccess(); + } + if (auto* w = dynamic_cast(work.get())) { + return w->getLocalSuccess(); + } + if (auto* w = dynamic_cast(work.get())) { + return w->getLocalSuccess(); + } + return false; +} + +int64_t getCurrentEpoch(c10::intrusive_ptr backend) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + return static_cast(mooncakeBackend->getCurrentEpoch()); +} + +/// Python-facing wrapper that extracts the raw TransferEngine* from a +/// mooncake.engine.TransferEngine Python object and makes it the process-wide +/// engine for all MooncakeBackend instances. The caller must ensure the +/// TransferEnginePy object outlives all MooncakeBackend instances. +void setTransferEnginePy(pybind11::object engine_obj) { + if (engine_obj.is_none()) { + checkResult(mooncakePgContextSetTransferEngine(getContext(), nullptr), + "mooncakePgContextSetTransferEngine"); + return; + } + auto get_engine_ptr = engine_obj.attr("get_engine_ptr"); + uintptr_t ptr = get_engine_ptr().cast(); + checkResult(mooncakePgContextSetTransferEngine( + getContext(), reinterpret_cast(ptr)), + "mooncakePgContextSetTransferEngine"); +} + +std::vector droppedRanks(const mooncakePgProposalResponse_t& response) { + const size_t count = std::min(response.droppedRankCount, + static_cast(MOONCAKE_PG_MAX_RANKS)); + return std::vector(response.droppedRanks, + response.droppedRanks + count); +} + +void shutdownProcessContext() { + auto context = g_ctx.handle; + if (!context) return; + // ContextDestroy rejects a parent-before-child teardown while a + // communicator is still alive. Keep the handle for the static-destructor + // fallback instead of losing ownership. + if (mooncakePgContextDestroy(context) == mooncakePgSuccess) { + g_ctx.handle = nullptr; + g_ctx.bootstrap_store.reset(); + } +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + // Python atexit handlers run while module globals still own an injected + // TransferEngine. Py_AtExit is too late: CPython may have already decref'd + // the Python TE wrapper before invoking its native exit handlers. + py::module_::import("atexit").attr("register")( + py::cpp_function(&shutdownProcessContext)); + m.def("createMooncakeBackend", &createMooncakeBackend); + m.def("createMooncakeCpuBackend", &createMooncakeCpuBackend); + m.def("set_host_ip", [](const std::string& host) { + checkResult(mooncakePgContextSetHostIp(getContext(), host.c_str()), + "mooncakePgContextSetHostIp"); + }); + m.def( + "set_collective_timeout_us", + [](size_t us) { + checkResult(mooncakePgContextSetCollectiveTimeout(getContext(), us), + "mooncakePgContextSetCollectiveTimeout"); + }, + py::arg("us"), + "Set the default peer-liveness probe timeout (microseconds) for " + "collective operations."); + m.def( + "set_p2p_timeout_us", + [](int64_t us) { + checkResult(mooncakePgContextSetP2PTimeout(getContext(), us), + "mooncakePgContextSetP2PTimeout"); + }, + py::arg("us"), "Set the default P2P transfer timeout (microseconds)."); + m.def( + "set_fault_reconciliation_window_us", + [](int64_t us) { + checkResult( + mooncakePgContextSetFaultReconciliationWindow(getContext(), us), + "mooncakePgContextSetFaultReconciliationWindow"); + }, + py::arg("us"), + "Set the coordinator fault reconciliation window (microseconds)."); + m.def("set_device_filter", [](std::vector filters) { + std::vector filter_pointers; + filter_pointers.reserve(filters.size()); + for (const auto& filter : filters) { + filter_pointers.push_back(filter.c_str()); + } + checkResult( + mooncakePgContextSetDeviceFilter( + getContext(), filter_pointers.data(), filter_pointers.size()), + "mooncakePgContextSetDeviceFilter"); + }); + m.def("set_transfer_engine", &setTransferEnginePy, py::arg("engine"), + "Set an external TransferEngine to be used by MooncakeBackend. " + "Must be called before init_process_group(). The engine must already " + "be initialized. Pass None to reset to default behavior. " + "The caller must ensure the TransferEngine object outlives all " + "MooncakeBackend instances."); + m.def("get_active_ranks", &getActiveRanks); + m.def("get_num_synced_ranks", &getNumSyncedRanks); + m.def("extend_group_size_to", &extendGroupSizeTo); + m.def("get_peer_state", &getPeerState); + m.def("recover_ranks", &recoverRanks); + m.def("activate_ranks", &activateRanks); + m.def("deactivate_ranks", &deactivateRanks, py::arg("backend"), + py::arg("ranks")); + m.def("join_group", &joinGroup); + m.def("get_failed_ranks_hint", &getFailedRanksHint, py::arg("work")); + m.def("get_local_success", &getLocalSuccess, py::arg("work"), + "Return True iff all locally-attempted peers succeeded in this " + "operation."); + m.def("get_current_epoch", &getCurrentEpoch, py::arg("backend"), + "Get the current GroupView epoch (monotonically increasing on " + "membership changes)."); + + m.def( + "sync_after_failure", + [](c10::intrusive_ptr backend) { + auto mooncakeBackend = + c10::static_intrusive_pointer_cast(backend); + return mooncakeBackend->syncAfterFailure(); + }, + py::arg("backend")); + + py::enum_(m, "SyncAfterFailureStatus") + .value("Reconciled", mooncakePgSyncReconciled) + .value("NoPending", mooncakePgSyncNoPending) + .value("Rejected", mooncakePgSyncRejected); + + auto proposal_status = + py::enum_(m, "ProposalStatus") + .value("Rejected", mooncakePgProposalRejected) + .value("Applied", mooncakePgProposalApplied) + .value("AppliedWithDroppedRanks", + mooncakePgProposalAppliedWithDroppedRanks); + // Keep existing Python callers source-compatible with the renamed enum. + m.attr("ViewUpdateStatus") = proposal_status; + + py::class_(m, + "SyncAfterFailureResponse") + .def_property_readonly( + "status", + [](const mooncakePgSyncAfterFailureResponse_t& value) { + return value.status; + }) + .def_property_readonly( + "reject_reason", + [](const mooncakePgSyncAfterFailureResponse_t& value) { + return std::string(value.rejectReason); + }); + + py::class_(m, "ProposeViewUpdateResponse") + .def_property_readonly("status", + [](const mooncakePgProposalResponse_t& value) { + return value.status; + }) + .def_property_readonly("new_epoch", + [](const mooncakePgProposalResponse_t& value) { + return value.newEpoch; + }) + .def_property_readonly("dropped_ranks", &droppedRanks) + .def_property_readonly("reject_reason", + [](const mooncakePgProposalResponse_t& value) { + return std::string(value.rejectReason); + }); + + py::class_>( + m, "MooncakeBackendOptions") + // IMPORTANT: these constructors with tensor MUST be registered + // before the (int, ...) constructors. Otherwise, when a 1-element + // Tensor is passed, pybind11 implicitly converts it to int and + // resolves to the wrong overload: + // e.g. MooncakeBackendOptions(tensor([1]), False) -> + // MooncakeBackendOptions(int maxGroupSize=1, + // bool isExtension=False) + // instead of the intended Tensor-based path. + .def(py::init(), py::arg("active_ranks")) + .def(py::init(), py::arg("active_ranks"), + py::arg("is_extension")) + .def(py::init(), py::arg("active_ranks"), + py::arg("is_extension"), py::arg("max_group_size")) + // Recommended constructors + .def(py::init(), py::arg("max_group_size")) + .def(py::init(), py::arg("max_group_size"), + py::arg("is_extension")) + .def(py::init(), py::arg("max_group_size"), + py::arg("is_extension"), py::arg("auto_deactivate_on_failure"), + py::arg("auto_sync_on_failure")); +} + +} // namespace mooncake diff --git a/mooncake-pg/torch/src/work_handles.cpp b/mooncake-pg/torch/src/work_handles.cpp new file mode 100644 index 0000000000..23bbb999a9 --- /dev/null +++ b/mooncake-pg/torch/src/work_handles.cpp @@ -0,0 +1,354 @@ +#include +#include +#include + +#include +#ifdef MOONCAKE_EP_USE_MUSA +#include +#else +#include +#endif + +#include +#include +#include +#include +#include + +namespace mooncake { +namespace { + +using CompletionHandle = std::shared_ptr<::mooncakePgCompletion>; + +CompletionHandle makeCompletionHandle(mooncakePgCompletion_t completion) { + TORCH_CHECK(completion, "Mooncake PG core returned a null completion"); + return CompletionHandle(completion, [](mooncakePgCompletion_t handle) { + if (handle) (void)mooncakePgCompletionDestroy(handle); + }); +} + +bool queryCompletion(mooncakePgCompletion_t completion) { + int completed = 0; + checkResult(mooncakePgCompletionIsCompleted(completion, &completed), + "mooncakePgCompletionIsCompleted"); + return completed != 0; +} + +bool waitCompletion(mooncakePgCompletion_t completion, int64_t timeout_us) { + const auto result = mooncakePgCompletionWait(completion, timeout_us); + if (result == mooncakePgTimeout) return false; + checkResult(result, "mooncakePgCompletionWait"); + return true; +} + +std::function makePostCompletionOnce( + std::function post_completion) { + if (!post_completion) return {}; + return [once = std::make_shared(), + post_completion = std::move(post_completion)]() mutable { + std::call_once(*once, [&] { post_completion(); }); + }; +} + +template +std::any makeTrackedResources(Resources&&... resources) { + return std::any(std::make_tuple(std::forward(resources)...)); +} + +} // namespace + +FailedRanksHint FailedRanksHint::allocate(int size) { + auto options = + torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU); + return FailedRanksHint(torch::zeros({size}, options)); +} + +bool FailedRanksHint::isLocalSuccess() const { + const auto* values = data(); + return std::all_of(values, values + tensor.numel(), + [](int32_t value) { return value == 0; }); +} + +MooncakeWorkTracker::MooncakeWorkTracker() = default; + +MooncakeWorkTracker::~MooncakeWorkTracker() { shutdown(); } + +void MooncakeWorkTracker::retire( + std::any resources, std::function ready_to_release) noexcept { + // Work destruction only transfers ownership. In particular, it does not + // run a post-completion callback that may perform a device copy. + std::lock_guard lock(mutex_); + if (is_shutdown_) return; + retired_.push_back({std::move(resources), std::move(ready_to_release)}); +} + +void MooncakeWorkTracker::retainUntilShutdown(std::any resources) noexcept { + std::lock_guard lock(mutex_); + if (is_shutdown_) return; + retained_until_shutdown_.push_back(std::move(resources)); +} + +void MooncakeWorkTracker::evictCompleted() noexcept { + std::vector candidates; + { + std::lock_guard lock(mutex_); + if (is_shutdown_ || retired_.empty()) return; + candidates.swap(retired_); + } + + std::vector pending; + pending.reserve(candidates.size()); + for (auto& candidate : candidates) { + try { + if (candidate.ready_to_release()) continue; + } catch (const std::exception& error) { + TORCH_WARN( + "MooncakeWorkTracker: failed to process retired work; " + "resources remain retained: ", + error.what()); + } catch (...) { + TORCH_WARN( + "MooncakeWorkTracker: failed to process retired work; " + "resources remain retained: unknown exception"); + } + pending.push_back(std::move(candidate)); + } + + std::lock_guard lock(mutex_); + if (is_shutdown_) return; + retired_.insert(retired_.end(), std::make_move_iterator(pending.begin()), + std::make_move_iterator(pending.end())); +} + +void MooncakeWorkTracker::shutdown() noexcept { + std::vector retired; + std::vector retained_until_shutdown; + { + std::lock_guard lock(mutex_); + if (is_shutdown_) return; + is_shutdown_ = true; + retired.swap(retired_); + retained_until_shutdown.swap(retained_until_shutdown_); + } + + // The core communicator is shut down before this method is called. Query + // once to run any ready CPU/P2P post-completion callbacks, then release + // all retained resources. Captured GPU operations have no host callback + // and are kept until this point solely to cover graph replay. + for (auto& resources : retired) { + try { + (void)resources.ready_to_release(); + } catch (const std::exception& error) { + TORCH_WARN( + "MooncakeWorkTracker: retired work callback failed during " + "shutdown: ", + error.what()); + } catch (...) { + TORCH_WARN( + "MooncakeWorkTracker: retired work callback failed during " + "shutdown: unknown exception"); + } + } +} + +MooncakeWorkCpu::MooncakeWorkCpu(c10d::OpType opType, + mooncakePgCompletion_t completion, + FailedRanksHint failedRanksHint, + std::shared_ptr tracker, + std::vector keepAlive, + std::function postCompletion) + : Work(-1, opType), + completion_(makeCompletionHandle(completion)), + failed_ranks_hint_(std::move(failedRanksHint)), + tracker_(std::move(tracker)), + keep_alive_(std::move(keepAlive)), + post_completion_(makePostCompletionOnce(std::move(postCompletion))) {} + +MooncakeWorkCpu::~MooncakeWorkCpu() { + if (!tracker_) return; + auto completion = std::move(completion_); + auto post_completion = std::move(post_completion_); + tracker_->retire(makeTrackedResources(std::move(failed_ranks_hint_), + std::move(keep_alive_)), + [completion = std::move(completion), + post_completion = std::move(post_completion)]() mutable { + if (!queryCompletion(completion.get())) return false; + if (post_completion) post_completion(); + return true; + }); +} + +bool MooncakeWorkCpu::isCompleted() { + const bool completed = queryCompletion(completion_.get()); + if (completed && post_completion_) post_completion_(); + return completed; +} + +bool MooncakeWorkCpu::wait(std::chrono::milliseconds) { + // Preserve the existing CPU Work behavior: its timeout argument is + // ignored and wait blocks until the operation completes. + if (!waitCompletion(completion_.get(), -1)) return false; + if (post_completion_) post_completion_(); + return true; +} + +at::Tensor MooncakeWorkCpu::getFailedRanksHint() const { + return failed_ranks_hint_.tensor; +} + +bool MooncakeWorkCpu::getLocalSuccess() const { + return failed_ranks_hint_.isLocalSuccess(); +} + +MooncakeWorkCuda::MooncakeWorkCuda(c10d::OpType opType, + std::shared_ptr event, + FailedRanksHint failedRanksHint, + std::shared_ptr tracker, + std::vector keepAlive) + : Work(-1, opType), + event_(std::move(event)), + is_captured_(at::cuda::currentStreamCaptureStatus() != + c10::cuda::CaptureStatus::None), + failed_ranks_hint_(std::move(failedRanksHint)), + tracker_(std::move(tracker)), + keep_alive_(std::move(keepAlive)) { + TORCH_CHECK(event_, "Mooncake PG Torch event is null"); +} + +MooncakeWorkCuda::~MooncakeWorkCuda() { + if (!tracker_) return; + if (is_captured_) { + tracker_->retainUntilShutdown(makeTrackedResources( + std::move(event_), std::move(failed_ranks_hint_), + std::move(keep_alive_))); + return; + } + + auto event = std::move(event_); + tracker_->retire(makeTrackedResources(std::move(failed_ranks_hint_), + std::move(keep_alive_)), + [event = std::move(event)] { return event->query(); }); +} + +bool MooncakeWorkCuda::wait(std::chrono::milliseconds) { + // Once all tasks have been submitted, use the event to synchronize + // the current stream and the enqueue stream, but do not wait on this + // event. + // + // See PyTorch docs for more details: + // https://docs.pytorch.org/docs/stable/distributed.html#synchronous-and-asynchronous-collective-operations + // "wait() - in the case of CPU collectives, will block the process + // until the operation is completed. In the case of CUDA collectives, + // will block the currently active CUDA stream until the operation + // is completed (but will not block the CPU)." + auto current_stream = at::cuda::getCurrentCUDAStream(); + event_->block(current_stream); + return true; +} + +at::Tensor MooncakeWorkCuda::getFailedRanksHint() const { + // Ensure the worker thread has completed the task and written + // the failed-ranks bitmap before returning the tensor. + if (event_ && at::cuda::currentStreamCaptureStatus() == + c10::cuda::CaptureStatus::None) { + event_->synchronize(); + } + return failed_ranks_hint_.tensor; +} + +bool MooncakeWorkCuda::getLocalSuccess() const { + if (event_ && at::cuda::currentStreamCaptureStatus() == + c10::cuda::CaptureStatus::None) { + event_->synchronize(); + } + return failed_ranks_hint_.isLocalSuccess(); +} + +bool MooncakeBarrierWorkCuda::wait(std::chrono::milliseconds timeout) { + // Skip host-side synchronization during CUDA graph capture. + // cudaEventSynchronize is not permitted while a stream is capturing. + if (at::cuda::currentStreamCaptureStatus() != + c10::cuda::CaptureStatus::None) { + // We still need stream-level synchronization so that subsequent + // operations on the capture stream are ordered after the barrier + // task on the enqueue stream. + auto current_stream = at::cuda::getCurrentCUDAStream(); + event_->block(current_stream); + return true; + } + + if (timeout == kNoTimeout) { + event_->synchronize(); + return true; + } + + BackoffWaiter waiter( + BackoffWaiterConfig::constantSleep(std::chrono::microseconds(10))); + return waiter.wait_for(timeout, [this] { return event_->query(); }); +} + +MooncakeP2PWork::MooncakeP2PWork(c10d::OpType opType, + mooncakePgCompletion_t completion, + FailedRanksHint failedRanksHint, + std::shared_ptr tracker, + std::vector keepAlive, + std::function postCompletion) + : Work(-1, opType), + completion_(makeCompletionHandle(completion)), + failed_ranks_hint_(std::move(failedRanksHint)), + tracker_(std::move(tracker)), + keep_alive_(std::move(keepAlive)), + post_completion_(makePostCompletionOnce(std::move(postCompletion))) {} + +MooncakeP2PWork::~MooncakeP2PWork() { + if (!tracker_) return; + auto completion = std::move(completion_); + auto failed_ranks_hint = failed_ranks_hint_; + auto post_completion = std::move(post_completion_); + tracker_->retire( + makeTrackedResources(std::move(failed_ranks_hint_), + std::move(keep_alive_)), + [completion = std::move(completion), + failed_ranks_hint = std::move(failed_ranks_hint), + post_completion = std::move(post_completion)]() mutable { + if (!queryCompletion(completion.get())) return false; + if (failed_ranks_hint.isLocalSuccess() && post_completion) { + post_completion(); + } + return true; + }); +} + +bool MooncakeP2PWork::isCompleted() { + const bool completed = queryCompletion(completion_.get()); + if (completed && failed_ranks_hint_.isLocalSuccess() && post_completion_) { + post_completion_(); + } + return completed; +} + +bool MooncakeP2PWork::isSuccess() const { + return queryCompletion(completion_.get()) && + failed_ranks_hint_.isLocalSuccess(); +} + +bool MooncakeP2PWork::wait(std::chrono::milliseconds timeout) { + const int64_t timeout_us = + timeout.count() > 0 + ? std::chrono::duration_cast(timeout) + .count() + : -1; + if (!waitCompletion(completion_.get(), timeout_us)) return false; + if (failed_ranks_hint_.isLocalSuccess() && post_completion_) { + post_completion_(); + } + return true; +} + +at::Tensor MooncakeP2PWork::getFailedRanksHint() const { + return failed_ranks_hint_.tensor; +} + +bool MooncakeP2PWork::getLocalSuccess() const { return isSuccess(); } + +} // namespace mooncake diff --git a/mooncake-store/AGENTS.md b/mooncake-store/AGENTS.md new file mode 100644 index 0000000000..7e2480637e --- /dev/null +++ b/mooncake-store/AGENTS.md @@ -0,0 +1,13 @@ +# Mooncake Store Instructions + +## Random Numbers + +- Production code in `include/` and `src/` must use `include/random.h`; do not + create random engines, seeds, or distributions at call sites. +- Extend `random.h` and its tests when an operation is missing. Keep bounded + sampling unbiased and reject invalid bounds. +- Tests and benchmarks may use explicit seeds for reproducibility, but should + use the shared sampling helpers. +- Do not migrate `src/cachelib_memory_allocator/` unless explicitly requested. +- The shared engine is not cryptographically secure; do not use it for secrets + or authentication tokens. diff --git a/mooncake-store/CMakeLists.txt b/mooncake-store/CMakeLists.txt index 5cf0f8a4ea..c913e5d731 100644 --- a/mooncake-store/CMakeLists.txt +++ b/mooncake-store/CMakeLists.txt @@ -1,5 +1,8 @@ project(MooncakeStore VERSION 2.0.0) +option(ENABLE_KV_EVENTS + "Build master KV events ZMQ publisher (requires libzmq when ON)" OFF) + # Extract version components for C++ usage set(MOONCAKE_STORE_VERSION ${PROJECT_VERSION}) @@ -68,6 +71,10 @@ include_directories( # Add subdirectories add_subdirectory(src) +if (STORE_USE_ETCD) + add_subdirectory(tools) +endif() + if (BUILD_UNIT_TESTS) add_subdirectory(tests) endif() diff --git a/mooncake-store/benchmarks/CMakeLists.txt b/mooncake-store/benchmarks/CMakeLists.txt index 76d23cda09..c1e5e2b810 100644 --- a/mooncake-store/benchmarks/CMakeLists.txt +++ b/mooncake-store/benchmarks/CMakeLists.txt @@ -3,6 +3,12 @@ add_executable(allocator_bench allocator_bench.cpp) target_link_libraries(allocator_bench PRIVATE cachelib_memory_allocator mooncake_store) +# Add focused OffsetAllocator concurrency benchmark +add_executable(offset_allocator_concurrency_bench + offset_allocator_concurrency_bench.cpp) +target_link_libraries(offset_allocator_concurrency_bench + PRIVATE mooncake_store gflags::gflags pthread) + # Add master benchmark executable add_executable(master_bench master_bench.cpp) target_link_libraries(master_bench PRIVATE cachelib_memory_allocator @@ -39,3 +45,16 @@ if(USE_NOF) target_link_libraries(nof_worker_pool_bench PRIVATE mooncake_store glog::glog gflags::gflags) endif() + +# Add BatchEvict benchmark executable +add_executable(batch_evict_bench batch_evict_bench.cpp) +target_link_libraries( + batch_evict_bench PRIVATE mooncake_store cachelib_memory_allocator + gflags::gflags glog::glog pthread) + +if(STORE_USE_ETCD) + add_executable(oplog_batch_bench oplog_batch_bench.cpp) + target_link_libraries( + oplog_batch_bench PRIVATE mooncake_store glog::glog gflags::gflags + JsonCpp::JsonCpp) +endif() diff --git a/mooncake-store/benchmarks/README.md b/mooncake-store/benchmarks/README.md new file mode 100644 index 0000000000..e461836dc6 --- /dev/null +++ b/mooncake-store/benchmarks/README.md @@ -0,0 +1,85 @@ +# Mooncake Store Benchmarks + +This directory contains benchmark tools for Mooncake Store internals. + +## Allocation Strategy Benchmark + +`allocation_strategy_bench` evaluates Store allocation behavior across segment +counts, replica counts, allocation strategies, and workload patterns. + +Build the benchmark from an existing CMake build directory: + +```bash +cmake --build build --target allocation_strategy_bench -j$(nproc) +``` + +### Size-Class Churn Fragmentation Benchmark + +The `size_class_churn` workload measures fragmentation under mixed-size +KVCache-like allocation pressure. It pre-fills the simulated cluster when +`--prefill_pct` is set, then repeatedly allocates objects from weighted size +classes. On allocation failure it randomly evicts a fraction of live objects and +retries. + +When prefill is enabled, the prefill attempt cap is auto-derived from target +utilization, total cluster capacity, weighted average object size, and replica +count, with a 5000-attempt minimum for small cases. + +This is an allocation-strategy-layer benchmark. It complements the existing +`dsa` workload by adding explicit fragmentation sampling and configurable +weighted size-class patterns. It is not a replacement for `allocator_bench`, +which remains the low-level `OffsetAllocator` microbenchmark. + +Run a small local validation: + +```bash +./build/mooncake-store/benchmarks/allocation_strategy_bench \ + --workload=size_class_churn \ + --segment_capacity=1024 \ + --num_allocations=10000 \ + --prefill_pct=70 +``` + +Run a larger baseline: + +```bash +./build/mooncake-store/benchmarks/allocation_strategy_bench \ + --workload=size_class_churn \ + --segment_capacity=1024 \ + --num_allocations=100000 \ + --prefill_pct=80 +``` + +Supported size-class patterns: + +- `kv_mixed`: 4KB at 70%, 256KB at 20%, and 3.12MB at 10%. +- `dsa_pair`: 3.12MB KV pages at 50% and 643KB indexer entries at 50%. +- `all`: run both patterns. + +Key output columns: + +- `Throughput`, `Avg(ns)`, `P50(ns)`, `P90(ns)`, and `P99(ns)` measure + allocation performance. +- `Frag_avg`, `Frag_p50`, `Frag_p90`, and `Frag_p99` summarize sampled + fragmentation ratios. +- `LargestFreeMB` shows the final largest contiguous free region. +- `Evictions` counts fail-triggered eviction rounds during measurement. +- `Full/Partial/Fail/Total` reports allocation outcomes. Only results with + `result->size() == replica_num` count as full success; shorter replica + results are counted as partial allocations. + +Fragmentation is computed per `OffsetBufferAllocator` and then averaged by free +space: + +```text +1 - largest_free_region / total_free_space +``` + +The weighted average avoids treating free space in different Store segments as +one mergeable region. `LargestFreeMB` still reports the final largest contiguous +free region across all segments. + +The benchmark also prints a one-line `Prefill summary`, `Fragmentation summary`, +and `Size-class breakdown` after each result row, so reviewers can read the +actual prefill utilization, fragmentation, and per-size-class latency numbers +without manually deriving them from the table. diff --git a/mooncake-store/benchmarks/allocation_strategy_bench.cpp b/mooncake-store/benchmarks/allocation_strategy_bench.cpp index 6e9155655e..adf8facb6f 100644 --- a/mooncake-store/benchmarks/allocation_strategy_bench.cpp +++ b/mooncake-store/benchmarks/allocation_strategy_bench.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -14,7 +15,7 @@ #include #include "types.h" -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "allocator.h" #include "allocation_strategy.h" @@ -23,14 +24,14 @@ DEFINE_int64(segment_capacity, 1024, "Per-segment capacity in MB (base capacity for skewed mode)"); DEFINE_int32(num_allocations, 10000, "Number of allocations to benchmark"); DEFINE_int32(convergence_sample_interval, 100, - "Sample utilization stddev every N allocations (scaleout only)"); -DEFINE_bool( - run_all, false, - "Also run the Scale-Out matrix in addition to the default Fillup matrix"); + "Sample utilization stddev or fragmentation every N allocations"); +DEFINE_bool(run_all, false, + "Run all workload matrices instead of only the selected workload"); // Scale-Out workload flags DEFINE_string(workload, "fillup", - "Workload type: fillup (default), scaleout, dsa"); + "Workload type: fillup (default), scaleout, dsa, " + "size_class_churn"); DEFINE_int32( scale_out_trigger_pct, 50, "Pre-fill cluster to this utilization % before injecting new nodes " @@ -73,6 +74,11 @@ DEFINE_double(dsa_evict_ratio, 0.05, "(default 0.05 = 5%). Lower values keep " "the steady-state cluster fill closer to the fragmentation " "ceiling but trigger evictions more frequently."); +DEFINE_string(size_class_pattern, "kv_mixed", + "Size-class churn pattern: kv_mixed, dsa_pair, or all"); +DEFINE_double(size_class_evict_ratio, 0.02, + "Fraction of live objects to evict on each Allocate failure " + "in size_class_churn mode."); using namespace mooncake; @@ -99,11 +105,18 @@ constexpr int kDsaMaxBatch = 128; // per-round batch upper bound constexpr int kDsaMaxRetries = 5; // Lower bound on the auto-derived allocation count per DSA case. constexpr int kDsaMinAllocs = 100; +constexpr int kSizeClassMaxRetries = 5; +constexpr size_t kSizeClassMinPrefillAttempts = 5000; +// The theoretical prefill budget assumes every allocation succeeds. Use a +// small multiplier to absorb partial allocations, retries, and fragmentation +// without turning unreachable targets into long-running cases. +constexpr double kSizeClassPrefillAttemptMultiplier = 2.0; enum class WorkloadType { FILL_UP, // Only allocate, measure throughput/latency SCALE_OUT, // Inject new nodes mid-run, measure adoption speed DSA, // DSA paired KV+indexer with random fail-triggered eviction + SIZE_CLASS_CHURN // Mixed size classes with steady-state churn }; struct BenchConfig { @@ -127,6 +140,10 @@ struct BenchConfig { // DSA workload knobs (only used when workload_type == DSA). bool dsa_paired = false; // false = KV-only, true = KV+indexer pair + + // Size-class churn knobs (only used when workload_type == + // SIZE_CLASS_CHURN). + std::string size_class_pattern = "kv_mixed"; }; struct UtilRatioStats { @@ -139,6 +156,63 @@ struct UtilRatioStats { bool valid = false; }; +struct DistributionStats { + double min = 0.0; + double p50 = 0.0; + double p90 = 0.0; + double p99 = 0.0; + double max = 0.0; + double avg = 0.0; + bool valid = false; +}; + +struct FragmentationSnapshot { + uint64_t total_free_space = 0; + uint64_t largest_free_region = 0; + uint64_t capacity = 0; + double fragmentation_ratio = 0.0; + bool valid = false; +}; + +struct SizeClassSpec { + std::string name; + size_t size; + int weight; +}; + +struct SizeClassStat { + std::string name; + size_t size = 0; + int weight = 0; + int success_count = 0; + int partial_count = 0; + int failed_count = 0; + int total_count = 0; + DistributionStats latency_stats; +}; + +enum class SizeClassAllocationStatus { + FAILED, + PARTIAL, + FULL, +}; + +struct SizeClassAllocationResult { + SizeClassAllocationStatus status = SizeClassAllocationStatus::FAILED; + size_t replica_count = 0; +}; + +struct SizeClassPrefillStats { + size_t attempts = 0; + size_t max_attempts = 0; + int full_count = 0; + int partial_count = 0; + int failed_count = 0; + int requested_pct = 0; + double achieved_util_pct = 0.0; + bool reached_target = false; +}; + /** * @brief Common base for all benchmark results. * @@ -162,6 +236,8 @@ struct BenchResultBase { int success_count = 0; // successful allocations int total_count = 0; // total attempted allocations + int partial_count = 0; // partial allocations, if reported + int failed_count = 0; // failed allocations, if reported double final_util_stddev; // utilization stddev at run end double final_avg_util; // average utilization at run end @@ -199,6 +275,15 @@ struct ScaleOutResult : BenchResultBase { std::vector new_node_util_over_time; }; +struct SizeClassChurnResult : BenchResultBase { + std::string pattern_name; + SizeClassPrefillStats prefill_stats; + DistributionStats fragmentation_stats; + DistributionStats largest_free_mb_stats; + FragmentationSnapshot final_fragmentation; + std::vector size_class_stats; +}; + static double computeClusterCapacityGB(int num_segments, size_t base_capacity, bool skewed) { double total = 0.0; @@ -391,6 +476,151 @@ static UtilRatioStats computeUtilRatioStats(std::vector& util_ratios) { return stats; } +static DistributionStats computeDistributionStats(std::vector& values) { + DistributionStats stats; + if (values.empty()) return stats; + + std::sort(values.begin(), values.end()); + + auto percentile = [&](double p) -> double { + size_t idx = static_cast(std::round(p * (values.size() - 1))); + return values[idx]; + }; + + stats.min = values.front(); + stats.p50 = percentile(0.50); + stats.p90 = percentile(0.90); + stats.p99 = percentile(0.99); + stats.max = values.back(); + stats.avg = + std::accumulate(values.begin(), values.end(), 0.0) / values.size(); + stats.valid = true; + return stats; +} + +static FragmentationSnapshot computeFragmentationSnapshot( + const AllocatorManager& manager) { + FragmentationSnapshot snapshot; + double weighted_fragmentation = 0.0; + + for (const auto& name : manager.getNames()) { + const auto* allocs = manager.getAllocators(name); + if (!allocs) continue; + + for (const auto& alloc : *allocs) { + auto offset_alloc = + std::dynamic_pointer_cast(alloc); + if (!offset_alloc) continue; + + auto allocator = offset_alloc->getOffsetAllocator(); + if (!allocator) continue; + + auto metrics = allocator->get_metrics(); + snapshot.total_free_space += metrics.total_free_space_; + snapshot.capacity += metrics.capacity; + snapshot.largest_free_region = std::max( + snapshot.largest_free_region, metrics.largest_free_region_); + if (metrics.total_free_space_ > 0) { + double local_fragmentation = + 1.0 - (static_cast(metrics.largest_free_region_) / + static_cast(metrics.total_free_space_)); + local_fragmentation = std::clamp(local_fragmentation, 0.0, 1.0); + weighted_fragmentation += + local_fragmentation * metrics.total_free_space_; + } + snapshot.valid = true; + } + } + + if (snapshot.valid && snapshot.total_free_space > 0) { + snapshot.fragmentation_ratio = + weighted_fragmentation / + static_cast(snapshot.total_free_space); + } + + return snapshot; +} + +static std::vector getSizeClassSpecs( + const std::string& pattern_name) { + if (pattern_name == "kv_mixed") { + return { + {"small", 4 * KiB, 70}, + {"medium", 256 * KiB, 20}, + {"large", kDsaKvSize, 10}, + }; + } + + if (pattern_name == "dsa_pair") { + return { + {"kv", kDsaKvSize, 50}, + {"indexer", kDsaIndexerSize, 50}, + }; + } + + return {}; +} + +static size_t chooseSizeClassIndex(const std::vector& specs, + std::mt19937& rng) { + if (specs.empty()) return 0; + + int total_weight = 0; + for (const auto& spec : specs) { + total_weight += spec.weight; + } + + if (total_weight <= 0) return 0; + + std::uniform_int_distribution dist(1, total_weight); + int pick = dist(rng); + for (size_t i = 0; i < specs.size(); ++i) { + pick -= specs[i].weight; + if (pick <= 0) return i; + } + + return specs.size() - 1; +} + +static double computeWeightedAverageObjectSize( + const std::vector& specs) { + double weighted_size = 0.0; + int total_weight = 0; + + for (const auto& spec : specs) { + if (spec.weight <= 0) continue; + weighted_size += static_cast(spec.size) * spec.weight; + total_weight += spec.weight; + } + + if (total_weight <= 0) return 0.0; + return static_cast(weighted_size / total_weight); +} + +static size_t deriveSizeClassPrefillMaxAttempts( + const AllocatorManager& manager, const BenchConfig& cfg, + const std::vector& specs) { + if (cfg.prefill_pct <= 0 || cfg.replica_num <= 0) return 0; + + const double avg_object_size = computeWeightedAverageObjectSize(specs); + if (avg_object_size <= 0.0) return kSizeClassMinPrefillAttempts; + + const double target_bytes = + static_cast(computeTotalCapacity(manager)) * cfg.prefill_pct / + 100.0; + const double bytes_per_attempt = + avg_object_size * static_cast(cfg.replica_num); + if (target_bytes <= 0.0 || bytes_per_attempt <= 0.0) { + return kSizeClassMinPrefillAttempts; + } + + const double derived_attempts = + std::ceil((target_bytes / bytes_per_attempt) * + kSizeClassPrefillAttemptMultiplier); + return std::max(kSizeClassMinPrefillAttempts, + static_cast(derived_attempts)); +} + static std::string strategyName(AllocationStrategyType type) { switch (type) { case AllocationStrategyType::RANDOM: @@ -798,6 +1028,87 @@ static bool dsaAllocateWithEvict( return false; } +static SizeClassAllocationResult sizeClassAllocateWithEvict( + const std::shared_ptr& strategy, + AllocatorManager& manager, size_t size, int replica_num, + std::vector>& live, std::mt19937& rng, + int& evict_count, double evict_ratio) { + for (int attempt = 0; attempt <= kSizeClassMaxRetries; ++attempt) { + auto result = strategy->Allocate(manager, size, replica_num); + if (result.has_value()) { + size_t replica_count = result->size(); + if (replica_count == 0) { + return {}; + } + live.push_back(std::move(result.value())); + return { + replica_count == static_cast(replica_num) + ? SizeClassAllocationStatus::FULL + : SizeClassAllocationStatus::PARTIAL, + replica_count, + }; + } + + if (live.empty()) return {}; + if (attempt == kSizeClassMaxRetries) return {}; + + evictRandomFraction(live, evict_ratio, rng); + ++evict_count; + } + + return {}; +} + +static SizeClassPrefillStats prefillSizeClassChurn( + const std::shared_ptr& strategy, + AllocatorManager& manager, const BenchConfig& cfg, + const std::vector& specs, + std::vector>& live_allocations, std::mt19937& rng) { + SizeClassPrefillStats stats; + if (cfg.prefill_pct <= 0 || specs.empty()) return stats; + stats.requested_pct = cfg.prefill_pct; + stats.max_attempts = deriveSizeClassPrefillMaxAttempts(manager, cfg, specs); + + int consec_failures = 0; + int evict_throwaway = 0; + const int kMaxConsecFailures = 10; + + while (stats.attempts < stats.max_attempts) { + if (stats.attempts % kPreFillSampleInterval == 0) { + stats.achieved_util_pct = computeAverageUtilAll(manager) * 100.0; + if (stats.achieved_util_pct >= cfg.prefill_pct) { + stats.reached_target = true; + break; + } + } + + size_t class_idx = chooseSizeClassIndex(specs, rng); + auto alloc_result = sizeClassAllocateWithEvict( + strategy, manager, specs[class_idx].size, cfg.replica_num, + live_allocations, rng, evict_throwaway, + FLAGS_size_class_evict_ratio); + ++stats.attempts; + if (alloc_result.status == SizeClassAllocationStatus::FULL) { + ++stats.full_count; + consec_failures = 0; + } else if (alloc_result.status == SizeClassAllocationStatus::PARTIAL) { + ++stats.partial_count; + consec_failures = 0; + } else { + ++stats.failed_count; + if (++consec_failures >= kMaxConsecFailures) { + break; + } + } + } + + stats.achieved_util_pct = computeAverageUtilAll(manager) * 100.0; + if (stats.achieved_util_pct >= cfg.prefill_pct) { + stats.reached_target = true; + } + return stats; +} + // Run DSA workload; the allocation count is derived from cluster capacity. static FillUpResult runDsaBenchmark(const BenchConfig& cfg) { AllocatorManager manager = @@ -922,6 +1233,127 @@ static FillUpResult runDsaBenchmark(const BenchConfig& cfg) { return res; } +static SizeClassChurnResult runSizeClassChurnBenchmark(const BenchConfig& cfg) { + AllocatorManager manager = + createCluster(cfg.num_segments, cfg.segment_capacity, cfg.skewed); + auto strategy = CreateAllocationStrategy(cfg.strategy_type); + auto specs = getSizeClassSpecs(cfg.size_class_pattern); + + std::vector per_class_stats; + per_class_stats.reserve(specs.size()); + std::vector> per_class_latencies(specs.size()); + for (const auto& spec : specs) { + SizeClassStat stat; + stat.name = spec.name; + stat.size = spec.size; + stat.weight = spec.weight; + per_class_stats.push_back(std::move(stat)); + } + + std::vector latencies; + latencies.reserve(cfg.num_allocations); + int sample_interval = std::max(1, FLAGS_convergence_sample_interval); + std::vector fragmentation_samples; + fragmentation_samples.reserve(cfg.num_allocations / sample_interval + 2); + std::vector largest_free_mb_samples; + largest_free_mb_samples.reserve(fragmentation_samples.capacity()); + + std::vector> live_allocations; + live_allocations.reserve(std::min(cfg.num_allocations, 1 << 20)); + + std::mt19937 rng(42); + SizeClassPrefillStats prefill_stats = prefillSizeClassChurn( + strategy, manager, cfg, specs, live_allocations, rng); + + int success_count = 0; + int partial_count = 0; + int failed_count = 0; + int total_count = 0; + int evict_count = 0; + double instrumentation_time_us = 0.0; + + auto total_start = std::chrono::high_resolution_clock::now(); + + for (int i = 0; i < cfg.num_allocations; ++i) { + size_t class_idx = chooseSizeClassIndex(specs, rng); + const auto& spec = specs[class_idx]; + + auto t0 = std::chrono::high_resolution_clock::now(); + auto alloc_result = sizeClassAllocateWithEvict( + strategy, manager, spec.size, cfg.replica_num, live_allocations, + rng, evict_count, FLAGS_size_class_evict_ratio); + auto t1 = std::chrono::high_resolution_clock::now(); + + double latency_ns = + std::chrono::duration(t1 - t0).count(); + latencies.push_back(latency_ns); + per_class_latencies[class_idx].push_back(latency_ns); + + ++total_count; + ++per_class_stats[class_idx].total_count; + if (alloc_result.status == SizeClassAllocationStatus::FULL) { + ++success_count; + ++per_class_stats[class_idx].success_count; + } else if (alloc_result.status == SizeClassAllocationStatus::PARTIAL) { + ++partial_count; + ++per_class_stats[class_idx].partial_count; + } else { + ++failed_count; + ++per_class_stats[class_idx].failed_count; + } + + if ((i + 1) % sample_interval == 0 || i == cfg.num_allocations - 1) { + auto s0 = std::chrono::high_resolution_clock::now(); + auto snapshot = computeFragmentationSnapshot(manager); + if (snapshot.valid) { + fragmentation_samples.push_back(snapshot.fragmentation_ratio); + largest_free_mb_samples.push_back( + static_cast(snapshot.largest_free_region) / MiB); + } + auto s1 = std::chrono::high_resolution_clock::now(); + instrumentation_time_us += + std::chrono::duration(s1 - s0).count(); + } + } + + auto total_end = std::chrono::high_resolution_clock::now(); + double total_us = + std::chrono::duration(total_end - total_start) + .count(); + total_us -= instrumentation_time_us; + total_us = std::max(total_us, 1.0); + + for (size_t i = 0; i < per_class_stats.size(); ++i) { + per_class_stats[i].latency_stats = + computeDistributionStats(per_class_latencies[i]); + } + + SizeClassChurnResult res; + res.strategy_name = cfg.strategy_name; + res.num_segments = cfg.num_segments; + res.alloc_size = 0; + res.replica_num = cfg.replica_num; + res.skewed = cfg.skewed; + res.cluster_capacity_gb = computeClusterCapacityGB( + cfg.num_segments, cfg.segment_capacity, cfg.skewed); + res.final_util_stddev = computeUtilizationStdDev(manager); + res.final_avg_util = computeAverageUtilAll(manager); + res.success_count = success_count; + res.partial_count = partial_count; + res.failed_count = failed_count; + res.total_count = total_count; + res.evict_count = evict_count; + res.pattern_name = cfg.size_class_pattern; + res.prefill_stats = prefill_stats; + res.fragmentation_stats = computeDistributionStats(fragmentation_samples); + res.largest_free_mb_stats = + computeDistributionStats(largest_free_mb_samples); + res.final_fragmentation = computeFragmentationSnapshot(manager); + res.size_class_stats = std::move(per_class_stats); + computeLatencyStats(latencies, total_us, total_count, res); + return res; +} + static void printFillUpHeader() { std::cout << std::string(184, '-') << std::endl; std::cout << std::left << std::setw(18) << "Strategy" << std::setw(9) @@ -1050,6 +1482,93 @@ static void printScaleOutResult(const ScaleOutResult& r) { << std::endl; } +static void printSizeClassChurnHeader() { + std::cout << std::string(260, '-') << std::endl; + std::cout << std::left << std::setw(18) << "Strategy" << std::setw(9) + << "Replica" << std::setw(10) << "Segments" << std::setw(14) + << "Pattern" << std::setw(12) << "Cluster(GB)" << std::setw(8) + << "Skewed" << std::right << std::setw(14) << "Throughput" + << std::setw(12) << "Avg(ns)" << std::setw(12) << "P50(ns)" + << std::setw(12) << "P90(ns)" << std::setw(12) << "P99(ns)" + << std::setw(12) << "Frag_avg" << std::setw(12) << "Frag_p50" + << std::setw(12) << "Frag_p90" << std::setw(12) << "Frag_p99" + << std::setw(15) << "LargestFreeMB" << std::setw(10) << "AvgUtil%" + << std::setw(24) << "Full/Partial/Fail/Total" << std::setw(14) + << "Evictions" << std::endl; + std::cout << std::string(260, '-') << std::endl; +} + +static void printSizeClassChurnResult(const SizeClassChurnResult& r) { + std::string alloc_ratio = std::to_string(r.success_count) + "/" + + std::to_string(r.partial_count) + "/" + + std::to_string(r.failed_count) + "/" + + std::to_string(r.total_count); + std::ostringstream cap_ss; + cap_ss << std::fixed << std::setprecision(1) << r.cluster_capacity_gb; + + double final_largest_free_mb = + static_cast(r.final_fragmentation.largest_free_region) / MiB; + + std::cout << std::left << std::setw(18) << r.strategy_name << std::setw(9) + << r.replica_num << std::setw(10) << r.num_segments + << std::setw(14) << r.pattern_name << std::setw(12) + << cap_ss.str() << std::setw(8) << (r.skewed ? "yes" : "no") + << std::right << std::fixed << std::setprecision(0) + << std::setw(14) << r.throughput << std::setw(12) << r.avg_ns + << std::setw(12) << r.p50_ns << std::setw(12) << r.p90_ns + << std::setw(12) << r.p99_ns << std::setprecision(4) + << std::setw(12) << r.fragmentation_stats.avg << std::setw(12) + << r.fragmentation_stats.p50 << std::setw(12) + << r.fragmentation_stats.p90 << std::setw(12) + << r.fragmentation_stats.p99 << std::setprecision(1) + << std::setw(15) << final_largest_free_mb << std::setprecision(2) + << std::setw(9) << (r.final_avg_util * 100.0) << "%" + << std::setw(24) << alloc_ratio << std::setw(14) << r.evict_count + << std::endl; + + std::cout << "Prefill summary [" << r.strategy_name + << ", pattern=" << r.pattern_name + << ", segments=" << r.num_segments + << ", replica=" << r.replica_num + << "]: requested_pct=" << std::fixed << std::setprecision(2) + << r.prefill_stats.requested_pct + << ", achieved_pct=" << r.prefill_stats.achieved_util_pct + << ", reached=" << (r.prefill_stats.reached_target ? "yes" : "no") + << ", attempts=" << r.prefill_stats.attempts << "/" + << r.prefill_stats.max_attempts + << ", full/partial/failed=" << r.prefill_stats.full_count << "/" + << r.prefill_stats.partial_count << "/" + << r.prefill_stats.failed_count << std::endl; + + std::cout << "Fragmentation summary [" << r.strategy_name + << ", pattern=" << r.pattern_name + << ", segments=" << r.num_segments + << ", replica=" << r.replica_num + << ", skewed=" << (r.skewed ? "yes" : "no") + << "]: avg=" << std::fixed << std::setprecision(4) + << r.fragmentation_stats.avg + << ", p50=" << r.fragmentation_stats.p50 + << ", p90=" << r.fragmentation_stats.p90 + << ", p99=" << r.fragmentation_stats.p99 + << ", max=" << r.fragmentation_stats.max + << ", final_largest_free=" << std::setprecision(1) + << final_largest_free_mb << " MB" << std::endl; + + std::cout << "Size-class breakdown:"; + for (const auto& stat : r.size_class_stats) { + std::string ratio = std::to_string(stat.success_count) + "/" + + std::to_string(stat.partial_count) + "/" + + std::to_string(stat.failed_count) + "/" + + std::to_string(stat.total_count); + std::cout << " " << stat.name << "(" << (stat.size / KiB) + << "KB,w=" << stat.weight + << ",full/partial/failed/total=" << ratio + << ",p99_ns=" << std::fixed << std::setprecision(0) + << stat.latency_stats.p99 << ")"; + } + std::cout << std::endl; +} + static void runFillupBenchmarks() { std::vector skewed_options = {false, true}; std::vector segment_counts = {1, 10, 100, 512, 1024}; @@ -1278,6 +1797,107 @@ static void runDsaMatrix() { } } +static void runSizeClassChurnMatrix() { + std::vector skewed_options = {false, true}; + std::vector segment_counts = {1, 10, 100}; + std::vector replica_nums = {1, 2, 3}; + std::vector strategies = { + AllocationStrategyType::RANDOM, + AllocationStrategyType::FREE_RATIO_FIRST, + }; + + std::vector patterns; + if (FLAGS_size_class_pattern == "all") { + patterns = {"kv_mixed", "dsa_pair"}; + } else { + patterns = {FLAGS_size_class_pattern}; + } + + if (FLAGS_size_class_evict_ratio <= 0.0 || + FLAGS_size_class_evict_ratio > 1.0) { + std::cout << "Invalid size_class_evict_ratio: " + << FLAGS_size_class_evict_ratio + << ". Use a value in the range (0.0, 1.0]." << std::endl; + return; + } + + for (const auto& pattern : patterns) { + if (getSizeClassSpecs(pattern).empty()) { + std::cout << "Invalid size_class_pattern: " << pattern + << ". Use --size_class_pattern=kv_mixed, dsa_pair, or " + "all." + << std::endl; + return; + } + } + + std::cout << "\n=== Size-Class Churn Fragmentation Benchmark Matrix ===\n" + << "Workload: prefill to --prefill_pct if set, then run " + << FLAGS_num_allocations + << " mixed-size allocation attempts with fail-triggered random " + "eviction and retry.\n" + << "Fragmentation: 1 - largest_free_region / total_free_space, " + "sampled every --convergence_sample_interval allocations.\n" + << "Config: segment_capacity=" << FLAGS_segment_capacity + << " MB, prefill_pct=" << FLAGS_prefill_pct + << ", evict_ratio=" << FLAGS_size_class_evict_ratio + << ", size_class_pattern=" << FLAGS_size_class_pattern << "\n" + << "Patterns: kv_mixed = 4KB:70%, 256KB:20%, 3198KB:10%; " + "dsa_pair = 3198KB:50%, 643KB:50%.\n" + << "Skewed setup: half nodes are (base + 50%) capacity, half are " + "(base - 50%)\n" + << std::endl; + + std::vector configs; + for (const auto& pattern : patterns) { + for (auto skew : skewed_options) { + for (auto strategy : strategies) { + for (auto segs : segment_counts) { + for (auto rep : replica_nums) { + if (rep > segs) continue; + BenchConfig cfg; + cfg.num_segments = segs; + cfg.segment_capacity = + static_cast(FLAGS_segment_capacity) * MiB; + cfg.alloc_size = 0; + cfg.replica_num = rep; + cfg.num_allocations = FLAGS_num_allocations; + cfg.skewed = skew; + cfg.strategy_type = strategy; + cfg.strategy_name = strategyName(strategy); + cfg.prefill_pct = FLAGS_prefill_pct; + cfg.workload_type = WorkloadType::SIZE_CLASS_CHURN; + cfg.size_class_pattern = pattern; + configs.push_back(cfg); + } + } + } + } + } + + bool first = true; + std::string prev_pattern; + AllocationStrategyType prev_strategy = AllocationStrategyType::RANDOM; + + for (const auto& cfg : configs) { + if (first || cfg.size_class_pattern != prev_pattern) { + std::cout << "\n--- Pattern: " << cfg.size_class_pattern << " ---" + << std::endl; + prev_pattern = cfg.size_class_pattern; + first = true; + } + + if (first || cfg.strategy_type != prev_strategy) { + printSizeClassChurnHeader(); + prev_strategy = cfg.strategy_type; + first = false; + } + + auto result = runSizeClassChurnBenchmark(cfg); + printSizeClassChurnResult(result); + } +} + int main(int argc, char* argv[]) { gflags::SetUsageMessage( "AllocationStrategy performance benchmark.\n" @@ -1289,16 +1909,19 @@ int main(int argc, char* argv[]) { runFillupBenchmarks(); runScaleOutMatrix(); runDsaMatrix(); + runSizeClassChurnMatrix(); } else if (FLAGS_workload == "fillup") { runFillupBenchmarks(); } else if (FLAGS_workload == "scaleout") { runScaleOutMatrix(); } else if (FLAGS_workload == "dsa") { runDsaMatrix(); + } else if (FLAGS_workload == "size_class_churn") { + runSizeClassChurnMatrix(); } else { std::cout << "Invalid workload type: " << FLAGS_workload - << ". Use --workload=fillup, --workload=scaleout, or " - "--workload=dsa." + << ". Use --workload=fillup, --workload=scaleout, " + "--workload=dsa, or --workload=size_class_churn." << std::endl; } diff --git a/mooncake-store/benchmarks/allocator_bench.cpp b/mooncake-store/benchmarks/allocator_bench.cpp index 8c4a8d908c..17cd58b2a7 100644 --- a/mooncake-store/benchmarks/allocator_bench.cpp +++ b/mooncake-store/benchmarks/allocator_bench.cpp @@ -6,7 +6,7 @@ #include #include -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" using namespace mooncake::offset_allocator; diff --git a/mooncake-store/benchmarks/batch_evict_bench.cpp b/mooncake-store/benchmarks/batch_evict_bench.cpp new file mode 100644 index 0000000000..dd44a22c61 --- /dev/null +++ b/mooncake-store/benchmarks/batch_evict_bench.cpp @@ -0,0 +1,462 @@ +#include "master_service.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gflags/gflags.h" +#include "glog/logging.h" +#include "types.h" + +DEFINE_uint64( + num_objects, 0, + "Run a single custom scale with this object count (0 = default scales)"); +DEFINE_double(evict_ratio_target, 0.50, "BatchEvict target eviction ratio"); +DEFINE_double(evict_ratio_lowerbound, 0.25, + "BatchEvict lower-bound eviction ratio"); + +namespace mooncake::benchmarks { + +class BatchEvictBench { + public: + static bool RunRealBatchEvictScales() { + std::vector scales = {10000, 100000}; + const char* large_mode = std::getenv("MOONCAKE_EVICT_BENCH_LARGE"); + if (large_mode != nullptr && std::string(large_mode) == "1") { + scales.push_back(1000000); + } + if (FLAGS_num_objects > 0) { + scales = {static_cast(FLAGS_num_objects)}; + } + + std::cout << "num_objects,total_us,objects_before,objects_after," + "evicted_count,freed_bytes" + << std::endl; + for (size_t scale : scales) { + if (!RunOneScale(scale)) { + return false; + } + } + return true; + } + + static bool RunSingleWaiterSnapshotMutexProbe() { + if (std::getenv("MOONCAKE_EVICT_BENCH_LOCK_PROBE") == nullptr) { + return true; + } + + const size_t num_objects = + ReadEnvSize("MOONCAKE_EVICT_BENCH_LOCK_OBJECTS", 1000000); + const size_t trials = std::max( + 1, ReadEnvSize("MOONCAKE_EVICT_BENCH_LOCK_TRIALS", 30)); + const auto waiter_delay = std::chrono::microseconds( + static_cast( + ReadEnvSize("MOONCAKE_EVICT_BENCH_LOCK_DELAY_US", 1000))); + + std::vector batch_evict_total_us; + std::vector unique_lock_wait_us; + + batch_evict_total_us.reserve(trials); + unique_lock_wait_us.reserve(trials); + + for (size_t trial = 0; trial < trials; ++trial) { + LockProbeTrialResult result; + if (!RunLockProbeTrial(num_objects, waiter_delay, result)) { + return false; + } + + batch_evict_total_us.push_back(result.batch_evict_total_us); + unique_lock_wait_us.push_back(result.unique_lock_wait_us); + } + + std::cout << "num_objects,trials,batch_evict_total_p50_us," + "unique_lock_wait_p50_us,unique_lock_wait_p95_us," + "unique_lock_wait_max_us" + << std::endl; + + std::cout << num_objects << "," << trials << "," + << PercentileValue(batch_evict_total_us, 0.50) << "," + << PercentileValue(unique_lock_wait_us, 0.50) << "," + << PercentileValue(unique_lock_wait_us, 0.95) << "," + << PercentileValue(unique_lock_wait_us, 1.00) << std::endl; + return true; + } + + private: + static constexpr const char* kSegmentName = "batch_evict_bench_segment"; + static constexpr size_t kSegmentBase = 0x300000000; + static constexpr uint64_t kObjectSize = 1024; + + struct MetadataStats { + size_t object_count{0}; + size_t completed_memory_replicas{0}; + size_t busy_memory_replicas{0}; + size_t non_memory_replicas{0}; + size_t incomplete_replicas{0}; + size_t unexpired_leases{0}; + }; + + struct LockProbeTrialResult { + uint64_t batch_evict_total_us{0}; + uint64_t unique_lock_wait_us{0}; + }; + + static MasterServiceConfig MakeConfig() { + return MasterServiceConfig::builder() + .set_memory_allocator(BufferAllocatorType::OFFSET) + .set_eviction_ratio(0.0) + .set_eviction_high_watermark_ratio(1.0) + .set_client_live_ttl_sec(3600) + .build(); + } + + static size_t SegmentSizeFor(size_t num_objects) { + constexpr size_t kMinSegmentSize = 16 * 1024 * 1024; + const size_t needed = num_objects * kObjectSize; + const size_t headroom = needed / 8 + 1024 * kObjectSize; + return std::max(kMinSegmentSize, needed + headroom); + } + + static Segment MakeSegment(size_t num_objects) { + Segment segment; + segment.id = generate_uuid(); + segment.name = kSegmentName; + segment.base = kSegmentBase; + segment.size = SegmentSizeFor(num_objects); + segment.te_endpoint = segment.name; + return segment; + } + + static std::string MakeKey(size_t index) { + return "batch_evict_bench_key_" + std::to_string(index); + } + + static size_t ReadEnvSize(const char* name, size_t default_value) { + const char* value = std::getenv(name); + if (value == nullptr || value[0] == '\0') { + return default_value; + } + char* end = nullptr; + const unsigned long long parsed = std::strtoull(value, &end, 10); + if (end == value) { + return default_value; + } + return static_cast(parsed); + } + + static uint64_t PercentileValue(std::vector values, + double percentile) { + if (values.empty()) { + return 0; + } + std::sort(values.begin(), values.end()); + const size_t rank = std::max( + 1, static_cast(std::ceil(percentile * values.size()))); + return values[std::min(rank - 1, values.size() - 1)]; + } + + static MetadataStats ExpireLeasesAndCollectStats(MasterService& service) { + MetadataStats stats; + auto now = std::chrono::system_clock::now(); + const auto base_expiration = now - std::chrono::hours(1); + size_t ordinal = 0; + + for (size_t shard_idx = 0; shard_idx < MasterService::kNumShards; + ++shard_idx) { + MasterService::MetadataShardAccessorRW shard(&service, shard_idx); + for (auto& [tenant_id, tenant_state] : shard->tenants) { + if (tenant_id != TenantId::Default()) { + continue; + } + for (auto& [key, metadata] : tenant_state.metadata) { + { + SpinLocker locker(&metadata.lock); + metadata.lease_timeout = + base_expiration + + std::chrono::nanoseconds(ordinal++); + } + + ++stats.object_count; + if (!metadata.IsLeaseExpired(now)) { + ++stats.unexpired_leases; + } + for (const auto& replica : metadata.GetAllReplicas()) { + if (replica.is_memory_replica()) { + if (replica.is_completed()) { + ++stats.completed_memory_replicas; + } else { + ++stats.incomplete_replicas; + } + if (replica.get_refcnt() != 0) { + ++stats.busy_memory_replicas; + } + } else { + ++stats.non_memory_replicas; + } + } + } + } + } + + return stats; + } + + static bool MountBenchSegment(MasterService& service, const UUID& client_id, + size_t num_objects) { + auto segment = MakeSegment(num_objects); + auto mount_result = service.MountSegment(segment, client_id); + if (!mount_result.has_value()) { + LOG(ERROR) << "MountSegment failed: " + << toString(mount_result.error()); + return false; + } + return true; + } + + static bool CreateCompletedMemoryObjects(MasterService& service, + const UUID& client_id, + size_t num_objects) { + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segment = kSegmentName; + + for (size_t i = 0; i < num_objects; ++i) { + const std::string key = MakeKey(i); + auto put_start = service.PutStart( + client_id, key, TenantId::Default(), kObjectSize, config); + if (!put_start.has_value()) { + LOG(ERROR) << "PutStart failed for i=" << i + << ", error=" << toString(put_start.error()); + return false; + } + if (put_start->size() != 1u) { + LOG(ERROR) << "PutStart returned " << put_start->size() + << " replicas for i=" << i; + return false; + } + + auto put_end = service.PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); + if (!put_end.has_value()) { + LOG(ERROR) << "PutEnd failed for i=" << i + << ", error=" << toString(put_end.error()); + return false; + } + } + return true; + } + + static bool UsedBytes(MasterService& service, uint64_t& used_bytes) { + auto segment_usage = service.QuerySegments(kSegmentName); + if (!segment_usage.has_value()) { + LOG(ERROR) << "QuerySegments failed: " + << toString(segment_usage.error()); + return false; + } + used_bytes = segment_usage->first; + return true; + } + + static uint64_t WaitForSnapshotUniqueLock( + MasterService& service, std::chrono::microseconds waiter_delay) { + std::this_thread::sleep_for(waiter_delay); + const auto wait_start = std::chrono::steady_clock::now(); + uint64_t wait_us = 0; + { + std::unique_lock lock(service.snapshot_mutex_); + const auto wait_end = std::chrono::steady_clock::now(); + wait_us = std::chrono::duration_cast( + wait_end - wait_start) + .count(); + } + return wait_us; + } + + static bool ValidatePopulatedObjects(const MasterService& service, + size_t num_objects, + const MetadataStats& stats) { + if (service.GetKeyCount() != num_objects) { + LOG(ERROR) << "GetKeyCount mismatch: expected=" << num_objects + << ", actual=" << service.GetKeyCount(); + return false; + } + if (stats.object_count != num_objects || + stats.completed_memory_replicas != num_objects || + stats.busy_memory_replicas != 0 || stats.non_memory_replicas != 0 || + stats.incomplete_replicas != 0 || stats.unexpired_leases != 0) { + LOG(ERROR) << "metadata validation failed: objects=" + << stats.object_count << ", completed_memory_replicas=" + << stats.completed_memory_replicas + << ", busy_memory_replicas=" + << stats.busy_memory_replicas + << ", non_memory_replicas=" << stats.non_memory_replicas + << ", incomplete_replicas=" << stats.incomplete_replicas + << ", unexpired_leases=" << stats.unexpired_leases; + return false; + } + return true; + } + + static bool ValidateEvictionResult(size_t objects_before, + size_t evicted_count, + uint64_t freed_bytes) { + const size_t lowerbound = static_cast( + std::ceil(objects_before * FLAGS_evict_ratio_lowerbound)); + if (evicted_count < lowerbound) { + LOG(ERROR) << "evicted_count below lowerbound: evicted=" + << evicted_count << ", lowerbound=" << lowerbound; + return false; + } + if (evicted_count * kObjectSize != freed_bytes) { + LOG(ERROR) << "freed_bytes mismatch: evicted_count=" + << evicted_count << ", object_size=" << kObjectSize + << ", freed_bytes=" << freed_bytes; + return false; + } + return true; + } + + static bool RunOneScale(size_t num_objects) { + MasterService service(MakeConfig()); + const UUID client_id = generate_uuid(); + + if (!MountBenchSegment(service, client_id, num_objects) || + !CreateCompletedMemoryObjects(service, client_id, num_objects)) { + return false; + } + + const MetadataStats stats = ExpireLeasesAndCollectStats(service); + if (!ValidatePopulatedObjects(service, num_objects, stats)) { + return false; + } + + const size_t objects_before = service.GetKeyCount(); + uint64_t used_before = 0; + if (!UsedBytes(service, used_before)) { + return false; + } + + const auto evict_start = std::chrono::steady_clock::now(); + service.BatchEvict(FLAGS_evict_ratio_target, + FLAGS_evict_ratio_lowerbound); + const auto total_us = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - evict_start) + .count(); + + const size_t objects_after = service.GetKeyCount(); + uint64_t used_after = 0; + if (!UsedBytes(service, used_after)) { + return false; + } + const size_t evicted_count = objects_before - objects_after; + const uint64_t freed_bytes = + used_before >= used_after ? used_before - used_after : 0; + + if (!ValidateEvictionResult(objects_before, evicted_count, + freed_bytes)) { + return false; + } + + std::cout << num_objects << "," << total_us << "," << objects_before + << "," << objects_after << "," << evicted_count << "," + << freed_bytes << std::endl; + return true; + } + + static bool RunLockProbeTrial(size_t num_objects, + std::chrono::microseconds waiter_delay, + LockProbeTrialResult& result) { + MasterService service(MakeConfig()); + const UUID client_id = generate_uuid(); + + if (!MountBenchSegment(service, client_id, num_objects) || + !CreateCompletedMemoryObjects(service, client_id, num_objects)) { + return false; + } + + const MetadataStats stats = ExpireLeasesAndCollectStats(service); + if (!ValidatePopulatedObjects(service, num_objects, stats)) { + return false; + } + + const size_t objects_before = service.GetKeyCount(); + uint64_t used_before = 0; + if (!UsedBytes(service, used_before)) { + return false; + } + + uint64_t unique_lock_wait_us = 0; + std::thread waiter([&]() { + unique_lock_wait_us = + WaitForSnapshotUniqueLock(service, waiter_delay); + }); + + const auto evict_start = std::chrono::steady_clock::now(); + service.BatchEvict(FLAGS_evict_ratio_target, + FLAGS_evict_ratio_lowerbound); + const auto batch_evict_total_us = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - evict_start) + .count(); + waiter.join(); + + const size_t objects_after = service.GetKeyCount(); + uint64_t used_after = 0; + if (!UsedBytes(service, used_after)) { + return false; + } + const size_t evicted_count = objects_before - objects_after; + const uint64_t freed_bytes = + used_before >= used_after ? used_before - used_after : 0; + + if (!ValidateEvictionResult(objects_before, evicted_count, + freed_bytes)) { + return false; + } + + result.batch_evict_total_us = + static_cast(batch_evict_total_us); + result.unique_lock_wait_us = unique_lock_wait_us; + return true; + } +}; + +} // namespace mooncake::benchmarks + +int main(int argc, char** argv) { + google::InitGoogleLogging("BatchEvictBench"); + FLAGS_logtostderr = true; + gflags::ParseCommandLineFlags(&argc, &argv, true); + + if (!(FLAGS_evict_ratio_lowerbound > 0.0 && + FLAGS_evict_ratio_lowerbound <= FLAGS_evict_ratio_target && + FLAGS_evict_ratio_target <= 1.0)) { + LOG(ERROR) << "Invalid eviction ratios: require 0 < lowerbound <= " + "target <= 1, got target=" + << FLAGS_evict_ratio_target + << ", lowerbound=" << FLAGS_evict_ratio_lowerbound; + google::ShutdownGoogleLogging(); + return 1; + } + + LOG(INFO) << "BatchEvict benchmark config: num_objects=" + << FLAGS_num_objects + << ", target_ratio=" << FLAGS_evict_ratio_target + << ", lowerbound_ratio=" << FLAGS_evict_ratio_lowerbound; + + using mooncake::benchmarks::BatchEvictBench; + const bool ok = BatchEvictBench::RunRealBatchEvictScales() && + BatchEvictBench::RunSingleWaiterSnapshotMutexProbe(); + + google::ShutdownGoogleLogging(); + return ok ? 0 : 1; +} diff --git a/mooncake-store/benchmarks/batch_get_replica_bench.cpp b/mooncake-store/benchmarks/batch_get_replica_bench.cpp index 60988592c3..f03228cc5f 100644 --- a/mooncake-store/benchmarks/batch_get_replica_bench.cpp +++ b/mooncake-store/benchmarks/batch_get_replica_bench.cpp @@ -247,8 +247,13 @@ WriteStats PutCompletedBatch(mooncake::MasterClient& client, return stats; } + std::vector object_metas; + object_metas.reserve(started_keys.size()); + for (const auto& key : started_keys) { + object_metas.emplace_back(mooncake::ObjectMeta{key, std::nullopt}); + } auto put_end_result = - client.BatchPutEnd(started_keys, mooncake::ReplicaType::MEMORY); + client.BatchPutEnd(object_metas, mooncake::ReplicaType::MEMORY); for (const auto& result : put_end_result) { if (result.has_value()) { ++stats.completed; diff --git a/mooncake-store/benchmarks/master_bench.cpp b/mooncake-store/benchmarks/master_bench.cpp index ad774a33b0..392376c3a6 100644 --- a/mooncake-store/benchmarks/master_bench.cpp +++ b/mooncake-store/benchmarks/master_bench.cpp @@ -195,8 +195,8 @@ class BenchClient { return false; } - auto put_end_result = - master_client_.PutEnd(key, mooncake::ReplicaType::MEMORY); + auto put_end_result = master_client_.PutEnd( + {key, std::nullopt}, mooncake::ReplicaType::MEMORY); if (!put_end_result.has_value()) { return false; } @@ -228,7 +228,12 @@ class BenchClient { return 0; } - auto put_end_result = master_client_.BatchPutEnd(started_keys); + std::vector object_metas; + object_metas.reserve(started_keys.size()); + for (const auto& key : started_keys) { + object_metas.emplace_back(mooncake::ObjectMeta{key, std::nullopt}); + } + auto put_end_result = master_client_.BatchPutEnd(object_metas); for (auto& result : put_end_result) { if (result.has_value()) { success_cnt++; @@ -455,7 +460,13 @@ int main(int argc, char** argv) { } if (!started_keys.empty()) { - auto put_end_result = prefill_client.BatchPutEnd(started_keys); + std::vector object_metas; + object_metas.reserve(started_keys.size()); + for (const auto& key : started_keys) { + object_metas.emplace_back( + mooncake::ObjectMeta{key, std::nullopt}); + } + auto put_end_result = prefill_client.BatchPutEnd(object_metas); for (auto& result : put_end_result) { if (result.has_value()) { filled_objects++; diff --git a/mooncake-store/benchmarks/nof_worker_pool_bench.cpp b/mooncake-store/benchmarks/nof_worker_pool_bench.cpp index ed3d011ef2..be127cb1b1 100644 --- a/mooncake-store/benchmarks/nof_worker_pool_bench.cpp +++ b/mooncake-store/benchmarks/nof_worker_pool_bench.cpp @@ -1,5 +1,6 @@ #include #include +#include #include #include @@ -23,7 +24,6 @@ #include "spdk/spdk_wrapper.h" #include "transfer_task.h" -#include "utils.h" namespace { @@ -714,7 +714,12 @@ int main(int argc, char **argv) { SetEnvU64IfRequested("MC_NOF_INFLIGHT_BYTES_LIMIT", FLAGS_nof_inflight_bytes_limit); - auto endpoint_strings = mooncake::splitString(FLAGS_endpoints); + std::vector endpoint_strings; + boost::split(endpoint_strings, FLAGS_endpoints, boost::is_any_of(","), + boost::token_compress_on); + for (auto &endpoint : endpoint_strings) { + boost::trim(endpoint); + } if (endpoint_strings.empty()) { LOG(ERROR) << "No valid endpoints parsed from --endpoints"; return 1; diff --git a/mooncake-store/benchmarks/offset_allocator_concurrency_bench.cpp b/mooncake-store/benchmarks/offset_allocator_concurrency_bench.cpp new file mode 100644 index 0000000000..1b14599194 --- /dev/null +++ b/mooncake-store/benchmarks/offset_allocator_concurrency_bench.cpp @@ -0,0 +1,506 @@ +// Benchmarks allocator and PutStart allocation stages for concurrent puts to +// one segment. QPS and latency use separate phases; RPC, metadata, and transfer +// work are intentionally excluded. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "allocation_strategy.h" +#include "allocator.h" +#include "offset_allocator/offset_allocator.h" + +DEFINE_uint32(threads, 16, "Number of concurrent worker threads"); +DEFINE_uint64(iterations, 100000, "Measured operations per worker and phase"); +DEFINE_uint64(warmup_iterations, 1000, + "Warmup operations per worker before each measured phase"); +DEFINE_uint64(pool_size_mb, 1024, "Allocator capacity in MiB"); +DEFINE_double(high_water_ratio, 0.90, + "Used-space ratio for the capacity-exhaustion scenario"); +DEFINE_uint64(failed_request_mb, 128, + "Request size in MiB for expected-failure scenarios"); +DEFINE_uint64(success_request_kb, 4, + "Request size in KiB for the successful allocate/free scenario"); +DEFINE_uint64(fragmentation_block_mb, 4, + "Block size in MiB used to create deterministic fragmentation"); +DEFINE_uint32(fragmentation_stride, 5, + "Free every Nth block; 5 leaves approximately 80% used"); +DEFINE_double(mixed_failure_ratio, 0.05, + "Expected-failure ratio for the high-water mixed scenario"); +DEFINE_string(layer, "all", + "Benchmark layer: all, offset_allocator, or put_allocation"); +DEFINE_string(scenario, "all", + "Scenario: all, capacity, fragmentation, mixed, or success"); + +namespace { + +using Clock = std::chrono::steady_clock; +using mooncake::AllocatedBuffer; +using mooncake::AllocatorManager; +using mooncake::OffsetBufferAllocator; +using mooncake::RandomAllocationStrategy; +using mooncake::ReplicaType; +using mooncake::offset_allocator::OffsetAllocationHandle; +using mooncake::offset_allocator::OffsetAllocator; +using mooncake::offset_allocator::OffsetAllocStorageReport; + +constexpr uint64_t kKiB = 1024; +constexpr uint64_t kMiB = 1024 * kKiB; +constexpr uint64_t kMixedPatternSize = 10000; +constexpr uintptr_t kBenchmarkBaseAddress = 0x100000000ULL; +constexpr char kSegmentName[] = "benchmark-segment"; + +struct PhaseResult { + uint64_t operations = 0; + uint64_t unexpected_results = 0; + double seconds = 0; + std::vector latencies_ns; +}; + +struct Result { + std::string layer; + std::string scenario; + uint64_t operations = 0; + uint64_t throughput_unexpected_results = 0; + uint64_t latency_unexpected_results = 0; + double seconds = 0; + std::vector latencies_ns; + uint64_t capacity = 0; + uint64_t total_free_space = 0; + uint64_t largest_free_region = 0; + uint64_t request_size = 0; + uint64_t success_request_size = 0; + double expected_failure_ratio = 0; +}; + +template +PhaseResult runConcurrentPhase(Operation& operation, bool collect_latency) { + std::atomic ready{0}; + std::atomic start{false}; + std::atomic unexpected_results{0}; + std::vector> per_thread_latencies(FLAGS_threads); + std::vector workers; + workers.reserve(FLAGS_threads); + + for (uint32_t thread_index = 0; thread_index < FLAGS_threads; + ++thread_index) { + workers.emplace_back([&, thread_index] { + auto& latencies = per_thread_latencies[thread_index]; + if (collect_latency) { + latencies.reserve(FLAGS_iterations); + } + + for (uint64_t i = 0; i < FLAGS_warmup_iterations; ++i) { + operation(thread_index, i); + } + + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + + if (collect_latency) { + for (uint64_t i = 0; i < FLAGS_iterations; ++i) { + const auto begin = Clock::now(); + const bool expected_result = operation(thread_index, i); + const auto end = Clock::now(); + if (!expected_result) { + unexpected_results.fetch_add(1, + std::memory_order_relaxed); + } + latencies.push_back( + std::chrono::duration_cast( + end - begin) + .count()); + } + } else { + for (uint64_t i = 0; i < FLAGS_iterations; ++i) { + if (!operation(thread_index, i)) { + unexpected_results.fetch_add(1, + std::memory_order_relaxed); + } + } + } + }); + } + + while (ready.load(std::memory_order_acquire) != FLAGS_threads) { + std::this_thread::yield(); + } + const auto begin = Clock::now(); + start.store(true, std::memory_order_release); + for (auto& worker : workers) { + worker.join(); + } + const auto end = Clock::now(); + + PhaseResult result; + result.operations = FLAGS_iterations * FLAGS_threads; + result.unexpected_results = + unexpected_results.load(std::memory_order_relaxed); + result.seconds = std::chrono::duration(end - begin).count(); + if (collect_latency) { + result.latencies_ns.reserve(result.operations); + for (auto& latencies : per_thread_latencies) { + result.latencies_ns.insert(result.latencies_ns.end(), + latencies.begin(), latencies.end()); + } + } + return result; +} + +template +Result runConcurrentBenchmark(const std::string& layer, + const std::string& scenario, + const OffsetAllocStorageReport& report, + uint64_t capacity, uint64_t request_size, + uint64_t success_request_size, + double expected_failure_ratio, + Operation operation) { + auto throughput = runConcurrentPhase(operation, false); + auto latency = runConcurrentPhase(operation, true); + + Result result; + result.layer = layer; + result.scenario = scenario; + result.operations = throughput.operations; + result.throughput_unexpected_results = throughput.unexpected_results; + result.latency_unexpected_results = latency.unexpected_results; + result.seconds = throughput.seconds; + result.latencies_ns = std::move(latency.latencies_ns); + result.capacity = capacity; + result.total_free_space = report.totalFreeSpace; + result.largest_free_region = report.largestFreeRegion; + result.request_size = request_size; + result.success_request_size = success_request_size; + result.expected_failure_ratio = expected_failure_ratio; + return result; +} + +uint64_t percentile(const std::vector& values, double quantile) { + const size_t index = + static_cast(quantile * static_cast(values.size() - 1)); + return values[index]; +} + +void printCsvHeader() { + std::cout + << "layer,scenario,threads,iterations_per_thread,operations," + "throughput_seconds,qps,p50_ns,p99_ns," + "throughput_unexpected_results,latency_unexpected_results," + "capacity_bytes,used_percent,total_free_bytes," + "largest_free_region_bytes,request_bytes,success_request_bytes," + "expected_failure_ratio" + << std::endl; +} + +void printResult(Result result) { + std::sort(result.latencies_ns.begin(), result.latencies_ns.end()); + const double qps = static_cast(result.operations) / result.seconds; + const double used_percent = + 100.0 * static_cast(result.capacity - result.total_free_space) / + static_cast(result.capacity); + + std::cout << result.layer << ',' << result.scenario << ',' << FLAGS_threads + << ',' << FLAGS_iterations << ',' << result.operations << ',' + << std::fixed << std::setprecision(6) << result.seconds << ',' + << std::setprecision(2) << qps << ',' + << percentile(result.latencies_ns, 0.50) << ',' + << percentile(result.latencies_ns, 0.99) << ',' + << result.throughput_unexpected_results << ',' + << result.latency_unexpected_results << ',' << result.capacity + << ',' << used_percent << ',' << result.total_free_space << ',' + << result.largest_free_region << ',' << result.request_size << ',' + << result.success_request_size << ',' + << result.expected_failure_ratio << std::endl; +} + +bool isSelected(const std::string& selected, const std::string& value) { + return selected == "all" || selected == value; +} + +template +bool fillCapacityPressure(uint64_t capacity, double used_ratio, + Allocate allocate, std::vector& held) { + auto allocation = allocate( + static_cast(static_cast(capacity) * used_ratio)); + if (!allocation) { + return false; + } + held.emplace_back(std::move(allocation)); + return true; +} + +template +bool createFragmentation(uint64_t block_size, uint32_t stride, + Allocate allocate, std::vector& held) { + while (auto allocation = allocate(block_size)) { + held.emplace_back(std::move(allocation)); + } + if (held.size() < stride) { + return false; + } + for (size_t i = 0; i < held.size(); i += stride) { + held[i].reset(); + } + return true; +} + +class OffsetAllocatorFixture { + public: + OffsetAllocatorFixture(uint64_t capacity, uint64_t /* unused */) + : capacity_(capacity), + allocator_(OffsetAllocator::create(0, capacity)) {} + + bool prepareCapacityPressure(double used_ratio) { + return fillCapacityPressure( + capacity_, used_ratio, + [&](uint64_t size) { return allocator_->allocate(size); }, held_); + } + + bool prepareFragmentation(uint64_t block_size, uint32_t stride) { + return createFragmentation( + block_size, stride, + [&](uint64_t size) { return allocator_->allocate(size); }, held_); + } + + bool expectAllocationFailure(uint64_t request_size) { + return !allocator_->allocate(request_size).has_value(); + } + + bool allocateAndFree(uint64_t request_size) { + return allocator_->allocate(request_size).has_value(); + } + + OffsetAllocStorageReport report() const { + return allocator_->storageReport(); + } + + private: + uint64_t capacity_; + std::shared_ptr allocator_; + std::vector> held_; +}; + +class PutAllocationFixture { + public: + PutAllocationFixture(uint64_t capacity, uint64_t /* unused */) + : capacity_(capacity), + allocator_(std::make_shared( + kSegmentName, kBenchmarkBaseAddress, capacity, + "benchmark-endpoint", ReplicaType::MEMORY)) { + allocator_manager_.addAllocator(kSegmentName, allocator_); + } + + bool prepareCapacityPressure(double used_ratio) { + return fillCapacityPressure( + capacity_, used_ratio, + [&](uint64_t size) { return allocator_->allocate(size); }, held_); + } + + bool prepareFragmentation(uint64_t block_size, uint32_t stride) { + return createFragmentation( + block_size, stride, + [&](uint64_t size) { return allocator_->allocate(size); }, held_); + } + + bool expectAllocationFailure(uint64_t request_size) { + auto result = strategy_.Allocate(allocator_manager_, request_size, 1); + return !result.has_value(); + } + + bool allocateAndFree(uint64_t request_size) { + auto result = strategy_.Allocate(allocator_manager_, request_size, 1); + return result.has_value(); + } + + OffsetAllocStorageReport report() const { + return allocator_->getOffsetAllocator()->storageReport(); + } + + private: + uint64_t capacity_; + std::shared_ptr allocator_; + AllocatorManager allocator_manager_; + RandomAllocationStrategy strategy_; + std::vector> held_; +}; + +template +bool runFixtureScenarios(const std::string& layer, uint64_t capacity, + uint64_t failed_request_size, + uint64_t success_request_size, + uint64_t fragmentation_block_size) { + if (isSelected(FLAGS_scenario, "capacity")) { + Fixture fixture(capacity, fragmentation_block_size); + if (!fixture.prepareCapacityPressure(FLAGS_high_water_ratio)) { + std::cerr << "failed to prepare capacity-pressure fixture for " + << layer << std::endl; + return false; + } + const auto report = fixture.report(); + if (report.largestFreeRegion >= failed_request_size) { + std::cerr << "capacity scenario request must exceed the largest " + "free region" + << std::endl; + return false; + } + auto operation = [&](uint32_t, uint64_t) { + return fixture.expectAllocationFailure(failed_request_size); + }; + printResult(runConcurrentBenchmark(layer, "capacity_failed", report, + capacity, failed_request_size, 0, + 1.0, operation)); + } + + if (isSelected(FLAGS_scenario, "fragmentation")) { + Fixture fixture(capacity, fragmentation_block_size); + if (!fixture.prepareFragmentation(fragmentation_block_size, + FLAGS_fragmentation_stride)) { + std::cerr << "failed to prepare fragmented fixture for " << layer + << std::endl; + return false; + } + const auto report = fixture.report(); + if (report.totalFreeSpace < failed_request_size || + report.largestFreeRegion >= failed_request_size) { + std::cerr << "fragmentation scenario requires total free space >= " + "request size and largest free region < request size" + << std::endl; + return false; + } + auto operation = [&](uint32_t, uint64_t) { + return fixture.expectAllocationFailure(failed_request_size); + }; + printResult(runConcurrentBenchmark( + layer, "fragmentation_failed", report, capacity, + failed_request_size, 0, 1.0, operation)); + } + + if (isSelected(FLAGS_scenario, "mixed")) { + Fixture fixture(capacity, fragmentation_block_size); + if (!fixture.prepareCapacityPressure(FLAGS_high_water_ratio)) { + std::cerr << "failed to prepare mixed high-water fixture for " + << layer << std::endl; + return false; + } + const auto report = fixture.report(); + if (report.largestFreeRegion >= failed_request_size) { + std::cerr << "mixed scenario failure request must exceed the " + "largest free region" + << std::endl; + return false; + } + + const uint64_t failure_slots = static_cast( + FLAGS_mixed_failure_ratio * static_cast(kMixedPatternSize) + + 0.5); + auto operation = [&](uint32_t thread_index, uint64_t operation_index) { + const uint64_t pattern_slot = + (operation_index * 7919 + + static_cast(thread_index) * 104729) % + kMixedPatternSize; + if (pattern_slot < failure_slots) { + return fixture.expectAllocationFailure(failed_request_size); + } + return fixture.allocateAndFree(success_request_size); + }; + printResult( + runConcurrentBenchmark(layer, "high_water_mixed", report, capacity, + failed_request_size, success_request_size, + static_cast(failure_slots) / + static_cast(kMixedPatternSize), + operation)); + } + + if (isSelected(FLAGS_scenario, "success")) { + Fixture fixture(capacity, fragmentation_block_size); + const auto report = fixture.report(); + auto operation = [&](uint32_t, uint64_t) { + return fixture.allocateAndFree(success_request_size); + }; + printResult(runConcurrentBenchmark( + layer, "successful_allocate_free", report, capacity, + success_request_size, success_request_size, 0.0, operation)); + } + return true; +} + +bool validateFlags() { + const bool valid_layer = FLAGS_layer == "all" || + FLAGS_layer == "offset_allocator" || + FLAGS_layer == "put_allocation"; + const bool valid_scenario = + FLAGS_scenario == "all" || FLAGS_scenario == "capacity" || + FLAGS_scenario == "fragmentation" || FLAGS_scenario == "mixed" || + FLAGS_scenario == "success"; + if (!valid_layer || !valid_scenario || FLAGS_threads == 0 || + FLAGS_iterations == 0 || FLAGS_pool_size_mb == 0 || + FLAGS_failed_request_mb == 0 || FLAGS_success_request_kb == 0 || + FLAGS_fragmentation_block_mb == 0 || FLAGS_fragmentation_stride < 2 || + FLAGS_high_water_ratio <= 0.0 || FLAGS_high_water_ratio >= 1.0) { + return false; + } + const uint64_t mixed_failure_slots = static_cast( + FLAGS_mixed_failure_ratio * static_cast(kMixedPatternSize) + + 0.5); + if (mixed_failure_slots == 0 || mixed_failure_slots >= kMixedPatternSize) { + return false; + } + constexpr uint64_t max_size = std::numeric_limits::max(); + return FLAGS_pool_size_mb <= max_size / kMiB && + FLAGS_failed_request_mb <= max_size / kMiB && + FLAGS_success_request_kb <= max_size / kKiB && + FLAGS_fragmentation_block_mb <= max_size / kMiB; +} + +} // namespace + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + if (!validateFlags()) { + std::cerr << "invalid benchmark flags; use --help for valid values" + << std::endl; + return 1; + } + + const uint64_t capacity = FLAGS_pool_size_mb * kMiB; + const uint64_t failed_request_size = FLAGS_failed_request_mb * kMiB; + const uint64_t success_request_size = FLAGS_success_request_kb * kKiB; + const uint64_t fragmentation_block_size = + FLAGS_fragmentation_block_mb * kMiB; + if (failed_request_size >= capacity || success_request_size >= capacity || + fragmentation_block_size >= capacity) { + std::cerr << "request and fragmentation block sizes must be smaller " + "than the allocator capacity" + << std::endl; + return 1; + } + + printCsvHeader(); + + if (isSelected(FLAGS_layer, "offset_allocator") && + !runFixtureScenarios( + "offset_allocator", capacity, failed_request_size, + success_request_size, fragmentation_block_size)) { + return 1; + } + if (isSelected(FLAGS_layer, "put_allocation") && + !runFixtureScenarios( + "put_allocation", capacity, failed_request_size, + success_request_size, fragmentation_block_size)) { + return 1; + } + return 0; +} diff --git a/mooncake-store/benchmarks/oplog_batch_bench.cpp b/mooncake-store/benchmarks/oplog_batch_bench.cpp new file mode 100644 index 0000000000..6f2870a1ad --- /dev/null +++ b/mooncake-store/benchmarks/oplog_batch_bench.cpp @@ -0,0 +1,341 @@ +#include +#include +#if __has_include() +#include // Ubuntu +#else +#include // CentOS +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "etcd_helper.h" +#include "ha/kv/etcd_ha_kv_backend.h" +#include "ha/oplog/oplog_batch_codec.h" +#include "ha/oplog/oplog_batch_storage.h" +#include "ha/oplog/oplog_types.h" +#include "ha/oplog/ordered_oplog_writer.h" + +DEFINE_string(endpoints, "127.0.0.1:2379", "Etcd endpoints"); +DEFINE_string(cluster_id, "", "Unique OpLog cluster ID"); +DEFINE_string(mode, "writer", "Benchmark mode: backend or writer"); +DEFINE_string(output_json, "", "Output JSON path"); +DEFINE_uint64(max_entries, 64, "Maximum entries per batch"); +DEFINE_uint64(entry_bytes, 128, "OpLog payload bytes per entry"); +DEFINE_uint64(producer_threads, 1, "Writer producer threads"); +DEFINE_uint64(warmup_sec, 10, "Warmup duration"); +DEFINE_uint64(duration_sec, 30, "Measurement duration"); + +namespace mooncake { +namespace { + +using Clock = std::chrono::steady_clock; + +struct Counters { + std::atomic entries{0}; + std::atomic batches{0}; + std::atomic txn_us{0}; + std::atomic callbacks{0}; + std::atomic commit_to_callback_us{0}; + std::atomic failures{0}; + std::atomic last_batch_id{0}; +}; + +struct Snapshot { + uint64_t entries; + uint64_t batches; + uint64_t txn_us; + uint64_t callbacks; + uint64_t commit_to_callback_us; + uint64_t failures; + uint64_t last_batch_id; +}; + +Snapshot TakeSnapshot(const Counters& counters) { + return {.entries = counters.entries.load(), + .batches = counters.batches.load(), + .txn_us = counters.txn_us.load(), + .callbacks = counters.callbacks.load(), + .commit_to_callback_us = counters.commit_to_callback_us.load(), + .failures = counters.failures.load(), + .last_batch_id = counters.last_batch_id.load()}; +} + +Snapshot Delta(const Snapshot& end, const Snapshot& begin) { + return {.entries = end.entries - begin.entries, + .batches = end.batches - begin.batches, + .txn_us = end.txn_us - begin.txn_us, + .callbacks = end.callbacks - begin.callbacks, + .commit_to_callback_us = + end.commit_to_callback_us - begin.commit_to_callback_us, + .failures = end.failures - begin.failures, + .last_batch_id = end.last_batch_id}; +} + +OpLogEntry MakeEntry(uint64_t id) { + OpLogEntry entry; + entry.timestamp_ms = 1; + entry.op_type = OpType::PUT_END; + entry.tenant_id = "default"; + entry.object_key = "bench-key-" + std::to_string(id); + entry.payload.assign(FLAGS_entry_bytes, 'x'); + entry.checksum = ComputeOpLogChecksum(entry.payload); + return entry; +} + +ErrorCode RunBackend(OpLogBatchStorage& storage, Counters& counters, + Clock::time_point deadline) { + DurablePrefix prefix; + ErrorCode err = storage.ReadDurablePrefix(prefix); + if (err != ErrorCode::OK) { + return err; + } + while (Clock::now() < deadline) { + OpLogBatchRecord batch; + batch.batch_id = prefix.batch_id + 1; + batch.first_seq = prefix.last_seq + 1; + for (uint64_t i = 0; i < FLAGS_max_entries; ++i) { + auto entry = MakeEntry(batch.first_seq + i); + entry.sequence_id = batch.first_seq + i; + batch.entries.push_back(std::move(entry)); + } + batch.last_seq = batch.entries.back().sequence_id; + const auto started = Clock::now(); + err = storage.WriteBatchAndAdvancePrefix(batch, prefix); + counters.txn_us += + std::chrono::duration_cast(Clock::now() - + started) + .count(); + if (err != ErrorCode::OK) { + ++counters.failures; + return err; + } + prefix = {.batch_id = batch.batch_id, .last_seq = batch.last_seq}; + counters.last_batch_id = batch.batch_id; + counters.entries += batch.entries.size(); + ++counters.batches; + } + return ErrorCode::OK; +} + +ErrorCode VerifyHistory(OpLogBatchStorage& storage, uint64_t expected_entries) { + DurablePrefix prefix; + ErrorCode err = storage.ReadDurablePrefix(prefix); + if (err != ErrorCode::OK || prefix.last_seq != expected_entries) { + return ErrorCode::INTERNAL_ERROR; + } + std::vector batches; + err = storage.ReadBatchesAfter(0, 0, batches); + if (err != ErrorCode::OK || batches.size() != prefix.batch_id) { + return ErrorCode::INTERNAL_ERROR; + } + return ErrorCode::OK; +} + +ErrorCode MeasureEncodedBytes(OpLogBatchStorage& storage, + uint64_t after_batch_id, uint64_t last_batch_id, + uint64_t* encoded_bytes) { + std::vector batches; + const size_t count = static_cast(last_batch_id - after_batch_id); + ErrorCode err = storage.ReadBatchesAfter(after_batch_id, count, batches); + if (err != ErrorCode::OK || batches.size() != count) { + return ErrorCode::INTERNAL_ERROR; + } + *encoded_bytes = 0; + for (const auto& batch : batches) { + *encoded_bytes += EncodeOpLogBatchRecord(batch).size(); + } + return ErrorCode::OK; +} + +void WriteResult(const Snapshot& measured, double seconds, + double cpu_cores_used, uint64_t encoded_bytes, + const DurablePrefix& prefix) { + Json::Value root; + root["schema_version"] = 1; + root["mode"] = FLAGS_mode; + root["entries"] = Json::UInt64(measured.entries); + root["batches"] = Json::UInt64(measured.batches); + root["entries_per_sec"] = measured.entries / seconds; + root["cpu_cores_used"] = cpu_cores_used; + root["transactions_per_sec"] = measured.batches / seconds; + root["batch_entries_mean"] = + measured.batches + ? static_cast(measured.entries) / measured.batches + : 0.0; + root["encoded_bytes"] = Json::UInt64(encoded_bytes); + root["failures"] = Json::UInt64(measured.failures); + root["durable_batch_id"] = Json::UInt64(prefix.batch_id); + root["durable_sequence"] = Json::UInt64(prefix.last_seq); + const double txn_us = + measured.batches + ? static_cast(measured.txn_us) / measured.batches + : 0.0; + const double commit_to_callback_us = + measured.callbacks + ? static_cast(measured.commit_to_callback_us) / + measured.callbacks + : 0.0; + root["commit_to_callback_us"] = commit_to_callback_us; + root["stage_latency_us"]["txn"] = txn_us; + root["stage_latency_us"]["post_txn"] = + std::max(0.0, commit_to_callback_us - txn_us); + Json::StreamWriterBuilder builder; + builder["indentation"] = " "; + std::ofstream output(FLAGS_output_json); + output << Json::writeString(builder, root) << '\n'; +} + +} // namespace +} // namespace mooncake + +int main(int argc, char** argv) { + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = true; + gflags::ParseCommandLineFlags(&argc, &argv, true); + if (FLAGS_cluster_id.empty() || FLAGS_output_json.empty() || + (FLAGS_mode != "backend" && FLAGS_mode != "writer") || + FLAGS_max_entries == 0 || FLAGS_entry_bytes == 0 || + FLAGS_producer_threads == 0 || FLAGS_duration_sec == 0 || + (FLAGS_mode == "backend" && FLAGS_producer_threads != 1)) { + LOG(ERROR) << "invalid benchmark arguments"; + return 1; + } + std::string cluster_id = FLAGS_cluster_id; + if (!mooncake::NormalizeAndValidateClusterId(cluster_id) || + cluster_id.empty()) { + LOG(ERROR) << "invalid cluster_id"; + return 1; + } + auto err = mooncake::EtcdHelper::ConnectToEtcdStoreClient(FLAGS_endpoints); + if (err != mooncake::ErrorCode::OK) { + LOG(ERROR) << "failed to connect to etcd: " << mooncake::toString(err); + return 1; + } + + mooncake::EtcdHaKvBackend backend; + mooncake::OpLogBatchStorage storage(cluster_id, backend); + mooncake::DurablePrefix initial_prefix; + err = storage.InitDurablePrefix(initial_prefix); + if (err != mooncake::ErrorCode::OK || initial_prefix.batch_id != 0 || + initial_prefix.last_seq != 0) { + LOG(ERROR) << "benchmark requires an empty cluster namespace"; + return 1; + } + + mooncake::Counters counters; + mooncake::Snapshot measured{}; + mooncake::Snapshot measure_begin{}; + mooncake::Snapshot measure_end{}; + std::clock_t cpu_begin = 0; + std::clock_t cpu_end = 0; + const auto warmup_deadline = + mooncake::Clock::now() + std::chrono::seconds(FLAGS_warmup_sec); + const auto measure_deadline = + warmup_deadline + std::chrono::seconds(FLAGS_duration_sec); + + if (FLAGS_mode == "backend") { + err = mooncake::RunBackend(storage, counters, warmup_deadline); + measure_begin = mooncake::TakeSnapshot(counters); + cpu_begin = std::clock(); + if (err == mooncake::ErrorCode::OK) { + err = mooncake::RunBackend(storage, counters, measure_deadline); + } + measure_end = mooncake::TakeSnapshot(counters); + cpu_end = std::clock(); + measured = mooncake::Delta(measure_end, measure_begin); + } else { + mooncake::OrderedOpLogWriter writer( + {.max_entries_per_batch = static_cast(FLAGS_max_entries)}, + [&](const mooncake::OpLogBatchRecord& batch, + const mooncake::DurablePrefix& prefix) { + const auto started = mooncake::Clock::now(); + const auto result = + storage.WriteBatchAndAdvancePrefix(batch, prefix); + counters.txn_us += + std::chrono::duration_cast( + mooncake::Clock::now() - started) + .count(); + if (result == mooncake::ErrorCode::OK) { + counters.entries += batch.entries.size(); + ++counters.batches; + counters.last_batch_id = batch.batch_id; + } else { + ++counters.failures; + } + return result; + }); + writer.Start(); + std::atomic stop{false}; + std::atomic next_id{1}; + std::vector producers; + for (uint64_t i = 0; i < FLAGS_producer_threads; ++i) { + producers.emplace_back([&] { + while (!stop.load(std::memory_order_relaxed)) { + auto reservation = writer.Reserve(); + if (!reservation.has_value()) { + std::this_thread::yield(); + continue; + } + const auto committed_at = mooncake::Clock::now(); + const auto committed = writer.Commit( + std::move(*reservation), + mooncake::MakeEntry(next_id.fetch_add(1)), + [&, committed_at](const mooncake::OpLogEntry&) { + counters.commit_to_callback_us += + std::chrono::duration_cast< + std::chrono::microseconds>( + mooncake::Clock::now() - committed_at) + .count(); + ++counters.callbacks; + }); + if (!committed.has_value()) { + ++counters.failures; + } + } + }); + } + std::this_thread::sleep_until(warmup_deadline); + measure_begin = mooncake::TakeSnapshot(counters); + cpu_begin = std::clock(); + std::this_thread::sleep_until(measure_deadline); + measure_end = mooncake::TakeSnapshot(counters); + cpu_end = std::clock(); + measured = mooncake::Delta(measure_end, measure_begin); + stop = true; + for (auto& producer : producers) { + producer.join(); + } + writer.Stop(); + err = writer.LastError(); + } + + const auto final = mooncake::TakeSnapshot(counters); + uint64_t encoded_bytes = 0; + if (err == mooncake::ErrorCode::OK) { + err = mooncake::VerifyHistory(storage, final.entries); + } + if (err == mooncake::ErrorCode::OK) { + err = mooncake::MeasureEncodedBytes( + storage, measure_begin.last_batch_id, measure_end.last_batch_id, + &encoded_bytes); + } + mooncake::DurablePrefix final_prefix; + if (err != mooncake::ErrorCode::OK || + storage.ReadDurablePrefix(final_prefix) != mooncake::ErrorCode::OK) { + LOG(ERROR) << "benchmark or history verification failed: " + << mooncake::toString(err); + return 3; + } + const double cpu_cores_used = static_cast(cpu_end - cpu_begin) / + CLOCKS_PER_SEC / FLAGS_duration_sec; + mooncake::WriteResult(measured, FLAGS_duration_sec, cpu_cores_used, + encoded_bytes, final_prefix); + return measured.failures == 0 ? 0 : 3; +} diff --git a/mooncake-store/benchmarks/report_oplog_batch.py b/mooncake-store/benchmarks/report_oplog_batch.py new file mode 100644 index 0000000000..100ac23adb --- /dev/null +++ b/mooncake-store/benchmarks/report_oplog_batch.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""Summarize repeatable batch-record OpLog benchmark samples.""" + +import argparse +import csv +import json +import statistics +from collections import defaultdict +from pathlib import Path + + +def _stats(values): + median = statistics.median(values) + mean = statistics.mean(values) + return { + "median": median, + "min": min(values), + "max": max(values), + "cv": statistics.pstdev(values) / mean if mean else None, + } + + +def _median_field(rows, name): + values = [row[name] for row in rows if isinstance(row.get(name), (int, float))] + return statistics.median(values) if len(values) == len(rows) else None + + +def _summarize_case(name, rows): + result = {"case": name, "repeats": len(rows), "status": "ok", "missing": []} + throughput = [ + row["entries_per_sec"] + for row in rows + if isinstance(row.get("entries_per_sec"), (int, float)) + ] + if throughput: + result["entries_per_sec"] = _stats(throughput) + else: + result["missing"].append("entries_per_sec") + + cpu_cores = [ + row["cpu_cores_used"] + for row in rows + if isinstance(row.get("cpu_cores_used"), (int, float)) + ] + if len(cpu_cores) == len(rows): + result["cpu_cores_used"] = _stats(cpu_cores) + + if len(rows) < 3: + result["missing"].append("at least 3 repeats") + + entries = _median_field(rows, "entries") + encoded_bytes = _median_field(rows, "encoded_bytes") + if entries and encoded_bytes is not None: + result["bytes_per_entry"] = encoded_bytes / entries + else: + result["missing"].append("encoded_bytes and entries") + + stage_names = sorted( + set().union(*(row.get("stage_latency_us", {}).keys() for row in rows)) + ) + stage_medians = {} + for stage in stage_names: + values = [row.get("stage_latency_us", {}).get(stage) for row in rows] + if all(isinstance(value, (int, float)) for value in values): + stage_medians[stage] = statistics.median(values) + total_stage_latency = sum(stage_medians.values()) + if total_stage_latency: + result["stage_latency_us"] = stage_medians + result["stage_share"] = { + stage: value / total_stage_latency for stage, value in stage_medians.items() + } + result["bottleneck_stage"] = max(stage_medians, key=stage_medians.get) + + txn_us = stage_medians.get("txn") + batch_entries = _median_field(rows, "batch_entries_mean") + if txn_us and batch_entries: + theoretical = batch_entries * 1_000_000 / txn_us + result["theoretical_entries_per_sec"] = theoretical + if throughput: + result["observed_efficiency"] = statistics.median(throughput) / theoretical + else: + result["missing"].append("stage_latency_us.txn") + if not batch_entries: + result["missing"].append("batch_entries_mean") + + if result["missing"]: + result["status"] = "insufficient_data" + return result + + +def _write_markdown(path, summary): + lines = [ + "# OpLog Batch Performance Report", + "", + "| Case | Status | Repeats | Median entries/s | CV | CPU cores | Bytes/entry | Bottleneck | Theoretical entries/s | Efficiency |", + "|---|---:|---:|---:|---:|---:|---:|---|---:|---:|", + ] + for case in summary["cases"]: + throughput = case.get("entries_per_sec", {}) + lines.append( + "| {case} | {status} | {repeats} | {median} | {cv} | {cpu} | {bytes_per_entry} | {bottleneck} | {theoretical} | {efficiency} |".format( + case=case["case"], + status=case["status"], + repeats=case["repeats"], + median=throughput.get("median", "n/a"), + cv=_format_number(throughput.get("cv")), + cpu=_format_number(case.get("cpu_cores_used", {}).get("median")), + bytes_per_entry=_format_number(case.get("bytes_per_entry")), + bottleneck=case.get("bottleneck_stage", "n/a"), + theoretical=_format_number(case.get("theoretical_entries_per_sec")), + efficiency=_format_number(case.get("observed_efficiency")), + ) + ) + path.write_text("\n".join(lines) + "\n") + + +def _format_number(value): + return "n/a" if value is None else f"{value:.4g}" + + +def generate_report(samples_path, output_dir): + samples_path = Path(samples_path) + output_dir = Path(output_dir) + grouped = defaultdict(list) + with samples_path.open() as stream: + for line_number, line in enumerate(stream, 1): + if not line.strip(): + continue + row = json.loads(line) + if row.get("schema_version") != 1 or not row.get("case"): + raise ValueError(f"invalid sample at line {line_number}") + grouped[row["case"]].append(row) + + summary = { + "schema_version": 1, + "sample_file": str(samples_path), + "cases": [_summarize_case(name, grouped[name]) for name in sorted(grouped)], + } + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + _write_markdown(output_dir / "report.md", summary) + with (output_dir / "cases.csv").open("w", newline="") as stream: + fields = [ + "case", + "status", + "repeats", + "median_entries_per_sec", + "cv", + "cpu_cores_used", + "bytes_per_entry", + "bottleneck_stage", + "theoretical_entries_per_sec", + "observed_efficiency", + ] + writer = csv.DictWriter(stream, fieldnames=fields) + writer.writeheader() + for case in summary["cases"]: + writer.writerow( + { + "case": case["case"], + "status": case["status"], + "repeats": case["repeats"], + "median_entries_per_sec": case.get("entries_per_sec", {}).get( + "median" + ), + "cv": case.get("entries_per_sec", {}).get("cv"), + "cpu_cores_used": case.get("cpu_cores_used", {}).get("median"), + "bytes_per_entry": case.get("bytes_per_entry"), + "bottleneck_stage": case.get("bottleneck_stage"), + "theoretical_entries_per_sec": case.get( + "theoretical_entries_per_sec" + ), + "observed_efficiency": case.get("observed_efficiency"), + } + ) + return summary + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--samples", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + args = parser.parse_args() + summary = generate_report(args.samples, args.output_dir) + print(json.dumps(summary, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mooncake-store/benchmarks/run_oplog_batch_sweep.py b/mooncake-store/benchmarks/run_oplog_batch_sweep.py new file mode 100644 index 0000000000..c6b7390ce7 --- /dev/null +++ b/mooncake-store/benchmarks/run_oplog_batch_sweep.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +"""Run a small, repeatable matrix against oplog_batch_bench.""" + +import argparse +import fcntl +import json +import os +import platform +import subprocess +import sys +import time +from pathlib import Path + + +PROFILES = { + "smoke": { + "cases": [ + { + "name": "writer-b1", + "mode": "writer", + "max_entries": 1, + "duration_sec": 5, + "warmup_sec": 1, + }, + { + "name": "writer-b64", + "mode": "writer", + "max_entries": 64, + "duration_sec": 5, + "warmup_sec": 1, + }, + { + "name": "writer-b1024", + "mode": "writer", + "max_entries": 1024, + "duration_sec": 5, + "warmup_sec": 1, + }, + ] + }, + "normal": { + "cases": [ + { + "name": f"writer-b{size}", + "mode": "writer", + "max_entries": size, + "duration_sec": 30, + "warmup_sec": 10, + } + for size in (1, 8, 64, 256, 1024) + ] + }, + "full": { + "cases": [ + { + "name": f"{mode}-b{size}", + "mode": mode, + "max_entries": size, + "duration_sec": 60, + "warmup_sec": 10, + } + for mode in ("backend", "writer") + for size in (1, 8, 64, 256, 1024, 4096) + ] + }, +} + + +def _load_matrix(matrix_path, profile): + if matrix_path: + return json.loads(Path(matrix_path).read_text()) + return PROFILES[profile] + + +def run_sweep( + benchmark, + output_dir, + repeat, + endpoints, + matrix_path=None, + profile="smoke", + resume=False, +): + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + with (output_dir / ".sweep.lock").open("w") as lock_stream: + try: + fcntl.flock(lock_stream, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as error: + raise RuntimeError(f"sweep already running in {output_dir}") from error + return _run_sweep_unlocked( + benchmark, + output_dir, + repeat, + endpoints, + matrix_path, + profile, + resume, + ) + + +def _run_sweep_unlocked( + benchmark, output_dir, repeat, endpoints, matrix_path, profile, resume +): + benchmark = Path(benchmark) + matrix = _load_matrix(matrix_path, profile) + manifest = { + "schema_version": 1, + "benchmark": str(benchmark), + "endpoints": endpoints, + "repeat": repeat, + "profile": profile if not matrix_path else None, + "matrix": matrix, + "created_unix_ns": time.time_ns(), + "host": { + "system": platform.system(), + "release": platform.release(), + "machine": platform.machine(), + "cpu_count": os.cpu_count(), + "python": sys.version.split()[0], + }, + } + samples_path = output_dir / "samples.jsonl" + manifest_path = output_dir / "manifest.json" + completed_keys = set() + if resume: + previous = json.loads(manifest_path.read_text()) + for key in ("benchmark", "endpoints", "repeat", "matrix"): + if previous.get(key) != manifest.get(key): + raise ValueError(f"resume manifest mismatch: {key}") + if samples_path.is_file(): + for line in samples_path.read_text().splitlines(): + row = json.loads(line) + if row.get("exit_code") == 0: + completed_keys.add((row["case"], row["repeat"])) + else: + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + + failed = False + with samples_path.open("a" if resume else "w") as samples: + for case in matrix.get("cases", []): + name = case["name"] + for repeat_index in range(repeat): + if (name, repeat_index) in completed_keys: + continue + run_id = f"{name}-r{repeat_index}-{time.time_ns()}" + result_path = output_dir / f"{run_id}.json" + command = [ + str(benchmark), + f"--endpoints={endpoints}", + f"--cluster_id=oplog-bench-{run_id}", + f"--output_json={result_path}", + ] + command.extend( + f"--{key}={str(value).lower() if isinstance(value, bool) else value}" + for key, value in case.items() + if key != "name" + ) + started_ns = time.time_ns() + completed = subprocess.run(command, capture_output=True, text=True) + row = {} + if result_path.is_file(): + try: + row = json.loads(result_path.read_text()) + except (json.JSONDecodeError, OSError): + row = {} + row.update( + { + "schema_version": 1, + "case": name, + "repeat": repeat_index, + "cluster_id": f"oplog-bench-{run_id}", + "exit_code": completed.returncode, + "wall_time_ns": time.time_ns() - started_ns, + "stdout_file": f"{run_id}.stdout.log", + "stderr_file": f"{run_id}.stderr.log", + } + ) + (output_dir / row["stdout_file"]).write_text(completed.stdout) + (output_dir / row["stderr_file"]).write_text(completed.stderr) + samples.write(json.dumps(row, sort_keys=True) + "\n") + samples.flush() + failed |= completed.returncode != 0 + return int(failed) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--benchmark", required=True, type=Path) + parser.add_argument("--output-dir", required=True, type=Path) + parser.add_argument("--endpoints", default="127.0.0.1:2379") + parser.add_argument("--repeat", type=int, default=3) + parser.add_argument("--profile", choices=PROFILES, default="smoke") + parser.add_argument("--matrix", type=Path) + parser.add_argument("--resume", action="store_true") + args = parser.parse_args() + if args.repeat < 1: + parser.error("--repeat must be at least 1") + raise SystemExit( + run_sweep( + benchmark=args.benchmark, + output_dir=args.output_dir, + repeat=args.repeat, + endpoints=args.endpoints, + matrix_path=args.matrix, + profile=args.profile, + resume=args.resume, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/mooncake-store/benchmarks/store_kv_bench.py b/mooncake-store/benchmarks/store_kv_bench.py index 68fa639bd6..f09151cfb8 100644 --- a/mooncake-store/benchmarks/store_kv_bench.py +++ b/mooncake-store/benchmarks/store_kv_bench.py @@ -5,6 +5,8 @@ import argparse import ctypes +import importlib +import json import logging import math import os @@ -15,7 +17,23 @@ from collections import Counter from dataclasses import dataclass, field from typing import Callable, Iterable, List, Optional -from mooncake.store import MooncakeDistributedStore, ReplicateConfig, get_alloc_func_addr, get_free_func_addr + +_store_module = importlib.import_module( + os.environ.get("MOONCAKE_STORE_MODULE", "mooncake.store") +) +MooncakeDistributedStore = _store_module.MooncakeDistributedStore +ReplicateConfig = _store_module.ReplicateConfig + +try: + get_alloc_func_addr = _store_module.get_alloc_func_addr + get_free_func_addr = _store_module.get_free_func_addr +except AttributeError: + + def get_alloc_func_addr(): + return None + + def get_free_func_addr(): + return None LOG = logging.getLogger("store_kv_bench") @@ -29,7 +47,17 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument( "--scenario", required=True, - choices=["verify_write", "fill", "write_perf", "read_perf", "mixed_rw"], + choices=[ + "verify_write", + "fill", + "write_perf", + "read_perf", + "mixed_rw", + "metadata_smoke", + "mixed_metadata", + "remove_perf", + "replay", + ], help="Benchmark scenario to execute.", ) @@ -50,7 +78,9 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--numjobs", type=int, default=1) parser.add_argument("--iodepth", type=int, default=1) parser.add_argument("--batch-size", type=int, default=1) - parser.add_argument("--runtime", type=int, default=0, help="Seconds. 0 means object-count based.") + parser.add_argument( + "--runtime", type=int, default=0, help="Seconds. 0 means object-count based." + ) parser.add_argument("--nr-objects", type=int, default=128) parser.add_argument("--write-objects", type=int, default=0) parser.add_argument( @@ -64,13 +94,27 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--key-size", type=int, default=20) parser.add_argument("--value-size", type=int, default=4096) parser.add_argument("--rand-seed", type=int, default=1) + parser.add_argument("--put-pct", type=int, default=40) + parser.add_argument("--get-pct", type=int, default=40) + parser.add_argument("--exist-pct", type=int, default=10) + parser.add_argument("--remove-pct", type=int, default=10) + parser.add_argument("--latency-sample-rate", type=int, default=100) + parser.add_argument("--output-dir", default="") + parser.add_argument( + "--journal", choices=["none", "failures", "all"], default="failures" + ) + parser.add_argument("--summary-json", default="") + parser.add_argument("--replay-file", default="") + parser.add_argument("--tenant-id", default="default") parser.add_argument("--memory-replica-num", type=int, default=1) parser.add_argument("--nof-replica-num", type=int, default=0) parser.add_argument("--verify", action="store_true") parser.add_argument("--pattern", default="") - parser.add_argument("--prepare-mode", choices=["auto", "none", "write"], default="auto") + parser.add_argument( + "--prepare-mode", choices=["auto", "none", "write"], default="auto" + ) parser.add_argument("--rwmixread", type=int, default=70) parser.add_argument( @@ -110,6 +154,7 @@ class PhaseStats: start_time: float = 0.0 end_time: float = 0.0 dataset_exhausted: bool = False + latency_seen: int = 0 @dataclass @@ -124,6 +169,115 @@ class RequestResult: error_counts: Counter = field(default_factory=Counter) +METADATA_OPERATIONS = ("put", "get", "exist", "remove") + + +def validate_operation_percentages(weights: dict) -> None: + if set(weights) != set(METADATA_OPERATIONS): + raise ValueError(f"operation percentages require {METADATA_OPERATIONS}") + if any(not isinstance(value, int) or value < 0 for value in weights.values()): + raise ValueError("operation percentages must be non-negative integers") + if sum(weights.values()) != 100: + raise ValueError("operation percentages must sum to 100") + + +def choose_metadata_operations( + seed: int, lane: int, count: int, weights: dict +) -> List[str]: + validate_operation_percentages(weights) + rng = random.Random(seed + lane) + boundaries = [] + total = 0 + for operation in METADATA_OPERATIONS: + total += weights[operation] + boundaries.append((total, operation)) + result = [] + for _ in range(count): + choice = rng.randrange(100) + result.append( + next(operation for limit, operation in boundaries if choice < limit) + ) + return result + + +class MetadataLaneModel: + def __init__(self): + self._present = set() + + def is_present(self, object_id: int) -> bool: + return object_id in self._present + + def record(self, object_id: int, operation: str, success: bool) -> None: + if not success: + return + if operation == "put": + self._present.add(object_id) + elif operation == "remove": + self._present.discard(object_id) + + +class LatencySampler: + def __init__(self, rate: int): + if rate <= 0: + raise ValueError("latency sample rate must be > 0") + self.rate = rate + self.seen = 0 + self.samples = [] + + def add(self, value: float) -> None: + if self.seen % self.rate == 0: + self.samples.append(value) + self.seen += 1 + + +def write_jsonl(path, records: Iterable[dict]) -> None: + os.makedirs(os.path.dirname(os.fspath(path)) or ".", exist_ok=True) + with open(path, "w", encoding="utf-8") as output: + for record in records: + output.write(json.dumps(record, sort_keys=True) + "\n") + + +def read_replay(path) -> List[dict]: + required = { + "schema_version", + "op_id", + "lane", + "op", + "object_id", + "key", + "invoke_ns", + "return_ns", + "result", + } + records = [] + with open(path, encoding="utf-8") as input_file: + for line_number, line in enumerate(input_file, 1): + if not line.strip(): + continue + record = json.loads(line) + missing = required - set(record) + if ( + missing + or record.get("schema_version") != 1 + or record.get("op") not in METADATA_OPERATIONS + ): + raise ValueError( + f"invalid replay record at line {line_number}: missing={sorted(missing)}" + ) + records.append(record) + return records + + +def write_json_atomic(path, value: dict) -> None: + path = os.fspath(path) + os.makedirs(os.path.dirname(path) or ".", exist_ok=True) + temporary = path + ".tmp" + with open(temporary, "w", encoding="utf-8") as output: + json.dump(value, output, sort_keys=True) + output.write("\n") + os.replace(temporary, path) + + class PayloadFactory: def __init__(self, value_size: int, pattern: bytes): self.value_size = value_size @@ -201,7 +355,9 @@ def next_read_ids( source: str = "written", ) -> List[int]: with self.ids_lock: - readable_ids = self.prepared_ids if source == "prepared" else self.written_ids + readable_ids = ( + self.prepared_ids if source == "prepared" else self.written_ids + ) if not readable_ids: return [] if sequential: @@ -232,7 +388,9 @@ def parse_pattern(pattern_text: str) -> bytes: def make_key(prefix: str, key_size: int, object_id: int) -> str: suffix = f"{object_id:016d}" if key_size < len(suffix): - raise ValueError(f"key_size={key_size} is smaller than suffix length {len(suffix)}") + raise ValueError( + f"key_size={key_size} is smaller than suffix length {len(suffix)}" + ) prefix_space = key_size - len(suffix) prefix_part = prefix[:prefix_space].ljust(prefix_space, "_") return f"{prefix_part}{suffix}" @@ -260,12 +418,17 @@ def close(self) -> None: self._zcopy = None def put_ids(self, object_ids: List[int]) -> RequestResult: - keys = [make_key(self.args.key_prefix, self.args.key_size, object_id) for object_id in object_ids] + keys = [ + make_key(self.args.key_prefix, self.args.key_size, object_id) + for object_id in object_ids + ] ret_codes = self._put_keys(keys, object_ids) errors: Counter = Counter() success_ids: List[int] = [] - if self.args.io_api == "plain" and not (len(object_ids) == 1 and self.args.batch_size == 1): + if self.args.io_api == "plain" and not ( + len(object_ids) == 1 and self.args.batch_size == 1 + ): ret = ret_codes[0] if ret == 0: success_ids.extend(object_ids) @@ -288,7 +451,10 @@ def put_ids(self, object_ids: List[int]) -> RequestResult: ) def get_ids(self, object_ids: List[int], verify: bool) -> RequestResult: - keys = [make_key(self.args.key_prefix, self.args.key_size, object_id) for object_id in object_ids] + keys = [ + make_key(self.args.key_prefix, self.args.key_size, object_id) + for object_id in object_ids + ] errors: Counter = Counter() kv_successes = 0 misses = 0 @@ -300,7 +466,9 @@ def get_ids(self, object_ids: List[int], verify: bool) -> RequestResult: misses += 1 errors["MISS"] += 1 continue - if verify and not self.payload_factory.verify_payload(object_id, payload): + if verify and not self.payload_factory.verify_payload( + object_id, payload + ): verify_failures += 1 errors["VERIFY_FAIL"] += 1 continue @@ -339,6 +507,62 @@ def get_ids(self, object_ids: List[int], verify: bool) -> RequestResult: error_counts=errors, ) + def metadata_operation( + self, operation: str, object_id: int, expected_present: bool + ): + key = make_key(self.args.key_prefix, self.args.key_size, object_id) + result_code = 0 + actual_present = expected_present + bytes_processed = 0 + success = False + verify_failures = 0 + + if operation == "put": + payload = self.payload_factory.build(object_id) + result_code = self.store.put(key, payload, self.config) + actual_present = result_code == 0 or self.store.isExist(key) == 1 + success = result_code == 0 or (expected_present and actual_present) + bytes_processed = len(payload) if success else 0 + elif operation == "get": + payload = self.store.get(key) + actual_present = payload not in (None, b"") + success = actual_present == expected_present + if actual_present and not self.payload_factory.verify_payload( + object_id, payload + ): + success = False + verify_failures = 1 + result_code = "VERIFY_FAIL" + elif not success: + result_code = "PRESENCE_MISMATCH" + bytes_processed = len(payload) if actual_present else 0 + elif operation == "exist": + result_code = self.store.isExist(key) + actual_present = result_code == 1 + success = result_code >= 0 and actual_present == expected_present + elif operation == "remove": + result_code = self.store.remove(key) + actual_present = self.store.isExist(key) == 1 + success = not actual_present and (result_code == 0 or not expected_present) + else: + raise ValueError(f"unsupported metadata operation: {operation}") + + errors = Counter() + if not success: + errors[result_code] += 1 + return ( + RequestResult( + request_ok=success, + kv_successes=int(success), + kv_failures=int(not success), + bytes_processed=bytes_processed, + verify_failures=verify_failures, + error_counts=errors, + ), + actual_present, + result_code, + ) + def _put_keys(self, keys: List[str], object_ids: List[int]) -> List[int]: values = [self.payload_factory.build(object_id) for object_id in object_ids] if self.args.io_api == "plain": @@ -381,14 +605,14 @@ def __init__(self, store_obj, value_size: int, slots: int): alloc_addr = get_alloc_func_addr() free_addr = get_free_func_addr() if alloc_addr is None or free_addr is None: - raise RuntimeError("store module does not expose hugepage alloc/free helpers") + raise RuntimeError( + "store module does not expose hugepage alloc/free helpers" + ) self._alloc_fn = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_size_t)( get_alloc_func_addr() ) - self._free_fn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)( - get_free_func_addr() - ) + self._free_fn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(get_free_func_addr()) raw_ptr = self._alloc_fn(self.total_size) self.base_ptr = ctypes.cast(raw_ptr, ctypes.c_void_p).value or 0 @@ -414,7 +638,11 @@ def close(self) -> None: try: self.store.unregister_buffer(self.base_ptr) except Exception: - LOG.debug("unregister_buffer failed for direct zcopy pool ptr=%s", self.base_ptr, exc_info=True) + LOG.debug( + "unregister_buffer failed for direct zcopy pool ptr=%s", + self.base_ptr, + exc_info=True, + ) self._registered = False if self._free_fn is not None: self._free_fn(ctypes.c_void_p(self.base_ptr)) @@ -434,7 +662,9 @@ def __init__(self, pool: ZcopyBufferPool, slot_offset: int, slots: int): def _slot_ptr(self, slot: int) -> int: if slot < 0 or slot >= self.slots: - raise IndexError(f"zcopy view slot {slot} is out of range [0, {self.slots})") + raise IndexError( + f"zcopy view slot {slot} is out of range [0, {self.slots})" + ) return self.pool.slot_ptr(self.slot_offset + slot) def fill_write_buffers(self, payloads: List[bytes]) -> List[int]: @@ -467,7 +697,6 @@ def read_bytes(self, slot: int, size: int) -> bytes: class StoreRuntime: def __init__(self, args: argparse.Namespace, lane_count: int): - self.lane_count = lane_count self.store = MooncakeDistributedStore() setup_ret = self.store.setup( @@ -485,9 +714,7 @@ def __init__(self, args: argparse.Namespace, lane_count: int): self.zcopy_pool: Optional[ZcopyBufferPool] = None if args.io_api == "zcopy": slots = max(1, args.batch_size) * lane_count - self.zcopy_pool = ZcopyBufferPool( - self.store, args.value_size, slots - ) + self.zcopy_pool = ZcopyBufferPool(self.store, args.value_size, slots) def make_session( self, @@ -529,10 +756,13 @@ def merge_stats(name: str, stats_list: List[PhaseStats]) -> PhaseStats: merged = PhaseStats(name=name) if not stats_list: return merged - merged.start_time = min((s.start_time for s in stats_list if s.start_time), default=0.0) + merged.start_time = min( + (s.start_time for s in stats_list if s.start_time), default=0.0 + ) merged.end_time = max((s.end_time for s in stats_list if s.end_time), default=0.0) for stats in stats_list: merged.request_latencies.extend(stats.request_latencies) + merged.latency_seen += stats.latency_seen merged.requests += stats.requests merged.successful_requests += stats.successful_requests merged.failed_requests += stats.failed_requests @@ -576,11 +806,16 @@ def summarize_stats(stats: PhaseStats) -> dict: "duration_sec": duration, "req_per_sec": (stats.requests / duration) if duration > 0 else 0.0, "kv_per_sec": (stats.kvs / duration) if duration > 0 else 0.0, - "MiB_per_sec": (stats.bytes_processed / duration / (1024 * 1024)) if duration > 0 else 0.0, - "lat_mean_ms": statistics.mean(stats.request_latencies) * 1000 if stats.request_latencies else 0.0, + "MiB_per_sec": (stats.bytes_processed / duration / (1024 * 1024)) + if duration > 0 + else 0.0, + "lat_mean_ms": statistics.mean(stats.request_latencies) * 1000 + if stats.request_latencies + else 0.0, "lat_p50_ms": percentile(stats.request_latencies, 0.50) * 1000, "lat_p95_ms": percentile(stats.request_latencies, 0.95) * 1000, "lat_p99_ms": percentile(stats.request_latencies, 0.99) * 1000, + "latency_samples": len(stats.request_latencies), "dataset_exhausted": stats.dataset_exhausted, "error_counts": dict(stats.error_counts), } @@ -629,6 +864,7 @@ def __init__(self, args: argparse.Namespace): self.lane_count = args.numjobs * args.iodepth self._sessions: Optional[List[StoreSession]] = None self._runtime: Optional[StoreRuntime] = None + self._journal_records: List[List[dict]] = [[] for _ in range(self.lane_count)] self._validate_args() def _validate_args(self) -> None: @@ -648,6 +884,20 @@ def _validate_args(self) -> None: raise ValueError("prepare-objects must be >= 0") if self.args.rwmixread < 0 or self.args.rwmixread > 100: raise ValueError("rwmixread must be within [0, 100]") + validate_operation_percentages( + { + "put": self.args.put_pct, + "get": self.args.get_pct, + "exist": self.args.exist_pct, + "remove": self.args.remove_pct, + } + ) + if self.args.latency_sample_rate <= 0: + raise ValueError("latency-sample-rate must be > 0") + if self.args.tenant_id != "default": + raise ValueError("only --tenant-id=default is currently supported") + if self.args.scenario == "replay" and not self.args.replay_file: + raise ValueError("replay requires --replay-file") if self.args.verify and not self.pattern: raise ValueError("verify mode currently requires --pattern") if self.args.memory_replica_num == 0 and self.args.nof_replica_num == 0: @@ -659,17 +909,36 @@ def _validate_args(self) -> None: if self.args.scenario == "mixed_rw" and self.args.runtime <= 0: raise ValueError("mixed_rw requires --runtime > 0") if self._scenario_has_write() and self.args.value_size % 512 != 0: - raise ValueError("write-involved scenarios require value-size to be 512B aligned") + raise ValueError( + "write-involved scenarios require value-size to be 512B aligned" + ) make_key(self.args.key_prefix, self.args.key_size, self.args.object_id_start) def _scenario_has_write(self) -> bool: - return self.args.scenario in {"verify_write", "fill", "write_perf", "mixed_rw"} + return self.args.scenario in { + "verify_write", + "fill", + "write_perf", + "mixed_rw", + "metadata_smoke", + "mixed_metadata", + "remove_perf", + "replay", + } def _write_budget(self) -> int: - return self.args.write_objects if self.args.write_objects > 0 else self.args.nr_objects + return ( + self.args.write_objects + if self.args.write_objects > 0 + else self.args.nr_objects + ) def _prepare_budget(self) -> int: - return self.args.prepare_objects if self.args.prepare_objects > 0 else self.args.nr_objects + return ( + self.args.prepare_objects + if self.args.prepare_objects > 0 + else self.args.nr_objects + ) def _make_sessions(self) -> List[StoreSession]: if self._sessions is None: @@ -699,17 +968,29 @@ def _phase_gap(self, label: str) -> None: time.sleep(self.args.phase_gap_sec) return if mode == "manual": - input(f"phase '{label}' is waiting, finish external operations then press Enter to continue...") + input( + f"phase '{label}' is waiting, finish external operations then press Enter to continue..." + ) return deadline = time.time() + self.args.phase_gap_timeout_sec while time.time() < deadline: if os.path.exists(self.args.phase_gap_file): - LOG.info("detected phase gap file %s, continuing to %s", self.args.phase_gap_file, label) + LOG.info( + "detected phase gap file %s, continuing to %s", + self.args.phase_gap_file, + label, + ) return time.sleep(1.0) - raise TimeoutError(f"timed out waiting for phase gap file {self.args.phase_gap_file}") + raise TimeoutError( + f"timed out waiting for phase gap file {self.args.phase_gap_file}" + ) - def _run_threads(self, phase_name: str, worker_builder: Callable[[StoreSession, int], Callable[[PhaseStats], None]]) -> PhaseStats: + def _run_threads( + self, + phase_name: str, + worker_builder: Callable[[StoreSession, int], Callable[[PhaseStats], None]], + ) -> PhaseStats: sessions = self._make_sessions() per_lane_stats: List[Optional[PhaseStats]] = [None] * self.lane_count threads: List[threading.Thread] = [] @@ -722,7 +1003,11 @@ def runner(index: int, session: StoreSession) -> None: per_lane_stats[index] = stats for lane_id, session in enumerate(sessions): - thread = threading.Thread(target=runner, args=(lane_id, session), name=f"{phase_name}-lane{lane_id}") + thread = threading.Thread( + target=runner, + args=(lane_id, session), + name=f"{phase_name}-lane{lane_id}", + ) threads.append(thread) thread.start() @@ -733,8 +1018,15 @@ def runner(index: int, session: StoreSession) -> None: log_phase_stats(merged) return merged - def _record(self, stats: PhaseStats, latency: float, request: RequestResult, kv_count: int) -> None: - stats.request_latencies.append(latency) + def _record( + self, stats: PhaseStats, latency: float, request: RequestResult, kv_count: int + ) -> None: + if ( + self.args.runtime <= 0 + or stats.latency_seen % self.args.latency_sample_rate == 0 + ): + stats.request_latencies.append(latency) + stats.latency_seen += 1 stats.requests += 1 stats.kvs += kv_count if request.request_ok: @@ -748,6 +1040,154 @@ def _record(self, stats: PhaseStats, latency: float, request: RequestResult, kv_ stats.bytes_processed += request.bytes_processed stats.error_counts.update(request.error_counts) + def _run_metadata_operations( + self, phase_name: str, operations_by_lane + ) -> PhaseStats: + def worker(session: StoreSession, lane_id: int) -> Callable[[PhaseStats], None]: + def run(stats: PhaseStats) -> None: + model = MetadataLaneModel() + for op_id, (operation, object_id) in enumerate( + operations_by_lane(lane_id) + ): + expected_present = model.is_present(object_id) + invoke_ns = time.time_ns() + start = time.perf_counter() + result, actual_present, result_code = session.metadata_operation( + operation, object_id, expected_present + ) + return_ns = time.time_ns() + self._record( + stats, + time.perf_counter() - start, + result, + 1, + ) + model.record(object_id, operation, result.request_ok) + if self.args.journal == "all" or ( + self.args.journal == "failures" and not result.request_ok + ): + self._journal_records[lane_id].append( + { + "schema_version": 1, + "op_id": op_id, + "lane": lane_id, + "op": operation, + "object_id": object_id, + "key": make_key( + self.args.key_prefix, + self.args.key_size, + object_id, + ), + "invoke_ns": invoke_ns, + "return_ns": return_ns, + "result": result_code, + "ok": result.request_ok, + "expected_present": expected_present, + "actual_present": actual_present, + } + ) + + return run + + return self._run_threads(phase_name, worker) + + def _metadata_smoke(self) -> PhaseStats: + sequence = ("put", "exist", "get", "remove", "exist") + + def operations(lane_id: int): + object_id = self.args.object_id_start + lane_id + if object_id >= self.args.object_id_start + self.args.nr_objects: + return [] + return [(operation, object_id) for operation in sequence] + + stats = self._run_metadata_operations("metadata_smoke", operations) + if stats.failed_kvs: + raise RuntimeError(f"metadata_smoke failed operations={stats.failed_kvs}") + return stats + + def _mixed_metadata(self) -> PhaseStats: + weights = { + "put": self.args.put_pct, + "get": self.args.get_pct, + "exist": self.args.exist_pct, + "remove": self.args.remove_pct, + } + + def operations(lane_id: int): + object_ids = list( + range( + self.args.object_id_start + lane_id, + self.args.object_id_start + self.args.nr_objects, + self.lane_count, + ) + ) + if not object_ids: + return [] + count = len(object_ids) + choices = choose_metadata_operations( + self.args.rand_seed, lane_id, count, weights + ) + return [ + (operation, object_ids[index % len(object_ids)]) + for index, operation in enumerate(choices) + ] + + return self._run_metadata_operations("mixed_metadata", operations) + + def _remove_prepared(self) -> PhaseStats: + object_ids = self.dataset.snapshot_written_ids() + + def operations(lane_id: int): + return [ + ("remove", object_id) + for index, object_id in enumerate(object_ids) + if index % self.lane_count == lane_id + ] + + return self._run_metadata_operations("remove_perf", operations) + + def _replay(self) -> PhaseStats: + records = read_replay(self.args.replay_file) + + def operations(lane_id: int): + return [ + (record["op"], int(record["object_id"])) + for record in records + if int(record["lane"]) == lane_id + ] + + return self._run_metadata_operations("replay", operations) + + def write_outputs(self, phases: List[PhaseStats]) -> None: + records = sorted( + (record for lane in self._journal_records for record in lane), + key=lambda record: ( + record["invoke_ns"], + record["lane"], + record["op_id"], + ), + ) + if self.args.output_dir: + os.makedirs(self.args.output_dir, exist_ok=True) + if self.args.journal != "none" and self.args.output_dir: + write_jsonl(os.path.join(self.args.output_dir, "journal.jsonl"), records) + summary_path = self.args.summary_json + if not summary_path and self.args.output_dir: + summary_path = os.path.join(self.args.output_dir, "summary.json") + if summary_path: + overall = merge_stats("overall", phases) + write_json_atomic( + summary_path, + { + "schema_version": 1, + "scenario": self.args.scenario, + "ok": overall.failed_kvs == 0 and overall.verify_failures == 0, + "phases": {phase.name: summarize_stats(phase) for phase in phases}, + "overall": summarize_stats(overall), + "journal_records": len(records), + }, + ) + def _run_fixed_write( self, phase_name: str, @@ -758,10 +1198,14 @@ def _run_fixed_write( ) -> PhaseStats: write_upper = self.dataset.next_write_id + total_objects - def worker(session: StoreSession, _lane_id: int) -> Callable[[PhaseStats], None]: + def worker( + session: StoreSession, _lane_id: int + ) -> Callable[[PhaseStats], None]: def run(stats: PhaseStats) -> None: while True: - object_ids = self.dataset.reserve_write_ids(self.args.batch_size, write_upper) + object_ids = self.dataset.reserve_write_ids( + self.args.batch_size, write_upper + ) if not object_ids: break start = time.perf_counter() @@ -772,7 +1216,10 @@ def run(stats: PhaseStats) -> None: if write_scope == "prepared": self.dataset.mark_prepared(result.successful_object_ids) else: - self.dataset.mark_runtime_written(result.successful_object_ids) + self.dataset.mark_runtime_written( + result.successful_object_ids + ) + return run stats = self._run_threads(phase_name, worker) @@ -791,10 +1238,14 @@ def _run_time_based_write(self, phase_name: str, total_objects: int) -> PhaseSta write_upper = self.dataset.next_write_id + total_objects stop_event = threading.Event() - def worker(session: StoreSession, _lane_id: int) -> Callable[[PhaseStats], None]: + def worker( + session: StoreSession, _lane_id: int + ) -> Callable[[PhaseStats], None]: def run(stats: PhaseStats) -> None: while time.time() < deadline and not stop_event.is_set(): - object_ids = self.dataset.reserve_write_ids(self.args.batch_size, write_upper) + object_ids = self.dataset.reserve_write_ids( + self.args.batch_size, write_upper + ) if not object_ids: stats.dataset_exhausted = True stop_event.set() @@ -805,6 +1256,7 @@ def run(stats: PhaseStats) -> None: self._record(stats, latency, result, len(object_ids)) if result.successful_object_ids: self.dataset.mark_runtime_written(result.successful_object_ids) + return run return self._run_threads(phase_name, worker) @@ -822,7 +1274,9 @@ def _run_read_phase( if runtime_sec > 0: deadline = time.time() + runtime_sec - def worker(session: StoreSession, lane_id: int) -> Callable[[PhaseStats], None]: + def worker( + session: StoreSession, lane_id: int + ) -> Callable[[PhaseStats], None]: rng = random.Random(seed_base + lane_id) def run(stats: PhaseStats) -> None: @@ -897,7 +1351,9 @@ def run(stats: PhaseStats) -> None: self._record(stats, latency, result, len(object_ids)) continue - object_ids = self.dataset.reserve_write_ids(self.args.batch_size, write_upper) + object_ids = self.dataset.reserve_write_ids( + self.args.batch_size, write_upper + ) if not object_ids: stats.dataset_exhausted = True stop_event.set() @@ -916,7 +1372,10 @@ def run(stats: PhaseStats) -> None: def _maybe_prepare_dataset(self) -> Optional[PhaseStats]: if self.args.prepare_mode == "none": return None - if self.args.prepare_mode == "write" or self.args.scenario in {"read_perf", "mixed_rw"}: + if self.args.prepare_mode == "write" or self.args.scenario in { + "read_perf", + "mixed_rw", + }: stats = self._run_fixed_write( "prepare_write", self._prepare_budget(), @@ -956,11 +1415,17 @@ def run(self) -> List[PhaseStats]: ) ) self._phase_gap("verify_read") - phases.append(self._run_read_phase("verify_read", verify=True, sequential=True, loop=False)) + phases.append( + self._run_read_phase( + "verify_read", verify=True, sequential=True, loop=False + ) + ) return phases if self.args.scenario == "fill": - phases.append(self._run_fixed_write("fill_write", self._write_budget(), strict=False)) + phases.append( + self._run_fixed_write("fill_write", self._write_budget(), strict=False) + ) return phases if self.args.scenario == "write_perf": @@ -968,7 +1433,33 @@ def run(self) -> List[PhaseStats]: if self.args.runtime > 0: phases.append(self._run_time_based_write("write_perf", total_objects)) else: - phases.append(self._run_fixed_write("write_perf", total_objects, strict=False)) + phases.append( + self._run_fixed_write("write_perf", total_objects, strict=False) + ) + return phases + + if self.args.scenario == "metadata_smoke": + phases.append(self._metadata_smoke()) + return phases + + if self.args.scenario == "mixed_metadata": + phases.append(self._mixed_metadata()) + return phases + + if self.args.scenario == "remove_perf": + phases.append( + self._run_fixed_write( + "prepare_write", + self._prepare_budget(), + strict=True, + write_scope="prepared", + ) + ) + phases.append(self._remove_prepared()) + return phases + + if self.args.scenario == "replay": + phases.append(self._replay()) return phases if self.args.scenario == "read_perf": @@ -1008,10 +1499,20 @@ def main() -> int: runner = BenchmarkRunner(args) phases = runner.run() log_overall_summary(phases) + runner.write_outputs(phases) if any(phase.verify_failures > 0 for phase in phases): return 20 - if args.verify and any(phase.misses > 0 for phase in phases if "read" in phase.name): + if args.verify and any( + phase.misses > 0 for phase in phases if "read" in phase.name + ): return 21 + if args.scenario in { + "metadata_smoke", + "mixed_metadata", + "remove_perf", + "replay", + } and any(phase.failed_kvs > 0 for phase in phases): + return 22 return 0 except KeyboardInterrupt: LOG.warning("benchmark interrupted") diff --git a/mooncake-store/benchmarks/test_report_oplog_batch.py b/mooncake-store/benchmarks/test_report_oplog_batch.py new file mode 100644 index 0000000000..612e4d80f7 --- /dev/null +++ b/mooncake-store/benchmarks/test_report_oplog_batch.py @@ -0,0 +1,87 @@ +import csv +import importlib.util +import json +import tempfile +import unittest +from pathlib import Path + + +MODULE_PATH = Path(__file__).with_name("report_oplog_batch.py") +SPEC = importlib.util.spec_from_file_location("report_oplog_batch", MODULE_PATH) +report = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(report) + + +class ReportTest(unittest.TestCase): + def test_summarizes_repeats_and_stage_limits(self): + with tempfile.TemporaryDirectory() as root: + root = Path(root) + samples = root / "samples.jsonl" + rows = [] + for repeat, entries_per_sec in enumerate((9000, 10000, 11000)): + rows.append( + { + "schema_version": 1, + "case": "normal-e64-b8", + "repeat": repeat, + "entries": 8000, + "entries_per_sec": entries_per_sec, + "transactions_per_sec": entries_per_sec / 8, + "cpu_cores_used": 0.75 + repeat * 0.25, + "batch_entries_mean": 8, + "encoded_bytes": 800000, + "stage_latency_us": { + "queue": 100, + "encode": 200, + "txn": 500, + "callback": 200, + }, + } + ) + samples.write_text("".join(json.dumps(row) + "\n" for row in rows)) + + summary = report.generate_report(samples, root / "out") + + case = summary["cases"][0] + self.assertEqual(case["entries_per_sec"]["median"], 10000) + self.assertEqual(case["entries_per_sec"]["min"], 9000) + self.assertEqual(case["entries_per_sec"]["max"], 11000) + self.assertAlmostEqual(case["bytes_per_entry"], 100.0) + self.assertEqual(case["cpu_cores_used"]["median"], 1.0) + self.assertAlmostEqual(case["stage_share"]["txn"], 0.5) + self.assertEqual(case["theoretical_entries_per_sec"], 16000) + self.assertAlmostEqual(case["observed_efficiency"], 0.625) + self.assertEqual(case["bottleneck_stage"], "txn") + self.assertTrue((root / "out" / "report.md").is_file()) + with (root / "out" / "cases.csv").open(newline="") as stream: + self.assertEqual( + list(csv.DictReader(stream))[0]["case"], "normal-e64-b8" + ) + + def test_marks_missing_measurements_as_insufficient(self): + with tempfile.TemporaryDirectory() as root: + root = Path(root) + samples = root / "samples.jsonl" + samples.write_text( + json.dumps( + { + "schema_version": 1, + "case": "partial", + "repeat": 0, + "entries_per_sec": 42, + } + ) + + "\n" + ) + + case = report.generate_report(samples, root / "out")["cases"][0] + + self.assertEqual(case["status"], "insufficient_data") + self.assertIn("at least 3 repeats", case["missing"]) + self.assertIn("encoded_bytes and entries", case["missing"]) + self.assertIn("stage_latency_us.txn", case["missing"]) + self.assertNotIn("theoretical_entries_per_sec", case) + + +if __name__ == "__main__": + unittest.main() diff --git a/mooncake-store/benchmarks/test_run_oplog_batch_sweep.py b/mooncake-store/benchmarks/test_run_oplog_batch_sweep.py new file mode 100644 index 0000000000..c641da96e2 --- /dev/null +++ b/mooncake-store/benchmarks/test_run_oplog_batch_sweep.py @@ -0,0 +1,109 @@ +import fcntl +import importlib.util +import json +import stat +import tempfile +import unittest +from pathlib import Path + + +MODULE_PATH = Path(__file__).with_name("run_oplog_batch_sweep.py") +SPEC = importlib.util.spec_from_file_location("run_oplog_batch_sweep", MODULE_PATH) +sweep = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(sweep) + + +class SweepTest(unittest.TestCase): + def test_expands_matrix_and_keeps_failed_runs(self): + with tempfile.TemporaryDirectory() as root: + root = Path(root) + fake = root / "fake_bench.py" + fake.write_text( + "#!/usr/bin/env python3\n" + "import json, pathlib, sys\n" + "args=dict(x[2:].split('=', 1) for x in sys.argv[1:])\n" + "attempts=pathlib.Path(__file__).with_name('attempts')\n" + "count=int(attempts.read_text()) if attempts.exists() else 0\n" + "attempts.write_text(str(count + 1))\n" + "result={'schema_version':1,'entries_per_sec':int(args['max_entries'])*100}\n" + "open(args['output_json'],'w').write(json.dumps(result))\n" + "raise SystemExit(7 if args['max_entries']=='8' and count < 4 else 0)\n" + ) + fake.chmod(fake.stat().st_mode | stat.S_IXUSR) + matrix = root / "matrix.json" + matrix.write_text( + json.dumps( + { + "cases": [ + {"name": "b1", "mode": "writer", "max_entries": 1}, + {"name": "b8", "mode": "writer", "max_entries": 8}, + ] + } + ) + ) + + exit_code = sweep.run_sweep( + benchmark=fake, + output_dir=root / "out", + repeat=2, + matrix_path=matrix, + endpoints="127.0.0.1:2379", + ) + + self.assertEqual(exit_code, 1) + rows = [ + json.loads(line) + for line in (root / "out/samples.jsonl").read_text().splitlines() + ] + self.assertEqual(len(rows), 4) + self.assertEqual([row["case"] for row in rows], ["b1", "b1", "b8", "b8"]) + self.assertEqual([row["repeat"] for row in rows], [0, 1, 0, 1]) + self.assertEqual([row["exit_code"] for row in rows], [0, 0, 7, 7]) + self.assertEqual(len({row["cluster_id"] for row in rows}), 4) + manifest = json.loads((root / "out/manifest.json").read_text()) + self.assertGreater(manifest["host"]["cpu_count"], 0) + self.assertTrue(manifest["host"]["machine"]) + + resume_exit_code = sweep.run_sweep( + benchmark=fake, + output_dir=root / "out", + repeat=2, + matrix_path=matrix, + endpoints="127.0.0.1:2379", + resume=True, + ) + self.assertEqual(resume_exit_code, 0) + resumed_rows = [ + json.loads(line) + for line in (root / "out/samples.jsonl").read_text().splitlines() + ] + self.assertEqual( + [ + (row["case"], row["repeat"], row["exit_code"]) + for row in resumed_rows + ], + [ + ("b1", 0, 0), + ("b1", 1, 0), + ("b8", 0, 7), + ("b8", 1, 7), + ("b8", 0, 0), + ("b8", 1, 0), + ], + ) + + with (root / "out/.sweep.lock").open("w") as lock_stream: + fcntl.flock(lock_stream, fcntl.LOCK_EX | fcntl.LOCK_NB) + with self.assertRaisesRegex(RuntimeError, "already running"): + sweep.run_sweep( + benchmark=fake, + output_dir=root / "out", + repeat=2, + matrix_path=matrix, + endpoints="127.0.0.1:2379", + resume=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/mooncake-store/benchmarks/test_store_kv_bench.py b/mooncake-store/benchmarks/test_store_kv_bench.py new file mode 100644 index 0000000000..2ede5f94f4 --- /dev/null +++ b/mooncake-store/benchmarks/test_store_kv_bench.py @@ -0,0 +1,174 @@ +import importlib.util +import json +import os +import sys +import types +import unittest +from pathlib import Path +from types import SimpleNamespace + + +store_module = types.ModuleType("mooncake.store") +store_module.MooncakeDistributedStore = object +store_module.ReplicateConfig = type("ReplicateConfig", (), {}) +store_module.get_alloc_func_addr = lambda: None +store_module.get_free_func_addr = lambda: None +sys.modules.setdefault("mooncake", types.ModuleType("mooncake")) +sys.modules["mooncake.store"] = store_module + +MODULE_PATH = Path(__file__).with_name("store_kv_bench.py") +SPEC = importlib.util.spec_from_file_location("store_kv_bench_under_test", MODULE_PATH) +bench = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = bench +SPEC.loader.exec_module(bench) + + +class MetadataWorkloadModelTest(unittest.TestCase): + def test_percentages_and_lane_choices_are_deterministic(self): + weights = {"put": 40, "get": 30, "exist": 20, "remove": 10} + bench.validate_operation_percentages(weights) + first = bench.choose_metadata_operations(7, 3, 20, weights) + second = bench.choose_metadata_operations(7, 3, 20, weights) + self.assertEqual(first, second) + self.assertNotEqual(first, bench.choose_metadata_operations(7, 4, 20, weights)) + with self.assertRaisesRegex(ValueError, "sum to 100"): + bench.validate_operation_percentages({**weights, "remove": 9}) + + def test_model_tracks_successful_transitions(self): + model = bench.MetadataLaneModel() + self.assertFalse(model.is_present(4)) + model.record(4, "put", True) + self.assertTrue(model.is_present(4)) + model.record(4, "remove", False) + self.assertTrue(model.is_present(4)) + model.record(4, "remove", True) + self.assertFalse(model.is_present(4)) + + +class OutputTest(unittest.TestCase): + def test_journal_replay_summary_and_sampling(self): + root = Path(f"/tmp/mooncake-kv-bench-test-{os.getpid()}") + root.mkdir(parents=True, exist_ok=True) + records = [ + { + "schema_version": 1, + "op_id": 1, + "lane": 0, + "op": "put", + "object_id": 9, + "key": "k9", + "invoke_ns": 10, + "return_ns": 20, + "result": 0, + } + ] + journal = root / "journal.jsonl" + bench.write_jsonl(journal, records) + self.assertEqual(bench.read_replay(journal), records) + + summary = root / "summary.json" + bench.write_json_atomic(summary, {"ok": True, "requests": 1}) + self.assertEqual(json.loads(summary.read_text()), {"ok": True, "requests": 1}) + + sampler = bench.LatencySampler(3) + for value in range(10): + sampler.add(float(value)) + self.assertEqual(sampler.samples, [0.0, 3.0, 6.0, 9.0]) + self.assertEqual(sampler.seen, 10) + + +class StoreSessionMetadataTest(unittest.TestCase): + def test_put_get_exist_remove(self): + class FakeStore: + def __init__(self): + self.values = {} + + def put(self, key, value, _config): + self.values[key] = value + return 0 + + def get(self, key): + return self.values.get(key, b"") + + def isExist(self, key): + return int(key in self.values) + + def remove(self, key): + return 0 if self.values.pop(key, None) is not None else -1 + + args = SimpleNamespace( + memory_replica_num=1, + nof_replica_num=0, + key_prefix="metadata", + key_size=24, + value_size=512, + io_api="plain", + batch_size=1, + ) + store = FakeStore() + session = bench.StoreSession(args, 0, bench.PayloadFactory(512, b"x"), store) + + for operation, expected_before, expected_after in [ + ("put", False, True), + ("exist", True, True), + ("get", True, True), + ("remove", True, False), + ("exist", False, False), + ]: + result, actual_present, _ = session.metadata_operation( + operation, 1, expected_before + ) + self.assertTrue(result.request_ok, operation) + self.assertEqual(actual_present, expected_after, operation) + + def test_metadata_smoke_writes_replayable_outputs(self): + class FakeStore: + def __init__(self): + self.values = {} + + def put(self, key, value, _config): + self.values[key] = value + return 0 + + def get(self, key): + return self.values.get(key, b"") + + def isExist(self, key): + return int(key in self.values) + + def remove(self, key): + return 0 if self.values.pop(key, None) is not None else -1 + + output_dir = Path(f"/tmp/mooncake-kv-smoke-test-{os.getpid()}") + args = bench.build_parser().parse_args( + [ + "--scenario=metadata_smoke", + "--nr-objects=1", + "--value-size=512", + "--key-size=24", + "--pattern=x", + f"--output-dir={output_dir}", + "--journal=all", + ] + ) + runner = bench.BenchmarkRunner(args) + runner._sessions = [ + bench.StoreSession(args, 0, runner.payload_factory, FakeStore()) + ] + phases = runner.run() + runner.write_outputs(phases) + + self.assertEqual(phases[0].requests, 5) + self.assertEqual(phases[0].failed_kvs, 0) + records = bench.read_replay(output_dir / "journal.jsonl") + self.assertEqual( + [record["op"] for record in records], + ["put", "exist", "get", "remove", "exist"], + ) + summary = json.loads((output_dir / "summary.json").read_text()) + self.assertTrue(summary["ok"]) + self.assertEqual(summary["journal_records"], 5) + + +if __name__ == "__main__": + unittest.main() diff --git a/mooncake-store/conf/master.json b/mooncake-store/conf/master.json index a376e72fcb..f55bdaf2b2 100644 --- a/mooncake-store/conf/master.json +++ b/mooncake-store/conf/master.json @@ -2,23 +2,32 @@ "enable_metric_reporting": true, "metrics_port": 9003, "rpc_port": 50051, - "rpc_thread_num": 4, + "rpc_thread_num": 16, "rpc_address": "0.0.0.0", "rpc_interface": "", "rpc_conn_timeout_seconds": 0, "rpc_enable_tcp_no_delay": true, - "default_kv_lease_ttl": 5000, + "default_kv_lease_ttl": 10000, "default_kv_soft_pin_ttl": 1800000, "allow_evict_soft_pinned_objects": true, "eviction_ratio": 0.1, "eviction_high_watermark_ratio": 1.0, "enable_ha": false, + "enable_oplog": false, "etcd_endpoints": "http://localhost:2379", "root_fs_dir": "", "cluster_id": "mooncake_cluster", "memory_allocator": "offset", - "client_live_ttl_sec": 60, + "client_live_ttl_sec": 60, "enable_http_metadata_server": false, "http_metadata_server_host": "0.0.0.0", - "http_metadata_server_port": 8080 + "http_metadata_server_port": 8080, + "enable_kv_events": false, + "kv_events_bind_endpoint": "tcp://0.0.0.0:5557", + "kv_events_model_name": "", + "kv_events_backend_id": "", + "kv_events_tenant_id": "default", + "kv_events_block_size": 0, + "kv_events_dp_rank": 0, + "kv_events_emit_object_key": true } diff --git a/mooncake-store/conf/master.yaml b/mooncake-store/conf/master.yaml index 02d06b9e05..545807008f 100644 --- a/mooncake-store/conf/master.yaml +++ b/mooncake-store/conf/master.yaml @@ -1,18 +1,24 @@ enable_metric_reporting: true metrics_port: 9003 rpc_port: 50051 -rpc_thread_num: 4 +rpc_thread_num: 16 rpc_address: "0.0.0.0" rpc_interface: "" rpc_conn_timeout_seconds: 0 rpc_enable_tcp_no_delay: true -default_kv_lease_ttl: 5000 +default_kv_lease_ttl: 10000 default_kv_soft_pin_ttl: 1800000 allow_evict_soft_pinned_objects: true eviction_ratio: 0.1 +# Overrides the 0.90 code default. A value of 1.0 disables proactive +# usage-ratio eviction; allocation-failure-triggered eviction remains enabled. eviction_high_watermark_ratio: 1.0 +enable_multi_tenants: false +tenant_quota_connector_type: "file" +tenant_quota_connector_uri: "" + enable_ha: false etcd_endpoints: "http://localhost:2379" root_fs_dir: "" @@ -20,6 +26,13 @@ cluster_id: "mooncake_cluster" memory_allocator: "offset" client_live_ttl_sec: 60 +# OpLog store configuration for HA hot-standby replication +enable_oplog: false +# Batch standby base polling/retry delay in milliseconds +oplog_poll_interval_ms: 1000 +# Maximum consecutive retryable batch-standby failure window in seconds +batch_oplog_retry_timeout_sec: 180 + enable_http_metadata_server: false http_metadata_server_host: "0.0.0.0" http_metadata_server_port: 8080 diff --git a/mooncake-store/go/build.sh b/mooncake-store/go/build.sh index 52304cc7e3..df9a9d6dd7 100755 --- a/mooncake-store/go/build.sh +++ b/mooncake-store/go/build.sh @@ -41,7 +41,8 @@ CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src" CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-transfer-engine/src/common/base" CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-common" CGO_LDFLAGS+=" -L${BUILD_DIR}/mooncake-common/src" -CGO_LDFLAGS+=" -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lasio -lmooncake_common -lxxhash" +CGO_LDFLAGS+=" -Wl,--start-group -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase -lmooncake_common -Wl,--end-group" +CGO_LDFLAGS+=" -lasio -lxxhash -lyaml-cpp" CGO_LDFLAGS+=" -lstdc++ -lnuma -lglog -lgflags -libverbs -lmlx5 -ljsoncpp -lzstd -lcurl -lm" if [ -d "/usr/local/cuda/lib64" ]; then @@ -50,6 +51,13 @@ fi CGO_LDFLAGS+=" -luring" +# KV events publisher (optional; linked when libzmq is installed). +if ldconfig -p 2>/dev/null | grep -q libzmq \ + || [ -f /usr/lib/x86_64-linux-gnu/libzmq.so ] \ + || [ -f /usr/lib/libzmq.so ]; then + CGO_LDFLAGS+=" -lzmq" +fi + if [ "$USE_ETCD" = "ON" ]; then if [ "$USE_ETCD_LEGACY" = "ON" ]; then CGO_LDFLAGS+=" -letcd-cpp-api -lprotobuf -lgrpc++ -lgrpc" diff --git a/mooncake-store/go/mooncakestore/errors.go b/mooncake-store/go/mooncakestore/errors.go index 6c673e281f..0fee47eab6 100644 --- a/mooncake-store/go/mooncakestore/errors.go +++ b/mooncake-store/go/mooncakestore/errors.go @@ -17,18 +17,18 @@ package mooncakestore import "errors" var ( - ErrStoreNil = errors.New("mooncakestore: store handle is nil") - ErrSetupFailed = errors.New("mooncakestore: setup failed") - ErrInitAllFailed = errors.New("mooncakestore: init_all failed") - ErrHealthCheck = errors.New("mooncakestore: health check failed") - ErrPut = errors.New("mooncakestore: put failed") - ErrGet = errors.New("mooncakestore: get failed") - ErrRemove = errors.New("mooncakestore: remove failed") - ErrExist = errors.New("mooncakestore: existence check failed") - ErrGetSize = errors.New("mooncakestore: get size failed") + ErrStoreNil = errors.New("mooncakestore: store handle is nil") + ErrSetupFailed = errors.New("mooncakestore: setup failed") + ErrInitAllFailed = errors.New("mooncakestore: init_all failed") + ErrHealthCheck = errors.New("mooncakestore: health check failed") + ErrPut = errors.New("mooncakestore: put failed") + ErrGet = errors.New("mooncakestore: get failed") + ErrRemove = errors.New("mooncakestore: remove failed") + ErrExist = errors.New("mooncakestore: existence check failed") + ErrGetSize = errors.New("mooncakestore: get size failed") ErrRegisterBuffer = errors.New("mooncakestore: register buffer failed") ErrUnregisterBuffer = errors.New("mooncakestore: unregister buffer failed") ErrBatchOp = errors.New("mooncakestore: batch operation failed") - ErrHostname = errors.New("mooncakestore: get hostname failed") - ErrInvalidArgument = errors.New("mooncakestore: invalid argument") + ErrHostname = errors.New("mooncakestore: get hostname failed") + ErrInvalidArgument = errors.New("mooncakestore: invalid argument") ) diff --git a/mooncake-store/include/aligned_client_buffer.hpp b/mooncake-store/include/aligned_client_buffer.h similarity index 98% rename from mooncake-store/include/aligned_client_buffer.hpp rename to mooncake-store/include/aligned_client_buffer.h index 436fc9bb92..1cfd44ecb7 100644 --- a/mooncake-store/include/aligned_client_buffer.hpp +++ b/mooncake-store/include/aligned_client_buffer.h @@ -1,6 +1,6 @@ #pragma once -#include "client_buffer.hpp" +#include "client_buffer.h" namespace mooncake { diff --git a/mooncake-store/include/allocation_strategy.h b/mooncake-store/include/allocation_strategy.h index e9f656126c..93f4dcaf68 100644 --- a/mooncake-store/include/allocation_strategy.h +++ b/mooncake-store/include/allocation_strategy.h @@ -2,7 +2,6 @@ #include #include -#include #include #include #include @@ -13,6 +12,7 @@ #include "allocator.h" // Contains BufferAllocator declaration #include "replica.h" #include "types.h" +#include "random.h" namespace mooncake { @@ -88,6 +88,33 @@ class AllocatorManager { return allocator_removed; } + struct Replacement { + std::string name; + std::shared_ptr expected; + std::shared_ptr replacement; + }; + + bool replaceAllocators(const std::vector& replacements) { + std::vector targets; + targets.reserve(replacements.size()); + for (const auto& replacement : replacements) { + auto it = allocators_.find(replacement.name); + if (it == allocators_.end() || !replacement.replacement) { + return false; + } + auto target = std::find(it->second.begin(), it->second.end(), + replacement.expected); + if (target == it->second.end()) { + return false; + } + targets.push_back(target); + } + for (size_t i = 0; i < replacements.size(); ++i) { + *targets[i] = replacements[i].replacement; + } + return true; + } + /** * @brief Get the names of all segments. This returns a vector of the * names so that we can randomly pick a segment without traversing. @@ -119,6 +146,14 @@ class AllocatorManager { friend class SegmentSerializer; // for fork serialize }; +class SsdMetricsProvider { + public: + virtual ~SsdMetricsProvider() = default; + virtual int64_t getSsdTotalCapacity( + const std::string& segment_name) const = 0; + virtual int64_t getSsdUsedBytes(const std::string& segment_name) const = 0; +}; + /** * @brief Abstract interface for allocation strategy, responsible for * allocating a slice (with one or more replicas) using available @@ -166,6 +201,18 @@ class AllocationStrategy { std::set(), const ReplicaType replica_type = ReplicaType::MEMORY) = 0; + virtual tl::expected, ErrorCode> Allocate( + const AllocatorManager& allocator_manager, const size_t slice_length, + const size_t replica_num, + const std::vector& preferred_segments, + const std::set& excluded_segments, + const ReplicaType replica_type, + const SsdMetricsProvider* ssd_provider) { + (void)ssd_provider; + return Allocate(allocator_manager, slice_length, replica_num, + preferred_segments, excluded_segments, replica_type); + } + /** * @brief Allocate one replica from the specified segment. * @@ -222,9 +269,6 @@ class RandomAllocationStrategy : public AllocationStrategy { return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } - // Random number generator. - static thread_local std::mt19937 generator(std::random_device{}()); - std::vector replicas; replicas.reserve(replica_num); @@ -234,8 +278,8 @@ class RandomAllocationStrategy : public AllocationStrategy { return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } - auto buffer = allocateSingle(allocator_manager, names[0], - slice_length, generator); + auto buffer = + allocateSingle(allocator_manager, names[0], slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -255,7 +299,7 @@ class RandomAllocationStrategy : public AllocationStrategy { } auto buffer = allocateSingle(allocator_manager, preferred_segment, - slice_length, generator); + slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -270,8 +314,7 @@ class RandomAllocationStrategy : public AllocationStrategy { // If replica_num is not satisfied, allocate the remaining replicas // randomly. - std::uniform_int_distribution distribution(0, names.size() - 1); - size_t start_idx = distribution(generator); + size_t start_idx = randomIndex(names.size()); const size_t max_retry = std::min(kMaxRetryLimit, names.size()); size_t try_count = 0; @@ -287,8 +330,8 @@ class RandomAllocationStrategy : public AllocationStrategy { continue; } - auto buffer = allocateSingle(allocator_manager, names[index], - slice_length, generator); + auto buffer = + allocateSingle(allocator_manager, names[index], slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -308,9 +351,6 @@ class RandomAllocationStrategy : public AllocationStrategy { tl::expected AllocateFrom( const AllocatorManager& allocator_manager, const size_t slice_length, const std::string& segment_name) { - // Random number generator. - static thread_local std::mt19937 generator(std::random_device{}()); - // Validate input parameters if (slice_length == 0) { return tl::make_unexpected(ErrorCode::INVALID_PARAMS); @@ -321,8 +361,8 @@ class RandomAllocationStrategy : public AllocationStrategy { return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND); } - auto buffer = allocateSingle(allocator_manager, segment_name, - slice_length, generator); + auto buffer = + allocateSingle(allocator_manager, segment_name, slice_length); if (buffer == nullptr) { return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } @@ -332,7 +372,7 @@ class RandomAllocationStrategy : public AllocationStrategy { std::unique_ptr allocateSingle( const AllocatorManager& allocator_manager, const std::string& name, - const size_t slice_length, std::mt19937& generator) { + const size_t slice_length) { const auto allocators = allocator_manager.getAllocators(name); if (allocators == nullptr || allocators->size() == 0) { return nullptr; @@ -346,9 +386,8 @@ class RandomAllocationStrategy : public AllocationStrategy { // Randomly select a start point to distribute // allocations across all segments - std::uniform_int_distribution dist(0, num_segs - 1); - size_t seg_offset = - dist(generator); // select a start segment to place replica + // Select a start segment to place the replica. + size_t seg_offset = randomIndex(num_segs); for (size_t i = 0; i < num_segs; i++) { // only allocate one replica auto& allocator = (*allocators)[(i + seg_offset) % num_segs]; if (auto buffer = allocator->allocate(slice_length)) { @@ -400,8 +439,6 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } - static thread_local std::mt19937 generator(std::random_device{}()); - std::vector replicas; replicas.reserve(replica_num); std::set used_segments; @@ -414,7 +451,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { } auto buffer = allocateSingle(allocator_manager, preferred_segment, - slice_length, generator); + slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -432,8 +469,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { size_t sample_count = std::min(kCandidateMultiplier * remaining, names.size()); - std::uniform_int_distribution start_dist(0, names.size() - 1); - size_t start_idx = start_dist(generator); + size_t start_idx = randomIndex(names.size()); struct Candidate { size_t name_idx; @@ -469,8 +505,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { continue; } - auto buffer = allocateSingle(allocator_manager, name, slice_length, - generator); + auto buffer = allocateSingle(allocator_manager, name, slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -483,8 +518,7 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { } // --- Fallback: Random allocation for any remaining replicas --- - std::uniform_int_distribution distribution(0, names.size() - 1); - size_t fallback_idx = distribution(generator); + size_t fallback_idx = randomIndex(names.size()); const size_t max_retry = std::min(kMaxRetryLimit, names.size()); size_t try_count = 0; @@ -499,8 +533,8 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { continue; } - auto buffer = allocateSingle(allocator_manager, names[index], - slice_length, generator); + auto buffer = + allocateSingle(allocator_manager, names[index], slice_length); if (buffer) { replicas.emplace_back(std::move(buffer), ReplicaStatus::PROCESSING, replica_type); @@ -538,6 +572,149 @@ class FreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { } }; +class SsdFreeRatioFirstAllocationStrategy : public RandomAllocationStrategy { + public: + SsdFreeRatioFirstAllocationStrategy() = default; + + tl::expected, ErrorCode> Allocate( + const AllocatorManager& allocator_manager, const size_t slice_length, + const size_t replica_num, + const std::vector& preferred_segments, + const std::set& excluded_segments, + const ReplicaType replica_type, + const SsdMetricsProvider* ssd_provider) override { + if (slice_length == 0 || replica_num == 0) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + const auto& names = allocator_manager.getNames(); + if (names.empty()) { + return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); + } + + std::vector replicas; + replicas.reserve(replica_num); + std::set used_segments; + + // Handle preferred segments first + for (const auto& preferred_segment : preferred_segments) { + if (excluded_segments.contains(preferred_segment) || + used_segments.contains(preferred_segment)) { + continue; + } + + auto buffer = allocateSingle(allocator_manager, preferred_segment, + slice_length); + if (buffer) { + replicas.emplace_back(std::move(buffer), + ReplicaStatus::PROCESSING, replica_type); + used_segments.insert(preferred_segment); + if (replicas.size() == replica_num) { + return replicas; + } + } + } + + const size_t remaining = replica_num - replicas.size(); + + // Sample candidates and sort by SSD free ratio + size_t sample_count = + std::min(kCandidateMultiplier * remaining, names.size()); + + size_t start_idx = randomIndex(names.size()); + + struct Candidate { + size_t name_idx; + double ssd_free_ratio; + }; + std::vector candidates; + candidates.reserve(sample_count); + + for (size_t i = 0; i < sample_count; ++i) { + size_t idx = (start_idx + i) % names.size(); + const auto& name = names[idx]; + + if (excluded_segments.contains(name) || + used_segments.contains(name)) { + continue; + } + + double ssd_free_ratio = getSegmentSsdFreeRatio(name, ssd_provider); + candidates.push_back({idx, ssd_free_ratio}); + } + + std::sort(candidates.begin(), candidates.end(), + [](const Candidate& a, const Candidate& b) { + return a.ssd_free_ratio > b.ssd_free_ratio; + }); + + for (const auto& candidate : candidates) { + if (replicas.size() >= replica_num) { + break; + } + + const auto& name = names[candidate.name_idx]; + auto buffer = allocateSingle(allocator_manager, name, slice_length); + if (buffer) { + replicas.emplace_back(std::move(buffer), + ReplicaStatus::PROCESSING, replica_type); + used_segments.insert(name); + } + } + + if (replicas.size() >= replica_num) { + return replicas; + } + + // Fallback: Random allocation for remaining replicas + size_t fallback_idx = randomIndex(names.size()); + const size_t max_retry = std::min(kMaxRetryLimit, names.size()); + size_t try_count = 0; + + while (replicas.size() < replica_num && try_count < max_retry) { + auto index = fallback_idx % names.size(); + fallback_idx++; + try_count++; + + const auto& name = names[index]; + + if (excluded_segments.contains(name) || + used_segments.contains(name)) { + continue; + } + + auto buffer = allocateSingle(allocator_manager, name, slice_length); + if (buffer) { + replicas.emplace_back(std::move(buffer), + ReplicaStatus::PROCESSING, replica_type); + used_segments.insert(name); + } + } + + if (replicas.empty()) { + return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); + } + return replicas; + } + + using AllocationStrategy::Allocate; + + private: + static constexpr size_t kMaxRetryLimit = 100; + static constexpr size_t kCandidateMultiplier = 6; + + double getSegmentSsdFreeRatio( + const std::string& name, const SsdMetricsProvider* ssd_provider) const { + if (!ssd_provider) return 1.0; + int64_t total = ssd_provider->getSsdTotalCapacity(name); + if (total <= 0) return 1.0; + int64_t used = ssd_provider->getSsdUsedBytes(name); + used = std::clamp(used, 0, total); + int64_t free_bytes = total - used; + return static_cast(free_bytes) / static_cast(total); + } +}; + class CxlAllocationStrategy : public AllocationStrategy { public: CxlAllocationStrategy() = default; @@ -611,6 +788,10 @@ inline std::shared_ptr CreateAllocationStrategy( return std::make_shared(); case AllocationStrategyType::CXL: return std::make_shared(); + case AllocationStrategyType::SSD_FREE_RATIO_FIRST: + return std::make_shared(); + case AllocationStrategyType::LOCAL_FIRST: + return std::make_shared(); default: return std::make_shared(); } diff --git a/mooncake-store/include/allocator.h b/mooncake-store/include/allocator.h index e28b7d5c54..e0321f5c6f 100644 --- a/mooncake-store/include/allocator.h +++ b/mooncake-store/include/allocator.h @@ -5,9 +5,10 @@ #include #include #include +#include #include "cachelib_memory_allocator/MemoryAllocator.h" -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "types.h" using facebook::cachelib::MemoryAllocator; @@ -49,6 +50,9 @@ class AllocatedBuffer { size_(size), offset_handle_(std::move(offset_handle)) {} + AllocatedBuffer(std::shared_ptr allocator, + const Descriptor& descriptor); + ~AllocatedBuffer(); AllocatedBuffer(const AllocatedBuffer&) = delete; @@ -126,6 +130,37 @@ class BufferAllocatorBase { virtual size_t getLargestFreeRegion() const = 0; }; +/** + * A no-op buffer allocator used only for keeping standby promotion metadata + * alive. It does not actually allocate memory - replicas constructed from + * this allocator are invalid for actual I/O but preserve endpoint info. + */ +class DummyBufferAllocator final : public BufferAllocatorBase { + public: + explicit DummyBufferAllocator(std::string segment_name, + std::string transport_endpoint) + : segment_name_(std::move(segment_name)), + transport_endpoint_(std::move(transport_endpoint)) {} + + std::unique_ptr allocate(size_t size) override { + return nullptr; + } + void deallocate(AllocatedBuffer* handle) override {} + size_t capacity() const override { return kAllocatorUnknownFreeSpace; } + size_t getLargestFreeRegion() const override { + return kAllocatorUnknownFreeSpace; + } + size_t size() const override { return 0; } + std::string getSegmentName() const override { return segment_name_; } + std::string getTransportEndpoint() const override { + return transport_endpoint_; + } + + private: + std::string segment_name_; + std::string transport_endpoint_; +}; + /** * CachelibBufferAllocator manages memory allocation using CacheLib's slab * allocation strategy. @@ -179,6 +214,8 @@ class CachelibBufferAllocator } private: + std::unique_ptr adoptImportedBuffer( + const AllocatedBuffer::Descriptor& descriptor); // metadata const std::string segment_name_; const size_t base_; @@ -194,8 +231,27 @@ class CachelibBufferAllocator size_t header_region_size_; std::unique_ptr memory_allocator_; facebook::cachelib::PoolId pool_id_; + + friend struct RestoredCachelibBufferAllocator; + friend std::optional + RestoreCachelibBufferAllocator( + std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + const std::vector& descriptors, + ReplicaType replica_type); }; +struct RestoredCachelibBufferAllocator { + std::shared_ptr allocator; + std::vector> buffers; +}; + +std::optional RestoreCachelibBufferAllocator( + std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + const std::vector& descriptors, + ReplicaType replica_type = ReplicaType::MEMORY); + /** * OffsetBufferAllocator manages memory allocation using the OffsetAllocator * strategy, which provides efficient memory allocation with bin-based @@ -245,9 +301,33 @@ class OffsetBufferAllocator // offset allocator implementation std::shared_ptr offset_allocator_; + // Keeps address gaps occupied after descriptor-based reconstruction. + std::vector> restored_gap_buffers_; + friend class Serializer; + friend struct RestoredOffsetBufferAllocator; + friend std::optional + RestoreOffsetBufferAllocator( + std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + const std::vector& descriptors, + ReplicaType replica_type); }; +struct RestoredOffsetBufferAllocator { + std::shared_ptr allocator; + std::vector> buffers; +}; + +// Reconstructs an empty OffsetBufferAllocator from final live descriptors. +// The returned buffers follow descriptor input order. No state is exposed on +// validation or allocation failure. +std::optional RestoreOffsetBufferAllocator( + std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + const std::vector& descriptors, + ReplicaType replica_type = ReplicaType::MEMORY); + // The main difference is that it allocates real memory and returns it, while // BufferAllocator allocates an address class SimpleAllocator { diff --git a/mooncake-store/include/cachelib_memory_allocator/AllocationClass.h b/mooncake-store/include/cachelib_memory_allocator/AllocationClass.h index 70ec9b568f..3da2a7965f 100644 --- a/mooncake-store/include/cachelib_memory_allocator/AllocationClass.h +++ b/mooncake-store/include/cachelib_memory_allocator/AllocationClass.h @@ -200,6 +200,11 @@ class AllocationClass { // @return new allocation. This cannot fail. void* addSlabAndAllocate(Slab* slab); + // Restores one slab on a new allocator. Listed chunk indexes stay + // occupied; every other chunk becomes normally allocatable. + bool importSlab(Slab* slab, + const std::vector& occupiedChunkIndexes); + // Releasing a slab is a two step process. // 1. Mark a slab for release, by calling `startSlabRelease`. // 2. Free all the activeAllocations diff --git a/mooncake-store/include/cachelib_memory_allocator/MemoryAllocator.h b/mooncake-store/include/cachelib_memory_allocator/MemoryAllocator.h index a2b5f28826..de990d5271 100644 --- a/mooncake-store/include/cachelib_memory_allocator/MemoryAllocator.h +++ b/mooncake-store/include/cachelib_memory_allocator/MemoryAllocator.h @@ -101,6 +101,15 @@ class MemoryAllocator { // allocation handed out by this allocator. void free(void* memory); + struct ImportedAllocation { + void* address; + uint32_t size; + }; + + // Imports final live allocations into a new, empty allocator. + bool importAllocations(PoolId id, + const std::vector& allocations); + // Memory pool interface. The memory pools must be established before the // first allocation happens. Currently we dont support adding / removing // pools dynamically. diff --git a/mooncake-store/include/cachelib_memory_allocator/MemoryPool.h b/mooncake-store/include/cachelib_memory_allocator/MemoryPool.h index d34197e1f9..4dd986fe8e 100644 --- a/mooncake-store/include/cachelib_memory_allocator/MemoryPool.h +++ b/mooncake-store/include/cachelib_memory_allocator/MemoryPool.h @@ -137,6 +137,11 @@ class MemoryPool { // @throw std::run_time_error if the slab class information is corrupted. void free(void* memory); + bool importSlab(Slab* slab, + ClassId classId, + const std::vector& occupiedChunkIndexes); + void importFreeSlab(Slab* slab); + // resize the memory pool. This only adjusts the Pool size. It does not // release the slabs back to the SlabAllocator if the new size is less than // the current size. The caller is responsible for doing that through diff --git a/mooncake-store/include/client_buffer.hpp b/mooncake-store/include/client_buffer.h similarity index 99% rename from mooncake-store/include/client_buffer.hpp rename to mooncake-store/include/client_buffer.h index b9f84089fe..f3fce555fe 100644 --- a/mooncake-store/include/client_buffer.hpp +++ b/mooncake-store/include/client_buffer.h @@ -5,7 +5,7 @@ #include #include -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "types.h" #include "replica.h" diff --git a/mooncake-store/include/client_metric.h b/mooncake-store/include/client_metric.h index 0674bea81e..87247f1770 100644 --- a/mooncake-store/include/client_metric.h +++ b/mooncake-store/include/client_metric.h @@ -14,8 +14,9 @@ #include #include #include -#include "utils.h" +#include "environ.h" #include "hybrid_metric.h" +#include "utils.h" namespace mooncake { @@ -23,22 +24,17 @@ namespace mooncake { // Tuned for RDMA: fine-grained in <1ms, with ms-scale tail up to 1s const std::vector kLatencyBucket = { // sub-ms to 1ms region - 125, 150, 200, 250, 300, 400, 500, 750, 1000, + 50, 75, 125, 150, 200, 250, 300, 400, 500, 750, 1000, // ms-level tail for batch/occasional spikes 1500, 2000, 3000, 5000, 7000, 15000, 20000, // safeguards for long tails - 50000, 100000, 200000, 500000, 1000000}; - -static inline std::string get_env_or_default( - const char* env_var, const std::string& default_val = "") { - const char* val = getenv(env_var); - return val ? val : default_val; -} + 50000, 100000, 200000, 500000, 1000000, 2000000, 5000000, 10000000, + 20000000}; // In production mode, more labels are needed for monitoring and troubleshooting // Static labels include but are not limited to machine address, cluster name, // etc. These labels remain constant during the lifetime of the application -const std::string kClusterID = get_env_or_default("MC_STORE_CLUSTER_ID"); +const std::string kClusterID = Environ::GetString("MC_STORE_CLUSTER_ID", ""); // Merge static labels with dynamic labels const inline std::map merge_labels( diff --git a/mooncake-store/include/client_service.h b/mooncake-store/include/client_service.h index 0094f4f14d..e89f911e6f 100644 --- a/mooncake-store/include/client_service.h +++ b/mooncake-store/include/client_service.h @@ -42,11 +42,15 @@ class QueryResult { const std::vector replicas; /** @brief Time point when the lease for this key expires */ const std::chrono::steady_clock::time_point lease_timeout; + /** @brief Optional full-object checksum */ + const std::optional object_checksum; QueryResult(std::vector&& replicas_param, - std::chrono::steady_clock::time_point lease_timeout_param) + std::chrono::steady_clock::time_point lease_timeout_param, + std::optional object_checksum_param = std::nullopt) : replicas(std::move(replicas_param)), - lease_timeout(lease_timeout_param) {} + lease_timeout(lease_timeout_param), + object_checksum(object_checksum_param) {} bool IsLeaseExpired() const { return std::chrono::steady_clock::now() >= lease_timeout; @@ -150,6 +154,10 @@ class Client { const std::vector& object_keys, const std::string& tenant_id); + tl::expected VerifyObjectChecksum( + const std::string& object_key, const std::vector& slices, + size_t object_size, std::optional expected_checksum); + /** * @brief Batch clear KV cache for specified object keys on a specific * segment for a given client. @@ -179,6 +187,9 @@ class Client { const QueryResult& query_result, std::vector& slices, uint64_t src_offset); + std::optional SubmitScatter( + const std::vector& transfers); + /** * @brief Transfers data using pre-queried object information * @param object_keys Keys of the objects @@ -425,6 +436,8 @@ class Client { bool enable_offloading, std::vector& offloading_objects); + tl::expected PollRemoveAll(); + tl::expected ReportSsdCapacity( int64_t ssd_total_capacity_bytes); @@ -659,6 +672,17 @@ class Client { const std::map& labels = {}, const std::string& tenant_id = "default"); + /** + * @brief Prepare and use the storage backend for persisting data. + * Exposed to subclasses for testing only. + * @return ErrorCode::OK on success. On failure no storage backend is + * retained, so persistence stays disabled. + */ + ErrorCode PrepareStorageBackend(const std::string& storage_root_dir, + const std::string& fsdir, + bool enable_eviction = true, + uint64_t quota_bytes = 0); + private: /** * @brief Internal helper functions for initialization and data transfer @@ -682,14 +706,9 @@ class Client { ErrorCode TransferReadRange(const Replica::Descriptor& replica_descriptor, std::vector& slices, uint64_t src_offset); - - /** - * @brief Prepare and use the storage backend for persisting data - */ - void PrepareStorageBackend(const std::string& storage_root_dir, - const std::string& fsdir, - bool enable_eviction = true, - uint64_t quota_bytes = 0); + tl::expected ComputeObjectChecksumForSlices( + const std::string& object_key, const std::vector& slices, + size_t object_size); void PutToLocalFile(const std::string& object_key, const std::vector& slices, @@ -761,6 +780,7 @@ class Client { const std::vector>& batched_slices); void StartBatchPut(std::vector& ops, const ReplicateConfig& config); + void ComputeBatchObjectChecksums(std::vector& ops); void SubmitTransfers(std::vector& ops); void WaitForTransfers(std::vector& ops); void FinalizeBatchPut(std::vector& ops); @@ -776,6 +796,7 @@ class Client { const std::vector& object_keys, const std::vector& query_results, std::unordered_map>& slices); + ReplicateConfig AttachHostId(const ReplicateConfig& config) const; // Client identification const UUID client_id_; @@ -817,8 +838,10 @@ class Client { // Configuration const std::string local_hostname_; + const std::string host_id_; const std::string metadata_connstring_; const std::string protocol_; + const bool object_checksum_enabled_; // Client persistent thread pool for async operations // Pinned host memory pool for GPU D2H staging (must outlive diff --git a/mooncake-store/include/crc32c.h b/mooncake-store/include/crc32c.h new file mode 100644 index 0000000000..ac76924538 --- /dev/null +++ b/mooncake-store/include/crc32c.h @@ -0,0 +1,53 @@ +#pragma once + +#include +#include +#include + +namespace mooncake { + +// CRC-32C (Castagnoli), reflected polynomial 0x82F63B78. +// Table-driven software implementation; incremental usage: +// Crc32c crc; +// crc.Extend(p1, n1); +// crc.Extend(p2, n2); +// uint32_t value = crc.Final(); +class Crc32c { + public: + Crc32c() : value_(kInit) {} + + void Extend(const void* data, size_t len) { + const auto* p = static_cast(data); + uint32_t v = value_; + for (size_t i = 0; i < len; ++i) { + v = kTable[(v ^ p[i]) & 0xFF] ^ (v >> 8); + } + value_ = v; + } + + uint32_t Final() const { return value_ ^ kInit; } + + private: + static constexpr uint32_t kInit = 0xFFFFFFFFu; + static constexpr std::array kTable = [] { + std::array t{}; + for (uint32_t i = 0; i < 256; ++i) { + uint32_t c = i; + for (int k = 0; k < 8; ++k) { + c = (c & 1) ? (0x82F63B78u ^ (c >> 1)) : (c >> 1); + } + t[i] = c; + } + return t; + }(); + + uint32_t value_; +}; + +inline uint32_t Crc32cValue(const void* data, size_t len) { + Crc32c crc; + crc.Extend(data, len); + return crc.Final(); +} + +} // namespace mooncake diff --git a/mooncake-store/include/device/accelerator_device.h b/mooncake-store/include/device/accelerator_device.h new file mode 100644 index 0000000000..cbe495ecd6 --- /dev/null +++ b/mooncake-store/include/device/accelerator_device.h @@ -0,0 +1,68 @@ +#pragma once + +#include +#include +#include + +#include "pinned_host_buffer.h" + +namespace mooncake { +namespace device { + +enum class AcceleratorVendor { + kNvidia, + kMusa, + kMaca, + kHygon, + kCorex, + kHip, + kAscend, + kSunrise, +}; + +enum class MemoryKind { + kHost, + kDevice, + kUnknown, +}; + +enum class CopyDirection { + kHostToHost, + kHostToDevice, + kDeviceToHost, + kDeviceToDevice, + kAuto, +}; + +struct PointerInfo { + MemoryKind kind = MemoryKind::kUnknown; + int32_t device_id = -1; +}; + +class AcceleratorDevice { + public: + virtual ~AcceleratorDevice() = default; + + virtual AcceleratorVendor Vendor() const = 0; + virtual bool Available(bool ensure = false) const = 0; + virtual PointerInfo QueryPointer(const void* ptr) const = 0; + virtual int32_t CurrentDeviceId() const = 0; + virtual void SetContext(int32_t device_id) const = 0; + virtual bool Copy(void* dst, const void* src, size_t size, + CopyDirection direction) const = 0; + virtual PinnedHostBuffer AllocatePinnedHost(size_t size) const = 0; +}; + +class ProbeCachedAcceleratorDevice : public AcceleratorDevice { + public: + bool Available(bool ensure = false) const override; + + protected: + virtual bool ProbeAvailable() const = 0; + + private: + mutable std::atomic available_state_{0}; +}; + +} // namespace device +} // namespace mooncake diff --git a/mooncake-store/include/device/accelerator_registry.h b/mooncake-store/include/device/accelerator_registry.h new file mode 100644 index 0000000000..e99ed97652 --- /dev/null +++ b/mooncake-store/include/device/accelerator_registry.h @@ -0,0 +1,32 @@ +#pragma once + +#include +#include + +#include "device/accelerator_device.h" +#include "device/runtime_accelerator.h" + +namespace mooncake { +namespace device { + +class AcceleratorRegistry { + public: + virtual ~AcceleratorRegistry() = default; + + virtual std::span RegisteredDevices() + const = 0; + virtual RuntimeAccelerator RuntimeAccelerators( + bool ensure = false) const = 0; + virtual const AcceleratorDevice* GetDevice( + AcceleratorVendor vendor) const = 0; +}; + +const AcceleratorRegistry& GetAcceleratorRegistry(); + +class AcceleratorDeviceRegistrar { + public: + explicit AcceleratorDeviceRegistrar(const AcceleratorDevice& device); +}; + +} // namespace device +} // namespace mooncake diff --git a/mooncake-store/include/device/cuda_ipc_buffer.h b/mooncake-store/include/device/cuda_ipc_buffer.h new file mode 100644 index 0000000000..3a6821482f --- /dev/null +++ b/mooncake-store/include/device/cuda_ipc_buffer.h @@ -0,0 +1,42 @@ +#pragma once + +#include +#include + +#include "device/cuda_ipc_buffer_handle.h" + +namespace mooncake { +namespace device { + +tl::expected ExportCudaIpcBuffer( + const void *ptr, size_t size); + +class CudaIpcBufferMapping { + public: + CudaIpcBufferMapping() = default; + ~CudaIpcBufferMapping(); + + CudaIpcBufferMapping(const CudaIpcBufferMapping &) = delete; + CudaIpcBufferMapping &operator=(const CudaIpcBufferMapping &) = delete; + + CudaIpcBufferMapping(CudaIpcBufferMapping &&other) noexcept; + CudaIpcBufferMapping &operator=(CudaIpcBufferMapping &&other) noexcept; + + static tl::expected Open( + const CudaIpcBufferHandle &handle); + + void *ptr() const { return ptr_; } + + private: + CudaIpcBufferMapping(void *base, void *ptr, int32_t device_id) + : base_(base), ptr_(ptr), device_id_(device_id) {} + + void Close(); + + void *base_ = nullptr; + void *ptr_ = nullptr; + int32_t device_id_ = -1; +}; + +} // namespace device +} // namespace mooncake diff --git a/mooncake-store/include/device/cuda_ipc_buffer_handle.h b/mooncake-store/include/device/cuda_ipc_buffer_handle.h new file mode 100644 index 0000000000..21c6951172 --- /dev/null +++ b/mooncake-store/include/device/cuda_ipc_buffer_handle.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include +#include + +#include "types.h" + +namespace mooncake { + +constexpr size_t kCudaIpcHandleSize = 64; + +struct CudaIpcBufferHandle { + std::array handle{}; + uint64_t offset = 0; + uint64_t size = 0; + int32_t device_id = -1; +}; + +struct CudaIpcShmBufferRef { + uint64_t ptr = 0; + uint64_t size = 0; +}; + +struct CudaIpcWriteRequest { + std::string key; + CudaIpcShmBufferRef metadata; + CudaIpcBufferHandle payload; +}; + +struct CudaIpcReadRequest { + std::string key; + CudaIpcBufferHandle destination; + uint64_t source_offset = 0; + uint64_t size = 0; +}; + +} // namespace mooncake + +YLT_REFL(mooncake::CudaIpcBufferHandle, handle, offset, size, device_id); +YLT_REFL(mooncake::CudaIpcShmBufferRef, ptr, size); +YLT_REFL(mooncake::CudaIpcWriteRequest, key, metadata, payload); +YLT_REFL(mooncake::CudaIpcReadRequest, key, destination, source_offset, size); diff --git a/mooncake-store/include/device/runtime_accelerator.h b/mooncake-store/include/device/runtime_accelerator.h new file mode 100644 index 0000000000..90ba28466a --- /dev/null +++ b/mooncake-store/include/device/runtime_accelerator.h @@ -0,0 +1,31 @@ +#pragma once + +#include +#include +#include + +#include "device/accelerator_device.h" + +namespace mooncake { +namespace device { + +class RuntimeAccelerator { + public: + RuntimeAccelerator() = default; + explicit RuntimeAccelerator(std::vector devices); + + std::span Devices() const; + + const AcceleratorDevice* FindDeviceForPointer( + const void* ptr, PointerInfo* out_info = nullptr) const; + + bool CopyToHost(void* dst, const void* src, size_t size) const; + + bool CopyFromHost(void* dst, const void* src, size_t size) const; + + private: + std::vector devices_; +}; + +} // namespace device +} // namespace mooncake diff --git a/mooncake-store/include/dummy_client.h b/mooncake-store/include/dummy_client.h index d538708828..7c642e91e8 100644 --- a/mooncake-store/include/dummy_client.h +++ b/mooncake-store/include/dummy_client.h @@ -8,7 +8,9 @@ #include #include "client_metric.h" +#include "device/cuda_ipc_buffer_handle.h" #include "pyclient.h" +#include "store_rpc_client_io_context.h" #include "shm_helper.h" #include @@ -30,7 +32,9 @@ class DummyClient : public PyClient { const std::string &ipc_socket_path, bool enable_ssd_offload = false, const std::string &ssd_offload_path = "", - const std::string &tenant_id = "default") { + const std::string &tenant_id = "default", + bool enable_client_http_server = false, + int client_http_port = DEFAULT_CLIENT_HTTP_PORT) { // Dummy client does not support real setup return -1; }; @@ -71,6 +75,9 @@ class DummyClient : public PyClient { const std::vector &buffers, const std::vector &sizes); + std::vector batch_get_into_cuda_ipc( + const std::vector &requests); + std::vector batch_get_into_multi_buffers( const std::vector &keys, const std::vector> &all_buffers, @@ -96,6 +103,10 @@ class DummyClient : public PyClient { const std::vector> &all_sizes, const ReplicateConfig &config = ReplicateConfig{}); + std::vector batch_put_from_cuda_ipc( + const std::vector &requests, + const ReplicateConfig &config = ReplicateConfig{}); + std::shared_ptr get_buffer(const std::string &key); std::vector> batch_get_buffer( @@ -176,6 +187,8 @@ class DummyClient : public PyClient { tl::expected query_task(const UUID &task_id); + std::optional allocate_client_buffer(size_t size) override; + private: ErrorCode connect(const std::string &server_address); @@ -235,39 +248,11 @@ class DummyClient : public PyClient { return to_py_ret(result); } - /** - * @brief Accessor for the coro_rpc_client pool. Since coro_rpc_client - * pool cannot reconnect to a different address, a new coro_rpc_client - * pool is created if the address is different from the current one. - */ - class RpcClientAccessor { - public: - void SetClientPool( - std::shared_ptr> - client_pool) { - std::lock_guard lock(client_mutex_); - client_pool_ = client_pool; - } - - std::shared_ptr> - GetClientPool() { - std::shared_lock lock(client_mutex_); - return client_pool_; - } - - private: - mutable std::shared_mutex client_mutex_; - std::shared_ptr> - client_pool_; - }; - RpcClientAccessor client_accessor_; + RpcClientPool client_accessor_; // The client identification. const UUID client_id_; - std::shared_ptr> - client_pools_; - // Mutex to insure the Connect function is atomic. mutable Mutex connect_mutex_; // The address which is passed to the coro_rpc_client @@ -276,6 +261,7 @@ class DummyClient : public PyClient { // For shared memory management ShmHelper *shm_helper_ = nullptr; std::string ipc_socket_path_; + void *local_buffer_base_ = nullptr; // Hot cache shm mapping (obtained from real client via IPC) void *hot_cache_base_ = nullptr; diff --git a/mooncake-store/include/etcd_helper.h b/mooncake-store/include/etcd_helper.h index 677bef9a91..b0d851c5b6 100644 --- a/mooncake-store/include/etcd_helper.h +++ b/mooncake-store/include/etcd_helper.h @@ -71,6 +71,31 @@ class EtcdHelper { static ErrorCode BatchCreate(const std::vector& keys, const std::vector& values); + enum class TxnCompareKind { + kValueEquals = 0, + kKeyNotExists = 1, + }; + + struct TxnCompare { + std::string key; + TxnCompareKind kind{TxnCompareKind::kValueEquals}; + std::string expected_value; + }; + + struct TxnPut { + std::string key; + std::string value; + }; + + /* + * @brief Execute an etcd transaction with compare predicates and put + * operations. + * @return: OK on success; ETCD_TRANSACTION_FAIL when compare predicates + * fail. + */ + static ErrorCode TxnCompareAndPut(const std::vector& compares, + const std::vector& puts); + /* * @brief Grant a lease from the etcd. * @param lease_ttl: The ttl of the lease, in seconds. diff --git a/mooncake-store/include/file_interface.h b/mooncake-store/include/file_interface.h index c2ab05ce6e..52c267451d 100644 --- a/mooncake-store/include/file_interface.h +++ b/mooncake-store/include/file_interface.h @@ -65,6 +65,11 @@ class StorageFile { * @brief Destructor * @note Automatically closes the file and releases resources */ + virtual tl::expected datasync() { return {}; } + + // Prevent destructor from unlinking the arena file on write failure. + void SetDeleteOnWriteFail(bool v) { delete_on_write_fail_ = v; } + virtual ~StorageFile() = default; /** @@ -150,6 +155,7 @@ class StorageFile { ErrorCode get_error_code() { return error_code_; } protected: + bool delete_on_write_fail_ = true; std::string filename_; int fd_; ErrorCode error_code_{ErrorCode::OK}; @@ -159,6 +165,8 @@ class StorageFile { class PosixFile : public StorageFile { public: PosixFile(const std::string &filename, int fd); + + tl::expected datasync() override; ~PosixFile() override; tl::expected write(const std::string &buffer, @@ -219,7 +227,7 @@ class UringFile : public StorageFile { // Flush data to stable storage via IORING_FSYNC_DATASYNC. // Must be called after write_aligned and before writing dependent metadata. - tl::expected datasync(); + tl::expected datasync() override; // Buffer registration — delegates to the shared ring (process-wide). // Static variant: no file instance needed. Must be called once from a diff --git a/mooncake-store/include/file_storage.h b/mooncake-store/include/file_storage.h index 6aa2ac3d56..501a80c66f 100644 --- a/mooncake-store/include/file_storage.h +++ b/mooncake-store/include/file_storage.h @@ -1,7 +1,7 @@ #pragma once #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "storage_backend.h" #include "pinned_buffer_pool.h" @@ -18,6 +18,8 @@ class FileStorage { tl::expected Init(); + void RemoveAll(); + /** * @brief Result of BatchGet operation containing batch_id and buffer * pointers. @@ -77,14 +79,36 @@ class FileStorage { tl::expected OffloadObjects( const std::vector& offloading_objects); + /** + * @brief Classifies a BatchOffload error as affecting only the current + * bucket rather than the whole offload cycle. + * + * Such an error means the bucket's keys simply cannot be persisted this + * round; OffloadObjects reports them back to the master as failed (so their + * offloading tasks and source-replica refcounts are released) and continues + * with the remaining buckets, instead of aborting the entire cycle. + * + * - INVALID_READ: source data for these keys could not be read/staged. + * - OBJECT_ALREADY_EXISTS: the key(s) were already offloaded, or are + * being offloaded concurrently. The bucket backend rejects the whole + * bucket atomically by design (see BucketStorageBackend::BatchOffload + * and its PrepareEviction duplicate guard). Treating this as fatal + * aborted the cycle, leaked offloading tasks, and left master/SSD + * metadata inconsistent, which surfaced as spurious INVALID_KEY on the + * read path (issue #2827). + */ + static bool IsPerBucketSoftOffloadError(ErrorCode error); + /** * @brief Performs a heartbeat operation for the FileStorage component. * 1. Sends object status (e.g., access frequency, size) to the master via * client. * 2. Receives feedback on which objects should be offloaded. * 3. Triggers asynchronous offloading of pending objects. - * 4. Pulls and processes any pending L2->L1 promotion tasks queued by the - * master (mirror of step 1+2 in the reverse direction). + * 4. If offload work was returned, pulls and processes any pending L2->L1 + * promotion tasks queued by the master (mirror of step 1+2 in the + * reverse direction). + * 5. Runs proactive local-disk watermark eviction. * @return tl::expected indicating operation status. */ tl::expected Heartbeat(); @@ -103,6 +127,11 @@ class FileStorage { tl::expected IsEnableOffloading(); + tl::expected RunDiskWatermarkEviction(); + + tl::expected NotifyEvictedDiskReplicas( + const std::vector& evicted_keys); + tl::expected BatchLoad( std::unordered_map& batch_object); diff --git a/mooncake-store/include/gpu_staging_utils.h b/mooncake-store/include/gpu_staging_utils.h deleted file mode 100644 index 07982aa015..0000000000 --- a/mooncake-store/include/gpu_staging_utils.h +++ /dev/null @@ -1,166 +0,0 @@ -#pragma once - -#include "cuda_alike.h" - -#if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) -#include -#endif - -#include -#include - -namespace mooncake { -namespace gpu_staging { - -// Detect whether ptr resides in accelerator device memory. -// If so, writes the device ID to *out_device_id for subsequent SetDevice. -inline bool IsDevicePointer(const void* ptr, int* out_device_id) { -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ - defined(USE_HYGON) || defined(USE_COREX) - cudaPointerAttributes attr{}; - if (cudaPointerGetAttributes(&attr, ptr) == cudaSuccess && - attr.type == cudaMemoryTypeDevice) { - if (out_device_id) *out_device_id = attr.device; - return true; - } -#elif defined(USE_HIP) - hipPointerAttribute_t attr{}; - if (hipPointerGetAttributes(&attr, ptr) == hipSuccess && - attr.type == hipMemoryTypeDevice) { - if (out_device_id) *out_device_id = attr.device; - return true; - } -#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) - aclrtPtrAttributes attr{}; - if (aclrtPointerGetAttributes(const_cast(ptr), &attr) == - ACL_SUCCESS && - attr.location.type == ACL_MEM_LOCATION_TYPE_DEVICE) { - if (out_device_id) *out_device_id = static_cast(attr.location.id); - return true; - } -#endif - (void)ptr; - (void)out_device_id; - return false; -} - -// Copy device memory to host. Caller must have called SetDevice first. -inline bool CopyDeviceToHost(void* dst, const void* src, size_t size) { -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ - defined(USE_HYGON) || defined(USE_COREX) - return cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost) == cudaSuccess; -#elif defined(USE_HIP) - return hipMemcpy(dst, src, size, hipMemcpyDeviceToHost) == hipSuccess; -#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) - return aclrtMemcpy(dst, size, src, size, ACL_MEMCPY_DEVICE_TO_HOST) == - ACL_SUCCESS; -#else - (void)dst; - (void)src; - (void)size; - return false; -#endif -} - -// Auto-direction copy: runtime determines the transfer direction from pointer -// attributes (cudaMemcpyDefault). Works for H2H, H2D, D2H, and D2D. -// Caller must have called SetDevice first when device memory is involved. -inline bool CopyAuto(void* dst, const void* src, size_t size) { -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ - defined(USE_HYGON) || defined(USE_COREX) - return cudaMemcpy(dst, src, size, cudaMemcpyDefault) == cudaSuccess; -#elif defined(USE_HIP) - return hipMemcpy(dst, src, size, hipMemcpyDefault) == hipSuccess; -#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) - aclrtPtrAttributes src_attr{}, dst_attr{}; - bool src_dev = aclrtPointerGetAttributes(const_cast(src), - &src_attr) == ACL_SUCCESS && - src_attr.location.type == ACL_MEM_LOCATION_TYPE_DEVICE; - bool dst_dev = aclrtPointerGetAttributes(dst, &dst_attr) == ACL_SUCCESS && - dst_attr.location.type == ACL_MEM_LOCATION_TYPE_DEVICE; - aclrtMemcpyKind kind = ACL_MEMCPY_HOST_TO_HOST; - if (src_dev && dst_dev) - kind = ACL_MEMCPY_DEVICE_TO_DEVICE; - else if (src_dev) - kind = ACL_MEMCPY_DEVICE_TO_HOST; - else if (dst_dev) - kind = ACL_MEMCPY_HOST_TO_DEVICE; - return aclrtMemcpy(dst, size, src, size, kind) == ACL_SUCCESS; -#else - (void)dst; - (void)src; - (void)size; - return false; -#endif -} - -// Bind the calling thread to the given device context. -inline void SetDevice(int device_id) { - if (device_id < 0) return; -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ - defined(USE_HYGON) || defined(USE_COREX) - cudaSetDevice(device_id); -#elif defined(USE_HIP) - hipSetDevice(device_id); -#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) - aclrtSetDevice(device_id); -#endif -} - -// Copy host memory to device. Caller must have called SetDevice first. -inline bool CopyHostToDevice(void* dst, const void* src, size_t size) { -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ - defined(USE_HYGON) || defined(USE_COREX) - return cudaMemcpy(dst, src, size, cudaMemcpyHostToDevice) == cudaSuccess; -#elif defined(USE_HIP) - return hipMemcpy(dst, src, size, hipMemcpyHostToDevice) == hipSuccess; -#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) - return aclrtMemcpy(dst, size, src, size, ACL_MEMCPY_HOST_TO_DEVICE) == - ACL_SUCCESS; -#else - (void)dst; - (void)src; - (void)size; - return false; -#endif -} - -// Detect whether ptr resides in host (CPU) memory. -// Used together with IsDevicePointer for safe pointer-type dispatching: -// if IsDevicePointer -> CopyHostToDevice / CopyDeviceToHost -// else if IsHostPointer -> memcpy -// else -> reject (unknown type, e.g. non-standard allocator) -// -// Pageable host memory (not tracked by CUDA runtime) is treated as host. -inline bool IsHostPointer(const void* ptr) { -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ - defined(USE_HYGON) || defined(USE_COREX) - cudaPointerAttributes attr{}; - if (cudaPointerGetAttributes(&attr, ptr) != cudaSuccess) { - // Query failed: pageable host memory not tracked by the runtime. - cudaGetLastError(); // clear sticky error - return true; - } - return attr.type != cudaMemoryTypeDevice; -#elif defined(USE_HIP) - hipPointerAttribute_t attr{}; - if (hipPointerGetAttributes(&attr, ptr) != hipSuccess) { - hipGetLastError(); // clear sticky error - return true; - } - return attr.type != hipMemoryTypeDevice; -#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) - aclrtPtrAttributes attr{}; - if (aclrtPointerGetAttributes(const_cast(ptr), &attr) != - ACL_SUCCESS) { - // Query failed: likely pageable host memory not tracked by the runtime. - return true; - } - return attr.location.type != ACL_MEM_LOCATION_TYPE_DEVICE; -#else - (void)ptr; - return true; // CPU-only build: all pointers are host -#endif -} -} // namespace gpu_staging -} // namespace mooncake diff --git a/mooncake-store/include/ha/ha_types.h b/mooncake-store/include/ha/ha_types.h index 1d85bc75e4..82a9181c87 100644 --- a/mooncake-store/include/ha/ha_types.h +++ b/mooncake-store/include/ha/ha_types.h @@ -69,7 +69,11 @@ inline ErrorCode ValidateHABackendAvailability(HABackendType type) { return ErrorCode::UNAVAILABLE_IN_CURRENT_MODE; #endif case HABackendType::K8S: +#ifdef STORE_USE_K8S_LEASE + return ErrorCode::OK; +#else return ErrorCode::UNAVAILABLE_IN_CURRENT_MODE; +#endif } return ErrorCode::INVALID_PARAMS; diff --git a/mooncake-store/include/ha/kv/etcd_ha_kv_backend.h b/mooncake-store/include/ha/kv/etcd_ha_kv_backend.h new file mode 100644 index 0000000000..ae58b5f93e --- /dev/null +++ b/mooncake-store/include/ha/kv/etcd_ha_kv_backend.h @@ -0,0 +1,17 @@ +#pragma once + +#include "ha/kv/ha_kv_backend.h" + +namespace mooncake { + +class EtcdHaKvBackend : public HaKvBackend { + public: + ErrorCode Get(std::string_view key, std::string& value) override; + ErrorCode Put(std::string_view key, std::string_view value) override; + ErrorCode Range(std::string_view begin_key, std::string_view end_key, + size_t limit, std::vector& kvs) override; + bool SupportsTxn() const override; + ErrorCode Txn(const KvTxn& txn) override; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/ha/kv/ha_kv_backend.h b/mooncake-store/include/ha/kv/ha_kv_backend.h new file mode 100644 index 0000000000..e48583774a --- /dev/null +++ b/mooncake-store/include/ha/kv/ha_kv_backend.h @@ -0,0 +1,45 @@ +#pragma once + +#include +#include +#include + +#include "types.h" + +namespace mooncake { + +struct KvPair { + std::string key; + std::string value; +}; + +enum class KvCompareKind { + kValueEquals, + kKeyNotExists, +}; + +struct KvCompare { + std::string key; + KvCompareKind kind{KvCompareKind::kValueEquals}; + std::string expected_value; +}; + +struct KvTxn { + std::vector compares; + std::vector puts; +}; + +class HaKvBackend { + public: + virtual ~HaKvBackend() = default; + + virtual ErrorCode Get(std::string_view key, std::string& value) = 0; + virtual ErrorCode Put(std::string_view key, std::string_view value) = 0; + virtual ErrorCode Range(std::string_view begin_key, + std::string_view end_key, size_t limit, + std::vector& kvs) = 0; + virtual bool SupportsTxn() const = 0; + virtual ErrorCode Txn(const KvTxn& txn) = 0; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/ha/leadership/leader_label_reconciler.h b/mooncake-store/include/ha/leadership/leader_label_reconciler.h new file mode 100644 index 0000000000..dc0c37cc9e --- /dev/null +++ b/mooncake-store/include/ha/leadership/leader_label_reconciler.h @@ -0,0 +1,99 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "types.h" + +namespace mooncake { +namespace ha { + +// Drives a single boolean "leader label" toward a desired state on a background +// thread. Callers only flip the desired flag (non-blocking); the worker runs +// the synchronous apply and retries until the actual state converges, so a +// transient failure cannot leave a stale leader label on a pod that is no +// longer serving. +class LeaderLabelReconciler { + public: + // Applies the label mutation: true sets the leader label, false clears it. + // Must return ErrorCode::OK only when the cluster state reflects `desired`. + using ApplyFn = std::function; + + LeaderLabelReconciler(bool enabled, ApplyFn apply, + std::chrono::milliseconds retry_interval) + : enabled_(enabled), + apply_(std::move(apply)), + retry_interval_(retry_interval) { + if (enabled_) { + worker_ = std::thread([this] { Run(); }); + } + } + + ~LeaderLabelReconciler() { + if (!enabled_) return; + { + std::lock_guard lock(mutex_); + stop_ = true; + } + cv_.notify_all(); + if (worker_.joinable()) worker_.join(); + } + + LeaderLabelReconciler(const LeaderLabelReconciler&) = delete; + LeaderLabelReconciler& operator=(const LeaderLabelReconciler&) = delete; + + void SetLeader(bool desired) { + if (!enabled_) return; + { + std::lock_guard lock(mutex_); + desired_leader_ = desired; + } + cv_.notify_all(); + } + + private: + void Run() { + std::optional applied; + std::unique_lock lock(mutex_); + while (true) { + cv_.wait(lock, [&] { return stop_ || applied != desired_leader_; }); + if (stop_) return; + bool target = desired_leader_; + lock.unlock(); + + ErrorCode err = apply_(target); + + lock.lock(); + if (err == ErrorCode::OK) { + applied = target; + continue; + } + applied.reset(); + LOG(WARNING) << "Failed to " << (target ? "set" : "clear") + << " leader label: " << toString(err); + cv_.wait_for(lock, retry_interval_, + [&] { return stop_ || desired_leader_ != target; }); + if (stop_) return; + } + } + + const bool enabled_; + const ApplyFn apply_; + const std::chrono::milliseconds retry_interval_; + + std::mutex mutex_; + std::condition_variable cv_; + bool desired_leader_ = false; + bool stop_ = false; + std::thread worker_; +}; + +} // namespace ha +} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/etcd_oplog_change_notifier.h b/mooncake-store/include/ha/oplog/etcd_oplog_change_notifier.h deleted file mode 100644 index 021fee22b5..0000000000 --- a/mooncake-store/include/ha/oplog/etcd_oplog_change_notifier.h +++ /dev/null @@ -1,103 +0,0 @@ -// mooncake-store/include/etcd_oplog_change_notifier.h -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "ha/oplog/etcd_oplog_store.h" -#include "ha/oplog/oplog_change_notifier.h" -#include "ha/oplog/oplog_manager.h" -#include "types.h" - -namespace mooncake { - -// Forward declaration -class EtcdOpLogChangeNotifier; - -// Shared control block for safe C-style watch callbacks. -// Because the etcd Watch goroutine can deliver callbacks after Stop() -// returns (or even after destruction), we cannot pass a raw `this` -// pointer as the callback context. -struct ChangeNotifierCallbackContext { - std::mutex mutex; - EtcdOpLogChangeNotifier* notifier{nullptr}; - - ChangeNotifierCallbackContext() = default; - ChangeNotifierCallbackContext(const ChangeNotifierCallbackContext&) = - delete; - ChangeNotifierCallbackContext& operator=( - const ChangeNotifierCallbackContext&) = delete; -}; - -// OpLogChangeNotifier implementation backed by etcd Watch. -// Migrated from OpLogReplicator's watch logic. -class EtcdOpLogChangeNotifier : public OpLogChangeNotifier { - public: - explicit EtcdOpLogChangeNotifier(const std::string& cluster_id, - EtcdOpLogStore* oplog_store); - ~EtcdOpLogChangeNotifier(); - - ErrorCode Start(uint64_t start_sequence_id, EntryCallback on_entry, - ErrorCallback on_error) override; - void Stop() override; - bool IsHealthy() const override; - - private: - // Read historical entries and return the etcd revision for watch resume. - bool ReadOpLogSince(uint64_t start_seq_id, std::vector& entries, - EtcdRevisionId& revision_id); - - // C-style callback for etcd Watch goroutine. - static void WatchCallback(void* context, const char* key, size_t key_size, - const char* value, size_t value_size, - int event_type, int64_t mod_revision); - - // Background thread running the watch loop. - void WatchLoop(); - - // Handle a single watch event (PUT/DELETE/BROKEN). - void HandleWatchEvent(const std::string& key, const std::string& value, - int event_type, int64_t mod_revision); - - // Reconnection with exponential backoff. - void TryReconnect(); - - // Sync missed entries after reconnection. - bool SyncMissedEntries(); - - // Read and deliver all entries since start_seq_id. Updates - // last_processed_sequence_id_ and next_watch_revision_. - // Returns the number of delivered entries, or -1 on read failure. - int64_t DeliverHistoricalEntries(uint64_t start_seq_id); - - std::string cluster_id_; - std::string watch_prefix_; // "/oplog/{cluster_id}/" - EtcdOpLogStore* oplog_store_; // Not owned - - EntryCallback on_entry_; - ErrorCallback on_error_; - - std::atomic running_{false}; - std::thread watch_thread_; - std::atomic last_processed_sequence_id_{0}; - std::atomic next_watch_revision_{0}; - - ChangeNotifierCallbackContext* callback_ctx_{nullptr}; - - std::atomic consecutive_errors_{0}; - std::atomic reconnect_count_{0}; - std::atomic watch_healthy_{false}; - - static constexpr int kMaxConsecutiveErrors = 10; - static constexpr int kReconnectDelayMs = 1000; - static constexpr int kMaxReconnectDelayMs = 30000; - static constexpr int kSyncBatchSize = 1000; -}; - -} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/etcd_oplog_store.h b/mooncake-store/include/ha/oplog/etcd_oplog_store.h deleted file mode 100644 index 1508c56f71..0000000000 --- a/mooncake-store/include/ha/oplog/etcd_oplog_store.h +++ /dev/null @@ -1,237 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "ha/oplog/oplog_manager.h" -#include "ha/oplog/oplog_store.h" -#include "types.h" - -namespace mooncake { - -/** - * @brief Store for OpLog entries in etcd. - * - * This class is responsible for writing OpLog entries to etcd and reading them - * back. OpLog entries are stored with keys in the format: - * /oplog/{cluster_id}/{sequence_id} - * - * The latest sequence_id is also stored at: - * /oplog/{cluster_id}/latest - */ -class EtcdOpLogStore : public OpLogStore { - public: - /** - * @brief Constructor. - * @param cluster_id: The cluster ID for this OpLog store. - * @param enable_latest_seq_batch_update: Whether to start background thread - * to batch-update `/latest`. Readers (Standby) should set this to - * false to avoid unnecessary thread creation. - * @param enable_batch_write: Whether to start the OpLog batch-write - * background thread. Readers (Standby) that only call Read* - * methods should set this to false to avoid unnecessary thread - * creation and /latest key initialization overhead. - */ - explicit EtcdOpLogStore(const std::string& cluster_id, - bool enable_latest_seq_batch_update = false, - bool enable_batch_write = false); - - ~EtcdOpLogStore(); - - /** - * @brief Initialize the store. - * Must be called after construction and before use. - * Performs necessary I/O (e.g. initializing /latest key) and starts - * background threads if enabled. - * @return: Error code. - */ - ErrorCode Init() override; - - /** - * @brief Write an OpLog entry to etcd. - * @param entry: The OpLog entry to write. - * @param sync: If true, wait until the entry is persisted to etcd. - * If false, buffer it and return immediately (Group Commit). - * @return: Error code. - */ - ErrorCode WriteOpLog(const OpLogEntry& entry, bool sync = true) override; - - /** - * @brief Read an OpLog entry from etcd by sequence_id. - * @param sequence_id: The sequence ID of the entry to read. - * @param entry: Output param, the OpLog entry. - * @return: Error code. - */ - ErrorCode ReadOpLog(uint64_t sequence_id, OpLogEntry& entry) override; - - /** - * @brief Read OpLog entries starting from a given sequence_id. - * @param start_sequence_id: The starting sequence ID (exclusive). - * @param limit: Maximum number of entries to read (default: 1000). - * @param entries: Output param, vector of OpLog entries. - * @return: Error code. - */ - ErrorCode ReadOpLogSince(uint64_t start_sequence_id, size_t limit, - std::vector& entries) override; - - // Like ReadOpLogSince, but also returns the etcd revision for consistent - // "read then watch(from revision+1)" startup. - ErrorCode ReadOpLogSinceWithRevision(uint64_t start_sequence_id, - size_t limit, - std::vector& entries, - EtcdRevisionId& revision_id); - - /** - * @brief Get the latest sequence_id from etcd. - * @param sequence_id: Output param, the latest sequence_id. - * @return: Error code. OPLOG_ENTRY_NOT_FOUND if no OpLog exists yet. - */ - ErrorCode GetLatestSequenceId(uint64_t& sequence_id) override; - - // Stronger (than `/latest`) best-effort query: return the maximum existing - // sequence_id by scanning etcd keys under /oplog/{cluster_id}/ with - // descending key order. - // Return OPLOG_ENTRY_NOT_FOUND if no OpLog exists yet. - ErrorCode GetMaxSequenceId(uint64_t& sequence_id) override; - - /** - * @brief Update the latest sequence_id in etcd. - * @param sequence_id: The latest sequence_id to update. - * @return: Error code. - */ - ErrorCode UpdateLatestSequenceId(uint64_t sequence_id) override; - - /** - * @brief Record the sequence_id corresponding to a snapshot. - * @param snapshot_id: The snapshot ID. - * @param sequence_id: The sequence_id at which the snapshot was taken. - * @return: Error code. - */ - ErrorCode RecordSnapshotSequenceId(const std::string& snapshot_id, - uint64_t sequence_id) override; - - /** - * @brief Get the sequence_id for a given snapshot. - * @param snapshot_id: The snapshot ID. - * @param sequence_id: Output param, the sequence_id. - * @return: Error code. OPLOG_ENTRY_NOT_FOUND if snapshot not found. - */ - ErrorCode GetSnapshotSequenceId(const std::string& snapshot_id, - uint64_t& sequence_id) override; - - /** - * @brief Clean up OpLog entries before a given sequence_id. - * @param before_sequence_id: All entries with sequence_id < - * before_sequence_id will be deleted. - * @return: Error code. - */ - ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id) override; - - // Create an EtcdOpLogChangeNotifier backed by this store. - std::unique_ptr CreateChangeNotifier( - const std::string& cluster_id) override; - - private: - /** - * @brief Build the etcd key for an OpLog entry. - * @param sequence_id: The sequence ID. - * @return: The etcd key. - */ - std::string BuildOpLogKey(uint64_t sequence_id) const; - - /** - * @brief Build the etcd key for the latest sequence_id. - * @return: The etcd key. - */ - std::string BuildLatestKey() const; - - /** - * @brief Build the etcd key for a snapshot sequence_id. - * @param snapshot_id: The snapshot ID. - * @return: The etcd key. - */ - std::string BuildSnapshotKey(const std::string& snapshot_id) const; - - // Best-effort: find the minimum existing OpLog sequence_id in etcd. - // Used for robust cleanup (Scheme 3) so we don't rely on a persisted - // "cleaned_upto" marker. - std::optional GetMinSequenceId() const; - - // Best-effort: find the maximum existing OpLog sequence_id in etcd. - std::optional GetMaxSequenceIdInternal() const; - - /** - * @brief Batch update thread function. - * Periodically updates latest_sequence_id in etcd. - */ - void BatchUpdateThread(); - - /** - * @brief Trigger immediate batch update if threshold is reached. - */ - void TriggerBatchUpdateIfNeeded(); - - /** - * @brief Perform the actual batch update to etcd. - */ - void DoBatchUpdate(); - - std::string cluster_id_; - static constexpr const char* kOpLogPrefix = "/oplog/"; - static constexpr const char* kLatestSuffix = "/latest"; - static constexpr const char* kSnapshotPrefix = "/oplog/"; - static constexpr const char* kSnapshotSuffix = "/snapshot/"; - - // Batch update mechanism for latest_sequence_id - const bool enable_latest_seq_batch_update_{false}; - const bool enable_batch_write_{false}; - std::atomic pending_latest_seq_id_{0}; - std::atomic pending_count_{0}; - std::atomic batch_update_running_{false}; - std::mutex batch_update_mutex_; - std::thread batch_update_thread_; - std::chrono::steady_clock::time_point last_update_time_; - - // Batch update configuration - static constexpr size_t kBatchSize = 100; // Update every 100 entries - static constexpr int kBatchIntervalMs = 1000; // Or every 1 second - - // Group Commit / Batch Write support - struct BatchEntry { - std::string key; - std::string value; - uint64_t sequence_id; - bool is_sync; // Track if entry requires sync - }; - - void BatchWriteThread(); - void FlushBatch(); - - mutable std::mutex batch_mutex_; - std::deque pending_batch_; - std::condition_variable cv_batch_updated_; // Notify background thread - std::condition_variable cv_sync_completed_; // Notify sync waiters - std::atomic batch_write_running_{false}; - std::thread batch_write_thread_; - std::atomic last_persisted_seq_id_{0}; - - // Configs for OpLog batching - static constexpr size_t kOpLogBatchSizeLimit = - 1 * 1024 * 1024; // 1MB payload limit (soft) - static constexpr size_t kOpLogBatchCountLimit = 100; // 100 entries - static constexpr int kOpLogBatchTimeoutMs = - 10; // 10ms max latency for Async - static constexpr int kSyncWaitTimeoutMs = - 3000; // 3s timeout for Sync writes - static constexpr int kFlushRetryCount = 3; // Retries for failed flush - static constexpr int kFlushRetryIntervalMs = 50; // Retry interval -}; - -} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/localfs_oplog_store.h b/mooncake-store/include/ha/oplog/localfs_oplog_store.h deleted file mode 100644 index 33ce61a6af..0000000000 --- a/mooncake-store/include/ha/oplog/localfs_oplog_store.h +++ /dev/null @@ -1,146 +0,0 @@ -// mooncake-store/include/localfs_oplog_store.h -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "ha/oplog/oplog_manager.h" -#include "ha/oplog/oplog_store.h" -#include "types.h" - -namespace mooncake { - -class LocalFsOpLogStore : public OpLogStore { - public: - explicit LocalFsOpLogStore(const std::string& cluster_id, - const std::string& root_dir, - bool enable_batch_write, - int poll_interval_ms = 1000); - ~LocalFsOpLogStore(); - - ErrorCode Init() override; - ErrorCode WriteOpLog(const OpLogEntry& entry, bool sync = true) override; - ErrorCode ReadOpLog(uint64_t sequence_id, OpLogEntry& entry) override; - ErrorCode ReadOpLogSince(uint64_t start_sequence_id, size_t limit, - std::vector& entries) override; - ErrorCode GetLatestSequenceId(uint64_t& sequence_id) override; - ErrorCode GetMaxSequenceId(uint64_t& sequence_id) override; - ErrorCode UpdateLatestSequenceId(uint64_t sequence_id) override; - ErrorCode RecordSnapshotSequenceId(const std::string& snapshot_id, - uint64_t sequence_id) override; - ErrorCode GetSnapshotSequenceId(const std::string& snapshot_id, - uint64_t& sequence_id) override; - ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id) override; - std::unique_ptr CreateChangeNotifier( - const std::string& cluster_id) override; - - private: - // Segment file format: 32-byte header + length-prefixed entries - static constexpr char kSegmentMagic[4] = {'M', 'C', 'S', 'G'}; - static constexpr uint32_t kSegmentVersion = 1; - static constexpr size_t kSegmentHeaderSize = 32; - - // Segment header layout (all little-endian): - // [0..3] magic "MCSG" - // [4..7] version uint32_t - // [8..15] min_seq uint64_t - // [16..23] max_seq uint64_t - // [24..27] count uint32_t - // [28..31] reserved uint32_t (0) - struct SegmentHeader { - char magic[4]; - uint32_t version; - uint64_t min_seq; - uint64_t max_seq; - uint32_t entry_count; - uint32_t reserved; - }; - static_assert(sizeof(SegmentHeader) == kSegmentHeaderSize, - "SegmentHeader must be 32 bytes"); - - // Segment file name info parsed from filename - struct SegmentInfo { - std::string filename; - uint64_t min_seq; - uint64_t max_seq; - }; - - // Group Commit batch entry - struct BatchEntry { - std::string serialized_value; - uint64_t sequence_id; - bool is_sync; - }; - - // Directory and path helpers - std::string SegmentsDir() const; - std::string SnapshotsDir() const; - std::string LatestFilePath() const; - std::string BuildSegmentFilename(uint64_t min_seq, uint64_t max_seq) const; - std::string BuildSnapshotPath(const std::string& snapshot_id) const; - - // Segment I/O - ErrorCode WriteSegmentFile(const std::vector& entries); - ErrorCode ReadSegmentEntries(const std::string& filepath, - std::vector& entries); - ErrorCode ReadSegmentHeader(const std::string& filepath, - SegmentHeader& header); - std::vector ListSegments() const; - static bool ParseSegmentFilename(const std::string& filename, - uint64_t& min_seq, uint64_t& max_seq); - ErrorCode NormalizeBatchEntries(std::vector& entries) const; - ErrorCode VerifyPersistedEntryMatches(uint64_t sequence_id, - const std::string& serialized_value); - ErrorCode RecoverPersistedState(); - - // Atomic file write: write to .tmp, fsync, rename - ErrorCode AtomicWriteFile(const std::string& target_path, - const std::string& content); - ErrorCode AtomicWriteFile(const std::string& target_path, const void* data, - size_t size); - - // Cleanup temp files from previous crash - void CleanupTempFiles(); - - // Snapshot ID validation - static bool ValidateSnapshotId(const std::string& snapshot_id); - - // Read a uint64 value from a single-value text file - ErrorCode ReadUint64FromFile(const std::string& filepath, - uint64_t& value) const; - - // Batch write thread - void BatchWriteThread(); - void FlushBatch(); - - // Members - std::string cluster_id_; - std::string root_dir_; - std::string cluster_dir_; // root_dir_/cluster_id_ - bool enable_batch_write_; - int poll_interval_ms_; - - // Group Commit state - mutable std::mutex batch_mutex_; - std::deque pending_batch_; - std::condition_variable cv_batch_updated_; - std::condition_variable cv_sync_completed_; - std::atomic batch_write_running_{false}; - std::thread batch_write_thread_; - std::atomic last_persisted_seq_id_{0}; - - // Batch write configs - static constexpr size_t kBatchCountLimit = 100; - static constexpr int kBatchTimeoutMs = 100; - static constexpr int kSyncWaitTimeoutMs = 3000; - static constexpr int kFlushRetryCount = 3; - static constexpr int kFlushRetryIntervalMs = 50; -}; - -} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_applier.h b/mooncake-store/include/ha/oplog/oplog_applier.h index 7bafb615a8..e3cb866b0d 100644 --- a/mooncake-store/include/ha/oplog/oplog_applier.h +++ b/mooncake-store/include/ha/oplog/oplog_applier.h @@ -1,22 +1,15 @@ #pragma once -#include #include #include -#include -#include -#include #include #include -#include "ha/oplog/oplog_manager.h" +#include "ha/oplog/oplog_types.h" #include "metadata_store.h" namespace mooncake { -// Forward declaration -class OpLogStore; - /** * @brief Apply OpLog entries to Standby metadata store with ordering guarantee * @@ -29,21 +22,12 @@ class OpLogApplier { * @brief Constructor * @param metadata_store Metadata store to apply changes to * @param cluster_id Cluster ID (for validation only) - * @param oplog_store Optional OpLogStore for requesting missing OpLog - * entries (caller owns the pointer) */ explicit OpLogApplier(MetadataStore* metadata_store, - const std::string& cluster_id = std::string(), - OpLogStore* oplog_store = nullptr); + const std::string& cluster_id = std::string()); ~OpLogApplier() = default; - /** - * @brief Set or replace the OpLogStore used for requesting missing entries - * @param oplog_store OpLogStore pointer (caller owns the pointer) - */ - void SetOpLogStore(OpLogStore* oplog_store) { oplog_store_ = oplog_store; } - /** * @brief Apply a single OpLog entry (with ordering checks) * @param entry OpLog entry to apply @@ -70,35 +54,18 @@ class OpLogApplier { */ void Recover(uint64_t last_applied_sequence_id); - /** - * @brief Process pending entries (entries with non-continuous sequence IDs) - * @return Number of entries processed - */ - size_t ProcessPendingEntries(); - - // Promotion helper: - // Try to resolve current gaps ONCE (no waiting) by fetching missing/skipped - // sequence_ids from etcd. If an entry arrives late: - // - REMOVE / PUT_REVOKE: delete the key - // - PUT_END: discard - // - // This is used during Standby promotion so we don't block promotion on - // gaps, but still best-effort clean up potentially stale metadata. - struct GapResolveResult { - size_t attempted{0}; - size_t fetched{0}; - size_t applied_deletes{0}; - }; - GapResolveResult TryResolveGapsOnceForPromotion(size_t max_ids = 1024); + const StandbySegmentRegistry& GetSegmentRegistry() const; + void ApplySegmentMount(const OpLogEntry& entry); + void ApplySegmentUnmount(const OpLogEntry& entry); + void ApplySegmentUpdate(const OpLogEntry& entry); - private: /** - * @brief Check if the entry's sequence order is valid - * @param entry OpLog entry - * @return true if order is valid, false otherwise + * @brief Load segment registry from snapshot baseline. + * Clears existing registry and replaces with given segments. */ - bool CheckSequenceOrder(const OpLogEntry& entry); + void LoadSegmentRegistry(const std::vector& segments); + private: /** * @brief Apply PUT_END operation * @param entry OpLog entry @@ -117,49 +84,16 @@ class OpLogApplier { */ void ApplyRemove(const OpLogEntry& entry); - /** - * @brief Request missing OpLog entry from the store - * @param missing_seq_id Missing sequence ID - * @return true if entry was found and applied, false otherwise - */ - bool RequestMissingOpLog(uint64_t missing_seq_id); - MetadataStore* metadata_store_; - // OpLogStore for requesting missing OpLog entries (optional, not owned) std::string cluster_id_; - OpLogStore* oplog_store_{nullptr}; - - // Note: key_sequence_map_ has been removed. - // Global sequence_id is sufficient for ordering guarantee. - - // Track pending entries (entries with non-continuous sequence IDs) - mutable std::mutex pending_mutex_; - std::map pending_entries_; - - // Track missing sequence IDs that we're waiting for - std::map - missing_sequence_ids_; - - // Sequence IDs we chose to skip (gap-timeout). If the late entry arrives: - // - REMOVE / PUT_REVOKE: delete the key (safe) - // - PUT_END: discard (do not resurrect potentially stale metadata) - std::map - skipped_sequence_ids_; // Next expected global sequence_id. Read frequently from monitoring thread, // updated by watch/apply thread. Use atomic to avoid data races. std::atomic expected_sequence_id_{1}; - // Constants for missing entry handling - // IMPORTANT: request must happen BEFORE skip, otherwise we will never - // request. - static constexpr int kMissingEntryRequestSeconds = - 1; // request from etcd after 1s - static constexpr int kMissingEntrySkipSeconds = - 3; // skip after 3s (avoid global stall) - static constexpr int kMaxPendingEntries = - 1000; // Max pending entries before giving up + // Standby segment registry + StandbySegmentRegistry segment_registry_; }; } // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_batch_codec.h b/mooncake-store/include/ha/oplog/oplog_batch_codec.h new file mode 100644 index 0000000000..32e9d0c7db --- /dev/null +++ b/mooncake-store/include/ha/oplog/oplog_batch_codec.h @@ -0,0 +1,17 @@ +#pragma once + +#include + +#include "ha/oplog/oplog_batch_types.h" + +namespace mooncake { + +std::string EncodeDurablePrefix(const DurablePrefix& prefix); +bool DecodeDurablePrefix(const std::string& value, DurablePrefix* prefix, + std::string* reason = nullptr); + +std::string EncodeOpLogBatchRecord(const OpLogBatchRecord& batch); +bool DecodeOpLogBatchRecord(const std::string& value, OpLogBatchRecord* batch, + std::string* reason = nullptr); + +} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_batch_standby_reader.h b/mooncake-store/include/ha/oplog/oplog_batch_standby_reader.h new file mode 100644 index 0000000000..94277049c2 --- /dev/null +++ b/mooncake-store/include/ha/oplog/oplog_batch_standby_reader.h @@ -0,0 +1,47 @@ +#pragma once + +#include +#include +#include +#include + +#include "ha/oplog/oplog_batch_storage.h" +#include "types.h" + +namespace mooncake { + +class HaKvBackend; +class OpLogApplier; + +enum class OpLogBatchStandbyPollDisposition { + OK, + RETRYABLE, + FATAL, +}; + +struct OpLogBatchStandbyPollResult { + OpLogBatchStandbyPollDisposition disposition{ + OpLogBatchStandbyPollDisposition::OK}; + ErrorCode error{ErrorCode::OK}; + bool durable_prefix_present{false}; + size_t applied_entries{0}; + DurablePrefix durable_prefix{}; +}; + +class OpLogBatchStandbyReader { + public: + OpLogBatchStandbyReader(std::string cluster_id, HaKvBackend& backend, + OpLogApplier& applier); + + OpLogBatchStandbyPollResult PollOnce(size_t max_batches = 1024); + + private: + OpLogBatchStorage storage_; + OpLogApplier& applier_; + bool batch_format_seen_{false}; + std::optional last_observed_prefix_; + std::optional last_scanned_batch_last_seq_; + uint64_t last_applied_batch_id_{0}; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_batch_storage.h b/mooncake-store/include/ha/oplog/oplog_batch_storage.h new file mode 100644 index 0000000000..eade16b397 --- /dev/null +++ b/mooncake-store/include/ha/oplog/oplog_batch_storage.h @@ -0,0 +1,34 @@ +#pragma once + +#include +#include + +#include "ha/kv/ha_kv_backend.h" +#include "ha/oplog/oplog_batch_types.h" +#include "types.h" + +namespace mooncake { + +class OpLogBatchStorage { + public: + OpLogBatchStorage(std::string cluster_id, HaKvBackend& backend); + + ErrorCode InitDurablePrefix(DurablePrefix& prefix); + ErrorCode ReadDurablePrefix(DurablePrefix& prefix); + ErrorCode WriteBatchAndAdvancePrefix(const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix); + ErrorCode ReadBatch(uint64_t batch_id, OpLogBatchRecord& batch); + ErrorCode ReadBatchesAfter(uint64_t after_batch_id, size_t limit, + std::vector& batches); + + private: + bool IsValidClusterId() const; + ErrorCode RejectLegacyLayout() const; + ErrorCode ValidateDurablePrefixAtStartup(const DurablePrefix& prefix); + + std::string cluster_id_; + HaKvBackend& backend_; + bool cluster_id_valid_{false}; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_batch_types.h b/mooncake-store/include/ha/oplog/oplog_batch_types.h new file mode 100644 index 0000000000..4e315383c1 --- /dev/null +++ b/mooncake-store/include/ha/oplog/oplog_batch_types.h @@ -0,0 +1,50 @@ +#pragma once + +#include +#include +#include + +#include "ha/oplog/oplog_types.h" + +namespace mooncake { + +static constexpr uint32_t kOpLogBatchRecordSchemaVersion = 1; +static constexpr uint32_t kDurablePrefixSchemaVersion = 1; +static constexpr int kOpLogBatchIdWidth = 20; + +struct DurablePrefix { + uint64_t batch_id{0}; + uint64_t last_seq{0}; +}; + +struct BatchRecordRange { + std::string begin_key; + std::string end_key; +}; + +struct OpLogBatchRecord { + uint32_t schema_version{kOpLogBatchRecordSchemaVersion}; + uint64_t batch_id{0}; + uint64_t first_seq{0}; + uint64_t last_seq{0}; + std::vector entries; + uint32_t checksum{0}; +}; + +bool ValidateOpLogBatchRecordShape(const OpLogBatchRecord& batch, + std::string* reason = nullptr); + +bool ValidateOpLogBatchEntry(const OpLogEntry& entry, + std::string* reason = nullptr); + +bool ValidateOpLogBatchClusterId(const std::string& cluster_id, + std::string* reason = nullptr); + +std::string FormatOpLogBatchId(uint64_t batch_id); +std::string BuildBatchRecordKey(const std::string& cluster_id, + uint64_t batch_id); +std::string BuildDurablePrefixKey(const std::string& cluster_id); +BatchRecordRange BuildBatchRecordRange(const std::string& cluster_id, + uint64_t after_batch_id); + +} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_change_notifier.h b/mooncake-store/include/ha/oplog/oplog_change_notifier.h deleted file mode 100644 index 2471c327b4..0000000000 --- a/mooncake-store/include/ha/oplog/oplog_change_notifier.h +++ /dev/null @@ -1,40 +0,0 @@ -// mooncake-store/include/oplog_change_notifier.h -#pragma once - -#include -#include - -#include "ha/oplog/oplog_manager.h" -#include "types.h" - -namespace mooncake { - -// Abstract interface for watching OpLog changes. -// -// Delivery semantics: -// The notifier provides **at-most-once delivery**. Its internal cursor -// tracks *fetch progress* (where to read from next), NOT consumer -// processing progress. The cursor advances unconditionally after entries -// are dispatched to EntryCallback, regardless of whether the callback -// processed them successfully. -// -// Consumers (e.g. OpLogApplier) must maintain their own cursor -// (expected_sequence_id_) and handle gaps, duplicates, and late arrivals -// independently. -// -// Implementations: EtcdOpLogChangeNotifier (push via Watch), -// PollingOpLogChangeNotifier (poll via ReadOpLogSince) -class OpLogChangeNotifier { - public: - virtual ~OpLogChangeNotifier() = default; - - using EntryCallback = std::function; - using ErrorCallback = std::function; - - virtual ErrorCode Start(uint64_t start_sequence_id, EntryCallback on_entry, - ErrorCallback on_error) = 0; - virtual void Stop() = 0; - virtual bool IsHealthy() const = 0; -}; - -} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_manager.h b/mooncake-store/include/ha/oplog/oplog_manager.h deleted file mode 100644 index da69bff13f..0000000000 --- a/mooncake-store/include/ha/oplog/oplog_manager.h +++ /dev/null @@ -1,154 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -#include - -#include "types.h" - -namespace mooncake { - -// Forward declaration -class OpLogStore; - -// Operation types for hot-standby replication. -// This is a minimal subset that can be extended later. -enum class OpType : uint8_t { - PUT_END = 1, - PUT_REVOKE = 2, - REMOVE = 3, - // Deprecated: LEASE_RENEW is intentionally not recorded in OpLog in the - // current etcd-based hot-standby design (Standby relies on Primary DELETE - // operations). - LEASE_RENEW = 4, -}; - -// A single operation log entry. -// Note: Payload contains JSON serialized MetadataPayload (defined in -// metadata_store.h) for PUT_END operations, allowing Standby to restore -// complete metadata. -struct OpLogEntry { - uint64_t sequence_id{0}; // Monotonically increasing global sequence - uint64_t timestamp_ms{0}; // Logical timestamp in milliseconds - OpType op_type{OpType::PUT_END}; - std::string object_key; // Target object key - std::string payload; // Serialized extra data (optional) - uint32_t checksum{0}; // Checksum of payload (implementation-defined) - uint32_t prefix_hash{ - 0}; // Hash of the entire key (for verification and optimization) -}; - -/** - * @brief In-memory operation log manager. - * - * This class is intentionally simple: it keeps a bounded deque of OpLogEntry - * and provides append / get-since primitives. It can later be extended to - * or to spill to disk if needed. OpLog entries are persisted to the - * configured OpLogStore backend (etcd, local filesystem, etc.). - */ -class OpLogManager { - public: - OpLogManager(); - - // Set the OpLogStore for writing OpLog to persistent storage (optional). - // If not set, OpLog will only be stored in memory buffer. - void SetOpLogStore(std::shared_ptr oplog_store); - - // Append a new entry and return the assigned sequence_id. - // This is a best-effort (async) path: the entry is buffered in memory - // and enqueued to etcd without waiting for persistence. - // Only suitable for idempotent, lag-tolerant operations (PUT_END). - // For operations that MUST be durable before returning (REMOVE, etc.), - // use AppendAndPersist() instead. - uint64_t Append(OpType type, const std::string& key, - const std::string& payload = std::string()); - - // Allocate a new OpLogEntry with a reserved sequence_id, append it to the - // in-memory buffer, and return the full entry. - // - // IMPORTANT: This will advance last_seq_id_ even if the caller later fails - // to persist it to etcd. This supports "seq pre-allocation" semantics where - // retries use the same (smaller) sequence_id. - OpLogEntry AllocateEntry(OpType type, const std::string& key, - const std::string& payload = std::string()); - - // Persist an already-allocated entry to the store using its sequence_id. - // Does NOT modify sequence counters. - ErrorCode PersistEntry(const OpLogEntry& entry) const; - - // Append a new entry and durably persist it to the store (if OpLogStore is - // set). - // - // This is intended for operations that may free/reuse memory (e.g. REMOVE), - // where best-effort replication is unsafe: Standby must observe the DELETE - // before promotion, otherwise it may return stale descriptors that point to - // reused memory and cause silent data corruption. - // - // Design (updated for seq pre-allocation): - // - sequence_id is allocated first and never reused. - // - If etcd write fails, caller may retry PersistEntry with the same - // entry (sequence_id fixed and "smaller" than later entries). - tl::expected AppendAndPersist( - OpType type, const std::string& key, - const std::string& payload = std::string()); - - // Get the latest assigned sequence id. Returns 0 if no entry exists. - uint64_t GetLastSequenceId() const; - - // Set the initial sequence ID (used when promoting Standby to Primary). - // This ensures the new Primary's OpLogManager continues from the correct - // sequence_id. - void SetInitialSequenceId(uint64_t sequence_id); - - // Current number of entries in the buffer. - size_t GetEntryCount() const; - - // Clean up OpLog entries in etcd before a given sequence_id. - // Delegates to EtcdOpLogStore::CleanupOpLogBefore. - ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id); - - // Verify checksum of an OpLogEntry payload. - // Returns true if checksum matches, false otherwise. - // This is public so OpLogReplicator and OpLogApplier can validate entries. - static bool VerifyChecksum(const OpLogEntry& entry); - - // Basic DoS protection for externally sourced OpLog entries (etcd watch / - // reads). Enforce conservative bounds on key/payload sizes before - // parsing/applying. - static constexpr size_t kMaxObjectKeySize = 4096; // 4 KiB - static constexpr size_t kMaxPayloadSize = 10 * 1024 * 1024; // 10 MiB - - // Validate OpLogEntry key/payload sizes. If invalid, returns false and - // optionally sets a human-readable reason. - static bool ValidateEntrySize(const OpLogEntry& entry, - std::string* reason = nullptr); - - private: - static uint64_t NowMs(); - static uint32_t ComputeChecksum(const std::string& data); - static uint32_t ComputePrefixHash(const std::string& key); - - mutable std::shared_mutex mutex_; - std::deque buffer_; - uint64_t first_seq_id_{1}; // sequence_id of buffer_.front() - uint64_t last_seq_id_{0}; // last assigned sequence_id - - // Note: We removed key_sequence_map_ and key_remove_time_map_. - // Global sequence_id is sufficient for ordering guarantee. - // All operations are applied in sequence_id order, which ensures - // consistency. - - // Optional OpLog store for persistent storage - std::shared_ptr oplog_store_; - - // Simple bounds to avoid unbounded memory growth. - static constexpr size_t kMaxBufferEntries_ = 100000; -}; - -} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_replicator.h b/mooncake-store/include/ha/oplog/oplog_replicator.h deleted file mode 100644 index 1b86fbeaf3..0000000000 --- a/mooncake-store/include/ha/oplog/oplog_replicator.h +++ /dev/null @@ -1,87 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include - -#include "ha/oplog/oplog_change_notifier.h" -#include "ha/oplog/oplog_manager.h" -#include "standby_state_machine.h" -#include "types.h" - -namespace mooncake { - -// Forward declarations -class OpLogApplier; - -// Callback type for state events -using ReplicatorStateCallback = std::function; - -/** - * @brief Replicate OpLog entries from a remote source and apply them locally. - * - * Delegates watch/notification to an OpLogChangeNotifier and applies - * received entries via OpLogApplier. This class is a thin orchestration - * layer; the actual watch implementation lives in OpLogChangeNotifier. - */ -class OpLogReplicator { - public: - /** - * @brief Constructor - * @param notifier Change notifier that delivers OpLog entries - * @param applier OpLog applier to process entries - */ - OpLogReplicator(OpLogChangeNotifier* notifier, OpLogApplier* applier); - - ~OpLogReplicator(); - - /** - * @brief Start replication from the beginning. - */ - void Start(); - - /** - * @brief Start from a known last-applied sequence_id. - */ - bool StartFromSequenceId(uint64_t start_seq_id); - - /** - * @brief Stop replication. - */ - void Stop(); - - /** - * @brief Get the last processed sequence ID. - */ - uint64_t GetLastProcessedSequenceId() const; - - /** - * @brief Set callback for state events. - */ - void SetStateCallback(ReplicatorStateCallback callback) { - state_callback_ = std::move(callback); - } - - /** - * @brief Check if replication is healthy. - */ - bool IsHealthy() const; - - private: - void NotifyStateEvent(StandbyEvent event) { - if (state_callback_) { - state_callback_(event); - } - } - - OpLogChangeNotifier* notifier_; - OpLogApplier* applier_; - std::atomic last_processed_sequence_id_{0}; - std::atomic running_{false}; - - ReplicatorStateCallback state_callback_; -}; - -} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_serializer.h b/mooncake-store/include/ha/oplog/oplog_serializer.h deleted file mode 100644 index 6ec09c2f7e..0000000000 --- a/mooncake-store/include/ha/oplog/oplog_serializer.h +++ /dev/null @@ -1,18 +0,0 @@ -// mooncake-store/include/oplog_serializer.h -#pragma once - -#include - -#include "ha/oplog/oplog_manager.h" - -namespace mooncake { - -// Serialize an OpLogEntry to JSON string (with base64-encoded payload). -// Format is backend-agnostic; all storage backends should use this. -std::string SerializeOpLogEntry(const OpLogEntry& entry); - -// Deserialize a JSON string to OpLogEntry. -// Returns true on success, false on parse error or size validation failure. -bool DeserializeOpLogEntry(const std::string& json_str, OpLogEntry& entry); - -} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_store.h b/mooncake-store/include/ha/oplog/oplog_store.h deleted file mode 100644 index ea4806de12..0000000000 --- a/mooncake-store/include/ha/oplog/oplog_store.h +++ /dev/null @@ -1,64 +0,0 @@ -// mooncake-store/include/oplog_store.h -#pragma once - -#include -#include -#include -#include - -#include "ha/oplog/oplog_change_notifier.h" -#include "ha/oplog/oplog_manager.h" -#include "types.h" - -namespace mooncake { - -// Normalize and validate cluster_id for OpLog key prefix construction. -// Strips trailing slashes, then validates the remaining string. -// Returns true if valid (or empty after normalization), false otherwise. -inline bool NormalizeAndValidateClusterId(std::string& cluster_id) { - while (!cluster_id.empty() && cluster_id.back() == '/') { - cluster_id.pop_back(); - } - return cluster_id.empty() || IsValidClusterIdComponent(cluster_id); -} - -// Abstract interface for OpLog persistent storage. -// Implementations: EtcdOpLogStore, (future) HdfsOpLogStore, etc. -class OpLogStore { - public: - virtual ~OpLogStore() = default; - virtual ErrorCode Init() = 0; - - // Write - virtual ErrorCode WriteOpLog(const OpLogEntry& entry, bool sync = true) = 0; - - // Read - virtual ErrorCode ReadOpLog(uint64_t sequence_id, OpLogEntry& entry) = 0; - virtual ErrorCode ReadOpLogSince(uint64_t start_sequence_id, size_t limit, - std::vector& entries) = 0; - - // Sequence ID management - virtual ErrorCode GetLatestSequenceId(uint64_t& sequence_id) = 0; - virtual ErrorCode GetMaxSequenceId(uint64_t& sequence_id) = 0; - virtual ErrorCode UpdateLatestSequenceId(uint64_t sequence_id) = 0; - - // Snapshot - virtual ErrorCode RecordSnapshotSequenceId(const std::string& snapshot_id, - uint64_t sequence_id) = 0; - virtual ErrorCode GetSnapshotSequenceId(const std::string& snapshot_id, - uint64_t& sequence_id) = 0; - - // Cleanup - virtual ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id) = 0; - - // Create a change notifier for this store. - // Each backend provides its own notifier (e.g., etcd watch, polling). - // Returns nullptr if the backend does not support change notification. - virtual std::unique_ptr CreateChangeNotifier( - const std::string& cluster_id) { - (void)cluster_id; - return nullptr; - } -}; - -} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_store_factory.h b/mooncake-store/include/ha/oplog/oplog_store_factory.h deleted file mode 100644 index 7b358ca513..0000000000 --- a/mooncake-store/include/ha/oplog/oplog_store_factory.h +++ /dev/null @@ -1,76 +0,0 @@ -// mooncake-store/include/oplog_store_factory.h -#pragma once - -#include -#include -#include -#include - -#include "ha/oplog/oplog_store.h" - -namespace mooncake { - -enum class OpLogStoreRole { - WRITER, // Primary: enable batch_write + batch_update - READER, // Standby: read-only -}; - -enum class OpLogStoreType { - ETCD, - LOCAL_FS, -}; - -#ifdef STORE_USE_ETCD -static constexpr OpLogStoreType kDefaultOpLogStoreType = OpLogStoreType::ETCD; -#else -static constexpr OpLogStoreType kDefaultOpLogStoreType = - OpLogStoreType::LOCAL_FS; -#endif - -// Parse string to OpLogStoreType (case-insensitive). -// Parse string to OpLogStoreType (case-insensitive). -// Returns kDefaultOpLogStoreType for unrecognized strings. -inline OpLogStoreType ParseOpLogStoreType(const std::string& type_str) { - std::string lower = type_str; - std::transform(lower.begin(), lower.end(), lower.begin(), - [](unsigned char c) -> char { return std::tolower(c); }); - if (lower == "localfs" || lower == "local_fs") { - return OpLogStoreType::LOCAL_FS; - } - if (lower == "etcd") { - return OpLogStoreType::ETCD; - } - return kDefaultOpLogStoreType; -} - -inline std::string OpLogStoreTypeToString(OpLogStoreType type) { - switch (type) { - case OpLogStoreType::LOCAL_FS: - return "localfs"; - case OpLogStoreType::ETCD: - default: - return "etcd"; - } -} - -// Default configuration values for LocalFS OpLog store -static constexpr const char* kDefaultOpLogRootDir = "/tmp/mooncake_oplog"; -static constexpr int kDefaultOpLogPollIntervalMs = 1000; - -class OpLogStoreFactory { - public: - /** - * @brief Create and initialize an OpLogStore instance. - * - * The returned instance is fully initialized (Init() has already been - * called internally). Callers must NOT call Init() again. - * Returns nullptr if the requested backend is unavailable or - * initialization fails. - */ - static std::unique_ptr Create( - OpLogStoreType type, const std::string& cluster_id, OpLogStoreRole role, - const std::string& oplog_root_dir = kDefaultOpLogRootDir, - int poll_interval_ms = kDefaultOpLogPollIntervalMs); -}; - -} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_test_failpoint.h b/mooncake-store/include/ha/oplog/oplog_test_failpoint.h new file mode 100644 index 0000000000..3080544dd4 --- /dev/null +++ b/mooncake-store/include/ha/oplog/oplog_test_failpoint.h @@ -0,0 +1,12 @@ +#pragma once + +#include + +namespace mooncake { + +class TestFailPoint { + public: + static bool Wait(std::string_view name); +}; + +} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/oplog_types.h b/mooncake-store/include/ha/oplog/oplog_types.h new file mode 100644 index 0000000000..04e1880531 --- /dev/null +++ b/mooncake-store/include/ha/oplog/oplog_types.h @@ -0,0 +1,71 @@ +#pragma once + +#include +#include +#include +#include + +#include "types.h" + +namespace mooncake { + +enum class OpType : uint8_t { + PUT_END = 1, + PUT_REVOKE = 2, + REMOVE = 3, + LEASE_RENEW = 4, + SEGMENT_MOUNT = 5, + SEGMENT_UNMOUNT = 6, + SEGMENT_UPDATE = 7, + OP_TYPE_MAX, +}; + +struct SegmentMountOp { + std::string segment_name; + std::string transport_endpoint; + uint64_t capacity{0}; + bool is_memory_segment{false}; + std::string file_path; + + YLT_REFL(SegmentMountOp, segment_name, transport_endpoint, capacity, + is_memory_segment, file_path); +}; + +struct SegmentUnmountOp { + std::string transport_endpoint; + + YLT_REFL(SegmentUnmountOp, transport_endpoint); +}; + +struct SegmentUpdateOp { + std::string segment_name; + std::string transport_endpoint; + uint64_t capacity{0}; + bool is_memory_segment{false}; + std::string file_path; + + YLT_REFL(SegmentUpdateOp, segment_name, transport_endpoint, capacity, + is_memory_segment, file_path); +}; + +struct OpLogEntry { + uint64_t sequence_id{0}; + uint64_t timestamp_ms{0}; + OpType op_type{OpType::PUT_END}; + std::string tenant_id{"default"}; + std::string object_key; + std::string payload; + uint32_t checksum{0}; + uint32_t prefix_hash{0}; +}; + +inline constexpr size_t kMaxOpLogObjectKeySize = 4096; +inline constexpr size_t kMaxOpLogPayloadSize = 10 * 1024 * 1024; + +bool NormalizeAndValidateClusterId(std::string& cluster_id); +uint32_t ComputeOpLogChecksum(std::string_view payload); +bool VerifyOpLogChecksum(const OpLogEntry& entry); +bool ValidateOpLogEntrySize(const OpLogEntry& entry, + std::string* reason = nullptr); + +} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/ordered_oplog_writer.h b/mooncake-store/include/ha/oplog/ordered_oplog_writer.h new file mode 100644 index 0000000000..1d4387d51a --- /dev/null +++ b/mooncake-store/include/ha/oplog/ordered_oplog_writer.h @@ -0,0 +1,75 @@ +#pragma once + +#include +#include +#include +#include + +#include + +#include "ha/oplog/oplog_batch_types.h" +#include "types.h" + +namespace mooncake { + +struct OrderedOpLogWriterConfig { + size_t max_entries_per_batch{1024}; + DurablePrefix initial_durable_prefix{}; +}; + +class OrderedOpLogWriter { + public: + using DurableCallback = std::function; + using WriteBatchFn = + std::function; + + class Reservation { + public: + Reservation(); + Reservation(Reservation&& other) noexcept; + Reservation& operator=(Reservation&& other) noexcept; + Reservation(const Reservation&) = delete; + Reservation& operator=(const Reservation&) = delete; + ~Reservation(); + + private: + friend class OrderedOpLogWriter; + Reservation(OrderedOpLogWriter* writer, uint64_t id); + + OrderedOpLogWriter* writer_{nullptr}; + uint64_t id_{0}; + }; + + class PendingHandle { + public: + PendingHandle(); + uint64_t sequence_id() const; + + private: + friend class OrderedOpLogWriter; + explicit PendingHandle(uint64_t sequence_id); + + uint64_t sequence_id_{0}; + }; + + OrderedOpLogWriter(OrderedOpLogWriterConfig config, + WriteBatchFn write_batch); + ~OrderedOpLogWriter(); + + tl::expected Reserve(); + tl::expected Commit(Reservation&& reservation, + OpLogEntry entry, + DurableCallback callback); + void Abort(Reservation&& reservation); + + bool IsAccepting() const; + ErrorCode LastError() const; + void Start(); + void Stop(); + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/ha/oplog/polling_oplog_change_notifier.h b/mooncake-store/include/ha/oplog/polling_oplog_change_notifier.h deleted file mode 100644 index 5422225bd2..0000000000 --- a/mooncake-store/include/ha/oplog/polling_oplog_change_notifier.h +++ /dev/null @@ -1,45 +0,0 @@ -// mooncake-store/include/polling_oplog_change_notifier.h -#pragma once - -#include -#include -#include -#include - -#include "ha/oplog/oplog_change_notifier.h" -#include "ha/oplog/oplog_store.h" - -namespace mooncake { - -class PollingOpLogChangeNotifier : public OpLogChangeNotifier { - public: - PollingOpLogChangeNotifier(OpLogStore* store, int poll_interval_ms); - ~PollingOpLogChangeNotifier(); - - ErrorCode Start(uint64_t start_sequence_id, EntryCallback on_entry, - ErrorCallback on_error) override; - void Stop() override; - bool IsHealthy() const override; - - private: - void PollLoop(); - - OpLogStore* store_; // Not owned - int poll_interval_ms_; - - EntryCallback on_entry_; - ErrorCallback on_error_; - - std::atomic running_{false}; - std::thread poll_thread_; - std::atomic last_sequence_id_{0}; - std::atomic healthy_{false}; - - // For interruptible sleep on Stop() - std::mutex stop_mutex_; - std::condition_variable stop_cv_; - - static constexpr size_t kPollBatchSize = 1000; -}; - -} // namespace mooncake diff --git a/mooncake-store/include/ha/snapshot/catalog/snapshot_catalog_store.h b/mooncake-store/include/ha/snapshot/catalog/snapshot_catalog_store.h index 763e242c34..5287531536 100644 --- a/mooncake-store/include/ha/snapshot/catalog/snapshot_catalog_store.h +++ b/mooncake-store/include/ha/snapshot/catalog/snapshot_catalog_store.h @@ -1,17 +1,16 @@ #pragma once -#include -#include #include #include #include #include -#include #include #include #include "ha/ha_types.h" +#include "ascii_string.h" +#include "integer_parser.h" namespace mooncake { namespace ha { @@ -36,9 +35,7 @@ inline std::string BuildSnapshotRoot(const std::string& cluster_id) { return root; } -inline bool IsAsciiDigit(char ch) { - return std::isdigit(static_cast(ch)) != 0; -} +inline bool IsAsciiDigit(char ch) { return ch >= '0' && ch <= '9'; } inline bool IsValidSnapshotId(std::string_view snapshot_id) { if (snapshot_id.size() != 19) { @@ -61,17 +58,6 @@ inline bool IsValidSnapshotId(std::string_view snapshot_id) { return true; } -inline std::string TrimAsciiWhitespace(std::string value) { - constexpr std::string_view kAsciiWhitespace = " \t\n\r\f\v"; - const auto first = value.find_first_not_of(kAsciiWhitespace); - if (first == std::string::npos) { - return ""; - } - - const auto last = value.find_last_not_of(kAsciiWhitespace); - return value.substr(first, last - first + 1); -} - inline std::string BuildSnapshotPrefix(const std::string& snapshot_root, const SnapshotId& snapshot_id) { return snapshot_root + snapshot_id + "/"; @@ -104,10 +90,12 @@ inline SnapshotDescriptor MakeSnapshotDescriptor( template inline bool ParseDecimal(std::string_view text, Integer& value) { - const char* begin = text.data(); - const char* end = begin + text.size(); - const auto result = std::from_chars(begin, end, value); - return result.ec == std::errc() && result.ptr == end; + const auto parsed = TryParseInteger(text); + if (!parsed.has_value()) { + return false; + } + value = *parsed; + return true; } inline std::string SerializeSnapshotDescriptor( diff --git a/mooncake-store/include/ha/snapshot/master_snapshot_codec.h b/mooncake-store/include/ha/snapshot/master_snapshot_codec.h new file mode 100644 index 0000000000..2fb533902a --- /dev/null +++ b/mooncake-store/include/ha/snapshot/master_snapshot_codec.h @@ -0,0 +1,153 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "types.h" + +namespace mooncake { + +// Forward declarations +class MasterService; +class SegmentManager; +class NoFSegmentManager; +class ClientTaskManager; + +namespace ha { + +/** + * @brief A view of the live master state for snapshot serialization. + * + * This struct holds references to the live state components that need to + * be serialized into a master snapshot. Non-const references are used because + * the underlying serializers require non-const pointers, avoiding const_cast. + */ +struct MasterSnapshotStateView { + MasterService& master_service; + SegmentManager& segment_manager; + NoFSegmentManager& nof_segment_manager; + ClientTaskManager& task_manager; + + MasterSnapshotStateView(MasterService& ms, SegmentManager& sm, + NoFSegmentManager& nsm, ClientTaskManager& tm) + : master_service(ms), + segment_manager(sm), + nof_segment_manager(nsm), + task_manager(tm) {} +}; + +/** + * @brief Container for serialized master snapshot payloads. + * + * This struct provides a type-safe, zero-overhead container for the three + * serialized payload buffers, eliminating map lookup overhead and potential + * runtime exceptions from using std::unordered_map. + */ +struct MasterSnapshotPayloads { + std::vector metadata; + std::vector segments; + std::vector task_manager; +}; + +/** + * @brief Encodes and decodes master snapshot payloads. + * + * This codec handles serialization of the complete master state bundle: + * - Metadata shards (objects, replicas, tenant state) + * - Segment manager state + * - Task manager state + * - Discarded replicas + * + * The current implementation preserves the existing snapshot format exactly + * to maintain backward compatibility with existing snapshots. + * + * Format details: + * - metadata: msgpack-encoded metadata shards (compressed per-shard with zstd) + * - segments: msgpack-encoded segment manager state + * - task_manager: msgpack-encoded task manager state + * - manifest.txt: format descriptor "||" + * (e.g., "messagepack|1.0.0|snapshot-000123") + */ +class MasterSnapshotCodec { + public: + MasterSnapshotCodec() = default; + ~MasterSnapshotCodec() = default; + + // Non-copyable, non-movable (contains no state, but enforce ownership + // semantics) + MasterSnapshotCodec(const MasterSnapshotCodec&) = delete; + MasterSnapshotCodec& operator=(const MasterSnapshotCodec&) = delete; + MasterSnapshotCodec(MasterSnapshotCodec&&) = delete; + MasterSnapshotCodec& operator=(MasterSnapshotCodec&&) = delete; + + /** + * @brief Encode master state into serialized buffers. + * + * @param state_view View of the live master state + * @return Structured payloads containing serialized data, or error + * + * The returned struct contains: + * - metadata: serialized metadata shards + * - segments: serialized segment manager state + * - task_manager: serialized task manager state + */ + tl::expected Encode( + MasterSnapshotStateView& state_view) const; + + /** + * @brief Decode snapshot payloads and restore into master service. + * + * @param master_service Target MasterService to restore state into + * @param payloads Structured payloads containing serialized data + * @return void on success, SerializationError on failure + */ + tl::expected Decode( + MasterService* master_service, + const MasterSnapshotPayloads& payloads) const; + + // Canonical serializer identifiers embedded in the snapshot manifest. + static constexpr const char* kSerializerType = "messagepack"; + static constexpr const char* kSerializerVersion = "1.0.0"; + + /** + * @brief Encode a snapshot manifest into its on-disk byte representation. + * + * The manifest is a "||" descriptor. Keeping + * the encoding here (rather than hand-crafting the format string at the + * call site) ensures the manifest layout stays owned by the codec. + * + * @param type Serializer/protocol type (e.g., "messagepack") + * @param version Snapshot format version (e.g., "1.0.0") + * @param snapshot_id Identifier of the snapshot being written + * @return Manifest bytes ready to upload + */ + static std::vector EncodeManifest(const std::string& type, + const std::string& version, + const std::string& snapshot_id); + + private: + // Metadata encoding/decoding (delegates to MetadataSerializer for now) + tl::expected, SerializationError> EncodeMetadata( + MasterService& master_service) const; + tl::expected DecodeMetadata( + MasterService* master_service, const std::vector& data) const; + + // Segment encoding/decoding + tl::expected, SerializationError> EncodeSegments( + SegmentManager& segment_manager, + NoFSegmentManager& nof_segment_manager) const; + tl::expected DecodeSegments( + MasterService* master_service, const std::vector& data) const; + + // Task manager encoding/decoding + tl::expected, SerializationError> EncodeTaskManager( + ClientTaskManager& task_manager) const; + tl::expected DecodeTaskManager( + MasterService* master_service, const std::vector& data) const; +}; + +} // namespace ha +} // namespace mooncake diff --git a/mooncake-store/include/ha/snapshot/snapshot_constants.h b/mooncake-store/include/ha/snapshot/snapshot_constants.h new file mode 100644 index 0000000000..da84a8b290 --- /dev/null +++ b/mooncake-store/include/ha/snapshot/snapshot_constants.h @@ -0,0 +1,27 @@ +#pragma once + +#include + +namespace mooncake::ha { + +// Snapshot file names +inline constexpr const char* kSnapshotMetadataFile = "metadata"; +inline constexpr const char* kSnapshotSegmentsFile = "segments"; +inline constexpr const char* kSnapshotTaskManagerFile = "task_manager"; +inline constexpr const char* kSnapshotManifestFile = "manifest.txt"; +inline constexpr const char* kSnapshotLatestFile = "latest.txt"; + +// Snapshot format +inline constexpr const char* kSnapshotSerializerType = "messagepack"; +inline constexpr const char* kSnapshotSerializerVersion = "1.0.0"; + +// Backup directories +inline constexpr const char* kSnapshotBackupSaveDir = + "mooncake_snapshot_save_backup"; +inline constexpr const char* kSnapshotBackupRestoreDir = + "mooncake_snapshot_restore_backup"; + +// List limit +inline constexpr std::size_t kUnlimitedSnapshotList = 0; + +} // namespace mooncake::ha diff --git a/mooncake-store/include/ha/snapshot/snapshot_provider.h b/mooncake-store/include/ha/snapshot/snapshot_provider.h index 66a74ba67f..df8680cb34 100644 --- a/mooncake-store/include/ha/snapshot/snapshot_provider.h +++ b/mooncake-store/include/ha/snapshot/snapshot_provider.h @@ -3,7 +3,6 @@ #include #include #include -#include #include #include @@ -15,7 +14,8 @@ namespace mooncake { struct LoadedSnapshot { std::string snapshot_id; uint64_t snapshot_sequence_id{0}; - std::vector> metadata; + std::vector metadata; + std::vector segments; }; /** diff --git a/mooncake-store/include/ha/standby_controller.h b/mooncake-store/include/ha/standby_controller.h index 33f26a2aa0..f6b738a713 100644 --- a/mooncake-store/include/ha/standby_controller.h +++ b/mooncake-store/include/ha/standby_controller.h @@ -3,14 +3,28 @@ #include #include #include +#include + +#include #include "ha/ha_types.h" #include "master_config.h" +#include "metadata_store.h" #include "types.h" namespace mooncake { namespace ha { +/** + * Context exported from standby at promotion time. + * Contains everything needed to restore primary's state. + */ +struct PromotionContext { + uint64_t applied_seq_id{0}; + std::vector objects; + std::vector segments; +}; + class StandbyController { public: using RuntimeStateCallback = std::function; @@ -24,6 +38,15 @@ class StandbyController { virtual ErrorCode PromoteStandby() = 0; + /** + * Promote standby and export complete context for new primary. + * This replaces the old PromoteStandby() for HA scenarios. + * + * @return PromotionContext on success, error code on failure + */ + virtual tl::expected + PromoteStandbyAndExport() = 0; + virtual void UpdateObservedLeader( const std::optional& observed_leader) = 0; diff --git a/mooncake-store/include/ha_metric_manager.h b/mooncake-store/include/ha_metric_manager.h index d10d040c13..59fc897970 100644 --- a/mooncake-store/include/ha_metric_manager.h +++ b/mooncake-store/include/ha_metric_manager.h @@ -146,6 +146,30 @@ class HAMetricManager { */ void observe_oplog_apply_latency_us(int64_t latency_us); + // ========== Batch-record OpLog Metrics ========== + + void inc_batch_record_durable_batches(int64_t val = 1); + int64_t get_batch_record_durable_batches_total(); + void inc_batch_record_durable_entries(int64_t val = 1); + int64_t get_batch_record_durable_entries_total(); + void inc_batch_record_retries(int64_t val = 1); + int64_t get_batch_record_retries_total(); + + void set_batch_record_committed_queue_depth(int64_t depth); + int64_t get_batch_record_committed_queue_depth(); + void set_batch_record_callback_queue_depth(int64_t depth); + int64_t get_batch_record_callback_queue_depth(); + void set_batch_record_last_batch_id(int64_t batch_id); + int64_t get_batch_record_last_batch_id(); + void set_batch_record_durable_sequence(int64_t sequence_id); + int64_t get_batch_record_durable_sequence(); + + void observe_batch_record_batch_entries(int64_t entries); + void observe_batch_record_batch_bytes(int64_t bytes); + void observe_batch_record_txn_latency_us(int64_t latency_us); + void observe_batch_record_commit_to_durable_us(int64_t latency_us); + void observe_batch_record_callback_latency_us(int64_t latency_us); + // ========== State Machine Metrics ========== /** @@ -207,6 +231,19 @@ class HAMetricManager { ylt::metric::histogram_t oplog_etcd_write_latency_us_; ylt::metric::histogram_t oplog_apply_latency_us_; + ylt::metric::counter_t batch_record_durable_batches_total_; + ylt::metric::counter_t batch_record_durable_entries_total_; + ylt::metric::counter_t batch_record_retry_total_; + ylt::metric::gauge_t batch_record_committed_queue_depth_; + ylt::metric::gauge_t batch_record_callback_queue_depth_; + ylt::metric::gauge_t batch_record_last_batch_id_; + ylt::metric::gauge_t batch_record_durable_sequence_; + ylt::metric::histogram_t batch_record_batch_entries_; + ylt::metric::histogram_t batch_record_batch_bytes_; + ylt::metric::histogram_t batch_record_txn_latency_us_; + ylt::metric::histogram_t batch_record_commit_to_durable_us_; + ylt::metric::histogram_t batch_record_callback_latency_us_; + // State Machine ylt::metric::gauge_t standby_state_; ylt::metric::counter_t state_transitions_total_; diff --git a/mooncake-store/include/hot_standby_service.h b/mooncake-store/include/hot_standby_service.h index 3013882279..f89bd76c0f 100644 --- a/mooncake-store/include/hot_standby_service.h +++ b/mooncake-store/include/hot_standby_service.h @@ -2,6 +2,7 @@ #include #include +#include #include #include #include @@ -14,9 +15,7 @@ #include "metadata_store.h" #include "ha/oplog/oplog_applier.h" -#include "ha/oplog/oplog_manager.h" -#include "ha/oplog/oplog_replicator.h" -#include "ha/oplog/oplog_store_factory.h" +#include "ha/oplog/oplog_types.h" #include "ha/snapshot/snapshot_provider.h" #include "standby_state_machine.h" #include "types.h" @@ -25,7 +24,9 @@ namespace mooncake { // Forward declarations class MasterService; -class ReplicationStream; +class HaKvBackend; +class OpLogBatchStandbyReader; +enum class OpLogBatchStandbyPollDisposition; /** * @brief Configuration for HotStandbyService @@ -47,10 +48,8 @@ struct HotStandbyConfig { // state. bool enable_oplog_following{true}; - // OpLog store configuration - OpLogStoreType oplog_store_type{kDefaultOpLogStoreType}; - std::string oplog_store_root_dir{kDefaultOpLogRootDir}; - int oplog_poll_interval_ms{kDefaultOpLogPollIntervalMs}; + int oplog_poll_interval_ms{1000}; + uint32_t batch_oplog_retry_timeout_sec{180}; }; /** @@ -65,6 +64,7 @@ struct StandbySyncStatus { bool is_connected{false}; StandbyState state{StandbyState::STOPPED}; std::chrono::milliseconds time_in_state{0}; + ErrorCode last_error{ErrorCode::OK}; }; /** @@ -127,6 +127,17 @@ class HotStandbyService { */ ErrorCode Promote(); + /** + * @brief Promote this standby to Primary and export a snapshot atomically. + * + * Holds mutex_ through the entire promotion + export process, ensuring + * the snapshot is captured before any state is released. + * + * @param out Output parameter to receive the snapshot + * @return ErrorCode::OK on success, other codes on failure + */ + ErrorCode PromoteAndExportSnapshot(StandbySnapshot& out); + /** * @brief Get the number of metadata entries in the local store */ @@ -136,7 +147,7 @@ class HotStandbyService { * @brief Get the latest applied sequence ID after promotion * * This should be called after Promote() to get the sequence_id - * that the new Primary's OpLogManager should start from. + * that the new Primary's ordered writer should start from. * * @return Latest applied sequence ID, or 0 if not available */ @@ -145,8 +156,18 @@ class HotStandbyService { // Export a point-in-time snapshot of all replicated metadata. // This is used by MasterServiceSupervisor to initialize the new Primary // after leader election (fast recovery). - bool ExportMetadataSnapshot( - std::vector>& out) const; + bool ExportMetadataSnapshot(std::vector& out) const; + + /** + * Export complete standby snapshot including: + * - Applied OpLog sequence ID + * - All object metadata + * - All registered segments (via OpLogApplier's segment registry) + * + * @param out Output parameter to receive the snapshot + * @return true on success, false on failure (e.g., service not running) + */ + bool ExportStandbySnapshot(StandbySnapshot& out) const; // Inject a snapshot provider (from external snapshot implementation). void SetSnapshotProvider(std::unique_ptr provider); @@ -155,6 +176,13 @@ class HotStandbyService { // from existing standby worker threads; no extra monitor thread is created. void SetSyncStatusCallback(SyncStatusCallback callback); + /** + * @brief Test seam: when set, promotion final catch-up first tries + * batch-record durable prefix/batches from this backend. + */ + void SetCatchUpBatchKvBackendForTesting( + std::shared_ptr backend); + /** * @brief Get current state from state machine */ @@ -167,20 +195,22 @@ class HotStandbyService { return state_machine_; } - /** - * @brief Callback for OpLogReplicator state changes - * @param event The event to process - */ - void OnWatcherEvent(StandbyEvent event); - private: ErrorCode PrepareBootstrapBaselineLocked(uint64_t& baseline_seq_id); ErrorCode LoadSnapshotBaselineLocked(uint64_t& baseline_seq_id); ErrorCode StartOplogFollowingLocked(uint64_t baseline_seq_id); void ActivateSnapshotOnlyStandbyLocked(uint64_t baseline_seq_id); uint64_t GetLocalLastAppliedSequenceIdLocked() const; - void ResolvePromotionGapsLocked(); ErrorCode FinalCatchUpForPromotionLocked(uint64_t current_applied_seq_id); + ErrorCode FinalCatchUpBatchRecordsLocked(HaKvBackend& backend); + void StopReplicationLoop(); + + // Shared body for Promote() and PromoteAndExportSnapshot(): runs the + // promotion sequence machine transitions + gap resolution + final + // catch-up + post catch-up gap check + success transition. Returns + // ErrorCode::OK on success or any fail-closed error code. + ErrorCode PromoteLockedInternal(uint64_t current_applied_seq_id); + void NotifySyncStatus(); /** @@ -193,54 +223,35 @@ class HotStandbyService { */ void VerificationLoop(); - /** - * @brief Apply a single OpLog entry to local metadata store - * @param entry The OpLog entry to apply - * @deprecated Use OpLogApplier instead - */ - void ApplyOpLogEntry(const OpLogEntry& entry); - - /** - * @brief Connect to Primary and establish replication stream - * @return true on success, false on failure - */ - bool ConnectToPrimary(); - - /** - * @brief Disconnect from Primary - */ - void DisconnectFromPrimary(); - - /** - * @brief Process a batch of OpLog entries received from Primary - * @param entries Batch of OpLog entries - */ - void ProcessOpLogBatch(const std::vector& entries); - HotStandbyConfig config_; // Simple in-memory metadata store implementation class StandbyMetadataStore : public MetadataStore { public: - bool PutMetadata(const std::string& key, + bool PutMetadata(const std::string& tenant_id, const std::string& key, const StandbyObjectMetadata& metadata) override; bool Put(const std::string& key, const std::string& payload = std::string()) override; std::optional GetMetadata( + const std::string& tenant_id, const std::string& key) const override; - bool Remove(const std::string& key) override; - bool Exists(const std::string& key) const override; + bool Remove(const std::string& tenant_id, + const std::string& key) override; + bool Exists(const std::string& tenant_id, + const std::string& key) const override; + size_t GetKeyCountForTenant( + const std::string& tenant_id) const override; size_t GetKeyCount() const override; void Clear(); // Snapshot for promotion/restore. - void Snapshot( - std::vector>& out) - const; + void Snapshot(std::vector& out) const; private: mutable std::mutex mutex_; - std::unordered_map store_; + std::unordered_map< + std::string, std::unordered_map> + store_; }; std::unique_ptr metadata_store_; std::unique_ptr snapshot_provider_{ @@ -248,18 +259,19 @@ class HotStandbyService { // OpLog replication components std::unique_ptr oplog_applier_; - std::shared_ptr watcher_oplog_store_; - std::unique_ptr oplog_change_notifier_; - std::unique_ptr oplog_replicator_; + std::shared_ptr batch_standby_kv_backend_; + std::unique_ptr batch_standby_reader_; + + std::shared_ptr catch_up_batch_kv_backend_for_testing_; // Configuration for OpLog sync std::string oplog_endpoints_; std::string cluster_id_; // Replication state - std::shared_ptr replication_stream_; std::atomic applied_seq_id_{0}; std::atomic primary_seq_id_{0}; + std::atomic last_error_{ErrorCode::OK}; // State machine for managing service lifecycle StandbyStateMachine state_machine_; @@ -271,6 +283,9 @@ class HotStandbyService { // Background threads std::thread replication_thread_; std::thread verification_thread_; + std::atomic replication_loop_running_{false}; + std::mutex replication_loop_mutex_; + std::condition_variable replication_loop_cv_; // Synchronization mutable std::mutex mutex_; diff --git a/mooncake-store/include/http_metadata_server.h b/mooncake-store/include/http_metadata_server.h index 411070c0fa..366ce4d779 100644 --- a/mooncake-store/include/http_metadata_server.h +++ b/mooncake-store/include/http_metadata_server.h @@ -5,6 +5,7 @@ #include #include #include +#include #include @@ -35,6 +36,14 @@ class HttpMetadataServer { // Check if the server is running bool is_running() const { return running_; } + // Remove a key from the metadata store (for internal use by MasterService) + // Returns true if key was found and removed, false if key did not exist + bool removeKey(const std::string& key); + + // Remove multiple keys from the metadata store + // Returns the number of keys that were successfully removed + size_t removeKeys(const std::vector& keys); + // Non-copyable HttpMetadataServer(const HttpMetadataServer&) = delete; HttpMetadataServer& operator=(const HttpMetadataServer&) = delete; diff --git a/mooncake-store/include/k8s_lease_helper.h b/mooncake-store/include/k8s_lease_helper.h index 16d42e04c5..4312f693e7 100644 --- a/mooncake-store/include/k8s_lease_helper.h +++ b/mooncake-store/include/k8s_lease_helper.h @@ -36,6 +36,14 @@ class K8sLeaseHelper { static ErrorCode CancelWatch(const std::string& ns, const std::string& lease); + static ErrorCode SetPodLabel(const std::string& ns, const std::string& pod, + const std::string& key, + const std::string& value); + + static ErrorCode ClearPodLabel(const std::string& ns, + const std::string& pod, + const std::string& key); + private: static std::mutex init_mutex_; static bool initialized_; diff --git a/mooncake-store/include/kv_event/key_util.h b/mooncake-store/include/kv_event/key_util.h new file mode 100644 index 0000000000..6c90900ae8 --- /dev/null +++ b/mooncake-store/include/kv_event/key_util.h @@ -0,0 +1,36 @@ +#pragma once + +#include +#include +#include + +namespace mooncake { + +// Parses object keys encoded as decimal or 0x-prefixed hex u64 hashes. +inline std::optional ParseSeqHashFromObjectKey( + const std::string& object_key) { + if (object_key.empty()) { + return std::nullopt; + } + try { + size_t idx = 0; + if (object_key.size() >= 2 && + (object_key[0] == '0' && + (object_key[1] == 'x' || object_key[1] == 'X'))) { + uint64_t value = std::stoull(object_key, &idx, 16); + if (idx == object_key.size()) { + return value; + } + return std::nullopt; + } + uint64_t value = std::stoull(object_key, &idx, 10); + if (idx == object_key.size()) { + return value; + } + } catch (const std::exception&) { + return std::nullopt; + } + return std::nullopt; +} + +} // namespace mooncake diff --git a/mooncake-store/include/kv_event/kv_event_config.h b/mooncake-store/include/kv_event/kv_event_config.h new file mode 100644 index 0000000000..794019f1ba --- /dev/null +++ b/mooncake-store/include/kv_event/kv_event_config.h @@ -0,0 +1,36 @@ +#pragma once + +#include +#include + +namespace mooncake { + +// Publisher transport/identity for the optional KV Events ZMQ socket (RFC +// #1527). Semantic block fields (model_name, block_size, lora_name, +// parent_hash, token_ids, dp_rank) belong on each event payload, not here — see +// https://docs.nvidia.com/dynamo/kv-managers/kv-events-for-custom-engines +struct KvEventConfig { + bool enabled{false}; + // ZMQ PUB bind address, e.g. "tcp://0.0.0.0:5557". + std::string bind_endpoint; + // Identifies the cache owner stream (storage daemon / pool node). + std::string backend_id; + // Emit legacy vLLM/SGLang field names alongside RFC #1527 fields. + bool emit_legacy_compat_fields{true}; + // Emit Mooncake object_key for consumers that match on store key format. + bool emit_object_key{true}; + // Max pending events in the async publisher queue; oldest dropped when + // full. + uint32_t queue_capacity{65536}; + + // Deprecated: not stamped on events. Indexer registration supplies model, + // block_size, dp_rank, and hash namespace for the Mooncake publisher. + std::string model_name; + std::string tenant_id{"default"}; + std::string additional_salt; + std::string lora_name; + uint32_t block_size{0}; + uint32_t dp_rank{0}; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/kv_event/kv_event_publisher.h b/mooncake-store/include/kv_event/kv_event_publisher.h new file mode 100644 index 0000000000..1db093dad5 --- /dev/null +++ b/mooncake-store/include/kv_event/kv_event_publisher.h @@ -0,0 +1,126 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "kv_event/kv_event_config.h" +#include "kv_event/key_util.h" +#include "tenant_id.h" + +namespace mooncake { + +#if defined(MOONCAKE_ENABLE_KV_EVENTS) && MOONCAKE_ENABLE_KV_EVENTS + +// Publishes standardized KV cache events (RFC #1527) over ZMQ for indexers. +class KvEventPublisher { + public: + explicit KvEventPublisher(KvEventConfig config); + ~KvEventPublisher(); + + KvEventPublisher(const KvEventPublisher&) = delete; + KvEventPublisher& operator=(const KvEventPublisher&) = delete; + + bool enabled() const { return config_.enabled; } + + // Non-blocking enqueue into a bounded queue; drops oldest when full. + void PublishStored(const std::string& object_key, const std::string& medium, + const TenantId& tenant_id = TenantId::Default(), + const std::string& group_id = ""); + void PublishRemoved(const std::string& object_key, + const std::string& medium, + const TenantId& tenant_id = TenantId::Default(), + const std::string& group_id = ""); + + struct Stats { + uint64_t published_batches{0}; + uint64_t published_events{0}; + uint64_t dropped_events{0}; + uint64_t skipped_unparsed_keys{0}; + }; + Stats GetStats() const; + + static std::optional ParseSeqHashFromObjectKey( + const std::string& object_key) { + return mooncake::ParseSeqHashFromObjectKey(object_key); + } + + private: + enum class EventKind { kStored, kRemoved }; + + struct PendingEvent { + EventKind kind; + std::string object_key; + std::string medium; + TenantId tenant_id; + std::string group_id; + }; + + void Enqueue(PendingEvent event); + void WorkerLoop(); + void PublishBatch(const std::vector& batch); + void DrainRemainingQueue(std::vector& batch); + + KvEventConfig config_; + void* zmq_context_{nullptr}; + void* zmq_socket_{nullptr}; + + mutable std::mutex queue_mutex_; + std::deque queue_; + std::condition_variable queue_cv_; + std::thread worker_; + std::atomic stop_{false}; + + std::atomic next_event_id_{1}; + std::atomic next_zmq_sequence_{1}; + + std::atomic published_batches_{0}; + std::atomic published_events_{0}; + std::atomic dropped_events_{0}; + std::atomic skipped_unparsed_keys_{0}; +}; + +#else + +// Stub when mooncake_store is built without libzmq (ENABLE_KV_EVENTS=OFF). +class KvEventPublisher { + public: + explicit KvEventPublisher(KvEventConfig config) + : config_(std::move(config)) {} + + bool enabled() const { return false; } + + void PublishStored(const std::string&, const std::string&, + const TenantId& = TenantId::Default(), + const std::string& = "") {} + void PublishRemoved(const std::string&, const std::string&, + const TenantId& = TenantId::Default(), + const std::string& = "") {} + + struct Stats { + uint64_t published_batches{0}; + uint64_t published_events{0}; + uint64_t dropped_events{0}; + uint64_t skipped_unparsed_keys{0}; + }; + Stats GetStats() const { return {}; } + + static std::optional ParseSeqHashFromObjectKey( + const std::string& object_key) { + return mooncake::ParseSeqHashFromObjectKey(object_key); + } + + private: + KvEventConfig config_; +}; + +#endif + +} // namespace mooncake diff --git a/mooncake-store/include/master_admin_service.h b/mooncake-store/include/master_admin_service.h index f10d040074..9d7dd60cb5 100644 --- a/mooncake-store/include/master_admin_service.h +++ b/mooncake-store/include/master_admin_service.h @@ -51,6 +51,8 @@ class MasterAdminServer { std::string BuildMetricsText() const; + std::string BuildTenantQuotaMetricsText() const; + std::string BuildMetricsSummaryText() const; std::shared_ptr GetActiveService() const; @@ -91,6 +93,16 @@ class MasterAdminServer { coro_http::coro_http_response& resp); void HandleBatchQueryKeys(coro_http::coro_http_request& req, coro_http::coro_http_response& resp); + void HandleKvEventsStatus(coro_http::coro_http_request& req, + coro_http::coro_http_response& resp); + void HandleGetTenantQuotas(coro_http::coro_http_request& req, + coro_http::coro_http_response& resp); + void HandleUpsertTenantQuota(coro_http::coro_http_request& req, + coro_http::coro_http_response& resp); + void HandleDeleteTenantQuota(coro_http::coro_http_request& req, + coro_http::coro_http_response& resp); + void HandleRemoveAll(coro_http::coro_http_request& req, + coro_http::coro_http_response& resp); void RegisterHandler(); diff --git a/mooncake-store/include/master_client.h b/mooncake-store/include/master_client.h index 7b54123044..0955bb4e99 100644 --- a/mooncake-store/include/master_client.h +++ b/mooncake-store/include/master_client.h @@ -18,6 +18,7 @@ #include "types.h" #include "rpc_types.h" #include "master_metric_manager.h" +#include "store_rpc_client_io_context.h" #include "task_manager.h" namespace mooncake { @@ -45,6 +46,26 @@ inline void MaybeEnableRdmaSocketConfig(SocketConfigVariant& socket_config) { } } +inline RpcClientPool::PoolConfig MakeMasterRpcClientPoolConfig() { + RpcClientPool::PoolConfig config; + const char* value = std::getenv("MC_RPC_PROTOCOL"); + if (value && std::string_view(value) == "rdma") { + MaybeEnableRdmaSocketConfig(config.client_config.socket_config); + } + + // Default request and connect timeouts remain coro_rpc's built-in 30s. + // A negative request timeout disables the per-request timer. + if (const char* timeout_ms = std::getenv("MC_RPC_TIMEOUT_MS")) { + config.client_config.request_timeout_duration = + std::chrono::milliseconds(std::atoll(timeout_ms)); + } + if (const char* connect_ms = std::getenv("MC_RPC_CONNECT_TIMEOUT_MS")) { + config.client_config.connect_timeout_duration = + std::chrono::milliseconds(std::atoll(connect_ms)); + } + return config; +} + } // namespace detail /** @@ -54,44 +75,14 @@ class MasterClient { public: MasterClient(const UUID& client_id, MasterClientMetric* metrics = nullptr, std::string tenant_id = "default") - : client_id_(client_id), - tenant_id_(NormalizeTenantId(std::move(tenant_id))), - metrics_(metrics) { - coro_io::client_pool::pool_config - pool_conf{}; - - // Disable alive_detect to prevent stale reconnection logs after HA - // failover. Old client_pool objects remain in client_pools_ map and - // would otherwise continue probing failed addresses indefinitely. See - // PR #1642. - pool_conf.host_alive_detect_duration = std::chrono::seconds(0); - const char* value = std::getenv("MC_RPC_PROTOCOL"); - if (value && std::string_view(value) == "rdma") { - detail::MaybeEnableRdmaSocketConfig( - pool_conf.client_config.socket_config); - } - - // Per-request timeout for all client->master RPCs. coro_rpc's - // send_request falls back to this config value when no per-call - // timeout is given, so setting it here covers every RPC method - // uniformly. Default stays at coro_rpc's built-in 30s if unset; - // a negative value disables the timeout (no timer is armed). - if (const char* timeout_ms = std::getenv("MC_RPC_TIMEOUT_MS")) { - pool_conf.client_config.request_timeout_duration = - std::chrono::milliseconds(std::atoll(timeout_ms)); - } - // Optional override for the TCP/RDMA connect timeout (default 30s). - if (const char* connect_ms = std::getenv("MC_RPC_CONNECT_TIMEOUT_MS")) { - pool_conf.client_config.connect_timeout_duration = - std::chrono::milliseconds(std::atoll(connect_ms)); - } - client_pools_ = - std::make_shared>( - pool_conf); - } + : client_accessor_(GetStoreRpcClientIoContextPool(), + detail::MakeMasterRpcClientPoolConfig()), + client_id_(client_id), + tenant_id_(std::move(tenant_id)), + metrics_(metrics) {} ~MasterClient(); - const std::string& tenant_id() const { return tenant_id_; } + const std::string& tenant_id() const { return tenant_id_.value(); } MasterClient(const MasterClient&) = delete; MasterClient& operator=(const MasterClient&) = delete; @@ -219,12 +210,12 @@ class MasterClient { /** * @brief Ends a put operation - * @param key Object key + * @param object_meta Object key and optional checksum * @param replica_type Type of replica (memory or disk) * @return tl::expected indicating success/failure */ [[nodiscard]] tl::expected PutEnd( - const std::string& key, ReplicaType replica_type); + const ObjectMeta& object_meta, ReplicaType replica_type); /** * @brief Ends a put operation for a batch of objects @@ -232,7 +223,7 @@ class MasterClient { * @return ErrorCode indicating success/failure */ [[nodiscard]] std::vector> BatchPutEnd( - const std::vector& keys, + const std::vector& object_metas, ReplicaType replica_type = ReplicaType::ALL); /** @@ -272,10 +263,10 @@ class MasterClient { const ReplicateConfig& config); [[nodiscard]] tl::expected UpsertEnd( - const std::string& key, ReplicaType replica_type); + const ObjectMeta& object_meta, ReplicaType replica_type); [[nodiscard]] std::vector> BatchUpsertEnd( - const std::vector& keys); + const std::vector& object_metas); [[nodiscard]] tl::expected UpsertRevoke( const std::string& key, ReplicaType replica_type); @@ -431,6 +422,12 @@ class MasterClient { [[nodiscard]] tl::expected, ErrorCode> OffloadObjectHeartbeat(const UUID& client_id, bool enable_offloading); + /** + * @brief Poll whether master has requested a full SSD clear. + * @return true if client should clear all SSD files + */ + [[nodiscard]] tl::expected PollRemoveAll(); + [[nodiscard]] tl::expected ReportSsdCapacity( const UUID& client_id, int64_t ssd_total_capacity_bytes); @@ -669,44 +666,16 @@ class MasterClient { [[nodiscard]] std::vector> invoke_batch_rpc(size_t input_size, Args&&... args); - /** - * @brief Accessor for the coro_rpc_client pool. Since coro_rpc_client pool - * cannot reconnect to a different address, a new coro_rpc_client pool is - * created if the address is different from the current one. - */ - class RpcClientAccessor { - public: - void SetClientPool( - std::shared_ptr> - client_pool) { - std::lock_guard lock(client_mutex_); - client_pool_ = client_pool; - } - - std::shared_ptr> - GetClientPool() { - std::shared_lock lock(client_mutex_); - return client_pool_; - } - - private: - mutable std::shared_mutex client_mutex_; - std::shared_ptr> - client_pool_; - }; - RpcClientAccessor client_accessor_; + RpcClientPool client_accessor_; // The client identification. const UUID client_id_; // Tenant identity for this client instance. - const std::string tenant_id_; + const TenantId tenant_id_; // Metrics for tracking RPC operations MasterClientMetric* metrics_; - std::shared_ptr> - client_pools_; - // Mutex to insure the Connect function is atomic. mutable Mutex connect_mutex_; // The address which is passed to the coro_rpc_client diff --git a/mooncake-store/include/master_config.h b/mooncake-store/include/master_config.h index 433c8872a2..2940d0e59d 100644 --- a/mooncake-store/include/master_config.h +++ b/mooncake-store/include/master_config.h @@ -11,6 +11,9 @@ namespace mooncake { +// Forwarded to the HA serve phase via MasterServiceSupervisorConfig. +class HttpMetadataServer; + inline std::string ResolveConfiguredHABackendConnstring( std::string_view ha_backend_type, std::string_view ha_backend_connstring, std::string_view etcd_endpoints) { @@ -52,6 +55,12 @@ struct MasterConfig { std::string ha_backend_connstring; std::string etcd_endpoints; + // OpLog store configuration + bool enable_oplog = false; + int oplog_poll_interval_ms = 1000; + uint32_t oplog_batch_max_entries = 1024; + uint32_t batch_oplog_retry_timeout_sec = 180; + std::string cluster_id; std::string root_fs_dir; int64_t global_file_segment_size; @@ -62,6 +71,15 @@ struct MasterConfig { bool enable_http_metadata_server; uint32_t http_metadata_server_port; std::string http_metadata_server_host; + // Enable cleanup of HTTP metadata (mooncake/ram/*, mooncake/rpc_meta/*) + // when client heartbeat times out. Works in two modes: (1) co-located + // (enable_http_metadata_server=true) via in-process removal, or + // (2) separately-deployed metadata server via async HTTP DELETE. + bool enable_metadata_cleanup_on_timeout; + + // Pod identity for K8s label-based routing + std::string pod_name; + std::string pod_namespace; uint64_t put_start_discard_timeout_sec; uint64_t put_start_release_timeout_sec; @@ -69,9 +87,9 @@ struct MasterConfig { // Storage backend eviction configuration bool enable_disk_eviction; uint64_t quota_bytes; - bool enable_tenant_quota = false; - uint64_t default_tenant_quota_bytes = 0; - uint64_t tenant_quota_pool_capacity_bytes = 0; + bool enable_multi_tenants = false; + std::string tenant_quota_connector_type = "file"; + std::string tenant_quota_connector_uri; bool enable_snapshot_restore; bool enable_snapshot; @@ -106,6 +124,8 @@ struct MasterConfig { // Offload-on-evict: defer LOCAL_DISK offload to eviction time bool offload_on_evict = false; bool offload_force_evict = false; + size_t offloading_queue_limit = 50000; + double offload_cap_ratio = 0.5; // Promotion-on-hit: when Get observes a LOCAL_DISK-only key, queue an // async copy back to MEMORY so the next Get is fast. @@ -118,6 +138,20 @@ struct MasterConfig { // liveness window. Default 1 is conservative; small-object or RDMA- // rich clusters may safely raise it. uint32_t promotion_max_per_heartbeat = 1; + + // KV Events publisher (RFC #1527) for cache-aware indexers. + bool enable_kv_events = false; + std::string kv_events_bind_endpoint; + std::string kv_events_model_name; + std::string kv_events_backend_id; + std::string kv_events_tenant_id = "default"; + std::string kv_events_additional_salt; + std::string kv_events_lora_name; + uint32_t kv_events_block_size = 0; + uint32_t kv_events_dp_rank = 0; + bool kv_events_emit_legacy_compat = true; + bool kv_events_emit_object_key = true; + uint32_t kv_events_queue_capacity = 65536; }; class MasterServiceSupervisorConfig { @@ -154,18 +188,25 @@ class MasterServiceSupervisorConfig { std::string ha_backend_type = "etcd"; std::string ha_backend_connstring; std::string etcd_endpoints = "0.0.0.0:2379"; + // OpLog store configuration + bool enable_oplog = false; + int oplog_poll_interval_ms = 1000; + uint32_t oplog_batch_max_entries = 1024; + uint32_t batch_oplog_retry_timeout_sec = 180; std::string local_hostname = "0.0.0.0:50051"; std::string cluster_id = DEFAULT_CLUSTER_ID; std::string root_fs_dir = DEFAULT_ROOT_FS_DIR; int64_t global_file_segment_size = DEFAULT_GLOBAL_FILE_SEGMENT_SIZE; BufferAllocatorType memory_allocator = BufferAllocatorType::OFFSET; + AllocationStrategyType allocation_strategy_type = + AllocationStrategyType::RANDOM; uint64_t put_start_discard_timeout_sec = DEFAULT_PUT_START_DISCARD_TIMEOUT; uint64_t put_start_release_timeout_sec = DEFAULT_PUT_START_RELEASE_TIMEOUT; bool enable_disk_eviction = true; uint64_t quota_bytes = 0; - bool enable_tenant_quota = false; - uint64_t default_tenant_quota_bytes = 0; - uint64_t tenant_quota_pool_capacity_bytes = 0; + bool enable_multi_tenants = false; + std::string tenant_quota_connector_type = "file"; + std::string tenant_quota_connector_uri; uint32_t max_total_finished_tasks = DEFAULT_MAX_TOTAL_FINISHED_TASKS; uint32_t max_total_pending_tasks = DEFAULT_MAX_TOTAL_PENDING_TASKS; uint32_t max_total_processing_tasks = DEFAULT_MAX_TOTAL_PROCESSING_TASKS; @@ -191,10 +232,36 @@ class MasterServiceSupervisorConfig { bool enable_cxl = false; bool offload_on_evict = false; bool offload_force_evict = false; + size_t offloading_queue_limit = 50000; + double offload_cap_ratio = 0.5; bool promotion_on_hit = false; uint32_t promotion_admission_threshold = 2; uint32_t promotion_queue_limit = 50000; uint32_t promotion_max_per_heartbeat = 1; + bool enable_kv_events = false; + std::string kv_events_bind_endpoint; + std::string kv_events_model_name; + std::string kv_events_backend_id; + std::string kv_events_tenant_id = "default"; + std::string kv_events_additional_salt; + std::string kv_events_lora_name; + uint32_t kv_events_block_size = 0; + uint32_t kv_events_dp_rank = 0; + bool kv_events_emit_legacy_compat = true; + bool kv_events_emit_object_key = true; + uint32_t kv_events_queue_capacity = 65536; + + // Pod identity for K8s label-based routing + std::string pod_name; + std::string pod_namespace; + + // Metadata cleanup on client timeout. Resolved in main() (not from + // MasterConfig) and forwarded to the serving primary's + // WrappedMasterService. Co-located: in-process server pointer; separate: + // derived http(s) URL. + HttpMetadataServer* http_metadata_server = nullptr; + std::string http_metadata_remote_url; + MasterServiceSupervisorConfig() = default; // From MasterConfig @@ -219,10 +286,24 @@ class MasterServiceSupervisorConfig { enable_offload = config.enable_offload; offload_on_evict = config.offload_on_evict; offload_force_evict = config.offload_force_evict; + offloading_queue_limit = config.offloading_queue_limit; + offload_cap_ratio = config.offload_cap_ratio; promotion_on_hit = config.promotion_on_hit; promotion_admission_threshold = config.promotion_admission_threshold; promotion_queue_limit = config.promotion_queue_limit; promotion_max_per_heartbeat = config.promotion_max_per_heartbeat; + enable_kv_events = config.enable_kv_events; + kv_events_bind_endpoint = config.kv_events_bind_endpoint; + kv_events_model_name = config.kv_events_model_name; + kv_events_backend_id = config.kv_events_backend_id; + kv_events_tenant_id = config.kv_events_tenant_id; + kv_events_additional_salt = config.kv_events_additional_salt; + kv_events_lora_name = config.kv_events_lora_name; + kv_events_block_size = config.kv_events_block_size; + kv_events_dp_rank = config.kv_events_dp_rank; + kv_events_emit_legacy_compat = config.kv_events_emit_legacy_compat; + kv_events_emit_object_key = config.kv_events_emit_object_key; + kv_events_queue_capacity = config.kv_events_queue_capacity; rpc_port = static_cast(config.rpc_port); rpc_thread_num = static_cast(config.rpc_thread_num); @@ -235,6 +316,10 @@ class MasterServiceSupervisorConfig { etcd_endpoints = config.etcd_endpoints; ha_backend_connstring = ResolveConfiguredHABackendConnstring( ha_backend_type, config.ha_backend_connstring, etcd_endpoints); + enable_oplog = config.enable_oplog; + oplog_poll_interval_ms = config.oplog_poll_interval_ms; + oplog_batch_max_entries = config.oplog_batch_max_entries; + batch_oplog_retry_timeout_sec = config.batch_oplog_retry_timeout_sec; local_hostname = rpc_address + ":" + std::to_string(rpc_port); cluster_id = config.cluster_id; root_fs_dir = config.root_fs_dir; @@ -247,14 +332,35 @@ class MasterServiceSupervisorConfig { memory_allocator = BufferAllocatorType::OFFSET; } + // Convert string allocation_strategy to AllocationStrategyType enum + if (config.allocation_strategy == "free_ratio_first") { + allocation_strategy_type = AllocationStrategyType::FREE_RATIO_FIRST; + } else if (config.allocation_strategy == "cxl") { + allocation_strategy_type = AllocationStrategyType::CXL; + } else if (config.allocation_strategy == "random") { + allocation_strategy_type = AllocationStrategyType::RANDOM; + } else if (config.allocation_strategy == "ssd_free_ratio_first") { + allocation_strategy_type = + AllocationStrategyType::SSD_FREE_RATIO_FIRST; + } else if (config.allocation_strategy == "local_first") { + allocation_strategy_type = AllocationStrategyType::LOCAL_FIRST; + } else { + LOG(WARNING) << "Unrecognized allocation_strategy value: '" + << config.allocation_strategy + << "'. Defaulting to 'random'. " + << "Valid options are: random, free_ratio_first, cxl, " + "ssd_free_ratio_first, local_first " + "(case-sensitive)"; + allocation_strategy_type = AllocationStrategyType::RANDOM; + } + put_start_discard_timeout_sec = config.put_start_discard_timeout_sec; put_start_release_timeout_sec = config.put_start_release_timeout_sec; enable_disk_eviction = config.enable_disk_eviction; quota_bytes = config.quota_bytes; - enable_tenant_quota = config.enable_tenant_quota; - default_tenant_quota_bytes = config.default_tenant_quota_bytes; - tenant_quota_pool_capacity_bytes = - config.tenant_quota_pool_capacity_bytes; + enable_multi_tenants = config.enable_multi_tenants; + tenant_quota_connector_type = config.tenant_quota_connector_type; + tenant_quota_connector_uri = config.tenant_quota_connector_uri; enable_snapshot_restore = config.enable_snapshot_restore; enable_snapshot = config.enable_snapshot; @@ -276,6 +382,9 @@ class MasterServiceSupervisorConfig { cxl_path = config.cxl_path; cxl_size = config.cxl_size; enable_cxl = config.enable_cxl; + + pod_name = config.pod_name; + pod_namespace = config.pod_namespace; validate(); } @@ -365,12 +474,30 @@ class WrappedMasterServiceConfig { bool enable_offload = false; bool offload_on_evict = false; bool offload_force_evict = false; + size_t offloading_queue_limit = 50000; + double offload_cap_ratio = 0.5; bool promotion_on_hit = false; uint32_t promotion_admission_threshold = 2; uint32_t promotion_queue_limit = 50000; uint32_t promotion_max_per_heartbeat = 1; + bool enable_kv_events = false; + std::string kv_events_bind_endpoint; + std::string kv_events_model_name; + std::string kv_events_backend_id; + std::string kv_events_tenant_id = "default"; + std::string kv_events_additional_salt; + std::string kv_events_lora_name; + uint32_t kv_events_block_size = 0; + uint32_t kv_events_dp_rank = 0; + bool kv_events_emit_legacy_compat = true; + bool kv_events_emit_object_key = true; + uint32_t kv_events_queue_capacity = 65536; std::string ha_backend_type = "etcd"; std::string ha_backend_connstring; + // OpLog store configuration + bool enable_oplog = false; + int oplog_poll_interval_ms = 1000; + uint32_t oplog_batch_max_entries = 1024; std::string cluster_id = DEFAULT_CLUSTER_ID; std::string root_fs_dir = DEFAULT_ROOT_FS_DIR; int64_t global_file_segment_size = DEFAULT_GLOBAL_FILE_SEGMENT_SIZE; @@ -381,9 +508,9 @@ class WrappedMasterServiceConfig { uint64_t put_start_release_timeout_sec = DEFAULT_PUT_START_RELEASE_TIMEOUT; bool enable_disk_eviction = true; uint64_t quota_bytes = 0; - bool enable_tenant_quota = false; - uint64_t default_tenant_quota_bytes = 0; - uint64_t tenant_quota_pool_capacity_bytes = 0; + bool enable_multi_tenants = false; + std::string tenant_quota_connector_type = "file"; + std::string tenant_quota_connector_uri; bool enable_snapshot_restore = false; bool enable_snapshot = false; @@ -436,23 +563,39 @@ class WrappedMasterServiceConfig { enable_offload = config.enable_offload; offload_on_evict = config.offload_on_evict; offload_force_evict = config.offload_force_evict; + offloading_queue_limit = config.offloading_queue_limit; + offload_cap_ratio = config.offload_cap_ratio; promotion_on_hit = config.promotion_on_hit; promotion_admission_threshold = config.promotion_admission_threshold; promotion_queue_limit = config.promotion_queue_limit; promotion_max_per_heartbeat = config.promotion_max_per_heartbeat; + enable_kv_events = config.enable_kv_events; + kv_events_bind_endpoint = config.kv_events_bind_endpoint; + kv_events_model_name = config.kv_events_model_name; + kv_events_backend_id = config.kv_events_backend_id; + kv_events_tenant_id = config.kv_events_tenant_id; + kv_events_additional_salt = config.kv_events_additional_salt; + kv_events_lora_name = config.kv_events_lora_name; + kv_events_block_size = config.kv_events_block_size; + kv_events_dp_rank = config.kv_events_dp_rank; + kv_events_emit_legacy_compat = config.kv_events_emit_legacy_compat; + kv_events_emit_object_key = config.kv_events_emit_object_key; + kv_events_queue_capacity = config.kv_events_queue_capacity; ha_backend_type = config.ha_backend_type; ha_backend_connstring = ResolveConfiguredHABackendConnstring( ha_backend_type, config.ha_backend_connstring, config.etcd_endpoints); + enable_oplog = config.enable_oplog; + oplog_poll_interval_ms = config.oplog_poll_interval_ms; + oplog_batch_max_entries = config.oplog_batch_max_entries; cluster_id = config.cluster_id; root_fs_dir = config.root_fs_dir; global_file_segment_size = config.global_file_segment_size; enable_disk_eviction = config.enable_disk_eviction; quota_bytes = config.quota_bytes; - enable_tenant_quota = config.enable_tenant_quota; - default_tenant_quota_bytes = config.default_tenant_quota_bytes; - tenant_quota_pool_capacity_bytes = - config.tenant_quota_pool_capacity_bytes; + enable_multi_tenants = config.enable_multi_tenants; + tenant_quota_connector_type = config.tenant_quota_connector_type; + tenant_quota_connector_uri = config.tenant_quota_connector_uri; // Convert string memory_allocator to BufferAllocatorType enum if (config.memory_allocator == "cachelib") { @@ -468,11 +611,17 @@ class WrappedMasterServiceConfig { allocation_strategy_type = AllocationStrategyType::CXL; } else if (config.allocation_strategy == "random") { allocation_strategy_type = AllocationStrategyType::RANDOM; + } else if (config.allocation_strategy == "ssd_free_ratio_first") { + allocation_strategy_type = + AllocationStrategyType::SSD_FREE_RATIO_FIRST; + } else if (config.allocation_strategy == "local_first") { + allocation_strategy_type = AllocationStrategyType::LOCAL_FIRST; } else { LOG(WARNING) << "Unrecognized allocation_strategy value: '" << config.allocation_strategy << "'. Defaulting to 'random'. " - << "Valid options are: random, free_ratio_first, cxl " + << "Valid options are: random, free_ratio_first, cxl, " + "ssd_free_ratio_first, local_first " "(case-sensitive)"; allocation_strategy_type = AllocationStrategyType::RANDOM; } @@ -521,29 +670,50 @@ class WrappedMasterServiceConfig { config.nof_eviction_high_watermark_ratio; view_version = view_version_param; client_live_ttl_sec = config.client_live_ttl_sec; + nof_heartbeat_interval_sec = config.nof_heartbeat_interval_sec; + nof_heartbeat_probe_timeout_ms = config.nof_heartbeat_probe_timeout_ms; + nof_heartbeat_failures_threshold = + config.nof_heartbeat_failures_threshold; enable_ha = true; // This is used in HA mode, so enable_ha should be true enable_offload = config.enable_offload; offload_on_evict = config.offload_on_evict; offload_force_evict = config.offload_force_evict; + offloading_queue_limit = config.offloading_queue_limit; + offload_cap_ratio = config.offload_cap_ratio; promotion_on_hit = config.promotion_on_hit; promotion_admission_threshold = config.promotion_admission_threshold; promotion_queue_limit = config.promotion_queue_limit; promotion_max_per_heartbeat = config.promotion_max_per_heartbeat; + enable_kv_events = config.enable_kv_events; + kv_events_bind_endpoint = config.kv_events_bind_endpoint; + kv_events_model_name = config.kv_events_model_name; + kv_events_backend_id = config.kv_events_backend_id; + kv_events_tenant_id = config.kv_events_tenant_id; + kv_events_additional_salt = config.kv_events_additional_salt; + kv_events_lora_name = config.kv_events_lora_name; + kv_events_block_size = config.kv_events_block_size; + kv_events_dp_rank = config.kv_events_dp_rank; + kv_events_emit_legacy_compat = config.kv_events_emit_legacy_compat; + kv_events_emit_object_key = config.kv_events_emit_object_key; + kv_events_queue_capacity = config.kv_events_queue_capacity; ha_backend_type = config.ha_backend_type; ha_backend_connstring = ResolveConfiguredHABackendConnstring( ha_backend_type, config.ha_backend_connstring, config.etcd_endpoints); + enable_oplog = config.enable_oplog; + oplog_poll_interval_ms = config.oplog_poll_interval_ms; + oplog_batch_max_entries = config.oplog_batch_max_entries; cluster_id = config.cluster_id; root_fs_dir = config.root_fs_dir; global_file_segment_size = config.global_file_segment_size; memory_allocator = config.memory_allocator; + allocation_strategy_type = config.allocation_strategy_type; enable_disk_eviction = config.enable_disk_eviction; quota_bytes = config.quota_bytes; - enable_tenant_quota = config.enable_tenant_quota; - default_tenant_quota_bytes = config.default_tenant_quota_bytes; - tenant_quota_pool_capacity_bytes = - config.tenant_quota_pool_capacity_bytes; + enable_multi_tenants = config.enable_multi_tenants; + tenant_quota_connector_type = config.tenant_quota_connector_type; + tenant_quota_connector_uri = config.tenant_quota_connector_uri; put_start_discard_timeout_sec = config.put_start_discard_timeout_sec; put_start_release_timeout_sec = config.put_start_release_timeout_sec; @@ -597,6 +767,10 @@ class MasterServiceConfigBuilder { bool enable_offload_ = false; std::string ha_backend_type_ = "etcd"; std::string ha_backend_connstring_; + // OpLog store configuration + bool enable_oplog_ = false; + int oplog_poll_interval_ms_ = 1000; + uint32_t oplog_batch_max_entries_ = 1024; std::string cluster_id_ = DEFAULT_CLUSTER_ID; std::string root_fs_dir_ = DEFAULT_ROOT_FS_DIR; int64_t global_file_segment_size_ = DEFAULT_GLOBAL_FILE_SEGMENT_SIZE; @@ -605,9 +779,9 @@ class MasterServiceConfigBuilder { AllocationStrategyType::RANDOM; bool enable_disk_eviction_ = true; uint64_t quota_bytes_ = 0; - bool enable_tenant_quota_ = false; - uint64_t default_tenant_quota_bytes_ = 0; - uint64_t tenant_quota_pool_capacity_bytes_ = 0; + bool enable_multi_tenants_ = false; + std::string tenant_quota_connector_type_ = "file"; + std::string tenant_quota_connector_uri_; uint64_t put_start_discard_timeout_sec_ = DEFAULT_PUT_START_DISCARD_TIMEOUT; uint64_t put_start_release_timeout_sec_ = DEFAULT_PUT_START_RELEASE_TIMEOUT; bool enable_snapshot_restore_ = false; @@ -721,6 +895,21 @@ class MasterServiceConfigBuilder { return *this; } + MasterServiceConfigBuilder& set_enable_oplog(bool enable) { + enable_oplog_ = enable; + return *this; + } + + MasterServiceConfigBuilder& set_oplog_poll_interval_ms(int interval_ms) { + oplog_poll_interval_ms_ = interval_ms; + return *this; + } + + MasterServiceConfigBuilder& set_oplog_batch_max_entries(uint32_t entries) { + oplog_batch_max_entries_ = entries; + return *this; + } + MasterServiceConfigBuilder& set_cluster_id(const std::string& id) { cluster_id_ = id; return *this; @@ -749,19 +938,20 @@ class MasterServiceConfigBuilder { return *this; } - MasterServiceConfigBuilder& set_enable_tenant_quota(bool enable) { - enable_tenant_quota_ = enable; + MasterServiceConfigBuilder& set_enable_multi_tenants(bool enable) { + enable_multi_tenants_ = enable; return *this; } - MasterServiceConfigBuilder& set_default_tenant_quota_bytes(uint64_t bytes) { - default_tenant_quota_bytes_ = bytes; + MasterServiceConfigBuilder& set_tenant_quota_connector_type( + const std::string& type) { + tenant_quota_connector_type_ = type; return *this; } - MasterServiceConfigBuilder& set_tenant_quota_pool_capacity_bytes( - uint64_t bytes) { - tenant_quota_pool_capacity_bytes_ = bytes; + MasterServiceConfigBuilder& set_tenant_quota_connector_uri( + const std::string& uri) { + tenant_quota_connector_uri_ = uri; return *this; } @@ -934,12 +1124,30 @@ class MasterServiceConfig { bool enable_offload = false; bool offload_on_evict = false; bool offload_force_evict = false; + size_t offloading_queue_limit = 50000; + double offload_cap_ratio = 0.5; bool promotion_on_hit = false; uint32_t promotion_admission_threshold = 2; uint32_t promotion_queue_limit = 50000; uint32_t promotion_max_per_heartbeat = 1; + bool enable_kv_events = false; + std::string kv_events_bind_endpoint; + std::string kv_events_model_name; + std::string kv_events_backend_id; + std::string kv_events_tenant_id = "default"; + std::string kv_events_additional_salt; + std::string kv_events_lora_name; + uint32_t kv_events_block_size = 0; + uint32_t kv_events_dp_rank = 0; + bool kv_events_emit_legacy_compat = true; + bool kv_events_emit_object_key = true; + uint32_t kv_events_queue_capacity = 65536; std::string ha_backend_type = "etcd"; std::string ha_backend_connstring; + // OpLog store configuration + bool enable_oplog = false; + int oplog_poll_interval_ms = 1000; + uint32_t oplog_batch_max_entries = 1024; std::string cluster_id = DEFAULT_CLUSTER_ID; std::string root_fs_dir = DEFAULT_ROOT_FS_DIR; int64_t global_file_segment_size = DEFAULT_GLOBAL_FILE_SEGMENT_SIZE; @@ -950,9 +1158,9 @@ class MasterServiceConfig { uint64_t put_start_release_timeout_sec = DEFAULT_PUT_START_RELEASE_TIMEOUT; bool enable_disk_eviction = true; uint64_t quota_bytes = 0; - bool enable_tenant_quota = false; - uint64_t default_tenant_quota_bytes = 0; - uint64_t tenant_quota_pool_capacity_bytes = 0; + bool enable_multi_tenants = false; + std::string tenant_quota_connector_type = "file"; + std::string tenant_quota_connector_uri; bool enable_snapshot_restore = false; bool enable_snapshot = false; @@ -1001,12 +1209,29 @@ class MasterServiceConfig { enable_offload = config.enable_offload; offload_on_evict = config.offload_on_evict; offload_force_evict = config.offload_force_evict; + offloading_queue_limit = config.offloading_queue_limit; + offload_cap_ratio = config.offload_cap_ratio; promotion_on_hit = config.promotion_on_hit; promotion_admission_threshold = config.promotion_admission_threshold; promotion_queue_limit = config.promotion_queue_limit; promotion_max_per_heartbeat = config.promotion_max_per_heartbeat; + enable_kv_events = config.enable_kv_events; + kv_events_bind_endpoint = config.kv_events_bind_endpoint; + kv_events_model_name = config.kv_events_model_name; + kv_events_backend_id = config.kv_events_backend_id; + kv_events_tenant_id = config.kv_events_tenant_id; + kv_events_additional_salt = config.kv_events_additional_salt; + kv_events_lora_name = config.kv_events_lora_name; + kv_events_block_size = config.kv_events_block_size; + kv_events_dp_rank = config.kv_events_dp_rank; + kv_events_emit_legacy_compat = config.kv_events_emit_legacy_compat; + kv_events_emit_object_key = config.kv_events_emit_object_key; + kv_events_queue_capacity = config.kv_events_queue_capacity; ha_backend_type = config.ha_backend_type; ha_backend_connstring = config.ha_backend_connstring; + enable_oplog = config.enable_oplog; + oplog_poll_interval_ms = config.oplog_poll_interval_ms; + oplog_batch_max_entries = config.oplog_batch_max_entries; cluster_id = config.cluster_id; root_fs_dir = config.root_fs_dir; global_file_segment_size = config.global_file_segment_size; @@ -1015,10 +1240,9 @@ class MasterServiceConfig { allocation_strategy_type = config.allocation_strategy_type; enable_disk_eviction = config.enable_disk_eviction; quota_bytes = config.quota_bytes; - enable_tenant_quota = config.enable_tenant_quota; - default_tenant_quota_bytes = config.default_tenant_quota_bytes; - tenant_quota_pool_capacity_bytes = - config.tenant_quota_pool_capacity_bytes; + enable_multi_tenants = config.enable_multi_tenants; + tenant_quota_connector_type = config.tenant_quota_connector_type; + tenant_quota_connector_uri = config.tenant_quota_connector_uri; put_start_discard_timeout_sec = config.put_start_discard_timeout_sec; put_start_release_timeout_sec = config.put_start_release_timeout_sec; @@ -1073,6 +1297,9 @@ inline MasterServiceConfig MasterServiceConfigBuilder::build() const { config.enable_offload = enable_offload_; config.ha_backend_type = ha_backend_type_; config.ha_backend_connstring = ha_backend_connstring_; + config.enable_oplog = enable_oplog_; + config.oplog_poll_interval_ms = oplog_poll_interval_ms_; + config.oplog_batch_max_entries = oplog_batch_max_entries_; config.cluster_id = cluster_id_; config.root_fs_dir = root_fs_dir_; config.global_file_segment_size = global_file_segment_size_; @@ -1082,9 +1309,9 @@ inline MasterServiceConfig MasterServiceConfigBuilder::build() const { config.put_start_release_timeout_sec = put_start_release_timeout_sec_; config.enable_disk_eviction = enable_disk_eviction_; config.quota_bytes = quota_bytes_; - config.enable_tenant_quota = enable_tenant_quota_; - config.default_tenant_quota_bytes = default_tenant_quota_bytes_; - config.tenant_quota_pool_capacity_bytes = tenant_quota_pool_capacity_bytes_; + config.enable_multi_tenants = enable_multi_tenants_; + config.tenant_quota_connector_type = tenant_quota_connector_type_; + config.tenant_quota_connector_uri = tenant_quota_connector_uri_; config.enable_snapshot_restore = enable_snapshot_restore_; config.enable_snapshot = enable_snapshot_; config.snapshot_backup_dir = snapshot_backup_dir_; diff --git a/mooncake-store/include/master_metric_manager.h b/mooncake-store/include/master_metric_manager.h index e6ce880a85..c98c43b177 100644 --- a/mooncake-store/include/master_metric_manager.h +++ b/mooncake-store/include/master_metric_manager.h @@ -33,8 +33,14 @@ class MasterMetricManager { void inc_mem_cache_hit_nums(int64_t val = 1); void inc_file_cache_hit_nums(int64_t val = 1); + void inc_mem_cache_hit_bytes(int64_t val = 1); + void inc_file_cache_hit_bytes(int64_t val = 1); + int64_t get_mem_cache_hit_bytes(); + int64_t get_file_cache_hit_bytes(); void inc_mem_cache_nums(int64_t val = 1); void inc_file_cache_nums(int64_t val = 1); + int64_t get_mem_cache_nums(); + int64_t get_file_cache_nums(); void dec_mem_cache_nums(int64_t val = 1); void dec_file_cache_nums(int64_t val = 1); void reset_cache_total_nums(); @@ -86,6 +92,11 @@ class MasterMetricManager { void reset_segment_total_mem_capacity(const std::string& segment); int64_t get_segment_allocated_mem_size(const std::string& segment); int64_t get_segment_total_mem_capacity(const std::string& segment); + // Remove all per-segment metric labels for the given segment. + // Called when a segment is unmounted to prevent stale 0-value entries + // from persisting in Prometheus output (e.g. after snapshot restore + // followed by client expiry / reaper cleanup). + void remove_segment_metrics(const std::string& segment); // NoF segment Metrics void inc_allocated_nof_size(int64_t val = 1); @@ -97,6 +108,8 @@ class MasterMetricManager { double get_segment_nof_used_ratio(const std::string& segment); int64_t get_segment_allocated_nof_size(const std::string& segment); int64_t get_segment_total_nof_capacity(const std::string& segment); + // Remove all per-segment NoF metric labels for the given segment. + void remove_nof_segment_metrics(const std::string& segment); // File Storage Metrics void inc_allocated_file_size(int64_t val = 1); @@ -134,6 +147,7 @@ class MasterMetricManager { void inc_put_start_requests(int64_t val = 1); void inc_put_start_failures(int64_t val = 1); void inc_put_start_alloc_failures(int64_t val = 1); + void inc_put_start_partial_allocations(int64_t val = 1); void inc_put_end_requests(int64_t val = 1); void inc_put_end_failures(int64_t val = 1); void inc_put_revoke_requests(int64_t val = 1); @@ -197,6 +211,7 @@ class MasterMetricManager { int64_t get_put_start_requests(); int64_t get_put_start_failures(); int64_t get_put_start_alloc_failures(); + int64_t get_put_start_partial_allocations(); int64_t get_put_end_requests(); int64_t get_put_end_failures(); int64_t get_put_revoke_requests(); @@ -309,6 +324,19 @@ class MasterMetricManager { void inc_promotion_rejected_watermark(int64_t val = 1); void inc_promotion_rejected_cap(int64_t val = 1); + // Tenant quota metrics + void inc_tenant_quota_reject(const std::string& tenant_id, + const std::string& reason, int64_t val = 1); + void inc_tenant_evict_bytes(const std::string& tenant_id, int64_t bytes); + + // Promotion retry candidate metrics + void inc_promotion_candidate_recorded(int64_t val = 1); + void inc_promotion_candidate_admitted(int64_t val = 1); + void inc_promotion_candidate_admission_rejected(int64_t val = 1); + void inc_promotion_candidate_expired_evaluated(int64_t val = 1); + void inc_promotion_candidate_expired_unevaluated(int64_t val = 1); + void inc_promotion_candidate_dropped_limit(int64_t val = 1); + // Promotion-on-hit Metrics Getters int64_t get_promotion_in_flight(); int64_t get_promotion_admitted(); @@ -320,6 +348,12 @@ class MasterMetricManager { int64_t get_promotion_rejected_frequency(); int64_t get_promotion_rejected_watermark(); int64_t get_promotion_rejected_cap(); + int64_t get_promotion_candidate_recorded(); + int64_t get_promotion_candidate_admitted(); + int64_t get_promotion_candidate_admission_rejected(); + int64_t get_promotion_candidate_expired_evaluated(); + int64_t get_promotion_candidate_expired_unevaluated(); + int64_t get_promotion_candidate_dropped_limit(); // CopyStart, CopyEnd, CopyRevoke, MoveStart, MoveEnd, MoveRevoke Metrics void inc_copy_start_requests(int64_t val = 1); @@ -407,6 +441,7 @@ class MasterMetricManager { int64_t put_starts = 0; int64_t put_start_fails = 0; int64_t put_start_alloc_fails = 0; + int64_t put_start_partial_allocs = 0; int64_t put_ends = 0; int64_t put_end_fails = 0; int64_t put_revoke_requests = 0; @@ -543,6 +578,7 @@ class MasterMetricManager { ylt::metric::counter_t put_start_requests_; ylt::metric::counter_t put_start_failures_; ylt::metric::counter_t put_start_alloc_failures_; + ylt::metric::counter_t put_start_partial_allocations_; ylt::metric::counter_t put_end_requests_; ylt::metric::counter_t put_end_failures_; ylt::metric::counter_t put_revoke_requests_; @@ -620,6 +656,8 @@ class MasterMetricManager { // end-to-end request/token-level cache hit ratio. ylt::metric::counter_t mem_cache_hit_nums_; ylt::metric::counter_t file_cache_hit_nums_; + ylt::metric::counter_t mem_cache_hit_bytes_; + ylt::metric::counter_t file_cache_hit_bytes_; ylt::metric::gauge_t mem_cache_nums_; ylt::metric::gauge_t file_cache_nums_; @@ -669,6 +707,16 @@ class MasterMetricManager { ylt::metric::counter_t promotion_rejected_frequency_; ylt::metric::counter_t promotion_rejected_watermark_; ylt::metric::counter_t promotion_rejected_cap_; + // Promotion retry candidate metrics + ylt::metric::counter_t promotion_candidate_recorded_; + ylt::metric::counter_t promotion_candidate_admitted_; + ylt::metric::counter_t promotion_candidate_admission_rejected_; + ylt::metric::counter_t promotion_candidate_expired_evaluated_; + ylt::metric::counter_t promotion_candidate_expired_unevaluated_; + ylt::metric::counter_t promotion_candidate_dropped_limit_; + + ylt::metric::dynamic_counter_2t tenant_quota_reject_total_; + ylt::metric::dynamic_counter_1t tenant_evict_bytes_total_; // Snapshot Metrics ylt::metric::histogram_t snapshot_duration_ms_; diff --git a/mooncake-store/include/master_service.h b/mooncake-store/include/master_service.h index fea1d61ef2..2988f21ba6 100644 --- a/mooncake-store/include/master_service.h +++ b/mooncake-store/include/master_service.h @@ -14,9 +14,11 @@ #include #include #include +#include #include #include #include +#include #include #include #include @@ -27,7 +29,8 @@ #include "master_metric_manager.h" #include "mutex.h" #include "segment.h" -#include "tenant_quota.h" +#include "tenant_quota_sharded.h" +#include "tenant_quota_policy_store.h" #include "types.h" #include "master_config.h" #include "rpc_types.h" @@ -35,17 +38,33 @@ #include "ha/ha_types.h" #include "ha/snapshot/object/snapshot_object_store.h" #include "task_manager.h" +#include "kv_event/kv_event_publisher.h" +#include "ha/oplog/oplog_types.h" +#include "ha/oplog/ordered_oplog_writer.h" +#include "allocator.h" +#include "metadata_store.h" namespace mooncake { + +// Forward declaration for MasterSnapshotManager +class MasterSnapshotManager; +class MasterSnapshotRepository; + namespace ha { class SnapshotCatalogStore; -} - -class EtcdOpLogStore; +class MasterSnapshotCodec; +struct MasterSnapshotPayloads; +class MasterSnapshotCodecTest; // test fixture, needs private state access +} // namespace ha // Forward declarations class AllocationStrategy; class EvictionStrategy; +class HaKvBackend; +class HttpMetadataServer; +class OpLogBatchStorage; +class OrderedOpLogWriter; +struct MetadataStoragePlugin; // Forward declarations for test classes namespace test { @@ -58,26 +77,59 @@ class SnapshotChildProcessTest; // exposing test-only accessors on MasterService itself. class PromotionOnHitTest; class MasterServiceTenantQuotaTest; +// Friended so the BatchEvict correctness tests can invoke the private +// BatchEvict entry point and seed lease timestamps directly, instead of +// relying on segment pressure plus the background eviction thread. +class BatchEvictTest; +class MasterServiceHATest; +// Friended so the processing_keys double-erase reproduction test can +// invalidate a segment allocator via PrepareUnmountSegment WITHOUT the +// ClearInvalidHandles sweep that MasterService::UnmountSegment performs. +class MasterServiceProcessingKeyDoubleEraseTest; } // namespace test +namespace benchmarks { +class BatchEvictBench; +} // namespace benchmarks /* * @brief MasterService is the main class for the master server. * Lock order: To avoid deadlocks, the following lock order should be followed: * 1. client_mutex_ - * 2. metadata_shards_[shard_idx_].mutex - * 3. tenant_quota_shards_[shard_idx_].mutex - * 4. segment_mutex_ + * 2. tenant_quota_policy_mutex_ + * 3. snapshot_mutex_ + * 4. metadata_shards_[shard_idx_].mutex + * 5. tenant_quota_recompute_mutex_ + * 6. ShardedTenantQuotaTable internal mutex or segment_mutex_ + * + * Strict tenant admission and policy mutation paths that need both + * tenant_quota_policy_mutex_ and snapshot_mutex_ must acquire the tenant + * policy mutex first, then snapshot_mutex_. + * tenant_quota_recompute_mutex_ serializes the capacity snapshot and the + * corresponding quota-table update. The segment mutex is released before + * entering ShardedTenantQuotaTable, so these two locks are never nested. */ class MasterService { // Test friend class for snapshot/restore testing friend class test::MasterServiceSnapshotTestBase; friend class test::SnapshotChildProcessTest; friend class test::PromotionOnHitTest; + friend class benchmarks::BatchEvictBench; friend class test::MasterServiceTenantQuotaTest; + friend class test::BatchEvictTest; + // double-erase processing_keys UAF repro (2026-08-03 prod segfault) + friend class test::MasterServiceProcessingKeyDoubleEraseTest; + friend class MasterSnapshotManager; // Allow access to internal state for + // snapshot + friend class ha::MasterSnapshotCodec; // Allow codec to access private + // members + friend class ha::MasterSnapshotCodecTest; // codec round-trip unit test + friend class test::MasterServiceHATest; public: using NoFProbeFn = std::function; + using DurableFinalizeCallback = + std::function; MasterService(); MasterService(const MasterServiceConfig& config); @@ -88,8 +140,28 @@ class MasterService { bool IsNoFSegmentMountedForTesting(const UUID& segment_id); std::optional GetNoFHeartbeatFailureCountForTesting( const UUID& segment_id); - std::optional GetTenantQuotaSnapshotForTesting( - const std::string& tenant_id) const; + bool IsTenantQuotaEnabled() const; + std::vector ListTenantQuotaSnapshots() const; + std::optional GetTenantQuotaSnapshot( + const TenantId& tenant_id) const; + tl::expected UpsertTenantQuotaPolicy( + const TenantId& tenant_id, uint64_t requested_quota_bytes); + tl::expected, ErrorCode> + DeleteTenantQuotaPolicy(const TenantId& tenant_id); + uint64_t GetTenantQuotaAllocatableCapacityBytes(); + + ErrorCode SetBatchOpLogBackendForTesting( + std::shared_ptr backend); + + /** + * @brief Test-only wrapper around BatchEvict / NoFBatchEvict so that + * unit tests can drive a single eviction cycle synchronously + * without standing up the periodic eviction thread. + */ + void RunBatchEvictForTesting(double evict_ratio_target, + double evict_ratio_lowerbound); + void RunNoFBatchEvictForTesting(double evict_ratio_target, + double evict_ratio_lowerbound); /** * @brief Mount a memory segment for buffer allocation. This function is @@ -170,17 +242,17 @@ class MasterService { * @brief Check if an object exists * @return ErrorCode::OK if exists, otherwise return other ErrorCode */ - auto ExistKey(const std::string& key, const std::string& tenant_id) + auto ExistKey(const std::string& key, const TenantId& tenant_id) -> tl::expected; std::vector> BatchExistKey( - const std::vector& keys, const std::string& tenant_id); + const std::vector& keys, const TenantId& tenant_id); /** * @brief Fetch all keys for a single tenant. * @return ErrorCode::OK if exists */ - auto GetAllKeys(const std::string& tenant_id) + auto GetAllKeys(const TenantId& tenant_id) -> tl::expected, ErrorCode>; /** @@ -266,6 +338,9 @@ class MasterService { std::unordered_map, boost::hash>, ErrorCode>; + bool KvEventsEnabled() const; + KvEventPublisher::Stats GetKvEventStats() const; + /** * @brief Batch clear KV cache replicas for specified object keys. * @param object_keys Vector of object key strings to clear. @@ -277,11 +352,20 @@ class MasterService { * keys on success, or an ErrorCode on failure. Only successfully * cleared keys are included in the result. */ + // Existing key-only overload (signature unchanged): kept for legacy + // callers; delegates with "default". auto BatchReplicaClear(const std::vector& object_keys, const UUID& client_id, const std::string& segment_name) -> tl::expected, ErrorCode>; + // New: tenant-aware overload + auto BatchReplicaClear(const std::vector& object_keys, + const UUID& client_id, + const std::string& segment_name, + const std::string& tenant_id) + -> tl::expected, ErrorCode>; + /** * @brief Retrieves replica lists for object keys that match a regex * pattern. @@ -290,7 +374,7 @@ class MasterService { * replica descriptors on success, or an ErrorCode on failure. */ auto GetReplicaListByRegex(const std::string& regex_pattern, - const std::string& tenant_id) + const TenantId& tenant_id) -> tl::expected< std::unordered_map>, ErrorCode>; @@ -301,7 +385,16 @@ class MasterService { * @return ErrorCode::OK on success, ErrorCode::REPLICA_IS_NOT_READY if not * ready */ - auto GetReplicaList(const std::string& key, const std::string& tenant_id) + auto GetReplicaList(const std::string& key, const TenantId& tenant_id) + -> tl::expected; + + /** + * @brief Read-only single-key replica list query for admin use. + * Unlike GetReplicaList, this does not grant leases, trigger + * promotion, or update cache-hit metrics. + */ + auto GetReplicaListForAdmin(const std::string& key, + const TenantId& tenant_id) -> tl::expected; /** @@ -309,7 +402,16 @@ class MasterService { */ std::vector> BatchGetReplicaList(const std::vector& keys, - const std::string& tenant_id); + const TenantId& tenant_id); + + /** + * @brief Read-only batch replica list query for admin use. + * Unlike BatchGetReplicaList, this does not grant leases, trigger + * promotion, or update cache-hit metrics. + */ + std::vector> + BatchGetReplicaListForAdmin(const std::vector& keys, + const TenantId& tenant_id); /** * @brief Start a put operation for an object @@ -320,7 +422,7 @@ class MasterService { * ErrorCode::INVALID_PARAMS if slice size is invalid */ auto PutStart(const UUID& client_id, const std::string& key, - const std::string& tenant_id, const uint64_t slice_length, + const TenantId& tenant_id, const uint64_t slice_length, const ReplicateConfig& config) -> tl::expected, ErrorCode>; @@ -330,16 +432,20 @@ class MasterService { * @return ErrorCode::OK on success, ErrorCode::OBJECT_NOT_FOUND if not * found, ErrorCode::INVALID_WRITE if replica status is invalid */ + auto PutEnd(const UUID& client_id, const ObjectMeta& object_meta, + const TenantId& tenant_id, ReplicaType replica_type) + -> tl::expected; + auto PutEnd(const UUID& client_id, const std::string& key, - const std::string& tenant_id, ReplicaType replica_type) + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected; /** * @brief Adds a replica instance associated with the given client and key. */ auto AddReplica(const UUID& client_id, const std::string& key, - const std::string& tenant_id, Replica& replica) - -> tl::expected; + const TenantId& tenant_id, Replica& replica) + -> tl::expected; /** * @brief Revoke a put operation, replica_type indicates the type of @@ -348,7 +454,7 @@ class MasterService { * found, ErrorCode::INVALID_WRITE if replica status is invalid */ auto PutRevoke(const UUID& client_id, const std::string& key, - const std::string& tenant_id, ReplicaType replica_type) + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected; /** @@ -357,9 +463,8 @@ class MasterService { * found, ErrorCode::INVALID_WRITE if replica status is invalid */ std::vector> BatchPutEnd( - const UUID& client_id, const std::vector& keys, - const std::string& tenant_id, - ReplicaType replica_type = ReplicaType::ALL); + const UUID& client_id, const std::vector& object_metas, + const TenantId& tenant_id, ReplicaType replica_type = ReplicaType::ALL); /** * @brief Revoke a batch of put operations @@ -368,8 +473,7 @@ class MasterService { */ std::vector> BatchPutRevoke( const UUID& client_id, const std::vector& keys, - const std::string& tenant_id, - ReplicaType replica_type = ReplicaType::ALL); + const TenantId& tenant_id, ReplicaType replica_type = ReplicaType::ALL); /** * @brief Start an upsert operation. If the key does not exist, behaves @@ -381,22 +485,26 @@ class MasterService { * progress), OBJECT_REPLICA_BUSY (replicas have non-zero refcnt). */ auto UpsertStart(const UUID& client_id, const std::string& key, - const std::string& tenant_id, const uint64_t slice_length, + const TenantId& tenant_id, const uint64_t slice_length, const ReplicateConfig& config) -> tl::expected, ErrorCode>; /** * @brief Complete an upsert operation. Delegates to PutEnd. */ + auto UpsertEnd(const UUID& client_id, const ObjectMeta& object_meta, + const TenantId& tenant_id, ReplicaType replica_type) + -> tl::expected; + auto UpsertEnd(const UUID& client_id, const std::string& key, - const std::string& tenant_id, ReplicaType replica_type) + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected; /** * @brief Revoke an upsert operation. Delegates to PutRevoke. */ auto UpsertRevoke(const UUID& client_id, const std::string& key, - const std::string& tenant_id, ReplicaType replica_type) + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected; /** @@ -405,7 +513,7 @@ class MasterService { std::vector, ErrorCode>> BatchUpsertStart(const UUID& client_id, const std::vector& keys, - const std::string& tenant_id, + const TenantId& tenant_id, const std::vector& slice_lengths, const ReplicateConfig& config); @@ -413,15 +521,15 @@ class MasterService { * @brief Complete a batch of upsert operations. Delegates to BatchPutEnd. */ std::vector> BatchUpsertEnd( - const UUID& client_id, const std::vector& keys, - const std::string& tenant_id); + const UUID& client_id, const std::vector& object_metas, + const TenantId& tenant_id); /** * @brief Revoke a batch of upsert operations. Delegates to BatchPutRevoke. */ std::vector> BatchUpsertRevoke( const UUID& client_id, const std::vector& keys, - const std::string& tenant_id); + const TenantId& tenant_id); /** * @brief Evict a disk replica for a key (triggered by client-side disk @@ -432,8 +540,7 @@ class MasterService { * @return ErrorCode::OK on success, OBJECT_NOT_FOUND if key missing */ auto EvictDiskReplica(const UUID& client_id, const std::string& key, - const std::string& tenant_id, - ReplicaType replica_type) + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected; /** @@ -445,7 +552,7 @@ class MasterService { */ std::vector> BatchEvictDiskReplica( const UUID& client_id, const std::vector& keys, - const std::string& tenant_id, ReplicaType replica_type); + const TenantId& tenant_id, ReplicaType replica_type); /** * @brief Start a copy operation @@ -462,16 +569,16 @@ class MasterService { */ tl::expected CopyStart( const UUID& client_id, const std::string& key, - const std::string& tenant_id, const std::string& src_segment, + const TenantId& tenant_id, const std::string& src_segment, const std::vector& tgt_segments); tl::expected CopyEnd(const UUID& client_id, const std::string& key, - const std::string& tenant_id); + const TenantId& tenant_id); tl::expected CopyRevoke(const UUID& client_id, const std::string& key, - const std::string& tenant_id); + const TenantId& tenant_id); /** * @brief Start a move operation @@ -488,16 +595,16 @@ class MasterService { */ tl::expected MoveStart( const UUID& client_id, const std::string& key, - const std::string& tenant_id, const std::string& src_segment, + const TenantId& tenant_id, const std::string& src_segment, const std::string& tgt_segment); tl::expected MoveEnd(const UUID& client_id, const std::string& key, - const std::string& tenant_id); + const TenantId& tenant_id); tl::expected MoveRevoke(const UUID& client_id, const std::string& key, - const std::string& tenant_id); + const TenantId& tenant_id); /** * @brief Remove an object and its replicas @@ -506,7 +613,7 @@ class MasterService { * @return ErrorCode::OK on success, ErrorCode::OBJECT_NOT_FOUND if not * found */ - auto Remove(const std::string& key, const std::string& tenant_id, + auto Remove(const std::string& key, const TenantId& tenant_id, bool force = false) -> tl::expected; /** @@ -516,7 +623,7 @@ class MasterService { * @return An expected object containing the number of removed objects on * success, or an ErrorCode on failure. */ - auto RemoveByRegex(const std::string& str, const std::string& tenant_id, + auto RemoveByRegex(const std::string& str, const TenantId& tenant_id, bool force = false) -> tl::expected; /** @@ -532,7 +639,7 @@ class MasterService { * @param force If true, skip lease and replication task checks. * @return return the number of objects removed */ - long RemoveAll(const std::string& tenant_id, bool force = false); + long RemoveAll(const TenantId& tenant_id, bool force = false); /** * @brief Batch remove objects and their replicas @@ -541,7 +648,7 @@ class MasterService { * @return Vector of expected results for each key. */ auto BatchRemove(const std::vector& keys, - const std::string& tenant_id, bool force = false) + const TenantId& tenant_id, bool force = false) -> std::vector>; /** @@ -589,6 +696,14 @@ class MasterService { auto OffloadObjectHeartbeat(const UUID& client_id, bool enable_offloading) -> tl::expected, ErrorCode>; + /** + * @brief Client polls whether master has requested a full SSD clear + * (triggered by RemoveAll). Atomically checks and clears the flag. + * @param client_id The client polling for the remove-all signal + * @return true if client should clear all SSD files, false otherwise + */ + auto PollRemoveAll(const UUID& client_id) -> tl::expected; + auto ReportSsdCapacity(const UUID& client_id, int64_t ssd_total_capacity_bytes) -> tl::expected; @@ -629,7 +744,7 @@ class MasterService { * arbitrary buffer size from a buggy or malicious caller. */ auto PromotionAllocStart(const UUID& client_id, const std::string& key, - const std::string& tenant_id, uint64_t size, + const TenantId& tenant_id, uint64_t size, const std::vector& preferred_segments) -> tl::expected; @@ -639,7 +754,7 @@ class MasterService { * NotifyOffloadSuccess. */ auto NotifyPromotionSuccess(const UUID& client_id, const std::string& key, - const std::string& tenant_id) + const TenantId& tenant_id) -> tl::expected; /** @@ -661,7 +776,7 @@ class MasterService { * holder's promotion_objects entry. */ auto NotifyPromotionFailure(const UUID& client_id, const std::string& key, - const std::string& tenant_id) + const TenantId& tenant_id) -> tl::expected; /** @@ -669,7 +784,7 @@ class MasterService { * @return Copy task ID on success, ErrorCode on failure */ tl::expected CreateCopyTask( - const std::string& key, const std::string& tenant_id, + const std::string& key, const TenantId& tenant_id, const std::vector& targets); /** @@ -678,7 +793,7 @@ class MasterService { * @return Move task ID on success, ErrorCode on failure */ tl::expected CreateMoveTask(const std::string& key, - const std::string& tenant_id, + const TenantId& tenant_id, const std::string& source, const std::string& target); @@ -710,6 +825,15 @@ class MasterService { tl::expected QuerySegmentStatusById( const UUID& segment_id); + /** + * @brief Restore primary state from standby promotion context. + * Called once at promotion time before serving requests. + */ + void RestoreFromStandbySnapshot( + const std::vector& objects, + uint64_t initial_oplog_sequence_id, + const std::vector& segments); + /** * @brief Query the status of a task * @return Task basic info @@ -732,45 +856,39 @@ class MasterService { tl::expected MarkTaskToComplete( const UUID& client_id, const TaskCompleteRequest& request); - private: - void SnapshotThreadFunc(); - - // Persist master state - tl::expected PersistState( - const std::string& snapshot_id); - tl::expected PersistState( - const ha::SnapshotDescriptor& descriptor); - tl::expected - BuildSnapshotDescriptor(const std::string& snapshot_id, - const std::string& manifest_path, - const std::string& object_prefix) const; - tl::expected - ResolveSnapshotSequenceId() const; -#ifdef STORE_USE_ETCD - tl::expected - GetSnapshotBoundaryOpLogStore() const; -#endif - - tl::expected UploadSnapshotPayloadFile( - const std::vector& data, const std::string& path, - const std::string& local_filename, const std::string& snapshot_id); + /** + * @brief Set the HttpMetadataServer pointer for cleanup on client timeout. + * @param server Pointer to HttpMetadataServer. If nullptr, cleanup is + * disabled. + */ + void setHttpMetadataServer(HttpMetadataServer* server); + /** + * @brief Configure cleanup against a separately-deployed HTTP metadata + * server (not co-located in the master process). The master sends HTTP + * DELETE requests to this endpoint when a client times out. Only http:// + * and https:// connection strings are supported; other schemes (etcd / + * redis / P2PHANDSHAKE) are ignored with a warning and leave cleanup + * disabled. + * @param metadata_connstring e.g. "http://host:8080/metadata". + */ + void setHttpMetadataRemoteUrl(const std::string& metadata_connstring); + + private: std::unique_ptr CreateSnapshotCatalogStore(); - void CleanupOldSnapshot(int keep_count, const std::string& snapshot_id); - ha::SnapshotCatalogStore* GetSnapshotCatalogStore(); // Restore master state void RestoreState(); - bool TryRestoreStateFromSnapshot( - const ha::SnapshotDescriptor& snapshot, - const std::chrono::system_clock::time_point& now); void ResetStateAfterFailedRestoreAttempt(); - void WaitForSnapshotChild(pid_t pid, const std::string& snapshot_id, - int log_pipe_fd); - - void HandleChildTimeout(pid_t pid, const std::string& snapshot_id); - void HandleChildExit(pid_t pid, int status, const std::string& snapshot_id); + /** + * @brief Apply decoded snapshot state to running master service + * @param payloads Decoded snapshot payloads + * @param now Current time for cleanup logic + * @return void on success, SerializationError on failure + */ + tl::expected ApplySnapshotState( + const std::chrono::system_clock::time_point& now); // BatchEvict evicts objects in a near-LRU way, i.e., prioritizes to evict // object with smaller lease timeout. It has two passes. The first pass only @@ -783,10 +901,18 @@ class MasterService { void BatchEvict(double evict_ratio_target, double evict_ratio_lowerbound); void NoFBatchEvict(double evict_ratio_target, double evict_ratio_lowerbound); + struct TenantQuotaEvictionResult { + uint64_t freed_bytes{0}; + uint64_t evicted_objects{0}; + }; + TenantQuotaEvictionResult EvictTenantMemoryForQuota( + const TenantId& tenant_id, uint64_t target_bytes); // Helper to get a snapshot of alive clients (under client_mutex_ shared // lock) std::unordered_set> getAliveClientsSnapshot() const; + void UpdateClientHostId(const UUID& client_id, const std::string& host_id); + std::string GetClientHostId(const UUID& client_id) const; // Clear invalid handles in all shards void ClearInvalidHandles(); @@ -802,7 +928,7 @@ class MasterService { // Internal data structures struct ObjectIdentity { - std::string tenant_id; + TenantId tenant_id; std::string user_key; }; @@ -823,7 +949,7 @@ class MasterService { size_t value_length, std::vector&& reps, bool enable_soft_pin, bool enable_hard_pin = false, ObjectDataType data_type_ = ObjectDataType::UNKNOWN, - std::string group_id_ = "", std::string tenant_id_ = "default", + std::string group_id_ = "", TenantId tenant_id_ = TenantId(), std::string user_key_ = {}) : client_id(client_id_), put_start_time(put_start_time_), @@ -854,9 +980,10 @@ class MasterService { // Updated by UpsertStart (Case B) to reset the discard timeout. std::chrono::system_clock::time_point put_start_time; const size_t size; + std::optional object_checksum; const ObjectDataType data_type{ObjectDataType::UNKNOWN}; const std::string group_id; - const std::string tenant_id; + const TenantId tenant_id; const std::string user_key; mutable SpinLock lock; @@ -1120,6 +1247,7 @@ class MasterService { } type; ReplicaID source_id; std::vector replica_ids; + uint64_t reserved_quota_charge_bytes{0}; }; struct OffloadingTask { @@ -1132,6 +1260,36 @@ class MasterService { // so it cannot be evicted. // // alloc_id pins down which staged PROCESSING MEMORY replica + enum class PromotionQueueResult { + kQueued, + kDisabled, + kFrequencyRejected, + kWatermarkRejected, + kQueueCapRejected, + kAlreadyInFlight, + kMemoryReplicaPresent, + kNoLocalDiskSource, + kNotFound, + kPushFailed, + }; + + enum class PromotionCandidateReason { + kWatermark, + kQueueCap, + kPushFailed, + }; + + struct PromotionCandidate { + uint8_t sketch_score{0}; + std::chrono::steady_clock::time_point first_seen; + std::chrono::steady_clock::time_point last_seen; + std::chrono::steady_clock::time_point retry_after; + PromotionCandidateReason last_reason{ + PromotionCandidateReason::kQueueCap}; + ErrorCode last_error{ErrorCode::OK}; + uint32_t retry_count{0}; + }; + // NotifyPromotionSuccess should commit, so a concurrent Put on the // same key cannot be confused with ours. 0 until // PromotionAllocStart records the new replica. @@ -1167,6 +1325,8 @@ class MasterService { replication_tasks; std::unordered_map offloading_tasks; std::unordered_map promotion_tasks; + std::unordered_map + promotion_candidates; std::unordered_map> group_members; // group_id → set of keys @@ -1174,14 +1334,20 @@ class MasterService { bool Empty() const { return metadata.empty() && processing_keys.empty() && replication_tasks.empty() && offloading_tasks.empty() && - promotion_tasks.empty() && group_members.empty(); + promotion_tasks.empty() && promotion_candidates.empty() && + group_members.empty(); } }; // Sharded metadata maps and their mutexes struct MetadataShard { mutable SharedMutex mutex; - std::unordered_map tenants GUARDED_BY(mutex); + std::unordered_map tenants + GUARDED_BY(mutex); + // Count of objects that have at least one completed LOCAL_DISK replica. + // Used to compute eviction_base = metadata.size() - disk_object_count, + // excluding disk-only objects from the eviction denominator. + long disk_object_count GUARDED_BY(mutex) = 0; }; std::array metadata_shards_; @@ -1197,15 +1363,8 @@ class MasterService { ObjectMetadata& metadata); size_t EraseReplicasWithCacheTotalAccounting( ObjectMetadata& metadata, - const std::function& pred_fn); - - static constexpr size_t kNumTenantQuotaShards = 1024; - struct TenantQuotaShard { - mutable std::mutex mutex; - std::unordered_map tenants - GUARDED_BY(mutex); - }; - std::array tenant_quota_shards_; + const std::function& pred_fn, + std::vector* erased_replica_ids = nullptr); std::unordered_map object_group_ids_ GUARDED_BY(group_routing_mutex_); @@ -1219,7 +1378,7 @@ class MasterService { std::unique_lock lock; }; - ObjectOperationLock AcquireObjectOperationLock(const std::string& tenant_id, + ObjectOperationLock AcquireObjectOperationLock(const TenantId& tenant_id, const std::string& key); std::array object_operation_locks_; @@ -1240,6 +1399,34 @@ class MasterService { const MetadataShard& get() const { return shard_; } + // Called after adding a LOCAL_DISK replica. Increments + // disk_object_count if this is the first completed LOCAL_DISK + // replica for the object (i.e., exactly 1 completed disk replica now). + void OnDiskReplicaAdded(const ObjectMetadata& metadata) { + size_t disk_count = metadata.CountReplicas([](const Replica& r) { + return r.is_local_disk_replica() && r.is_completed(); + }); + if (disk_count == 1) shard_.disk_object_count++; + } + + // Called after removing a LOCAL_DISK replica, or when erasing an + // object that had one. Pass had_completed_disk=true if the object + // had at least one completed LOCAL_DISK replica before the removal. + // When the entire object is being erased, call the one-arg overload. + void OnDiskReplicaRemoved(bool had_completed_disk, + const ObjectMetadata& metadata) { + if (!had_completed_disk) return; + bool still_has_disk = metadata.HasReplica([](const Replica& r) { + return r.is_local_disk_replica() && r.is_completed(); + }); + if (!still_has_disk) shard_.disk_object_count--; + } + + // Overload for full object erasure — no metadata needed. + void OnDiskReplicaRemoved(bool had_completed_disk) { + if (had_completed_disk) shard_.disk_object_count--; + } + private: MetadataShard& shard_; SharedMutexLocker lock_; @@ -1263,29 +1450,26 @@ class MasterService { }; static ObjectIdentity MakeObjectIdentity(const std::string& user_key, - const std::string& tenant_id) { - return {NormalizeTenantId(tenant_id), user_key}; - } - - static std::string MakeTenantScopedKey(const std::string& tenant_id, - const std::string& key) { - const auto normalized_tenant = NormalizeTenantId(tenant_id); - std::string scoped_key; - scoped_key.reserve(normalized_tenant.size() + key.size() + 1); - scoped_key.append(normalized_tenant); - scoped_key.push_back('\0'); - scoped_key.append(key); - return scoped_key; + TenantId tenant_id) { + return {std::move(tenant_id), user_key}; } + const TenantId& ResolveRequestTenantId(const TenantId& tenant_id) const; + ObjectIdentity MakeObjectIdentityForRequest( + const std::string& user_key, const TenantId& tenant_id) const; + tl::expected ResolveTenantIdForWrite( + const TenantId& tenant_id) const; + tl::expected ResolveTenantIdForWriteLocked( + const TenantId& tenant_id) const; + bool IsTenantRegistered(const TenantId& tenant_id) const; + bool TenantHasObjects(const TenantId& tenant_id) const; // Helper to get shard index from tenant-scoped object identity. - size_t getShardIndex(const std::string& tenant_id, + size_t getShardIndex(const TenantId& tenant_id, const std::string& user_key) const { - const auto normalized_tenant = NormalizeTenantId(tenant_id); - if (normalized_tenant == "default") { + if (tenant_id.IsDefault()) { return std::hash{}(user_key) % kNumShards; } - size_t seed = std::hash{}(normalized_tenant); + size_t seed = std::hash{}(tenant_id.value()); boost::hash_combine(seed, user_key); return seed % kNumShards; } @@ -1295,23 +1479,22 @@ class MasterService { return std::hash{}(key) % kNumShards; } - size_t getMetadataShardIndex(const std::string& tenant_id, + size_t getMetadataShardIndex(const TenantId& tenant_id, const std::string& key) const; - size_t getTenantQuotaShardIndex(const std::string& tenant_id) const; - std::optional GetGroupRoute(const std::string& tenant_id, + std::optional GetGroupRoute(const TenantId& tenant_id, const std::string& key) const; void RegisterGroupMember(TenantState& tenant_state, - const std::string& tenant_id, - const std::string& key, + const TenantId& tenant_id, const std::string& key, const std::string& group_id); void UnregisterGroupMember(TenantState& tenant_state, - const std::string& tenant_id, + const TenantId& tenant_id, const std::string& key, const std::string& group_id); std::unordered_map::iterator EraseMetadata( TenantState& tenant_state, std::unordered_map::iterator it, - const std::string& tenant_id); + const TenantId& tenant_id); + void ReleaseLocalDiskUsage(const std::vector& replicas); enum class QuotaEraseMode { kFull, kPreserveOld, @@ -1320,21 +1503,56 @@ class MasterService { std::unordered_map::iterator EraseMetadata( TenantState& tenant_state, std::unordered_map::iterator it, - const std::string& tenant_id, QuotaEraseMode quota_mode); + const TenantId& tenant_id, QuotaEraseMode quota_mode); uint64_t CompletedMemoryQuotaCharge(const ObjectMetadata& metadata) const; uint64_t RequestedMemoryQuotaCharge(uint64_t value_length, const ReplicateConfig& config) const; - tl::expected ReserveTenantQuota( - const std::string& tenant_id, uint64_t bytes); - void CommitTenantQuota(const std::string& tenant_id, uint64_t bytes); - void AbortTenantQuota(const std::string& tenant_id, uint64_t bytes); - void ReleaseTenantQuota(const std::string& tenant_id, uint64_t bytes); - void ReleaseTenantQuotaPartial(const std::string& tenant_id, - uint64_t bytes); + bool ShouldProtectZeroChargeMetadataCreate( + uint64_t requested_quota_charge) const; + tl::expected ReserveTenantQuota(const TenantId& tenant_id, + uint64_t bytes); + void CommitTenantQuota(const TenantId& tenant_id, uint64_t bytes); + void AbortTenantQuota(const TenantId& tenant_id, uint64_t bytes); + void ReleaseTenantQuota(const TenantId& tenant_id, uint64_t bytes); + void ReleaseTenantQuotaPartial(const TenantId& tenant_id, uint64_t bytes); + void CommitAdditionalTenantQuota(const TenantId& tenant_id, uint64_t bytes); + void IncrementTenantMetadataObjectCount(const TenantId& tenant_id); + void DecrementTenantMetadataObjectCount(const TenantId& tenant_id); void ReleaseCommittedQuotaCharge(ObjectMetadata& metadata, uint64_t bytes); void RecomputeTenantEffectiveQuotas(); void RebuildTenantQuotaUsageFromMetadata(); - uint64_t GetTenantQuotaCapacityBytes(); + void LoadTenantQuotaPoliciesFromStoreOrThrow(); + void ApplyTenantQuotaPolicies(const TenantQuotaPolicySnapshot& snapshot); + TenantQuotaPolicySnapshot BuildTenantQuotaPolicySnapshot() const; + std::unordered_map::iterator EraseMetadata( + TenantState& tenant_state, + std::unordered_map::iterator it, + const TenantId& tenant_id, QuotaEraseMode quota_mode, + MetadataShardAccessorRW* shard); + void FinalizeRemovedReplicasAfterDurable( + const OpLogEntry& durable_entry, + const std::vector& replica_ids, QuotaEraseMode quota_mode); + void FinalizeMetadataEraseAfterDurable(const OpLogEntry& durable_entry, + QuotaEraseMode quota_mode); + void FinalizeExpiredProcessingReplicasAfterDurable( + const OpLogEntry& durable_entry, + const std::chrono::system_clock::time_point& ttl); + void FinalizeExpiredReplicationTaskAfterDurable( + const OpLogEntry& durable_entry, ReplicaID source_id, + const std::vector& target_ids, + const std::chrono::system_clock::time_point& ttl); + struct StaleHandleCleanupPlan { + std::vector removed_ids; + std::vector remaining; + bool would_invalidate{false}; + }; + StaleHandleCleanupPlan BuildStaleHandleCleanupPlan( + const ObjectMetadata& metadata, + const std::unordered_set>& alive_clients) const; + tl::expected PersistStaleHandleCleanupForHA( + const std::string& why, const TenantId& tenant_id, + const std::string& key, ObjectMetadata& metadata, + const StaleHandleCleanupPlan& plan); void RebuildGroupRoutingIndex(); void GrantLeaseForGroup(const TenantState& tenant_state, const std::string& key, @@ -1343,8 +1561,9 @@ class MasterService { // Helper to clean up stale handles pointing to unmounted segments // or local_disk replicas whose owner client is no longer alive. bool CleanupStaleHandles( - ObjectMetadata& metadata, - const std::unordered_set>& alive_clients); + TenantState& tenant_state, ObjectMetadata& metadata, + const std::unordered_set>& alive_clients, + MetadataShardAccessorRW* shard = nullptr); // Helper: allocate replicas, create ObjectMetadata, insert into shard, // and return descriptor list. Shared by PutStart and UpsertStart. @@ -1352,7 +1571,7 @@ class MasterService { MetadataShardAccessorRW& shard, const UUID& client_id, const std::string& key, uint64_t value_length, const ReplicateConfig& config, const std::string& group_id, - const std::string& tenant_id, + const TenantId& tenant_id, const std::chrono::system_clock::time_point& now) -> tl::expected, ErrorCode>; @@ -1407,14 +1626,32 @@ class MasterService { * map. Acquires its own RW shard accessor; safe to call after * GetReplicaList's RO accessor has been released. */ - void TryPushPromotionQueue(const ObjectIdentity& object_id); + PromotionQueueResult TryPushPromotionQueue(const ObjectIdentity& object_id, + bool record_candidate = true); + void RecordOrUpdateCandidate(TenantState& tenant_state, + const std::string& key, uint8_t sketch_score, + PromotionCandidateReason reason, + ErrorCode last_error); + void EraseCandidate(TenantState& tenant_state, const std::string& key); + void EraseCandidate(const ObjectIdentity& object_id); + void DecrementCandidateCount(); + void BackoffCandidate(const ObjectIdentity& object_id, + PromotionQueueResult result); + void ClearCandidatesForReload(); + std::chrono::milliseconds CandidateBackoff(uint32_t retry_count) const; + bool IsTransientResult(PromotionQueueResult result) const; + size_t RunPromotionCandidateRetry(size_t max_shards_to_scan); + size_t RunPromotionCandidateRetry(); + size_t RunPromotionCandidateRetryForTesting(); + size_t CountCandidatesForTesting(const TenantId& tenant_id); + void ResetCandidateBackoffsForTesting(); // Erase any in-flight PromotionTask for `key`, abort any staged promotion // quota reservation, and decrement the cluster-wide in-flight counter. Safe // no-op if no task exists. void ErasePromotionTaskIfPresent( TenantState& tenant_state, const std::string& key, - const std::string& tenant_id) NO_THREAD_SAFETY_ANALYSIS { + const TenantId& tenant_id) NO_THREAD_SAFETY_ANALYSIS { auto task_it = tenant_state.promotion_tasks.find(key); if (task_it != tenant_state.promotion_tasks.end()) { AbortTenantQuota(tenant_id, @@ -1425,6 +1662,10 @@ class MasterService { MasterMetricManager::instance().inc_promotion_cancelled(); } } + void CancelPromotionTaskForRemovedReplicas( + TenantState& tenant_state, ObjectMetadata& metadata, + const std::vector& removed_replica_ids) + NO_THREAD_SAFETY_ANALYSIS; // Lease related members const uint64_t default_kv_lease_ttl_; // in milliseconds @@ -1447,8 +1688,9 @@ class MasterService { static constexpr uint64_t kEvictionThreadSleepMs = 10; // 10 ms sleep between eviction checks - std::thread snapshot_thread_; - std::atomic snapshot_running_{false}; + // Snapshot manager handles snapshot lifecycle orchestration + std::unique_ptr snapshot_manager_; + // Task cleanup thread related members std::thread task_cleanup_thread_; std::atomic task_cleanup_running_{false}; @@ -1462,10 +1704,9 @@ class MasterService { // Helper class for accessing metadata with automatic locking and cleanup class MetadataAccessorRW { public: - MetadataAccessorRW(MasterService* service, - const ObjectIdentity& object_id) + MetadataAccessorRW(MasterService* service, ObjectIdentity object_id) : service_(service), - object_id_(object_id), + object_id_(std::move(object_id)), shard_idx_(service_->getMetadataShardIndex(object_id_.tenant_id, object_id_.user_key)), shard_guard_(service_, shard_idx_), @@ -1489,17 +1730,23 @@ class MasterService { // violation (client_mutex_ must be acquired before metadata shard). // local_disk replicas are cleaned up by ClearInvalidHandles() in // ClientMonitorFunc. - if (tenant_state_ != nullptr && + if (!(service_->enable_ha_ && service_->enable_oplog_) && + tenant_state_ != nullptr && it_ != tenant_state_->metadata.end()) { // Erase invalid memory replicas (those with unmounted // segments). No client_mutex_ needed since we only check memory // replicas. const uint64_t before_charge = service_->CompletedMemoryQuotaCharge(it_->second); + std::vector removed_replica_ids; service_->EraseReplicasWithCacheTotalAccounting( - it_->second, [](const Replica& replica) { + it_->second, + [](const Replica& replica) { return replica.has_invalid_mem_handle(); - }); + }, + &removed_replica_ids); + service_->CancelPromotionTaskForRemovedReplicas( + *tenant_state_, it_->second, removed_replica_ids); const uint64_t after_charge = service_->CompletedMemoryQuotaCharge(it_->second); if (before_charge > after_charge) { @@ -1508,12 +1755,12 @@ class MasterService { } // If no valid replicas remain, delete the whole object. if (!it_->second.IsValid()) { - const bool had_processing = - processing_it_ != tenant_state_->processing_keys.end(); + // NOTE: Erase() -> EraseMetadata() already removes the key + // from processing_keys (by key), so calling + // EraseFromProcessing() here would re-erase the same node + // via the now-dangling processing_it_ iterator + // (use-after-free, prod segfault 2026-08-03). this->Erase(); - if (tenant_state_ != nullptr && had_processing) { - this->EraseFromProcessing(); - } if (tenant_state_ != nullptr) { service_->ErasePromotionTaskIfPresent( *tenant_state_, object_id_.user_key, @@ -1560,7 +1807,8 @@ class MasterService { // Delete current metadata (for PutRevoke or Remove operations) void Erase() NO_THREAD_SAFETY_ANALYSIS { - service_->EraseMetadata(*tenant_state_, it_, object_id_.tenant_id); + service_->EraseMetadata(*tenant_state_, it_, object_id_.tenant_id, + QuotaEraseMode::kFull, &shard_guard_); it_ = tenant_state_->metadata.end(); MaybeEraseEmptyTenant(); } @@ -1595,6 +1843,10 @@ class MasterService { enable_soft_pin, enable_hard_pin, data_type, group_id, object_id_.tenant_id, object_id_.user_key)); it_ = result.first; + if (result.second) { + service_->IncrementTenantMetadataObjectCount( + object_id_.tenant_id); + } } private: @@ -1629,7 +1881,8 @@ class MasterService { ObjectIdentity object_id_; size_t shard_idx_; MetadataShardAccessorRW shard_guard_; - std::unordered_map::iterator tenant_it_; + std::unordered_map::iterator + tenant_it_; TenantState* tenant_state_; ObjectMetadataIterator it_; ProcessingIterator processing_it_; @@ -1681,9 +1934,9 @@ class MasterService { class MetadataAccessorRO { public: MetadataAccessorRO(const MasterService* service, - const ObjectIdentity& object_id) + ObjectIdentity object_id) : service_(service), - object_id_(object_id), + object_id_(std::move(object_id)), shard_idx_(service_->getMetadataShardIndex(object_id_.tenant_id, object_id_.user_key)), shard_guard_(service_, shard_idx_), @@ -1734,7 +1987,8 @@ class MasterService { const ObjectIdentity object_id_; const size_t shard_idx_; MetadataShardAccessorRO shard_guard_; - std::unordered_map::const_iterator tenant_it_; + std::unordered_map::const_iterator + tenant_it_; const TenantState* tenant_state_; ObjectMetadataConstIterator it_; ProcessingConstIterator processing_it_; @@ -1749,6 +2003,7 @@ class MasterService { mutable std::shared_mutex client_mutex_; std::unordered_set> ok_client_; // client with ok status + std::unordered_map> client_host_id_; void ClientMonitorFunc(); std::thread client_monitor_thread_; std::atomic client_monitor_running_{false}; @@ -1812,6 +2067,19 @@ class MasterService { // promotion task reaper after the task entry is erased. Relaxed memory // order is safe — the value is an advisory soft cap, not a barrier. std::atomic promotion_in_flight_{0}; + // Promotion retry candidate state. + std::atomic promotion_candidate_count_{0}; + std::atomic promotion_retry_cursor_{0}; + static constexpr size_t kPromotionCandidateLimit = 50000; + static constexpr uint32_t kPromotionCandidateMaxRetries = 8; + static constexpr size_t kPromotionRetryBatchSize = 128; + static constexpr size_t kPromotionRetryShardBatch = 64; + static constexpr std::chrono::milliseconds kPromotionCandidateTtl{60000}; + static constexpr std::chrono::milliseconds + kPromotionCandidateInitialBackoff{10}; + static constexpr std::chrono::milliseconds kPromotionCandidateMaxBackoff{ + 1000}; + // Master-side frequency sketch. Constructed only when promotion_on_hit_ is // true. CountMinSketch is mutex-protected internally so we can call into it // from any GetReplicaList caller without additional locking. @@ -1820,6 +2088,8 @@ class MasterService { const std::string ha_backend_type_; const std::string ha_backend_connstring_; + const bool enable_oplog_; + const uint32_t oplog_batch_max_entries_; // cluster id for persistent sub directory const std::string cluster_id_; @@ -1830,9 +2100,39 @@ class MasterService { // storage backend eviction configuration const bool enable_disk_eviction_; const uint64_t quota_bytes_; - const bool enable_tenant_quota_; - const uint64_t default_tenant_quota_bytes_; - const uint64_t tenant_quota_pool_capacity_bytes_; + const bool enable_multi_tenants_; + const std::string tenant_quota_connector_type_; + const std::string tenant_quota_connector_uri_; + std::unique_ptr tenant_quota_policy_store_; + mutable std::mutex tenant_quota_policy_mutex_; + mutable std::mutex tenant_quota_recompute_mutex_; + ShardedTenantQuotaTable<1024> tenant_quota_table_; + + // HTTP metadata server pointer for cleanup on client timeout + // nullptr means cleanup is disabled + HttpMetadataServer* http_metadata_server_{nullptr}; + + // Remote HTTP metadata client, used when the metadata server is deployed + // separately. nullptr = no remote cleanup (co-located prefers the pointer). + std::shared_ptr http_metadata_remote_; + + // Cached HTTP metadata key prefix (initialized once at startup) + std::string http_metadata_prefix_; + + // Async worker for remote cleanup: segments are enqueued from the client + // monitor thread so a slow/unreachable server never blocks heartbeats. + std::thread http_metadata_cleanup_thread_; + std::atomic http_metadata_cleanup_running_{false}; + std::mutex http_metadata_cleanup_mutex_; + std::condition_variable http_metadata_cleanup_cv_; + std::vector http_metadata_cleanup_queue_; + + void HttpMetadataCleanupThreadFunc(); + + // Clean up HTTP metadata (mooncake/ram/*, mooncake/rpc_meta/*) for a + // segment. For the co-located case this is synchronous (no network I/O); + // for the remote case it enqueues to the async cleanup worker. + void cleanupHttpMetadata(const std::string& segment_name); bool use_disk_replica_{false}; @@ -1840,6 +2140,7 @@ class MasterService { SegmentManager segment_manager_; NoFSegmentManager nof_segment_manager_; BufferAllocatorType memory_allocator_type_; + const AllocationStrategyType allocation_strategy_type_; std::shared_ptr allocation_strategy_; bool enable_snapshot_restore_ = false; @@ -1855,11 +2156,9 @@ class MasterService { std::string snapshot_catalog_store_connstring_; std::unique_ptr snapshot_object_store_; std::unique_ptr snapshot_catalog_store_; + std::unique_ptr snapshot_repository_; + std::unique_ptr snapshot_codec_; mutable std::shared_mutex snapshot_mutex_; -#ifdef STORE_USE_ETCD - mutable std::mutex snapshot_boundary_oplog_store_mutex_; - mutable std::unique_ptr snapshot_boundary_oplog_store_; -#endif // Discarded replicas management const std::chrono::seconds put_start_discard_timeout_sec_; @@ -1905,13 +2204,14 @@ class MasterService { std::list discarded_replicas_ GUARDED_BY(discarded_replicas_mutex_); size_t offloading_queue_limit_ = 50000; + double offload_cap_ratio_ = 0.5; // Task manager ClientTaskManager task_manager_; struct ActiveDrainTask { UUID task_id; - std::string tenant_id; + TenantId tenant_id; std::string key; std::string source_segment; std::string target_segment; @@ -1953,7 +2253,7 @@ class MasterService { std::optional SelectDrainTargetForKey( const ObjectMetadata& metadata, const std::string& source_segment, const std::vector& requested_targets); - std::string MakeDrainUnitKey(const std::string& tenant_id, + std::string MakeDrainUnitKey(const TenantId& tenant_id, const std::string& key, const std::string& source_segment) const; @@ -1963,6 +2263,102 @@ class MasterService { std::mutex job_mutex_; std::unordered_map, boost::hash> drain_jobs_ GUARDED_BY(job_mutex_); + + std::unique_ptr kv_event_publisher_; + + static KvEventConfig BuildKvEventConfig(const MasterServiceConfig& config); + static std::string MediumForReplicaType(ReplicaType replica_type); + static std::string MediumForMetadata(const ObjectMetadata& metadata); + void PublishKvStored(const std::string& key, ReplicaType replica_type, + const ObjectMetadata& metadata, + const TenantId& tenant_id); + void PublishKvRemoved(const std::string& key, + const ObjectMetadata& metadata, + const TenantId& tenant_id); + void PublishKvRemoved(const std::string& key, const std::string& medium, + const TenantId& tenant_id, + const std::string& group_id); + void PublishKvRemovedAfterEvict(const std::string& key, + uint64_t freed_bytes, + const std::string& medium, + const ObjectMetadata& metadata, + const TenantId& tenant_id); + + // OpLog publishing + std::shared_ptr batch_oplog_kv_backend_; + std::unique_ptr batch_oplog_storage_; + std::unique_ptr ordered_oplog_writer_; + + // OpLog publishing helpers + std::string SerializeMetadataForOpLog(const ObjectMetadata& metadata) const; + std::string SerializeMetadataForOpLogWithoutMemReplicas( + const ObjectMetadata& metadata) const; + std::string SerializeMetadataForOpLogFromReplicaDescriptors( + const UUID& client_id, uint64_t size, + const std::vector& replicas, + const std::string& group_id = "", + ObjectDataType data_type = ObjectDataType::UNKNOWN) const; + ErrorCode InitializeBatchOpLogWriter(std::shared_ptr backend); + tl::expected AppendOpLogVisibleBeforeDurable( + OpType type, const std::string& tenant_id, const std::string& key, + const std::string& payload); + tl::expected AppendOpLogWithDurableFinalize( + OpType type, const std::string& tenant_id, const std::string& key, + const std::string& payload, DurableFinalizeCallback callback); + tl::expected + ReserveBatchOpLogSlot(); + tl::expected AppendReservedOpLogWithDurableFinalize( + OrderedOpLogWriter::Reservation&& reservation, OpType type, + const std::string& tenant_id, const std::string& key, + const std::string& payload, DurableFinalizeCallback callback); + + // Invalid endpoints from standby that don't exist locally + std::unordered_set invalid_replica_endpoints_; + + // Keep DummyBufferAllocator alive after standby restore. + // Key: transport_endpoint, Value: allocator. + std::unordered_map> + standby_allocator_keepalive_; + std::vector standby_memory_segments_; + std::unordered_map standby_accounted_memory_bytes_; + + ErrorCode ValidateStandbyRemountSegment(const Segment& segment) const; + + bool IsReplicaReadable(const Replica& replica) const; + + /** + * Segment lifecycle persist helper. Tries to durably persist the + * SEGMENT_MOUNT / SEGMENT_UNMOUNT entry up-front; on failure enqueues + * the same OpLogEntry (with its already-allocated sequence_id) for + * background retry so the standby segment registry eventually + * converges. Suitable for paths where the local segment commit has + * already happened (UnmountSegment) and rolling back is impossible. + */ + void PersistSegmentOpForHAOrEnqueue(const char* why, OpType type, + const std::string& key, + const std::string& payload); + void PersistSegmentOpForHAOrEnqueue(const char* why, OpType type, + const TenantId& tenant_id, + const std::string& key, + const std::string& payload); + + /** + * Helper to persist REMOVE OpLog for a key with strong-consistency. + * @return OK on success, error on persist failure (caller must skip erase) + */ + tl::expected PersistRemoveForHA(const char* why, + const std::string& key); + tl::expected PersistRemoveForHA(const char* why, + const TenantId& tenant_id, + const std::string& key); + + /** + * Build replica descriptors after removing replicas matching pred_fn. + * Returns empty if no complete replicas remain. + */ + std::vector BuildRemainingReplicaDescriptors( + const ObjectMetadata& metadata, + const std::function& should_remove) const; }; } // namespace mooncake diff --git a/mooncake-store/include/master_snapshot_manager.h b/mooncake-store/include/master_snapshot_manager.h new file mode 100644 index 0000000000..bd5c04249e --- /dev/null +++ b/mooncake-store/include/master_snapshot_manager.h @@ -0,0 +1,119 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "types.h" +#include "ha/ha_types.h" +#include "ha/snapshot/master_snapshot_codec.h" + +namespace mooncake { + +// Forward declarations +class MasterService; +class SnapshotObjectStore; +class MasterSnapshotRepository; + +namespace ha { +class SnapshotCatalogStore; +} + +namespace test { +class MasterServiceSnapshotTestBase; +class SnapshotChildProcessTest; +} // namespace test + +struct MasterSnapshotManagerOptions { + bool enable_snapshot{false}; + uint64_t snapshot_interval_seconds{0}; + uint64_t snapshot_child_timeout_seconds{0}; + uint32_t snapshot_retention_count{0}; + std::string snapshot_backup_dir; + bool use_snapshot_backup_dir{false}; + std::string snapshot_catalog_store_type; + std::string snapshot_catalog_store_connstring; + std::string ha_backend_type; + std::string ha_backend_connstring; + std::string cluster_id; + bool enable_ha{false}; +}; + +/** + * @brief MasterSnapshotManager handles snapshot lifecycle orchestration for + * MasterService. This includes periodic snapshot scheduling, snapshot ID + * generation, descriptor construction, child process lifecycle management, + * timeout handling, payload upload, catalog publish, retention cleanup, and + * snapshot metrics updates. + * + * This is a behavior-preserving refactor that moves snapshot orchestration + * logic out of MasterService without changing snapshot format, restore + * behavior, storage layout, flags, or locking semantics. + */ +class MasterSnapshotManager { + friend class test::MasterServiceSnapshotTestBase; // Allow test access to + // private methods + friend class test::SnapshotChildProcessTest; // Allow test access to + // private methods + + public: + MasterSnapshotManager(MasterService* master_service, + MasterSnapshotManagerOptions options, + std::shared_mutex& snapshot_mutex, + SnapshotObjectStore* snapshot_object_store, + ha::SnapshotCatalogStore* snapshot_catalog_store); + + ~MasterSnapshotManager(); + + void Start(); + void Stop(); + + private: + void SnapshotThreadFunc(); + void WaitForSnapshotChild(pid_t pid, const std::string& snapshot_id, + int log_pipe_fd); + void HandleChildTimeout(pid_t pid, const std::string& snapshot_id); + void HandleChildExit(pid_t pid, int status, const std::string& snapshot_id); + + tl::expected PersistState( + const std::string& snapshot_id); + tl::expected PersistState( + const ha::SnapshotDescriptor& descriptor); + tl::expected + BuildSnapshotDescriptor(const std::string& snapshot_id, + const std::string& manifest_path, + const std::string& object_prefix) const; + tl::expected + ResolveSnapshotSequenceId() const; + + tl::expected UploadSnapshotPayloadFile( + const std::vector& data, const std::string& path, + const std::string& local_filename, const std::string& snapshot_id); + + void CleanupOldSnapshot(size_t keep_count, const std::string& snapshot_id); + + std::string FormatTimestamp( + const std::chrono::system_clock::time_point& tp); + + MasterService* master_service_; + MasterSnapshotManagerOptions options_; + + std::shared_mutex& snapshot_mutex_; + SnapshotObjectStore* snapshot_object_store_; + ha::SnapshotCatalogStore* snapshot_catalog_store_; + + std::unique_ptr repository_; + + std::thread snapshot_thread_; + std::atomic snapshot_running_{false}; + std::mutex snapshot_thread_mutex_; + std::condition_variable snapshot_thread_cv_; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/master_snapshot_repository.h b/mooncake-store/include/master_snapshot_repository.h new file mode 100644 index 0000000000..fb550eab09 --- /dev/null +++ b/mooncake-store/include/master_snapshot_repository.h @@ -0,0 +1,118 @@ +#pragma once + +#include +#include +#include + +#include + +#include "types.h" +#include "ha/ha_types.h" +#include "ha/snapshot/master_snapshot_codec.h" + +namespace mooncake { + +// Forward declarations +class SnapshotObjectStore; + +namespace ha { +class SnapshotCatalogStore; +} + +/** + * @brief MasterSnapshotRepository handles storage and catalog operations for + * snapshots. This includes uploading payload files to object storage, + * publishing snapshots to the catalog, listing snapshots, deleting snapshots, + * and enforcing retention policies. + * + * This class encapsulates all interactions with SnapshotObjectStore and + * SnapshotCatalogStore, separating storage concerns from snapshot orchestration + * logic in MasterSnapshotManager. + */ +class MasterSnapshotRepository { + public: + MasterSnapshotRepository(SnapshotObjectStore* object_store, + ha::SnapshotCatalogStore* catalog_store, + const std::string& backup_dir, + bool use_backup_dir); + + /** + * @brief Upload a single snapshot payload file to object storage + * @param data Binary data to upload + * @param path Storage path/key + * @param local_filename Filename for logging and local backup + * @param snapshot_id Snapshot ID for logging + * @return Empty on success, SerializationError on failure + */ + tl::expected UploadPayloadFile( + const std::vector& data, const std::string& path, + const std::string& local_filename, const std::string& snapshot_id); + + /** + * @brief Publish snapshot descriptor to catalog store + * @param descriptor Snapshot descriptor to publish + * @return ErrorCode::OK on success, error code on failure + */ + ErrorCode PublishSnapshot(const ha::SnapshotDescriptor& descriptor); + + /** + * @brief Cleanup old snapshots based on retention policy + * @param keep_count Number of recent snapshots to keep + * @param current_snapshot_id Current snapshot ID (for logging) + */ + void CleanupOldSnapshots(size_t keep_count, + const std::string& current_snapshot_id); + + /** + * @brief List all snapshots from catalog store + * @param limit Maximum number of snapshots to return (0 = unlimited) + * @return Vector of snapshot descriptors on success, error code on failure + */ + tl::expected, ErrorCode> ListSnapshots( + size_t limit); + + /** + * @brief Delete a specific snapshot from catalog store + * @param snapshot_id Snapshot ID to delete + * @return ErrorCode::OK on success, error code on failure + */ + ErrorCode DeleteSnapshot(const ha::SnapshotId& snapshot_id); + + /** + * @brief Get object store connection info for logging + * @return Connection info string + */ + std::string GetObjectStoreConnectionInfo() const; + + /** + * @brief Load the latest snapshot descriptor from catalog + * @return Snapshot descriptor on success, error code on failure + */ + tl::expected LoadLatestSnapshot(); + + /** + * @brief Load all restorable snapshot descriptors + * @param latest_snapshot Optional latest snapshot descriptor to filter + * candidates + * @return Vector of candidate snapshots in chronological order, or error + */ + tl::expected, ErrorCode> + LoadRestoreCandidates( + const std::optional& latest_snapshot); + + /** + * @brief Download snapshot payloads from object storage + * @param descriptor Snapshot descriptor with object paths + * @return Structured payloads ready for decoding, or error + */ + tl::expected + DownloadSnapshotPayloads(const ha::SnapshotDescriptor& descriptor); + + private: + SnapshotObjectStore* object_store_; + ha::SnapshotCatalogStore* catalog_store_; + std::string backup_dir_; + bool use_backup_dir_; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/metadata_store.h b/mooncake-store/include/metadata_store.h index 57b3ce222c..97c886ddf7 100644 --- a/mooncake-store/include/metadata_store.h +++ b/mooncake-store/include/metadata_store.h @@ -2,7 +2,10 @@ #include #include +#include #include +#include +#include #include #include "replica.h" @@ -10,6 +13,10 @@ namespace mooncake { +inline std::string NormalizeTenantId(std::string_view tenant_id) { + return TenantId(std::string(tenant_id)).value(); +} + /** * @brief Metadata structure for Standby to store and restore object information * @@ -25,7 +32,10 @@ struct StandbyObjectMetadata { // 2. After promotion, new Primary should grant fresh leases, not restore // old ones uint64_t last_sequence_id{ - 0}; // Last OpLog sequence ID that modified this key + 0}; // Last OpLog sequence ID that modified this key + std::string group_id; // Tenant group identifier + ObjectDataType data_type{ + ObjectDataType::UNKNOWN}; // Data type classification StandbyObjectMetadata() = default; @@ -33,6 +43,48 @@ struct StandbyObjectMetadata { bool HasReplicas() const { return !replicas.empty(); } }; +/** + * Segment info stored in standby's segment registry. + * Used for recovering segment view after promotion. + */ +struct StandbySegmentInfo { + std::string segment_name; + std::string transport_endpoint; + uint64_t capacity{0}; + bool is_memory_segment{false}; + std::string file_path; // empty for memory segments + + YLT_REFL(StandbySegmentInfo, segment_name, transport_endpoint, capacity, + is_memory_segment, file_path); +}; + +/** + * @brief Standby object entry with tenant-aware key + * + * Replaces std::pair for cleaner + * struct_pack serialization and explicit tenant_id support. + */ +struct StandbyObjectEntry { + std::string tenant_id{"default"}; + std::string key; + StandbyObjectMetadata metadata; + + YLT_REFL(StandbyObjectEntry, tenant_id, key, metadata); +}; + +/** + * Complete snapshot exported from standby at promotion time. + * Includes applied OpLog sequence ID, all object metadata, + * and all registered segments. + */ +struct StandbySnapshot { + uint64_t oplog_sequence_id{0}; + std::vector segments; + std::vector objects; + + YLT_REFL(StandbySnapshot, oplog_sequence_id, segments, objects); +}; + /** * @brief Payload structure for struct_pack serialization (msgpack binary * format) @@ -43,9 +95,10 @@ struct MetadataPayload { UUID client_id{0, 0}; uint64_t size{0}; std::vector replicas; - // NOTE: Lease information removed - not needed by Standby + struct_pack::compatible group_id; // Tenant group + struct_pack::compatible data_type; // Data type - YLT_REFL(MetadataPayload, client_id, size, replicas); + YLT_REFL(MetadataPayload, client_id, size, replicas, group_id, data_type); // Convert to StandbyObjectMetadata StandbyObjectMetadata ToStandbyMetadata(uint64_t sequence_id) const { @@ -54,10 +107,38 @@ struct MetadataPayload { meta.size = size; meta.replicas = replicas; meta.last_sequence_id = sequence_id; + meta.group_id = group_id.value_or(""); + meta.data_type = data_type.value_or(ObjectDataType::UNKNOWN); return meta; } }; +/** + * Thread-safe registry of segments known to standby. + * Maintained by applying SEGMENT_MOUNT/UNMOUNT/UPDATE OpLog events. + * Used to reconstruct segment view after promotion. + */ +class StandbySegmentRegistry { + public: + StandbySegmentRegistry() = default; + + // Segment lifecycle events + void OnSegmentMount(const StandbySegmentInfo& info); + void OnSegmentUnmount(const std::string& transport_endpoint); + void OnSegmentUpdate(const StandbySegmentInfo& info); + + // Queries + bool HasSegment(const std::string& transport_endpoint) const; + std::optional GetSegment( + const std::string& transport_endpoint) const; + std::vector GetAllSegments() const; + void Clear(); + + private: + mutable std::shared_mutex mutex_; + std::unordered_map segments_by_endpoint_; +}; + /** * @brief Abstract interface for metadata storage on Standby * @@ -69,51 +150,40 @@ class MetadataStore { public: virtual ~MetadataStore() = default; - /** - * @brief Put or update metadata for a key with structured metadata - * @param key Object key - * @param metadata Structured metadata object - * @return true on success, false on failure - */ - virtual bool PutMetadata(const std::string& key, + // NEW: tenant-aware methods (primary API) + virtual bool PutMetadata(const std::string& tenant_id, + const std::string& key, const StandbyObjectMetadata& metadata) = 0; + virtual std::optional GetMetadata( + const std::string& tenant_id, const std::string& key) const = 0; + virtual bool Remove(const std::string& tenant_id, + const std::string& key) = 0; + virtual bool Exists(const std::string& tenant_id, + const std::string& key) const = 0; + virtual size_t GetKeyCountForTenant(const std::string& tenant_id) const = 0; + + // DEPRECATED: key-only overloads delegate to tenant-aware with "default" + virtual bool PutMetadata(const std::string& key, + const StandbyObjectMetadata& metadata) { + return PutMetadata("default", key, metadata); + } + virtual std::optional GetMetadata( + const std::string& key) const { + return GetMetadata("default", key); + } + virtual bool Remove(const std::string& key) { + return Remove("default", key); + } + virtual bool Exists(const std::string& key) const { + return Exists("default", key); + } - /** - * @brief Put or update metadata for a key (legacy interface for backward - * compatibility) - * @param key Object key - * @param payload Optional payload data (JSON serialized metadata) - * @return true on success, false on failure - */ + // NOTE: legacy Put(key, payload) remains as default-tenant delegate. + // StandbyMetadataStore and MockMetadataStore continue to implement it. virtual bool Put(const std::string& key, const std::string& payload = std::string()) = 0; - /** - * @brief Get metadata for a key - * @param key Object key - * @return Copy of metadata if found, std::nullopt otherwise - */ - virtual std::optional GetMetadata( - const std::string& key) const = 0; - - /** - * @brief Remove metadata for a key - * @param key Object key - * @return true if key was found and removed, false otherwise - */ - virtual bool Remove(const std::string& key) = 0; - - /** - * @brief Check if a key exists - * @param key Object key - * @return true if key exists, false otherwise - */ - virtual bool Exists(const std::string& key) const = 0; - - /** - * @brief Get the count of keys in the store - * @return Number of keys - */ + // GetKeyCount semantics unchanged - returns total across ALL tenants virtual size_t GetKeyCount() const = 0; }; diff --git a/mooncake-store/include/offset_allocator/offset_allocator.hpp b/mooncake-store/include/offset_allocator/offset_allocator.h similarity index 72% rename from mooncake-store/include/offset_allocator/offset_allocator.hpp rename to mooncake-store/include/offset_allocator/offset_allocator.h index 20052ee1ca..d6740119e9 100644 --- a/mooncake-store/include/offset_allocator/offset_allocator.hpp +++ b/mooncake-store/include/offset_allocator/offset_allocator.h @@ -2,13 +2,15 @@ // (C) Sebastian Aaltonen 2023 // MIT License (see file: LICENSE) +#include #include #include #include +#include #include #include "mutex.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" namespace mooncake::offset_allocator { typedef unsigned char uint8; @@ -43,6 +45,7 @@ struct OffsetAllocation { friend class __Allocator; friend class Serializer; + friend class OffsetAllocator; // for createHandleAtNode during recovery }; struct OffsetAllocStorageReport { @@ -156,6 +159,26 @@ class OffsetAllocator : public std::enable_shared_from_this { [[nodiscard]] std::optional allocate(size_t size); + // ===== Recovery helpers ===== + + template + void visit_used_nodes(Func&& callback) const; + + [[nodiscard]] std::optional createHandleAtNode( + uint32_t node_index, uint64_t real_offset, uint64_t requested_size); + + // Returns the actual region size consumed by allocate(size), or zero when + // the request cannot be represented by this allocator. + [[nodiscard]] uint64_t normalizedAllocationSize(size_t size) const; + + // Returns a mutex-free atomic upper-bound hint for the largest allocatable + // free region. The hint may be larger than the current value after a + // successful allocation, is tightened on a locked allocation failure, and + // is raised when free creates a larger region. + [[nodiscard]] uint64_t getLargestFreeRegion() const noexcept { + return m_largest_free_region.load(std::memory_order_relaxed); + } + // Get storage report (thread-safe) [[nodiscard]] OffsetAllocStorageReport storageReport() const; @@ -186,6 +209,8 @@ class OffsetAllocator : public std::enable_shared_from_this { [[nodiscard]] OffsetAllocatorMetrics get_metrics_internal() const REQUIRES(m_mutex); + void refreshLargestFreeRegion() REQUIRES(m_mutex); + std::unique_ptr<__Allocator> m_allocator GUARDED_BY(m_mutex); uint64_t m_base; // The real offset and size of the allocated memory need to be multiplied by @@ -197,6 +222,8 @@ class OffsetAllocator : public std::enable_shared_from_this { // Lightweight metrics maintained during allocation/deallocation uint64_t m_allocated_size GUARDED_BY(m_mutex) = 0; uint64_t m_allocated_num GUARDED_BY(m_mutex) = 0; + std::atomic m_largest_free_region{0}; + bool m_largest_free_region_tightened GUARDED_BY(m_mutex) = false; // Private constructor - use create() factory method instead OffsetAllocator(uint64_t base, size_t size, uint32 init_capacity, @@ -225,7 +252,9 @@ class __Allocator { void reset(); OffsetAllocation allocate(uint32 size); - void free(OffsetAllocation allocation); + // Returns the size of the newly merged free region, or zero if allocator + // metadata capacity still prevents another allocation. + uint32 free(OffsetAllocation allocation); uint32 allocationSize(OffsetAllocation allocation) const; OffsetAllocStorageReport storageReport() const; @@ -266,6 +295,7 @@ class __Allocator { friend class OffsetAllocatorTest; // for unit tests friend class mooncake::Serializer<__Allocator>; + friend class OffsetAllocator; // for visit_used_nodes / createHandleAtNode }; // Template method implementations @@ -303,7 +333,26 @@ OffsetAllocator::OffsetAllocator(T& serializer) { serializer.read(&m_capacity, sizeof(m_capacity)); serializer.read(&m_allocated_size, sizeof(m_allocated_size)); serializer.read(&m_allocated_num, sizeof(m_allocated_num)); + // Sanity-check the fields that drive later bit shifts and + // allocations; corrupt values must fail loudly here (caught by + // the caller as "corrupt meta") instead of causing UB or OOM. + if (m_multiplier_bits >= 32 || m_capacity == 0) { + LOG(ERROR) << "Deserializing OffsetAllocator failed: corrupt " + "header fields (multiplier_bits=" + << m_multiplier_bits << ", capacity=" << m_capacity + << ")"; + throw std::runtime_error( + "Deserializing OffsetAllocator failed: corrupt header"); + } m_allocator = std::make_unique<__Allocator>(serializer); + const uint64_t largest_free_region = + m_allocator->storageReport().largestFreeRegion << m_multiplier_bits; + m_largest_free_region.store(largest_free_region, + std::memory_order_relaxed); + const uint64_t allocator_capacity = + static_cast(m_allocator->m_size) << m_multiplier_bits; + m_largest_free_region_tightened = + largest_free_region < allocator_capacity; } catch (const std::exception& e) { LOG(ERROR) << "Deserializing OffsetAllocator failed, error=" << e.what(); @@ -345,6 +394,26 @@ __Allocator::__Allocator(T& serializer) { serializer.read(&m_binIndices, sizeof(m_binIndices)); serializer.read(&m_freeOffset, sizeof(m_freeOffset)); + // Sanity-check the values that drive the allocations below. A + // corrupt-but-parseable meta could otherwise request billions of + // nodes and trigger the OOM killer before bad_alloc is ever + // thrown, defeating the "corrupt meta -> fresh start" fallback. + // 1<<24 (16.7M) stays above every legitimate configuration (the + // storage backend clamps node capacity to ~9.6M). + static constexpr uint32 kMaxSerializedNodes = 1u << 24; + if (m_max_capacity == 0 || m_max_capacity > kMaxSerializedNodes || + m_current_capacity > m_max_capacity || + m_freeOffset > m_current_capacity || m_size == 0) { + LOG(ERROR) << "Deserializing __Allocator failed: corrupt " + "capacity fields (max_capacity=" + << m_max_capacity + << ", current_capacity=" << m_current_capacity + << ", freeOffset=" << m_freeOffset << ", size=" << m_size + << ")"; + throw std::runtime_error( + "Deserializing __Allocator failed: corrupt capacities"); + } + // Allocate memory for nodes and freeNodes m_nodes.reserve(m_max_capacity); m_freeNodes.reserve(m_max_capacity); @@ -362,4 +431,36 @@ __Allocator::__Allocator(T& serializer) { } } +// Out-of-line template definition for visit_used_nodes (must be +// after __Allocator is fully defined to access m_nodes etc.) +template +void OffsetAllocator::visit_used_nodes(Func&& callback) const { + struct NodeInfo { + uint64_t real_offset; + uint64_t alloc_size; + uint32_t node_index; + }; + std::vector infos; + { + MutexLocker guard(&m_mutex); + if (!m_allocator) return; + infos.reserve(std::min( + static_cast(m_allocator->m_current_capacity) / 4, + m_allocated_num * 2ULL + 16)); + for (uint32_t i = 0; i < m_allocator->m_current_capacity; ++i) { + const auto& node = m_allocator->m_nodes[i]; + if (node.used) { + infos.push_back( + {m_base + (static_cast(node.dataOffset) + << m_multiplier_bits), + static_cast(node.dataSize) << m_multiplier_bits, + i}); + } + } + } + for (const auto& info : infos) { + callback(info.real_offset, info.alloc_size, info.node_index); + } +} + } // namespace mooncake::offset_allocator diff --git a/mooncake-store/include/pinned_buffer_pool.h b/mooncake-store/include/pinned_buffer_pool.h index 11cdb9d815..25461f7d1f 100644 --- a/mooncake-store/include/pinned_buffer_pool.h +++ b/mooncake-store/include/pinned_buffer_pool.h @@ -1,25 +1,19 @@ #pragma once +#include #include +#include #include -#include -#include "cuda_alike.h" -// Ascend CANN is not covered by cuda_alike.h -#if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) -#include -#endif +#include "device/accelerator_registry.h" +#include "pinned_host_buffer.h" namespace mooncake { /** * PinnedBufferPool: Thread-safe pool of reusable pinned host memory buffers. * - * Platform pinned alloc APIs: - * CUDA / MUSA / MACA : cudaMallocHost (mapped via cuda_alike.h) - * HIP : hipHostMalloc (not mapped in hip.h, native API) - * Ascend : aclrtMallocHost - * Other : new char[] (pageable fallback) + * Platform-specific pinned host allocation is delegated to AcceleratorDevice. * * Pinned memory provides 10x~100x higher D2H bandwidth than pageable memory. * Falls back to new char[] if pinned allocation fails. @@ -33,9 +27,46 @@ class PinnedBufferPool { static constexpr size_t kDefaultMaxPoolSize = 32; struct Buffer { + PinnedHostBuffer pinned_host; + std::unique_ptr pageable_host; char* data = nullptr; size_t capacity = 0; - bool is_pinned = false; // Selects correct free API in FreeBuffer + + Buffer() = default; + explicit Buffer(PinnedHostBuffer pinned_host) + : pinned_host(std::move(pinned_host)), + data(static_cast(this->pinned_host.addr)), + capacity(this->pinned_host.size) {} + + static Buffer Pageable(size_t size) { + Buffer buf; + buf.pageable_host = std::make_unique(size); + buf.data = buf.pageable_host.get(); + buf.capacity = size; + return buf; + } + + Buffer(const Buffer&) = delete; + Buffer& operator=(const Buffer&) = delete; + Buffer(Buffer&& other) noexcept + : pinned_host(std::move(other.pinned_host)), + pageable_host(std::move(other.pageable_host)), + data(other.data), + capacity(other.capacity) { + other.data = nullptr; + other.capacity = 0; + } + Buffer& operator=(Buffer&& other) noexcept { + if (this != &other) { + pinned_host = std::move(other.pinned_host); + pageable_host = std::move(other.pageable_host); + data = other.data; + capacity = other.capacity; + other.data = nullptr; + other.capacity = 0; + } + return *this; + } }; explicit PinnedBufferPool(size_t max_pool_size = kDefaultMaxPoolSize) @@ -48,9 +79,11 @@ class PinnedBufferPool { std::lock_guard lk(mutex_); for (size_t i = 0; i < pool_.size(); ++i) { if (pool_[i].capacity >= size) { - Buffer buf = pool_[i]; + Buffer buf = std::move(pool_[i]); // O(1) erase: swap with back then pop - pool_[i] = pool_.back(); + if (i != pool_.size() - 1) { + pool_[i] = std::move(pool_.back()); + } pool_.pop_back(); return buf; } @@ -62,9 +95,9 @@ class PinnedBufferPool { void Release(Buffer buf) { std::lock_guard lk(mutex_); if (pool_.size() < max_pool_size_) { - pool_.push_back(buf); + pool_.push_back(std::move(buf)); } else { - // Pool full — free immediately to bound pinned memory usage + // Pool full: free immediately to bound pinned memory usage. FreeBuffer(buf); } } @@ -79,57 +112,19 @@ class PinnedBufferPool { private: static Buffer AllocNew(size_t size) { - Buffer buf; - buf.capacity = size; - buf.is_pinned = false; - -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ - defined(USE_HYGON) || defined(USE_COREX) - if (cudaMallocHost(reinterpret_cast(&buf.data), size) == - cudaSuccess) { - buf.is_pinned = true; - } else { - buf.data = new char[size]; + const auto& registry = device::GetAcceleratorRegistry(); + auto runtime_accelerator = registry.RuntimeAccelerators(); + for (auto* accelerator : runtime_accelerator.Devices()) { + auto host = accelerator->AllocatePinnedHost(size); + if (host.addr) return Buffer(std::move(host)); } - -#elif defined(USE_HIP) - if (hipHostMalloc(reinterpret_cast(&buf.data), size, 0) == - hipSuccess) { - buf.is_pinned = true; - } else { - buf.data = new char[size]; - } - -#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) - if (aclrtMallocHost(reinterpret_cast(&buf.data), size) == - ACL_SUCCESS) { - buf.is_pinned = true; - } else { - buf.data = new char[size]; - } - -#else - buf.data = new char[size]; -#endif - return buf; + return Buffer::Pageable(size); } static void FreeBuffer(Buffer& buf) { - if (!buf.data) return; - if (!buf.is_pinned) { - delete[] buf.data; - return; - } -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ - defined(USE_HYGON) || defined(USE_COREX) - cudaFreeHost(buf.data); -#elif defined(USE_HIP) - hipHostFree(buf.data); -#elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) - aclrtFreeHost(buf.data); -#else - delete[] buf.data; -#endif + buf.pinned_host.reset(); + buf.pageable_host.reset(); + buf = {}; } const size_t max_pool_size_; diff --git a/mooncake-store/include/pinned_host_buffer.h b/mooncake-store/include/pinned_host_buffer.h new file mode 100644 index 0000000000..b58f86b0a8 --- /dev/null +++ b/mooncake-store/include/pinned_host_buffer.h @@ -0,0 +1,51 @@ +#pragma once + +#include +#include + +namespace mooncake { + +using PinnedHostBufferDeleter = void (*)(void* addr); + +struct PinnedHostBuffer { + void* addr = nullptr; + size_t size = 0; + PinnedHostBufferDeleter deleter = nullptr; + + PinnedHostBuffer() = default; + PinnedHostBuffer(void* addr, size_t size, PinnedHostBufferDeleter deleter) + : addr(addr), size(size), deleter(deleter) {} + + PinnedHostBuffer(const PinnedHostBuffer&) = delete; + PinnedHostBuffer& operator=(const PinnedHostBuffer&) = delete; + + PinnedHostBuffer(PinnedHostBuffer&& other) noexcept + : addr(other.addr), size(other.size), deleter(other.deleter) { + other.addr = nullptr; + other.size = 0; + other.deleter = nullptr; + } + PinnedHostBuffer& operator=(PinnedHostBuffer&& other) noexcept { + if (this != &other) { + reset(); + addr = other.addr; + size = other.size; + deleter = other.deleter; + other.addr = nullptr; + other.size = 0; + other.deleter = nullptr; + } + return *this; + } + + ~PinnedHostBuffer() { reset(); } + + void reset() { + if (addr && deleter) deleter(addr); + addr = nullptr; + size = 0; + deleter = nullptr; + } +}; + +} // namespace mooncake diff --git a/mooncake-store/include/pyclient.h b/mooncake-store/include/pyclient.h index c2e2201c5d..4219ae6202 100644 --- a/mooncake-store/include/pyclient.h +++ b/mooncake-store/include/pyclient.h @@ -5,13 +5,14 @@ #include #include #include +#include #include #include #include #include #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "mutex.h" #include "utils.h" #include "file_storage.h" @@ -59,8 +60,8 @@ build_ranged_read_results_like( ? all_dst_offsets[i][j].size() : 1; std::vector fragments; - fragments.reserve(std::max(fragment_count, 1)); - for (size_t k = 0; k < std::max(fragment_count, 1); ++k) { + fragments.reserve(fragment_count); + for (size_t k = 0; k < fragment_count; ++k) { fragments.push_back(make_error()); } key_rows.emplace_back(std::move(fragments)); @@ -225,7 +226,9 @@ class PyClient { const std::shared_ptr &transfer_engine, const std::string &ipc_socket_path, bool enable_ssd_offload = false, const std::string &ssd_offload_path = "", - const std::string &tenant_id = "default") = 0; + const std::string &tenant_id = "default", + bool enable_client_http_server = false, + int client_http_port = DEFAULT_CLIENT_HTTP_PORT) = 0; virtual int setup_dummy(size_t mem_pool_size, size_t local_buffer_size, const std::string &server_address, @@ -368,6 +371,13 @@ class PyClient { virtual tl::expected query_task( const UUID &task_id) = 0; + virtual std::optional allocate_client_buffer(size_t size) { + if (!client_buffer_allocator_) { + return std::nullopt; + } + return client_buffer_allocator_->allocate(size); + } + std::shared_ptr client_ = nullptr; std::shared_ptr client_requester_ = nullptr; std::shared_ptr file_storage_ = nullptr; @@ -398,7 +408,8 @@ inline CachedQueryResultResponse to_cached_query_result_response( return CachedQueryResultResponse(GetReplicaListResponse( std::vector(query_result->replicas.begin(), query_result->replicas.end()), - remaining_lease_ttl_ms(*query_result, now))); + remaining_lease_ttl_ms(*query_result, now), + query_result->object_checksum)); } inline tl::expected from_cached_query_result_response( @@ -410,7 +421,8 @@ inline tl::expected from_cached_query_result_response( return QueryResult( std::vector(cached_result.value.replicas.begin(), cached_result.value.replicas.end()), - now + std::chrono::milliseconds(cached_result.value.lease_ttl_ms)); + now + std::chrono::milliseconds(cached_result.value.lease_ttl_ms), + cached_result.value.object_checksum); } inline PyClient::QueryResultCache build_query_result_cache_from_cached_results( diff --git a/mooncake-store/include/random.h b/mooncake-store/include/random.h new file mode 100644 index 0000000000..b6e8e67a88 --- /dev/null +++ b/mooncake-store/include/random.h @@ -0,0 +1,75 @@ +#pragma once + +#include +#include +#include +#include + +namespace mooncake { + +using RandomEngine = std::mt19937_64; + +namespace detail { + +template +Result sampleUniform(Result lower_bound, Result upper_bound, + Generator& generator) { + std::uniform_int_distribution distribution( + static_cast(lower_bound), + static_cast(upper_bound)); + return static_cast(distribution(generator)); +} + +} // namespace detail + +// Returns the pseudo-random engine shared by random helpers on this thread. +// The engine is seeded once per thread and is not safe to use from another +// thread. +inline RandomEngine& threadLocalRandomEngine() { + thread_local RandomEngine engine = [] { + std::random_device device; + std::seed_seq seed{device(), device(), device(), device(), + device(), device(), device(), device()}; + return RandomEngine(seed); + }(); + return engine; +} + +template +size_t randomIndex(size_t upper_bound, Generator& generator) { + if (upper_bound == 0) { + throw std::invalid_argument("randomIndex upper bound must be positive"); + } + std::uniform_int_distribution distribution(0, upper_bound - 1); + return distribution(generator); +} + +// Returns an unbiased random index in [0, upper_bound). +inline size_t randomIndex(size_t upper_bound) { + return randomIndex(upper_bound, threadLocalRandomEngine()); +} + +template +Integer randomUniform(Integer lower_bound, Integer upper_bound, + Generator& generator) { + if (lower_bound > upper_bound) { + throw std::invalid_argument( + "randomUniform lower bound must not exceed upper bound"); + } + if constexpr (std::signed_integral) { + return detail::sampleUniform( + lower_bound, upper_bound, generator); + } else { + return detail::sampleUniform( + lower_bound, upper_bound, generator); + } +} + +// Returns an unbiased random integer in [lower_bound, upper_bound]. +template +Integer randomUniform(Integer lower_bound, Integer upper_bound) { + return randomUniform(lower_bound, upper_bound, threadLocalRandomEngine()); +} + +} // namespace mooncake diff --git a/mooncake-store/include/real_client.h b/mooncake-store/include/real_client.h index 0678964497..89bf4a4c4e 100644 --- a/mooncake-store/include/real_client.h +++ b/mooncake-store/include/real_client.h @@ -14,10 +14,14 @@ #include "pyclient.h" #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" +#include "device/cuda_ipc_buffer_handle.h" #include "mutex.h" #include "utils.h" #include "rpc_types.h" +#if defined(USE_SUNRISE) +#include "sunrise_allocator.h" +#endif #include #include #include @@ -26,6 +30,9 @@ namespace mooncake { class RealClient; +class RegisteredPinnedRegion; +class UdsAcceptor; +class UdsConnection; // Global resource tracker to handle cleanup on abnormal termination class ResourceTracker { @@ -84,7 +91,9 @@ class RealClient : public PyClient { const std::string &ipc_socket_path = "", bool enable_ssd_offload = false, const std::string &ssd_offload_path = "", - const std::string &tenant_id = "default"); + const std::string &tenant_id = "default", + bool enable_client_http_server = false, + int client_http_port = DEFAULT_CLIENT_HTTP_PORT); int setup_dummy(size_t mem_pool_size, size_t local_buffer_size, const std::string &server_address, @@ -381,6 +390,9 @@ class RealClient : public PyClient { batch_acquire_buffer_dummy(const std::vector &keys, const UUID &client_id); + tl::expected, ErrorCode> allocate_buffer_dummy( + size_t size, const UUID &client_id); + tl::expected put_dummy_helper( const std::string &key, std::span value, const ReplicateConfig &config, const UUID &client_id); @@ -437,6 +449,15 @@ class RealClient : public PyClient { const ReplicateConfig &config, int32_t device_id, const UUID &client_id); + std::vector> + batch_put_from_cuda_ipc_dummy_helper( + const std::vector &requests, + const ReplicateConfig &config, const UUID &client_id); + + std::vector> + batch_get_into_cuda_ipc_dummy_helper( + const std::vector &requests, const UUID &client_id); + std::vector> batch_get_into_multi_buffers_dummy_helper( const std::vector &keys, @@ -447,7 +468,8 @@ class RealClient : public PyClient { tl::expected get_into_range_shm_helper( const std::string &key, uint64_t buffer, size_t dst_offset, - size_t src_offset, size_t size, const UUID &client_id); + size_t src_offset, size_t size, bool size_is_buffer_capacity, + bool verify_checksum, const UUID &client_id); std::vector>>> get_into_ranges_shm_helper( @@ -506,7 +528,9 @@ class RealClient : public PyClient { const std::string &ipc_socket_path = "", int local_rpc_port = 50052, bool enable_ssd_offload = false, bool start_offload_rpc_server = false, const std::string &ssd_offload_path = "", - const std::string &tenant_id = "default"); + const std::string &tenant_id = "default", + bool enable_client_http_server = false, + int client_http_port = DEFAULT_CLIENT_HTTP_PORT); // Overload that accepts a configuration dictionary tl::expected setup_internal(const ConfigDict &config); @@ -538,16 +562,18 @@ class RealClient : public PyClient { tl::expected query_result); tl::expected resolve_ranged_read_metadata( - const std::string &key); + const std::string &key, + const QueryResultCache *query_result_cache = nullptr); tl::expected execute_ranged_read( const std::string &key, void *buffer, size_t dst_offset, size_t src_offset, size_t size, const RangedReadMetadata &metadata, - bool size_is_buffer_capacity = false); + bool size_is_buffer_capacity, bool verify_checksum); tl::expected get_into_range_internal( const std::string &key, void *buffer, size_t dst_offset, - size_t src_offset, size_t size, bool size_is_buffer_capacity = false); + size_t src_offset, size_t size, bool size_is_buffer_capacity, + bool verify_checksum); std::vector>>> get_into_ranges_internal( @@ -557,10 +583,6 @@ class RealClient : public PyClient { const std::vector>> &all_src_offsets, const std::vector>> &all_sizes, const std::vector *buffer_capacities = nullptr, - std::vector>>> - *prepared_results = nullptr, - const std::vector>> *valid_fragments = - nullptr, const QueryResultCache *query_result_cache = nullptr); std::vector> batch_get_into_internal( @@ -747,8 +769,11 @@ class RealClient : public PyClient { void *base = nullptr; size_t size = 0; std::string protocol; + std::shared_ptr pinned_region; }; + void FreeAllocatedStoreSegment(AllocatedSegmentRecord &record); + std::unique_ptr port_binder_ = nullptr; struct SegmentDeleter { @@ -787,12 +812,42 @@ class RealClient : public PyClient { } }; +#ifdef USE_VRAM_SEGMENT + struct VRAMSegmentDeleter { + void operator()(void *ptr) { + if (ptr) { + free_memory("vram", ptr); + } + } + }; +#endif + +#if defined(USE_SUNRISE) + struct SunriseSegmentDeleter { + void operator()(void *ptr) { + if (ptr) { + sunrise_free_memory(ptr); + } + } + }; +#endif + std::vector> hugepage_segment_ptrs_; std::vector> segment_ptrs_; std::vector> ascend_segment_ptrs_; std::vector> ub_segment_ptrs_; +#ifdef USE_VRAM_SEGMENT + std::vector> vram_segment_ptrs_; +#endif +#if defined(USE_SUNRISE) + std::vector> + sunrise_segment_ptrs_; +#endif + std::vector> + setup_segment_pinned_regions_; + bool setup_segment_memory_must_leak_ = false; std::string protocol; std::string device_name; std::string local_hostname; @@ -840,7 +895,8 @@ class RealClient : public PyClient { bool map_dummy_buffer_to_real(const ShmContext &shm_ctx, uint64_t dummy_addr, size_t buf_size, const MappedShm *&last_hit_shm, - void *&out_real) const; + void *&out_real, + size_t *out_capacity = nullptr) const; bool map_dummy_buffer_range_to_real(const ShmContext &shm_ctx, uint64_t dummy_addr, size_t dst_offset, @@ -848,7 +904,8 @@ class RealClient : public PyClient { tl::expected, ErrorCode> map_dummy_addrs_to_real_ptrs( const ShmContext &context, const std::vector &dummy_addrs, - const std::vector &sizes, const UUID &client_id) const; + const std::vector &sizes, const UUID &client_id, + std::vector *capacities = nullptr) const; tl::expected>, ErrorCode> map_dummy_nested_addrs_to_real_ptrs( @@ -885,18 +942,16 @@ class RealClient : public PyClient { // IPC Server members for receiving FD from Dummy Clients std::string ipc_socket_path_; - std::jthread ipc_thread_; - std::atomic ipc_running_{false}; + std::unique_ptr uds_acceptor_; int start_ipc_server(); int stop_ipc_server(); - void ipc_server_func(); // Embedded HTTP server for health-check / metrics std::unique_ptr http_server_; - int start_http_server(); + int start_http_server(int port); void stop_http_server(); - void handle_ipc_shm_register(int client_sock); - void handle_ipc_shm_fd_request(int client_sock); + void handle_ipc_shm_register(UdsConnection &connection); + void handle_ipc_shm_fd_request(UdsConnection &connection); void teardown_ascend_shm_buffer(MappedShm &shm); tl::expected setup_ascend_internal( diff --git a/mooncake-store/include/replica.h b/mooncake-store/include/replica.h index 677e77b333..3aa42ecf98 100644 --- a/mooncake-store/include/replica.h +++ b/mooncake-store/include/replica.h @@ -91,6 +91,7 @@ struct ReplicateConfig { preferred_nof_segments{}; // Preferred NoF segments for allocation bool prefer_alloc_in_same_node{false}; ObjectDataType data_type{ObjectDataType::UNKNOWN}; + std::string host_id{}; // Optional per-key routing group IDs. Empty string keeps that key // ungrouped. Grouped keys share metadata routing, coalesced lease refresh, // and memory eviction behavior. @@ -130,6 +131,9 @@ struct ReplicateConfig { os << ", prefer_alloc_in_same_node: " << config.prefer_alloc_in_same_node << ", data_type: " << config.data_type; + if (!config.host_id.empty()) { + os << ", host_id: " << config.host_id; + } if (config.group_ids.has_value()) { os << ", group_ids: ["; for (size_t i = 0; i < config.group_ids->size(); ++i) { @@ -372,6 +376,14 @@ class Replica { return false; // DiskReplicaData does not have handles } + bool replace_memory_buffer(std::unique_ptr buffer) { + if (!buffer || !is_memory_replica()) { + return false; + } + std::get(data_).buffer = std::move(buffer); + return true; + } + [[nodiscard]] bool has_invalid_nof_handle() const { if (is_nof_replica()) { const auto& nof_data = std::get(data_); @@ -441,6 +453,17 @@ class Replica { } } + void mark_removed() { + if (status_ == ReplicaStatus::COMPLETE || + status_ == ReplicaStatus::PROCESSING) { + status_ = ReplicaStatus::REMOVED; + } else if (status_ == ReplicaStatus::REMOVED) { + LOG(WARNING) << "Replica already marked as removed"; + } else { + LOG(ERROR) << "Cannot mark_removed from status: " << status_; + } + } + void inc_refcnt() { refcnt_.fetch_add(1); } void dec_refcnt() { refcnt_.fetch_sub(1); } diff --git a/mooncake-store/include/replica_selection.h b/mooncake-store/include/replica_selection.h new file mode 100644 index 0000000000..ed7a49c2af --- /dev/null +++ b/mooncake-store/include/replica_selection.h @@ -0,0 +1,168 @@ +// Copyright 2025 Mooncake Authors +// +// Replica selection for reads: given the list of replicas the master returned +// for a key, choose which one to actually fetch from. +// +// The base policy is type + locality based (local MEMORY > local NOF_SSD > +// remote MEMORY > remote NOF_SSD > LOCAL_DISK > DISK). On top of that, when a +// key has more than one *remote* MEMORY replica, the historical code kept the +// first one master happened to return, which is effectively arbitrary. This +// header adds an opt-in scoring hook so a better remote replica can be picked +// (issue #2516). +// +// Design constraints: +// * Disabled by default — behaviour is byte-identical to the historical +// "first remote MEMORY" pick unless the operator sets +// MC_STORE_REPLICA_SCORING=1 or a scorer is injected. +// * This layer only sees what a replica descriptor carries (endpoint, +// protocol). Richer signals (NIC role, NUMA distance, live load) live in +// the transfer engine; they can be fed in via SetRemoteReplicaScorer() +// without mooncake-store growing a dependency on that layer. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "replica.h" + +namespace mooncake { + +// Scores a remote replica candidate; lower score == more preferred. +using ReplicaScorer = std::function; + +namespace detail { + +inline std::shared_mutex &ScorerMutex() { + static std::shared_mutex mu; + return mu; +} + +inline ReplicaScorer &ScorerStorage() { + static ReplicaScorer scorer; + return scorer; +} + +} // namespace detail + +// Return a snapshot (copy) of the current scorer. The copy is taken under a +// shared lock so concurrent reads are non-blocking. Callers invoke the +// returned copy outside the lock, avoiding both the data race and any +// potential deadlock from calling user code under a lock. +inline ReplicaScorer GetRemoteReplicaScorer() { + std::shared_lock lk(detail::ScorerMutex()); + return detail::ScorerStorage(); +} + +// Inject a topology-/load-aware scorer (e.g. from the transfer engine). +// Takes a unique lock; expected to be called once during initialization. +inline void SetRemoteReplicaScorer(ReplicaScorer scorer) { + std::unique_lock lk(detail::ScorerMutex()); + detail::ScorerStorage() = std::move(scorer); +} + +// Built-in remote-replica score: prefer RDMA over TCP. Purely a function of +// info the replica already carries (protocol_), so it needs no topology. +inline double BuiltinRemoteReplicaScore(const Replica::Descriptor &r) { + if (!r.is_memory_replica()) return 100.0; + const std::string &proto = + r.get_memory_descriptor().buffer_descriptor.protocol_; + if (proto == "rdma") return 0.0; + if (proto == "tcp") return 1.0; + return 2.0; // unknown protocol — least preferred, but still usable +} + +// Whether remote-replica scoring is active. Opt-in via env; always on if a +// scorer has been injected. Env is read once; the injected-scorer check is +// live so tests / late injection take effect. +inline bool RemoteReplicaScoringEnabled() { + static const bool env_enabled = [] { + const char *env = std::getenv("MC_STORE_REPLICA_SCORING"); + return env && std::string(env) == "1"; + }(); + if (env_enabled) return true; + std::shared_lock lk(detail::ScorerMutex()); + return static_cast(detail::ScorerStorage()); +} + +// Among the remote MEMORY replicas, return the lowest-scoring one (nullptr if +// none). Ties keep master return order (strictly-less comparison), so a +// symmetric cluster degrades to the historical "first remote MEMORY" pick. +inline const Replica::Descriptor *PickBestRemoteMemory( + const std::vector &replicas, + const std::unordered_set &local_endpoints) { + ReplicaScorer scorer = GetRemoteReplicaScorer(); + const Replica::Descriptor *best = nullptr; + double best_score = std::numeric_limits::max(); + for (const auto &r : replicas) { + if (r.status != ReplicaStatus::COMPLETE) continue; + if (!r.is_memory_replica()) continue; + if (local_endpoints.count(r.get_memory_descriptor() + .buffer_descriptor.transport_endpoint_)) + continue; // local replicas are handled by the caller + double score = scorer ? scorer(r) : BuiltinRemoteReplicaScore(r); + if (score < best_score) { + best_score = score; + best = &r; + } + } + return best; +} + +// Select the best replica from a list: prefer local MEMORY, then any MEMORY, +// then LOCAL_DISK, then DISK. Master may return replicas in any order, so we +// always scan. When scoring is enabled and there are multiple remote MEMORY +// replicas, the best-scoring one is chosen instead of the first encountered. +inline const Replica::Descriptor *SelectBestReplica( + const std::vector &replicas, + const std::unordered_set &local_endpoints) { + const Replica::Descriptor *first_memory = nullptr; + const Replica::Descriptor *first_nof = nullptr; + for (const auto &r : replicas) { + if (r.status != ReplicaStatus::COMPLETE) continue; + if (r.is_memory_replica()) { + if (local_endpoints.count( + r.get_memory_descriptor() + .buffer_descriptor.transport_endpoint_)) { + return &r; // local MEMORY — best case + } + if (!first_memory) first_memory = &r; + } else if (r.is_nof_replica()) { + if (local_endpoints.count( + r.get_nof_descriptor() + .buffer_descriptor.transport_endpoint_)) { + return &r; // local NOF_SSD — also good + } + if (!first_nof) first_nof = &r; + } + } + // No local replica. Among remote MEMORY replicas, optionally pick the + // best-scoring one instead of the first encountered (issue #2516). + if (first_memory && RemoteReplicaScoringEnabled()) { + if (const auto *scored = + PickBestRemoteMemory(replicas, local_endpoints)) { + return scored; + } + } + if (first_memory) return first_memory; + if (first_nof) return first_nof; + + const Replica::Descriptor *best = nullptr; + for (const auto &r : replicas) { + if (r.status != ReplicaStatus::COMPLETE) continue; + if (r.is_local_disk_replica()) { + best = &r; // LOCAL_DISK always overrides DISK + } else if (r.is_disk_replica() && !best) { + best = &r; + } + } + return best; +} + +} // namespace mooncake diff --git a/mooncake-store/include/rpc_service.h b/mooncake-store/include/rpc_service.h index 3f7de060dd..53e7212452 100644 --- a/mooncake-store/include/rpc_service.h +++ b/mooncake-store/include/rpc_service.h @@ -12,13 +12,25 @@ #include "types.h" #include "rpc_types.h" #include "master_config.h" +#include "kv_event/kv_event_publisher.h" #include "segment.h" namespace mooncake { +// Forward declaration +class HttpMetadataServer; class WrappedMasterService { public: - WrappedMasterService(const WrappedMasterServiceConfig& config); + // Constructor with optional metadata-cleanup-on-timeout configuration. + // - http_metadata_server: in-process pointer used when the HTTP metadata + // server is co-located in the master process (nullptr = not co-located). + // - http_metadata_remote_url: http(s) connection string used when the + // metadata server is deployed separately (empty = none). Only consulted + // when http_metadata_server is nullptr. If both are unset, cleanup is + // disabled. + WrappedMasterService(const WrappedMasterServiceConfig& config, + HttpMetadataServer* http_metadata_server = nullptr, + const std::string& http_metadata_remote_url = ""); ~WrappedMasterService(); @@ -54,13 +66,22 @@ class WrappedMasterService { BatchGetReplicaList(const std::vector& keys, const std::string& tenant_id = "default"); + // Read-only admin variants: no lease grants, no promotion, no metric + // updates. + std::vector> + BatchGetReplicaListForAdmin(const std::vector& keys, + const std::string& tenant_id = "default"); + + tl::expected GetReplicaListForAdmin( + const std::string& key, const std::string& tenant_id = "default"); + tl::expected, ErrorCode> PutStart( const UUID& client_id, const std::string& key, const uint64_t slice_length, const ReplicateConfig& config, const std::string& tenant_id = "default"); tl::expected PutEnd( - const UUID& client_id, const std::string& key, + const UUID& client_id, const ObjectMeta& object_meta, ReplicaType replica_type = ReplicaType::ALL, const std::string& tenant_id = "default"); @@ -76,7 +97,7 @@ class WrappedMasterService { const std::string& tenant_id = "default"); std::vector> BatchPutEnd( - const UUID& client_id, const std::vector& keys, + const UUID& client_id, const std::vector& object_metas, ReplicaType replica_type = ReplicaType::ALL, const std::string& tenant_id = "default"); @@ -91,7 +112,7 @@ class WrappedMasterService { const std::string& tenant_id = "default"); tl::expected UpsertEnd( - const UUID& client_id, const std::string& key, + const UUID& client_id, const ObjectMeta& object_meta, ReplicaType replica_type = ReplicaType::ALL, const std::string& tenant_id = "default"); @@ -108,7 +129,7 @@ class WrappedMasterService { const std::string& tenant_id = "default"); std::vector> BatchUpsertEnd( - const UUID& client_id, const std::vector& keys, + const UUID& client_id, const std::vector& object_metas, const std::string& tenant_id = "default"); std::vector> BatchUpsertRevoke( @@ -166,6 +187,16 @@ class WrappedMasterService { tl::expected ServiceReady(); + tl::expected, ErrorCode> + ListTenantQuotaSnapshots(); + tl::expected GetTenantQuotaSnapshot( + const std::string& tenant_id); + tl::expected UpsertTenantQuotaPolicy( + const std::string& tenant_id, uint64_t requested_quota_bytes); + tl::expected, ErrorCode> + DeleteTenantQuotaPolicy(const std::string& tenant_id); + tl::expected GetTenantQuotaAllocatableCapacityBytes(); + tl::expected, ErrorCode> GetAllKeysForAdmin(); tl::expected, ErrorCode> GetAllSegmentsForAdmin(); @@ -182,6 +213,8 @@ class WrappedMasterService { tl::expected, ErrorCode> OffloadObjectHeartbeat(const UUID& client_id, bool enable_offloading); + tl::expected PollRemoveAll(const UUID& client_id); + tl::expected ReportSsdCapacity( const UUID& client_id, int64_t ssd_total_capacity_bytes); @@ -217,6 +250,13 @@ class WrappedMasterService { const std::string& segment_name); tl::expected QuerySegmentStatusById( const UUID& segment_id); + + // Internal method called by supervisor during promotion; NOT an RPC + // endpoint. + void RestoreFromStandby(const std::vector& objects, + uint64_t initial_oplog_sequence_id, + const std::vector& segments); + tl::expected CreateCopyTask( const std::string& key, const std::string& tenant_id, const std::vector& targets); @@ -269,6 +309,9 @@ class WrappedMasterService { const UUID& client_id, const std::vector& keys, const std::string& tenant_id, ReplicaType replica_type); + bool KvEventsEnabled() const; + KvEventPublisher::Stats GetKvEventStats() const; + private: MasterService master_service_; }; diff --git a/mooncake-store/include/rpc_types.h b/mooncake-store/include/rpc_types.h index fbf4fe3779..242cc0edfd 100644 --- a/mooncake-store/include/rpc_types.h +++ b/mooncake-store/include/rpc_types.h @@ -1,11 +1,19 @@ #pragma once +#include + #include "types.h" #include "replica.h" #include "task_manager.h" namespace mooncake { +struct ObjectMeta { + std::string key; + std::optional object_checksum; +}; +YLT_REFL(ObjectMeta, key, object_checksum); + /** * @brief Response structure for Ping operation */ @@ -32,14 +40,18 @@ YLT_REFL(PingResponse, view_version_id, client_status); struct GetReplicaListResponse { std::vector replicas; uint64_t lease_ttl_ms; + std::optional object_checksum; GetReplicaListResponse() : lease_ttl_ms(0) {} - GetReplicaListResponse(std::vector&& replicas_param, - uint64_t lease_ttl_ms_param) + GetReplicaListResponse( + std::vector&& replicas_param, + uint64_t lease_ttl_ms_param, + std::optional object_checksum_param = std::nullopt) : replicas(std::move(replicas_param)), - lease_ttl_ms(lease_ttl_ms_param) {} + lease_ttl_ms(lease_ttl_ms_param), + object_checksum(object_checksum_param) {} }; -YLT_REFL(GetReplicaListResponse, replicas, lease_ttl_ms); +YLT_REFL(GetReplicaListResponse, replicas, lease_ttl_ms, object_checksum); struct CachedQueryResultResponse { bool success; diff --git a/mooncake-store/include/segment.h b/mooncake-store/include/segment.h index 0d8aa1c1db..b0085ab9e2 100644 --- a/mooncake-store/include/segment.h +++ b/mooncake-store/include/segment.h @@ -2,7 +2,9 @@ #include #include +#include #include +#include #include #include #include @@ -15,6 +17,9 @@ #include "types.h" namespace mooncake { +using HostSegmentIndex = + std::map>>; + /** * @brief Status of a mounted segment in master */ @@ -86,6 +91,7 @@ struct LocalDiskSegment { mutable Mutex offloading_mutex_; bool enable_offloading; int64_t ssd_total_capacity_bytes = 0; // last reported by client heartbeat + std::atomic ssd_used_bytes{0}; std::unordered_map GUARDED_BY( offloading_mutex_) offloading_objects; // Promotion-on-hit pending work for this client. Populated by master's @@ -94,6 +100,11 @@ struct LocalDiskSegment { // offloading_objects (offloading_mutex_). std::unordered_map GUARDED_BY( offloading_mutex_) promotion_objects; + // Set by master's RemoveAll. When the client sees this flag via + // PollRemoveAll, it calls FileStorage::RemoveAll() to physically + // delete all SSD files. Same locking as offloading_objects + // (offloading_mutex_). + bool GUARDED_BY(offloading_mutex_) pending_remove_all = false; explicit LocalDiskSegment(bool enable_offloading) : enable_offloading(enable_offloading) {} @@ -137,6 +148,22 @@ class ScopedSegmentAccess { ErrorCode ReMountSegment(const std::vector& segments, const UUID& client_id); + ErrorCode ValidateRemountSegment(const Segment& segment, + const UUID& client_id) const; + + bool GetSegment(const UUID& segment_id, Segment& segment) const; + + struct AllocatorReplacement { + UUID segment_id; + std::shared_ptr expected; + std::shared_ptr replacement; + }; + bool ReplaceAllocators( + const std::vector& replacements); + + std::shared_ptr GetAllocator( + const UUID& segment_id) const; + /** * @brief Prepare to unmount a segment by deleting its allocator */ @@ -173,6 +200,9 @@ class ScopedSegmentAccess { ErrorCode GetAllSegments( std::vector>& all_segments); + std::vector GetHostOrderedSegments( + const std::string& writer_host_id, const std::string& key) const; + ErrorCode GetAllSegmentNames(std::vector& all_segment_names); /** @@ -316,10 +346,21 @@ class ScopedAllocatorAccess { std::shared_mutex& mutex) : allocator_manager_(allocator_manager), lock_(mutex) {} + explicit ScopedAllocatorAccess(const AllocatorManager& allocator_manager, + const HostSegmentIndex& segments_by_host, + std::shared_mutex& mutex) + : allocator_manager_(allocator_manager), + segments_by_host_(&segments_by_host), + lock_(mutex) {} + const AllocatorManager& getAllocatorManager() { return allocator_manager_; } + std::vector GetHostOrderedSegments( + const std::string& writer_host_id, const std::string& key) const; + private: const AllocatorManager& allocator_manager_; + const HostSegmentIndex* segments_by_host_{nullptr}; std::shared_lock lock_; }; @@ -327,7 +368,7 @@ class ScopedAllocatorAccess { * @brief RAII-style access to LocalDiskOffloadingQueues for thread-safe * LocalDiskOffloadingQueue usage */ -class ScopedLocalDiskSegmentAccess { +class ScopedLocalDiskSegmentAccess : public SsdMetricsProvider { public: explicit ScopedLocalDiskSegmentAccess( std::unordered_map& client_by_name, @@ -348,6 +389,9 @@ class ScopedLocalDiskSegmentAccess { return client_local_disk_segment_; } + int64_t getSsdTotalCapacity(const std::string& segment_name) const override; + int64_t getSsdUsedBytes(const std::string& segment_name) const override; + private: const std::unordered_map& client_by_name_; // segment name -> client_id @@ -407,6 +451,18 @@ class SegmentManager { bool enable_cxl = false) : memory_allocator_(memory_allocator), enable_cxl_(enable_cxl) {} + /** + * @brief Releases the capacity metric contribution of segments that + * are still mounted. Intended to be called when the owning + * MasterService is torn down: MasterMetricManager outlives + * MasterService instances (e.g. across HA leadership changes). + * This is deliberately not done in the destructor, because other + * SegmentManager instances (such as the temporary snapshot + * readers) hold deserialized records that never contributed to + * the metrics and must not release them. + */ + void releaseCapacityMetrics(); + /** * @brief Get RAII-style access to segment management operations * @return ScopedSegmentAccess object that holds the lock @@ -420,7 +476,8 @@ class SegmentManager { * @return ScopedAllocatorAccess object that holds the lock */ ScopedAllocatorAccess getAllocatorAccess() { - return ScopedAllocatorAccess(allocator_manager_, segment_mutex_); + return ScopedAllocatorAccess(allocator_manager_, segments_by_host_, + segment_mutex_); } ScopedLocalDiskSegmentAccess getLocalDiskSegmentAccess() { @@ -433,6 +490,11 @@ class SegmentManager { void initializeCxlAllocator(const std::string& cxl_path, const size_t cxl_size); + // Endpoint-based segment queries (for standby restore) + bool HasSegmentByEndpoint(const std::string& endpoint) const; + bool GetSegmentBasicInfo(const UUID& segment_id, std::string& segment_name, + std::string& te_endpoint) const; + private: mutable std::shared_mutex segment_mutex_; std::shared_ptr allocation_strategy_; @@ -451,7 +513,9 @@ class SegmentManager { std::unordered_map client_by_name_; // segment name -> client_id std::unordered_map - segment_id_by_name_; // segment name -> segment_id + segment_id_by_name_; // segment name -> segment_id + HostSegmentIndex segments_by_host_; // host_id -> segment name -> segment + // ids for allocatable segments std::unordered_map, boost::hash> client_local_disk_segment_; // client_id -> local_disk_segment diff --git a/mooncake-store/include/serialize/serializer.hpp b/mooncake-store/include/serialize/serializer.h similarity index 99% rename from mooncake-store/include/serialize/serializer.hpp rename to mooncake-store/include/serialize/serializer.h index cb153b1a84..135e345ee4 100644 --- a/mooncake-store/include/serialize/serializer.hpp +++ b/mooncake-store/include/serialize/serializer.h @@ -141,4 +141,4 @@ class SerializationHelper { } }; -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/include/shm_helper.h b/mooncake-store/include/shm_helper.h index e79c8a38ef..b7069ee056 100644 --- a/mooncake-store/include/shm_helper.h +++ b/mooncake-store/include/shm_helper.h @@ -9,24 +9,12 @@ namespace mooncake { -/** - * @brief Send a file descriptor + data payload over a Unix socket (SCM_RIGHTS). - * @return bytes sent on success, -1 on error. - */ -int ipc_send_fd(int socket, int fd, void *data, size_t data_len); - -/** - * @brief Receive a file descriptor + data payload from a Unix socket. - * @return the received fd on success, -1 on error. - */ -int ipc_recv_fd(int socket, void *data, size_t data_len); - /** * @brief Manages anonymous shared memory segments backed by memfd. * - * Each segment is created via memfd_create + mmap. The fd can be passed - * to other processes (via Unix socket SCM_RIGHTS) for cross-process - * zero-copy sharing. + * Each segment is created via memfd_create + mmap. The fd can be passed to + * other processes via UdsConnection::sendFd() for cross-process zero-copy + * sharing. * * Thread-safe singleton; all operations are mutex-protected. */ diff --git a/mooncake-store/include/storage/distributed/distributed_storage_backend.h b/mooncake-store/include/storage/distributed/distributed_storage_backend.h index d4c935bb1f..fe733b7fde 100644 --- a/mooncake-store/include/storage/distributed/distributed_storage_backend.h +++ b/mooncake-store/include/storage/distributed/distributed_storage_backend.h @@ -37,8 +37,7 @@ class DistributedStorageBackend : public StorageBackendInterface { std::function& keys, std::vector& metadatas)> complete_handler, - std::function& evicted_keys)> - eviction_handler = nullptr) override; + EvictionHandler eviction_handler = nullptr) override; tl::expected BatchLoad( std::unordered_map& batched_slices) override; diff --git a/mooncake-store/include/storage_backend.h b/mooncake-store/include/storage_backend.h index 5fd071b59b..83afc3c700 100644 --- a/mooncake-store/include/storage_backend.h +++ b/mooncake-store/include/storage_backend.h @@ -2,18 +2,22 @@ #include +#include #include +#include #include +#include #include #include #include #include #include +#include #include #include "file_interface.h" #include "mutex.h" -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "types.h" namespace mooncake { @@ -200,6 +204,92 @@ struct BucketBackendConfig { static BucketBackendConfig FromEnvironment(); }; +enum class OffsetEvictionPolicy { + NONE, // No eviction + FIFO, // Evict oldest key first (by insertion order) + LRU, // Approximate LRU via cross-shard sampling (phase 2) +}; + +enum class OffsetPersistMode { + kDisabled, // No persistence (default) + kRelaxed, // Periodic checkpoint + kStrict, // Every BatchOffload is durable +}; + +struct OffsetAllocatorBackendConfig { + OffsetEvictionPolicy eviction_policy = OffsetEvictionPolicy::NONE; + + // Watermark thresholds: eviction triggers when total_size_ exceeds high, + // drives down to low. 0 = auto-resolved in Init() from ratios. + int64_t high_watermark_bytes = 0; + int64_t low_watermark_bytes = 0; + double high_ratio = 0.90; + double low_ratio = 0.80; + + // Key-count watermarks (symmetric with byte watermarks). + // high triggers eviction, drives down to low. + int64_t high_watermark_keys = 0; + int64_t low_watermark_keys = 0; + double keys_high_ratio = 0.95; + double keys_low_ratio = 0.90; + + // Eviction caps + size_t max_evict_per_offload = 4096; + size_t fallback_evict_batch = 16; + + // Allocator node capacity override. + // 0 = auto-derived from capacity_ / kMinObjectSize (capped at RAM budget). + // Must be <= UINT32_MAX (OffsetAllocator::create takes uint32 + // max_capacity). + int64_t max_capacity_nodes = 0; + + bool Validate() const; + + static OffsetAllocatorBackendConfig FromEnvironment(); + + // ---- Persistence settings ---- + OffsetPersistMode persist_mode = OffsetPersistMode::kDisabled; + int64_t persist_interval_seconds = 60; + + // ---- Record integrity ---- + // When true (default), every written record carries a CRC-32C over + // header-prefix + key + value (RecordHeader::kFlagHasCrc), verified + // once on recovery. Disable only when torn writes are otherwise + // impossible (kStrict mode on storage that honors fsync ordering, + // e.g. power-loss-protected NVMe) or when values never pass through + // the CPU (future DMA/GDS writers): unchecksummed records are then + // validated by checkpoint ordering (seq guard) alone. + bool enable_record_crc = true; +}; + +// ===== Persistence metadata structures ===== + +struct PersistedFifoEntry { + uint64_t seq; + std::string key; +}; +YLT_REFL(PersistedFifoEntry, seq, key); + +struct OffsetAllocatorPersistedMetadata { + uint32_t version = 1; + std::string allocator_state; + uint64_t insert_seq = 0; + std::vector fifo_entries; + std::vector evicted_keys_this_batch; +}; +YLT_REFL(OffsetAllocatorPersistedMetadata, version, allocator_state, insert_seq, + fifo_entries, evicted_keys_this_batch); + +// Current on-disk format version of OffsetAllocatorPersistedMetadata. +// v2: RecordHeader grew from 8 to 20 bytes (added per-record seq + CRC-32C). +// v3: RecordHeader is 24 bytes (added `flags`; CRC-32C is now optional per +// record) and the value region is aligned to 4 KiB within the record +// (zero padding derived from key_len), so that DMA writers (e.g. GDS) +// can share the layout. +// Older metadata is rejected on load (fresh start) because its data-file +// records cannot be parsed with the current record layout. +inline constexpr uint32_t kOffsetAllocatorPersistVersion = 3; + struct FileStorageConfig { // type of the storage backend StorageBackendType storage_backend_type = StorageBackendType::kBucket; @@ -229,6 +319,12 @@ struct FileStorageConfig { // Use io_uring for file I/O instead of POSIX pread/pwrite bool use_uring = false; + // Proactively evict local disk objects from the heartbeat thread once + // backend usage crosses the high watermark. + bool enable_disk_watermark_eviction = true; + double disk_eviction_high_watermark_ratio = 0.90; + double disk_eviction_low_watermark_ratio = 0.80; + // Validates the configuration for correctness and consistency bool Validate() const; @@ -250,6 +346,9 @@ class StorageBackendInterface { public: StorageBackendInterface(const FileStorageConfig& file_storage_config); + using EvictionHandler = std::function( + const std::vector& evicted_keys)>; + virtual tl::expected Init() = 0; virtual tl::expected BatchOffload( @@ -257,8 +356,7 @@ class StorageBackendInterface { std::function& keys, std::vector& metadatas)> complete_handler, - std::function& evicted_keys)> - eviction_handler = nullptr) = 0; + EvictionHandler eviction_handler = nullptr) = 0; virtual tl::expected BatchLoad( std::unordered_map& batched_slices) = 0; @@ -286,6 +384,17 @@ class StorageBackendInterface { // Default: no-op (no test failures injected) } + // Remove all persisted objects from disk. Called during RemoveAll to + // clean up physical SSD files alongside master metadata deletion. + virtual void RemoveAll() {} + + virtual tl::expected, ErrorCode> + EvictAboveDiskWatermark(double /* high_watermark_ratio */, + double /* low_watermark_ratio */, + EvictionHandler /* eviction_handler */ = nullptr) { + return std::vector{}; + } + FileStorageConfig file_storage_config_; }; @@ -318,28 +427,39 @@ class StorageBackend { * @param fsdir subdirectory name * @param enable_eviction Whether to enable disk eviction feature (default: * true) Note: Eviction is controlled by the enable_eviction parameter - * @return shared_ptr to new instance or nullptr if directory is invalid + * @return shared_ptr to new instance, or INVALID_PARAMS if the + * configuration is invalid * * Performs validation of the root directory before creating the instance: * - Verifies directory exists * - Verifies path is actually a directory + * - Verifies fsdir is not empty */ - static std::shared_ptr Create(const std::string& root_dir, - const std::string& fsdir, - bool enable_eviction = true) { + static tl::expected, ErrorCode> Create( + const std::string& root_dir, const std::string& fsdir, + bool enable_eviction = true) { namespace fs = std::filesystem; - if (!fs::exists(root_dir)) { - LOG(INFO) << "Root directory does not exist: " << root_dir; - return nullptr; - } else if (!fs::is_directory(root_dir)) { - LOG(INFO) << "Root path is not a directory: " << root_dir; - return nullptr; - } else if (fsdir.empty()) { - LOG(INFO) << "FSDIR cannot be empty"; - return nullptr; + std::error_code ec; + const auto root_status = fs::status(root_dir, ec); + if (ec) { + LOG(ERROR) << "Failed to access root directory: " << root_dir + << " (error: " << ec.message() << ")"; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + if (!fs::exists(root_status)) { + LOG(ERROR) << "Root directory does not exist: " << root_dir + << ". Please create it first or fix the configured " + "storage root directory."; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + if (!fs::is_directory(root_status)) { + LOG(ERROR) << "Root path is not a directory: " << root_dir; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + if (fsdir.empty()) { + LOG(ERROR) << "FSDIR cannot be empty"; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } - - fs::path root_path(root_dir); std::string real_fsdir = "moon_" + fsdir; return std::make_shared(root_dir, real_fsdir, @@ -384,7 +504,8 @@ class StorageBackend { */ tl::expected, ErrorCode> StoreObject( const std::string& path, const std::vector& slices, - const std::string& key = ""); + const std::string& key = "", + StorageBackendInterface::EvictionHandler eviction_handler = nullptr); /** * @brief Stores an object from a string @@ -395,7 +516,8 @@ class StorageBackend { */ tl::expected, ErrorCode> StoreObject( const std::string& path, const std::string& str, - const std::string& key = ""); + const std::string& key = "", + StorageBackendInterface::EvictionHandler eviction_handler = nullptr); /** * @brief Stores an object from a span of data @@ -406,7 +528,14 @@ class StorageBackend { */ tl::expected, ErrorCode> StoreObject( const std::string& path, std::span data, - const std::string& key = ""); + const std::string& key = "", + StorageBackendInterface::EvictionHandler eviction_handler = nullptr); + + tl::expected, ErrorCode> EvictAboveDiskWatermark( + double high_watermark_ratio, double low_watermark_ratio, + StorageBackendInterface::EvictionHandler eviction_handler = nullptr); + + void UpdateFileRecordKey(const std::string& path, const std::string& key); /** * @brief Loads an object into slices @@ -464,8 +593,11 @@ class StorageBackend { std::list file_write_queue_; std::unordered_map::iterator> file_queue_map_; + std::unordered_set pending_eviction_paths_; mutable std::shared_mutex file_queue_mutex_; // Mutex to protect file queue operations + static constexpr size_t kFilePathLockCount = 64; + std::array file_path_mutexes_; // Storage space tracking variables mutable std::shared_mutex @@ -497,6 +629,12 @@ class StorageBackend { */ FileRecord EvictFile(); + FileRecord PopFileToEvictByFIFO(); + + void RestoreFileToWriteQueueFront(const FileRecord& record); + + tl::expected DeleteEvictedFile(const FileRecord& record); + /** * @brief Add file to write queue for FIFO tracking * @param path Path of the file to add to queue @@ -518,13 +656,6 @@ class StorageBackend { */ bool CheckDiskSpace(size_t required_size); - /** - * @brief Select a file to evict based on FIFO order (earliest written - * first) - * @return The file to evict, or empty structure if no file found - */ - FileRecord SelectFileToEvictByFIFO(); - /** * @brief Ensures that a specified amount of disk space is available, * performing evictions if necessary. @@ -536,7 +667,8 @@ class StorageBackend { * attempting evictions up to the maximum attempt limit. */ tl::expected, ErrorCode> EnsureDiskSpace( - size_t required_size); + size_t required_size, + StorageBackendInterface::EvictionHandler eviction_handler = nullptr); /** * @brief Releases a specified amount of disk space and updates internal @@ -564,6 +696,10 @@ class StorageBackend { */ bool IsEvictionEnabled() const; + Mutex& GetFilePathMutex(const std::string& path); + + bool IsFilePendingEviction(const std::string& path) const; + /** * @brief Helper: Creates a file for writing and handles errors * @param path File path @@ -628,8 +764,7 @@ class StorageBackendAdaptor : public StorageBackendInterface { std::function& keys, std::vector& metadatas)> complete_handler, - std::function& evicted_keys)> - eviction_handler = nullptr) override; + EvictionHandler eviction_handler = nullptr) override; tl::expected BatchLoad( std::unordered_map& batched_slices) override; @@ -651,6 +786,12 @@ class StorageBackendAdaptor : public StorageBackendInterface { test_failure_predicate_ = std::move(predicate); } + void RemoveAll() override; + + tl::expected, ErrorCode> EvictAboveDiskWatermark( + double high_watermark_ratio, double low_watermark_ratio, + EvictionHandler eviction_handler = nullptr) override; + private: const FilePerKeyConfig file_per_key_config_; @@ -703,8 +844,7 @@ class BucketStorageBackend : public StorageBackendInterface { std::function& keys, std::vector& metadatas)> complete_handler, - std::function& evicted_keys)> - eviction_handler = nullptr) override; + EvictionHandler eviction_handler = nullptr) override; /** * @brief Retrieves metadata for multiple objects in a single batch @@ -775,6 +915,8 @@ class BucketStorageBackend : public StorageBackendInterface { */ tl::expected IsEnableOffloading() override; + void RemoveAll() override; + /** * @brief 根据后端 bucket 限制(keys/size)将 offloading_objects 分桶。 * @param offloading_objects Input map of object keys and their sizes @@ -835,6 +977,10 @@ class BucketStorageBackend : public StorageBackendInterface { */ tl::expected DeleteBucket(int64_t bucket_id); + tl::expected, ErrorCode> EvictAboveDiskWatermark( + double high_watermark_ratio, double low_watermark_ratio, + EvictionHandler eviction_handler = nullptr) override; + private: tl::expected, ErrorCode> BuildBucket( int64_t bucket_id, @@ -874,20 +1020,38 @@ class BucketStorageBackend : public StorageBackendInterface { tl::expected HasNext(); /** - * @brief Cleanup orphaned bucket files (data + metadata) for a given bucket - * ID. Called when BatchOffload fails due to duplicate keys after files were - * written. + * @brief Remove any remaining data and metadata files for a bucket. + * Used by write rollback and startup recovery of incomplete buckets. * @param bucket_id The bucket ID whose files should be deleted. */ void CleanupOrphanedBucket(int64_t bucket_id); + /** + * @brief Rollback a committed bucket from the local index when + * NotifyOffloadSuccess fails after local commit. Removes keys from + * object_bucket_map_, removes the bucket from buckets_ and lru_index_, + * waits for inflight reads to drain, then cleans up on-disk files. + * + * Called from BatchOffload when complete_handler fails after the local + * index has already been committed. + * + * @param bucket_id The bucket ID to roll back. + * @param keys The keys that were committed. + */ + void RollbackCommittedBucket(int64_t bucket_id, + const std::vector& keys); + // Holds eviction state between PrepareEviction and FinalizeEviction. // PrepareEviction removes buckets from metadata maps and returns this. - // FinalizeEviction waits for in-flight reads and deletes the files. + // FinalizeEviction removes persisted metadata, waits for in-flight reads, + // and then deletes the data files. struct PendingEviction { std::vector keys; // All keys in evicted buckets std::vector>> buckets; // (bucket_id, metadata) for file deletion + std::vector write_keys; + int64_t evicted_size = 0; + int64_t write_size = 0; }; /** @@ -898,7 +1062,18 @@ class BucketStorageBackend : public StorageBackendInterface { * @param required_size Size of the incoming bucket to be written. * @return PendingEviction with all keys and bucket metadata removed. */ - PendingEviction PrepareEviction(int64_t required_size); + tl::expected PrepareEviction( + int64_t required_size, const std::vector& write_keys = {}); + + void RestorePreparedEviction(PendingEviction&& pending); + + void RestorePreparedEvictionLocked(PendingEviction&& pending); + + void CommitPreparedEviction(const PendingEviction& pending); + + void ReleasePreparedWrite(const PendingEviction& pending); + + void ReleasePreparedWriteLocked(const PendingEviction& pending); /** * @brief Select the next bucket to evict according to the configured @@ -910,12 +1085,16 @@ class BucketStorageBackend : public StorageBackendInterface { SelectEvictionCandidate(); /** - * @brief Phase 2 of eviction: wait for in-flight reads on each evicted - * bucket to drain, then delete the data and metadata files. + * @brief Phase 2 of eviction: delete persisted metadata for each evicted + * bucket, wait for in-flight reads to drain, then delete the data file. + * When metadata removal succeeds, doing it first prevents a later read + * timeout or data-file deletion failure from leaving a bucket that Init() + * could recover. * Must be called AFTER master has been notified via eviction_handler. * @param pending The result of a prior PrepareEviction call. */ - void FinalizeEviction(const PendingEviction& pending); + tl::expected FinalizeEviction( + const PendingEviction& pending); public: /** @@ -963,6 +1142,10 @@ class BucketStorageBackend : public StorageBackendInterface { int64_t total_size_ GUARDED_BY(mutex_) = 0; std::unordered_map GUARDED_BY(mutex_) object_bucket_map_; + std::unordered_set GUARDED_BY(mutex_) pending_eviction_keys_; + std::unordered_set GUARDED_BY(mutex_) pending_write_keys_; + int64_t pending_eviction_size_ GUARDED_BY(mutex_) = 0; + int64_t pending_write_size_ GUARDED_BY(mutex_) = 0; std::map> GUARDED_BY( mutex_) buckets_; // LRU eviction index: ordered set of {last_access_ns_, bucket_id}. @@ -992,7 +1175,17 @@ class BucketStorageBackend : public StorageBackendInterface { class OffsetAllocatorStorageBackend : public StorageBackendInterface { public: OffsetAllocatorStorageBackend( - const FileStorageConfig& file_storage_config_); + const FileStorageConfig& file_storage_config_, + const OffsetAllocatorBackendConfig& offset_backend_config = {}); + + ~OffsetAllocatorStorageBackend(); + OffsetAllocatorStorageBackend(OffsetAllocatorStorageBackend&&) = default; + OffsetAllocatorStorageBackend& operator=(OffsetAllocatorStorageBackend&&) = + default; + OffsetAllocatorStorageBackend(const OffsetAllocatorStorageBackend&) = + delete; + OffsetAllocatorStorageBackend& operator=( + const OffsetAllocatorStorageBackend&) = delete; /** * @brief Initializes the offset allocator storage backend. @@ -1015,8 +1208,7 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { std::function& keys, std::vector& metadatas)> complete_handler, - std::function& evicted_keys)> - eviction_handler = nullptr) override; + EvictionHandler eviction_handler = nullptr) override; /** * @brief Loads data for multiple objects in a batch operation. @@ -1062,19 +1254,127 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { test_failure_predicate_ = std::move(predicate); } - private: - // On-disk record header: [u32 key_len][u32 value_len] (8 bytes total) + // Returns the number of keys skipped after fallback eviction + // could not make enough room (fragmentation, extents pinned by + // in-flight reads, or allocator node exhaustion). Monotonically + // increasing; useful for distinguishing "watermark working" from + // "thrashing but unable to free space". + int64_t GetEvictionSkips() const { + return eviction_skips_.load(std::memory_order_relaxed); + } + + void RemoveAll() override; + + // On-disk record layout v3 (single definition, shared by the write, + // read and recovery paths of this backend, and by future DMA writers + // such as GDS): + // + // [u32 key_len][u32 value_len][u64 seq][u32 flags][u32 crc32] + // [key bytes][zero padding][value bytes] + // + // The value region always starts at a kValueAlignment boundary within + // the record so that DMA engines (e.g. cuFile) operate on aligned file + // offsets. The padding is a pure function of key_len, so writer, + // reader and recovery derive the same layout independently. + // + // `seq` is the write's insert_seq_ stamp; on recovery any record with + // seq >= the checkpoint's insert_seq was written after that checkpoint + // and is dropped (its extent may hold a torn write). + // + // `flags` bit kFlagHasCrc: when set, `crc32` is a CRC-32C over the + // header prefix (everything before crc32), the key and the value, + // verified once on recovery so torn/stale records are detected and + // skipped instead of being served as valid data. When clear (records + // whose value never touched the CPU, or CRC disabled via config), + // recovery skips the checksum and trusts checkpoint ordering alone. struct RecordHeader { // Length of key in bytes uint32_t key_len; - // Length of value in bytes + // Length of value in bytes. Currently assumes max object size is + // 4GB. If we need to support larger objects, change this to 8 bytes. uint32_t value_len; - // Header size: 8 bytes (2 * uint32_t). Currently assumes max object - // size is 4GB. If we need to support larger objects, change this to 16 - // bytes. - static constexpr size_t SIZE = sizeof(uint32_t) * 2; + // insert_seq_ stamp of this write (monotonic per BatchOffload entry) + uint64_t seq; + + // Record flags; see kFlag* constants below. + uint32_t flags; + + // CRC-32C over [key_len|value_len|seq|flags] + key + value. + // Valid only when (flags & kFlagHasCrc). + uint32_t crc32; + + // flags: crc32 field carries a valid CRC-32C of this record. + static constexpr uint32_t kFlagHasCrc = 1u << 0; + + // All currently defined flag bits; recovery drops records with + // unknown bits set (written by a newer format). + static constexpr uint32_t kKnownFlags = kFlagHasCrc; + + // File-offset alignment of the value region within a record. + // 4 KiB covers the logical block size of currently supported NVMe + // devices (cuFile requirement for DMA). + static constexpr uint32_t kValueAlignment = 4096; + + // Header size: 24 bytes on disk (fields are (de)serialized + // field-by-field; do NOT use sizeof(RecordHeader), which includes + // padding). + static constexpr size_t SIZE = + sizeof(uint32_t) * 2 + sizeof(uint64_t) + sizeof(uint32_t) * 2; + + // Size of the crc-covered header prefix (everything before crc32). + static constexpr size_t PREFIX_SIZE = + sizeof(uint32_t) * 2 + sizeof(uint64_t) + sizeof(uint32_t); + + // Zero-padding between key and value for the given key length. + static constexpr uint32_t ValuePadding(uint32_t key_len) { + const uint64_t head = SIZE + key_len; + return static_cast( + (kValueAlignment - head % kValueAlignment) % kValueAlignment); + } + + // Offset of the value region relative to the record start. + static constexpr uint64_t ValueOffsetInRecord(uint32_t key_len) { + return SIZE + key_len + ValuePadding(key_len); + } + + // Total on-disk record size including padding. + static constexpr uint64_t RecordSize(uint32_t key_len, + uint32_t value_len) { + return ValueOffsetInRecord(key_len) + value_len; + } + + bool HasCrc() const { return (flags & kFlagHasCrc) != 0; } + + void WritePrefixTo(char* out) const { + std::memcpy(out, &key_len, sizeof(key_len)); + std::memcpy(out + sizeof(key_len), &value_len, sizeof(value_len)); + std::memcpy(out + sizeof(key_len) + sizeof(value_len), &seq, + sizeof(seq)); + std::memcpy(out + sizeof(key_len) + sizeof(value_len) + sizeof(seq), + &flags, sizeof(flags)); + } + + void WriteTo(char* out) const { + WritePrefixTo(out); + std::memcpy(out + PREFIX_SIZE, &crc32, sizeof(crc32)); + } + + static RecordHeader ReadFrom(const char* buf) { + RecordHeader h{}; + size_t off = 0; + std::memcpy(&h.key_len, buf + off, sizeof(h.key_len)); + off += sizeof(h.key_len); + std::memcpy(&h.value_len, buf + off, sizeof(h.value_len)); + off += sizeof(h.value_len); + std::memcpy(&h.seq, buf + off, sizeof(h.seq)); + off += sizeof(h.seq); + std::memcpy(&h.flags, buf + off, sizeof(h.flags)); + off += sizeof(h.flags); + std::memcpy(&h.crc32, buf + off, sizeof(h.crc32)); + return h; + } // Validate header against expected metadata bool ValidateAgainstMetadata(uint32_t expected_value_len) const { @@ -1099,6 +1399,12 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { } }; + private: + // Maximum key length accepted by BatchOffload and trusted on recovery. + // Write-side enforcement (BatchOffload) and recovery-side validation + // (RebuildShardMapsFromAllocator) must agree on this bound. + static constexpr uint32_t kMaxKeyLen = 1024 * 1024; + // Refcounted wrapper for move-only OffsetAllocationHandle. Physical extent // freed when last shared_ptr reference drops. struct RefCountedAllocationHandle { @@ -1124,7 +1430,8 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { // Byte offset in data file where record is stored uint64_t offset; - // Total record size: header (8) + key + value + // Total record size: header + key + padding + value + // (see RecordHeader::RecordSize) uint32_t total_size; // Value size only (excluding header and key) @@ -1132,12 +1439,25 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { // Refcounted handle keeps physical extent alive during reads AllocationPtr allocation; + + // Monotonic insertion sequence number. Points back to the slot in + // fifo_index_ (seq -> key). Used during eviction to detect stale + // index entries (lazy-repair) and to remove old slots on overwrite. + uint64_t fifo_seq = 0; + ObjectEntry(uint64_t off, uint32_t total, uint32_t val, - AllocationPtr alloc_ptr) + AllocationPtr alloc_ptr, uint64_t seq = 0) : offset(off), total_size(total), value_size(val), - allocation(std::move(alloc_ptr)) {} + allocation(std::move(alloc_ptr)), + fifo_seq(seq) {} + }; + + // Keeps evicted metadata and allocation handles alive until the master + // accepts the replica-removal notification. + struct PendingEviction { + std::vector> objects; }; // Returns full path to data file: {storage_path_}/kv_cache.data @@ -1186,8 +1506,12 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { // Thread-safe allocator managing free space within [0, capacity_) range std::shared_ptr allocator_; - // File handle wrapper for I/O operations using preadv/pwritev - std::unique_ptr data_file_; + // File handle wrapper for I/O operations using preadv/pwritev. Held as a + // shared_ptr so that an in-flight BatchLoad (which copies it into its + // ReadPlan under the shard lock) keeps the old file alive while RemoveAll + // rebinds this member to a freshly rebuilt file. Avoids use-after-free + // when RemoveAll runs concurrently with a reader on another thread. + std::shared_ptr data_file_; // Sharded metadata maps: one map per shard with its own lock (prevents data // races) @@ -1201,12 +1525,131 @@ class OffsetAllocatorStorageBackend : public StorageBackendInterface { // counting) std::atomic total_keys_{0}; + // ===== Eviction-related members ===== + OffsetAllocatorBackendConfig cfg_; + + // Counter for keys skipped due to fallback eviction exhaustion. + // See GetEvictionSkips() for the public accessor. + std::atomic eviction_skips_{0}; + + // Mutex protecting fifo_index_ and insert_seq_. Must be acquired BEFORE + // any shard mutex (shards_[i].mutex) when both are held. + mutable Mutex eviction_mutex_; + + // Global FIFO index: insertion sequence number -> key. + // begin() = oldest key, the default eviction victim. + // Entries allowed to be stale; lazy-repair at eviction time. + std::map fifo_index_; + + // Monotonic sequence number source for fifo_index_. + std::atomic insert_seq_{0}; + + // Resolved watermark thresholds (bytes), computed in Init(). + int64_t high_watermark_bytes_ = 0; + int64_t low_watermark_bytes_ = 0; + + // Resolved watermark thresholds (key count), computed in Init(). + int64_t high_watermark_keys_ = 0; + int64_t low_watermark_keys_ = 0; + + // Evict keys from the FIFO index until both byte and key-count watermarks + // are satisfied (or until the eviction cap is reached). Allocations remain + // pinned in out_pending until notification succeeds. + void EvictToMakeRoom(int64_t required_bytes, size_t min_victims, + const std::unordered_set& batch_keys, + PendingEviction& out_pending); + + // Restore prepared victims when the master rejects their removal. + void RestorePreparedEviction(PendingEviction&& pending); + + // Notify the master, then release prepared allocations on success or + // restore their metadata on failure. + tl::expected NotifyAndCommitPreparedEviction( + const EvictionHandler& eviction_handler, PendingEviction& pending); + + // Record restart-persistence tombstones for prepared victims whose + // eviction has become final. No-op when persistence is disabled. + void RecordEvictionTombstones(const PendingEviction& pending); + + // ===== Persistence methods ===== + + std::string GetMetaFilePath() const; + + bool ShouldPersistNow() const; + + tl::expected SaveMetadata( + const std::unordered_set& evicted_keys_this_batch); + + tl::expected LoadMetadata(); + + // Outcome of a recovery attempt. Distinguishes "safe to start fresh" + // (kNoMeta / kCorrupt) from "must not touch the persisted data" + // (kTransientError, e.g. fd exhaustion or OOM — retrying Init later may + // succeed, while a fresh start would destroy a recoverable cache). + enum class RecoveryResult { + kRecovered, // persisted state fully restored + kNoMeta, // no metadata file: genuine first boot + kCorrupt, // meta/data missing, incompatible or corrupt + kTransientError, // temporary resource error: do NOT wipe data + }; + + RecoveryResult TryRecoverFromMetadata(); + + // Rebuilds shard maps by scanning extents marked used in the + // deserialized allocator. Records stamped with + // seq >= checkpoint_insert_seq were written after the checkpoint and + // are dropped (their extents may hold torn writes). + void RebuildShardMapsFromAllocator(uint64_t checkpoint_insert_seq); + + void RestoreAndRepairFifoIndex( + const OffsetAllocatorPersistedMetadata& meta); + // Test-only: Predicate to determine which keys should fail in BatchOffload. // Used for deterministic testing of partial success behavior. std::function test_failure_predicate_; + + private: + // ---- Persistence state ---- + std::atomic last_persist_time_us_{0}; + std::unordered_set all_evicted_this_batch_; + std::atomic metadata_dirty_{false}; + + // ---- Persistence metrics ---- + std::atomic last_save_metadata_cost_us_{0}; + std::atomic metadata_save_failures_{0}; + std::atomic metadata_load_fallbacks_{0}; + std::atomic metadata_consecutive_failures_{0}; + + // ---- Test-only hooks ---- + std::atomic test_metadata_write_failure_step_{0}; + // Skip the destructor's final checkpoint (simulates an abrupt crash). + std::atomic test_skip_final_checkpoint_{false}; + + public: + // ---- Test accessors ---- + void SetMetadataWriteFailure(int step) { + test_metadata_write_failure_step_ = step; + } + void SetSkipFinalCheckpointForTest() { + test_skip_final_checkpoint_.store(true, std::memory_order_relaxed); + } + size_t GetAllEvictedThisBatchSizeForTest() const { + return all_evicted_this_batch_.size(); + } + int64_t GetMetadataSaveFailures() const { + return metadata_save_failures_.load(std::memory_order_relaxed); + } + int64_t GetMetadataLoadFallbacks() const { + return metadata_load_fallbacks_.load(std::memory_order_relaxed); + } + int64_t GetMetadataConsecutiveFailures() const { + return metadata_consecutive_failures_.load(std::memory_order_relaxed); + } + + private: }; tl::expected, ErrorCode> CreateStorageBackend(const FileStorageConfig& config); -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/include/store_rpc_client_io_context.h b/mooncake-store/include/store_rpc_client_io_context.h new file mode 100644 index 0000000000..bbe5c39eed --- /dev/null +++ b/mooncake-store/include/store_rpc_client_io_context.h @@ -0,0 +1,23 @@ +#pragma once + +#include "environ.h" +#include "rpc_client_io_context.h" + +namespace mooncake { + +inline uint32_t GetStoreRpcClientIoThreads() { + return Environ::Get().GetStoreRpcClientIoThreads(); +} + +namespace detail { +struct StoreRpcClientIoContextPoolTag {}; +} // namespace detail + +inline coro_io::io_context_pool& GetStoreRpcClientIoContextPool() { + static auto& io_pool = + GetRpcClientIoContextPool( + GetStoreRpcClientIoThreads()); + return io_pool; +} + +} // namespace mooncake diff --git a/mooncake-store/include/tenant_id.h b/mooncake-store/include/tenant_id.h new file mode 100644 index 0000000000..27e639a9b5 --- /dev/null +++ b/mooncake-store/include/tenant_id.h @@ -0,0 +1,87 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace mooncake { + +class TenantId final { + public: + static constexpr std::string_view kDefaultValue = "default"; + + TenantId() : value_(kDefaultValue) {} + + explicit TenantId(std::string raw) + : value_(raw.empty() ? std::string(kDefaultValue) : std::move(raw)) {} + + static const TenantId& Default() { + static const TenantId kDefaultTenant; + return kDefaultTenant; + } + + const std::string& value() const noexcept { return value_; } + + bool IsDefault() const noexcept { return value_ == kDefaultValue; } + + bool IsValid() const noexcept { + if (value_.empty() || value_.front() == kReservedPrefix) { + return false; + } + for (unsigned char c : value_) { + if (c < kFirstPrintableAscii || c == kDeleteAscii) { + return false; + } + } + return true; + } + + std::string MakeScopedKey(std::string_view local_key) const { + std::string scoped_key; + scoped_key.reserve(value_.size() + local_key.size() + 1); + scoped_key.append(value_); + scoped_key.push_back(kScopedKeySeparator); + scoped_key.append(local_key); + return scoped_key; + } + + static std::pair ParseScopedKey( + std::string_view scoped_key) { + const auto separator = scoped_key.find(kScopedKeySeparator); + if (separator == std::string_view::npos) { + return {TenantId::Default(), std::string(scoped_key)}; + } + return {TenantId(std::string(scoped_key.substr(0, separator))), + std::string(scoped_key.substr(separator + 1))}; + } + + friend bool operator==(const TenantId&, const TenantId&) = default; + + friend bool operator<(const TenantId& lhs, const TenantId& rhs) noexcept { + return lhs.value_ < rhs.value_; + } + + friend std::ostream& operator<<(std::ostream& os, + const TenantId& tenant_id) { + return os << tenant_id.value_; + } + + private: + static constexpr char kReservedPrefix = '_'; + static constexpr unsigned char kFirstPrintableAscii = 0x20; + static constexpr unsigned char kDeleteAscii = 0x7f; + static constexpr char kScopedKeySeparator = '\0'; + + std::string value_; +}; + +struct TenantIdHash { + size_t operator()(const TenantId& tenant_id) const noexcept { + return std::hash{}(tenant_id.value()); + } +}; + +} // namespace mooncake diff --git a/mooncake-store/include/tenant_quota.h b/mooncake-store/include/tenant_quota.h index 301e92d6c3..6621e54974 100644 --- a/mooncake-store/include/tenant_quota.h +++ b/mooncake-store/include/tenant_quota.h @@ -1,73 +1,118 @@ #pragma once +#include #include #include #include #include +#include #include +#include "tenant_id.h" + #include namespace mooncake { -struct TenantQuotaState { +struct TenantQuotaSnapshot { + TenantId tenant_id; uint64_t requested_quota_bytes = 0; uint64_t effective_quota_bytes = 0; uint64_t used_bytes = 0; uint64_t reserved_bytes = 0; uint64_t committed_count = 0; + uint64_t metadata_object_count = 0; bool has_explicit_policy = false; bool over_quota = false; }; -struct TenantQuotaSnapshot { - std::string tenant_id; - uint64_t requested_quota_bytes = 0; - uint64_t effective_quota_bytes = 0; +struct TenantQuotaUsage { uint64_t used_bytes = 0; - uint64_t reserved_bytes = 0; uint64_t committed_count = 0; - bool has_explicit_policy = false; - bool over_quota = false; + uint64_t metadata_object_count = 0; }; +using TenantQuotaPolicyMap = std::map; +using TenantQuotaUsageMap = + std::unordered_map; + enum class TenantQuotaError { kQuotaExceeded, kInvalidArgument, kAccountingMismatch, + kTenantNotRegistered, + kTenantNotEmpty, + kTenantNotFound, }; using TenantQuotaResult = tl::expected; +using TenantQuotaPolicyResult = tl::expected; + +template +class ShardedTenantQuotaTable; +// Single-threaded tenant quota state machine. This class owns quota rules and +// accounting invariants, but deliberately contains no locking or sharding. class TenantQuotaTable { public: - void SetDefaultRequestedQuota(uint64_t bytes); - uint64_t GetDefaultRequestedQuota() const; - - TenantQuotaResult UpsertTenantPolicy(std::string tenant_id, + TenantQuotaResult UpsertTenantPolicy(const TenantId& tenant_id, uint64_t requested_quota_bytes); - void EraseTenantPolicy(std::string tenant_id); + TenantQuotaPolicyResult DisableTenantPolicyIfEmpty( + const TenantId& tenant_id); + void ApplyTenantPolicies(const TenantQuotaPolicyMap& policies); + TenantQuotaPolicyMap GetTenantPolicies() const; void RecomputeEffectiveQuotas(uint64_t allocatable_capacity_bytes); + bool IsTenantRegistered(const TenantId& tenant_id) const; std::optional GetTenantSnapshot( - std::string tenant_id) const; + const TenantId& tenant_id) const; std::vector ListTenantSnapshots() const; + uint64_t ComputeDeficit(const TenantId& tenant_id, + uint64_t incoming_bytes) const; + + TenantQuotaResult Reserve(const TenantId& tenant_id, uint64_t bytes); + TenantQuotaResult Commit(const TenantId& tenant_id, uint64_t bytes); + TenantQuotaResult CommitAdditional(const TenantId& tenant_id, + uint64_t bytes); + TenantQuotaResult Abort(const TenantId& tenant_id, uint64_t bytes); + TenantQuotaResult Release(const TenantId& tenant_id, uint64_t bytes); + TenantQuotaResult ReleasePartial(const TenantId& tenant_id, uint64_t bytes); - TenantQuotaResult Reserve(std::string tenant_id, uint64_t bytes); - TenantQuotaResult Commit(std::string tenant_id, uint64_t bytes); - TenantQuotaResult Abort(std::string tenant_id, uint64_t bytes); - TenantQuotaResult Release(std::string tenant_id, uint64_t bytes); - TenantQuotaResult ReleasePartial(std::string tenant_id, uint64_t bytes); + void IncrementMetadataObjectCount(const TenantId& tenant_id); + TenantQuotaResult DecrementMetadataObjectCount(const TenantId& tenant_id); + void RebuildUsage(const TenantQuotaUsageMap& usage); private: - TenantQuotaState& GetOrCreateState(const std::string& tenant_id); - TenantQuotaSnapshot MakeSnapshot(const std::string& tenant_id, + template + friend class ShardedTenantQuotaTable; + + struct TenantQuotaState { + uint64_t requested_quota_bytes = 0; + uint64_t effective_quota_bytes = 0; + uint64_t used_bytes = 0; + uint64_t reserved_bytes = 0; + uint64_t committed_count = 0; + uint64_t metadata_object_count = 0; + bool has_explicit_policy = false; + bool over_quota = false; + }; + + using StateMap = std::map; + + TenantQuotaState& GetOrCreateState(const TenantId& tenant_id); + TenantQuotaSnapshot MakeSnapshot(const TenantId& tenant_id, const TenantQuotaState& state) const; - void RefreshOverQuota(TenantQuotaState* state) const; + static bool IsLazyEmptyTenant(const TenantQuotaState& state); + static void RefreshOverQuota(TenantQuotaState* state); + static std::map BuildEffectiveQuotaAssignments( + const std::vector& tenants, + uint64_t allocatable_capacity_bytes); + void ApplyEffectiveQuotas( + const std::map& effective_quotas); + void EraseIfLazyEmpty(StateMap::iterator it); - uint64_t default_requested_quota_bytes_ = 0; - std::map tenants_; + StateMap tenants_; }; } // namespace mooncake diff --git a/mooncake-store/include/tenant_quota_policy_store.h b/mooncake-store/include/tenant_quota_policy_store.h new file mode 100644 index 0000000000..ad6d2ae9c0 --- /dev/null +++ b/mooncake-store/include/tenant_quota_policy_store.h @@ -0,0 +1,68 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include + +namespace mooncake { + +struct TenantQuotaPolicySnapshot { + std::map tenant_quotas; +}; + +tl::expected ParseTenantQuotaBytes( + const std::string& value); + +tl::expected ParseTenantQuotaPolicyYaml( + const std::string& yaml); + +std::string FormatTenantQuotaPolicyYaml( + const TenantQuotaPolicySnapshot& snapshot); + +class TenantQuotaPolicyStore { + public: + virtual ~TenantQuotaPolicyStore() = default; + + virtual tl::expected Load() = 0; + virtual tl::expected Save( + const TenantQuotaPolicySnapshot& snapshot) = 0; +}; + +class YamlTenantQuotaPolicyStore final : public TenantQuotaPolicyStore { + public: + explicit YamlTenantQuotaPolicyStore(std::string path); + + tl::expected Load() override; + tl::expected Save( + const TenantQuotaPolicySnapshot& snapshot) override; + + private: + std::string path_; + std::mutex mutex_; +}; + +#ifdef STORE_USE_ETCD +class EtcdTenantQuotaPolicyStore final : public TenantQuotaPolicyStore { + public: + EtcdTenantQuotaPolicyStore(const std::string& endpoints, + const std::string& cluster_id); + + tl::expected Load() override; + tl::expected Save( + const TenantQuotaPolicySnapshot& snapshot) override; + + private: + std::string key_; + std::mutex mutex_; +}; +#endif + +tl::expected, std::string> +CreateTenantQuotaPolicyStore(const std::string& type, const std::string& uri, + const std::string& cluster_id); + +} // namespace mooncake diff --git a/mooncake-store/include/tenant_quota_sharded.h b/mooncake-store/include/tenant_quota_sharded.h new file mode 100644 index 0000000000..afe0e47e06 --- /dev/null +++ b/mooncake-store/include/tenant_quota_sharded.h @@ -0,0 +1,68 @@ +#pragma once + +#include +#include +#include + +#include "tenant_quota.h" + +namespace mooncake { + +// Thread-safe production wrapper around TenantQuotaTable. Per-tenant +// operations lock only one shard; cross-shard policy, usage, and recompute +// operations are serialized by recompute_mutex_. +template +class ShardedTenantQuotaTable { + public: + static_assert(NumShards > 0, "tenant quota table needs at least one shard"); + static constexpr size_t kNumShards = NumShards; + + TenantQuotaResult UpsertTenantPolicy(const TenantId& tenant_id, + uint64_t requested_quota_bytes, + uint64_t allocatable_capacity_bytes); + TenantQuotaPolicyResult DisableTenantPolicyIfEmpty( + const TenantId& tenant_id); + void ApplyTenantPolicies(const TenantQuotaPolicyMap& policies, + uint64_t allocatable_capacity_bytes); + TenantQuotaPolicyMap GetTenantPolicies() const; + + void RecomputeEffectiveQuotas(uint64_t allocatable_capacity_bytes); + + bool IsTenantRegistered(const TenantId& tenant_id) const; + std::optional GetTenantSnapshot( + const TenantId& tenant_id) const; + std::vector ListTenantSnapshots() const; + uint64_t ComputeDeficit(const TenantId& tenant_id, + uint64_t incoming_bytes) const; + + TenantQuotaResult Reserve(const TenantId& tenant_id, uint64_t bytes); + TenantQuotaResult Commit(const TenantId& tenant_id, uint64_t bytes); + TenantQuotaResult CommitAdditional(const TenantId& tenant_id, + uint64_t bytes); + TenantQuotaResult Abort(const TenantId& tenant_id, uint64_t bytes); + TenantQuotaResult Release(const TenantId& tenant_id, uint64_t bytes); + TenantQuotaResult ReleasePartial(const TenantId& tenant_id, uint64_t bytes); + + void IncrementMetadataObjectCount(const TenantId& tenant_id); + TenantQuotaResult DecrementMetadataObjectCount(const TenantId& tenant_id); + void RebuildUsage(const TenantQuotaUsageMap& usage, + uint64_t allocatable_capacity_bytes); + + private: + struct Shard { + mutable std::mutex mutex; + TenantQuotaTable table; + }; + + size_t GetShardIndex(const TenantId& tenant_id) const; + Shard& GetShard(const TenantId& tenant_id); + const Shard& GetShard(const TenantId& tenant_id) const; + void RecomputeEffectiveQuotasLocked(uint64_t allocatable_capacity_bytes); + + std::array shards_; + mutable std::mutex recompute_mutex_; +}; + +} // namespace mooncake + +#include "tenant_quota_sharded_impl.h" diff --git a/mooncake-store/include/tenant_quota_sharded_impl.h b/mooncake-store/include/tenant_quota_sharded_impl.h new file mode 100644 index 0000000000..d74c272a4d --- /dev/null +++ b/mooncake-store/include/tenant_quota_sharded_impl.h @@ -0,0 +1,243 @@ +#pragma once + +#include + +#include "tenant_quota_sharded.h" + +namespace mooncake { + +template +TenantQuotaResult ShardedTenantQuotaTable::UpsertTenantPolicy( + const TenantId& tenant_id, uint64_t requested_quota_bytes, + uint64_t allocatable_capacity_bytes) { + std::lock_guard recompute_lock(recompute_mutex_); + auto& shard = GetShard(tenant_id); + { + std::lock_guard lock(shard.mutex); + auto result = + shard.table.UpsertTenantPolicy(tenant_id, requested_quota_bytes); + if (!result) { + return result; + } + } + RecomputeEffectiveQuotasLocked(allocatable_capacity_bytes); + return {}; +} + +template +TenantQuotaPolicyResult +ShardedTenantQuotaTable::DisableTenantPolicyIfEmpty( + const TenantId& tenant_id) { + std::lock_guard recompute_lock(recompute_mutex_); + auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.DisableTenantPolicyIfEmpty(tenant_id); +} + +template +void ShardedTenantQuotaTable::ApplyTenantPolicies( + const TenantQuotaPolicyMap& policies, uint64_t allocatable_capacity_bytes) { + std::lock_guard recompute_lock(recompute_mutex_); + std::array grouped_policies; + for (const auto& [tenant_id, requested_quota_bytes] : policies) { + grouped_policies[GetShardIndex(tenant_id)].emplace( + tenant_id, requested_quota_bytes); + } + + for (size_t i = 0; i < kNumShards; ++i) { + auto& shard = shards_[i]; + std::lock_guard lock(shard.mutex); + shard.table.ApplyTenantPolicies(grouped_policies[i]); + } + RecomputeEffectiveQuotasLocked(allocatable_capacity_bytes); +} + +template +TenantQuotaPolicyMap ShardedTenantQuotaTable::GetTenantPolicies() + const { + TenantQuotaPolicyMap policies; + for (const auto& shard : shards_) { + std::lock_guard lock(shard.mutex); + auto shard_policies = shard.table.GetTenantPolicies(); + policies.insert(shard_policies.begin(), shard_policies.end()); + } + return policies; +} + +template +void ShardedTenantQuotaTable::RecomputeEffectiveQuotas( + uint64_t allocatable_capacity_bytes) { + std::lock_guard recompute_lock(recompute_mutex_); + RecomputeEffectiveQuotasLocked(allocatable_capacity_bytes); +} + +template +bool ShardedTenantQuotaTable::IsTenantRegistered( + const TenantId& tenant_id) const { + const auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.IsTenantRegistered(tenant_id); +} + +template +std::optional +ShardedTenantQuotaTable::GetTenantSnapshot( + const TenantId& tenant_id) const { + const auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.GetTenantSnapshot(tenant_id); +} + +template +std::vector +ShardedTenantQuotaTable::ListTenantSnapshots() const { + std::vector snapshots; + for (const auto& shard : shards_) { + std::lock_guard lock(shard.mutex); + auto shard_snapshots = shard.table.ListTenantSnapshots(); + snapshots.insert(snapshots.end(), shard_snapshots.begin(), + shard_snapshots.end()); + } + std::sort(snapshots.begin(), snapshots.end(), + [](const auto& lhs, const auto& rhs) { + return lhs.tenant_id < rhs.tenant_id; + }); + return snapshots; +} + +template +uint64_t ShardedTenantQuotaTable::ComputeDeficit( + const TenantId& tenant_id, uint64_t incoming_bytes) const { + const auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.ComputeDeficit(tenant_id, incoming_bytes); +} + +template +TenantQuotaResult ShardedTenantQuotaTable::Reserve( + const TenantId& tenant_id, uint64_t bytes) { + auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.Reserve(tenant_id, bytes); +} + +template +TenantQuotaResult ShardedTenantQuotaTable::Commit( + const TenantId& tenant_id, uint64_t bytes) { + auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.Commit(tenant_id, bytes); +} + +template +TenantQuotaResult ShardedTenantQuotaTable::CommitAdditional( + const TenantId& tenant_id, uint64_t bytes) { + auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.CommitAdditional(tenant_id, bytes); +} + +template +TenantQuotaResult ShardedTenantQuotaTable::Abort( + const TenantId& tenant_id, uint64_t bytes) { + auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.Abort(tenant_id, bytes); +} + +template +TenantQuotaResult ShardedTenantQuotaTable::Release( + const TenantId& tenant_id, uint64_t bytes) { + auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.Release(tenant_id, bytes); +} + +template +TenantQuotaResult ShardedTenantQuotaTable::ReleasePartial( + const TenantId& tenant_id, uint64_t bytes) { + auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.ReleasePartial(tenant_id, bytes); +} + +template +void ShardedTenantQuotaTable::IncrementMetadataObjectCount( + const TenantId& tenant_id) { + auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + shard.table.IncrementMetadataObjectCount(tenant_id); +} + +template +TenantQuotaResult +ShardedTenantQuotaTable::DecrementMetadataObjectCount( + const TenantId& tenant_id) { + auto& shard = GetShard(tenant_id); + std::lock_guard lock(shard.mutex); + return shard.table.DecrementMetadataObjectCount(tenant_id); +} + +template +void ShardedTenantQuotaTable::RebuildUsage( + const TenantQuotaUsageMap& usage, uint64_t allocatable_capacity_bytes) { + std::lock_guard recompute_lock(recompute_mutex_); + std::array grouped_usage; + for (const auto& [tenant_id, tenant_usage] : usage) { + grouped_usage[GetShardIndex(tenant_id)].emplace(tenant_id, + tenant_usage); + } + + for (size_t i = 0; i < kNumShards; ++i) { + auto& shard = shards_[i]; + std::lock_guard lock(shard.mutex); + shard.table.RebuildUsage(grouped_usage[i]); + } + RecomputeEffectiveQuotasLocked(allocatable_capacity_bytes); +} + +template +size_t ShardedTenantQuotaTable::GetShardIndex( + const TenantId& tenant_id) const { + return TenantIdHash{}(tenant_id) % kNumShards; +} + +template +typename ShardedTenantQuotaTable::Shard& +ShardedTenantQuotaTable::GetShard(const TenantId& tenant_id) { + return shards_[GetShardIndex(tenant_id)]; +} + +template +const typename ShardedTenantQuotaTable::Shard& +ShardedTenantQuotaTable::GetShard(const TenantId& tenant_id) const { + return shards_[GetShardIndex(tenant_id)]; +} + +template +void ShardedTenantQuotaTable::RecomputeEffectiveQuotasLocked( + uint64_t allocatable_capacity_bytes) { + std::vector snapshots; + for (const auto& shard : shards_) { + std::lock_guard lock(shard.mutex); + auto shard_snapshots = shard.table.ListTenantSnapshots(); + snapshots.insert(snapshots.end(), shard_snapshots.begin(), + shard_snapshots.end()); + } + + auto effective_quotas = TenantQuotaTable::BuildEffectiveQuotaAssignments( + snapshots, allocatable_capacity_bytes); + std::array, kNumShards> grouped_quotas; + for (const auto& [tenant_id, effective_quota_bytes] : effective_quotas) { + grouped_quotas[GetShardIndex(tenant_id)].emplace(tenant_id, + effective_quota_bytes); + } + + for (size_t i = 0; i < kNumShards; ++i) { + auto& shard = shards_[i]; + std::lock_guard lock(shard.mutex); + shard.table.ApplyEffectiveQuotas(grouped_quotas[i]); + } +} + +} // namespace mooncake diff --git a/mooncake-store/include/transfer_task.h b/mooncake-store/include/transfer_task.h index 63b061914d..5fe938e09e 100644 --- a/mooncake-store/include/transfer_task.h +++ b/mooncake-store/include/transfer_task.h @@ -560,6 +560,9 @@ class TransferSubmitter { const Replica::Descriptor& replica, std::vector& slices, uint64_t src_offset); + TransferEngine::ScatterTransferOperation submitScatter( + const std::vector& transfers); + std::optional submit_batch( const std::vector& replicas, std::vector>& all_slices, diff --git a/mooncake-store/include/types.h b/mooncake-store/include/types.h index 9c3f47f28a..c85eb8288c 100644 --- a/mooncake-store/include/types.h +++ b/mooncake-store/include/types.h @@ -11,6 +11,8 @@ #include #include +#include "tenant_id.h" + #include "Slab.h" #include "ylt/struct_json/json_reader.h" #include "ylt/struct_json/json_writer.h" @@ -83,14 +85,14 @@ inline bool IsValidClusterIdComponent(const std::string& cluster_id) { return true; } static constexpr uint64_t DEFAULT_DEFAULT_KV_LEASE_TTL = - 5000; // in milliseconds + 10000; // in milliseconds static constexpr uint64_t DEFAULT_KV_SOFT_PIN_TTL_MS = 30 * 60 * 1000; // 30 minutes static constexpr bool DEFAULT_ALLOW_EVICT_SOFT_PINNED_OBJECTS = true; static constexpr double DEFAULT_EVICTION_RATIO = 0.05; -static constexpr double DEFAULT_EVICTION_HIGH_WATERMARK_RATIO = 0.95; +static constexpr double DEFAULT_EVICTION_HIGH_WATERMARK_RATIO = 0.90; static constexpr double DEFAULT_NOF_EVICTION_RATIO = 0.05; -static constexpr double DEFAULT_NOF_EVICTION_HIGH_WATERMARK_RATIO = 0.95; +static constexpr double DEFAULT_NOF_EVICTION_HIGH_WATERMARK_RATIO = 0.90; static constexpr int64_t DEFAULT_MASTER_VIEW_LEASE_TTL_SEC = 5; // in seconds static constexpr int64_t DEFAULT_CLIENT_LIVE_TTL_SEC = 10; // in seconds static constexpr int64_t DEFAULT_NOF_HEARTBEAT_INTERVAL_SEC = 10; @@ -216,37 +218,16 @@ constexpr const char* CONFIG_KEY_RDMA_DEVICES = "rdma_devices"; constexpr const char* CONFIG_KEY_MASTER_SERVER_ADDR = "master_server_addr"; constexpr const char* CONFIG_KEY_IPC_SOCKET_PATH = "ipc_socket_path"; constexpr const char* CONFIG_KEY_TENANT_ID = "tenant_id"; +constexpr const char* CONFIG_KEY_ENABLE_CLIENT_HTTP_SERVER = + "enable_client_http_server"; +constexpr const char* CONFIG_KEY_CLIENT_HTTP_PORT = "client_http_port"; // Store client configuration defaults static constexpr size_t DEFAULT_GLOBAL_SEGMENT_SIZE = 1024 * 1024 * 16; // 16MB static constexpr size_t DEFAULT_LOCAL_BUFFER_SIZE = 1024 * 1024 * 16; // 16MB constexpr const char* DEFAULT_PROTOCOL = "tcp"; constexpr const char* DEFAULT_MASTER_SERVER_ADDR = "127.0.0.1:50051"; - -inline std::string NormalizeTenantId(const std::string& tenant_id) { - return tenant_id.empty() ? "default" : tenant_id; -} - -inline std::string MakeTenantScopedStorageKey(const std::string& tenant_id, - const std::string& key) { - const auto normalized_tenant = NormalizeTenantId(tenant_id); - std::string scoped_key; - scoped_key.reserve(normalized_tenant.size() + key.size() + 1); - scoped_key.append(normalized_tenant); - scoped_key.push_back('\0'); - scoped_key.append(key); - return scoped_key; -} - -inline std::pair ParseTenantScopedStorageKey( - const std::string& storage_key) { - const auto separator = storage_key.find('\0'); - if (separator == std::string::npos) { - return {"default", storage_key}; - } - return {NormalizeTenantId(storage_key.substr(0, separator)), - storage_key.substr(separator + 1)}; -} +static constexpr int DEFAULT_CLIENT_HTTP_PORT = 9300; struct OffloadTaskItem { std::string tenant_id; @@ -355,6 +336,8 @@ enum class ErrorCode : int32_t { // Transfer errors (Range: -800 to -899) TRANSFER_FAIL = -800, ///< Transfer operation failed. + /// Store checksum verification failed. + CHECKSUM_MISMATCH = -801, // RPC errors (Range: -900 to -999) RPC_FAIL = -900, ///< RPC operation failed. @@ -369,6 +352,10 @@ enum class ErrorCode : int32_t { -1004, ///< OpLog entry not found (backend-agnostic). K8S_LEASE_OPERATION_ERROR = -1005, ///< K8s Lease operation failed. K8S_LEASE_NOT_FOUND = -1006, ///< K8s Lease not found. + INCOMPLETE_OPLOG_CATCH_UP = + -1007, ///< Promotion catch-up could not prove all durable OpLog + ///< entries were applied, or unresolved skipped/missing + ///< gaps remain after final catch-up + second gap resolution. UNAVAILABLE_IN_CURRENT_STATUS = -1010, ///< Request cannot be done in current status. UNAVAILABLE_IN_CURRENT_MODE = @@ -408,6 +395,8 @@ enum class ErrorCode : int32_t { DFS_STALE_HANDLE = -1604, ///< DFS file handle expired. DFS_PARTIAL_WRITE = -1605, ///< DFS partial write success. TENANT_QUOTA_EXCEEDED = -1700, ///< Tenant memory quota exceeded. + TENANT_NOT_REGISTERED = -1701, ///< Tenant has no quota policy. + TENANT_NOT_EMPTY = -1702, ///< Tenant still owns objects or quota. }; int32_t toInt(ErrorCode errorCode) noexcept; @@ -454,17 +443,20 @@ struct Segment { // TE p2p endpoint (ip:port) for transport-only addressing std::string te_endpoint{}; std::string protocol; + std::string host_id{}; Segment() = default; }; -YLT_REFL(Segment, id, name, base, size, te_endpoint, protocol); +YLT_REFL(Segment, id, name, base, size, te_endpoint, protocol, host_id); /** * @brief Allocation strategy type for segment allocation */ enum class AllocationStrategyType { - RANDOM = 0, // Pure random allocation - FREE_RATIO_FIRST, // Free-ratio-first allocation - CXL, // CXL-specific allocation + RANDOM = 0, // Pure random allocation + FREE_RATIO_FIRST, // Free-ratio-first allocation + CXL, // CXL-specific allocation + SSD_FREE_RATIO_FIRST, // SSD free-ratio-first allocation + LOCAL_FIRST // Prefer local host before ordered remote fallback }; /** diff --git a/mooncake-store/include/uds_transport.h b/mooncake-store/include/uds_transport.h new file mode 100644 index 0000000000..a4bb5d42d3 --- /dev/null +++ b/mooncake-store/include/uds_transport.h @@ -0,0 +1,83 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace mooncake { + +// Lightweight Unix domain socket helper for one-shot, performance-insensitive +// control-plane exchanges such as local fd passing. This opens short-lived +// blocking sockets and handles one connection at a time; use asio-based +// transport instead for latency-sensitive, high-throughput, or long-lived I/O. +class UdsConnection { + public: + UdsConnection() = default; + explicit UdsConnection(int fd); + ~UdsConnection(); + + UdsConnection(const UdsConnection &) = delete; + UdsConnection &operator=(const UdsConnection &) = delete; + + UdsConnection(UdsConnection &&other) noexcept; + UdsConnection &operator=(UdsConnection &&other) noexcept; + + bool valid() const; + int fd() const; + int release(); + void close(); + tl::expected setRecvTimeout( + std::chrono::seconds timeout); + + int sendRaw(const void *data, size_t len); + int recvRaw(void *data, size_t len); + int sendFd(int fd, void *data, size_t data_len); + int recvFd(void *data, size_t data_len); + + private: + int fd_ = -1; +}; + +class UdsConnector { + public: + explicit UdsConnector( + std::string socket_name, + std::chrono::milliseconds connect_timeout = std::chrono::seconds(5)); + tl::expected, std::string> connect(); + + private: + std::string socket_name_; + std::chrono::milliseconds connect_timeout_; +}; + +class UdsAcceptor { + public: + using Handler = std::function; + + explicit UdsAcceptor(std::string socket_name); + ~UdsAcceptor(); + + void registerHandler(Handler handler); + tl::expected start(); + void stop(); + + private: + void acceptLoop(); + void wakeAccept(); + + std::string socket_name_; + Handler handler_; + int listen_fd_ = -1; + std::atomic active_client_fd_{-1}; + std::atomic running_{false}; + std::jthread thread_; +}; + +} // namespace mooncake diff --git a/mooncake-store/include/utils.h b/mooncake-store/include/utils.h index 160292c2e9..41d6e3d205 100644 --- a/mooncake-store/include/utils.h +++ b/mooncake-store/include/utils.h @@ -258,45 +258,6 @@ std::string expected_to_str(const tl::expected& expected) { return parsed.value_or(0); } -/** - * @brief Convert a boolean-like string to a bool - * @param str String representation ("1"/"true"/"yes"/"on" or - * "0"/"false"/"no"/"off") - * @return std::optional Parsed value, or std::nullopt if parsing fails - */ -[[nodiscard]] inline std::optional string_to_bool(std::string str) { - if (str.empty()) { - return std::nullopt; - } - - str.erase(0, str.find_first_not_of(" \t\r\n")); - str.erase(str.find_last_not_of(" \t\r\n") + 1); - std::transform(str.begin(), str.end(), str.begin(), - [](unsigned char c) { return std::tolower(c); }); - - if (str == "1" || str == "true" || str == "yes" || str == "on") { - return true; - } - if (str == "0" || str == "false" || str == "no" || str == "off") { - return false; - } - - return std::nullopt; -} - -/** - * @brief Split a string by delimiter into a vector of strings - * @param str The string to split - * @param delimiter The delimiter to split by (default is comma) - * @param trim_spaces Whether to trim leading/trailing spaces from each token - * @param keep_empty Whether to keep empty tokens in the result - * @return Vector of split strings - */ -std::vector splitString(const std::string& str, - char delimiter = ',', - bool trim_spaces = true, - bool keep_empty = false); - // Buffer allocator functions constexpr size_t SZ_2MB = 2 * 1024 * 1024; @@ -377,6 +338,24 @@ inline size_t align_up(size_t size, size_t alignment) { return size; } +/** + * @brief Fault in a fresh HugeTLB mapping with parallel CPU writes. + * + * Touches one byte per configured hugepage. Call this only for a newly + * allocated mapping whose contents may be zeroed. + */ +void populate_hugetlb_mapping(void* ptr, size_t total_size); + +/** + * @brief Fault in an mbind-partitioned HugeTLB mapping with NUMA-local workers. + * + * The mapping is divided into equal regions in the same order as numa_nodes. + * Workers are scheduled on the corresponding node before touching that + * region. + */ +void populate_hugetlb_numa_mapping(void* ptr, size_t total_size, + const std::vector& numa_nodes); + /** * Allocate mmap-backed buffer memory for host KV / transfer buffers. * @@ -394,6 +373,24 @@ inline size_t align_up(size_t size, size_t alignment) { */ void* allocate_buffer_mmap_memory(size_t total_size, size_t alignment); +/** + * Allocate mmap-backed memory, optionally deferring direct HugeTLB population. + * + * When defer_hugetlb_population is true, a direct HugeTLB mmap omits + * MAP_POPULATE so the caller can populate the mapping later. Arena allocations + * retain their existing eager-population behavior. + */ +void* allocate_buffer_mmap_memory(size_t total_size, size_t alignment, + bool defer_hugetlb_population); + +/** + * @brief Return whether ptr is backed by the global mmap arena. + * + * Intended for callers that need to distinguish an eagerly populated arena + * allocation from a direct mmap fallback. + */ +[[nodiscard]] bool is_mmap_arena_allocation(const void* ptr); + /** * Release memory previously returned by allocate_buffer_mmap_memory(). * @@ -411,8 +408,9 @@ void free_buffer_mmap_memory(void* ptr, size_t total_size); * * Reserves a single VMA via mmap, divides it into N equal regions, * binds each region to the corresponding NUMA node via mbind(MPOL_BIND). - * No explicit prefault — ibv_reg_mr() will fault and pin pages respecting - * the mbind policy, allocating directly on the target NUMA node. + * The mapping remains lazy after allocation. The caller may populate it with + * NUMA-local workers or let ibv_reg_mr() fault and pin pages while respecting + * the mbind policy. * * @param total_size Total buffer size in bytes * @param numa_nodes NUMA node IDs to bind regions to (e.g., {1,3,5,7}) @@ -424,7 +422,8 @@ void* allocate_buffer_numa_segments(size_t total_size, const std::vector& numa_nodes, size_t page_size = 0); -void free_memory(const std::string& protocol, void* ptr); +void free_memory(const std::string& protocol, void* ptr, + bool use_spdk_dma = false); // Network utility functions @@ -476,26 +475,7 @@ std::vector getFreeTcpPorts(int count); int64_t time_gen(); -// Helper: Get integer from environment variable, fallback to default -template -T GetEnvOr(const char* name, T default_value) { - const char* env_val = std::getenv(name); - if (!env_val || std::string(env_val).empty()) { - return default_value; - } - try { - long long value = std::stoll(env_val); - // Check range for unsigned types - if constexpr (std::is_same_v) { - if (value < 0 || value > UINT32_MAX) throw std::out_of_range(""); - } - return static_cast(value); - } catch (...) { - return default_value; - } -} - -std::string GetEnvStringOr(const char* name, const std::string& default_value); +std::string ResolveMooncakeHostId(const std::string& local_hostname); std::string ResolvePathFromKey(const std::string& key, const std::string& root_dir, diff --git a/mooncake-store/include/utils/type_util.h b/mooncake-store/include/utils/type_util.h deleted file mode 100644 index b97667b644..0000000000 --- a/mooncake-store/include/utils/type_util.h +++ /dev/null @@ -1,41 +0,0 @@ -#pragma once - -#include -#include - -namespace mooncake { - -/** - * @brief Type conversion utility class providing conversion between strings and - * basic types Can be used for various string and data type conversion - * operations - */ -class TypeUtil { - public: - /** - * @brief Parse string to boolean - * @param value String to parse (supports "true"/"false", "1"/"0", - * "yes"/"no", case-insensitive) - * @param out Output parameter, stores the result on success - * @return Returns true on success, false otherwise - */ - static bool ParseBool(std::string_view value, bool& out); - - /** - * @brief Parse string to uint64_t - * @param value String to parse - * @param out Output parameter, stores the result on success - * @return Returns true on success, false otherwise - */ - static bool ParseUint64(std::string_view value, uint64_t& out); - - /** - * @brief Parse string to int64_t - * @param value String to parse - * @param out Output parameter, stores the result on success - * @return Returns true on success, false otherwise - */ - static bool ParseInt64(std::string_view value, int64_t& out); -}; - -} // namespace mooncake diff --git a/mooncake-store/rust/CMakeLists.txt b/mooncake-store/rust/CMakeLists.txt index 49fa328fd5..820f92deb8 100644 --- a/mooncake-store/rust/CMakeLists.txt +++ b/mooncake-store/rust/CMakeLists.txt @@ -49,3 +49,16 @@ add_custom_command( WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} COMMENT "Building mooncake_store Rust tests" VERBATIM) + +# Maintainer target: regenerate the committed dlopen bindings from store_c.h +# (run after the C ABI changes). CI should run this then `git diff --exit-code`. +add_custom_target(generate_store_rust_dlopen_bindings) +add_custom_command( + TARGET generate_store_rust_dlopen_bindings + COMMAND + ${CMAKE_COMMAND} -E env CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR} + MOONCAKE_STORE_INCLUDE_DIR=${CMAKE_CURRENT_SOURCE_DIR}/../include cargo run + --example generate_dlopen_bindings + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + COMMENT "Regenerating committed mooncake_store Rust dlopen bindings" + VERBATIM) diff --git a/mooncake-store/rust/Cargo.lock b/mooncake-store/rust/Cargo.lock index 01eb6cbfa9..3c349152b3 100644 --- a/mooncake-store/rust/Cargo.lock +++ b/mooncake-store/rust/Cargo.lock @@ -124,6 +124,7 @@ version = "0.1.0" dependencies = [ "bindgen", "libc", + "libloading", "thiserror", ] diff --git a/mooncake-store/rust/Cargo.toml b/mooncake-store/rust/Cargo.toml index eefd85a09d..1dd0d39d6f 100644 --- a/mooncake-store/rust/Cargo.toml +++ b/mooncake-store/rust/Cargo.toml @@ -14,9 +14,35 @@ crate-type = ["rlib"] name = "basic_usage" path = "examples/basic_usage.rs" +# Maintainer tool: regenerates src/generated/ffi_dlopen_bindings.rs from +# store_c.h. Gated on `link` so it (and its bindgen dep) are not built for a +# plain `dlopen` build. Run via `cargo run --example generate_dlopen_bindings`. +[[example]] +name = "generate_dlopen_bindings" +path = "examples/generate_dlopen_bindings.rs" +required-features = ["link"] + +[features] +default = ["link"] +# Statically link `libmooncake_store` at build time. build.rs generates the FFI +# with bindgen and links the Mooncake C++ static-library graph, so a built +# Mooncake tree (libraries + `store_c.h`) must be present at build time. This is +# the historical behavior. +link = ["dep:bindgen"] +# Load `libmooncake_store.so` at run time via `dlopen`. Uses committed, +# pre-generated bindings (src/generated/ffi_dlopen_bindings.rs), so a consumer +# needs only `libloading` — no bindgen, header, or libclang, and no C++ linking. +dlopen = ["dep:libloading"] + [dependencies] thiserror = "2.0" libc = "0.2" +libloading = { version = "0.8", optional = true } +# bindgen: build-dep for the `link` backend's static bindings; dev-dep for the +# generate_dlopen_bindings example. The `dlopen` backend needs neither. [build-dependencies] +bindgen = { version = "0.70", optional = true } + +[dev-dependencies] bindgen = "0.70" diff --git a/mooncake-store/rust/README.md b/mooncake-store/rust/README.md index 65d4912525..cdd57b9e91 100644 --- a/mooncake-store/rust/README.md +++ b/mooncake-store/rust/README.md @@ -1,8 +1,44 @@ # Mooncake Store Rust Bindings This directory contains the Rust bindings, examples, benchmarks, and smoke tests -for Mooncake Store. The Rust crate links against the C++ Mooncake Store build, so -run a CMake build before using standalone Cargo commands. +for Mooncake Store. + +## Backends (features) + +The crate has two backends. Enable at least one; if both are enabled, `link` +takes precedence: + +- **`link`** (default): statically links `libmooncake_store` at build time. + `build.rs` generates the FFI with bindgen and links the C++ dependency graph, + so a built Mooncake tree (libraries + `store_c.h`) must be present at build + time. Everything below assumes this backend. +- **`dlopen`**: loads `libmooncake_store.so` at run time via `libloading`, using + **committed, pre-generated** bindings (`src/generated/ffi_dlopen_bindings.rs`). + A consumer needs only `libloading` — no bindgen, no `store_c.h`, no libclang, + no C++ linking — and only the shared library at run time. Maintainers + regenerate the bindings after the C ABI changes (see below): + + ```toml + mooncake_store = { version = "0.1", default-features = false, features = ["dlopen"] } + ``` + + Build the loadable library with `cmake -DWITH_STORE_C_SHARED=ON` (emits a + `libmooncake_store.so` exporting only the `store_c.h` C ABI). At run time it is + located via the `MOONCAKE_STORE_LIBRARY` environment variable (default + `libmooncake_store.so`, resolved through the OS loader search path), or pin an + explicit path with `mooncake_store::load_library(path)` before creating a + store. The public API is identical to the `link` backend. + + Regenerate the committed bindings after the C ABI (`store_c.h`) changes: + + ```bash + cargo run --example generate_dlopen_bindings # or: cmake --build build --target generate_store_rust_dlopen_bindings + ``` + + A pre-commit hook (on `store_c.h`/generator changes) and CI both enforce this + via `git diff --exit-code`, so stale bindings cannot merge. CI additionally + builds the packaged `.crate` with `--features dlopen` to confirm a published + consumer stays independent of `store_c.h`. ## Prerequisites diff --git a/mooncake-store/rust/build.rs b/mooncake-store/rust/build.rs index 5622d9550a..df77b0b1a7 100644 --- a/mooncake-store/rust/build.rs +++ b/mooncake-store/rust/build.rs @@ -12,6 +12,11 @@ // See the License for the specific language governing permissions and // limitations under the License. +// Everything below is used only by the `link` build path (see the two cfg'd +// main()s). Scope the allowances to non-`link` builds, where the helpers/imports +// are unused, so the `link` build keeps normal warning hygiene. +#![cfg_attr(not(feature = "link"), allow(dead_code, unused_imports))] + use std::env; use std::fs; use std::path::PathBuf; @@ -163,6 +168,13 @@ fn add_compiler_runtime_search_dir(search_dirs: &mut Vec, file_name: &s false } +// A pure `dlopen` build uses committed, pre-generated bindings +// (src/generated/ffi_dlopen_bindings.rs), so build.rs has nothing to do. Also +// covers the no-feature case (lib.rs emits a compile_error! there). +#[cfg(not(feature = "link"))] +fn main() {} + +#[cfg(feature = "link")] fn main() { // ----------------------------------------------------------------------- // Library search path @@ -203,19 +215,26 @@ fn main() { // common/base library (contains mooncake::Status etc.) println!( "cargo:rustc-link-search=native={}", - build_dir.join("mooncake-transfer-engine/src/common/base").display() + build_dir + .join("mooncake-transfer-engine/src/common/base") + .display() ); // CUDA runtime libraries (needed by transfer_engine RDMA transport). let cuda_home = env::var("CUDA_HOME") .or_else(|_| env::var("CUDA_PATH")) .unwrap_or_else(|_| "/usr/local/cuda".to_string()); - println!("cargo:rustc-link-search=native={}/targets/x86_64-linux/lib", cuda_home); + println!( + "cargo:rustc-link-search=native={}/targets/x86_64-linux/lib", + cuda_home + ); // cachelib_memory_allocator is a static library built alongside mooncake_store. println!( "cargo:rustc-link-search=native={}", - build_dir.join("mooncake-store/src/cachelib_memory_allocator").display() + build_dir + .join("mooncake-store/src/cachelib_memory_allocator") + .display() ); println!("cargo:rustc-link-lib=mooncake_store"); @@ -231,16 +250,18 @@ fn main() { println!("cargo:rustc-link-lib=stdc++"); println!("cargo:rustc-link-lib=glog"); println!("cargo:rustc-link-lib=gflags"); - println!("cargo:rustc-link-lib=numa"); // NUMA binding - println!("cargo:rustc-link-lib=curl"); // HTTP metadata plugin + println!("cargo:rustc-link-lib=numa"); // NUMA binding + println!("cargo:rustc-link-lib=curl"); // HTTP metadata plugin println!("cargo:rustc-link-lib=ibverbs"); // RDMA transport + println!("cargo:rustc-link-lib=yaml-cpp"); // tenant quota policy connector println!("cargo:rustc-link-lib=pthread"); println!("cargo:rustc-link-lib=xxhash"); // ----------------------------------------------------------------------- // Header path for bindgen // ----------------------------------------------------------------------- - let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").expect("missing CARGO_MANIFEST_DIR")); + let manifest_dir = + PathBuf::from(env::var("CARGO_MANIFEST_DIR").expect("missing CARGO_MANIFEST_DIR")); let mut search_dirs = Vec::new(); let explicit_lib_dir = env::var("MOONCAKE_STORE_LIB_DIR").ok().map(PathBuf::from); @@ -339,6 +360,7 @@ fn main() { "numa", "ibverbs", "jsoncpp", + "yaml-cpp", "zstd", "m", ] { @@ -353,6 +375,7 @@ fn main() { ("cudart", &["cudart"]), ("mlx5", &["mlx5"]), // IBGDA device transport (mlx5 DevX) pulled into transfer_engine, CUDA-only ("uring", &["uring"]), + ("zmq", &["zmq"]), ] { if has_library(&search_dirs, candidates) { println!("cargo:rustc-link-lib={link_name}"); @@ -363,8 +386,8 @@ fn main() { println!("cargo:rustc-link-lib=gcov"); } - let include_dir = env::var("MOONCAKE_STORE_INCLUDE_DIR") - .unwrap_or_else(|_| "../include".to_string()); + let include_dir = + env::var("MOONCAKE_STORE_INCLUDE_DIR").unwrap_or_else(|_| "../include".to_string()); let header = format!("{include_dir}/store_c.h"); diff --git a/mooncake-store/rust/examples/basic_usage.rs b/mooncake-store/rust/examples/basic_usage.rs index 4159b36978..c6f2e1b442 100644 --- a/mooncake-store/rust/examples/basic_usage.rs +++ b/mooncake-store/rust/examples/basic_usage.rs @@ -67,8 +67,8 @@ fn main() { // In CI or environments without a running metadata server the setup // call is expected to fail – this is not an error in the bindings // themselves. - let metadata_server = - std::env::var("MC_METADATA_SERVER").unwrap_or_else(|_| "http://127.0.0.1:8080/metadata".to_string()); + let metadata_server = std::env::var("MC_METADATA_SERVER") + .unwrap_or_else(|_| "http://127.0.0.1:8080/metadata".to_string()); println!("Connecting to metadata server: {metadata_server}"); @@ -78,7 +78,7 @@ fn main() { 512 << 20, // global_segment_size = 512 MiB 128 << 20, // local_buffer_size = 128 MiB "tcp", - "", // device_name (auto-select) + "", // device_name (auto-select) "127.0.0.1:50051", ) { eprintln!( @@ -98,13 +98,16 @@ fn main() { eprintln!("[FAIL] put() failed: {e}"); std::process::exit(1); } - println!("[OK] put(\"{key}\", {:?})", std::str::from_utf8(value).unwrap()); + println!( + "[OK] put(\"{key}\", {:?})", + std::str::from_utf8(value).unwrap() + ); // Step 4: Check existence. match store.is_exist(key) { - Ok(true) => println!("[OK] is_exist(\"{key}\") = true"), + Ok(true) => println!("[OK] is_exist(\"{key}\") = true"), Ok(false) => println!("[WARN] is_exist(\"{key}\") = false (unexpected)"), - Err(e) => eprintln!("[FAIL] is_exist() failed: {e}"), + Err(e) => eprintln!("[FAIL] is_exist() failed: {e}"), } // Step 5: Get size. diff --git a/mooncake-store/rust/examples/generate_dlopen_bindings.rs b/mooncake-store/rust/examples/generate_dlopen_bindings.rs new file mode 100644 index 0000000000..150c43667e --- /dev/null +++ b/mooncake-store/rust/examples/generate_dlopen_bindings.rs @@ -0,0 +1,50 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Maintainer tool: regenerates the committed `dlopen` bindings from +//! `store_c.h`. Run after the C ABI changes: +//! +//! ```bash +//! cargo run --example generate_dlopen_bindings +//! ``` +//! +//! The header is taken from `MOONCAKE_STORE_INCLUDE_DIR` (default `../include`). +//! CI should run this and `git diff --exit-code` to catch drift. + +use std::path::PathBuf; + +fn main() { + let manifest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + let include_dir = std::env::var("MOONCAKE_STORE_INCLUDE_DIR") + .map(PathBuf::from) + .unwrap_or_else(|_| manifest_dir.join("../include")); + let header = include_dir.join("store_c.h"); + let out = manifest_dir.join("src/generated/ffi_dlopen_bindings.rs"); + + let bindings = bindgen::Builder::default() + .header(header.to_string_lossy()) + .allowlist_function("mooncake_store_.*") + .allowlist_type("mooncake_.*") + .dynamic_library_name("MooncakeStoreLib") + .dynamic_link_require_all(true) + .raw_line("// @generated by `cargo run --example generate_dlopen_bindings` from store_c.h.") + .raw_line("// Do not edit by hand; see mooncake-store/rust/README.md to regenerate.") + .generate() + .expect("failed to generate dlopen bindings from store_c.h"); + + bindings + .write_to_file(&out) + .unwrap_or_else(|e| panic!("failed to write {}: {e}", out.display())); + println!("wrote {}", out.display()); +} diff --git a/mooncake-store/rust/src/error.rs b/mooncake-store/rust/src/error.rs index a854a3aadd..d8e6067aa7 100644 --- a/mooncake-store/rust/src/error.rs +++ b/mooncake-store/rust/src/error.rs @@ -15,7 +15,11 @@ use std::ffi::NulError; /// Errors returned by Mooncake Store operations. +/// +/// `#[non_exhaustive]` so adding variants (e.g. backend-specific ones) is not a +/// breaking change; downstream matches must include a wildcard arm. #[derive(Debug, thiserror::Error)] +#[non_exhaustive] pub enum StoreError { /// A required pointer argument was null (e.g. the store handle has not /// been initialised yet, or an internal allocation failed). @@ -44,4 +48,19 @@ pub enum StoreError { /// One or more arguments are invalid (e.g. mismatched array lengths). #[error("invalid argument: {0}")] InvalidArgument(String), + + /// The Mooncake shared library could not be loaded, or it does not export + /// the expected `store_c.h` C ABI (a required symbol was missing). + /// + /// Only ever produced by the `dlopen` backend (see `load_library`), but + /// always present so `StoreError` is identical across backends. + #[error("failed to load Mooncake shared library: {0}")] + LibraryLoad(String), + + /// `load_library` was called after the library had already been loaded; it + /// must be called before creating any store. + /// + /// Only ever produced by the `dlopen` backend. + #[error("Mooncake shared library already loaded; call load_library() before creating a store")] + LibraryAlreadyLoaded, } diff --git a/mooncake-store/rust/src/ffi_dlopen.rs b/mooncake-store/rust/src/ffi_dlopen.rs new file mode 100644 index 0000000000..d145bb719b --- /dev/null +++ b/mooncake-store/rust/src/ffi_dlopen.rs @@ -0,0 +1,249 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Runtime `dlopen` bindings for the Mooncake Store C API (`store_c.h`). +//! +//! The raw loader (types, layout checks, all `mooncake_store_*` symbols) is +//! bindgen-generated from `store_c.h` (`dynamic_library_name`, see build.rs), so +//! the ABI is never hand-maintained; this module wraps it in a process-global +//! plus free-function shims so `crate::store` is backend-agnostic. The library +//! loads lazily on first store creation from `MOONCAKE_STORE_LIBRARY` (default +//! `libmooncake_store.so`), or eagerly via [`load_library`]. + +use std::ffi::{OsStr, OsString}; +use std::os::raw::{c_char, c_int, c_void}; +use std::path::Path; +use std::sync::{Mutex, OnceLock}; + +use crate::error::StoreError; + +/// Bindgen-generated dynamic bindings (`MooncakeStoreLib` + the C types), +/// committed and regenerated via the `generate_dlopen_bindings` example. +mod sys { + #![allow(non_camel_case_types, non_snake_case, non_upper_case_globals)] + #![allow(dead_code)] + #![allow(clippy::all)] // generated code + include!("generated/ffi_dlopen_bindings.rs"); +} + +// Re-export the C types so `crate::store` can name them backend-agnostically. +pub use sys::{mooncake_replicate_config_t, mooncake_store_t}; + +/// Environment variable naming the shared library to load. +const LIBRARY_ENV: &str = "MOONCAKE_STORE_LIBRARY"; +/// Default library name, resolved via the OS loader search path. Platform-aware +/// so a non-Linux consumer that built its own library still finds it by default +/// (the `WITH_STORE_C_SHARED` producer is Linux-only, but the loader is not). +#[cfg(target_os = "linux")] +const DEFAULT_LIBRARY: &str = "libmooncake_store.so"; +#[cfg(target_os = "macos")] +const DEFAULT_LIBRARY: &str = "libmooncake_store.dylib"; +#[cfg(target_os = "windows")] +const DEFAULT_LIBRARY: &str = "mooncake_store.dll"; +#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))] +const DEFAULT_LIBRARY: &str = "libmooncake_store.so"; + +/// Process-wide, load-once handle to the generated loader. +static API: OnceLock = OnceLock::new(); +/// Serializes initialization so a concurrent race dlopen()s at most once. +static INIT_LOCK: Mutex<()> = Mutex::new(()); + +fn library_path() -> OsString { + std::env::var_os(LIBRARY_ENV).unwrap_or_else(|| OsString::from(DEFAULT_LIBRARY)) +} + +/// Open the library and resolve every symbol up front (`dynamic_link_require_all` +/// makes this fail if any `mooncake_store_*` symbol is missing). +fn load(path: &OsStr) -> Result { + unsafe { sys::MooncakeStoreLib::new(path) }.map_err(|e| StoreError::LibraryLoad(e.to_string())) +} + +/// Load `libmooncake_store.so` from an explicit `path`. Optional; otherwise the +/// library loads lazily on first store creation. Must be called before any store +/// exists, and returns [`StoreError::LibraryAlreadyLoaded`] if one is loaded. +pub fn load_library(path: impl AsRef) -> Result<(), StoreError> { + // Recover on poison rather than panic a library call (the section only loads). + let _guard = INIT_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if API.get().is_some() { + return Err(StoreError::LibraryAlreadyLoaded); + } + let lib = load(path.as_ref().as_os_str())?; + // Cannot fail: set under INIT_LOCK after the is_some() check above. + if API.set(lib).is_err() { + unreachable!("Mooncake API set while holding INIT_LOCK"); + } + Ok(()) +} + +/// Load the library from the default path on first use. Called by +/// `MooncakeStore::new()` before any other C call. +pub fn ensure_loaded() -> Result<(), StoreError> { + if API.get().is_some() { + return Ok(()); + } + // Recover on poison rather than panic a library call (the section only loads). + let _guard = INIT_LOCK + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if API.get().is_some() { + return Ok(()); + } + let lib = load(&library_path())?; + // Cannot fail: set under INIT_LOCK after the is_some() check above. + if API.set(lib).is_err() { + unreachable!("Mooncake API set while holding INIT_LOCK"); + } + Ok(()) +} + +/// The loaded loader. `MooncakeStore::new()` calls [`ensure_loaded`] first, so +/// this never fires in practice. +#[inline] +fn api() -> &'static sys::MooncakeStoreLib { + API.get() + .expect("Mooncake library not loaded; create a MooncakeStore first") +} + +/// Emits a free-function shim per entry that forwards to the generated loader +/// method of the same name. The shim signatures are compile-checked against the +/// bindgen-generated methods, so they cannot silently drift from `store_c.h`. +macro_rules! shims { + ( $( fn $name:ident ( $( $arg:ident : $argty:ty ),* $(,)? ) $( -> $ret:ty )? ; )* ) => { + $( + /// # Safety + /// Same contract as the C function in `store_c.h`. + #[inline] + #[allow(clippy::too_many_arguments)] + pub unsafe fn $name( $( $arg: $argty ),* ) $( -> $ret )? { + api().$name( $( $arg ),* ) + } + )* + }; +} + +shims! { + fn mooncake_store_create() -> mooncake_store_t; + fn mooncake_store_destroy(store: mooncake_store_t); + fn mooncake_store_setup( + store: mooncake_store_t, + local_hostname: *const c_char, + metadata_server: *const c_char, + global_segment_size: u64, + local_buffer_size: u64, + protocol: *const c_char, + device_name: *const c_char, + master_server_addr: *const c_char, + ) -> c_int; + fn mooncake_store_health_check(store: mooncake_store_t) -> c_int; + fn mooncake_store_put( + store: mooncake_store_t, + key: *const c_char, + value: *const c_void, + size: usize, + config: *const mooncake_replicate_config_t, + ) -> c_int; + fn mooncake_store_put_from( + store: mooncake_store_t, + key: *const c_char, + buffer: *mut c_void, + size: usize, + config: *const mooncake_replicate_config_t, + ) -> c_int; + fn mooncake_store_batch_put_from( + store: mooncake_store_t, + keys: *mut *const c_char, + buffers: *mut *mut c_void, + sizes: *const usize, + count: usize, + config: *const mooncake_replicate_config_t, + results_out: *mut c_int, + ) -> c_int; + fn mooncake_store_get_into( + store: mooncake_store_t, + key: *const c_char, + buffer: *mut c_void, + size: usize, + ) -> i64; + fn mooncake_store_batch_get_into( + store: mooncake_store_t, + keys: *mut *const c_char, + buffers: *mut *mut c_void, + sizes: *const usize, + count: usize, + results_out: *mut i64, + ) -> c_int; + fn mooncake_store_is_exist(store: mooncake_store_t, key: *const c_char) -> c_int; + fn mooncake_store_batch_is_exist( + store: mooncake_store_t, + keys: *mut *const c_char, + count: usize, + results_out: *mut c_int, + ) -> c_int; + fn mooncake_store_get_size(store: mooncake_store_t, key: *const c_char) -> i64; + fn mooncake_store_get_hostname( + store: mooncake_store_t, + buf_out: *mut c_char, + buf_len: usize, + ) -> c_int; + fn mooncake_store_remove(store: mooncake_store_t, key: *const c_char, force: c_int) -> c_int; + fn mooncake_store_remove_by_regex( + store: mooncake_store_t, + pattern: *const c_char, + force: c_int, + ) -> i64; + fn mooncake_store_remove_all(store: mooncake_store_t, force: c_int) -> i64; + fn mooncake_store_register_buffer( + store: mooncake_store_t, + buffer: *mut c_void, + size: usize, + ) -> c_int; + fn mooncake_store_unregister_buffer(store: mooncake_store_t, buffer: *mut c_void) -> c_int; +} + +#[cfg(test)] +mod tests { + use super::*; + + // A missing library must surface StoreError::LibraryLoad, not panic. `load` + // fails before touching the global, so this needs no real .so. + #[test] + fn load_library_missing_reports_library_load_error() { + let err = load_library("/nonexistent/does-not-exist/libmooncake_store.so") + .expect_err("loading a missing library must fail"); + assert!(matches!(err, StoreError::LibraryLoad(_)), "got {err:?}"); + } + + // new() must surface the load failure, not panic. Safe despite the shared + // env/global: no unit test loads a real library, so API stays unset. + #[test] + fn new_with_missing_library_reports_library_load_error() { + let prev = std::env::var_os(LIBRARY_ENV); + std::env::set_var( + LIBRARY_ENV, + "/nonexistent/does-not-exist/libmooncake_store.so", + ); + let result = crate::MooncakeStore::new(); + match prev { + Some(v) => std::env::set_var(LIBRARY_ENV, v), + None => std::env::remove_var(LIBRARY_ENV), + } + // MooncakeStore isn't Debug, so match rather than expect_err. + assert!( + matches!(result, Err(StoreError::LibraryLoad(_))), + "new() should report LibraryLoad when the library is missing" + ); + } +} diff --git a/mooncake-store/rust/src/generated/ffi_dlopen_bindings.rs b/mooncake-store/rust/src/generated/ffi_dlopen_bindings.rs new file mode 100644 index 0000000000..94d08f062b --- /dev/null +++ b/mooncake-store/rust/src/generated/ffi_dlopen_bindings.rs @@ -0,0 +1,387 @@ +/* automatically generated by rust-bindgen 0.70.1 */ + +// @generated by `cargo run --example generate_dlopen_bindings` from store_c.h. +// Do not edit by hand; see mooncake-store/rust/README.md to regenerate. + +pub type mooncake_store_t = *mut ::std::os::raw::c_void; +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct mooncake_replicate_config { + pub replica_num: usize, + pub with_soft_pin: ::std::os::raw::c_int, + pub with_hard_pin: ::std::os::raw::c_int, + pub preferred_segments: *mut *const ::std::os::raw::c_char, + pub preferred_segments_count: usize, +} +#[allow(clippy::unnecessary_operation, clippy::identity_op)] +const _: () = { + ["Size of mooncake_replicate_config"] + [::std::mem::size_of::() - 32usize]; + ["Alignment of mooncake_replicate_config"] + [::std::mem::align_of::() - 8usize]; + ["Offset of field: mooncake_replicate_config::replica_num"] + [::std::mem::offset_of!(mooncake_replicate_config, replica_num) - 0usize]; + ["Offset of field: mooncake_replicate_config::with_soft_pin"] + [::std::mem::offset_of!(mooncake_replicate_config, with_soft_pin) - 8usize]; + ["Offset of field: mooncake_replicate_config::with_hard_pin"] + [::std::mem::offset_of!(mooncake_replicate_config, with_hard_pin) - 12usize]; + ["Offset of field: mooncake_replicate_config::preferred_segments"] + [::std::mem::offset_of!(mooncake_replicate_config, preferred_segments) - 16usize]; + ["Offset of field: mooncake_replicate_config::preferred_segments_count"] + [::std::mem::offset_of!(mooncake_replicate_config, preferred_segments_count) - 24usize]; +}; +pub type mooncake_replicate_config_t = mooncake_replicate_config; +pub struct MooncakeStoreLib { + __library: ::libloading::Library, + pub mooncake_store_create: unsafe extern "C" fn() -> mooncake_store_t, + pub mooncake_store_destroy: unsafe extern "C" fn(store: mooncake_store_t), + pub mooncake_store_setup: unsafe extern "C" fn( + store: mooncake_store_t, + local_hostname: *const ::std::os::raw::c_char, + metadata_server: *const ::std::os::raw::c_char, + global_segment_size: u64, + local_buffer_size: u64, + protocol: *const ::std::os::raw::c_char, + device_name: *const ::std::os::raw::c_char, + master_server_addr: *const ::std::os::raw::c_char, + ) -> ::std::os::raw::c_int, + pub mooncake_store_init_all: unsafe extern "C" fn( + store: mooncake_store_t, + protocol: *const ::std::os::raw::c_char, + device_name: *const ::std::os::raw::c_char, + mount_segment_size: u64, + ) -> ::std::os::raw::c_int, + pub mooncake_store_health_check: + unsafe extern "C" fn(store: mooncake_store_t) -> ::std::os::raw::c_int, + pub mooncake_store_put: unsafe extern "C" fn( + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + value: *const ::std::os::raw::c_void, + size: usize, + config: *const mooncake_replicate_config_t, + ) -> ::std::os::raw::c_int, + pub mooncake_store_put_from: unsafe extern "C" fn( + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + buffer: *mut ::std::os::raw::c_void, + size: usize, + config: *const mooncake_replicate_config_t, + ) -> ::std::os::raw::c_int, + pub mooncake_store_batch_put_from: unsafe extern "C" fn( + store: mooncake_store_t, + keys: *mut *const ::std::os::raw::c_char, + buffers: *mut *mut ::std::os::raw::c_void, + sizes: *const usize, + count: usize, + config: *const mooncake_replicate_config_t, + results_out: *mut ::std::os::raw::c_int, + ) -> ::std::os::raw::c_int, + pub mooncake_store_get_into: unsafe extern "C" fn( + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + buffer: *mut ::std::os::raw::c_void, + size: usize, + ) -> i64, + pub mooncake_store_batch_get_into: unsafe extern "C" fn( + store: mooncake_store_t, + keys: *mut *const ::std::os::raw::c_char, + buffers: *mut *mut ::std::os::raw::c_void, + sizes: *const usize, + count: usize, + results_out: *mut i64, + ) -> ::std::os::raw::c_int, + pub mooncake_store_is_exist: unsafe extern "C" fn( + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + ) -> ::std::os::raw::c_int, + pub mooncake_store_batch_is_exist: unsafe extern "C" fn( + store: mooncake_store_t, + keys: *mut *const ::std::os::raw::c_char, + count: usize, + results_out: *mut ::std::os::raw::c_int, + ) -> ::std::os::raw::c_int, + pub mooncake_store_get_size: + unsafe extern "C" fn(store: mooncake_store_t, key: *const ::std::os::raw::c_char) -> i64, + pub mooncake_store_get_hostname: unsafe extern "C" fn( + store: mooncake_store_t, + buf_out: *mut ::std::os::raw::c_char, + buf_len: usize, + ) -> ::std::os::raw::c_int, + pub mooncake_store_remove: unsafe extern "C" fn( + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + force: ::std::os::raw::c_int, + ) -> ::std::os::raw::c_int, + pub mooncake_store_remove_by_regex: unsafe extern "C" fn( + store: mooncake_store_t, + pattern: *const ::std::os::raw::c_char, + force: ::std::os::raw::c_int, + ) -> i64, + pub mooncake_store_remove_all: + unsafe extern "C" fn(store: mooncake_store_t, force: ::std::os::raw::c_int) -> i64, + pub mooncake_store_register_buffer: unsafe extern "C" fn( + store: mooncake_store_t, + buffer: *mut ::std::os::raw::c_void, + size: usize, + ) -> ::std::os::raw::c_int, + pub mooncake_store_unregister_buffer: unsafe extern "C" fn( + store: mooncake_store_t, + buffer: *mut ::std::os::raw::c_void, + ) -> ::std::os::raw::c_int, +} +impl MooncakeStoreLib { + pub unsafe fn new

(path: P) -> Result + where + P: AsRef<::std::ffi::OsStr>, + { + let library = ::libloading::Library::new(path)?; + Self::from_library(library) + } + pub unsafe fn from_library(library: L) -> Result + where + L: Into<::libloading::Library>, + { + let __library = library.into(); + let mooncake_store_create = __library.get(b"mooncake_store_create\0").map(|sym| *sym)?; + let mooncake_store_destroy = __library.get(b"mooncake_store_destroy\0").map(|sym| *sym)?; + let mooncake_store_setup = __library.get(b"mooncake_store_setup\0").map(|sym| *sym)?; + let mooncake_store_init_all = __library + .get(b"mooncake_store_init_all\0") + .map(|sym| *sym)?; + let mooncake_store_health_check = __library + .get(b"mooncake_store_health_check\0") + .map(|sym| *sym)?; + let mooncake_store_put = __library.get(b"mooncake_store_put\0").map(|sym| *sym)?; + let mooncake_store_put_from = __library + .get(b"mooncake_store_put_from\0") + .map(|sym| *sym)?; + let mooncake_store_batch_put_from = __library + .get(b"mooncake_store_batch_put_from\0") + .map(|sym| *sym)?; + let mooncake_store_get_into = __library + .get(b"mooncake_store_get_into\0") + .map(|sym| *sym)?; + let mooncake_store_batch_get_into = __library + .get(b"mooncake_store_batch_get_into\0") + .map(|sym| *sym)?; + let mooncake_store_is_exist = __library + .get(b"mooncake_store_is_exist\0") + .map(|sym| *sym)?; + let mooncake_store_batch_is_exist = __library + .get(b"mooncake_store_batch_is_exist\0") + .map(|sym| *sym)?; + let mooncake_store_get_size = __library + .get(b"mooncake_store_get_size\0") + .map(|sym| *sym)?; + let mooncake_store_get_hostname = __library + .get(b"mooncake_store_get_hostname\0") + .map(|sym| *sym)?; + let mooncake_store_remove = __library.get(b"mooncake_store_remove\0").map(|sym| *sym)?; + let mooncake_store_remove_by_regex = __library + .get(b"mooncake_store_remove_by_regex\0") + .map(|sym| *sym)?; + let mooncake_store_remove_all = __library + .get(b"mooncake_store_remove_all\0") + .map(|sym| *sym)?; + let mooncake_store_register_buffer = __library + .get(b"mooncake_store_register_buffer\0") + .map(|sym| *sym)?; + let mooncake_store_unregister_buffer = __library + .get(b"mooncake_store_unregister_buffer\0") + .map(|sym| *sym)?; + Ok(MooncakeStoreLib { + __library, + mooncake_store_create, + mooncake_store_destroy, + mooncake_store_setup, + mooncake_store_init_all, + mooncake_store_health_check, + mooncake_store_put, + mooncake_store_put_from, + mooncake_store_batch_put_from, + mooncake_store_get_into, + mooncake_store_batch_get_into, + mooncake_store_is_exist, + mooncake_store_batch_is_exist, + mooncake_store_get_size, + mooncake_store_get_hostname, + mooncake_store_remove, + mooncake_store_remove_by_regex, + mooncake_store_remove_all, + mooncake_store_register_buffer, + mooncake_store_unregister_buffer, + }) + } + pub unsafe fn mooncake_store_create(&self) -> mooncake_store_t { + (self.mooncake_store_create)() + } + pub unsafe fn mooncake_store_destroy(&self, store: mooncake_store_t) { + (self.mooncake_store_destroy)(store) + } + pub unsafe fn mooncake_store_setup( + &self, + store: mooncake_store_t, + local_hostname: *const ::std::os::raw::c_char, + metadata_server: *const ::std::os::raw::c_char, + global_segment_size: u64, + local_buffer_size: u64, + protocol: *const ::std::os::raw::c_char, + device_name: *const ::std::os::raw::c_char, + master_server_addr: *const ::std::os::raw::c_char, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_setup)( + store, + local_hostname, + metadata_server, + global_segment_size, + local_buffer_size, + protocol, + device_name, + master_server_addr, + ) + } + pub unsafe fn mooncake_store_init_all( + &self, + store: mooncake_store_t, + protocol: *const ::std::os::raw::c_char, + device_name: *const ::std::os::raw::c_char, + mount_segment_size: u64, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_init_all)(store, protocol, device_name, mount_segment_size) + } + pub unsafe fn mooncake_store_health_check( + &self, + store: mooncake_store_t, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_health_check)(store) + } + pub unsafe fn mooncake_store_put( + &self, + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + value: *const ::std::os::raw::c_void, + size: usize, + config: *const mooncake_replicate_config_t, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_put)(store, key, value, size, config) + } + pub unsafe fn mooncake_store_put_from( + &self, + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + buffer: *mut ::std::os::raw::c_void, + size: usize, + config: *const mooncake_replicate_config_t, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_put_from)(store, key, buffer, size, config) + } + pub unsafe fn mooncake_store_batch_put_from( + &self, + store: mooncake_store_t, + keys: *mut *const ::std::os::raw::c_char, + buffers: *mut *mut ::std::os::raw::c_void, + sizes: *const usize, + count: usize, + config: *const mooncake_replicate_config_t, + results_out: *mut ::std::os::raw::c_int, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_batch_put_from)( + store, + keys, + buffers, + sizes, + count, + config, + results_out, + ) + } + pub unsafe fn mooncake_store_get_into( + &self, + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + buffer: *mut ::std::os::raw::c_void, + size: usize, + ) -> i64 { + (self.mooncake_store_get_into)(store, key, buffer, size) + } + pub unsafe fn mooncake_store_batch_get_into( + &self, + store: mooncake_store_t, + keys: *mut *const ::std::os::raw::c_char, + buffers: *mut *mut ::std::os::raw::c_void, + sizes: *const usize, + count: usize, + results_out: *mut i64, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_batch_get_into)(store, keys, buffers, sizes, count, results_out) + } + pub unsafe fn mooncake_store_is_exist( + &self, + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_is_exist)(store, key) + } + pub unsafe fn mooncake_store_batch_is_exist( + &self, + store: mooncake_store_t, + keys: *mut *const ::std::os::raw::c_char, + count: usize, + results_out: *mut ::std::os::raw::c_int, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_batch_is_exist)(store, keys, count, results_out) + } + pub unsafe fn mooncake_store_get_size( + &self, + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + ) -> i64 { + (self.mooncake_store_get_size)(store, key) + } + pub unsafe fn mooncake_store_get_hostname( + &self, + store: mooncake_store_t, + buf_out: *mut ::std::os::raw::c_char, + buf_len: usize, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_get_hostname)(store, buf_out, buf_len) + } + pub unsafe fn mooncake_store_remove( + &self, + store: mooncake_store_t, + key: *const ::std::os::raw::c_char, + force: ::std::os::raw::c_int, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_remove)(store, key, force) + } + pub unsafe fn mooncake_store_remove_by_regex( + &self, + store: mooncake_store_t, + pattern: *const ::std::os::raw::c_char, + force: ::std::os::raw::c_int, + ) -> i64 { + (self.mooncake_store_remove_by_regex)(store, pattern, force) + } + pub unsafe fn mooncake_store_remove_all( + &self, + store: mooncake_store_t, + force: ::std::os::raw::c_int, + ) -> i64 { + (self.mooncake_store_remove_all)(store, force) + } + pub unsafe fn mooncake_store_register_buffer( + &self, + store: mooncake_store_t, + buffer: *mut ::std::os::raw::c_void, + size: usize, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_register_buffer)(store, buffer, size) + } + pub unsafe fn mooncake_store_unregister_buffer( + &self, + store: mooncake_store_t, + buffer: *mut ::std::os::raw::c_void, + ) -> ::std::os::raw::c_int { + (self.mooncake_store_unregister_buffer)(store, buffer) + } +} diff --git a/mooncake-store/rust/src/lib.rs b/mooncake-store/rust/src/lib.rs index 683f5fbf54..70177dacb3 100644 --- a/mooncake-store/rust/src/lib.rs +++ b/mooncake-store/rust/src/lib.rs @@ -41,8 +41,29 @@ //! store.remove("hello", false)?; //! ``` +//! ## Backends (features) +//! +//! - `link` (default): statically link `libmooncake_store` at build time +//! (bindgen; requires a built Mooncake C++ tree + `store_c.h`). +//! - `dlopen`: load `libmooncake_store.so` at run time from committed bindings; +//! builds with only `libloading` (no bindgen, header, or C++ toolchain). See +//! the `load_library` function (exported under this feature). +//! +//! At least one must be enabled. If both are enabled (e.g. via Cargo feature +//! unification), `link` takes precedence. + +#[cfg(not(any(feature = "link", feature = "dlopen")))] +compile_error!("enable at least one of the `link` (default) or `dlopen` features"); + pub mod error; pub mod store; +// Runtime dlopen backend; `link` wins if both are enabled. +#[cfg(all(feature = "dlopen", not(feature = "link")))] +mod ffi_dlopen; + pub use error::StoreError; pub use store::{MooncakeStore, ReplicateConfig}; + +#[cfg(all(feature = "dlopen", not(feature = "link")))] +pub use ffi_dlopen::load_library; diff --git a/mooncake-store/rust/src/store.rs b/mooncake-store/rust/src/store.rs index cb59439c7f..6c54952eb6 100644 --- a/mooncake-store/rust/src/store.rs +++ b/mooncake-store/rust/src/store.rs @@ -18,7 +18,13 @@ use std::ffi::{c_void, CString}; use crate::error::StoreError; -// Raw FFI bindings generated by build.rs / bindgen. +// Raw FFI layer. Exactly one backend is active and both expose the same items +// (`mooncake_store_t`, `mooncake_replicate_config_t`, `mooncake_store_*`), so +// the safe wrapper below is backend-agnostic. +// +// `link`: bindgen bindings, statically linked (build.rs). `dlopen`: runtime +// libloading bindings (ffi_dlopen.rs). +#[cfg(feature = "link")] #[allow(dead_code)] mod ffi { #![allow(non_snake_case)] @@ -28,6 +34,9 @@ mod ffi { include!(concat!(env!("OUT_DIR"), "/bindings.rs")); } +#[cfg(all(feature = "dlopen", not(feature = "link")))] +use crate::ffi_dlopen as ffi; + // --------------------------------------------------------------------------- // ReplicateConfig // --------------------------------------------------------------------------- @@ -57,7 +66,11 @@ impl ReplicateConfig { fn to_ffi( &self, ) -> Result< - (ffi::mooncake_replicate_config_t, Vec, Vec<*const libc::c_char>), + ( + ffi::mooncake_replicate_config_t, + Vec, + Vec<*const libc::c_char>, + ), StoreError, > { let strings: Vec = self @@ -120,6 +133,10 @@ impl MooncakeStore { /// /// Call [`MooncakeStore::setup`] before performing any data operations. pub fn new() -> Result { + // dlopen backend: resolve libmooncake_store.so before the first C call. + #[cfg(all(feature = "dlopen", not(feature = "link")))] + ffi::ensure_loaded()?; + let handle = unsafe { ffi::mooncake_store_create() }; if handle.is_null() { return Err(StoreError::NullHandle); @@ -277,8 +294,7 @@ impl MooncakeStore { pub fn get(&self, key: &str) -> Result, StoreError> { let size = self.get_size(key)?; let mut buf = vec![0u8; size as usize]; - let written = - unsafe { self.get_into(key, buf.as_mut_ptr() as *mut c_void, buf.len())? }; + let written = unsafe { self.get_into(key, buf.as_mut_ptr() as *mut c_void, buf.len())? }; buf.truncate(written as usize); Ok(buf) } @@ -345,9 +361,8 @@ impl MooncakeStore { /// read by another client. pub fn remove(&self, key: &str, force: bool) -> Result<(), StoreError> { let key_c = CString::new(key)?; - let rc = unsafe { - ffi::mooncake_store_remove(self.handle, key_c.as_ptr(), i32::from(force)) - }; + let rc = + unsafe { ffi::mooncake_store_remove(self.handle, key_c.as_ptr(), i32::from(force)) }; if rc != 0 { return Err(StoreError::OperationFailed(rc)); } @@ -451,8 +466,7 @@ impl MooncakeStore { .iter() .map(|k| CString::new(*k)) .collect::>()?; - let key_ptrs: Vec<*const libc::c_char> = - key_strings.iter().map(|s| s.as_ptr()).collect(); + let key_ptrs: Vec<*const libc::c_char> = key_strings.iter().map(|s| s.as_ptr()).collect(); let (_c_config, _strings, _ptrs) = Self::prepare_config(config)?; let cfg_ptr = _c_config @@ -506,8 +520,7 @@ impl MooncakeStore { .iter() .map(|k| CString::new(*k)) .collect::>()?; - let key_ptrs: Vec<*const libc::c_char> = - key_strings.iter().map(|s| s.as_ptr()).collect(); + let key_ptrs: Vec<*const libc::c_char> = key_strings.iter().map(|s| s.as_ptr()).collect(); let mut results = vec![0i64; count]; @@ -528,10 +541,7 @@ impl MooncakeStore { } /// Batch check existence of multiple keys. - pub fn batch_is_exist( - &self, - keys: &[&str], - ) -> Result, StoreError> { + pub fn batch_is_exist(&self, keys: &[&str]) -> Result, StoreError> { let count = keys.len(); if count == 0 { return Ok(Vec::new()); @@ -540,8 +550,7 @@ impl MooncakeStore { .iter() .map(|k| CString::new(*k)) .collect::>()?; - let key_ptrs: Vec<*const libc::c_char> = - key_strings.iter().map(|s| s.as_ptr()).collect(); + let key_ptrs: Vec<*const libc::c_char> = key_strings.iter().map(|s| s.as_ptr()).collect(); let mut results = vec![0i32; count]; @@ -665,15 +674,13 @@ mod tests { preferred_segments: vec!["bad\0segment".to_string()], }; - assert!(matches!( - config.to_ffi(), - Err(StoreError::InvalidString(_)) - )); + assert!(matches!(config.to_ffi(), Err(StoreError::InvalidString(_)))); } #[test] fn prepare_config_none_returns_null_config() { - let (c_cfg, strings, ptrs) = MooncakeStore::prepare_config(None).expect("prepare should succeed"); + let (c_cfg, strings, ptrs) = + MooncakeStore::prepare_config(None).expect("prepare should succeed"); assert!(c_cfg.is_none()); assert!(strings.is_empty()); assert!(ptrs.is_empty()); @@ -703,13 +710,19 @@ mod tests { // Batch API tests // ----------------------------------------------------------------------- + // These construct a live store handle, so they need the library: statically + // linked under `link`, or present at run time under `dlopen`. + #[cfg(feature = "link")] #[test] fn batch_is_exist_empty() { let store = MooncakeStore::new().expect("new should succeed"); - let results = store.batch_is_exist(&[]).expect("empty batch should succeed"); + let results = store + .batch_is_exist(&[]) + .expect("empty batch should succeed"); assert!(results.is_empty()); } + #[cfg(feature = "link")] #[test] fn batch_put_from_rejects_mismatched_lengths() { let store = MooncakeStore::new().expect("new should succeed"); @@ -718,12 +731,10 @@ mod tests { let sizes = &[100usize, 200]; let result = unsafe { store.batch_put_from(keys, &buffers, sizes, None) }; - assert!(matches!( - result, - Err(StoreError::InvalidArgument(_)) - )); + assert!(matches!(result, Err(StoreError::InvalidArgument(_)))); } + #[cfg(feature = "link")] #[test] fn batch_get_into_rejects_mismatched_lengths() { let store = MooncakeStore::new().expect("new should succeed"); @@ -732,9 +743,6 @@ mod tests { let sizes = &[100usize]; let result = unsafe { store.batch_get_into(keys, &buffers, sizes) }; - assert!(matches!( - result, - Err(StoreError::InvalidArgument(_)) - )); + assert!(matches!(result, Err(StoreError::InvalidArgument(_)))); } } diff --git a/mooncake-store/src/CMakeLists.txt b/mooncake-store/src/CMakeLists.txt index b86e93846a..f20c7cbf7d 100644 --- a/mooncake-store/src/CMakeLists.txt +++ b/mooncake-store/src/CMakeLists.txt @@ -4,6 +4,8 @@ add_subdirectory(cachelib_memory_allocator) set(MOONCAKE_STORE_SOURCES allocator.cpp master_service.cpp + master_snapshot_manager.cpp + master_snapshot_repository.cpp client_service.cpp client_metric.cpp types.cpp @@ -17,6 +19,7 @@ set(MOONCAKE_STORE_SOURCES segment.cpp transfer_task.cpp tenant_quota.cpp + tenant_quota_policy_store.cpp rpc_service.cpp master_admin_service.cpp offset_allocator.cpp @@ -24,45 +27,55 @@ set(MOONCAKE_STORE_SOURCES client_buffer.cpp aligned_client_buffer.cpp real_client.cpp + registered_pinned_memory.cpp dummy_client.cpp + uds_transport.cpp shm_helper.cpp http_metadata_server.cpp file_storage.cpp serialize/serializer.cpp storage/distributed/distributed_storage_backend.cpp + device/accelerator_device.cpp + device/accelerator_registry.cpp + device/runtime_accelerator.cpp + device/cuda_like_accelerator_device.cpp + device/cuda_ipc_buffer.cpp + device/hip_accelerator_device.cpp + device/ascend_accelerator_device.cpp + device/sunrise_accelerator_device.cpp ha/leadership/leader_coordinator_factory.cpp ha/leadership/backends/etcd/etcd_leader_coordinator.cpp ha/common/redis/redis_connection.cpp ha/leadership/backends/redis/redis_leader_coordinator.cpp ha/leadership/master_service_supervisor.cpp + ha/kv/etcd_ha_kv_backend.cpp ha/standby_controller.cpp ha/snapshot/catalog_backed_snapshot_provider.cpp + ha/snapshot/master_snapshot_codec.cpp ha/snapshot/object/snapshot_object_store.cpp ha/snapshot/object/backends/local/local_file_snapshot_object_store.cpp ha/snapshot/object/backends/s3/s3_snapshot_object_store.cpp ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.cpp ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.cpp - utils/type_util.cpp utils/file_util.cpp task_manager.cpp local_hot_cache.cpp - ha/oplog/oplog_manager.cpp - ha/oplog/etcd_oplog_store.cpp - ha/oplog/etcd_oplog_change_notifier.cpp - ha/oplog/oplog_serializer.cpp - ha/oplog/oplog_store_factory.cpp - ha/oplog/oplog_replicator.cpp + ha/oplog/oplog_types.cpp + ha/oplog/oplog_batch_codec.cpp + ha/oplog/oplog_batch_storage.cpp + ha/oplog/oplog_batch_standby_reader.cpp + ha/oplog/oplog_batch_types.cpp + ha/oplog/ordered_oplog_writer.cpp + ha/oplog/oplog_test_failpoint.cpp ha/oplog/oplog_applier.cpp - ha/oplog/localfs_oplog_store.cpp - ha/oplog/polling_oplog_change_notifier.cpp hot_standby_service.cpp + metadata_store.cpp standby_state_machine.cpp ha_metric_manager.cpp store_c.cpp memory_alloc.cpp ssd_register_client.cpp - engram/engram_store.cpp -) + engram/engram_store.cpp) set(EXTRA_LIBS "") set(SPDK_STATIC_LIBS "") @@ -106,6 +119,31 @@ find_library( XXHASH_LIBRARY NAMES xxhash libxxhash PATHS /usr/lib /usr/local/lib /usr/lib64) +set(KV_EVENTS_ZMQ_FOUND FALSE) +if(ENABLE_KV_EVENTS) + find_library( + ZMQ_LIBRARY + NAMES zmq libzmq + PATHS /usr/lib /usr/local/lib /usr/lib64) + find_path( + ZMQ_INCLUDE_DIR + NAMES zmq.h + PATHS /usr/include /usr/local/include) + if(ZMQ_INCLUDE_DIR AND ZMQ_LIBRARY) + message(STATUS "Found ZMQ: include=${ZMQ_INCLUDE_DIR} lib=${ZMQ_LIBRARY}") + list(APPEND MASTER_EXTRA_INCS ${ZMQ_INCLUDE_DIR}) + list(APPEND EXTRA_LIBS ${ZMQ_LIBRARY}) + list(APPEND MOONCAKE_STORE_SOURCES kv_event/kv_event_publisher.cpp) + set(KV_EVENTS_ZMQ_FOUND TRUE) + else() + message( + FATAL_ERROR + "ENABLE_KV_EVENTS is ON but libzmq was not found. " + "Install libzmq3-dev or pass -DENABLE_KV_EVENTS=OFF to configure without ZMQ." + ) + endif() +endif() + if(XXHASH_INCLUDE_DIR AND XXHASH_LIBRARY) message( STATUS "Found xxHash: include=${XXHASH_INCLUDE_DIR} lib=${XXHASH_LIBRARY}") @@ -147,11 +185,16 @@ if(STORE_USE_REDIS) list(APPEND EXTRA_LIBS ${MOONCAKE_STORE_HIREDIS_LIBRARY}) endif() +list(APPEND MOONCAKE_STORE_SOURCES k8s_lease_helper.cpp) if(STORE_USE_K8S_LEASE) list(APPEND MOONCAKE_STORE_SOURCES - k8s_lease_helper.cpp - ha/leadership/backends/k8s/k8s_leader_coordinator.cpp) + ha/leadership/backends/k8s/k8s_leader_coordinator.cpp) list(APPEND EXTRA_LIBS ${K8S_LEASE_WRAPPER_LIB}) + set_source_files_properties( + k8s_lease_helper.cpp + PROPERTIES + OBJECT_DEPENDS + "${CMAKE_BINARY_DIR}/mooncake-common/k8s-lease/libk8s_lease_wrapper.h") endif() if(USE_NOF) @@ -223,6 +266,7 @@ if(USE_NOF) crypto aio z + isal elf ibverbs rdmacm @@ -233,6 +277,13 @@ endif() # The cache_allocator library include_directories(${Python3_INCLUDE_DIRS}) add_library(mooncake_store ${MOONCAKE_STORE_SOURCES}) +if(KV_EVENTS_ZMQ_FOUND) + target_compile_definitions(mooncake_store PRIVATE MOONCAKE_ENABLE_KV_EVENTS=1) +endif() +if(MOONCAKE_ENABLE_OPLOG_PERF_METRICS) + target_compile_definitions(mooncake_store + PRIVATE MOONCAKE_ENABLE_OPLOG_PERF_METRICS) +endif() target_include_directories(mooncake_store PUBLIC ${XXHASH_INCLUDE_DIR}) if(USE_NOF) target_include_directories(mooncake_store PRIVATE ${SPDK_INCLUDE_DIR} @@ -248,14 +299,23 @@ endif() # mooncake_master). Targets that need transfer_engine should link it explicitly. target_link_libraries( mooncake_store - PUBLIC cachelib_memory_allocator ${ETCD_WRAPPER_LIB} glog::glog gflags::gflags - ${EXTRA_LIBS} ${SPDK_STATIC_LIBS} asio_shared - PRIVATE transfer_engine) + PUBLIC cachelib_memory_allocator + ${ETCD_WRAPPER_LIB} + glog::glog + gflags::gflags + ${EXTRA_LIBS} + ${SPDK_STATIC_LIBS} + mooncake_common + yaml-cpp + asio_shared + mooncake_common + yalantinglibs::yalantinglibs + PRIVATE transfer_engine JsonCpp::JsonCpp) if(STORE_USE_ETCD) add_dependencies(mooncake_store build_etcd_wrapper) endif() if(STORE_USE_K8S_LEASE) - add_dependencies(mooncake_store build_k8s_lease_wrapper) + add_dependencies(mooncake_store build_k8s_lease_wrapper) endif() if(URING_LIB AND URING_INCLUDE) @@ -268,6 +328,9 @@ if(BUILD_SHARED_LIBS) endif() # Master binary +include(CheckPIESupported) +check_pie_supported(LANGUAGES CXX) + add_executable(mooncake_master master.cpp) set(MASTER_EXTRA_INCS) @@ -288,9 +351,17 @@ target_link_libraries( pthread ibverbs mooncake_common + JsonCpp::JsonCpp ${ETCD_WRAPPER_LIB} ${MASTER_EXTRA_LIBS} asio_shared) +# mooncake_store is static and propagates optional runtime dependencies to its +# consumers. The metadata-only master does not use accelerator staging, so drop +# those dependencies when they do not satisfy any symbols in the final binary. +target_link_options(mooncake_master PRIVATE "LINKER:--as-needed") +if(USE_SUNRISE) + target_link_directories(mooncake_master PRIVATE ${MC_TANGRT_ROOT}/lib) +endif() if(STORE_USE_ETCD) add_dependencies(mooncake_master build_etcd_wrapper) @@ -301,26 +372,28 @@ add_executable(mooncake_client real_client_main.cpp) # Client needs transfer_engine for data transfer operations target_link_libraries(mooncake_client PRIVATE mooncake_store transfer_engine asio_shared) +set_target_properties(mooncake_master mooncake_client + PROPERTIES POSITION_INDEPENDENT_CODE ON) # Optimize binary sizes only in Release mode string(TOUPPER "${CMAKE_BUILD_TYPE}" CMAKE_BUILD_TYPE_UPPER) -if (CMAKE_BUILD_TYPE_UPPER STREQUAL "RELEASE") +if(CMAKE_BUILD_TYPE_UPPER STREQUAL "RELEASE") target_compile_options(mooncake_master PRIVATE -Os) target_link_options(mooncake_master PRIVATE -Os -s) target_compile_options(mooncake_client PRIVATE -Os) target_link_options(mooncake_client PRIVATE -Os -s) endif() - # GPU runtime library for D2H staging in PutToLocalFile / OffloadObjects. -# transfer_engine is PRIVATE-linked, so its CUDA/HIP/Ascend dependencies -# are not propagated; we must detect and link them independently. +# transfer_engine is PRIVATE-linked, so its CUDA/HIP/Ascend dependencies are not +# propagated; we must detect and link them independently. # -# Auto-detect each toolkit regardless of global USE_CUDA/USE_HIP flags, -# because USE_CUDA may be OFF even when GPU pointers are present -# (e.g. WITH_NVIDIA_PEERMEM env var uses nvidia-peermem for RDMA without cudart). -# Each detected toolkit gets both link libraries AND compile definitions, -# so that gpu_staging_utils.h / pinned_buffer_pool.h pick the correct backend. +# Auto-detect each toolkit regardless of global USE_CUDA/USE_HIP flags, because +# USE_CUDA may be OFF even when GPU pointers are present (e.g. +# WITH_NVIDIA_PEERMEM env var uses nvidia-peermem for RDMA without cudart). Each +# detected toolkit gets both link libraries AND compile definitions, so that +# AcceleratorDevice implementations and pinned_buffer_pool.h pick the correct +# backend. # # NOTE: mooncake_store is a static library (.a). External consumers (Go CGo, # Python pybind) that link it must also link the GPU runtime (e.g. -lcudart). @@ -328,30 +401,103 @@ endif() find_package(CUDAToolkit QUIET) if(CUDAToolkit_FOUND) - message(STATUS "mooncake_store: CUDAToolkit detected, enabling D2H staging") - target_compile_definitions(mooncake_store PRIVATE USE_CUDA) - target_compile_definitions(mooncake_client PRIVATE USE_CUDA) - target_include_directories(mooncake_store PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) - target_include_directories(mooncake_client PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) - target_link_libraries(mooncake_store PRIVATE CUDA::cudart) - target_link_libraries(mooncake_client PRIVATE CUDA::cudart) + message(STATUS "mooncake_store: CUDAToolkit detected, enabling D2H staging") + target_compile_definitions(mooncake_store PRIVATE USE_CUDA) + target_compile_definitions(mooncake_client PRIVATE USE_CUDA) + target_include_directories(mooncake_store PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + target_include_directories(mooncake_client + PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + target_link_libraries(mooncake_store PRIVATE CUDA::cudart) + target_link_libraries(mooncake_client PRIVATE CUDA::cudart) endif() if(NOT CUDAToolkit_FOUND) - find_package(hip QUIET) - if(hip_FOUND) - message(STATUS "mooncake_store: HIP detected, enabling D2H staging") - target_compile_definitions(mooncake_store PRIVATE USE_HIP) - target_compile_definitions(mooncake_client PRIVATE USE_HIP) - target_link_libraries(mooncake_store PRIVATE hip::host) - target_link_libraries(mooncake_client PRIVATE hip::host) - endif() + find_package(hip QUIET) + if(hip_FOUND) + message(STATUS "mooncake_store: HIP detected, enabling D2H staging") + target_compile_definitions(mooncake_store PRIVATE USE_HIP) + target_compile_definitions(mooncake_client PRIVATE USE_HIP) + target_link_libraries(mooncake_store PRIVATE hip::host) + target_link_libraries(mooncake_client PRIVATE hip::host) + endif() +endif() + +if(USE_ASCEND + OR USE_ASCEND_DIRECT + OR USE_UBSHMEM) + target_include_directories(mooncake_store PRIVATE $ENV{ASCEND_PATH}/include) + target_link_libraries(mooncake_store PRIVATE ascendcl) + target_include_directories(mooncake_client PRIVATE $ENV{ASCEND_PATH}/include) + target_link_libraries(mooncake_client PRIVATE ascendcl) endif() -if(USE_ASCEND OR USE_ASCEND_DIRECT OR USE_UBSHMEM) - target_include_directories(mooncake_store PRIVATE $ENV{ASCEND_PATH}/include) - target_link_libraries(mooncake_store PRIVATE ascendcl) - target_include_directories(mooncake_client PRIVATE $ENV{ASCEND_PATH}/include) - target_link_libraries(mooncake_client PRIVATE ascendcl) +if(USE_SUNRISE) + target_link_libraries( + mooncake_store PRIVATE ${MC_TANGRT_ROOT}/lib/libtangrt_shared.so + ${MC_TANGRT_ROOT}/lib/libptml_shared.so dl) + target_link_libraries( + mooncake_client PRIVATE ${MC_TANGRT_ROOT}/lib/libtangrt_shared.so + ${MC_TANGRT_ROOT}/lib/libptml_shared.so dl) endif() + install(TARGETS mooncake_master mooncake_client DESTINATION bin) + +# --------------------------------------------------------------------------- +# Optional: a self-contained libmooncake_store.so exporting only the store_c.h +# C ABI, for consumers that dlopen the store (e.g. the Rust crate's `dlopen` +# feature). Whole-archives the static mooncake_store + asio_static (objects are +# -fPIC, see mooncake-common/common.cmake) and localizes everything else via the +# version script; only system/toolkit libs (and the Go etcd/k8s backends when +# enabled) stay DT_NEEDED. +# --------------------------------------------------------------------------- +if(WITH_STORE_C_SHARED) + # ELF/GNU-ld features (--whole-archive, --as-needed, version script, $ORIGIN); + # supported by GNU ld, gold, and lld. Linux only. + if(NOT CMAKE_SYSTEM_NAME STREQUAL "Linux") + message( + FATAL_ERROR + "WITH_STORE_C_SHARED currently supports Linux/ELF only " + "(got CMAKE_SYSTEM_NAME=${CMAKE_SYSTEM_NAME}).") + endif() + + # We whole-archive mooncake_store and reuse its name, so require it to be static + # (add_library(mooncake_store) follows BUILD_SHARED_LIBS, which some backends + # force ON). + get_target_property(_mooncake_store_type mooncake_store TYPE) + if(NOT _mooncake_store_type STREQUAL "STATIC_LIBRARY") + message( + FATAL_ERROR + "WITH_STORE_C_SHARED requires mooncake_store to be a STATIC_LIBRARY " + "(got ${_mooncake_store_type}); configure with BUILD_SHARED_LIBS=OFF.") + endif() + + add_library(mooncake_store_c SHARED store_c_shared.cpp) + set_target_properties( + mooncake_store_c + PROPERTIES OUTPUT_NAME mooncake_store + # Separate dir so this .so does not shadow the static + # libmooncake_store.a on the Rust `link` backend's -L path. + LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/shared" + # $ORIGIN so the installed .so finds co-installed deps (CMake + # strips the build RUNPATH on install). + INSTALL_RPATH "$ORIGIN" + # Relink when the version script (a bare link flag) changes. + LINK_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/store_c.map") + add_dependencies(mooncake_store_c mooncake_store asio_static) + # Whole-archive the static lib + asio_static via file paths (no interface-dep + # propagation); re-list the mooncake_store target for its transitive deps. + # --as-needed drops the now-redundant transitive libasio.so. + target_link_libraries( + mooncake_store_c + PRIVATE -Wl,--whole-archive "$" + "$" -Wl,--no-whole-archive -Wl,--as-needed + mooncake_store) + # Export only the C ABI; the version script localizes every other symbol + # pulled from the static archives. Strip debug info in Release (it dominates + # the .so size). + target_link_options( + mooncake_store_c PRIVATE + "-Wl,--version-script=${CMAKE_CURRENT_SOURCE_DIR}/store_c.map" + "$<$:-Wl,-s>") + install(TARGETS mooncake_store_c DESTINATION lib) +endif() diff --git a/mooncake-store/src/aligned_client_buffer.cpp b/mooncake-store/src/aligned_client_buffer.cpp index 232539d89a..7569868e42 100644 --- a/mooncake-store/src/aligned_client_buffer.cpp +++ b/mooncake-store/src/aligned_client_buffer.cpp @@ -1,4 +1,4 @@ -#include "aligned_client_buffer.hpp" +#include "aligned_client_buffer.h" #include #include @@ -8,13 +8,19 @@ #include "utils.h" +#if defined(USE_SUNRISE) +#include "sunrise_allocator.h" +#endif + namespace mooncake { namespace { constexpr std::string_view kAscendProtocol = "ascend"; constexpr std::string_view kUbshmemProtocol = "ubshmem"; +constexpr std::string_view kSunriseLinkProtocol = "sunrise_link"; bool UseProtocolAllocator(const std::string& protocol) { - return protocol == kAscendProtocol || protocol == kUbshmemProtocol; + return protocol == kAscendProtocol || protocol == kUbshmemProtocol || + protocol == kSunriseLinkProtocol; } void FreeAlignedBuffer(void* buffer, size_t size, const std::string& protocol, @@ -30,6 +36,18 @@ void FreeAlignedBuffer(void* buffer, size_t size, const std::string& protocol, void* AllocateProtocolAlignedBuffer(size_t aligned_size, const std::string& protocol) { +#if defined(USE_SUNRISE) + if (protocol == kSunriseLinkProtocol) { + void* buf = sunrise_allocate_memory( + aligned_size, AlignedClientBufferAllocator::kDirectIOAlignment, + false); + if (!buf) { + LOG(ERROR) << "AlignedClientBufferAllocator: failed to allocate " + << "sunrise_link host memory of size " << aligned_size; + } + return buf; + } +#endif void* aligned_buffer = allocate_buffer_allocator_memory( aligned_size, protocol, AlignedClientBufferAllocator::kDirectIOAlignment); diff --git a/mooncake-store/src/allocator.cpp b/mooncake-store/src/allocator.cpp index 23311b8334..6af07242ea 100644 --- a/mooncake-store/src/allocator.cpp +++ b/mooncake-store/src/allocator.cpp @@ -3,7 +3,9 @@ #include +#include #include +#include #include "master_metric_manager.h" @@ -28,6 +30,17 @@ AllocatedBuffer::~AllocatedBuffer() { } } +AllocatedBuffer::AllocatedBuffer(std::shared_ptr allocator, + const AllocatedBuffer::Descriptor& descriptor) + : allocator_(std::move(allocator)), + buffer_ptr_(reinterpret_cast(descriptor.buffer_address_)), + size_(descriptor.size_), + protocol(descriptor.protocol_) { + if (protocol == "cxl") { + segment_name_ = descriptor.transport_endpoint_; + } +} + // Implementation of get_descriptor AllocatedBuffer::Descriptor AllocatedBuffer::get_descriptor() const { auto alloc = allocator_.lock(); @@ -180,6 +193,64 @@ void CachelibBufferAllocator::deallocate(AllocatedBuffer* handle) { } } +std::unique_ptr CachelibBufferAllocator::adoptImportedBuffer( + const AllocatedBuffer::Descriptor& descriptor) { + cur_size_.fetch_add(descriptor.size_); + if (replica_type_ == ReplicaType::MEMORY) { + MasterMetricManager::instance().inc_allocated_mem_size( + segment_name_, descriptor.size_); + } else if (replica_type_ == ReplicaType::NOF_SSD) { + MasterMetricManager::instance().inc_allocated_nof_size( + segment_name_, descriptor.size_); + } + return std::make_unique(shared_from_this(), descriptor); +} + +std::optional RestoreCachelibBufferAllocator( + std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + const std::vector& descriptors, + ReplicaType replica_type) { + if (replica_type != ReplicaType::MEMORY || + base % facebook::cachelib::Slab::kSize != 0 || + size < facebook::cachelib::Slab::kSize || + size % facebook::cachelib::Slab::kSize != 0 || + base > std::numeric_limits::max() - size) { + return std::nullopt; + } + const size_t end = base + size; + std::vector imports; + imports.reserve(descriptors.size()); + for (const auto& descriptor : descriptors) { + if (descriptor.protocol_ == "cxl" || + descriptor.transport_endpoint_ != transport_endpoint || + descriptor.size_ == 0 || descriptor.size_ > UINT32_MAX || + descriptor.buffer_address_ < base || + descriptor.buffer_address_ >= end || + descriptor.size_ > end - descriptor.buffer_address_) { + return std::nullopt; + } + imports.push_back({reinterpret_cast(descriptor.buffer_address_), + static_cast(std::max( + descriptor.size_, kMinSliceSize))}); + } + + auto allocator = std::make_shared( + std::move(segment_name), base, size, transport_endpoint, replica_type); + if (!allocator->memory_allocator_->importAllocations(allocator->pool_id_, + imports)) { + return std::nullopt; + } + + std::vector> buffers; + buffers.reserve(descriptors.size()); + for (const auto& descriptor : descriptors) { + buffers.push_back(allocator->adoptImportedBuffer(descriptor)); + } + return RestoredCachelibBufferAllocator{std::move(allocator), + std::move(buffers)}; +} + // OffsetBufferAllocator implementation OffsetBufferAllocator::OffsetBufferAllocator(std::string segment_name, size_t base, size_t size, @@ -321,6 +392,94 @@ size_t OffsetBufferAllocator::getLargestFreeRegion() const { } } +std::optional RestoreOffsetBufferAllocator( + std::string segment_name, size_t base, size_t size, + std::string transport_endpoint, + const std::vector& descriptors, + ReplicaType replica_type) { + if (base > std::numeric_limits::max() - size) { + return std::nullopt; + } + const size_t end = base + size; + auto allocator = std::make_shared( + std::move(segment_name), base, size, transport_endpoint, replica_type); + const auto offset_allocator = allocator->getOffsetAllocator(); + + std::vector order(descriptors.size()); + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), [&](size_t lhs, size_t rhs) { + return descriptors[lhs].buffer_address_ < + descriptors[rhs].buffer_address_; + }); + + std::vector> buffers(descriptors.size()); + std::vector> gaps; + size_t cursor = base; + + auto fill_gap = [&](size_t gap) { + while (gap != 0) { + size_t low = 1; + size_t high = gap; + size_t request = 0; + while (low <= high) { + const size_t mid = low + (high - low) / 2; + const uint64_t normalized = + offset_allocator->normalizedAllocationSize(mid); + if (normalized != 0 && normalized <= gap) { + request = mid; + low = mid + 1; + } else { + high = mid - 1; + } + } + if (request == 0) { + return false; + } + const size_t occupied = + offset_allocator->normalizedAllocationSize(request); + if (occupied == 0 || occupied > gap) { + return false; + } + auto filler = allocator->allocate(request); + if (!filler || reinterpret_cast(filler->data()) != cursor) { + return false; + } + cursor += occupied; + gap -= occupied; + gaps.push_back(std::move(filler)); + } + return true; + }; + + for (const size_t index : order) { + const auto& descriptor = descriptors[index]; + if (descriptor.transport_endpoint_ != transport_endpoint || + descriptor.size_ == 0 || descriptor.buffer_address_ < cursor || + descriptor.buffer_address_ < base || + descriptor.buffer_address_ >= end || + descriptor.size_ > end - descriptor.buffer_address_) { + return std::nullopt; + } + const uint64_t occupied = + offset_allocator->normalizedAllocationSize(descriptor.size_); + if (occupied == 0 || occupied > end - descriptor.buffer_address_ || + !fill_gap(descriptor.buffer_address_ - cursor)) { + return std::nullopt; + } + auto buffer = allocator->allocate(descriptor.size_); + if (!buffer || reinterpret_cast(buffer->data()) != + descriptor.buffer_address_) { + return std::nullopt; + } + cursor = descriptor.buffer_address_ + occupied; + buffers[index] = std::move(buffer); + } + + allocator->restored_gap_buffers_ = std::move(gaps); + return RestoredOffsetBufferAllocator{std::move(allocator), + std::move(buffers)}; +} + SimpleAllocator::SimpleAllocator(size_t size) { LOG(INFO) << "initializing_simple_allocator size=" << size; diff --git a/mooncake-store/src/cachelib_memory_allocator/AllocationClass.cpp b/mooncake-store/src/cachelib_memory_allocator/AllocationClass.cpp index af69877d94..bef2bc0153 100644 --- a/mooncake-store/src/cachelib_memory_allocator/AllocationClass.cpp +++ b/mooncake-store/src/cachelib_memory_allocator/AllocationClass.cpp @@ -99,6 +99,38 @@ void* AllocationClass::addSlabAndAllocate(Slab* slab) { })(); } +bool AllocationClass::importSlab( + Slab* slab, const std::vector& occupiedChunkIndexes) { + XDCHECK_NE(nullptr, slab); + std::unique_lock l(lock_); + const uint32_t count = getAllocsPerSlab(); + if (!std::is_sorted(occupiedChunkIndexes.begin(), + occupiedChunkIndexes.end()) || + std::adjacent_find(occupiedChunkIndexes.begin(), + occupiedChunkIndexes.end()) != + occupiedChunkIndexes.end() || + (!occupiedChunkIndexes.empty() && + occupiedChunkIndexes.back() >= count)) { + return false; + } + auto* header = slabAlloc_.getSlabHeader(slab); + header->classId = classId_; + header->allocSize = allocationSize_; + allocatedSlabs_.push_back(slab); + + size_t occupiedPos = 0; + for (uint32_t index = 0; index < count; ++index) { + if (occupiedPos < occupiedChunkIndexes.size() && + occupiedChunkIndexes[occupiedPos] == index) { + ++occupiedPos; + } else { + freedAllocations_.push_back(getAllocForIdx(slab, index)); + } + } + canAllocate_ = !freedAllocations_.empty(); + return true; +} + void* AllocationClass::allocateFromCurrentSlabLocked() noexcept { XDCHECK(canAllocateFromCurrentSlabLocked()); void* ret = currSlab_->memoryAtOffset(currOffset_); diff --git a/mooncake-store/src/cachelib_memory_allocator/MemoryAllocator.cpp b/mooncake-store/src/cachelib_memory_allocator/MemoryAllocator.cpp index a69fff45c4..b6c9a5bca0 100644 --- a/mooncake-store/src/cachelib_memory_allocator/MemoryAllocator.cpp +++ b/mooncake-store/src/cachelib_memory_allocator/MemoryAllocator.cpp @@ -16,6 +16,9 @@ #include "MemoryAllocator.h" +#include +#include + using namespace facebook::cachelib; namespace { @@ -84,6 +87,86 @@ void MemoryAllocator::free(void* memory) { mp.free(memory); } +bool MemoryAllocator::importAllocations( + PoolId id, const std::vector& allocations) { + auto& pool = memoryPoolManager_.getPoolById(id); + if (pool.getCurrentUsedSize() != 0 || pool.getCurrentAllocSize() != 0) { + return false; + } + if (allocations.empty()) { + return true; + } + + struct SlabImport { + ClassId classId{Slab::kInvalidClassId}; + std::vector occupied; + }; + std::map imports; + const uintptr_t base = + reinterpret_cast(slabAllocator_.getSlabForIdx(0)); + const size_t slabCount = slabAllocator_.getNumUsableSlabs(); + if (slabCount > std::numeric_limits::max() / Slab::kSize || + base > std::numeric_limits::max() - slabCount * Slab::kSize) { + return false; + } + const uintptr_t end = base + slabCount * Slab::kSize; + + for (const auto& allocation : allocations) { + const uintptr_t address = reinterpret_cast(allocation.address); + if (allocation.size == 0 || address < base || address >= end) { + return false; + } + ClassId classId; + uint32_t allocSize; + try { + classId = pool.getAllocationClassId(allocation.size); + allocSize = getAllocSize(id, classId); + } catch (...) { + return false; + } + const size_t slabIndex = (address - base) / Slab::kSize; + const size_t slabOffset = (address - base) % Slab::kSize; + const size_t chunkIndex = slabOffset / allocSize; + if (slabOffset % allocSize != 0 || + chunkIndex >= Slab::kSize / allocSize || + allocSize > Slab::kSize - slabOffset) { + return false; + } + auto& import = imports[slabIndex]; + if (import.classId != Slab::kInvalidClassId && import.classId != classId) { + return false; + } + import.classId = classId; + import.occupied.push_back(static_cast(chunkIndex)); + } + + for (auto& [slabIndex, import] : imports) { + std::sort(import.occupied.begin(), import.occupied.end()); + if (std::adjacent_find(import.occupied.begin(), import.occupied.end()) != + import.occupied.end()) { + return false; + } + } + + const size_t highest = imports.rbegin()->first; + for (size_t slabIndex = 0; slabIndex <= highest; ++slabIndex) { + Slab* slab = slabAllocator_.makeNewSlab(id); + if (slab == nullptr || + reinterpret_cast(slab) != base + slabIndex * Slab::kSize) { + return false; + } + auto it = imports.find(slabIndex); + if (it == imports.end()) { + pool.importFreeSlab(slab); + } else { + if (!pool.importSlab(slab, it->second.classId, it->second.occupied)) { + return false; + } + } + } + return true; +} + MemoryPool& MemoryAllocator::getMemoryPool(const void* memory) const { const auto* header = slabAllocator_.getSlabHeader(memory); if (header == nullptr) { diff --git a/mooncake-store/src/cachelib_memory_allocator/MemoryPool.cpp b/mooncake-store/src/cachelib_memory_allocator/MemoryPool.cpp index 3649d62f81..bb422f60b7 100644 --- a/mooncake-store/src/cachelib_memory_allocator/MemoryPool.cpp +++ b/mooncake-store/src/cachelib_memory_allocator/MemoryPool.cpp @@ -175,7 +175,7 @@ ClassId MemoryPool::getAllocationClassId(const void* memory) const { const auto classId = header->classId; if (classId >= static_cast(ac_.size()) || classId < 0) { // at this point, the slab indicates that it belongs to a bogus classId and - // things are corrupt and the caller cant do anything about it. so throw an + // things are corrupt and the caller can't do anything about it. so throw an // exception to abort. throw std::runtime_error(fmt::format( "corrupt slab header/memory pool with class id {}", classId)); @@ -265,6 +265,23 @@ void MemoryPool::free(void* alloc) { currAllocSize_ -= ac.getAllocSize(); } +bool MemoryPool::importSlab(Slab* slab, + ClassId classId, + const std::vector& occupiedChunkIndexes) { + auto& ac = getAllocationClassFor(classId); + if (!ac.importSlab(slab, occupiedChunkIndexes)) { + return false; + } + currSlabAllocSize_ += Slab::kSize; + currAllocSize_ += occupiedChunkIndexes.size() * ac.getAllocSize(); + return true; +} + +void MemoryPool::importFreeSlab(Slab* slab) { + std::unique_lock l(lock_); + freeSlabs_.push_back(slab); +} + void MemoryPool::releaseSlab(SlabReleaseMode mode, const Slab* slab, ClassId receiverClassId) { diff --git a/mooncake-store/src/client_buffer.cpp b/mooncake-store/src/client_buffer.cpp index e846d4c0c6..7ac1dc9212 100644 --- a/mooncake-store/src/client_buffer.cpp +++ b/mooncake-store/src/client_buffer.cpp @@ -1,4 +1,4 @@ -#include "client_buffer.hpp" +#include "client_buffer.h" #include #include @@ -10,6 +10,10 @@ #include "utils.h" +#if defined(USE_SUNRISE) +#include "sunrise_allocator.h" +#endif + namespace mooncake { std::shared_ptr ClientBufferAllocator::create( @@ -43,8 +47,15 @@ ClientBufferAllocator::ClientBufferAllocator(size_t size, if (use_hugepage_) { buffer_ = allocate_buffer_mmap_memory(size, alignment); } else { - buffer_ = allocate_buffer_allocator_memory(size, protocol, alignment, - use_spdk_dma_); +#if defined(USE_SUNRISE) + if (protocol == "sunrise_link") { + buffer_ = sunrise_allocate_memory(size, alignment, false); + } else +#endif + { + buffer_ = allocate_buffer_allocator_memory( + size, protocol, alignment, use_spdk_dma_); + } } if (!buffer_) { throw std::bad_alloc(); @@ -69,7 +80,7 @@ ClientBufferAllocator::~ClientBufferAllocator() { if (use_hugepage_) { free_buffer_mmap_memory(buffer_, buffer_size_); } else { - free_memory(protocol, buffer_); + free_memory(protocol, buffer_, use_spdk_dma_); } } } diff --git a/mooncake-store/src/client_metric.cpp b/mooncake-store/src/client_metric.cpp index 306d479de9..4fbaef026f 100644 --- a/mooncake-store/src/client_metric.cpp +++ b/mooncake-store/src/client_metric.cpp @@ -1,31 +1,23 @@ #include "client_metric.h" #include -#include -#include #include #include #include +#include "bool_parser.h" +#include "integer_parser.h" + namespace mooncake { namespace { -std::string toLower(const std::string& str) { - std::string result = str; - std::transform(result.begin(), result.end(), result.begin(), - [](unsigned char c) { return std::tolower(c); }); - return result; -} - bool parseMetricsEnabled() { const char* metric_env = std::getenv("MC_STORE_CLIENT_METRIC"); if (!metric_env) { return true; } - std::string value = toLower(metric_env); - return (value == "1" || value == "true" || value == "yes" || - value == "on" || value == "enable"); + return TryParseBool(metric_env).value_or(false); } bool parseBoolEnv(const char* env_name, bool default_value) { @@ -34,14 +26,9 @@ bool parseBoolEnv(const char* env_name, bool default_value) { return default_value; } - std::string value = toLower(env_value); - if (value == "1" || value == "true" || value == "yes" || value == "on" || - value == "enable") { - return true; - } - if (value == "0" || value == "false" || value == "no" || value == "off" || - value == "disable") { - return false; + const auto parsed = TryParseBool(env_value); + if (parsed.has_value()) { + return *parsed; } LOG(WARNING) << "Failed to parse " << env_name << ": " << env_value @@ -56,21 +43,22 @@ uint64_t parseMetricsInterval() { return 0; } - try { - uint64_t interval = std::stoull(interval_env); - if (interval == 0) { - LOG(INFO) << "Client metrics reporting disabled (interval=0) via " - "MC_STORE_CLIENT_METRIC_INTERVAL"; - } else { - LOG(INFO) << "Client metrics interval set to " << interval - << "s via MC_STORE_CLIENT_METRIC_INTERVAL"; - } - return interval; - } catch (const std::exception& e) { + const auto interval = TryParseInteger( + interval_env, + {.trim_ascii_whitespace = true, .allow_leading_plus = true}); + if (!interval.has_value()) { LOG(WARNING) << "Failed to parse MC_STORE_CLIENT_METRIC_INTERVAL: " << interval_env << ", disabling metrics reporting"; return 0; } + if (*interval == 0) { + LOG(INFO) << "Client metrics reporting disabled (interval=0) via " + "MC_STORE_CLIENT_METRIC_INTERVAL"; + } else { + LOG(INFO) << "Client metrics interval set to " << *interval + << "s via MC_STORE_CLIENT_METRIC_INTERVAL"; + } + return *interval; } } // anonymous namespace diff --git a/mooncake-store/src/client_service.cpp b/mooncake-store/src/client_service.cpp index 7210e0d52b..137d4b434c 100644 --- a/mooncake-store/src/client_service.cpp +++ b/mooncake-store/src/client_service.cpp @@ -1,7 +1,9 @@ #include "client_service.h" +#include #include +#include "ascii_string.h" #include "allocator.h" #include "segment.h" @@ -24,6 +26,8 @@ #include #include #include +#include +#include #include #include "transfer_engine.h" @@ -33,20 +37,37 @@ #include "config.h" #include "ha/leadership/leader_coordinator_factory.h" #include "types.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "utils.h" #include "rpc_types.h" #include "local_hot_cache.h" -#include "gpu_staging_utils.h" +#include "device/accelerator_registry.h" +#ifdef USE_INTRA_NVLINK +#include "gpu_vendor/intra_nvlink.h" +#endif +#include "crc_checksum.h" +#include "environ.h" namespace mooncake { -using gpu_staging::CopyDeviceToHost; -using gpu_staging::IsDevicePointer; -using gpu_staging::SetDevice; - namespace { +constexpr size_t kObjectChecksumD2HChunkSize = 8 * 1024 * 1024; + +class ScopedObjectChecksumBuffer { + public: + ScopedObjectChecksumBuffer(PinnedBufferPool& pool, size_t size) + : pool_(pool), buffer_(pool_.Acquire(size)) {} + + ~ScopedObjectChecksumBuffer() { pool_.Release(std::move(buffer_)); } + + char* data() const { return buffer_.data; } + + private: + PinnedBufferPool& pool_; + PinnedBufferPool::Buffer buffer_; +}; + #ifdef USE_NOF std::optional GetConfiguredNumaSocketId() { const char* raw_value = std::getenv("MC_STORE_NUMA_SOCKET_ID"); @@ -280,12 +301,20 @@ Client::Client(const std::string& local_hostname, metrics_ ? &metrics_->master_client_metric : nullptr, tenant_id), local_hostname_(local_hostname), + host_id_(ResolveMooncakeHostId(local_hostname)), metadata_connstring_(metadata_connstring), protocol_(protocol), + object_checksum_enabled_(Environ::Get().GetStoreChecksumEnabled()), pinned_buffer_pool_(std::make_unique()), write_thread_pool_(2), task_thread_pool_(4) { LOG(INFO) << "client_id=" << client_id_; + if (!host_id_.empty()) { + LOG(INFO) << "client_id=" << client_id_ << ", host_id=" << host_id_; + } + if (object_checksum_enabled_) { + LOG(INFO) << "Object checksum validation is enabled"; + } if (metrics_) { if (metrics_->GetReportingInterval() > 0) { @@ -390,61 +419,62 @@ Client::~Client() { hot_cache_.reset(); } +ReplicateConfig Client::AttachHostId(const ReplicateConfig& config) const { + ReplicateConfig client_cfg = config; + if (!host_id_.empty()) { + client_cfg.host_id = host_id_; + } + return client_cfg; +} + static std::optional get_auto_discover() { const char* ev_ad = std::getenv("MC_MS_AUTO_DISC"); if (ev_ad) { - int iv = std::stoi(ev_ad); - if (iv == 1) { - LOG(INFO) << "auto discovery set by env MC_MS_AUTO_DISC"; - return true; - } else if (iv == 0) { - LOG(INFO) << "auto discovery not set by env MC_MS_AUTO_DISC"; - return false; - } else { - LOG(WARNING) - << "invalid MC_MS_AUTO_DISC value: " << ev_ad - << ", should be 0 or 1, using default: auto discovery not set"; + try { + int iv = std::stoi(ev_ad); + if (iv == 1) { + LOG(INFO) << "auto discovery set by env MC_MS_AUTO_DISC"; + return true; + } else if (iv == 0) { + LOG(INFO) << "auto discovery not set by env MC_MS_AUTO_DISC"; + return false; + } + } catch (const std::exception&) { + // A non-numeric or out-of-range value makes std::stoi throw; fall + // through to the warning below and use the default instead of + // letting the exception abort client initialization. } + LOG(WARNING) + << "invalid MC_MS_AUTO_DISC value: " << ev_ad + << ", should be 0 or 1, using default: auto discovery not set"; } return std::nullopt; } -static inline void ltrim(std::string& s) { - s.erase(s.begin(), std::find_if(s.begin(), s.end(), [](unsigned char ch) { - return !std::isspace(ch); - })); -} +static std::vector get_auto_discover_filters() { + const char* raw_filters = std::getenv("MC_MS_FILTERS"); + if (raw_filters == nullptr) { + return {}; + } -static inline void rtrim(std::string& s) { - s.erase(std::find_if(s.rbegin(), s.rend(), - [](unsigned char ch) { return !std::isspace(ch); }) - .base(), - s.end()); + LOG(INFO) << "whitelist filters: " << raw_filters; + std::vector filters; + boost::split(filters, std::string(raw_filters), boost::is_any_of(","), + boost::token_compress_off); + for (auto& filter : filters) { + filter = std::string(TrimAsciiWhitespace(filter)); + } + return filters; } -static std::vector get_auto_discover_filters() { - std::vector whitelst_filters; - char* ev_ad = std::getenv("MC_MS_FILTERS"); - if (ev_ad) { - LOG(INFO) << "whitelist filters: " << ev_ad; - char delimiter = ','; - char* end = ev_ad + std::strlen(ev_ad); - char *start = ev_ad, *pos = ev_ad; - while ((pos = std::find(start, end, delimiter)) != end) { - std::string str(start, pos); - ltrim(str); - rtrim(str); - whitelst_filters.emplace_back(std::move(str)); - start = pos + 1; - } - if (start != (end + 1)) { - std::string str(start, end); - ltrim(str); - rtrim(str); - whitelst_filters.emplace_back(std::move(str)); - } - } - return whitelst_filters; +static std::vector ParseDeviceNames(std::string_view value) { + std::vector devices; + boost::split(devices, std::string(value), boost::is_any_of(","), + boost::token_compress_on); + for (auto& device : devices) { + device = std::string(TrimAsciiWhitespace(device)); + } + return devices; } tl::expected, ErrorCode> ParseHABackendSpec( @@ -648,9 +678,10 @@ ErrorCode Client::InitTransferEngine( } // Check if using TENT mode - TENT handles transport configuration - // internally - bool use_tent = (std::getenv("MC_USE_TENT") != nullptr) || - (std::getenv("MC_USE_TEV1") != nullptr); + // internally. Use the engine's own check rather than the raw env var: + // if Mooncake was built without USE_TENT, the env var has no effect on + // the TransferEngine, so the store must still install transports. + bool use_tent = transfer_engine_->isUsingTent(); bool auto_discover = false; if (!use_tent) { @@ -660,17 +691,16 @@ ErrorCode Client::InitTransferEngine( // Use user-specified auto-discover setting auto_discover = env_auto_discover.value(); } else { - // Enable auto-discover for RDMA if no devices are specified + // Enable auto-discover for RDMA/EFA if no devices are specified if ((protocol == "rdma" || protocol == "efa") && !device_names.has_value()) { - LOG(INFO) - << "Set auto discovery ON by default for RDMA protocol, " - "since no " - "device names provided"; + LOG(INFO) << "Set auto discovery ON by default for " << protocol + << " protocol, since no device names provided"; auto_discover = true; } } - transfer_engine_->setAutoDiscover(auto_discover); + transfer_engine_->setAutoDiscover( + {.enabled = auto_discover, .protocol = protocol}); // Honor filters when auto-discovery is enabled; otherwise warn once if (auto_discover) { @@ -696,6 +726,15 @@ ErrorCode Client::InitTransferEngine( globalConfig().ascend_use_fabric_mem = true; } } + if (protocol == "sunrise_link") { + const char* sunrise_use_device_mem_env = + std::getenv("SUNRISE_USE_DEVICE_MEM"); + if (sunrise_use_device_mem_env) { + globalConfig().sunrise_use_device_mem = true; + LOG(INFO) << "SUNRISE_USE_DEVICE_MEM enabled: segments will be " + << "allocated in device memory via tangMalloc"; + } + } auto [hostname, port] = parseHostNameWithPort(local_hostname); int rc = transfer_engine_->init(metadata_connstring, local_hostname, hostname, port); @@ -735,7 +774,7 @@ ErrorCode Client::InitTransferEngine( << device_names.value(); std::vector devices = - splitString(device_names.value(), ',', /*skip_empty=*/true); + ParseDeviceNames(device_names.value()); // Manually discover topology with specified devices only auto topology = transfer_engine_->getLocalTopology(); @@ -770,7 +809,8 @@ ErrorCode Client::InitTransferEngine( LOG(ERROR) << "Failed to install TCP transport"; return ErrorCode::INTERNAL_ERROR; } - } else if (protocol == "ascend" || protocol == "ubshmem") { + } else if (protocol == "ascend" || protocol == "ubshmem" || + protocol == "sunrise_link") { if (device_names.has_value()) { LOG(WARNING) << protocol << " protocol does not use device names, ignoring"; @@ -806,10 +846,17 @@ ErrorCode Client::InitTransferEngine( LOG(ERROR) << "Failed to install CXL transport"; return ErrorCode::INTERNAL_ERROR; } + } else if (protocol == "cxi") { + try { + transport = transfer_engine_->installTransport("cxi", nullptr); + } catch (std::exception& e) { + LOG(ERROR) << "cxi_transport_install_failed error_message=\"" + << e.what() << "\""; + } } else if (protocol == "ub") { auto deviceName = device_names.value_or("bonding_dev_0"); LOG(ERROR) << "ub protocol entable device names is " << deviceName; - auto devices = splitString(deviceName, ',', true); + auto devices = ParseDeviceNames(deviceName); auto topology = transfer_engine_->getLocalTopology(); if (topology) { topology->discover(devices); @@ -823,6 +870,21 @@ ErrorCode Client::InitTransferEngine( "devices"; return ErrorCode::INTERNAL_ERROR; } + } else if (protocol == "nvlink_intra") { +#ifdef USE_INTRA_NVLINK + LOG(INFO) << "Using intra-NVLink protocol."; + transport = + transfer_engine_->installTransport("nvlink_intra", nullptr); + if (!transport) { + LOG(ERROR) << "Failed to install nvlink_intra transport."; + return ErrorCode::INTERNAL_ERROR; + } +#else + LOG(ERROR) + << "--protocol=nvlink_intra requires USE_INTRA_NVLINK=ON, " + "please rebuild mooncake from source."; + return ErrorCode::INVALID_PARAMS; +#endif } else { LOG(ERROR) << "unsupported_protocol protocol=" << protocol; return ErrorCode::INVALID_PARAMS; @@ -884,10 +946,18 @@ std::optional> Client::Create( LOG(INFO) << "Storage root directory is: " << storage_root_dir; LOG(INFO) << "Fs subdir is: " << fs_subdir; // Initialize storage backend with default eviction settings - client->PrepareStorageBackend(storage_root_dir, fs_subdir, true, - 0); + auto prep_err = client->PrepareStorageBackend( + storage_root_dir, fs_subdir, true, 0); + if (prep_err != ErrorCode::OK) { + LOG(ERROR) + << "Failed to initialize storage backend: " << prep_err + << ". Persistence was requested via fsdir but is " + "unavailable."; + return std::nullopt; + } } else { LOG(ERROR) << "Invalid fsdir format: " << dir_string; + return std::nullopt; } } } else { @@ -907,11 +977,19 @@ std::optional> Client::Create( << config.enable_disk_eviction; LOG(INFO) << "Quota bytes: " << config.quota_bytes; // Initialize storage backend with config from master - client->PrepareStorageBackend(storage_root_dir, fs_subdir, - config.enable_disk_eviction, - config.quota_bytes); + auto prep_err = client->PrepareStorageBackend( + storage_root_dir, fs_subdir, config.enable_disk_eviction, + config.quota_bytes); + if (prep_err != ErrorCode::OK) { + LOG(ERROR) + << "Failed to initialize storage backend: " << prep_err + << ". Persistence was requested via storage config " + "but is unavailable."; + return std::nullopt; + } } else { LOG(ERROR) << "Invalid fsdir format: " << config.fsdir; + return std::nullopt; } } } @@ -1029,7 +1107,8 @@ tl::expected Client::Query( } return QueryResult( std::move(result.value().replicas), - start_time + std::chrono::milliseconds(result.value().lease_ttl_ms)); + start_time + std::chrono::milliseconds(result.value().lease_ttl_ms), + result.value().object_checksum); } std::vector> Client::BatchQuery( @@ -1061,8 +1140,9 @@ std::vector> Client::BatchQuery( if (response[i]) { results.emplace_back(QueryResult( std::move(response[i].value().replicas), - start_time + std::chrono::milliseconds( - response[i].value().lease_ttl_ms))); + start_time + + std::chrono::milliseconds(response[i].value().lease_ttl_ms), + response[i].value().object_checksum)); } else { results.emplace_back(tl::unexpected(response[i].error())); } @@ -1070,6 +1150,86 @@ std::vector> Client::BatchQuery( return results; } +tl::expected Client::ComputeObjectChecksumForSlices( + const std::string& object_key, const std::vector& slices, + size_t object_size) { + CrcChecksum checksum; + size_t remaining = object_size; + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); + std::unique_ptr staging; + + for (const auto& slice : slices) { + const size_t bytes = std::min(slice.size, remaining); + if (bytes == 0) { + continue; + } + if (slice.ptr == nullptr) { + LOG(ERROR) << "object_checksum_null_slice key=" << object_key; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + if (!runtime_accelerator.FindDeviceForPointer(slice.ptr)) { + checksum.Update(slice.ptr, bytes); + remaining -= bytes; + continue; + } + + if (!staging) { + staging = std::make_unique( + *pinned_buffer_pool_, + std::min(kObjectChecksumD2HChunkSize, object_size)); + } + size_t offset = 0; + while (offset < bytes) { + const size_t chunk = + std::min(kObjectChecksumD2HChunkSize, bytes - offset); + const auto* source = static_cast(slice.ptr) + offset; + if (!runtime_accelerator.CopyToHost(staging->data(), source, + chunk)) { + LOG(ERROR) << "object_checksum_d2h_failed key=" << object_key + << " size=" << chunk; + return tl::unexpected(ErrorCode::TRANSFER_FAIL); + } + checksum.Update(staging->data(), chunk); + offset += chunk; + } + remaining -= bytes; + } + + if (remaining != 0) { + LOG(ERROR) << "object_checksum_slices_too_small key=" << object_key + << " missing_bytes=" << remaining; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + return checksum.Finalize(); +} + +tl::expected Client::VerifyObjectChecksum( + const std::string& object_key, const std::vector& slices, + size_t object_size, std::optional expected_checksum) { + if (!object_checksum_enabled_) { + return {}; + } + if (!expected_checksum.has_value()) { + VLOG(1) << "object_checksum_absent key=" << object_key; + return {}; + } + + auto actual_checksum = + ComputeObjectChecksumForSlices(object_key, slices, object_size); + if (!actual_checksum) { + return tl::unexpected(actual_checksum.error()); + } + if (*actual_checksum != *expected_checksum) { + LOG(ERROR) << "object_checksum_mismatch key=" << object_key + << " expected=" << *expected_checksum + << " actual=" << *actual_checksum; + return tl::unexpected(ErrorCode::CHECKSUM_MISMATCH); + } + return {}; +} + tl::expected, ErrorCode> Client::BatchReplicaClear( const std::vector& object_keys, const UUID& client_id, const std::string& segment_name) { @@ -1117,6 +1277,13 @@ tl::expected Client::Get(const std::string& object_key, return tl::unexpected(err); } + auto checksum_result = + VerifyObjectChecksum(object_key, slices, calculate_total_size(replica), + query_result.object_checksum); + if (!checksum_result) { + return tl::unexpected(checksum_result.error()); + } + // Frequency admission: only promote frequently accessed keys to hot cache. // Skip when cache_used — data was already served from local cache, no need // to re-promote or increment the CMS counter. @@ -1177,6 +1344,15 @@ tl::expected Client::Get(const std::string& object_key, return {}; } +std::optional Client::SubmitScatter( + const std::vector& transfers) { + if (!transfer_submitter_) { + LOG(ERROR) << "TransferSubmitter not initialized"; + return std::nullopt; + } + return transfer_submitter_->submitScatter(transfers); +} + struct BatchGetOperation { std::vector replicas; std::vector> batched_slices; @@ -1277,7 +1453,6 @@ std::vector> Client::BatchGetWhenPreferSameNode( auto index = op.key_indexes[idx]; VLOG(1) << "Transfer completed successfully for key: " << object_keys[index]; - results[index] = {}; // Release the cache block after transfer completes (memcpy is // done) @@ -1286,6 +1461,16 @@ std::vector> Client::BatchGetWhenPreferSameNode( hot_cache_->ReleaseHotKey(object_keys[index]); } + auto checksum_result = VerifyObjectChecksum( + object_keys[index], op.batched_slices[idx], + calculate_total_size(op.replicas[idx]), + query_results[index].object_checksum); + if (!checksum_result) { + results[index] = tl::unexpected(checksum_result.error()); + continue; + } + results[index] = {}; + // Frequency admission: only promote frequently accessed keys. // Skip when cache was used (data served from local cache). if (idx < op.replicas.size() && @@ -1427,6 +1612,19 @@ std::vector> Client::BatchGet( results[index] = tl::unexpected(result); } else { VLOG(1) << "Transfer completed successfully for key: " << key; + + auto slices_it = slices.find(key); + if (slices_it == slices.end()) { + results[index] = tl::unexpected(ErrorCode::INVALID_PARAMS); + continue; + } + auto checksum_result = VerifyObjectChecksum( + key, slices_it->second, calculate_total_size(stored_replica), + query_results[index].object_checksum); + if (!checksum_result) { + results[index] = tl::unexpected(checksum_result.error()); + continue; + } results[index] = {}; // Frequency admission: only promote frequently accessed keys. @@ -1498,13 +1696,23 @@ bool Client::RedirectToHotCache(const std::string& key, tl::expected Client::Put(const ObjectKey& key, std::vector& slices, const ReplicateConfig& config) { + std::optional object_checksum; + if (object_checksum_enabled_) { + auto checksum_result = ComputeObjectChecksumForSlices( + key, slices, CalculateSliceSize(slices)); + if (!checksum_result) { + return tl::unexpected(checksum_result.error()); + } + object_checksum = *checksum_result; + } + // Prepare slice lengths std::vector slice_lengths; for (size_t i = 0; i < slices.size(); ++i) { slice_lengths.emplace_back(slices[i].size); } - ReplicateConfig client_cfg = config; + ReplicateConfig client_cfg = AttachHostId(config); if (protocol_ == "cxl") { client_cfg.preferred_segment = local_hostname_; } @@ -1580,8 +1788,8 @@ tl::expected Client::Put(const ObjectKey& key, DetermineFinalizeDecision(config, transfer_summary); if (finalize_decision.end_type.has_value()) { - auto end_result = - master_client_.PutEnd(key, *finalize_decision.end_type); + auto end_result = master_client_.PutEnd( + ObjectMeta{key, object_checksum}, *finalize_decision.end_type); if (!end_result) { ErrorCode err = end_result.error(); LOG(ERROR) << "Failed to end put operation: " << err; @@ -1608,13 +1816,23 @@ tl::expected Client::Put(const ObjectKey& key, tl::expected Client::Upsert(const ObjectKey& key, std::vector& slices, const ReplicateConfig& config) { + std::optional object_checksum; + if (object_checksum_enabled_) { + auto checksum_result = ComputeObjectChecksumForSlices( + key, slices, CalculateSliceSize(slices)); + if (!checksum_result) { + return tl::unexpected(checksum_result.error()); + } + object_checksum = *checksum_result; + } + // Prepare slice lengths std::vector slice_lengths; for (size_t i = 0; i < slices.size(); ++i) { slice_lengths.emplace_back(slices[i].size); } - ReplicateConfig client_cfg = config; + ReplicateConfig client_cfg = AttachHostId(config); if (protocol_ == "cxl") { client_cfg.preferred_segment = local_hostname_; } @@ -1678,7 +1896,8 @@ tl::expected Client::Upsert(const ObjectKey& key, } // End upsert operation - auto end_result = master_client_.UpsertEnd(key, ReplicaType::MEMORY); + auto end_result = master_client_.UpsertEnd(ObjectMeta{key, object_checksum}, + ReplicaType::MEMORY); if (!end_result) { ErrorCode err = end_result.error(); LOG(ERROR) << "Failed to end upsert operation: " << err; @@ -1700,7 +1919,7 @@ std::vector> Client::BatchUpsert( const std::vector& keys, std::vector>& batched_slices, const ReplicateConfig& config) { - ReplicateConfig client_cfg = config; + ReplicateConfig client_cfg = AttachHostId(config); if (protocol_ == "cxl") { client_cfg.preferred_segment = local_hostname_; } @@ -1711,6 +1930,7 @@ std::vector> Client::BatchUpsert( } std::vector ops = CreatePutOperations(keys, batched_slices); + ComputeBatchObjectChecksums(ops); StartBatchUpsert(ops, client_cfg); auto t0 = std::chrono::steady_clock::now(); @@ -1755,6 +1975,7 @@ class PutOperation { std::string key; std::vector slices; + std::optional object_checksum; std::vector> batched_slices; // Enhanced state tracking @@ -1851,13 +2072,32 @@ std::vector Client::CreatePutOperations( return ops; } +void Client::ComputeBatchObjectChecksums(std::vector& ops) { + if (!object_checksum_enabled_) { + return; + } + for (auto& op : ops) { + auto checksum_result = ComputeObjectChecksumForSlices( + op.key, op.slices, CalculateSliceSize(op.slices)); + if (!checksum_result) { + op.SetTerminalError(checksum_result.error(), + PutOperationState::MASTER_FAILED, + "Object checksum calculation failed"); + continue; + } + op.object_checksum = *checksum_result; + } +} + void Client::StartBatchPut(std::vector& ops, const ReplicateConfig& config) { std::vector keys; std::vector> slice_lengths; + std::vector active_indices; keys.reserve(ops.size()); slice_lengths.reserve(ops.size()); + active_indices.reserve(ops.size()); if (hot_cache_) { std::vector hot_keys; @@ -1866,7 +2106,12 @@ void Client::StartBatchPut(std::vector& ops, hot_cache_->RemoveHotKeys(hot_keys); } - for (const auto& op : ops) { + for (size_t i = 0; i < ops.size(); ++i) { + const auto& op = ops[i]; + if (op.IsResolved()) { + continue; + } + active_indices.emplace_back(i); keys.emplace_back(op.key); std::vector slice_sizes; @@ -1877,14 +2122,20 @@ void Client::StartBatchPut(std::vector& ops, slice_lengths.emplace_back(std::move(slice_sizes)); } + if (active_indices.empty()) { + return; + } + auto start_responses = master_client_.BatchPutStart(keys, slice_lengths, config); // Ensure response size matches request size - if (start_responses.size() != ops.size()) { + if (start_responses.size() != active_indices.size()) { LOG(ERROR) << "BatchPutStart response size mismatch: expected " - << ops.size() << ", got " << start_responses.size(); - for (auto& op : ops) { + << active_indices.size() << ", got " + << start_responses.size(); + for (size_t index : active_indices) { + auto& op = ops[index]; op.SetError(ErrorCode::RPC_FAIL, "BatchPutStart response size mismatch"); } @@ -1892,27 +2143,27 @@ void Client::StartBatchPut(std::vector& ops, } // Process individual responses with robust error handling - for (size_t i = 0; i < ops.size(); ++i) { - ops[i].InitializeRequestedReplicas(config); + for (size_t i = 0; i < active_indices.size(); ++i) { + auto& op = ops[active_indices[i]]; + op.InitializeRequestedReplicas(config); if (!start_responses[i]) { - ops[i].SetTerminalError(start_responses[i].error(), - PutOperationState::MASTER_FAILED, - "Master failed to start put operation"); + op.SetTerminalError(start_responses[i].error(), + PutOperationState::MASTER_FAILED, + "Master failed to start put operation"); } else { - ops[i].replicas = start_responses[i].value(); - ops[i].RecordAllocatedReplicas(); - if (!HasExpectedReplicaAllocation(config, - ops[i].transfer_summary)) { - ops[i].SetTerminalError(ErrorCode::NO_AVAILABLE_HANDLE, - PutOperationState::MASTER_FAILED, - "Allocated replicas do not satisfy " - "requested replica policy"); + op.replicas = start_responses[i].value(); + op.RecordAllocatedReplicas(); + if (!HasExpectedReplicaAllocation(config, op.transfer_summary)) { + op.SetTerminalError(ErrorCode::NO_AVAILABLE_HANDLE, + PutOperationState::MASTER_FAILED, + "Allocated replicas do not satisfy " + "requested replica policy"); continue; } // Operation continues to next stage - result remains INTERNAL_ERROR // until fully successful - VLOG(1) << "Successfully started put for key " << ops[i].key - << " with " << ops[i].replicas.size() << " replicas"; + VLOG(1) << "Successfully started put for key " << op.key << " with " + << op.replicas.size() << " replicas"; } } } @@ -1921,9 +2172,11 @@ void Client::StartBatchUpsert(std::vector& ops, const ReplicateConfig& config) { std::vector keys; std::vector> slice_lengths; + std::vector active_indices; keys.reserve(ops.size()); slice_lengths.reserve(ops.size()); + active_indices.reserve(ops.size()); if (hot_cache_) { std::vector hot_keys; @@ -1932,7 +2185,12 @@ void Client::StartBatchUpsert(std::vector& ops, hot_cache_->RemoveHotKeys(hot_keys); } - for (const auto& op : ops) { + for (size_t i = 0; i < ops.size(); ++i) { + const auto& op = ops[i]; + if (op.IsResolved()) { + continue; + } + active_indices.emplace_back(i); keys.emplace_back(op.key); std::vector slice_sizes; @@ -1943,14 +2201,20 @@ void Client::StartBatchUpsert(std::vector& ops, slice_lengths.emplace_back(std::move(slice_sizes)); } + if (active_indices.empty()) { + return; + } + auto start_responses = master_client_.BatchUpsertStart(keys, slice_lengths, config); // Ensure response size matches request size - if (start_responses.size() != ops.size()) { + if (start_responses.size() != active_indices.size()) { LOG(ERROR) << "BatchUpsertStart response size mismatch: expected " - << ops.size() << ", got " << start_responses.size(); - for (auto& op : ops) { + << active_indices.size() << ", got " + << start_responses.size(); + for (size_t index : active_indices) { + auto& op = ops[index]; op.SetError(ErrorCode::RPC_FAIL, "BatchUpsertStart response size mismatch"); } @@ -1958,22 +2222,30 @@ void Client::StartBatchUpsert(std::vector& ops, } // Process individual responses with robust error handling - for (size_t i = 0; i < ops.size(); ++i) { + for (size_t i = 0; i < active_indices.size(); ++i) { + auto& op = ops[active_indices[i]]; if (!start_responses[i]) { - ops[i].SetError(start_responses[i].error(), - "Master failed to start upsert operation"); + op.SetError(start_responses[i].error(), + "Master failed to start upsert operation"); } else { - ops[i].replicas = start_responses[i].value(); - VLOG(1) << "Successfully started upsert for key " << ops[i].key - << " with " << ops[i].replicas.size() << " replicas"; + op.replicas = start_responses[i].value(); + VLOG(1) << "Successfully started upsert for key " << op.key + << " with " << op.replicas.size() << " replicas"; } } } void Client::SubmitTransfers(std::vector& ops) { + if (std::all_of(ops.begin(), ops.end(), + [](const PutOperation& op) { return op.IsResolved(); })) { + return; + } if (!transfer_submitter_) { LOG(ERROR) << "TransferSubmitter not initialized"; for (auto& op : ops) { + if (op.IsResolved()) { + continue; + } op.SetTerminalError(ErrorCode::INVALID_PARAMS, PutOperationState::TRANSFER_FAILED, "TransferSubmitter not initialized"); @@ -2179,7 +2451,13 @@ void Client::FinalizeBatchPut(std::vector& ops) { if (group.keys.empty()) { return; } - auto responses = master_client_.BatchPutEnd(group.keys, replica_type); + std::vector object_metas; + object_metas.reserve(group.indices.size()); + for (size_t i = 0; i < group.indices.size(); ++i) { + object_metas.emplace_back(ObjectMeta{ + group.keys[i], ops[group.indices[i]].object_checksum}); + } + auto responses = master_client_.BatchPutEnd(object_metas, replica_type); if (responses.size() != group.keys.size()) { for (size_t idx : group.indices) { finalize_rpc_errors[idx] = ErrorCode::RPC_FAIL; @@ -2285,12 +2563,12 @@ void Client::FinalizeBatchPut(std::vector& ops) { } void Client::FinalizeBatchUpsert(std::vector& ops) { - std::vector successful_keys; + std::vector successful_object_metas; std::vector successful_indices; std::vector failed_keys; std::vector failed_indices; - successful_keys.reserve(ops.size()); + successful_object_metas.reserve(ops.size()); successful_indices.reserve(ops.size()); failed_keys.reserve(ops.size()); failed_indices.reserve(ops.size()); @@ -2300,7 +2578,8 @@ void Client::FinalizeBatchUpsert(std::vector& ops) { if (!op.IsResolved() && !op.replicas.empty() && !op.pending_transfers.empty()) { - successful_keys.emplace_back(op.key); + successful_object_metas.emplace_back( + ObjectMeta{op.key, op.object_checksum}); successful_indices.emplace_back(i); } else if (op.state != PutOperationState::PENDING && !op.replicas.empty()) { @@ -2311,12 +2590,13 @@ void Client::FinalizeBatchUpsert(std::vector& ops) { // Process successful operations std::vector finalized_keys; - if (!successful_keys.empty()) { - finalized_keys.reserve(successful_keys.size()); - auto end_responses = master_client_.BatchUpsertEnd(successful_keys); - if (end_responses.size() != successful_keys.size()) { + if (!successful_object_metas.empty()) { + finalized_keys.reserve(successful_object_metas.size()); + auto end_responses = + master_client_.BatchUpsertEnd(successful_object_metas); + if (end_responses.size() != successful_object_metas.size()) { LOG(ERROR) << "BatchUpsertEnd response size mismatch: expected " - << successful_keys.size() << ", got " + << successful_object_metas.size() << ", got " << end_responses.size(); for (size_t idx : successful_indices) { ops[idx].SetError(ErrorCode::RPC_FAIL, @@ -2327,15 +2607,15 @@ void Client::FinalizeBatchUpsert(std::vector& ops) { const size_t op_idx = successful_indices[i]; if (!end_responses[i]) { LOG(ERROR) << "Failed to finalize upsert for key " - << successful_keys[i] << ": " + << successful_object_metas[i].key << ": " << toString(end_responses[i].error()); ops[op_idx].SetError(end_responses[i].error(), "BatchUpsertEnd failed"); } else { ops[op_idx].SetSuccess(); - finalized_keys.emplace_back(successful_keys[i]); + finalized_keys.emplace_back(successful_object_metas[i].key); VLOG(1) << "Successfully completed upsert for key " - << successful_keys[i]; + << successful_object_metas[i].key; } } } @@ -2536,11 +2816,12 @@ std::vector> Client::BatchPut( const std::vector& keys, std::vector>& batched_slices, const ReplicateConfig& config) { - ReplicateConfig client_cfg = config; + ReplicateConfig client_cfg = AttachHostId(config); if (protocol_ == "cxl") { client_cfg.preferred_segment = local_hostname_; } std::vector ops = CreatePutOperations(keys, batched_slices); + ComputeBatchObjectChecksums(ops); if (client_cfg.prefer_alloc_in_same_node) { if (client_cfg.nof_replica_num > 0) { LOG(ERROR) << "prefer_alloc_in_same_node is not supported with " @@ -2619,8 +2900,10 @@ tl::expected Client::RemoveAll(bool force) { } auto result = master_client_.RemoveAll(force); - if (result && storage_backend_) { - storage_backend_->RemoveAll(); + if (result) { + if (storage_backend_) { + storage_backend_->RemoveAll(); + } } if (result && result.value() > 0 && hot_cache_) { hot_cache_->RemoveAllHotKeys(); @@ -2786,6 +3069,7 @@ tl::expected Client::MountSegmentAndGetId( segment.base = reinterpret_cast(buffer); segment.size = size; segment.protocol = protocol; + segment.host_id = host_id_; if (metadata_connstring_ == P2PHANDSHAKE) { segment.te_endpoint = transfer_engine_->getLocalIpAndPort(); } else { @@ -3009,6 +3293,10 @@ tl::expected Client::OffloadObjectHeartbeat( return {}; } +tl::expected Client::PollRemoveAll() { + return master_client_.PollRemoveAll(); +} + tl::expected Client::ReportSsdCapacity( int64_t ssd_total_capacity_bytes) { auto response = @@ -3308,20 +3596,32 @@ tl::expected Client::MarkTaskToComplete( return master_client_.MarkTaskToComplete(update_request); } -void Client::PrepareStorageBackend(const std::string& storage_root_dir, - const std::string& fsdir, - bool enable_eviction, uint64_t quota_bytes) { +ErrorCode Client::PrepareStorageBackend(const std::string& storage_root_dir, + const std::string& fsdir, + bool enable_eviction, + uint64_t quota_bytes) { // Initialize storage backend - storage_backend_ = + auto backend_result = StorageBackend::Create(storage_root_dir, fsdir, enable_eviction); - if (!storage_backend_) { - LOG(INFO) << "Failed to initialize storage backend"; - } - auto init_result = storage_backend_->Init(quota_bytes); + if (!backend_result) { + LOG(ERROR) << "Failed to create storage backend: " + << backend_result.error(); + return backend_result.error(); + } + // Initialize into a local first and only publish to storage_backend_ + // after a successful Init(): users of storage_backend_ only null-check + // it, so it must never point to a backend whose Init() failed (using it + // before successful Init() is undefined behavior). If Init() throws, + // stack unwinding destroys the local and storage_backend_ stays clean. + auto backend = std::move(backend_result.value()); + auto init_result = backend->Init(quota_bytes); if (!init_result) { LOG(ERROR) << "Failed to initialize StorageBackend. Error: " << init_result.error() << ". The backend will be unusable."; + return init_result.error(); } + storage_backend_ = std::move(backend); + return ErrorCode::OK; } void Client::PutToLocalFile(const std::string& key, @@ -3341,16 +3641,21 @@ void Client::PutToLocalFile(const std::string& key, // (BatchPut has not yet returned to Python, so blocks are not reused). std::string value; value.reserve(total_size); + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); for (const auto& slice : slices) { - int device_id = -1; - if (IsDevicePointer(slice.ptr, &device_id)) { - SetDevice(device_id); + device::PointerInfo info{}; + auto* device = + runtime_accelerator.FindDeviceForPointer(slice.ptr, &info); + if (device) { + device->SetContext(info.device_id); auto buf = pinned_buffer_pool_->Acquire(slice.size); - if (!CopyDeviceToHost(buf.data, slice.ptr, slice.size)) { + if (!device->Copy(buf.data, slice.ptr, slice.size, + device::CopyDirection::kDeviceToHost)) { LOG(ERROR) << "D2H copy failed for key: " << key << ", triggering PutRevoke for disk replica"; - pinned_buffer_pool_->Release(buf); + pinned_buffer_pool_->Release(std::move(buf)); // Must revoke to avoid phantom replica in master auto revoke_result = master_client_.PutRevoke(key, ReplicaType::DISK); @@ -3361,7 +3666,7 @@ void Client::PutToLocalFile(const std::string& key, return; } value.append(buf.data, slice.size); - pinned_buffer_pool_->Release(buf); + pinned_buffer_pool_->Release(std::move(buf)); } else { value.append(static_cast(slice.ptr), slice.size); } @@ -3370,9 +3675,34 @@ void Client::PutToLocalFile(const std::string& key, // Async StoreObject + PutEnd (unchanged from original) write_thread_pool_.enqueue([this, backend = storage_backend_, key, value = std::move(value), path] { - // Store the object - auto store_result = backend->StoreObject(path, value, key); ReplicaType replica_type = ReplicaType::DISK; + // Store the object + auto store_result = backend->StoreObject( + path, value, key, + [this, replica_type](const std::vector& evicted_keys) + -> tl::expected { + auto evict_results = master_client_.BatchEvictDiskReplica( + evicted_keys, replica_type); + if (evict_results.size() != evicted_keys.size()) { + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + for (size_t i = 0; i < evict_results.size(); ++i) { + if (!evict_results[i]) { + if (evict_results[i].error() == + ErrorCode::OBJECT_NOT_FOUND) { + VLOG(1) << "Master no longer tracks evicted key: " + << evicted_keys[i]; + continue; + } + LOG(WARNING) + << "Failed to notify master about evicted key: " + << evicted_keys[i] + << ", error: " << evict_results[i].error(); + return tl::make_unexpected(evict_results[i].error()); + } + } + return {}; + }); if (!store_result) { // If storage failed, revoke the put operation @@ -3384,23 +3714,9 @@ void Client::PutToLocalFile(const std::string& key, return; } - // Notify master about any evicted disk replicas (batch) - if (!store_result.value().empty()) { - const auto& evicted_keys = store_result.value(); - auto evict_results = master_client_.BatchEvictDiskReplica( - evicted_keys, replica_type); - for (size_t i = 0; i < evict_results.size(); ++i) { - if (!evict_results[i]) { - LOG(WARNING) - << "Failed to notify master about evicted key: " - << evicted_keys[i] - << ", error: " << evict_results[i].error(); - } - } - } - // If storage succeeded, end the put operation - auto end_result = master_client_.PutEnd(key, replica_type); + auto end_result = + master_client_.PutEnd(ObjectMeta{key, std::nullopt}, replica_type); if (!end_result) { LOG(ERROR) << "Failed to end put operation for key: " << key; } @@ -3967,6 +4283,10 @@ ErrorCode Client::InitLocalHotCache() { hot_cache_.reset(); admission_sketch_.reset(); + if (object_checksum_enabled_) { + return ErrorCode::OK; + } + // Defaults: hot cache is disabled unless MC_STORE_LOCAL_HOT_CACHE_SIZE is // set to a positive value; when enabled, default block size is 16MB and // thread_num is 2. diff --git a/mooncake-store/src/device/accelerator_device.cpp b/mooncake-store/src/device/accelerator_device.cpp new file mode 100644 index 0000000000..4f6653b28c --- /dev/null +++ b/mooncake-store/src/device/accelerator_device.cpp @@ -0,0 +1,25 @@ +#include "device/accelerator_device.h" + +namespace mooncake { +namespace device { + +namespace { + +constexpr uint8_t kAvailableProbed = 1 << 0; +constexpr uint8_t kAvailable = 1 << 1; + +} // namespace + +bool ProbeCachedAcceleratorDevice::Available(bool ensure) const { + uint8_t state = available_state_.load(std::memory_order_acquire); + if (!ensure) { + return (state & kAvailableProbed) ? (state & kAvailable) : true; + } + const bool available = ProbeAvailable(); + state = kAvailableProbed | (available ? kAvailable : 0); + available_state_.store(state, std::memory_order_release); + return available; +} + +} // namespace device +} // namespace mooncake diff --git a/mooncake-store/src/device/accelerator_registry.cpp b/mooncake-store/src/device/accelerator_registry.cpp new file mode 100644 index 0000000000..6375f4c2f0 --- /dev/null +++ b/mooncake-store/src/device/accelerator_registry.cpp @@ -0,0 +1,112 @@ +#include "device/accelerator_registry.h" + +#include +#include +#include +#include + +namespace mooncake { +namespace device { + +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ + defined(USE_HYGON) || defined(USE_COREX) +void EnsureCudaLikeAcceleratorDeviceLinked(); +#endif + +namespace { + +void RegisterStaticAcceleratorDevice(const AcceleratorDevice& device); + +class AcceleratorRegistryImpl final : public AcceleratorRegistry { + using DeviceList = std::vector; + using DeviceListPtr = std::shared_ptr; + + public: + std::span RegisteredDevices() + const override { + return std::span( + registered_devices_.data(), registered_devices_.size()); + } + + RuntimeAccelerator RuntimeAccelerators(bool ensure = false) const override { + auto available_devices = std::atomic_load(&available_devices_); + if (ShouldRefresh(ensure, available_devices)) { + std::lock_guard lock(refresh_mutex_); + available_devices = std::atomic_load(&available_devices_); + if (ShouldRefresh(ensure, available_devices)) { + available_devices = BuildAvailableDevices(); + std::atomic_store(&available_devices_, available_devices); + } + } + return RuntimeAccelerator(*available_devices); + } + + const AcceleratorDevice* GetDevice( + AcceleratorVendor vendor) const override { + for (auto* device : registered_devices_) { + if (device->Vendor() == vendor) return device; + } + return nullptr; + } + + private: + void Register(const AcceleratorDevice& device) { + for (auto*& registered_device : registered_devices_) { + if (registered_device->Vendor() == device.Vendor()) { + registered_device = &device; + return; + } + } + registered_devices_.push_back(&device); + std::atomic_store(&available_devices_, DeviceListPtr()); + } + + static bool ShouldRefresh(bool ensure, + const DeviceListPtr& available_devices) { + return ensure || !available_devices || available_devices->empty(); + } + + DeviceListPtr BuildAvailableDevices() const { + auto available_devices = std::make_shared(); + for (auto* device : registered_devices_) { + if (device->Available(true)) { + available_devices->push_back(device); + } + } + return available_devices; + } + + DeviceList registered_devices_; + mutable std::mutex refresh_mutex_; + mutable std::shared_ptr available_devices_; + + friend void RegisterStaticAcceleratorDevice( + const AcceleratorDevice& device); +}; + +AcceleratorRegistryImpl& MutableRegistry() { + static AcceleratorRegistryImpl registry; + return registry; +} + +void RegisterStaticAcceleratorDevice(const AcceleratorDevice& device) { + MutableRegistry().Register(device); +} + +} // namespace + +const AcceleratorRegistry& GetAcceleratorRegistry() { +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ + defined(USE_HYGON) || defined(USE_COREX) + EnsureCudaLikeAcceleratorDeviceLinked(); +#endif + return MutableRegistry(); +} + +AcceleratorDeviceRegistrar::AcceleratorDeviceRegistrar( + const AcceleratorDevice& device) { + RegisterStaticAcceleratorDevice(device); +} + +} // namespace device +} // namespace mooncake diff --git a/mooncake-store/src/device/ascend_accelerator_device.cpp b/mooncake-store/src/device/ascend_accelerator_device.cpp new file mode 100644 index 0000000000..f084732478 --- /dev/null +++ b/mooncake-store/src/device/ascend_accelerator_device.cpp @@ -0,0 +1,106 @@ +#include "device/accelerator_registry.h" +#include "pinned_host_buffer.h" + +#if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) +#include +#endif + +#if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) + +namespace mooncake { +namespace device { +namespace { + +void FreeAscendPinnedHostBuffer(void* addr) { aclrtFreeHost(addr); } + +class AscendAcceleratorDevice final : public ProbeCachedAcceleratorDevice { + public: + AcceleratorVendor Vendor() const override { + return AcceleratorVendor::kAscend; + } + + bool ProbeAvailable() const override { + uint32_t count = 0; + return aclrtGetDeviceCount(&count) == ACL_SUCCESS && count > 0; + } + + PointerInfo QueryPointer(const void* ptr) const override { + aclrtPtrAttributes attr{}; + if (aclrtPointerGetAttributes(const_cast(ptr), &attr) == + ACL_SUCCESS && + attr.location.type == ACL_MEM_LOCATION_TYPE_DEVICE) { + return PointerInfo{ + .kind = MemoryKind::kDevice, + .device_id = static_cast(attr.location.id), + }; + } + return PointerInfo{.kind = MemoryKind::kHost, .device_id = -1}; + } + + int32_t CurrentDeviceId() const override { + int32_t logic_dev = 0; + if (aclrtGetDevice(&logic_dev) != ACL_SUCCESS) { + return -1; + } + return logic_dev; + } + + void SetContext(int32_t device_id) const override { + if (device_id >= 0) aclrtSetDevice(device_id); + } + + bool Copy(void* dst, const void* src, size_t size, + CopyDirection direction) const override { + aclrtMemcpyKind kind = ACL_MEMCPY_HOST_TO_HOST; + switch (direction) { + case CopyDirection::kHostToDevice: + kind = ACL_MEMCPY_HOST_TO_DEVICE; + break; + case CopyDirection::kDeviceToHost: + kind = ACL_MEMCPY_DEVICE_TO_HOST; + break; + case CopyDirection::kDeviceToDevice: + kind = ACL_MEMCPY_DEVICE_TO_DEVICE; + break; + case CopyDirection::kHostToHost: + kind = ACL_MEMCPY_HOST_TO_HOST; + break; + case CopyDirection::kAuto: { + auto src_info = QueryPointer(src); + auto dst_info = QueryPointer(dst); + if (src_info.kind == MemoryKind::kDevice && + dst_info.kind == MemoryKind::kDevice) { + kind = ACL_MEMCPY_DEVICE_TO_DEVICE; + } else if (src_info.kind == MemoryKind::kDevice) { + kind = ACL_MEMCPY_DEVICE_TO_HOST; + } else if (dst_info.kind == MemoryKind::kDevice) { + kind = ACL_MEMCPY_HOST_TO_DEVICE; + } + break; + } + } + return aclrtMemcpy(dst, size, src, size, kind) == ACL_SUCCESS; + } + + PinnedHostBuffer AllocatePinnedHost(size_t size) const override { + void* addr = nullptr; + if (aclrtMallocHost(&addr, size) != ACL_SUCCESS) { + return PinnedHostBuffer(); + } + return PinnedHostBuffer(addr, size, FreeAscendPinnedHostBuffer); + } +}; + +const AcceleratorDevice& AscendDeviceInstance() { + static AscendAcceleratorDevice device; + return device; +} + +const AcceleratorDeviceRegistrar registered_ascend_device( + AscendDeviceInstance()); + +} // namespace +} // namespace device +} // namespace mooncake + +#endif diff --git a/mooncake-store/src/device/cuda_ipc_buffer.cpp b/mooncake-store/src/device/cuda_ipc_buffer.cpp new file mode 100644 index 0000000000..5953ff57d4 --- /dev/null +++ b/mooncake-store/src/device/cuda_ipc_buffer.cpp @@ -0,0 +1,214 @@ +#include "device/cuda_ipc_buffer.h" + +#include +#include +#include + +#if defined(USE_CUDA) +#include + +#include "cuda_alike.h" +#endif + +namespace mooncake { +namespace device { +namespace { + +tl::expected UnsupportedCudaIpc() { + return tl::unexpected(ErrorCode::INVALID_PARAMS); +} + +#if defined(USE_CUDA) +bool AddOverflows(uint64_t a, uint64_t b) { + return a > std::numeric_limits::max() - b; +} + +void ClearCudaError() { cudaGetLastError(); } + +using CuMemGetAddressRangeFn = CUresult (*)(CUdeviceptr *, size_t *, + CUdeviceptr); + +CuMemGetAddressRangeFn LoadCuMemGetAddressRange() { + static CuMemGetAddressRangeFn fn = [] { + void *handle = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); + if (handle == nullptr) { + handle = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL); + } + void *symbol = handle != nullptr + ? dlsym(handle, "cuMemGetAddressRange_v2") + : dlsym(RTLD_DEFAULT, "cuMemGetAddressRange_v2"); + if (symbol == nullptr) { + symbol = handle != nullptr + ? dlsym(handle, "cuMemGetAddressRange") + : dlsym(RTLD_DEFAULT, "cuMemGetAddressRange"); + } + return reinterpret_cast(symbol); + }(); + return fn; +} + +CUresult GetCudaAllocationRange(CUdeviceptr *base, size_t *size, + CUdeviceptr ptr) { + auto fn = LoadCuMemGetAddressRange(); + if (fn == nullptr) return CUDA_ERROR_NOT_FOUND; + return fn(base, size, ptr); +} +#endif + +} // namespace + +tl::expected ExportCudaIpcBuffer( + const void *ptr, size_t size) { +#if defined(USE_CUDA) + static_assert(sizeof(cudaIpcMemHandle_t) == kCudaIpcHandleSize, + "Unexpected CUDA IPC handle size"); + + if (ptr == nullptr || size == 0) return UnsupportedCudaIpc(); + + cudaPointerAttributes attr{}; + if (cudaPointerGetAttributes(&attr, ptr) != cudaSuccess || + attr.type != cudaMemoryTypeDevice || attr.devicePointer == nullptr) { + ClearCudaError(); + return UnsupportedCudaIpc(); + } + + // PyTorch may hand out suballocated pointers; CUDA IPC must export the + // allocation base and carry the caller pointer offset separately. + const auto ptr_addr = reinterpret_cast(ptr); + CUdeviceptr base_ptr = 0; + size_t allocation_size = 0; + CUresult cu_ret = + GetCudaAllocationRange(&base_ptr, &allocation_size, (CUdeviceptr)ptr); + if (cu_ret != CUDA_SUCCESS || base_ptr == 0 || allocation_size == 0) { + ClearCudaError(); + return UnsupportedCudaIpc(); + } + + const auto base_addr = static_cast(base_ptr); + if (ptr_addr < base_addr) return UnsupportedCudaIpc(); + + const uint64_t offset = static_cast(ptr_addr - base_addr); + const uint64_t payload_size = static_cast(size); + if (AddOverflows(offset, payload_size) || + offset + payload_size > allocation_size) { + return UnsupportedCudaIpc(); + } + + int current_device = -1; + cudaGetDevice(¤t_device); + if (cudaSetDevice(attr.device) != cudaSuccess) { + ClearCudaError(); + return tl::unexpected(ErrorCode::INTERNAL_ERROR); + } + + cudaIpcMemHandle_t ipc_handle{}; + cudaError_t ret = cudaIpcGetMemHandle(&ipc_handle, (void *)base_ptr); + if (current_device >= 0) cudaSetDevice(current_device); + if (ret != cudaSuccess) { + ClearCudaError(); + return UnsupportedCudaIpc(); + } + + CudaIpcBufferHandle exported; + std::memcpy(exported.handle.data(), &ipc_handle, sizeof(ipc_handle)); + exported.offset = offset; + exported.size = payload_size; + exported.device_id = attr.device; + return exported; +#else + (void)ptr; + (void)size; + return UnsupportedCudaIpc(); +#endif +} + +CudaIpcBufferMapping::~CudaIpcBufferMapping() { Close(); } + +CudaIpcBufferMapping::CudaIpcBufferMapping( + CudaIpcBufferMapping &&other) noexcept + : base_(std::exchange(other.base_, nullptr)), + ptr_(std::exchange(other.ptr_, nullptr)), + device_id_(std::exchange(other.device_id_, -1)) {} + +CudaIpcBufferMapping &CudaIpcBufferMapping::operator=( + CudaIpcBufferMapping &&other) noexcept { + if (this != &other) { + Close(); + base_ = std::exchange(other.base_, nullptr); + ptr_ = std::exchange(other.ptr_, nullptr); + device_id_ = std::exchange(other.device_id_, -1); + } + return *this; +} + +tl::expected CudaIpcBufferMapping::Open( + const CudaIpcBufferHandle &handle) { +#if defined(USE_CUDA) + static_assert(sizeof(cudaIpcMemHandle_t) == kCudaIpcHandleSize, + "Unexpected CUDA IPC handle size"); + + if (handle.size == 0 || handle.device_id < 0 || + AddOverflows(handle.offset, handle.size)) { + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + int current_device = -1; + cudaGetDevice(¤t_device); + if (cudaSetDevice(handle.device_id) != cudaSuccess) { + ClearCudaError(); + return tl::unexpected(ErrorCode::INTERNAL_ERROR); + } + + cudaIpcMemHandle_t ipc_handle{}; + std::memcpy(&ipc_handle, handle.handle.data(), sizeof(ipc_handle)); + + void *base = nullptr; + cudaError_t ret = + cudaIpcOpenMemHandle(&base, ipc_handle, cudaIpcMemLazyEnablePeerAccess); + CUdeviceptr allocation_base = 0; + size_t allocation_size = 0; + if (ret == cudaSuccess && base != nullptr) { + CUresult cu_ret = GetCudaAllocationRange( + &allocation_base, &allocation_size, (CUdeviceptr)base); + if (cu_ret != CUDA_SUCCESS || allocation_base != (CUdeviceptr)base || + allocation_size == 0 || AddOverflows(handle.offset, handle.size) || + handle.offset + handle.size > allocation_size) { + cudaIpcCloseMemHandle(base); + base = nullptr; + ret = cudaErrorInvalidValue; + } + } + if (current_device >= 0) cudaSetDevice(current_device); + if (ret != cudaSuccess || base == nullptr) { + ClearCudaError(); + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + auto *ptr = static_cast(base) + handle.offset; + return CudaIpcBufferMapping(base, ptr, handle.device_id); +#else + (void)handle; + return tl::unexpected(ErrorCode::INVALID_PARAMS); +#endif +} + +void CudaIpcBufferMapping::Close() { +#if defined(USE_CUDA) + if (base_ != nullptr) { + int current_device = -1; + const bool restore_device = + cudaGetDevice(¤t_device) == cudaSuccess; + if (device_id_ >= 0 && cudaSetDevice(device_id_) != cudaSuccess) { + ClearCudaError(); + } + cudaIpcCloseMemHandle(base_); + if (restore_device) cudaSetDevice(current_device); + } +#endif + base_ = nullptr; + ptr_ = nullptr; + device_id_ = -1; +} + +} // namespace device +} // namespace mooncake diff --git a/mooncake-store/src/device/cuda_like_accelerator_device.cpp b/mooncake-store/src/device/cuda_like_accelerator_device.cpp new file mode 100644 index 0000000000..385b548542 --- /dev/null +++ b/mooncake-store/src/device/cuda_like_accelerator_device.cpp @@ -0,0 +1,114 @@ +#include "device/accelerator_registry.h" +#include "pinned_host_buffer.h" + +#include "cuda_alike.h" + +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA) || \ + defined(USE_HYGON) || defined(USE_COREX) + +namespace mooncake { +namespace device { + +void EnsureCudaLikeAcceleratorDeviceLinked() {} + +namespace { + +void FreeCudaLikePinnedHostBuffer(void* addr) { cudaFreeHost(addr); } + +class CudaLikeAcceleratorDevice final : public ProbeCachedAcceleratorDevice { + public: + explicit CudaLikeAcceleratorDevice(AcceleratorVendor vendor) + : vendor_(vendor) {} + + AcceleratorVendor Vendor() const override { return vendor_; } + + bool ProbeAvailable() const override { + int count = 0; + return cudaGetDeviceCount(&count) == cudaSuccess && count > 0; + } + + PointerInfo QueryPointer(const void* ptr) const override { + cudaPointerAttributes attr{}; + if (cudaPointerGetAttributes(&attr, ptr) == cudaSuccess && + attr.type == cudaMemoryTypeDevice) { + return PointerInfo{.kind = MemoryKind::kDevice, + .device_id = attr.device}; + } + cudaGetLastError(); + return PointerInfo{.kind = MemoryKind::kHost, .device_id = -1}; + } + + int32_t CurrentDeviceId() const override { + int device_id = -1; + return cudaGetDevice(&device_id) == cudaSuccess ? device_id : -1; + } + + void SetContext(int32_t device_id) const override { + if (device_id >= 0) cudaSetDevice(device_id); + } + + bool Copy(void* dst, const void* src, size_t size, + CopyDirection direction) const override { + cudaMemcpyKind kind = cudaMemcpyDefault; + switch (direction) { + case CopyDirection::kHostToDevice: + kind = cudaMemcpyHostToDevice; + break; + case CopyDirection::kDeviceToHost: + kind = cudaMemcpyDeviceToHost; + break; + case CopyDirection::kDeviceToDevice: + kind = cudaMemcpyDeviceToDevice; + break; + case CopyDirection::kHostToHost: + case CopyDirection::kAuto: + kind = cudaMemcpyDefault; + break; + } + return cudaMemcpy(dst, src, size, kind) == cudaSuccess; + } + + PinnedHostBuffer AllocatePinnedHost(size_t size) const override { + void* addr = nullptr; + if (cudaMallocHost(&addr, size) != cudaSuccess) { + cudaGetLastError(); + return PinnedHostBuffer(); + } + return PinnedHostBuffer(addr, size, FreeCudaLikePinnedHostBuffer); + } + + private: + AcceleratorVendor vendor_; +}; + +#define REGISTER_CUDA_LIKE_ACCELERATOR_DEVICE(name, vendor) \ + const CudaLikeAcceleratorDevice name##_device(vendor); \ + const AcceleratorDeviceRegistrar name##_registrar(name##_device) + +#if defined(USE_CUDA) +REGISTER_CUDA_LIKE_ACCELERATOR_DEVICE(nvidia, AcceleratorVendor::kNvidia); +#endif + +#if defined(USE_MUSA) +REGISTER_CUDA_LIKE_ACCELERATOR_DEVICE(musa, AcceleratorVendor::kMusa); +#endif + +#if defined(USE_MACA) +REGISTER_CUDA_LIKE_ACCELERATOR_DEVICE(maca, AcceleratorVendor::kMaca); +#endif + +#if defined(USE_HYGON) +REGISTER_CUDA_LIKE_ACCELERATOR_DEVICE(hygon, AcceleratorVendor::kHygon); +#endif + +#if defined(USE_COREX) +REGISTER_CUDA_LIKE_ACCELERATOR_DEVICE(corex, AcceleratorVendor::kCorex); +#endif + +#undef REGISTER_CUDA_LIKE_ACCELERATOR_DEVICE + +} // namespace +} // namespace device +} // namespace mooncake + +#endif diff --git a/mooncake-store/src/device/hip_accelerator_device.cpp b/mooncake-store/src/device/hip_accelerator_device.cpp new file mode 100644 index 0000000000..49bbdc313a --- /dev/null +++ b/mooncake-store/src/device/hip_accelerator_device.cpp @@ -0,0 +1,87 @@ +#include "device/accelerator_registry.h" +#include "pinned_host_buffer.h" + +#include "cuda_alike.h" + +#if defined(USE_HIP) + +namespace mooncake { +namespace device { +namespace { + +void FreeHipPinnedHostBuffer(void* addr) { hipHostFree(addr); } + +class HipAcceleratorDevice final : public ProbeCachedAcceleratorDevice { + public: + AcceleratorVendor Vendor() const override { + return AcceleratorVendor::kHip; + } + + bool ProbeAvailable() const override { + int count = 0; + return hipGetDeviceCount(&count) == hipSuccess && count > 0; + } + + PointerInfo QueryPointer(const void* ptr) const override { + hipPointerAttribute_t attr{}; + if (hipPointerGetAttributes(&attr, ptr) == hipSuccess && + attr.type == hipMemoryTypeDevice) { + return PointerInfo{.kind = MemoryKind::kDevice, + .device_id = attr.device}; + } + hipGetLastError(); + return PointerInfo{.kind = MemoryKind::kHost, .device_id = -1}; + } + + int32_t CurrentDeviceId() const override { + int device_id = -1; + return hipGetDevice(&device_id) == hipSuccess ? device_id : -1; + } + + void SetContext(int32_t device_id) const override { + if (device_id >= 0) hipSetDevice(device_id); + } + + bool Copy(void* dst, const void* src, size_t size, + CopyDirection direction) const override { + hipMemcpyKind kind = hipMemcpyDefault; + switch (direction) { + case CopyDirection::kHostToDevice: + kind = hipMemcpyHostToDevice; + break; + case CopyDirection::kDeviceToHost: + kind = hipMemcpyDeviceToHost; + break; + case CopyDirection::kDeviceToDevice: + kind = hipMemcpyDeviceToDevice; + break; + case CopyDirection::kHostToHost: + case CopyDirection::kAuto: + kind = hipMemcpyDefault; + break; + } + return hipMemcpy(dst, src, size, kind) == hipSuccess; + } + + PinnedHostBuffer AllocatePinnedHost(size_t size) const override { + void* addr = nullptr; + if (hipHostMalloc(&addr, size, 0) != hipSuccess) { + hipGetLastError(); + return PinnedHostBuffer(); + } + return PinnedHostBuffer(addr, size, FreeHipPinnedHostBuffer); + } +}; + +const AcceleratorDevice& HipDeviceInstance() { + static HipAcceleratorDevice device; + return device; +} + +const AcceleratorDeviceRegistrar registered_hip_device(HipDeviceInstance()); + +} // namespace +} // namespace device +} // namespace mooncake + +#endif diff --git a/mooncake-store/src/device/runtime_accelerator.cpp b/mooncake-store/src/device/runtime_accelerator.cpp new file mode 100644 index 0000000000..7161fd469d --- /dev/null +++ b/mooncake-store/src/device/runtime_accelerator.cpp @@ -0,0 +1,55 @@ +#include "device/runtime_accelerator.h" + +#include +#include + +namespace mooncake { +namespace device { + +RuntimeAccelerator::RuntimeAccelerator( + std::vector devices) + : devices_(std::move(devices)) {} + +std::span RuntimeAccelerator::Devices() const { + return std::span(devices_.data(), + devices_.size()); +} + +const AcceleratorDevice* RuntimeAccelerator::FindDeviceForPointer( + const void* ptr, PointerInfo* out_info) const { + if (!ptr) return nullptr; + for (auto* accelerator : devices_) { + auto info = accelerator->QueryPointer(ptr); + if (info.kind != MemoryKind::kDevice) continue; + if (out_info) *out_info = info; + return accelerator; + } + return nullptr; +} + +bool RuntimeAccelerator::CopyToHost(void* dst, const void* src, + size_t size) const { + PointerInfo pointer_info; + auto* accelerator = FindDeviceForPointer(src, &pointer_info); + if (!accelerator) { + std::memcpy(dst, src, size); + return true; + } + accelerator->SetContext(pointer_info.device_id); + return accelerator->Copy(dst, src, size, CopyDirection::kDeviceToHost); +} + +bool RuntimeAccelerator::CopyFromHost(void* dst, const void* src, + size_t size) const { + PointerInfo pointer_info; + auto* accelerator = FindDeviceForPointer(dst, &pointer_info); + if (!accelerator) { + std::memcpy(dst, src, size); + return true; + } + accelerator->SetContext(pointer_info.device_id); + return accelerator->Copy(dst, src, size, CopyDirection::kHostToDevice); +} + +} // namespace device +} // namespace mooncake diff --git a/mooncake-store/src/device/sunrise_accelerator_device.cpp b/mooncake-store/src/device/sunrise_accelerator_device.cpp new file mode 100644 index 0000000000..fc5ee76017 --- /dev/null +++ b/mooncake-store/src/device/sunrise_accelerator_device.cpp @@ -0,0 +1,160 @@ +#include "device/accelerator_registry.h" +#include "pinned_host_buffer.h" +#include "sunrise_allocator.h" + +#if defined(USE_SUNRISE) +#include +#include +#include +#include + +namespace mooncake { +namespace device { +namespace { + +struct SavedTangDevice { + int dev{-1}; + SavedTangDevice() { tangGetDevice(&dev); } + ~SavedTangDevice() { + if (dev >= 0) tangSetDevice(dev); + } + SavedTangDevice(const SavedTangDevice&) = delete; + SavedTangDevice& operator=(const SavedTangDevice&) = delete; +}; + +void FreeSunrisePinnedHostBuffer(void* addr) { + if (!addr) return; + { + std::lock_guard guard( + mooncake::sunrise_alloc_detail::tangAllocMutex()); + mooncake::sunrise_alloc_detail::tangHostAllocatedSet().erase(addr); + } + mooncake::sunrise_alloc_detail::removeStoreMemRange(addr); + tangFreeHost(addr); +} + +class SunriseAcceleratorDevice final : public ProbeCachedAcceleratorDevice { + public: + AcceleratorVendor Vendor() const override { + return AcceleratorVendor::kSunrise; + } + + bool ProbeAvailable() const override { + int dev = -1; + return tangGetDevice(&dev) == tangSuccess; + } + + PointerInfo QueryPointer(const void* ptr) const override { + if (sunrise_is_device_memory_range(const_cast(ptr))) { + return PointerInfo{.kind = MemoryKind::kDevice, .device_id = 0}; + } + return PointerInfo{.kind = MemoryKind::kHost, .device_id = -1}; + } + + int32_t CurrentDeviceId() const override { + int dev = -1; + return tangGetDevice(&dev) == tangSuccess ? dev : -1; + } + + void SetContext(int32_t device_id) const override { + if (device_id >= 0) tangSetDevice(device_id); + } + + bool Copy(void* dst, const void* src, size_t size, + CopyDirection direction) const override { + switch (direction) { + case CopyDirection::kHostToHost: + std::memcpy(dst, src, size); + return true; + + case CopyDirection::kDeviceToHost: { + bool dst_host_alloc = sunrise_is_host_allocated(dst); + bool dst_dev = sunrise_is_device_memory_range(dst); + + if (dst_host_alloc && !dst_dev) { + SavedTangDevice saved; + tangSetDevice(0); + std::vector staging(size); + if (tangMemcpy(staging.data(), src, size, + tangMemcpyDeviceToHost) != tangSuccess) + return false; + tangDeviceSynchronize(); + std::memcpy(dst, staging.data(), size); + return true; + } + if (tangMemcpy(dst, src, size, tangMemcpyDeviceToHost) != + tangSuccess) + return false; + tangDeviceSynchronize(); + return true; + } + + case CopyDirection::kHostToDevice: { + bool src_host_alloc = + sunrise_is_host_allocated(const_cast(src)); + bool src_dev = + sunrise_is_device_memory_range(const_cast(src)); + + if (src_host_alloc && !src_dev) { + SavedTangDevice saved; + tangSetDevice(0); + std::vector staging(size); + std::memcpy(staging.data(), src, size); + return tangMemcpy(dst, staging.data(), size, + tangMemcpyHostToDevice) == tangSuccess; + } + return tangMemcpy(dst, src, size, tangMemcpyHostToDevice) == + tangSuccess; + } + + case CopyDirection::kDeviceToDevice: + return tangMemcpy(dst, src, size, tangMemcpyDeviceToDevice) == + tangSuccess; + + case CopyDirection::kAuto: { + bool src_dev = + sunrise_is_device_memory_range(const_cast(src)); + bool dst_dev = sunrise_is_device_memory_range(dst); + if (!src_dev && !dst_dev) { + std::memcpy(dst, src, size); + return true; + } + if (src_dev && dst_dev) + return Copy(dst, src, size, CopyDirection::kDeviceToDevice); + if (src_dev) + return Copy(dst, src, size, CopyDirection::kDeviceToHost); + return Copy(dst, src, size, CopyDirection::kHostToDevice); + } + } + return false; + } + + PinnedHostBuffer AllocatePinnedHost(size_t size) const override { + void* addr = nullptr; + if (tangHostAlloc(&addr, size, tangHostAllocDefault) == tangSuccess) { + { + std::lock_guard guard( + mooncake::sunrise_alloc_detail::tangAllocMutex()); + mooncake::sunrise_alloc_detail::tangHostAllocatedSet().insert( + addr); + } + mooncake::sunrise_alloc_detail::addStoreMemRange(addr, size); + return PinnedHostBuffer(addr, size, FreeSunrisePinnedHostBuffer); + } + return PinnedHostBuffer(); + } +}; + +const AcceleratorDevice& SunriseDeviceInstance() { + static SunriseAcceleratorDevice device; + return device; +} + +const AcceleratorDeviceRegistrar registered_sunrise_device( + SunriseDeviceInstance()); + +} // namespace +} // namespace device +} // namespace mooncake + +#endif diff --git a/mooncake-store/src/dummy_client.cpp b/mooncake-store/src/dummy_client.cpp index c1a84a5c2a..bc726515d2 100644 --- a/mooncake-store/src/dummy_client.cpp +++ b/mooncake-store/src/dummy_client.cpp @@ -12,12 +12,14 @@ #include "real_client.h" #include "dummy_client.h" +#include "uds_transport.h" #include "utils.h" #include "utils/scoped_vlog_timer.h" #include "rpc_types.h" #include "types.h" #include "default_config.h" #include "config.h" +#include "device/cuda_ipc_buffer.h" #ifdef USE_ASCEND_DIRECT #include "acl/acl_rt.h" #include "ascend_allocator.h" @@ -63,6 +65,34 @@ size_t sum_successful_nested_sizes( return total; } +size_t sum_successful_cuda_ipc_sizes( + const std::vector& results, + const std::vector& requests) { + size_t total = 0; + for (size_t i = 0; i < results.size() && i < requests.size(); ++i) { + if (results[i] == 0) { + total += requests[i].metadata.size + requests[i].payload.size; + } + } + return total; +} + +std::optional> +try_export_cuda_ipc_buffers(const std::vector& buffers, + const std::vector& sizes) { + if (buffers.size() != sizes.size() || buffers.empty()) return std::nullopt; + + std::vector payloads; + payloads.reserve(buffers.size()); + for (size_t i = 0; i < buffers.size(); ++i) { + auto payload = + mooncake::device::ExportCudaIpcBuffer(buffers[i], sizes[i]); + if (!payload) return std::nullopt; + payloads.push_back(*payload); + } + return payloads; +} + size_t sum_positive_results(const std::vector& results) { size_t total = 0; for (int64_t result : results) { @@ -94,6 +124,17 @@ size_t sum_positive_ranges( return total; } +template +std::vector expected_results_to_py( + const std::vector>& internal_results) { + std::vector results; + results.reserve(internal_results.size()); + for (const auto& result : internal_results) { + results.push_back(static_cast(mooncake::to_py_ret(result))); + } + return results; +} + std::vector void_ptrs_to_u64(const std::vector& ptrs) { std::vector out; out.reserve(ptrs.size()); @@ -213,16 +254,12 @@ std::vector> DummyClient::invoke_batch_rpc( } DummyClient::DummyClient() - : client_id_(generate_uuid()), + : client_accessor_(GetStoreRpcClientIoContextPool()), + client_id_(generate_uuid()), metrics_(ClientMetric::Create(merge_labels({{"client_mode", "dummy"}}), false)) { // Initialize logging severity (leave as before) mooncake::init_ylt_log_level(); - // Initialize client pools - coro_io::client_pool::pool_config pool_conf{}; - client_pools_ = - std::make_shared>( - pool_conf); } DummyClient::~DummyClient() { tearDownAll(); } @@ -257,13 +294,9 @@ ErrorCode DummyClient::connect(const std::string& server_address) { MutexLocker lock(&connect_mutex_); if (client_addr_param_ != server_address) { - // WARNING: The existing client pool cannot be erased. So if there are a - // lot of different addresses, there will be resource leak problems. - auto client_pool = client_pools_->at(server_address); - client_accessor_.SetClientPool(client_pool); + client_accessor_.GetOrCreateClientPool(server_address); client_addr_param_ = server_address; } - auto pool = client_accessor_.GetClientPool(); // The client pool does not have native connection check method, so we need // to use custom ServiceReady API. auto result = invoke_rpc<&RealClient::service_ready_internal, void>(); @@ -377,40 +410,21 @@ int DummyClient::register_shm_via_ipc(const ShmHelper::ShmSegment* shm, return -1; } - int sock_fd = socket(AF_UNIX, SOCK_STREAM, 0); - if (sock_fd < 0) { - LOG(ERROR) << "Failed to create IPC socket: " << strerror(errno); - return -1; - } - - struct sockaddr_un addr; - memset(&addr, 0, sizeof(addr)); - addr.sun_family = AF_UNIX; - - // Use abstract namespace - std::string abstract_name = ipc_socket_path_; - if (abstract_name.size() > sizeof(addr.sun_path) - 2) { - LOG(ERROR) << "IPC socket path too long"; - close(sock_fd); - return -1; - } - addr.sun_path[0] = '\0'; - strncpy(addr.sun_path + 1, abstract_name.c_str(), - sizeof(addr.sun_path) - 2); - socklen_t addr_len = sizeof(sa_family_t) + 1 + abstract_name.length(); - LOG(INFO) << "Connecting to IPC socket: " << abstract_name; - - if (::connect(sock_fd, (struct sockaddr*)&addr, addr_len) < 0) { + UdsConnector connector(ipc_socket_path_); + LOG(INFO) << "Connecting to IPC socket: " << ipc_socket_path_; + auto connection_result = connector.connect(); + if (!connection_result) { + LOG(ERROR) << "Failed to connect IPC socket '" << ipc_socket_path_ + << "': " << connection_result.error(); // This is expected if RealClient is down - close(sock_fd); return -1; } + auto connection = std::move(connection_result.value()); // Send request type first IpcRequestType type = IPC_SHM_REGISTER; - if (::send(sock_fd, &type, sizeof(type), 0) < 0) { + if (connection->sendRaw(&type, sizeof(type)) < 0) { LOG(ERROR) << "Failed to send IPC request type: " << strerror(errno); - close(sock_fd); return -1; } @@ -423,21 +437,17 @@ int DummyClient::register_shm_via_ipc(const ShmHelper::ShmSegment* shm, : kInvalidPhysicalDeviceId; req.is_local_buffer = is_local; - if (ipc_send_fd(sock_fd, shm->fd, &req, sizeof(req)) < 0) { + if (connection->sendFd(shm->fd, &req, sizeof(req)) < 0) { LOG(ERROR) << "Failed to send FD to RealClient: " << strerror(errno); - close(sock_fd); return -1; } int status = -1; - if (recv(sock_fd, &status, sizeof(status), 0) < 0) { + if (connection->recvRaw(&status, sizeof(status)) < 0) { LOG(ERROR) << "Failed to receive response from RealClient"; - close(sock_fd); return -1; } - close(sock_fd); - if (status != 0) { LOG(ERROR) << "RealClient failed to map shared memory, error code: " << status; @@ -489,6 +499,7 @@ int DummyClient::setup_dummy(size_t mem_pool_size, size_t local_buffer_size, if (local_buffer_size > 0) { try { base_addr = shm_helper_->allocate(local_buffer_size); + local_buffer_base_ = base_addr; } catch (const std::exception& e) { LOG(ERROR) << "Failed to allocate shared memory: " << e.what(); return -1; @@ -498,6 +509,7 @@ int DummyClient::setup_dummy(size_t mem_pool_size, size_t local_buffer_size, if (!local_buffer_shm) { LOG(ERROR) << "Failed to get shm segment for base address"; shm_helper_->free(base_addr); + local_buffer_base_ = nullptr; return -1; } @@ -506,6 +518,7 @@ int DummyClient::setup_dummy(size_t mem_pool_size, size_t local_buffer_size, LOG(ERROR) << "Failed to register SHM via IPC"; // Register failed, cleanup shm_helper_->free(local_buffer_shm->base_addr); + local_buffer_base_ = nullptr; return -1; } } else { @@ -513,6 +526,7 @@ int DummyClient::setup_dummy(size_t mem_pool_size, size_t local_buffer_size, LOG(ERROR) << "Failed to register SHM via IPC"; // Register failed, cleanup shm_helper_->free(local_buffer_shm->base_addr); + local_buffer_base_ = nullptr; return -1; } } @@ -532,7 +546,14 @@ int DummyClient::setup_dummy(size_t mem_pool_size, size_t local_buffer_size, } int DummyClient::tearDownAll() { + void* local_buffer_base = local_buffer_base_; unregister_shm(); + if (local_buffer_base && shm_helper_ && + shm_helper_->get_shm(local_buffer_base) && + shm_helper_->free(local_buffer_base) != 0) { + LOG(ERROR) << "Failed to free dummy local shared memory"; + } + local_buffer_base_ = nullptr; // Cleanup hot cache shm mapping if (hot_cache_base_) { @@ -562,6 +583,23 @@ int DummyClient::tearDownAll() { return 0; } +std::optional DummyClient::allocate_client_buffer(size_t size) { + auto result = invoke_rpc<&RealClient::allocate_buffer_dummy, + std::tuple>(size, client_id_); + if (!result.has_value()) { + return std::nullopt; + } + + auto [dummy_addr, allocated_size] = result.value(); + void* local_ptr = reinterpret_cast(dummy_addr); + auto release = [this, dummy_addr]() { + (void)invoke_rpc<&RealClient::release_buffer_dummy, void>(dummy_addr, + client_id_); + }; + return std::make_optional(local_ptr, allocated_size, + std::move(release)); +} + int64_t DummyClient::unregister_shm() { LOG(INFO) << "[unregister_shm] client_id=" << client_id_; #if defined(USE_ASCEND_DIRECT) @@ -1018,11 +1056,24 @@ std::vector> DummyClient::batch_get_buffer( int64_t DummyClient::get_into(const std::string& key, void* buffer, size_t size) { + if (auto dst_buffer = try_export_cuda_ipc_buffers({buffer}, {size})) { + std::vector requests{ + CudaIpcReadRequest{ + .key = key, + .destination = (*dst_buffer)[0], + .source_offset = 0, + .size = static_cast(size), + }, + }; + auto results = batch_get_into_cuda_ipc(requests); + return results.empty() ? toInt(ErrorCode::INVALID_PARAMS) : results[0]; + } + uint64_t buf_addr = reinterpret_cast(buffer); const auto start_time = std::chrono::steady_clock::now(); auto result = invoke_rpc<&RealClient::get_into_range_shm_helper, tl::expected>( - key, buf_addr, 0, 0, size, client_id_); + key, buf_addr, 0, 0, size, true, true, client_id_); if (!result) { return static_cast(toInt(result.error())); } @@ -1088,15 +1139,8 @@ std::vector> DummyClient::batch_query( results.reserve(keys.size()); const auto now = std::chrono::steady_clock::now(); for (const auto& cached_result : *cached_results) { - if (!cached_result.success) { - results.emplace_back(tl::unexpected(cached_result.error)); - continue; - } - results.emplace_back(QueryResult( - std::vector( - cached_result.value.replicas.begin(), - cached_result.value.replicas.end()), - now + std::chrono::milliseconds(cached_result.value.lease_ttl_ms))); + results.emplace_back( + from_cached_query_result_response(cached_result, now)); } return results; } @@ -1109,6 +1153,20 @@ std::string DummyClient::get_hostname() const { std::vector DummyClient::batch_put_from( const std::vector& keys, const std::vector& buffer_ptrs, const std::vector& sizes, const ReplicateConfig& config) { + if (auto payloads = try_export_cuda_ipc_buffers(buffer_ptrs, sizes); + payloads && keys.size() == payloads->size()) { + std::vector requests; + requests.reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + requests.push_back(CudaIpcWriteRequest{ + .key = keys[i], + .metadata = CudaIpcShmBufferRef{}, + .payload = (*payloads)[i], + }); + } + return batch_put_from_cuda_ipc(requests, config); + } + std::vector buffers = void_ptrs_to_u64(buffer_ptrs); const auto start_time = std::chrono::steady_clock::now(); auto internal_results = @@ -1133,24 +1191,34 @@ std::vector DummyClient::batch_put_from( int DummyClient::put_from(const std::string& key, void* buffer, size_t size, const ReplicateConfig& config) { - // TODO: implement this function - return -1; + auto results = batch_put_from({key}, {buffer}, {size}, config); + return results.empty() ? -1 : results[0]; } std::vector DummyClient::batch_get_into( const std::vector& keys, const std::vector& buffer_ptrs, const std::vector& sizes) { + if (auto dst_buffers = try_export_cuda_ipc_buffers(buffer_ptrs, sizes); + dst_buffers && keys.size() == dst_buffers->size()) { + std::vector requests; + requests.reserve(keys.size()); + for (size_t i = 0; i < keys.size(); ++i) { + requests.push_back(CudaIpcReadRequest{ + .key = keys[i], + .destination = (*dst_buffers)[i], + .source_offset = 0, + .size = static_cast(sizes[i]), + }); + } + return batch_get_into_cuda_ipc(requests); + } + std::vector buffers = void_ptrs_to_u64(buffer_ptrs); const auto start_time = std::chrono::steady_clock::now(); auto internal_results = invoke_batch_rpc<&RealClient::batch_get_into_dummy_helper, int64_t>( keys.size(), keys, buffers, sizes, device_id_, client_id_); - std::vector results; - results.reserve(internal_results.size()); - - for (const auto& result : internal_results) { - results.push_back(to_py_ret(result)); - } + auto results = expected_results_to_py(internal_results); const size_t total_bytes = sum_positive_results(results); if (total_bytes > 0) { @@ -1161,12 +1229,31 @@ std::vector DummyClient::batch_get_into( return results; } +std::vector DummyClient::batch_get_into_cuda_ipc( + const std::vector& requests) { + const auto start_time = std::chrono::steady_clock::now(); + auto internal_results = + invoke_batch_rpc<&RealClient::batch_get_into_cuda_ipc_dummy_helper, + int64_t>(requests.size(), requests, client_id_); + auto results = expected_results_to_py(internal_results); + + const size_t total_bytes = sum_positive_results(results); + if (total_bytes > 0) { + ObserveTransferMetric(TransferOperationKind::kRead, + "batch_get_into_cuda_ipc", total_bytes, + elapsed_us_since(start_time), true); + } + + return results; +} + int DummyClient::put_from_with_metadata(const std::string& key, void* buffer, void* metadata_buffer, size_t size, size_t metadata_size, const ReplicateConfig& config) { - // TODO: implement this function - return -1; + auto results = batch_put_from_multi_buffers( + {key}, {{metadata_buffer, buffer}}, {{metadata_size, size}}, config); + return results.empty() ? -1 : results[0]; } std::vector DummyClient::batch_put_from_multi_buffers( @@ -1181,11 +1268,7 @@ std::vector DummyClient::batch_put_from_multi_buffers( invoke_batch_rpc<&RealClient::batch_put_from_multi_buffers_dummy_helper, void>(keys.size(), keys, dummy_nested, all_sizes, config, device_id_, client_id_); - std::vector results; - results.reserve(internal_results.size()); - for (const auto& result : internal_results) { - results.push_back(to_py_ret(result)); - } + auto results = expected_results_to_py(internal_results); const size_t successful_bytes = sum_successful_nested_sizes(results, all_sizes); if (successful_bytes > 0) { @@ -1196,6 +1279,24 @@ std::vector DummyClient::batch_put_from_multi_buffers( return results; } +std::vector DummyClient::batch_put_from_cuda_ipc( + const std::vector& requests, + const ReplicateConfig& config) { + const auto start_time = std::chrono::steady_clock::now(); + auto internal_results = + invoke_batch_rpc<&RealClient::batch_put_from_cuda_ipc_dummy_helper, + void>(requests.size(), requests, config, client_id_); + auto results = expected_results_to_py(internal_results); + const size_t successful_bytes = + sum_successful_cuda_ipc_sizes(results, requests); + if (successful_bytes > 0) { + ObserveTransferMetric(TransferOperationKind::kWrite, + "batch_put_from_cuda_ipc", successful_bytes, + elapsed_us_since(start_time), true); + } + return results; +} + std::vector DummyClient::batch_get_into_multi_buffers( const std::vector& keys, const std::vector>& all_buffer_ptrs, @@ -1383,30 +1484,19 @@ int DummyClient::health_check() { } int DummyClient::request_hot_cache_fd() { - int sock_fd = socket(AF_UNIX, SOCK_STREAM, 0); - if (sock_fd < 0) { - LOG(ERROR) << "Failed to create IPC socket: " << strerror(errno); - return -1; - } - - struct sockaddr_un addr; - memset(&addr, 0, sizeof(addr)); - addr.sun_family = AF_UNIX; - addr.sun_path[0] = '\0'; - strncpy(addr.sun_path + 1, ipc_socket_path_.c_str(), - sizeof(addr.sun_path) - 2); - socklen_t addr_len = sizeof(sa_family_t) + 1 + ipc_socket_path_.length(); - - if (::connect(sock_fd, (struct sockaddr*)&addr, addr_len) < 0) { - close(sock_fd); + UdsConnector connector(ipc_socket_path_); + auto connection_result = connector.connect(); + if (!connection_result) { + LOG(ERROR) << "Failed to connect IPC socket '" << ipc_socket_path_ + << "': " << connection_result.error(); return -1; } + auto connection = std::move(connection_result.value()); // Send request type IpcRequestType type = IPC_SHM_FD_REQUEST; - if (::send(sock_fd, &type, sizeof(type), 0) < 0) { + if (connection->sendRaw(&type, sizeof(type)) < 0) { LOG(ERROR) << "Failed to send IPC request type"; - close(sock_fd); return -1; } @@ -1415,16 +1505,14 @@ int DummyClient::request_hot_cache_fd() { req.client_id_first = client_id_.first; req.client_id_second = client_id_.second; req.segment_type = SHM_SEG_HOT_CACHE; - if (::send(sock_fd, &req, sizeof(req), 0) < 0) { + if (connection->sendRaw(&req, sizeof(req)) < 0) { LOG(ERROR) << "Failed to send ShmFdRequest"; - close(sock_fd); return -1; } // Receive fd + response ShmFdResponse resp; - int fd = ipc_recv_fd(sock_fd, &resp, sizeof(resp)); - close(sock_fd); + int fd = connection->recvFd(&resp, sizeof(resp)); if (fd < 0 || resp.status != 0) { LOG(ERROR) << "Failed to receive hot cache fd, status=" << resp.status; diff --git a/mooncake-store/src/etcd_helper.cpp b/mooncake-store/src/etcd_helper.cpp index ba096c720e..69cd19ad7a 100644 --- a/mooncake-store/src/etcd_helper.cpp +++ b/mooncake-store/src/etcd_helper.cpp @@ -151,6 +151,67 @@ ErrorCode EtcdHelper::BatchCreate(const std::vector& keys, return ErrorCode::OK; } +ErrorCode EtcdHelper::TxnCompareAndPut(const std::vector& compares, + const std::vector& puts) { + std::vector compare_keys; + std::vector compare_key_sizes; + std::vector compare_kinds; + std::vector compare_values; + std::vector compare_value_sizes; + compare_keys.reserve(compares.size()); + compare_key_sizes.reserve(compares.size()); + compare_kinds.reserve(compares.size()); + compare_values.reserve(compares.size()); + compare_value_sizes.reserve(compares.size()); + for (const auto& compare : compares) { + compare_keys.push_back(const_cast(compare.key.data())); + compare_key_sizes.push_back(static_cast(compare.key.size())); + compare_kinds.push_back(static_cast(compare.kind)); + compare_values.push_back( + const_cast(compare.expected_value.data())); + compare_value_sizes.push_back( + static_cast(compare.expected_value.size())); + } + + std::vector put_keys; + std::vector put_key_sizes; + std::vector put_values; + std::vector put_value_sizes; + put_keys.reserve(puts.size()); + put_key_sizes.reserve(puts.size()); + put_values.reserve(puts.size()); + put_value_sizes.reserve(puts.size()); + for (const auto& put : puts) { + put_keys.push_back(const_cast(put.key.data())); + put_key_sizes.push_back(static_cast(put.key.size())); + put_values.push_back(const_cast(put.value.data())); + put_value_sizes.push_back(static_cast(put.value.size())); + } + + char* err_msg = nullptr; + int ret = EtcdStoreTxnCompareAndPutWrapper( + compare_keys.data(), compare_key_sizes.data(), compare_kinds.data(), + compare_values.data(), compare_value_sizes.data(), + static_cast(compares.size()), put_keys.data(), + put_key_sizes.data(), put_values.data(), put_value_sizes.data(), + static_cast(puts.size()), &err_msg); + if (ret == -2) { + if (err_msg != nullptr) { + free(err_msg); + } + return ErrorCode::ETCD_TRANSACTION_FAIL; + } + if (ret != 0) { + LOG(ERROR) << "TxnCompareAndPut failed: " + << (err_msg == nullptr ? "" : err_msg); + if (err_msg != nullptr) { + free(err_msg); + } + return ErrorCode::ETCD_OPERATION_ERROR; + } + return ErrorCode::OK; +} + ErrorCode EtcdHelper::GrantLease(int64_t lease_ttl, EtcdLeaseId& lease_id) { char* err_msg = nullptr; if (0 != EtcdStoreGrantLeaseWrapper(lease_ttl, &lease_id, &err_msg)) { @@ -527,6 +588,14 @@ ErrorCode EtcdHelper::Create(const char* key, const size_t key_size, return ErrorCode::ETCD_OPERATION_ERROR; } +ErrorCode EtcdHelper::TxnCompareAndPut(const std::vector& compares, + const std::vector& puts) { + (void)compares; + (void)puts; + LOG(FATAL) << "Etcd is not enabled in compilation"; + return ErrorCode::ETCD_OPERATION_ERROR; +} + ErrorCode EtcdHelper::GetRangeAsJson(const char* start_key, const size_t start_key_size, const char* end_key, diff --git a/mooncake-store/src/file_storage.cpp b/mooncake-store/src/file_storage.cpp index 4830c13eb3..f4ccfbc0a0 100644 --- a/mooncake-store/src/file_storage.cpp +++ b/mooncake-store/src/file_storage.cpp @@ -1,35 +1,79 @@ #include "file_storage.h" +#include +#include #include +#include +#include +#include #include -#include "aligned_client_buffer.hpp" +#include "aligned_client_buffer.h" +#include "bool_parser.h" +#include "environ.h" #include "storage_backend.h" #include "client_metric.h" #include "utils.h" -#include "gpu_staging_utils.h" +#include "device/accelerator_registry.h" #ifdef USE_URING #include "file_interface.h" #endif namespace mooncake { -using gpu_staging::CopyDeviceToHost; -using gpu_staging::IsDevicePointer; -using gpu_staging::SetDevice; - namespace { +double ParseEnvRatioOr(const std::string& raw_value, double default_value) { + if (raw_value.empty()) { + return default_value; + } + + std::istringstream stream(raw_value); + stream.imbue(std::locale::classic()); + + double value = 0.0; + stream >> value; + if (stream.fail()) { + return default_value; + } + if (!stream.eof() || !std::isfinite(value) || value <= 0.0 || value > 1.0) { + return default_value; + } + return value; +} + +double GetEnvRatioOr(const char* name, double default_value) { + const auto raw_value = Environ::GetString(name, ""); + return ParseEnvRatioOr(raw_value, default_value); +} + +double GetEnvRatioOr(const char* preferred_name, const char* fallback_name, + double default_value) { + const auto preferred_value = Environ::GetString(preferred_name, ""); + if (!preferred_value.empty()) { + return ParseEnvRatioOr(preferred_value, default_value); + } + return GetEnvRatioOr(fallback_name, default_value); +} + +bool GetEnvBoolStringOr(const char* name, bool default_value) { + const auto raw_value = + Environ::GetString(name, default_value ? "true" : "false"); + return TryParseBool(raw_value, {.token_set = BoolTokenSet::kTrueFalse, + .trim_ascii_whitespace = false}) + .value_or(default_value); +} + std::vector BuildOffloadTasksFromStorageKeys( const std::vector& storage_keys, const std::vector& metadatas) { std::vector tasks; tasks.reserve(storage_keys.size()); for (size_t i = 0; i < storage_keys.size(); ++i) { - auto [tenant_id, key] = ParseTenantScopedStorageKey(storage_keys[i]); + auto [tenant_id, key] = TenantId::ParseScopedKey(storage_keys[i]); const int64_t size = i < metadatas.size() ? metadatas[i].data_size : int64_t{0}; - tasks.push_back(OffloadTaskItem{.tenant_id = std::move(tenant_id), + tasks.push_back(OffloadTaskItem{.tenant_id = tenant_id.value(), .key = std::move(key), .size = size}); } @@ -42,8 +86,8 @@ FileStorageConfig FileStorageConfig::FromEnvironment() { FileStorageConfig config; auto storage_backend_descriptor = - GetEnvStringOr("MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR", - "bucket_storage_backend"); + Environ::GetString("MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR", + "bucket_storage_backend"); if (storage_backend_descriptor == "bucket_storage_backend") { config.storage_backend_type = StorageBackendType::kBucket; @@ -58,38 +102,53 @@ FileStorageConfig FileStorageConfig::FromEnvironment() { LOG(ERROR) << "Unknown storage backend."; } - config.storage_filepath = GetEnvStringOr( + config.storage_filepath = Environ::GetString( "MOONCAKE_OFFLOAD_FILE_STORAGE_PATH", config.storage_filepath); - config.local_buffer_size = GetEnvOr( + config.local_buffer_size = Environ::GetInt64( "MOONCAKE_OFFLOAD_LOCAL_BUFFER_SIZE_BYTES", config.local_buffer_size); - config.scanmeta_iterator_keys_limit = GetEnvOr( + config.scanmeta_iterator_keys_limit = Environ::GetInt64( "MOONCAKE_OFFLOAD_SCANMETA_ITERATOR_KEYS_LIMIT", - GetEnvOr("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT", + Environ::GetInt64("MOONCAKE_SCANMETA_ITERATOR_KEYS_LIMIT", config.scanmeta_iterator_keys_limit)); - config.total_keys_limit = GetEnvOr( + config.total_keys_limit = Environ::GetInt64( "MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT", config.total_keys_limit); - config.total_size_limit = GetEnvOr( + config.total_size_limit = Environ::GetInt64( "MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES", config.total_size_limit); config.heartbeat_interval_seconds = - GetEnvOr("MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS", + Environ::GetUInt32("MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS", config.heartbeat_interval_seconds); config.client_buffer_gc_interval_seconds = - GetEnvOr("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_INTERVAL_SECONDS", + Environ::GetUInt32("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_INTERVAL_SECONDS", config.client_buffer_gc_interval_seconds); config.client_buffer_gc_ttl_ms = - GetEnvOr("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS", + Environ::GetUInt64("MOONCAKE_OFFLOAD_CLIENT_BUFFER_GC_TTL_MS", config.client_buffer_gc_ttl_ms); - auto use_uring_str = - GetEnvStringOr("MOONCAKE_OFFLOAD_USE_URING", - GetEnvStringOr("MOONCAKE_USE_URING", "false")); - config.use_uring = (use_uring_str == "true" || use_uring_str == "1"); + config.enable_disk_watermark_eviction = + GetEnvBoolStringOr("MOONCAKE_OFFLOAD_ENABLE_DISK_WATERMARK_EVICTION", + config.enable_disk_watermark_eviction); + config.disk_eviction_high_watermark_ratio = + GetEnvRatioOr("MOONCAKE_OFFLOAD_DISK_EVICTION_HIGH_WATERMARK_RATIO", + "MOONCAKE_DISK_EVICTION_HIGH_WATERMARK_RATIO", + config.disk_eviction_high_watermark_ratio); + config.disk_eviction_low_watermark_ratio = + GetEnvRatioOr("MOONCAKE_OFFLOAD_DISK_EVICTION_LOW_WATERMARK_RATIO", + "MOONCAKE_DISK_EVICTION_LOW_WATERMARK_RATIO", + config.disk_eviction_low_watermark_ratio); + + const auto use_uring_str = + Environ::GetString("MOONCAKE_OFFLOAD_USE_URING", + Environ::GetString("MOONCAKE_USE_URING", "false")); + config.use_uring = + TryParseBool(use_uring_str, {.token_set = BoolTokenSet::kTrueFalse, + .trim_ascii_whitespace = false}) + .value_or(false); return config; } @@ -171,6 +230,25 @@ bool FileStorageConfig::Validate() const { LOG(ERROR) << "FileStorageConfig: heartbeat_interval_seconds must > 0"; return false; } + if (disk_eviction_low_watermark_ratio <= 0.0 || + disk_eviction_low_watermark_ratio > 1.0) { + LOG(ERROR) << "FileStorageConfig: " + << "disk_eviction_low_watermark_ratio must be in (0, 1]"; + return false; + } + if (disk_eviction_high_watermark_ratio <= 0.0 || + disk_eviction_high_watermark_ratio > 1.0) { + LOG(ERROR) << "FileStorageConfig: " + << "disk_eviction_high_watermark_ratio must be in (0, 1]"; + return false; + } + if (disk_eviction_low_watermark_ratio >= + disk_eviction_high_watermark_ratio) { + LOG(ERROR) << "FileStorageConfig: " + << "disk_eviction_low_watermark_ratio must be lower than " + << "disk_eviction_high_watermark_ratio"; + return false; + } return true; } @@ -361,6 +439,11 @@ tl::expected FileStorage::BatchGet( return batch_result; } +bool FileStorage::IsPerBucketSoftOffloadError(ErrorCode error) { + return error == ErrorCode::INVALID_READ || + error == ErrorCode::OBJECT_ALREADY_EXISTS; +} + tl::expected FileStorage::OffloadObjects( const std::vector& offloading_objects) { if (offloading_objects.empty()) { @@ -372,7 +455,7 @@ tl::expected FileStorage::OffloadObjects( task_by_storage_key.reserve(offloading_objects.size()); for (const auto& task : offloading_objects) { const auto storage_key = - MakeTenantScopedStorageKey(task.tenant_id, task.key); + TenantId(task.tenant_id).MakeScopedKey(task.key); storage_object_sizes.emplace(storage_key, task.size); task_by_storage_key.emplace(storage_key, task); } @@ -416,14 +499,25 @@ tl::expected FileStorage::OffloadObjects( } auto result = client_->NotifyOffloadSuccess(tasks, metadatas); if (!result) { - LOG(ERROR) << "NotifyOffloadSuccess failed with error: " - << result.error(); + LOG(ERROR) << "[OFFLOAD] NotifyOffloadSuccess failed with error: " + << result.error() << " keys count: " << keys.size(); return result.error(); } return ErrorCode::OK; }; + // Collect keys drained from master queue but not actually offloaded. + // Report them back with data_size=-1 sentinel so the master can clean up + // orphaned offloading_tasks and release source replica refcounts. + std::vector failed_tasks; + std::unordered_set all_bucket_keys; + // Set when a whole-cycle error aborts the bucket loop early. We still fall + // through to the NACK flush below before returning it, so no drained key is + // left waiting on the TTL reaper. + std::optional abort_error; + for (const auto& keys : buckets_keys) { + for (const auto& k : keys) all_bucket_keys.insert(k); std::unordered_map> batch_object; std::unordered_map> storage_keys_by_tenant; @@ -438,22 +532,22 @@ tl::expected FileStorage::OffloadObjects( std::vector user_keys; user_keys.reserve(storage_keys.size()); for (const auto& storage_key : storage_keys) { - user_keys.push_back(task_by_storage_key[storage_key].key); + user_keys.push_back(task_by_storage_key.at(storage_key).key); } std::unordered_map> user_batch_object; - auto query_result = BatchQuerySegmentSlices(user_keys, tenant_id, - user_batch_object); - if (!query_result) { - LOG(ERROR) << "BatchQuerySlices failed with error: " - << query_result.error(); - continue; - } - for (size_t i = 0; i < storage_keys.size(); ++i) { - auto it = user_batch_object.find(user_keys[i]); + [[maybe_unused]] auto query_result = BatchQuerySegmentSlices( + user_keys, tenant_id, user_batch_object); + // BatchQuerySegmentSlices is now best-effort: it always returns + // OK. Keys present in user_batch_object go to batch_object; the + // rest are reported as failed. + for (const auto& storage_key : storage_keys) { + const auto& task = task_by_storage_key.at(storage_key); + auto it = user_batch_object.find(task.key); if (it != user_batch_object.end()) { - batch_object.emplace(storage_keys[i], - std::move(it->second)); + batch_object.emplace(storage_key, std::move(it->second)); + } else { + failed_tasks.push_back(task); } } } @@ -461,53 +555,34 @@ tl::expected FileStorage::OffloadObjects( continue; } - auto eviction_handler = [this](const std::vector& - evicted_keys) { - if (evicted_keys.empty()) return; - std::unordered_map> - keys_by_tenant; - for (const auto& storage_key : evicted_keys) { - auto [tenant_id, key] = - ParseTenantScopedStorageKey(storage_key); - keys_by_tenant[tenant_id].push_back(key); - } - for (const auto& [tenant_id, keys] : keys_by_tenant) { - auto results = client_->BatchEvictDiskReplica( - keys, tenant_id, ReplicaType::LOCAL_DISK); - for (size_t i = 0; i < results.size(); ++i) { - if (!results[i]) { - LOG(WARNING) - << "Failed to notify master about evicted local " - "disk key: " - << keys[i] << ", tenant_id=" << tenant_id - << ", error: " << results[i].error(); - } - } - } - }; - // D2H staging: replace device slices with host memory slices // so that storage_backend (ConcatSlicesToString / BuildBucket / // WriteBucket) always receives host pointers. std::unordered_map> host_batch_object; std::vector staging_bufs; + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); for (auto& [obj_key, slices] : batch_object) { std::vector host_slices; bool obj_success = true; for (const auto& slice : slices) { - int device_id = -1; - if (IsDevicePointer(slice.ptr, &device_id)) { - SetDevice(device_id); + device::PointerInfo info{}; + auto* device = + runtime_accelerator.FindDeviceForPointer(slice.ptr, &info); + if (device) { + device->SetContext(info.device_id); auto buf = pinned_buffer_pool_->Acquire(slice.size); - if (!CopyDeviceToHost(buf.data, slice.ptr, slice.size)) { + if (!device->Copy(buf.data, slice.ptr, slice.size, + device::CopyDirection::kDeviceToHost)) { LOG(ERROR) << "D2H staging failed for key: " << obj_key; - pinned_buffer_pool_->Release(buf); + pinned_buffer_pool_->Release(std::move(buf)); obj_success = false; + failed_tasks.push_back(task_by_storage_key.at(obj_key)); break; } host_slices.emplace_back(Slice{buf.data, slice.size}); - staging_bufs.push_back(buf); + staging_bufs.push_back(std::move(buf)); } else { host_slices.push_back(slice); } @@ -517,6 +592,21 @@ tl::expected FileStorage::OffloadObjects( } } + // If every object in this bucket failed D2H staging, host_batch_object + // is empty (those keys are already in failed_tasks). Skip BatchOffload, + // which rejects an empty map as INVALID_KEY and would otherwise trip + // the whole-cycle abort below for a bucket that has nothing left to + // persist. staging_bufs can still be non-empty here (an object whose + // first slices copied fine but a later one failed), so hand those + // buffers back before continuing: the release loop after BatchOffload + // is unreachable on this path. + if (host_batch_object.empty()) { + for (auto& buf : staging_bufs) { + pinned_buffer_pool_->Release(std::move(buf)); + } + continue; + } + auto offload_start = std::chrono::steady_clock::now(); auto bucket_complete_handler = [this, offload_start, complete_handler]( @@ -544,25 +634,140 @@ tl::expected FileStorage::OffloadObjects( return res; }; auto offload_res = storage_backend_->BatchOffload( - host_batch_object, bucket_complete_handler, eviction_handler); + host_batch_object, bucket_complete_handler, + [this](const std::vector& evicted_keys) { + return NotifyEvictedDiskReplicas(evicted_keys); + }); - // Release staging buffers back to pool (Buffer is POD, no destructor) + // Release staging buffers back to pool. for (auto& buf : staging_bufs) { - pinned_buffer_pool_->Release(buf); + pinned_buffer_pool_->Release(std::move(buf)); } if (!offload_res) { LOG(ERROR) << "Failed to store objects with error: " << offload_res.error(); + // This bucket did not persist, so report its keys back to the + // master as failed regardless of whether we continue or abort. + // Doing it here (rather than only on the soft path) keeps their + // offloading tasks and source-replica refcounts from leaking until + // the put_start_release_timeout_sec_ TTL reaper fires. + for (const auto& [key, _] : host_batch_object) { + failed_tasks.push_back(task_by_storage_key.at(key)); + } if (offload_res.error() == ErrorCode::KEYS_ULTRA_LIMIT) { + // Disk is over the key-count limit: stop offloading entirely. MutexLocker locker(&offloading_mutex_); enable_offloading_ = false; - return tl::make_unexpected(offload_res.error()); } - if (offload_res.error() != ErrorCode::INVALID_READ) { - return tl::make_unexpected(offload_res.error()); + if (!IsPerBucketSoftOffloadError(offload_res.error())) { + // Whole-cycle error (KEYS_ULTRA_LIMIT or any hard failure): + // stop processing further buckets, but fall through to the NACK + // flush below so every drained key is released. Unvisited + // buckets are not yet in all_bucket_keys, so the sweep NACKs + // them too; this bucket's keys were just pushed above. + abort_error = offload_res.error(); + break; } + // Soft per-bucket error: keep processing the remaining buckets. } } + + // Keys skipped by GroupOffloadingKeysByBucket don't appear in any bucket, + // so they never reach BatchOffload or complete_handler. + for (const auto& [storage_key, task] : task_by_storage_key) { + if (all_bucket_keys.find(storage_key) == all_bucket_keys.end()) { + failed_tasks.push_back(task); + } + } + + if (!failed_tasks.empty()) { + std::vector failed_metadatas; + failed_metadatas.reserve(failed_tasks.size()); + for (size_t i = 0; i < failed_tasks.size(); ++i) { + failed_metadatas.push_back(StorageObjectMetadata{-1, 0, 0, -1, ""}); + } + auto result = + client_->NotifyOffloadSuccess(failed_tasks, failed_metadatas); + if (!result) { + LOG(WARNING) << "[OFFLOAD] NotifyOffloadSuccess for failed tasks " + "returned error: " + << result.error() << " count: " << failed_tasks.size(); + } + } + + if (abort_error) { + return tl::make_unexpected(*abort_error); + } + return {}; +} + +tl::expected FileStorage::NotifyEvictedDiskReplicas( + const std::vector& evicted_keys) { + if (evicted_keys.empty()) return {}; + + std::optional first_error; + std::unordered_map, TenantIdHash> + keys_by_tenant; + for (const auto& storage_key : evicted_keys) { + auto [tenant_id, key] = TenantId::ParseScopedKey(storage_key); + keys_by_tenant[tenant_id].push_back(key); + } + + for (const auto& [tenant_id, keys] : keys_by_tenant) { + auto results = client_->BatchEvictDiskReplica(keys, tenant_id.value(), + ReplicaType::LOCAL_DISK); + if (results.size() != keys.size()) { + LOG(ERROR) << "BatchEvictDiskReplica returned " << results.size() + << " result(s) for " << keys.size() + << " key(s), tenant_id=" << tenant_id; + if (!first_error.has_value()) { + first_error = ErrorCode::INTERNAL_ERROR; + } + continue; + } + + for (size_t i = 0; i < results.size(); ++i) { + if (!results[i]) { + if (results[i].error() == ErrorCode::OBJECT_NOT_FOUND) { + VLOG(1) + << "Master no longer tracks evicted local disk key: " + << keys[i] << ", tenant_id=" << tenant_id; + continue; + } + if (!first_error.has_value()) { + first_error = results[i].error(); + } + LOG(WARNING) + << "Failed to notify master about evicted local disk key: " + << keys[i] << ", tenant_id=" << tenant_id + << ", error: " << results[i].error(); + } + } + } + if (first_error.has_value()) { + return tl::make_unexpected(first_error.value()); + } + return {}; +} + +tl::expected FileStorage::RunDiskWatermarkEviction() { + if (!config_.enable_disk_watermark_eviction) { + return {}; + } + + auto eviction_result = storage_backend_->EvictAboveDiskWatermark( + config_.disk_eviction_high_watermark_ratio, + config_.disk_eviction_low_watermark_ratio, + [this](const std::vector& evicted_keys) { + return NotifyEvictedDiskReplicas(evicted_keys); + }); + if (!eviction_result) { + return tl::make_unexpected(eviction_result.error()); + } + if (!eviction_result.value().empty()) { + LOG(INFO) << "Disk watermark eviction removed " + << eviction_result.value().size() << " LOCAL_DISK key(s)"; + } return {}; } @@ -626,6 +831,19 @@ tl::expected FileStorage::Heartbeat() { auto remount_result = client_->MountLocalDiskSegment(enable_offloading_); if (remount_result) { + // Report configured SSD capacity so the Master can + // restore file_total_capacity_ (the denominator in + // "SSD Storage: X / Y"). This was lost on restart; + // re-reporting it here avoids the 0 B display. + if (config_.total_size_limit > 0) { + auto cap_result = client_->ReportSsdCapacity( + config_.total_size_limit); + if (!cap_result) { + LOG(WARNING) + << "ReportSsdCapacity failed during " + << "heartbeat recovery: " << cap_result.error(); + } + } heartbeat_result = client_->OffloadObjectHeartbeat( enable_offloading_, offloading_objects); if (!heartbeat_result) { @@ -663,27 +881,56 @@ tl::expected FileStorage::Heartbeat() { } } - if (offloading_objects.empty()) { - return {}; + // === STEP 2: Poll whether master requested a full SSD clear === + auto remove_all_result = client_->PollRemoveAll(); + if (remove_all_result && remove_all_result.value()) { + RemoveAll(); } - // === STEP 2: Persist offloaded objects (trigger actual data migration) === - auto offload_result = OffloadObjects(offloading_objects); - if (!offload_result) { - LOG(ERROR) << "Failed to persist objects with error: " - << offload_result.error(); - return offload_result; + + // === STEP 3: Persist offloaded objects (trigger actual data migration) === + if (!offloading_objects.empty()) { + auto offload_result = OffloadObjects(offloading_objects); + if (!offload_result) { + LOG(ERROR) << "Failed to persist objects with error: " + << offload_result.error(); + return offload_result; + } } + VLOG(1) << "Completed heartbeat with offloaded objects count: " + << offloading_objects.size(); + // Drive any pending L2->L1 promotion work for this client. Failures // inside ProcessPromotionTasks are logged per-key and do not propagate; // promotion is best-effort and must never break offload. (void)ProcessPromotionTasks(); - // TODO(eviction): Implement an LRU eviction mechanism to manage local - // storage capacity. + // Proactive disk watermarks keep LOCAL_DISK usage below the configured + // low watermark even when no new write arrives to trigger reactive + // eviction. + auto disk_eviction_result = RunDiskWatermarkEviction(); + if (!disk_eviction_result) { + LOG(WARNING) << "Disk watermark eviction failed: " + << disk_eviction_result.error(); + } return {}; } +void FileStorage::RemoveAll() { + // TODO(tenant-isolation): This performs a tenant-UNAWARE global wipe of the + // storage directory. Storage backends store physical files without a + // tenant dimension, so a tenant-scoped master RemoveAll("tenant_A") that + // signals this client will also delete tenant_B's SSD files here, while + // master still holds valid metadata for tenant_B (subsequent reads get + // OBJECT_NOT_FOUND on this node). Safe for the global RemoveAll(force) and + // for single-tenant / shared-nothing deployments. Proper per-tenant + // physical isolation needs backend-level tenant-scoped layout (follow-up). + if (storage_backend_) { + storage_backend_->RemoveAll(); + } + LOG(INFO) << "FileStorage::RemoveAll: cleared storage backend"; +} + tl::expected FileStorage::ProcessPromotionTasks() { if (client_ == nullptr) { return tl::make_unexpected(ErrorCode::INVALID_PARAMS); @@ -724,7 +971,7 @@ tl::expected FileStorage::ProcessPromotionTasks() { const auto& key = task.key; const auto& tenant_id = task.tenant_id; const int64_t size = task.size; - const auto storage_key = MakeTenantScopedStorageKey(tenant_id, key); + const auto storage_key = TenantId(tenant_id).MakeScopedKey(key); if (size <= 0) { LOG(WARNING) << "Skipping promotion for key=" << key << " with non-positive size=" << size; @@ -866,8 +1113,9 @@ tl::expected FileStorage::BatchQuerySegmentSlices( const std::vector& keys, const std::string& tenant_id, std::unordered_map>& batched_slices) { auto batched_query_results = client_->BatchQuery(keys, tenant_id); - if (batched_query_results.empty()) - return tl::make_unexpected(ErrorCode::INVALID_REPLICA); + if (batched_query_results.empty()) { + return {}; + } for (size_t i = 0; i < batched_query_results.size(); ++i) { if (batched_query_results[i]) { for (const auto& descriptor : @@ -884,22 +1132,21 @@ tl::expected FileStorage::BatchQuerySegmentSlices( break; } } - if (batched_slices.find(keys[i]) == batched_slices.end()) { - LOG(ERROR) << "Key not found: " << keys[i]; - return tl::make_unexpected(ErrorCode::INVALID_KEY); - } - } else { - LOG(ERROR) << "Key not found: " << keys[i]; - return tl::make_unexpected(batched_query_results[i].error()); } } return {}; } tl::expected FileStorage::RegisterLocalMemory() { + // The buffer pool backs SSD-offload read results that are fetched by + // remote peers via RDMA READ. It must therefore be registered with + // remote_accessible=true so its BufferDesc publishes an rkey; otherwise + // every remote read of an offloaded object fails with "No rkey for MR + // access" (the pool is only reachable through the address-range lookup, + // and with remote_accessible=false the rkey array is left empty). auto error_code = client_->RegisterLocalMemory( client_buffer_allocator_->getBase(), config_.local_buffer_size, - kWildcardLocation, false, true); + kWildcardLocation, true, true); if (!error_code) { LOG(ERROR) << "Failed to register local memory: " << error_code.error(); return error_code; diff --git a/mooncake-store/src/ha/kv/etcd_ha_kv_backend.cpp b/mooncake-store/src/ha/kv/etcd_ha_kv_backend.cpp new file mode 100644 index 0000000000..530d62c449 --- /dev/null +++ b/mooncake-store/src/ha/kv/etcd_ha_kv_backend.cpp @@ -0,0 +1,94 @@ +#include "ha/kv/etcd_ha_kv_backend.h" + +#include +#include + +#include + +#if __has_include() +#include +#else +#include +#endif + +#include "etcd_helper.h" + +namespace mooncake { +namespace { + +EtcdHelper::TxnCompareKind ToEtcdCompareKind(KvCompareKind kind) { + switch (kind) { + case KvCompareKind::kValueEquals: + return EtcdHelper::TxnCompareKind::kValueEquals; + case KvCompareKind::kKeyNotExists: + return EtcdHelper::TxnCompareKind::kKeyNotExists; + } + return EtcdHelper::TxnCompareKind::kValueEquals; +} + +} // namespace + +ErrorCode EtcdHaKvBackend::Get(std::string_view key, std::string& value) { + EtcdRevisionId revision_id = 0; + return EtcdHelper::Get(key.data(), key.size(), value, revision_id); +} + +ErrorCode EtcdHaKvBackend::Put(std::string_view key, std::string_view value) { + return EtcdHelper::Put(key.data(), key.size(), value.data(), value.size()); +} + +ErrorCode EtcdHaKvBackend::Range(std::string_view begin_key, + std::string_view end_key, size_t limit, + std::vector& kvs) { + kvs.clear(); + std::string json; + EtcdRevisionId revision_id = 0; + ErrorCode err = EtcdHelper::GetRangeAsJson( + begin_key.data(), begin_key.size(), end_key.data(), end_key.size(), + limit, json, revision_id); + if (err != ErrorCode::OK) { + return err; + } + + Json::Value root; + Json::CharReaderBuilder reader; + std::string errors; + std::istringstream stream(json); + if (!Json::parseFromStream(reader, stream, &root, &errors) || + !root.isArray()) { + LOG(ERROR) << "Failed to parse etcd range JSON: " << errors; + return ErrorCode::INTERNAL_ERROR; + } + + kvs.reserve(root.size()); + for (const auto& item : root) { + if (!item.isObject() || !item["key"].isString() || + !item["value"].isString()) { + return ErrorCode::INTERNAL_ERROR; + } + kvs.push_back( + {.key = item["key"].asString(), .value = item["value"].asString()}); + } + return ErrorCode::OK; +} + +bool EtcdHaKvBackend::SupportsTxn() const { return true; } + +ErrorCode EtcdHaKvBackend::Txn(const KvTxn& txn) { + std::vector compares; + compares.reserve(txn.compares.size()); + for (const auto& compare : txn.compares) { + compares.push_back({.key = compare.key, + .kind = ToEtcdCompareKind(compare.kind), + .expected_value = compare.expected_value}); + } + + std::vector puts; + puts.reserve(txn.puts.size()); + for (const auto& put : txn.puts) { + puts.push_back({.key = put.key, .value = put.value}); + } + return EtcdHelper::TxnCompareAndPut(compares, puts); +} + +} // namespace mooncake diff --git a/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp b/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp index 95b09a9094..26bc7b23e7 100644 --- a/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp +++ b/mooncake-store/src/ha/leadership/backends/etcd/etcd_leader_coordinator.cpp @@ -31,6 +31,10 @@ constexpr auto kViewChangeFallbackPollInterval = std::chrono::milliseconds(200); // watch context is destroyed. Mirrors the value used by the oplog notifier. constexpr int kWatchStopTimeoutMs = 5000; +// event_type delivered by the etcd watch goroutine when the watch itself ends +// (0 = PUT, 1 = DELETE, 2 = WATCH_BROKEN; see EtcdHelper::WatchCallbackFn). +constexpr int kWatchEventBroken = 2; + // Shared state between WaitForViewChange and the etcd watch callback. The // callback only flips `changed` and wakes the waiter; the waiter decides what // the change means by re-reading the view. @@ -38,6 +42,7 @@ struct ViewChangeWatchState { std::mutex mutex; std::condition_variable cv; bool changed = false; + bool broken = false; }; // C-style trampoline invoked by the etcd watch goroutine for every event on the @@ -47,7 +52,7 @@ struct ViewChangeWatchState { // CancelWatchWithPrefix + WaitWatchWithPrefixStopped before it is destroyed. void ViewChangeWatchCallback(void* context, const char* /*key*/, size_t /*key_size*/, const char* /*value*/, - size_t /*value_size*/, int /*event_type*/, + size_t /*value_size*/, int event_type, int64_t /*mod_revision*/) { auto* state = static_cast(context); if (state == nullptr) { @@ -55,6 +60,9 @@ void ViewChangeWatchCallback(void* context, const char* /*key*/, } { std::lock_guard lock(state->mutex); + if (event_type == kWatchEventBroken) { + state->broken = true; + } state->changed = true; } state->cv.notify_all(); @@ -362,7 +370,53 @@ EtcdLeaderCoordinator::WaitForViewChange( // store client. Revisit (e.g. per-waiter watch IDs) if multiple // same-process clients ever need to watch the same master_view // concurrently. + // + // Arm ONE watch for the whole wait instead of a fresh one per loop + // iteration. The old per-iteration cycle (new state + arm + guard-destruct + // with a kWatchStopTimeoutMs stop budget) leaked a ViewChangeWatchState and + // its watch goroutine every time the goroutine missed that budget, which + // accumulated OS threads under load in HA mode (issue #3059). A persistent + // watch serves every wait below and is re-armed only when etcd reports it + // broken, so the bounded one-time cleanup in PrefixWatchGuard at function + // exit is the only cancel/stop cycle left. + auto* state = new ViewChangeWatchState(); + PrefixWatchGuard guard(master_view_key_, state); + + // Arm the watch BEFORE the first read, preserving the ordering invariant: + // the watch is established at some etcd revision R_watch and the read + // observes R_read >= R_watch, so a change at revision C is caught by the + // read (C < R_read) or the watch (C >= R_watch), never missed in between. + // The watch starts from the current revision (start_revision = 0), which + // also avoids depending on a possibly-compacted historical revision. Later + // iterations re-read under the same persistent watch, so the invariant + // holds for the whole wait. + auto arm_watch = [this, state, &guard]() { + // Defensively clear any lingering watch on this key, then arm a new + // one. Both calls are no-ops when nothing is registered. On the re-arm + // path the old watch is already broken, so the stop wait returns fast. + EtcdHelper::CancelWatchWithPrefix(master_view_key_.c_str(), + master_view_key_.size()); + EtcdHelper::WaitWatchWithPrefixStopped(master_view_key_.c_str(), + master_view_key_.size(), + kWatchStopTimeoutMs); + { + std::lock_guard lock(state->mutex); + state->changed = false; + state->broken = false; + } + auto watch_err = EtcdHelper::WatchWithPrefixFromRevision( + master_view_key_.c_str(), master_view_key_.size(), + /*start_revision=*/0, state, &ViewChangeWatchCallback); + if (watch_err == ErrorCode::OK) { + guard.Arm(); + return true; + } + return false; + }; + const auto deadline = std::chrono::steady_clock::now() + timeout; + bool watching = timeout > kViewChangeFallbackPollInterval && arm_watch(); + while (true) { if (timeout <= std::chrono::milliseconds::zero() || std::chrono::steady_clock::now() >= deadline) { @@ -377,41 +431,6 @@ EtcdLeaderCoordinator::WaitForViewChange( std::chrono::duration_cast( deadline - std::chrono::steady_clock::now()); - // Arm the watch BEFORE reading the view. The watch is established at - // some etcd revision R_watch; the subsequent read observes a revision - // R_read >= R_watch. Any change to master_view at revision C is then - // caught by exactly one of the two: C < R_read (the read sees it) or - // C >= R_watch (the watch delivers it), and the two ranges overlap, so - // a change happening between read and watch cannot be missed. The watch - // starts from the current revision (start_revision = 0), which also - // avoids depending on a possibly-compacted historical revision. - // - // The watch state is heap-allocated and owned by `guard`. The guard - // cancels the watch and waits for the goroutine to exit before - // releasing the state in its destructor; if the goroutine fails to stop - // in time it leaks the state instead of freeing it, so an in-flight - // callback can never reference freed memory (see PrefixWatchGuard). - auto* state = new ViewChangeWatchState(); - PrefixWatchGuard guard(master_view_key_, state); - - bool watching = false; - if (remaining > kViewChangeFallbackPollInterval) { - // Defensively clear any lingering watch on this key, then arm a new - // one. Both calls are no-ops when nothing is registered. - EtcdHelper::CancelWatchWithPrefix(master_view_key_.c_str(), - master_view_key_.size()); - EtcdHelper::WaitWatchWithPrefixStopped(master_view_key_.c_str(), - master_view_key_.size(), - kWatchStopTimeoutMs); - auto watch_err = EtcdHelper::WatchWithPrefixFromRevision( - master_view_key_.c_str(), master_view_key_.size(), - /*start_revision=*/0, state, &ViewChangeWatchCallback); - if (watch_err == ErrorCode::OK) { - watching = true; - guard.Arm(); - } - } - auto current_view = ReadCurrentView(); if (!current_view) { return tl::make_unexpected(current_view.error()); @@ -426,19 +445,35 @@ EtcdLeaderCoordinator::WaitForViewChange( if (!watching) { // Could not arm a watch (RPC failed, or too little time left). - // Fall back to a short poll so a change is still picked up. + // Retry the arm once there is room again, otherwise fall back to a + // short poll so a change is still picked up. + if (remaining > kViewChangeFallbackPollInterval && arm_watch()) { + watching = true; + continue; + } std::this_thread::sleep_for( std::min(kViewChangeFallbackPollInterval, remaining)); continue; } // Block until the watch reports an event or the caller's deadline - // elapses. Either way we loop and re-read: an event tells us the view - // changed (re-read returns it), a timeout falls through to the deadline - // check above and returns timed_out. - std::unique_lock lock(state->mutex); - state->cv.wait_for(lock, remaining, - [state]() { return state->changed; }); + // elapses. An event just clears the flag and the same persistent watch + // keeps serving the next wait; only a broken watch needs a new + // goroutine. A timeout falls through to the deadline check above and + // returns timed_out. + bool should_rearm = false; + { + std::unique_lock lock(state->mutex); + state->cv.wait_for(lock, remaining, [state]() { + return state->changed || state->broken; + }); + should_rearm = state->broken; + state->changed = false; + } + if (should_rearm) { + watching = + remaining > kViewChangeFallbackPollInterval && arm_watch(); + } } } diff --git a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp b/mooncake-store/src/ha/leadership/master_service_supervisor.cpp index c2caadc10b..a60995a353 100644 --- a/mooncake-store/src/ha/leadership/master_service_supervisor.cpp +++ b/mooncake-store/src/ha/leadership/master_service_supervisor.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -12,7 +13,9 @@ #include #include "ha/leadership/leader_coordinator_factory.h" +#include "ha/leadership/leader_label_reconciler.h" #include "ha/standby_controller.h" +#include "k8s_lease_helper.h" #include "master_admin_service.h" #include "rpc_service.h" @@ -24,6 +27,27 @@ namespace { constexpr auto kAcquireRetryInterval = std::chrono::seconds(1); constexpr auto kRenewCheckInterval = std::chrono::seconds(1); constexpr auto kSupervisorRetryInterval = std::chrono::seconds(1); +constexpr auto kLabelReconcileRetryInterval = std::chrono::seconds(1); +constexpr char kLeaderLabelKey[] = "mooncake.io/store-role"; +constexpr char kLeaderLabelValue[] = "leader"; + +bool HasPodIdentity(const MasterServiceSupervisorConfig& config) { + return !config.pod_name.empty() && !config.pod_namespace.empty() && + config.ha_backend_type == "k8s"; +} + +LeaderLabelReconciler MakeLeaderLabelReconciler( + const MasterServiceSupervisorConfig& config) { + return LeaderLabelReconciler( + HasPodIdentity(config), + [ns = config.pod_namespace, pod = config.pod_name](bool desired) { + return desired ? K8sLeaseHelper::SetPodLabel( + ns, pod, kLeaderLabelKey, kLeaderLabelValue) + : K8sLeaseHelper::ClearPodLabel(ns, pod, + kLeaderLabelKey); + }, + kLabelReconcileRetryInterval); +} std::string ResolveHABackendConnstring( const MasterServiceSupervisorConfig& config) { @@ -131,17 +155,20 @@ void SetRuntimeState(MasterAdminServer& admin_server, << ", role=" << MasterRuntimeRoleToString(state); } -void ActivateServingState( - MasterAdminServer& admin_server, - const std::shared_ptr& service) { +void ActivateServingState(MasterAdminServer& admin_server, + const std::shared_ptr& service, + LeaderLabelReconciler& label_reconciler) { admin_server.SetServiceDelegate(service); admin_server.SetServiceAvailable(true); SetRuntimeState(admin_server, MasterRuntimeState::kServing); + label_reconciler.SetLeader(true); } -void DeactivateServingState(MasterAdminServer& admin_server) { +void DeactivateServingState(MasterAdminServer& admin_server, + LeaderLabelReconciler& label_reconciler) { admin_server.SetServiceAvailable(false); admin_server.SetServiceDelegate(nullptr); + label_reconciler.SetLeader(false); } void StopLeadershipMonitor(std::unique_ptr& monitor) { @@ -191,6 +218,8 @@ void EnterStandbyMode(MasterAdminServer& admin_server, int RunSupervisorLoop(const HABackendSpec& spec, const MasterServiceSupervisorConfig& config, MasterAdminServer& admin_server) { + auto label_reconciler = MakeLeaderLabelReconciler(config); + label_reconciler.SetLeader(false); SetRuntimeState(admin_server, MasterRuntimeState::kStarting); auto standby_controller = CreateStandbyController(spec, config); std::atomic accept_standby_runtime_updates{false}; @@ -302,13 +331,13 @@ int RunSupervisorLoop(const HABackendSpec& spec, } accept_standby_runtime_updates.store(false, std::memory_order_release); - auto promote_standby = standby_controller->PromoteStandby(); - if (promote_standby != ErrorCode::OK) { + auto promotion_ctx = standby_controller->PromoteStandbyAndExport(); + if (!promotion_ctx) { EnterStandbyMode(admin_server, *standby_controller, accept_standby_runtime_updates, leadership_session->view); if (HandleSupervisorError("promote standby for serve", - promote_standby, spec.type)) { + promotion_ctx.error(), spec.type)) { return -1; } continue; @@ -350,15 +379,32 @@ int RunSupervisorLoop(const HABackendSpec& spec, server.init_ibv(); } + mooncake::WrappedMasterServiceConfig wrapped_config( + config, leadership_session->view.view_version); + // In HA serving-primary mode, snapshot bootstrap belongs to standby. + // The new primary must restore from PromotionContext only. + wrapped_config.enable_snapshot_restore = false; + // The serving primary handles heartbeats/unmounts, so forward the + // metadata cleanup config here like the non-HA path does. auto wrapped_master_service = std::make_shared( - mooncake::WrappedMasterServiceConfig( - config, leadership_session->view.view_version)); + wrapped_config, config.http_metadata_server, + config.http_metadata_remote_url); + + // Restore from standby if we have context + if (promotion_ctx->applied_seq_id > 0 || + !promotion_ctx->objects.empty() || + !promotion_ctx->segments.empty()) { + wrapped_master_service->RestoreFromStandby( + promotion_ctx->objects, promotion_ctx->applied_seq_id, + promotion_ctx->segments); + } + mooncake::RegisterRpcService(server, *wrapped_master_service); auto serve_preflight = leader_coordinator.RenewLeadership(*leadership_session); if (!serve_preflight) { - DeactivateServingState(admin_server); + DeactivateServingState(admin_server, label_reconciler); EnterStandbyMode(admin_server, *standby_controller, accept_standby_runtime_updates, leadership_session->view); @@ -371,7 +417,7 @@ int RunSupervisorLoop(const HABackendSpec& spec, continue; } if (!serve_preflight.value()) { - DeactivateServingState(admin_server); + DeactivateServingState(admin_server, label_reconciler); EnterStandbyMode(admin_server, *standby_controller, accept_standby_runtime_updates, std::nullopt); LogLeadershipReleaseWarning( @@ -385,16 +431,18 @@ int RunSupervisorLoop(const HABackendSpec& spec, std::atomic serve_shutdown_requested{false}; auto leadership_monitor = leader_coordinator.StartLeadershipMonitor( *leadership_session, - [&server, &admin_server, &serve_shutdown_requested](auto reason) { + [&server, &admin_server, &serve_shutdown_requested, + &label_reconciler](auto reason) { serve_shutdown_requested.store(true, std::memory_order_release); admin_server.SetServiceAvailable(false); + label_reconciler.SetLeader(false); SetRuntimeState(admin_server, MasterRuntimeState::kStandby); LOG(INFO) << "Trying to stop server, reason=" << LeadershipLossReasonToString(reason); server.stop(); }); if (!leadership_monitor) { - DeactivateServingState(admin_server); + DeactivateServingState(admin_server, label_reconciler); EnterStandbyMode(admin_server, *standby_controller, accept_standby_runtime_updates, leadership_session->view); @@ -413,7 +461,7 @@ int RunSupervisorLoop(const HABackendSpec& spec, LOG(ERROR) << "Failed to start master service: " << ec.result().value(); StopLeadershipMonitor(leadership_monitor_handle); - DeactivateServingState(admin_server); + DeactivateServingState(admin_server, label_reconciler); EnterStandbyMode(admin_server, *standby_controller, accept_standby_runtime_updates, leadership_session->view); @@ -426,14 +474,15 @@ int RunSupervisorLoop(const HABackendSpec& spec, } if (!serve_shutdown_requested.load(std::memory_order_acquire)) { - ActivateServingState(admin_server, wrapped_master_service); + ActivateServingState(admin_server, wrapped_master_service, + label_reconciler); } auto server_err = std::move(ec).get(); LOG(ERROR) << "Master service stopped: " << server_err; StopLeadershipMonitor(leadership_monitor_handle); - DeactivateServingState(admin_server); + DeactivateServingState(admin_server, label_reconciler); auto err = leader_coordinator.ReleaseLeadership(*leadership_session); LOG(INFO) << "Release leadership: " << toString(err); auto current_view = leader_coordinator.ReadCurrentView(); diff --git a/mooncake-store/src/ha/oplog/etcd_oplog_change_notifier.cpp b/mooncake-store/src/ha/oplog/etcd_oplog_change_notifier.cpp deleted file mode 100644 index 5437976829..0000000000 --- a/mooncake-store/src/ha/oplog/etcd_oplog_change_notifier.cpp +++ /dev/null @@ -1,430 +0,0 @@ -#include "ha/oplog/etcd_oplog_change_notifier.h" - -#include -#include -#include -#include -#include - -#ifdef STORE_USE_ETCD -#include "etcd_helper.h" -#include "ha_metric_manager.h" -#include "ha/oplog/oplog_manager.h" -#include "ha/oplog/oplog_serializer.h" - -namespace mooncake { - -EtcdOpLogChangeNotifier::EtcdOpLogChangeNotifier(const std::string& cluster_id, - EtcdOpLogStore* oplog_store) - : cluster_id_(cluster_id), oplog_store_(oplog_store) { - if (!NormalizeAndValidateClusterId(cluster_id_)) { - LOG(FATAL) << "Invalid cluster_id for EtcdOpLogChangeNotifier: '" - << cluster_id_ - << "'. Allowed chars: [A-Za-z0-9_.-], max_len=128."; - } - watch_prefix_ = "/oplog/" + cluster_id_ + "/"; - - callback_ctx_ = new ChangeNotifierCallbackContext(); - callback_ctx_->notifier = this; -} - -EtcdOpLogChangeNotifier::~EtcdOpLogChangeNotifier() { - Stop(); - // If Stop() returned early (Change Notifier was never started), - // callback_ctx_ was never freed. It is safe to delete here because no - // goroutine / watch thread was ever launched, so no callbacks can be - // in-flight. In all other paths, Stop() already sets callback_ctx_ to - // nullptr, so `delete nullptr` is a harmless no-op. - delete callback_ctx_; - callback_ctx_ = nullptr; -} - -ErrorCode EtcdOpLogChangeNotifier::Start(uint64_t start_sequence_id, - EntryCallback on_entry, - ErrorCallback on_error) { - if (running_.load()) { - LOG(WARNING) << "EtcdOpLogChangeNotifier is already running"; - return ErrorCode::OK; - } - - on_entry_ = std::move(on_entry); - on_error_ = std::move(on_error); - last_processed_sequence_id_.store(start_sequence_id); - - // Initial sync: read and deliver historical entries - if (oplog_store_) { - int64_t delivered = DeliverHistoricalEntries(start_sequence_id); - if (delivered < 0) { - next_watch_revision_.store(0); - } - LOG(INFO) << "EtcdOpLogChangeNotifier initial sync done" - << ", last_seq=" << last_processed_sequence_id_.load() - << ", next_watch_revision=" << next_watch_revision_.load(); - } - - running_.store(true); - watch_thread_ = std::thread(&EtcdOpLogChangeNotifier::WatchLoop, this); - LOG(INFO) << "EtcdOpLogChangeNotifier started for cluster_id=" - << cluster_id_; - return ErrorCode::OK; -} - -void EtcdOpLogChangeNotifier::Stop() { - if (!running_.load()) { - return; - } - - running_.store(false); - - // Invalidate the callback context under the mutex. - if (callback_ctx_) { - std::lock_guard lock(callback_ctx_->mutex); - callback_ctx_->notifier = nullptr; - } - - if (watch_thread_.joinable()) { - watch_thread_.join(); - } - - // Cancel the Go goroutine and wait for it to fully exit. - ErrorCode err = EtcdHelper::CancelWatchWithPrefix(watch_prefix_.c_str(), - watch_prefix_.size()); - if (err != ErrorCode::OK) { - LOG(WARNING) << "Failed to cancel watch for prefix " << watch_prefix_ - << ", error=" << static_cast(err); - } - - ErrorCode wait_err = EtcdHelper::WaitWatchWithPrefixStopped( - watch_prefix_.c_str(), watch_prefix_.size(), /*timeout_ms=*/5000); - - if (wait_err == ErrorCode::OK) { - delete callback_ctx_; - } else { - LOG(WARNING) << "Watch goroutine did not stop in time for prefix " - << watch_prefix_ - << "; leaking ChangeNotifierCallbackContext to avoid UAF"; - } - callback_ctx_ = nullptr; - - LOG(INFO) << "EtcdOpLogChangeNotifier stopped"; -} - -bool EtcdOpLogChangeNotifier::IsHealthy() const { - return watch_healthy_.load(); -} - -bool EtcdOpLogChangeNotifier::ReadOpLogSince(uint64_t start_seq_id, - std::vector& entries, - EtcdRevisionId& revision_id) { - if (!oplog_store_) { - return false; - } - ErrorCode err = oplog_store_->ReadOpLogSinceWithRevision( - start_seq_id, kSyncBatchSize, entries, revision_id); - if (err != ErrorCode::OK) { - LOG(ERROR) << "Failed to read OpLog since sequence_id=" << start_seq_id - << ", error=" << static_cast(err); - return false; - } - return true; -} - -void EtcdOpLogChangeNotifier::WatchCallback(void* context, const char* key, - size_t key_size, const char* value, - size_t value_size, int event_type, - int64_t mod_revision) { - auto* ctx = static_cast(context); - if (ctx == nullptr) { - return; - } - - std::lock_guard lock(ctx->mutex); - EtcdOpLogChangeNotifier* notifier = ctx->notifier; - if (notifier == nullptr) { - return; - } - - if (!notifier->running_.load(std::memory_order_acquire)) { - return; - } - - std::string key_str; - if (key != nullptr && key_size > 0) { - key_str.assign(key, key_size); - } - std::string value_str; - if (value != nullptr && value_size > 0) { - value_str = std::string(value, value_size); - } - notifier->HandleWatchEvent(key_str, value_str, event_type, mod_revision); -} - -void EtcdOpLogChangeNotifier::WatchLoop() { - LOG(INFO) << "OpLog watch thread started for cluster_id=" << cluster_id_; - - while (running_.load()) { - // Cancel any existing watch before starting a new one - (void)EtcdHelper::CancelWatchWithPrefix(watch_prefix_.c_str(), - watch_prefix_.size()); - (void)EtcdHelper::WaitWatchWithPrefixStopped(watch_prefix_.c_str(), - watch_prefix_.size(), - /*timeout_ms=*/5000); - - EtcdRevisionId start_rev = - static_cast(next_watch_revision_.load()); - ErrorCode err = EtcdHelper::WatchWithPrefixFromRevision( - watch_prefix_.c_str(), watch_prefix_.size(), start_rev, - callback_ctx_, WatchCallback); - - if (err != ErrorCode::OK) { - LOG(ERROR) << "Failed to start watch for prefix " << watch_prefix_ - << ", error=" << static_cast(err); - watch_healthy_.store(false); - if (on_error_) { - on_error_(err); - } - - std::this_thread::sleep_for(std::chrono::milliseconds(500)); - TryReconnect(); - continue; - } - - LOG(INFO) << "Watch started for prefix " << watch_prefix_; - watch_healthy_.store(true); - consecutive_errors_.store(0); - - while (running_.load() && watch_healthy_.load()) { - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - - if (consecutive_errors_.load() >= kMaxConsecutiveErrors) { - LOG(WARNING) - << "Too many consecutive errors (" - << consecutive_errors_.load() << "), reconnecting watch..."; - watch_healthy_.store(false); - break; - } - } - - if (running_.load() && !watch_healthy_.load()) { - (void)EtcdHelper::CancelWatchWithPrefix(watch_prefix_.c_str(), - watch_prefix_.size()); - (void)EtcdHelper::WaitWatchWithPrefixStopped(watch_prefix_.c_str(), - watch_prefix_.size(), - /*timeout_ms=*/5000); - TryReconnect(); - } - } - - LOG(INFO) << "OpLog watch thread stopped"; -} - -void EtcdOpLogChangeNotifier::HandleWatchEvent(const std::string& key, - const std::string& value, - int event_type, - int64_t mod_revision) { - // event_type: 0 = PUT, 1 = DELETE, 2 = WATCH_BROKEN - if (event_type == 2) { - LOG(WARNING) << "OpLog watch broken, will reconnect. cluster_id=" - << cluster_id_ - << ", next_watch_revision=" << next_watch_revision_.load() - << ", last_seq=" << last_processed_sequence_id_.load(); - watch_healthy_.store(false); - consecutive_errors_.fetch_add(1); - return; - } - - if (mod_revision > 0) { - int64_t candidate = mod_revision + 1; - int64_t cur = next_watch_revision_.load(); - while (candidate > cur && - !next_watch_revision_.compare_exchange_weak(cur, candidate)) { - } - } - - if (event_type == 1) { - VLOG(1) << "OpLog entry deleted: " << key; - consecutive_errors_.store(0); - return; - } - - if (event_type != 0) { - LOG(WARNING) << "Unknown event type: " << event_type - << " for key: " << key; - consecutive_errors_.fetch_add(1); - return; - } - - // Skip the "latest" key and snapshot keys - if (key.find("/latest") != std::string::npos || - key.find("/snapshot/") != std::string::npos) { - return; - } - - // Parse the OpLog entry (DeserializeOpLogEntry validates entry size) - OpLogEntry entry; - if (!DeserializeOpLogEntry(value, entry)) { - LOG(ERROR) << "Failed to deserialize OpLog entry from key: " << key; - consecutive_errors_.fetch_add(1); - return; - } - - // Verify checksum at the trust boundary (data from etcd) - if (!OpLogManager::VerifyChecksum(entry)) { - LOG(ERROR) << "OpLog entry checksum mismatch: sequence_id=" - << entry.sequence_id << ", key=" << entry.object_key - << ". Possible data corruption. Discarding entry."; - consecutive_errors_.fetch_add(1); - HAMetricManager::instance().inc_oplog_checksum_failures(); - return; - } - - // Deliver to callback - if (on_entry_) { - on_entry_(entry); - } - - // Update last processed (monotonic) - uint64_t cur = last_processed_sequence_id_.load(); - while (IsSequenceNewer(entry.sequence_id, cur) && - !last_processed_sequence_id_.compare_exchange_weak( - cur, entry.sequence_id)) { - } - - consecutive_errors_.store(0); - reconnect_count_.store(0); - VLOG(2) << "Delivered OpLog entry: sequence_id=" << entry.sequence_id - << ", op_type=" << static_cast(entry.op_type) - << ", key=" << entry.object_key; -} - -void EtcdOpLogChangeNotifier::TryReconnect() { - if (!running_.load()) { - return; - } - - int reconnect_attempt = reconnect_count_.fetch_add(1) + 1; - int delay_ms = - std::min(kReconnectDelayMs * reconnect_attempt, kMaxReconnectDelayMs); - - LOG(INFO) << "Attempting to reconnect watch (attempt #" << reconnect_attempt - << "), waiting " << delay_ms << "ms..."; - - std::this_thread::sleep_for(std::chrono::milliseconds(delay_ms)); - - if (SyncMissedEntries()) { - LOG(INFO) << "Successfully synced missed OpLog entries"; - } else { - LOG(WARNING) - << "Failed to sync missed OpLog entries, continuing anyway"; - } -} - -bool EtcdOpLogChangeNotifier::SyncMissedEntries() { - uint64_t last_seq = last_processed_sequence_id_.load(); - if (last_seq == 0) { - return true; - } - LOG(INFO) << "Syncing missed OpLog entries since sequence_id=" << last_seq; - int64_t delivered = DeliverHistoricalEntries(last_seq); - if (delivered < 0) { - return false; - } - LOG(INFO) << "Synced " << delivered << " missed OpLog entries"; - return true; -} - -int64_t EtcdOpLogChangeNotifier::DeliverHistoricalEntries( - uint64_t start_seq_id) { - uint64_t read_seq_id = start_seq_id; - EtcdRevisionId last_read_rev = 0; - int64_t total_delivered = 0; - - for (;;) { - std::vector batch; - EtcdRevisionId rev = 0; - if (!ReadOpLogSince(read_seq_id, batch, rev)) { - return -1; - } - if (rev > 0) { - last_read_rev = rev; - } - for (const auto& entry : batch) { - if (on_entry_) { - on_entry_(entry); - } - last_processed_sequence_id_.store(entry.sequence_id); - read_seq_id = entry.sequence_id; - total_delivered++; - } - if (batch.size() < kSyncBatchSize) { - break; - } - } - - if (last_read_rev > 0) { - next_watch_revision_.store(static_cast(last_read_rev + 1)); - } - - return total_delivered; -} - -} // namespace mooncake - -#else // STORE_USE_ETCD not defined - -namespace mooncake { - -EtcdOpLogChangeNotifier::EtcdOpLogChangeNotifier(const std::string& cluster_id, - EtcdOpLogStore* oplog_store) - : cluster_id_(cluster_id), oplog_store_(oplog_store) { - callback_ctx_ = new ChangeNotifierCallbackContext(); - callback_ctx_->notifier = this; -} - -EtcdOpLogChangeNotifier::~EtcdOpLogChangeNotifier() { - Stop(); - delete callback_ctx_; - callback_ctx_ = nullptr; -} - -ErrorCode EtcdOpLogChangeNotifier::Start(uint64_t /*start_sequence_id*/, - EntryCallback /*on_entry*/, - ErrorCallback /*on_error*/) { - LOG(ERROR) << "EtcdOpLogChangeNotifier requires STORE_USE_ETCD"; - return ErrorCode::ETCD_OPERATION_ERROR; -} - -void EtcdOpLogChangeNotifier::Stop() {} - -bool EtcdOpLogChangeNotifier::IsHealthy() const { return false; } - -bool EtcdOpLogChangeNotifier::ReadOpLogSince( - uint64_t /*start_seq_id*/, std::vector& /*entries*/, - EtcdRevisionId& /*revision_id*/) { - return false; -} - -void EtcdOpLogChangeNotifier::WatchCallback( - void* /*context*/, const char* /*key*/, size_t /*key_size*/, - const char* /*value*/, size_t /*value_size*/, int /*event_type*/, - int64_t /*mod_revision*/) {} - -void EtcdOpLogChangeNotifier::WatchLoop() {} - -void EtcdOpLogChangeNotifier::HandleWatchEvent(const std::string& /*key*/, - const std::string& /*value*/, - int /*event_type*/, - int64_t /*mod_revision*/) {} - -void EtcdOpLogChangeNotifier::TryReconnect() {} - -bool EtcdOpLogChangeNotifier::SyncMissedEntries() { return false; } - -int64_t EtcdOpLogChangeNotifier::DeliverHistoricalEntries( - uint64_t /*start_seq_id*/) { - return -1; -} - -} // namespace mooncake - -#endif // STORE_USE_ETCD diff --git a/mooncake-store/src/ha/oplog/etcd_oplog_store.cpp b/mooncake-store/src/ha/oplog/etcd_oplog_store.cpp deleted file mode 100644 index 8a7cc1e1ef..0000000000 --- a/mooncake-store/src/ha/oplog/etcd_oplog_store.cpp +++ /dev/null @@ -1,702 +0,0 @@ -#include "ha/oplog/etcd_oplog_store.h" - -#include -#include -#include - -#include "ha_metric_manager.h" -#include "ha/oplog/oplog_serializer.h" -#include "ha/oplog/etcd_oplog_change_notifier.h" -#include "utils/base64.h" - -#if __has_include() -#include // Ubuntu -#else -#include // CentOS -#endif - -#include "etcd_helper.h" - -namespace mooncake { - -EtcdOpLogStore::EtcdOpLogStore(const std::string& cluster_id, - bool enable_latest_seq_batch_update, - bool enable_batch_write) - : cluster_id_(cluster_id), - enable_latest_seq_batch_update_(enable_latest_seq_batch_update), - enable_batch_write_(enable_batch_write), - last_update_time_(std::chrono::steady_clock::now()) { - if (!NormalizeAndValidateClusterId(cluster_id_)) { - LOG(FATAL) - << "Invalid cluster_id for EtcdOpLogStore: '" << cluster_id_ - << "'. Allowed chars: [A-Za-z0-9_.-], max_len=128, no slashes."; - } -} - -ErrorCode EtcdOpLogStore::Init() { - // Initialize /latest key to 0 if it doesn't exist (first startup). - // This avoids "key not found" errors when querying the latest sequence ID. - // Important: Only initialize if the key doesn't exist to avoid overwriting - // existing data. - // Note: We allow any instance (reader or writer) to ensure the key exists, - // because snapshot resolution depends on it regardless of the HA role. - // The Create call uses CAS semantics so concurrent initializers are safe. - if (!cluster_id_.empty()) { - std::string latest_key = BuildLatestKey(); - std::string existing_value; - EtcdRevisionId revision_id; - ErrorCode get_err = EtcdHelper::Get( - latest_key.c_str(), latest_key.size(), existing_value, revision_id); - if (get_err == ErrorCode::ETCD_KEY_NOT_EXIST) { - // Key doesn't exist, safe to initialize to 0 - std::string initial_value = "0"; - ErrorCode create_err = - EtcdHelper::Create(latest_key.c_str(), latest_key.size(), - initial_value.c_str(), initial_value.size()); - if (create_err == ErrorCode::OK) { - LOG(INFO) << "Initialized /latest key to 0 for cluster_id=" - << cluster_id_; - } else if (create_err == ErrorCode::ETCD_TRANSACTION_FAIL) { - // Race condition: another instance created it between Get and - // Create - LOG(INFO) << "/latest key was created by another instance for " - "cluster_id=" - << cluster_id_; - } else { - // Other errors (e.g., etcd not connected) are logged but don't - // fail initialization. The key will be created when the first - // OpLog entry is written - LOG(WARNING) - << "Failed to initialize /latest key (error=" << create_err - << "), will be created on first OpLog write"; - } - } else if (get_err == ErrorCode::OK) { - // Key already exists, do nothing - preserve existing value - LOG(INFO) << "/latest key already exists (value=" << existing_value - << ") for cluster_id=" << cluster_id_; - } else { - // Other errors (e.g., etcd not connected) are logged but don't fail - // initialization - LOG(WARNING) << "Failed to check /latest key existence (error=" - << get_err - << "), will be created on first OpLog write"; - } - } - - // Start batch update thread only for writers. - if (enable_latest_seq_batch_update_) { - // Prevent double start - if (!batch_update_running_.exchange(true)) { - batch_update_thread_ = - std::thread(&EtcdOpLogStore::BatchUpdateThread, this); - } - } - - // Start OpLog batch write thread only for writers. - if (enable_batch_write_) { - // Prevent double start - if (!batch_write_running_.exchange(true)) { - batch_write_thread_ = - std::thread(&EtcdOpLogStore::BatchWriteThread, this); - } - } - - return ErrorCode::OK; -} - -EtcdOpLogStore::~EtcdOpLogStore() { - // Stop OpLog batch write thread (only started for writers) - if (enable_batch_write_) { - batch_write_running_.store(false); - cv_batch_updated_.notify_all(); - if (batch_write_thread_.joinable()) { - batch_write_thread_.join(); - } - - // Attempt final flush (FlushBatch manages its own locking) - FlushBatch(); - } - - if (!enable_latest_seq_batch_update_) { - return; - } - - // Stop batch update thread - batch_update_running_.store(false); - if (batch_update_thread_.joinable()) { - batch_update_thread_.join(); - } - - // Perform final update if there are pending updates - if (pending_count_.load() > 0) { - DoBatchUpdate(); - } -} - -ErrorCode EtcdOpLogStore::WriteOpLog(const OpLogEntry& entry, bool sync) { - if (!enable_batch_write_) { - LOG(ERROR) << "WriteOpLog called on a read-only EtcdOpLogStore " - << "(enable_batch_write=false), cluster_id=" << cluster_id_; - return ErrorCode::INVALID_PARAMS; - } - std::string key = BuildOpLogKey(entry.sequence_id); - std::string value = mooncake::SerializeOpLogEntry(entry); - - { - std::unique_lock lock(batch_mutex_); - pending_batch_.push_back( - {std::move(key), std::move(value), entry.sequence_id, sync}); - - bool should_notify = false; - if (sync) { - // Strategy 2+: Sync writes (DELETE) trigger immediate flush - should_notify = true; - } else { - // Async writes (PUT_END): trigger if threshold reached - if (pending_batch_.size() >= kOpLogBatchCountLimit) { - should_notify = true; - } - } - - if (should_notify) { - cv_batch_updated_.notify_one(); - } - - if (sync) { - // Wait for persistence - uint64_t target_seq = entry.sequence_id; - bool success = cv_sync_completed_.wait_for( - lock, std::chrono::milliseconds(kSyncWaitTimeoutMs), - [&] { return last_persisted_seq_id_.load() >= target_seq; }); - if (!success) { - LOG(ERROR) << "Timeout waiting for OpLog persistence, seq=" - << target_seq; - return ErrorCode::ETCD_OPERATION_ERROR; - } - } - } - - // Update /latest pointer logic - // We defer this to the batch flush or just queue it up here? - // Original logic: - if (!enable_latest_seq_batch_update_) { - // Direct update (may be slow, but it's what config asked for) - // Warning: This is now done AFTER op log write, which is correct order. - return UpdateLatestSequenceId(entry.sequence_id); - } - - // For batch update, we update the pending counter - pending_latest_seq_id_.store(entry.sequence_id); - size_t count = pending_count_.fetch_add(1) + 1; - if (count >= kBatchSize) { - DoBatchUpdate(); - } - - return ErrorCode::OK; -} - -void EtcdOpLogStore::BatchWriteThread() { - while (batch_write_running_.load()) { - { - std::unique_lock lock(batch_mutex_); - if (pending_batch_.empty()) { - // Wait for signal or timeout (Group Commit time window) - cv_batch_updated_.wait_for( - lock, std::chrono::milliseconds(kOpLogBatchTimeoutMs)); - } - - if (!batch_write_running_.load() && pending_batch_.empty()) { - break; - } - } - - FlushBatch(); - } -} - -void EtcdOpLogStore::FlushBatch() { - // Step 1: Take pending batch under lock. - std::deque batch_to_write; - { - std::lock_guard lock(batch_mutex_); - batch_to_write.swap(pending_batch_); - } - - if (batch_to_write.empty()) { - return; - } - - // Step 2: Perform IO without holding the lock. - std::vector keys; - std::vector values; - keys.reserve(batch_to_write.size()); - values.reserve(batch_to_write.size()); - - uint64_t max_seq = 0; - bool has_sync_entry = false; - for (const auto& entry : batch_to_write) { - keys.push_back(entry.key); - if (entry.is_sync) { - has_sync_entry = true; - } - values.push_back(entry.value); - if (entry.sequence_id > max_seq) { - max_seq = entry.sequence_id; - } - } - - ErrorCode err = ErrorCode::OK; - for (int i = 0; i <= kFlushRetryCount; ++i) { - err = EtcdHelper::BatchCreate(keys, values); - if (err == ErrorCode::OK) { - break; - } - if (err == ErrorCode::ETCD_TRANSACTION_FAIL) { - // BatchCreate uses Txn(If all keys CreateRevision==0). - // Transaction failure means some keys already exist — likely - // from a previous attempt that timed out but actually succeeded - // on the etcd side. Since OpLog entries are idempotent (same - // sequence_id -> same key/value), we can safely fall back to - // individual Put (overwrite) for the remaining keys. - LOG(WARNING) - << "BatchCreate transaction failed (keys already exist), " - << "falling back to per-key Put for " << keys.size() - << " entries"; - bool all_ok = true; - for (size_t j = 0; j < keys.size(); ++j) { - ErrorCode put_err = - EtcdHelper::Put(keys[j].c_str(), keys[j].size(), - values[j].c_str(), values[j].size()); - if (put_err != ErrorCode::OK) { - LOG(ERROR) << "Fallback Put failed for key=" << keys[j]; - all_ok = false; - } - } - if (all_ok) { - err = ErrorCode::OK; - } - break; // Do not retry further; fallback already handled it. - } - if (i < kFlushRetryCount) { - LOG(WARNING) << "Failed to flush OpLog batch (attempt " << i + 1 - << "/" << kFlushRetryCount + 1 << "), retrying..."; - std::this_thread::sleep_for( - std::chrono::milliseconds(kFlushRetryIntervalMs)); - } - } - - // Step 3: Update state under lock. - { - std::lock_guard lock(batch_mutex_); - - if (err == ErrorCode::OK) { - if (max_seq > last_persisted_seq_id_.load()) { - last_persisted_seq_id_.store(max_seq); - } - - // Update HA metrics - HAMetricManager::instance().inc_oplog_batch_commits(); - if (has_sync_entry) { - HAMetricManager::instance().inc_oplog_sync_batch_commits(); - } - - if (batch_to_write.size() > 1) { - LOG(INFO) - << "HA Strategy: Group Commit flush success. batch_size=" - << batch_to_write.size() << ", max_seq=" << max_seq; - } else { - VLOG(3) - << "HA Strategy: Group Commit flush success. batch_size=1, " - "max_seq=" - << max_seq; - if (!has_sync_entry) { - LOG_EVERY_N(INFO, 1000) - << "Note: Frequent single-entry async " - "flushes detected (sample)."; - } - } - } else { - LOG(ERROR) << "Failed to flush OpLog batch, count=" - << batch_to_write.size(); - } - } - - // Wake up all waiting threads (Strategy 2+: DELETE waiters) - cv_sync_completed_.notify_all(); -} - -ErrorCode EtcdOpLogStore::ReadOpLog(uint64_t sequence_id, OpLogEntry& entry) { - std::string key = BuildOpLogKey(sequence_id); - std::string value; - EtcdRevisionId revision_id; - ErrorCode err = - EtcdHelper::Get(key.c_str(), key.size(), value, revision_id); - if (err != ErrorCode::OK) { - // Translate etcd-specific "key not found" to the generic OpLogStore - // error. - if (err == ErrorCode::ETCD_KEY_NOT_EXIST) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - return err; - } - - if (!mooncake::DeserializeOpLogEntry(value, entry)) { - LOG(ERROR) << "Failed to deserialize OpLog entry, sequence_id=" - << sequence_id; - return ErrorCode::INTERNAL_ERROR; - } - - return ErrorCode::OK; -} - -ErrorCode EtcdOpLogStore::ReadOpLogSince(uint64_t start_sequence_id, - size_t limit, - std::vector& entries) { - EtcdRevisionId rev = 0; - return ReadOpLogSinceWithRevision(start_sequence_id, limit, entries, rev); -} - -ErrorCode EtcdOpLogStore::ReadOpLogSinceWithRevision( - uint64_t start_sequence_id, size_t limit, std::vector& entries, - EtcdRevisionId& revision_id) { - entries.clear(); - entries.reserve(limit); - - // Range is limited to OpLog entry keys only. - const std::string prefix = std::string(kOpLogPrefix) + cluster_id_ + "/"; - std::string current_start_key = BuildOpLogKey(start_sequence_id + 1); - - // Compute prefix range end (etcd prefix end). - auto prefix_end = [](std::string p) -> std::string { - for (int i = static_cast(p.size()) - 1; i >= 0; --i) { - unsigned char c = static_cast(p[i]); - if (c < 0xFF) { - p[i] = static_cast(c + 1); - p.resize(i + 1); - return p; - } - } - return std::string(1, '\0'); - }; - const std::string end_key = prefix_end(prefix); - - // Pagination: - // - Use range-get with limit - // - Start next page from lastKey + '\0' (lexicographically just after - // lastKey) This avoids repeating the last key without adding new Go/C++ - // APIs. - revision_id = 0; - while (entries.size() < limit) { - const size_t page_limit = limit - entries.size(); - std::string json; - EtcdRevisionId page_rev = 0; - ErrorCode err = EtcdHelper::GetRangeAsJson( - current_start_key.c_str(), current_start_key.size(), - end_key.c_str(), end_key.size(), page_limit, json, page_rev); - if (err != ErrorCode::OK) { - return err; - } - if (page_rev > revision_id) { - revision_id = page_rev; - } - - // Parse kv list: [{"key":"...","value":"..."}] - Json::Value root; - Json::CharReaderBuilder reader; - std::string errs; - std::istringstream s(json); - if (!Json::parseFromStream(reader, s, &root, &errs)) { - LOG(ERROR) << "Failed to parse range JSON: " << errs; - return ErrorCode::INTERNAL_ERROR; - } - if (!root.isArray()) { - return ErrorCode::INTERNAL_ERROR; - } - if (root.empty()) { - break; // no more data - } - - std::string last_key_in_page; - for (const auto& kv : root) { - const std::string key = kv.get("key", "").asString(); - last_key_in_page = key; - if (key.empty() || key.find("/latest") != std::string::npos || - key.find("/snapshot/") != std::string::npos) { - continue; - } - - // Parse seq from key suffix and filter (handles legacy keys too). - size_t pos = key.rfind('/'); - if (pos == std::string::npos || pos + 1 >= key.size()) { - continue; - } - uint64_t seq = 0; - try { - seq = static_cast(std::stoull(key.substr(pos + 1))); - } catch (...) { - continue; - } - if (IsSequenceOlderOrEqual(seq, start_sequence_id)) { - continue; - } - - OpLogEntry entry; - const std::string value = kv.get("value", "").asString(); - if (!mooncake::DeserializeOpLogEntry(value, entry)) { - LOG(ERROR) << "Failed to deserialize OpLog entry from key=" - << key; - return ErrorCode::INTERNAL_ERROR; - } - entries.push_back(std::move(entry)); - if (entries.size() >= limit) { - break; - } - } - - // Advance start key for next page. - if (last_key_in_page.empty()) { - break; - } - current_start_key = last_key_in_page; - current_start_key.push_back('\0'); - } - - return ErrorCode::OK; -} - -ErrorCode EtcdOpLogStore::GetLatestSequenceId(uint64_t& sequence_id) { - std::string key = BuildLatestKey(); - std::string value; - EtcdRevisionId revision_id; - ErrorCode err = - EtcdHelper::Get(key.c_str(), key.size(), value, revision_id); - if (err != ErrorCode::OK) { - if (err == ErrorCode::ETCD_KEY_NOT_EXIST) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - return err; - } - - try { - sequence_id = std::stoull(value); - } catch (const std::exception& e) { - LOG(ERROR) << "Failed to parse latest sequence_id: " << e.what(); - return ErrorCode::INTERNAL_ERROR; - } - - return ErrorCode::OK; -} - -ErrorCode EtcdOpLogStore::GetMaxSequenceId(uint64_t& sequence_id) { - auto max_seq_opt = GetMaxSequenceIdInternal(); - if (!max_seq_opt.has_value()) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - sequence_id = max_seq_opt.value(); - return ErrorCode::OK; -} - -ErrorCode EtcdOpLogStore::UpdateLatestSequenceId(uint64_t sequence_id) { - std::string key = BuildLatestKey(); - std::string value = std::to_string(sequence_id); - return EtcdHelper::Put(key.c_str(), key.size(), value.c_str(), - value.size()); -} - -ErrorCode EtcdOpLogStore::RecordSnapshotSequenceId( - const std::string& snapshot_id, uint64_t sequence_id) { - std::string key = BuildSnapshotKey(snapshot_id); - std::string value = std::to_string(sequence_id); - return EtcdHelper::Put(key.c_str(), key.size(), value.c_str(), - value.size()); -} - -ErrorCode EtcdOpLogStore::GetSnapshotSequenceId(const std::string& snapshot_id, - uint64_t& sequence_id) { - std::string key = BuildSnapshotKey(snapshot_id); - std::string value; - EtcdRevisionId revision_id; - ErrorCode err = - EtcdHelper::Get(key.c_str(), key.size(), value, revision_id); - if (err != ErrorCode::OK) { - if (err == ErrorCode::ETCD_KEY_NOT_EXIST) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - return err; - } - - try { - sequence_id = std::stoull(value); - } catch (const std::exception& e) { - LOG(ERROR) << "Failed to parse snapshot sequence_id: " << e.what(); - return ErrorCode::INTERNAL_ERROR; - } - - return ErrorCode::OK; -} - -ErrorCode EtcdOpLogStore::CleanupOpLogBefore(uint64_t before_sequence_id) { - // Robust cleanup (Scheme 3): - // - Determine current minimum sequence_id in etcd - // - DeleteRange [min_key, before_key) - // - // IMPORTANT: This relies on lexicographical ordering of keys, so the - // sequence_id portion MUST be fixed-width (zero-padded). - auto min_seq_opt = GetMinSequenceId(); - if (!min_seq_opt.has_value()) { - return ErrorCode::OK; // nothing to cleanup - } - - uint64_t min_seq = min_seq_opt.value(); - if (before_sequence_id <= min_seq) { - return ErrorCode::OK; - } - - std::string start_key = BuildOpLogKey(min_seq); - std::string end_key = - BuildOpLogKey(before_sequence_id); // delete < before_sequence_id - - return EtcdHelper::DeleteRange(start_key.c_str(), start_key.size(), - end_key.c_str(), end_key.size()); -} - -std::string EtcdOpLogStore::BuildOpLogKey(uint64_t sequence_id) const { - std::ostringstream oss; - // Fixed-width encoding for correct etcd lexicographical range operations. - // 20 digits is enough for uint64_t max (18446744073709551615). - oss << kOpLogPrefix << cluster_id_ << "/" << std::setw(20) - << std::setfill('0') << sequence_id; - return oss.str(); -} - -std::optional EtcdOpLogStore::GetMinSequenceId() const { - std::string prefix = std::string(kOpLogPrefix) + cluster_id_ + "/"; - std::string first_key; - ErrorCode err = EtcdHelper::GetFirstKeyWithPrefix(prefix.c_str(), - prefix.size(), first_key); - if (err != ErrorCode::OK) { - return std::nullopt; - } - - // Skip non-entry keys if any (e.g. "/latest" or "/snapshot/..."). - // Entries are expected to be ".../<20-digit-seq>". - // If the first key isn't an entry key, fall back to nullopt (safe no-op). - if (first_key.find("/latest") != std::string::npos || - first_key.find("/snapshot/") != std::string::npos) { - return std::nullopt; - } - - size_t pos = first_key.rfind('/'); - if (pos == std::string::npos || pos + 1 >= first_key.size()) { - return std::nullopt; - } - std::string seq_str = first_key.substr(pos + 1); - try { - return static_cast(std::stoull(seq_str)); - } catch (...) { - return std::nullopt; - } -} - -std::optional EtcdOpLogStore::GetMaxSequenceIdInternal() const { - // Entry keys are fixed-width 20-digit numbers, which (in practice) start - // with '0'. Use "/0" to avoid picking up "/latest" which is - // lexicographically after digits. - std::string prefix = std::string(kOpLogPrefix) + cluster_id_ + "/0"; - std::string last_key; - ErrorCode err = EtcdHelper::GetLastKeyWithPrefix(prefix.c_str(), - prefix.size(), last_key); - if (err != ErrorCode::OK) { - return std::nullopt; - } - - size_t pos = last_key.rfind('/'); - if (pos == std::string::npos || pos + 1 >= last_key.size()) { - return std::nullopt; - } - std::string seq_str = last_key.substr(pos + 1); - try { - return static_cast(std::stoull(seq_str)); - } catch (...) { - return std::nullopt; - } -} - -std::string EtcdOpLogStore::BuildLatestKey() const { - std::ostringstream oss; - oss << kOpLogPrefix << cluster_id_ << kLatestSuffix; - return oss.str(); -} - -std::string EtcdOpLogStore::BuildSnapshotKey( - const std::string& snapshot_id) const { - std::ostringstream oss; - oss << kOpLogPrefix << cluster_id_ << kSnapshotSuffix << snapshot_id - << "/sequence_id"; - return oss.str(); -} - -std::unique_ptr EtcdOpLogStore::CreateChangeNotifier( - const std::string& cluster_id) { - return std::make_unique(cluster_id, this); -} - -void EtcdOpLogStore::BatchUpdateThread() { - if (!enable_latest_seq_batch_update_) { - return; - } - while (batch_update_running_.load()) { - std::this_thread::sleep_for( - std::chrono::milliseconds(kBatchIntervalMs)); - - // Check if we need to update based on time interval - auto now = std::chrono::steady_clock::now(); - auto elapsed = std::chrono::duration_cast( - now - last_update_time_) - .count(); - - if (pending_count_.load() > 0 && elapsed >= kBatchIntervalMs) { - DoBatchUpdate(); - } - } -} - -void EtcdOpLogStore::TriggerBatchUpdateIfNeeded() { - // This method is kept for potential future use (e.g., manual trigger) - // Currently, DoBatchUpdate() is called directly from WriteOpLog - // when batch size threshold is reached - if (pending_count_.load() >= kBatchSize) { - DoBatchUpdate(); - } -} - -void EtcdOpLogStore::DoBatchUpdate() { - if (!enable_latest_seq_batch_update_) { - return; - } - std::lock_guard lock(batch_update_mutex_); - - // Get the pending sequence_id and reset counters - uint64_t seq_id_to_update = pending_latest_seq_id_.load(); - size_t count = pending_count_.exchange(0); - - if (count == 0) { - return; // Nothing to update - } - - // Update latest_sequence_id in etcd - ErrorCode err = UpdateLatestSequenceId(seq_id_to_update); - if (err != ErrorCode::OK) { - LOG(WARNING) << "Failed to batch update latest_sequence_id=" - << seq_id_to_update << ", error=" << err - << ". Will retry in next batch."; - // Restore the count so it will be retried - pending_count_.fetch_add(count); - } else { - last_update_time_ = std::chrono::steady_clock::now(); - VLOG(2) << "Batch updated latest_sequence_id=" << seq_id_to_update - << " (count=" << count << " entries)"; - } -} - -} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/localfs_oplog_store.cpp b/mooncake-store/src/ha/oplog/localfs_oplog_store.cpp deleted file mode 100644 index 26bef9ac3a..0000000000 --- a/mooncake-store/src/ha/oplog/localfs_oplog_store.cpp +++ /dev/null @@ -1,788 +0,0 @@ -#include "ha/oplog/localfs_oplog_store.h" - -#include - -#include -#include -#include -#include - -#include -#include - -#include "ha/oplog/oplog_serializer.h" -#include "ha/oplog/polling_oplog_change_notifier.h" - -namespace fs = std::filesystem; - -namespace mooncake { - -// ============================================================ -// Constructor / Destructor -// ============================================================ - -LocalFsOpLogStore::LocalFsOpLogStore(const std::string& cluster_id, - const std::string& root_dir, - bool enable_batch_write, - int poll_interval_ms) - : cluster_id_(cluster_id), - root_dir_(root_dir), - cluster_dir_(root_dir + "/" + cluster_id), - enable_batch_write_(enable_batch_write), - poll_interval_ms_(poll_interval_ms) { - if (!NormalizeAndValidateClusterId(cluster_id_)) { - LOG(FATAL) << "Invalid cluster_id for LocalFsOpLogStore: '" - << cluster_id - << "'. Allowed chars: [A-Za-z0-9_.-], max_len=128."; - } - // Rebuild cluster_dir_ after normalization. - cluster_dir_ = root_dir_ + "/" + cluster_id_; -} - -LocalFsOpLogStore::~LocalFsOpLogStore() { - if (batch_write_running_.load()) { - batch_write_running_.store(false); - cv_batch_updated_.notify_all(); - if (batch_write_thread_.joinable()) { - batch_write_thread_.join(); - } - } - // Final flush of any remaining entries - if (!pending_batch_.empty()) { - FlushBatch(); - } - // Notify any sync waiters so they don't block forever - cv_sync_completed_.notify_all(); -} - -// ============================================================ -// Init -// ============================================================ - -ErrorCode LocalFsOpLogStore::Init() { - try { - fs::create_directories(SegmentsDir()); - } catch (const fs::filesystem_error& e) { - LOG(ERROR) << "LocalFsOpLogStore::Init: failed to create directories: " - << e.what(); - return ErrorCode::INTERNAL_ERROR; - } - - CleanupTempFiles(); - - if (enable_batch_write_) { - // Writer: initialize latest file if not exists - std::string latest_path = LatestFilePath(); - if (!fs::exists(latest_path)) { - auto err = AtomicWriteFile(latest_path, "0"); - if (err != ErrorCode::OK) { - LOG(ERROR) << "LocalFsOpLogStore::Init: failed to create " - "latest file"; - return err; - } - } - - auto recover_err = RecoverPersistedState(); - if (recover_err != ErrorCode::OK) { - LOG(ERROR) << "LocalFsOpLogStore::Init: failed to recover " - "persisted state"; - return recover_err; - } - - // Start batch write thread - batch_write_running_.store(true); - batch_write_thread_ = - std::thread(&LocalFsOpLogStore::BatchWriteThread, this); - } - - return ErrorCode::OK; -} - -// ============================================================ -// Path helpers -// ============================================================ - -std::string LocalFsOpLogStore::SegmentsDir() const { - return cluster_dir_ + "/segments"; -} - -std::string LocalFsOpLogStore::SnapshotsDir() const { - return cluster_dir_ + "/snapshots"; -} - -std::string LocalFsOpLogStore::LatestFilePath() const { - return cluster_dir_ + "/latest"; -} - -std::string LocalFsOpLogStore::BuildSegmentFilename(uint64_t min_seq, - uint64_t max_seq) const { - char buf[128]; - snprintf(buf, sizeof(buf), "seg_%020lu_%020lu", min_seq, max_seq); - return SegmentsDir() + "/" + buf; -} - -std::string LocalFsOpLogStore::BuildSnapshotPath( - const std::string& snapshot_id) const { - return SnapshotsDir() + "/" + snapshot_id; -} - -// ============================================================ -// Atomic file write -// ============================================================ - -ErrorCode LocalFsOpLogStore::AtomicWriteFile(const std::string& target_path, - const std::string& content) { - return AtomicWriteFile(target_path, content.data(), content.size()); -} - -ErrorCode LocalFsOpLogStore::AtomicWriteFile(const std::string& target_path, - const void* data, size_t size) { - std::string tmp_path = target_path + ".tmp"; - int fd = ::open(tmp_path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644); - if (fd < 0) { - LOG(ERROR) << "AtomicWriteFile: open failed: " << tmp_path - << ", errno=" << errno; - return ErrorCode::INTERNAL_ERROR; - } - - const char* ptr = static_cast(data); - size_t remaining = size; - while (remaining > 0) { - ssize_t written = ::write(fd, ptr, remaining); - if (written < 0) { - if (errno == EINTR) continue; - LOG(ERROR) << "AtomicWriteFile: write failed: " << tmp_path - << ", errno=" << errno; - ::close(fd); - ::unlink(tmp_path.c_str()); - return ErrorCode::INTERNAL_ERROR; - } - ptr += written; - remaining -= static_cast(written); - } - - if (::fsync(fd) != 0) { - LOG(ERROR) << "AtomicWriteFile: fsync failed: " << tmp_path - << ", errno=" << errno; - ::close(fd); - ::unlink(tmp_path.c_str()); - return ErrorCode::INTERNAL_ERROR; - } - ::close(fd); - - if (::rename(tmp_path.c_str(), target_path.c_str()) != 0) { - LOG(ERROR) << "AtomicWriteFile: rename failed: " << tmp_path << " -> " - << target_path << ", errno=" << errno; - ::unlink(tmp_path.c_str()); - return ErrorCode::INTERNAL_ERROR; - } - - // fsync parent directory to ensure the new directory entry is durable. - // Critical for DFS mounts where rename visibility is not guaranteed - // without an explicit directory sync. - std::string parent_dir = fs::path(target_path).parent_path().string(); - int dir_fd = ::open(parent_dir.c_str(), O_RDONLY | O_DIRECTORY); - if (dir_fd >= 0) { - ::fsync(dir_fd); - ::close(dir_fd); - } - - return ErrorCode::OK; -} - -// ============================================================ -// Cleanup temp files -// ============================================================ - -void LocalFsOpLogStore::CleanupTempFiles() { - try { - for (auto& dir : {SegmentsDir(), cluster_dir_}) { - if (!fs::exists(dir)) continue; - for (auto& entry : fs::directory_iterator(dir)) { - if (entry.path().extension() == ".tmp") { - fs::remove(entry.path()); - LOG(INFO) << "Cleaned up temp file: " << entry.path(); - } - } - } - } catch (const fs::filesystem_error& e) { - LOG(WARNING) << "CleanupTempFiles: " << e.what(); - } -} - -// ============================================================ -// File read helper -// ============================================================ - -ErrorCode LocalFsOpLogStore::ReadUint64FromFile(const std::string& filepath, - uint64_t& value) const { - std::ifstream f(filepath); - if (!f) { - return ErrorCode::INTERNAL_ERROR; - } - std::string content; - f >> content; - try { - value = std::stoull(content); - } catch (...) { - LOG(ERROR) << "ReadUint64FromFile: invalid content in " << filepath - << ": " << content; - return ErrorCode::INTERNAL_ERROR; - } - return ErrorCode::OK; -} - -// ============================================================ -// Snapshot ID validation -// ============================================================ - -bool LocalFsOpLogStore::ValidateSnapshotId(const std::string& snapshot_id) { - if (snapshot_id.empty()) return false; - if (snapshot_id.find('/') != std::string::npos) return false; - if (snapshot_id.find("..") != std::string::npos) return false; - if (snapshot_id.find('\0') != std::string::npos) return false; - return true; -} - -// ============================================================ -// Segment filename parsing -// ============================================================ - -bool LocalFsOpLogStore::ParseSegmentFilename(const std::string& filename, - uint64_t& min_seq, - uint64_t& max_seq) { - // Expected format: seg_XXXXXXXXXXXXXXXXXXXX_XXXXXXXXXXXXXXXXXXXX - if (filename.size() < 45) return false; - if (filename.substr(0, 4) != "seg_") return false; - try { - min_seq = std::stoull(filename.substr(4, 20)); - max_seq = std::stoull(filename.substr(25, 20)); - return true; - } catch (...) { - return false; - } -} - -std::vector LocalFsOpLogStore::ListSegments() - const { - std::vector segments; - std::string seg_dir = SegmentsDir(); - if (!fs::exists(seg_dir)) return segments; - - try { - for (auto& entry : fs::directory_iterator(seg_dir)) { - if (entry.path().extension() == ".tmp") continue; - std::string filename = entry.path().filename().string(); - uint64_t min_seq = 0, max_seq = 0; - if (ParseSegmentFilename(filename, min_seq, max_seq)) { - segments.push_back({filename, min_seq, max_seq}); - } - } - } catch (const fs::filesystem_error& e) { - LOG(WARNING) << "ListSegments: " << e.what(); - } - - // Sort by min_seq - std::sort(segments.begin(), segments.end(), - [](const SegmentInfo& a, const SegmentInfo& b) { - return a.min_seq < b.min_seq; - }); - return segments; -} - -ErrorCode LocalFsOpLogStore::NormalizeBatchEntries( - std::vector& entries) const { - if (entries.empty()) { - return ErrorCode::OK; - } - - std::sort(entries.begin(), entries.end(), - [](const BatchEntry& a, const BatchEntry& b) { - return a.sequence_id < b.sequence_id; - }); - - std::vector normalized; - normalized.reserve(entries.size()); - for (auto& entry : entries) { - if (normalized.empty() || - normalized.back().sequence_id != entry.sequence_id) { - normalized.push_back(std::move(entry)); - continue; - } - - if (normalized.back().serialized_value != entry.serialized_value) { - LOG(ERROR) << "NormalizeBatchEntries: fencing violation for seq=" - << entry.sequence_id << " inside pending batch"; - return ErrorCode::INTERNAL_ERROR; - } - - normalized.back().is_sync = normalized.back().is_sync || entry.is_sync; - } - - entries = std::move(normalized); - return ErrorCode::OK; -} - -ErrorCode LocalFsOpLogStore::VerifyPersistedEntryMatches( - uint64_t sequence_id, const std::string& serialized_value) { - OpLogEntry existing_entry; - ErrorCode err = ReadOpLog(sequence_id, existing_entry); - if (err != ErrorCode::OK) { - if (err != ErrorCode::OPLOG_ENTRY_NOT_FOUND) { - LOG(ERROR) << "VerifyPersistedEntryMatches: seq=" << sequence_id - << " marked persisted but read failed, err=" - << static_cast(err); - } - return err; - } - - if (SerializeOpLogEntry(existing_entry) != serialized_value) { - LOG(ERROR) << "VerifyPersistedEntryMatches: fencing violation for " - << "persisted seq=" << sequence_id; - return ErrorCode::INTERNAL_ERROR; - } - - return ErrorCode::OK; -} - -ErrorCode LocalFsOpLogStore::RecoverPersistedState() { - uint64_t max_seq = 0; - ErrorCode err = GetMaxSequenceId(max_seq); - if (err == ErrorCode::OPLOG_ENTRY_NOT_FOUND) { - last_persisted_seq_id_.store(0); - return UpdateLatestSequenceId(0); - } - if (err != ErrorCode::OK) { - LOG(ERROR) << "RecoverPersistedState: failed to inspect segments, err=" - << static_cast(err); - return err; - } - - last_persisted_seq_id_.store(max_seq); - - uint64_t latest_seq = 0; - ErrorCode latest_err = GetLatestSequenceId(latest_seq); - if (latest_err != ErrorCode::OK || latest_seq != max_seq) { - ErrorCode update_err = UpdateLatestSequenceId(max_seq); - if (update_err != ErrorCode::OK) { - LOG(ERROR) << "RecoverPersistedState: failed to refresh latest " - "file to seq=" - << max_seq; - return update_err; - } - } - - return ErrorCode::OK; -} - -// ============================================================ -// Segment file I/O -// ============================================================ - -ErrorCode LocalFsOpLogStore::WriteSegmentFile( - const std::vector& entries) { - if (entries.empty()) return ErrorCode::OK; - - uint64_t min_seq = entries.front().sequence_id; - uint64_t max_seq = entries.back().sequence_id; - - // Build file content: header + length-prefixed entries - std::string content; - - // Reserve approximate size - size_t estimated_size = kSegmentHeaderSize; - for (const auto& e : entries) { - estimated_size += sizeof(uint32_t) + e.serialized_value.size(); - } - content.reserve(estimated_size); - - // Header - SegmentHeader header; - std::memcpy(header.magic, kSegmentMagic, 4); - header.version = kSegmentVersion; - header.min_seq = min_seq; - header.max_seq = max_seq; - header.entry_count = static_cast(entries.size()); - header.reserved = 0; - - content.append(reinterpret_cast(&header), sizeof(header)); - - // Entries: [uint32_t length][serialized_value] - for (const auto& e : entries) { - uint32_t len = static_cast(e.serialized_value.size()); - content.append(reinterpret_cast(&len), sizeof(len)); - content.append(e.serialized_value); - } - - std::string filepath = BuildSegmentFilename(min_seq, max_seq); - return AtomicWriteFile(filepath, content.data(), content.size()); -} - -ErrorCode LocalFsOpLogStore::ReadSegmentHeader(const std::string& filepath, - SegmentHeader& header) { - std::ifstream f(filepath, std::ios::binary); - if (!f) { - LOG(ERROR) << "ReadSegmentHeader: cannot open " << filepath; - return ErrorCode::INTERNAL_ERROR; - } - - f.read(reinterpret_cast(&header), sizeof(header)); - if (!f || static_cast(f.gcount()) < sizeof(header)) { - LOG(ERROR) << "ReadSegmentHeader: truncated header in " << filepath; - return ErrorCode::INTERNAL_ERROR; - } - - if (std::memcmp(header.magic, kSegmentMagic, 4) != 0) { - LOG(ERROR) << "ReadSegmentHeader: bad magic in " << filepath; - return ErrorCode::INTERNAL_ERROR; - } - - if (header.version != kSegmentVersion) { - LOG(ERROR) << "ReadSegmentHeader: unsupported version " - << header.version << " in " << filepath; - return ErrorCode::INTERNAL_ERROR; - } - - return ErrorCode::OK; -} - -ErrorCode LocalFsOpLogStore::ReadSegmentEntries( - const std::string& filepath, std::vector& entries) { - SegmentHeader header; - auto err = ReadSegmentHeader(filepath, header); - if (err != ErrorCode::OK) return err; - - // Re-open and skip past header to read entries - std::ifstream f(filepath, std::ios::binary); - if (!f) { - LOG(ERROR) << "ReadSegmentEntries: cannot open " << filepath; - return ErrorCode::INTERNAL_ERROR; - } - f.seekg(sizeof(SegmentHeader)); - - for (uint32_t i = 0; i < header.entry_count; ++i) { - uint32_t len = 0; - f.read(reinterpret_cast(&len), sizeof(len)); - if (!f) { - LOG(ERROR) << "ReadSegmentEntries: truncated entry length at " - << "entry " << i << " in " << filepath; - return ErrorCode::INTERNAL_ERROR; - } - - std::string buf(len, '\0'); - f.read(buf.data(), len); - if (!f) { - LOG(ERROR) << "ReadSegmentEntries: truncated entry data at " - << "entry " << i << " in " << filepath; - return ErrorCode::INTERNAL_ERROR; - } - - OpLogEntry entry; - if (!DeserializeOpLogEntry(buf, entry)) { - LOG(ERROR) << "ReadSegmentEntries: failed to deserialize entry " - << i << " in " << filepath; - return ErrorCode::INTERNAL_ERROR; - } - entries.push_back(std::move(entry)); - } - - return ErrorCode::OK; -} - -// ============================================================ -// Write path (Group Commit) -// ============================================================ - -ErrorCode LocalFsOpLogStore::WriteOpLog(const OpLogEntry& entry, bool sync) { - if (!enable_batch_write_) { - LOG(ERROR) << "WriteOpLog called on READER instance"; - return ErrorCode::INVALID_PARAMS; - } - - if (!batch_write_running_.load()) { - LOG(ERROR) << "WriteOpLog called after shutdown"; - return ErrorCode::INTERNAL_ERROR; - } - - // Serialize entry - std::string serialized = SerializeOpLogEntry(entry); - - // Idempotent retry / fencing check for already-persisted sequence IDs. - // Note: last_persisted_seq_id_ is only a high watermark. With pre- - // allocated sequence IDs, larger seq may flush before a smaller retry. - // So a cache hit here must still verify the concrete persisted record. - if (last_persisted_seq_id_.load() >= entry.sequence_id) { - ErrorCode verify_err = - VerifyPersistedEntryMatches(entry.sequence_id, serialized); - if (verify_err == ErrorCode::OK) { - return ErrorCode::OK; - } - if (verify_err != ErrorCode::OPLOG_ENTRY_NOT_FOUND) { - return verify_err; - } - } - - { - std::lock_guard lock(batch_mutex_); - pending_batch_.push_back({serialized, entry.sequence_id, sync}); - } - - // Always notify — spurious wakeup is cheap, data race is not - cv_batch_updated_.notify_one(); - - if (sync) { - // Wait for the specific sequence to become durably readable. - std::unique_lock lock(batch_mutex_); - bool flushed = cv_sync_completed_.wait_for( - lock, std::chrono::milliseconds(kSyncWaitTimeoutMs), [&] { - if (last_persisted_seq_id_.load() < entry.sequence_id) { - return false; - } - return VerifyPersistedEntryMatches(entry.sequence_id, - serialized) == ErrorCode::OK; - }); - if (!flushed) { - LOG(ERROR) << "WriteOpLog sync wait timed out for seq=" - << entry.sequence_id; - return ErrorCode::INTERNAL_ERROR; - } - } - - return ErrorCode::OK; -} - -void LocalFsOpLogStore::BatchWriteThread() { - while (batch_write_running_.load()) { - { - std::unique_lock lock(batch_mutex_); - if (pending_batch_.empty()) { - cv_batch_updated_.wait_for( - lock, std::chrono::milliseconds(kBatchTimeoutMs)); - } - if (!batch_write_running_.load() && pending_batch_.empty()) { - break; - } - } - FlushBatch(); - } -} - -void LocalFsOpLogStore::FlushBatch() { - std::deque batch_to_write; - { - std::lock_guard lock(batch_mutex_); - if (pending_batch_.empty()) return; - batch_to_write.swap(pending_batch_); - } - - std::vector entries( - std::make_move_iterator(batch_to_write.begin()), - std::make_move_iterator(batch_to_write.end())); - ErrorCode normalize_err = NormalizeBatchEntries(entries); - if (normalize_err != ErrorCode::OK) { - LOG(ERROR) << "FlushBatch: invalid duplicate sequence detected in " - "pending batch"; - cv_sync_completed_.notify_all(); - return; - } - - // Retry on failure - ErrorCode err = ErrorCode::INTERNAL_ERROR; - for (int retry = 0; retry < kFlushRetryCount; ++retry) { - err = WriteSegmentFile(entries); - if (err == ErrorCode::OK) break; - LOG(WARNING) << "FlushBatch: WriteSegmentFile failed, retry " - << (retry + 1) << "/" << kFlushRetryCount; - if (retry + 1 < kFlushRetryCount) { - std::this_thread::sleep_for( - std::chrono::milliseconds(kFlushRetryIntervalMs)); - } - } - - if (err != ErrorCode::OK) { - LOG(ERROR) << "FlushBatch: all retries failed, " << entries.size() - << " entries lost (seq " << entries.front().sequence_id - << " - " << entries.back().sequence_id << ")"; - // Notify sync waiters so they don't block forever - cv_sync_completed_.notify_all(); - return; - } - - // Update last_persisted_seq_id_ - uint64_t max_seq = entries.back().sequence_id; - last_persisted_seq_id_.store(max_seq); - - // Update latest file - AtomicWriteFile(LatestFilePath(), std::to_string(max_seq)); - - // Notify sync waiters - cv_sync_completed_.notify_all(); -} - -// ============================================================ -// Read path -// ============================================================ - -ErrorCode LocalFsOpLogStore::ReadOpLog(uint64_t sequence_id, - OpLogEntry& entry) { - auto segments = ListSegments(); - - // Binary search for the segment containing sequence_id - auto it = std::lower_bound( - segments.begin(), segments.end(), sequence_id, - [](const SegmentInfo& seg, uint64_t seq) { return seg.max_seq < seq; }); - - if (it == segments.end()) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - - // Check if sequence_id is within this segment's range - if (sequence_id < it->min_seq || sequence_id > it->max_seq) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - - std::string filepath = SegmentsDir() + "/" + it->filename; - std::vector entries; - auto err = ReadSegmentEntries(filepath, entries); - if (err != ErrorCode::OK) return err; - - for (auto& e : entries) { - if (e.sequence_id == sequence_id) { - entry = std::move(e); - return ErrorCode::OK; - } - } - - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; -} - -ErrorCode LocalFsOpLogStore::ReadOpLogSince(uint64_t start_sequence_id, - size_t limit, - std::vector& entries) { - entries.clear(); - auto segments = ListSegments(); - - // Find first segment that could contain entries after start_sequence_id - auto it = - std::lower_bound(segments.begin(), segments.end(), start_sequence_id, - [](const SegmentInfo& seg, uint64_t seq) { - return seg.max_seq <= seq; - }); - - for (; it != segments.end() && entries.size() < limit; ++it) { - std::string filepath = SegmentsDir() + "/" + it->filename; - std::vector seg_entries; - auto err = ReadSegmentEntries(filepath, seg_entries); - if (err != ErrorCode::OK) { - LOG(ERROR) << "ReadOpLogSince: corrupted segment " << it->filename - << ", aborting read to prevent oplog gaps"; - return err; - } - - for (auto& e : seg_entries) { - if (e.sequence_id > start_sequence_id) { - entries.push_back(std::move(e)); - if (entries.size() >= limit) break; - } - } - } - - return ErrorCode::OK; -} - -// ============================================================ -// Sequence ID management -// ============================================================ - -ErrorCode LocalFsOpLogStore::GetLatestSequenceId(uint64_t& sequence_id) { - std::string latest_path = LatestFilePath(); - auto err = ReadUint64FromFile(latest_path, sequence_id); - if (err != ErrorCode::OK) { - // File doesn't exist yet — default to 0 - sequence_id = 0; - return ErrorCode::OK; - } - return ErrorCode::OK; -} - -ErrorCode LocalFsOpLogStore::GetMaxSequenceId(uint64_t& sequence_id) { - auto segments = ListSegments(); - if (segments.empty()) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - - sequence_id = segments.back().max_seq; - return ErrorCode::OK; -} - -ErrorCode LocalFsOpLogStore::UpdateLatestSequenceId(uint64_t sequence_id) { - return AtomicWriteFile(LatestFilePath(), std::to_string(sequence_id)); -} - -// ============================================================ -// Snapshot operations -// ============================================================ - -ErrorCode LocalFsOpLogStore::RecordSnapshotSequenceId( - const std::string& snapshot_id, uint64_t sequence_id) { - if (!ValidateSnapshotId(snapshot_id)) { - return ErrorCode::INVALID_PARAMS; - } - // Create snapshots directory lazily on first write - try { - fs::create_directories(SnapshotsDir()); - } catch (const fs::filesystem_error& e) { - LOG(ERROR) << "RecordSnapshotSequenceId: failed to create dir: " - << e.what(); - return ErrorCode::INTERNAL_ERROR; - } - return AtomicWriteFile(BuildSnapshotPath(snapshot_id), - std::to_string(sequence_id)); -} - -ErrorCode LocalFsOpLogStore::GetSnapshotSequenceId( - const std::string& snapshot_id, uint64_t& sequence_id) { - if (!ValidateSnapshotId(snapshot_id)) { - return ErrorCode::INVALID_PARAMS; - } - - std::string path = BuildSnapshotPath(snapshot_id); - auto err = ReadUint64FromFile(path, sequence_id); - if (err != ErrorCode::OK) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - - return ErrorCode::OK; -} - -// ============================================================ -// Cleanup -// ============================================================ - -ErrorCode LocalFsOpLogStore::CleanupOpLogBefore(uint64_t before_sequence_id) { - auto segments = ListSegments(); - for (const auto& seg : segments) { - if (seg.max_seq < before_sequence_id) { - std::string filepath = SegmentsDir() + "/" + seg.filename; - try { - fs::remove(filepath); - } catch (const fs::filesystem_error& e) { - LOG(WARNING) << "CleanupOpLogBefore: failed to remove " - << filepath << ": " << e.what(); - } - } - } - return ErrorCode::OK; -} - -// ============================================================ -// Change notifier -// ============================================================ - -std::unique_ptr LocalFsOpLogStore::CreateChangeNotifier( - const std::string& /*cluster_id*/) { - return std::make_unique(this, - poll_interval_ms_); -} - -} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_applier.cpp b/mooncake-store/src/ha/oplog/oplog_applier.cpp index 205e84f1d3..b715b30d06 100644 --- a/mooncake-store/src/ha/oplog/oplog_applier.cpp +++ b/mooncake-store/src/ha/oplog/oplog_applier.cpp @@ -2,22 +2,16 @@ #include -#include -#include - -#include "ha/oplog/oplog_store.h" #include "ha_metric_manager.h" #include "metadata_store.h" -#include "ha/oplog/oplog_manager.h" +#include "ha/oplog/oplog_types.h" namespace mooncake { OpLogApplier::OpLogApplier(MetadataStore* metadata_store, - const std::string& cluster_id, - OpLogStore* oplog_store) + const std::string& cluster_id) : metadata_store_(metadata_store), cluster_id_(cluster_id), - oplog_store_(oplog_store), expected_sequence_id_(1) { if (metadata_store_ == nullptr) { LOG(FATAL) << "OpLogApplier: metadata_store cannot be null"; @@ -31,7 +25,7 @@ OpLogApplier::OpLogApplier(MetadataStore* metadata_store, bool OpLogApplier::ApplyOpLogEntry(const OpLogEntry& entry) { // Basic DoS protection: validate key/payload sizes before parsing/applying. std::string size_reason; - if (!OpLogManager::ValidateEntrySize(entry, &size_reason)) { + if (!ValidateOpLogEntrySize(entry, &size_reason)) { LOG(ERROR) << "OpLogApplier: entry size rejected, sequence_id=" << entry.sequence_id << ", key=" << entry.object_key << ", reason=" << size_reason; @@ -39,7 +33,7 @@ bool OpLogApplier::ApplyOpLogEntry(const OpLogEntry& entry) { } // Verify checksum to detect data corruption or tampering. - if (!OpLogManager::VerifyChecksum(entry)) { + if (!VerifyOpLogChecksum(entry)) { LOG(ERROR) << "OpLogApplier: checksum mismatch, sequence_id=" << entry.sequence_id << ", key=" << entry.object_key @@ -50,70 +44,18 @@ bool OpLogApplier::ApplyOpLogEntry(const OpLogEntry& entry) { // Global ordering only. // - // IMPORTANT: - // - Watch callbacks / retries may deliver duplicate or already-applied - // entries. - // - Those must be treated as no-op, not as "out-of-order pending", - // otherwise - // pending_entries_ can grow and the applier may appear stuck. + // Retries may deliver duplicate or already-applied entries. const uint64_t expected = expected_sequence_id_.load(); if (IsSequenceOlder(entry.sequence_id, expected)) { - // Late arrival of a previously-skipped gap entry: apply only if it's a - // delete/revoke. - bool was_skipped = false; - { - std::lock_guard lock(pending_mutex_); - auto it = skipped_sequence_ids_.find(entry.sequence_id); - if (it != skipped_sequence_ids_.end()) { - was_skipped = true; - skipped_sequence_ids_.erase(it); - } - } - if (was_skipped) { - if (entry.op_type == OpType::REMOVE || - entry.op_type == OpType::PUT_REVOKE) { - // Safe: ensure we don't keep stale metadata. - if (entry.op_type == OpType::REMOVE) { - ApplyRemove(entry); - } else { - ApplyPutRevoke(entry); - } - return true; - } - // PUT_END (or others): discard to avoid resurrecting stale state. - if (entry.op_type == OpType::PUT_END) { - HAMetricManager::instance().inc_oplog_dropped_put_end(); - } - VLOG(1) << "OpLogApplier: discard late skipped entry, op_type=" - << static_cast(entry.op_type) - << ", sequence_id=" << entry.sequence_id - << ", key=" << entry.object_key; - return true; - } - VLOG(2) << "OpLogApplier: skip already-applied entry, sequence_id=" << entry.sequence_id << ", expected=" << expected << ", key=" << entry.object_key; - return true; // consumed (no-op) + return true; } if (IsSequenceNewer(entry.sequence_id, expected)) { - // Future entry - store into pending, wait for the gap to be filled. - std::lock_guard lock(pending_mutex_); - - if (pending_entries_.size() >= - static_cast(kMaxPendingEntries)) { - LOG(ERROR) << "OpLogApplier: too many pending entries (" - << pending_entries_.size() - << "), discarding entry sequence_id=" - << entry.sequence_id << ", key=" << entry.object_key; - return false; - } - - pending_entries_[entry.sequence_id] = entry; - VLOG(1) << "OpLogApplier: future entry buffered, sequence_id=" - << entry.sequence_id << ", expected=" << expected - << ", key=" << entry.object_key - << ", pending_entries=" << pending_entries_.size(); + LOG(ERROR) << "OpLogApplier: future entry rejected, sequence_id=" + << entry.sequence_id << ", expected=" << expected + << ", key=" << entry.object_key; return false; } @@ -128,6 +70,15 @@ bool OpLogApplier::ApplyOpLogEntry(const OpLogEntry& entry) { case OpType::REMOVE: ApplyRemove(entry); break; + case OpType::SEGMENT_MOUNT: + ApplySegmentMount(entry); + break; + case OpType::SEGMENT_UNMOUNT: + ApplySegmentUnmount(entry); + break; + case OpType::SEGMENT_UPDATE: + ApplySegmentUpdate(entry); + break; default: LOG(ERROR) << "OpLogApplier: unsupported op_type=" << static_cast(entry.op_type) @@ -144,9 +95,6 @@ bool OpLogApplier::ApplyOpLogEntry(const OpLogEntry& entry) { HAMetricManager::instance().set_oplog_applied_sequence_id( static_cast(entry.sequence_id)); - // Try to process pending entries - ProcessPendingEntries(); - return true; } @@ -171,257 +119,6 @@ void OpLogApplier::Recover(uint64_t last_applied_sequence_id) { << expected_sequence_id_.load(); } -size_t OpLogApplier::ProcessPendingEntries() { - // Check for missing sequence IDs, possibly skip after timeout, and/or - // request them. - uint64_t missing_seq_to_request = 0; - uint64_t skipped_count = 0; - { - std::lock_guard lock(pending_mutex_); - auto now = std::chrono::steady_clock::now(); - for (;;) { - if (pending_entries_.empty()) { - break; - } - const uint64_t first_pending_seq = pending_entries_.begin()->first; - const uint64_t expected = expected_sequence_id_.load(); - if (IsSequenceOlderOrEqual(first_pending_seq, expected)) { - break; - } - - // There's a gap: expected is missing. - const uint64_t missing_seq = expected; - auto it = missing_sequence_ids_.find(missing_seq); - if (it == missing_sequence_ids_.end()) { - missing_sequence_ids_[missing_seq] = now; - VLOG(1) << "OpLogApplier: scheduling wait for missing " - "sequence_id=" - << missing_seq << ", will request after " - << kMissingEntryRequestSeconds << " seconds"; - break; - } - - const auto waited = - std::chrono::duration_cast(now - - it->second); - - // Skip after timeout to avoid global stall (user requested - // behavior). - if (waited.count() >= kMissingEntrySkipSeconds) { - skipped_sequence_ids_[missing_seq] = now; - missing_sequence_ids_.erase(missing_seq); - expected_sequence_id_.store(missing_seq + 1); - skipped_count++; - HAMetricManager::instance().inc_oplog_skipped_entries(); - LOG(WARNING) - << "OpLogApplier: skipped missing entry seq=" << missing_seq - << " after " << waited.count() << "s timeout"; - continue; // may skip multiple consecutive gaps - } - - // Best-effort request from etcd (before skip triggers). - if (waited.count() >= kMissingEntryRequestSeconds) { - missing_seq_to_request = missing_seq; - break; - } - break; - } - } - - // Request missing OpLog if needed (outside the lock to avoid deadlock) - bool retrieved_missing = false; - if (missing_seq_to_request > 0) { - retrieved_missing = RequestMissingOpLog(missing_seq_to_request); - if (retrieved_missing) { - std::lock_guard lock(pending_mutex_); - missing_sequence_ids_.erase(missing_seq_to_request); - } - } - - size_t processed_count = 0; - for (;;) { - OpLogEntry entry_copy; - bool has_entry = false; - - { - std::lock_guard lock(pending_mutex_); - if (pending_entries_.empty()) { - break; - } - - auto it = pending_entries_.begin(); - const uint64_t expected = expected_sequence_id_.load(); - if (!IsSequenceEqual(it->first, expected)) { - break; // still waiting for earlier sequence_id - } - - entry_copy = it->second; - pending_entries_.erase(it); - has_entry = true; - } - - if (!has_entry) { - break; - } - - // Apply outside lock. - switch (entry_copy.op_type) { - case OpType::PUT_END: - ApplyPutEnd(entry_copy); - break; - case OpType::PUT_REVOKE: - ApplyPutRevoke(entry_copy); - break; - case OpType::REMOVE: - ApplyRemove(entry_copy); - break; - default: - LOG(ERROR) - << "OpLogApplier: unsupported op_type in pending entry"; - break; - } - - expected_sequence_id_.store(entry_copy.sequence_id + 1); - - { - std::lock_guard lock(pending_mutex_); - missing_sequence_ids_.erase(entry_copy.sequence_id); - } - - processed_count++; - } - - // Clean up old missing sequence IDs (older than 1 minute) - { - std::lock_guard lock(pending_mutex_); - auto now = std::chrono::steady_clock::now(); - for (auto it = missing_sequence_ids_.begin(); - it != missing_sequence_ids_.end();) { - auto age = std::chrono::duration_cast( - now - it->second); - if (age.count() > 60) { - LOG(WARNING) - << "OpLogApplier: giving up on missing sequence_id=" - << it->first << " after " << age.count() << " seconds"; - it = missing_sequence_ids_.erase(it); - } else { - ++it; - } - } - - // Clean up old skipped sequence IDs too (avoid unbounded growth). - for (auto it = skipped_sequence_ids_.begin(); - it != skipped_sequence_ids_.end();) { - auto age = std::chrono::duration_cast( - now - it->second); - if (age.count() > 60) { - it = skipped_sequence_ids_.erase(it); - } else { - ++it; - } - } - } - - if (skipped_count > 0) { - LOG(WARNING) << "OpLogApplier: skipped " << skipped_count - << " missing sequence_id(s) after timeout, " - "expected_sequence_id now=" - << expected_sequence_id_.load(); - } - - if (processed_count > 0) { - LOG(INFO) << "OpLogApplier: processed " << processed_count - << " pending entries, expected_sequence_id now=" - << expected_sequence_id_.load(); - } - - // Update pending entries metric - { - std::lock_guard lock(pending_mutex_); - HAMetricManager::instance().set_oplog_pending_entries( - static_cast(pending_entries_.size())); - } - - return processed_count; -} - -OpLogApplier::GapResolveResult OpLogApplier::TryResolveGapsOnceForPromotion( - size_t max_ids) { - GapResolveResult r; - if (oplog_store_ == nullptr) { - return r; - } - - std::vector gap_ids; - gap_ids.reserve(max_ids); - { - std::lock_guard lock(pending_mutex_); - for (const auto& kv : missing_sequence_ids_) { - if (gap_ids.size() >= max_ids) break; - gap_ids.push_back(kv.first); - } - for (const auto& kv : skipped_sequence_ids_) { - if (gap_ids.size() >= max_ids) break; - gap_ids.push_back(kv.first); - } - } - - if (gap_ids.empty()) { - return r; - } - - std::sort(gap_ids.begin(), gap_ids.end()); - gap_ids.erase(std::unique(gap_ids.begin(), gap_ids.end()), gap_ids.end()); - - r.attempted = gap_ids.size(); - std::vector successfully_processed; - for (uint64_t seq : gap_ids) { - OpLogEntry e; - ErrorCode err = oplog_store_->ReadOpLog(seq, e); - if (err != ErrorCode::OK) { - // Log failed gap for monitoring, but don't clear it so it can be - // retried later. - LOG(WARNING) << "Promotion gap resolve: failed to fetch seq=" << seq - << ", err=" << static_cast(err); - continue; - } - r.fetched++; - - // Apply policy: only delete/revoke; drop PUT_END. - if (e.op_type == OpType::REMOVE) { - ApplyRemove(e); - r.applied_deletes++; - successfully_processed.push_back(seq); - } else if (e.op_type == OpType::PUT_REVOKE) { - ApplyPutRevoke(e); - r.applied_deletes++; - successfully_processed.push_back(seq); - } else { - // PUT_END or others: mark as processed (dropped) so we don't retry. - successfully_processed.push_back(seq); - } - } - - // Only clear gaps we successfully fetched and processed. - // Failed gaps remain in missing_sequence_ids_/skipped_sequence_ids_ for - // potential retry or monitoring. - if (!successfully_processed.empty()) { - std::lock_guard lock(pending_mutex_); - for (uint64_t seq : successfully_processed) { - missing_sequence_ids_.erase(seq); - skipped_sequence_ids_.erase(seq); - } - } - return r; -} - -bool OpLogApplier::CheckSequenceOrder(const OpLogEntry& entry) { - // Only check global sequence order. - // Use IsSequenceEqual for wrap-around safety (though equality check doesn't - // need special handling, we use it for consistency). - return IsSequenceEqual(entry.sequence_id, expected_sequence_id_.load()); -} - void OpLogApplier::ApplyPutEnd(const OpLogEntry& entry) { // Payload contains serialized metadata (replicas, size, etc.) in JSON // format. Deserialize the payload immediately and store structured @@ -435,7 +132,8 @@ void OpLogApplier::ApplyPutEnd(const OpLogEntry& entry) { << ", sequence_id=" << entry.sequence_id; StandbyObjectMetadata empty_metadata; empty_metadata.last_sequence_id = entry.sequence_id; - if (!metadata_store_->PutMetadata(entry.object_key, empty_metadata)) { + if (!metadata_store_->PutMetadata(entry.tenant_id, entry.object_key, + empty_metadata)) { LOG(ERROR) << "OpLogApplier: failed to PutMetadata key=" << entry.object_key << ", sequence_id=" << entry.sequence_id; @@ -454,7 +152,8 @@ void OpLogApplier::ApplyPutEnd(const OpLogEntry& entry) { // Fallback to empty metadata if parsing fails StandbyObjectMetadata empty_metadata; empty_metadata.last_sequence_id = entry.sequence_id; - metadata_store_->PutMetadata(entry.object_key, empty_metadata); + metadata_store_->PutMetadata(entry.tenant_id, entry.object_key, + empty_metadata); return; } @@ -462,7 +161,8 @@ void OpLogApplier::ApplyPutEnd(const OpLogEntry& entry) { StandbyObjectMetadata metadata = payload.ToStandbyMetadata(entry.sequence_id); - if (!metadata_store_->PutMetadata(entry.object_key, metadata)) { + if (!metadata_store_->PutMetadata(entry.tenant_id, entry.object_key, + metadata)) { LOG(ERROR) << "OpLogApplier: failed to PutMetadata key=" << entry.object_key << ", sequence_id=" << entry.sequence_id; } else { @@ -478,7 +178,7 @@ void OpLogApplier::ApplyPutRevoke(const OpLogEntry& entry) { // (but the key itself may still exist if there are other replicas). // Current implementation removes the entire key; if we later support // partial replica revocation this logic will need to be refined. - if (!metadata_store_->Remove(entry.object_key)) { + if (!metadata_store_->Remove(entry.tenant_id, entry.object_key)) { LOG(WARNING) << "OpLogApplier: failed to Remove key=" << entry.object_key << " in PUT_REVOKE, sequence_id=" << entry.sequence_id @@ -490,7 +190,7 @@ void OpLogApplier::ApplyPutRevoke(const OpLogEntry& entry) { } void OpLogApplier::ApplyRemove(const OpLogEntry& entry) { - if (!metadata_store_->Remove(entry.object_key)) { + if (!metadata_store_->Remove(entry.tenant_id, entry.object_key)) { LOG(WARNING) << "OpLogApplier: failed to Remove key=" << entry.object_key << ", sequence_id=" << entry.sequence_id @@ -501,62 +201,64 @@ void OpLogApplier::ApplyRemove(const OpLogEntry& entry) { } } -bool OpLogApplier::RequestMissingOpLog(uint64_t missing_seq_id) { - HAMetricManager::instance().inc_oplog_gap_resolve_attempts(); - - if (oplog_store_ == nullptr) { - LOG(WARNING) - << "OpLogApplier: cannot request missing OpLog, no store set"; - return false; - } +const StandbySegmentRegistry& OpLogApplier::GetSegmentRegistry() const { + return segment_registry_; +} - OpLogEntry entry; - ErrorCode err = oplog_store_->ReadOpLog(missing_seq_id, entry); - if (err == ErrorCode::OPLOG_ENTRY_NOT_FOUND) { - LOG(INFO) << "OpLogApplier: missing OpLog entry not found in store, " - "sequence_id=" - << missing_seq_id; - return false; - } - if (err != ErrorCode::OK) { - LOG(ERROR) << "OpLogApplier: failed to read missing OpLog from store, " - "sequence_id=" - << missing_seq_id << ", error=" << static_cast(err); - return false; +void OpLogApplier::LoadSegmentRegistry( + const std::vector& segments) { + segment_registry_.Clear(); + for (const auto& seg : segments) { + segment_registry_.OnSegmentMount(seg); } +} - std::string size_reason; - if (!OpLogManager::ValidateEntrySize(entry, &size_reason)) { - LOG(ERROR) << "OpLogApplier: missing entry size rejected, sequence_id=" - << missing_seq_id << ", key=" << entry.object_key - << ", reason=" << size_reason; - return false; +void OpLogApplier::ApplySegmentMount(const OpLogEntry& entry) { + SegmentMountOp op; + if (struct_pack::deserialize_to(op, entry.payload) != + struct_pack::errc::ok) { + LOG(ERROR) << "Failed to deserialize SEGMENT_MOUNT payload for key " + << entry.object_key; + return; } + StandbySegmentInfo info; + info.segment_name = op.segment_name; + info.transport_endpoint = op.transport_endpoint; + info.capacity = op.capacity; + info.is_memory_segment = op.is_memory_segment; + info.file_path = op.file_path; + segment_registry_.OnSegmentMount(info); + HAMetricManager::instance().inc_oplog_applied_entries(); +} - // Verify checksum before adding to pending entries. - if (!OpLogManager::VerifyChecksum(entry)) { - LOG(ERROR) << "OpLogApplier: checksum mismatch for retrieved missing " - "entry, sequence_id=" - << missing_seq_id << ", key=" << entry.object_key - << ". Possible data corruption. Discarding entry."; - HAMetricManager::instance().inc_oplog_checksum_failures(); - return false; +void OpLogApplier::ApplySegmentUnmount(const OpLogEntry& entry) { + SegmentUnmountOp op; + if (struct_pack::deserialize_to(op, entry.payload) != + struct_pack::errc::ok) { + LOG(ERROR) << "Failed to deserialize SEGMENT_UNMOUNT payload for key " + << entry.object_key; + return; } + segment_registry_.OnSegmentUnmount(op.transport_endpoint); + HAMetricManager::instance().inc_oplog_applied_entries(); +} - // Successfully retrieved the missing OpLog entry - LOG(INFO) << "OpLogApplier: retrieved missing OpLog entry, sequence_id=" - << missing_seq_id - << ", op_type=" << static_cast(entry.op_type) - << ", key=" << entry.object_key; - HAMetricManager::instance().inc_oplog_gap_resolve_success(); - - // Add to pending entries - { - std::lock_guard lock(pending_mutex_); - pending_entries_[entry.sequence_id] = entry; +void OpLogApplier::ApplySegmentUpdate(const OpLogEntry& entry) { + SegmentUpdateOp op; + if (struct_pack::deserialize_to(op, entry.payload) != + struct_pack::errc::ok) { + LOG(ERROR) << "Failed to deserialize SEGMENT_UPDATE payload for key " + << entry.object_key; + return; } - - return true; + StandbySegmentInfo info; + info.segment_name = op.segment_name; + info.transport_endpoint = op.transport_endpoint; + info.capacity = op.capacity; + info.is_memory_segment = op.is_memory_segment; + info.file_path = op.file_path; + segment_registry_.OnSegmentUpdate(info); + HAMetricManager::instance().inc_oplog_applied_entries(); } } // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_batch_codec.cpp b/mooncake-store/src/ha/oplog/oplog_batch_codec.cpp new file mode 100644 index 0000000000..5c5a644bfa --- /dev/null +++ b/mooncake-store/src/ha/oplog/oplog_batch_codec.cpp @@ -0,0 +1,271 @@ +#include "ha/oplog/oplog_batch_codec.h" + +#include +#include + +#include + +#if __has_include() +#include +#else +#include +#endif + +#include "utils/base64.h" + +namespace mooncake { + +namespace { + +void SetReason(std::string* reason, const std::string& value) { + if (reason != nullptr) { + *reason = value; + } +} + +std::string WriteJson(const Json::Value& root) { + Json::StreamWriterBuilder builder; + builder["indentation"] = ""; + std::unique_ptr writer(builder.newStreamWriter()); + std::ostringstream oss; + writer->write(root, &oss); + return oss.str(); +} + +bool ParseJson(const std::string& value, Json::Value* root, + std::string* reason) { + Json::CharReaderBuilder builder; + std::unique_ptr reader(builder.newCharReader()); + std::string errors; + if (!reader->parse(value.data(), value.data() + value.size(), root, + &errors)) { + SetReason(reason, errors.empty() ? "malformed json" : errors); + return false; + } + return true; +} + +bool GetUIntField(const Json::Value& root, const char* field, uint32_t* out, + std::string* reason) { + if (!root.isMember(field)) { + SetReason(reason, std::string("missing field: ") + field); + return false; + } + if (!root[field].isUInt()) { + SetReason(reason, std::string("field must be uint32: ") + field); + return false; + } + *out = root[field].asUInt(); + return true; +} + +bool GetUInt64Field(const Json::Value& root, const char* field, uint64_t* out, + std::string* reason) { + if (!root.isMember(field)) { + SetReason(reason, std::string("missing field: ") + field); + return false; + } + if (!root[field].isUInt64()) { + SetReason(reason, std::string("field must be uint64: ") + field); + return false; + } + *out = root[field].asUInt64(); + return true; +} + +Json::Value BatchRecordToJson(const OpLogBatchRecord& batch, + bool include_checksum) { + Json::Value root(Json::arrayValue); + root.append(static_cast(batch.schema_version)); + root.append(static_cast(batch.batch_id)); + root.append(static_cast(batch.first_seq)); + root.append(static_cast(batch.last_seq)); + Json::Value entries(Json::arrayValue); + for (const auto& entry : batch.entries) { + Json::Value entry_root(Json::arrayValue); + entry_root.append(static_cast(entry.op_type)); + entry_root.append(entry.tenant_id); + entry_root.append(entry.object_key); + entry_root.append(base64::Encode(entry.payload)); + entries.append(entry_root); + } + root.append(std::move(entries)); + if (include_checksum) { + root.append(static_cast(batch.checksum)); + } + return root; +} + +uint32_t ComputeJsonChecksum(const Json::Value& value) { + const std::string stable_payload = WriteJson(value); + return static_cast( + XXH32(stable_payload.data(), stable_payload.size(), 0)); +} + +uint32_t ComputeBatchChecksum(const OpLogBatchRecord& batch) { + return ComputeJsonChecksum( + BatchRecordToJson(batch, /*include_checksum=*/false)); +} + +bool JsonEntryToOpLogEntry(const Json::Value& root, uint64_t sequence_id, + OpLogEntry* entry, std::string* reason) { + if (!root.isArray() || root.size() != 4) { + SetReason(reason, "oplog entry must be a four-element array"); + return false; + } + if (!root[0].isUInt() || !root[1].isString() || !root[2].isString() || + !root[3].isString()) { + SetReason(reason, "oplog entry has invalid field types"); + return false; + } + const uint32_t op_type = root[0].asUInt(); + if (op_type == 0 || op_type >= static_cast(OpType::OP_TYPE_MAX)) { + SetReason(reason, "oplog entry op_type is outside the enum range"); + return false; + } + + const std::string encoded_payload = root[3].asString(); + std::string payload = base64::Decode(encoded_payload); + if (base64::Encode(payload) != encoded_payload) { + SetReason(reason, "oplog entry payload is not canonical base64"); + return false; + } + + entry->sequence_id = sequence_id; + entry->timestamp_ms = 0; + entry->op_type = static_cast(op_type); + entry->tenant_id = root[1].asString(); + entry->object_key = root[2].asString(); + entry->payload = std::move(payload); + entry->checksum = ComputeOpLogChecksum(entry->payload); + entry->prefix_hash = 0; + if (!ValidateOpLogBatchEntry(*entry, reason)) { + return false; + } + return true; +} + +} // namespace + +std::string EncodeDurablePrefix(const DurablePrefix& prefix) { + Json::Value root; + root["schema_version"] = + static_cast(kDurablePrefixSchemaVersion); + root["batch_id"] = static_cast(prefix.batch_id); + root["last_seq"] = static_cast(prefix.last_seq); + return WriteJson(root); +} + +bool DecodeDurablePrefix(const std::string& value, DurablePrefix* prefix, + std::string* reason) { + if (reason != nullptr) { + reason->clear(); + } + if (prefix == nullptr) { + SetReason(reason, "prefix output is null"); + return false; + } + + Json::Value root; + if (!ParseJson(value, &root, reason)) { + return false; + } + if (!root.isObject()) { + SetReason(reason, "durable prefix must be a JSON object"); + return false; + } + uint32_t schema_version = 0; + if (!GetUIntField(root, "schema_version", &schema_version, reason)) { + return false; + } + if (schema_version != kDurablePrefixSchemaVersion) { + SetReason(reason, "unsupported durable prefix schema_version"); + return false; + } + if (!GetUInt64Field(root, "batch_id", &prefix->batch_id, reason)) { + return false; + } + if (!GetUInt64Field(root, "last_seq", &prefix->last_seq, reason)) { + return false; + } + return true; +} + +std::string EncodeOpLogBatchRecord(const OpLogBatchRecord& batch) { + OpLogBatchRecord encoded = batch; + encoded.schema_version = kOpLogBatchRecordSchemaVersion; + encoded.checksum = ComputeBatchChecksum(encoded); + return WriteJson(BatchRecordToJson(encoded, /*include_checksum=*/true)); +} + +bool DecodeOpLogBatchRecord(const std::string& value, OpLogBatchRecord* batch, + std::string* reason) { + if (reason != nullptr) { + reason->clear(); + } + if (batch == nullptr) { + SetReason(reason, "batch output is null"); + return false; + } + + Json::Value root; + if (!ParseJson(value, &root, reason)) { + return false; + } + if (!root.isArray() || root.size() != 6) { + SetReason(reason, "batch record must be a six-element array"); + return false; + } + if (!root[0].isUInt() || !root[1].isUInt64() || !root[2].isUInt64() || + !root[3].isUInt64() || !root[4].isArray() || !root[5].isUInt()) { + SetReason(reason, "batch record has invalid field types"); + return false; + } + const uint32_t schema_version = root[0].asUInt(); + if (schema_version != kOpLogBatchRecordSchemaVersion) { + SetReason(reason, "unsupported batch record schema_version"); + return false; + } + + OpLogBatchRecord decoded; + decoded.schema_version = schema_version; + decoded.batch_id = root[1].asUInt64(); + decoded.first_seq = root[2].asUInt64(); + decoded.last_seq = root[3].asUInt64(); + decoded.checksum = root[5].asUInt(); + const auto& encoded_entries = root[4]; + if (encoded_entries.empty() || decoded.first_seq == 0 || + decoded.last_seq < decoded.first_seq || + decoded.last_seq - decoded.first_seq != encoded_entries.size() - 1) { + SetReason(reason, "batch sequence range does not match entry count"); + return false; + } + + Json::Value checksum_payload(Json::arrayValue); + for (Json::ArrayIndex i = 0; i < 5; ++i) { + checksum_payload.append(root[i]); + } + const uint32_t expected = ComputeJsonChecksum(checksum_payload); + if (decoded.checksum != expected) { + SetReason(reason, "batch record checksum mismatch"); + return false; + } + + decoded.entries.reserve(encoded_entries.size()); + for (Json::ArrayIndex i = 0; i < encoded_entries.size(); ++i) { + OpLogEntry entry; + if (!JsonEntryToOpLogEntry(encoded_entries[i], decoded.first_seq + i, + &entry, reason)) { + return false; + } + decoded.entries.push_back(std::move(entry)); + } + + if (!ValidateOpLogBatchRecordShape(decoded, reason)) { + return false; + } + *batch = std::move(decoded); + return true; +} + +} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_batch_standby_reader.cpp b/mooncake-store/src/ha/oplog/oplog_batch_standby_reader.cpp new file mode 100644 index 0000000000..1c1d9295bd --- /dev/null +++ b/mooncake-store/src/ha/oplog/oplog_batch_standby_reader.cpp @@ -0,0 +1,143 @@ +#include "ha/oplog/oplog_batch_standby_reader.h" + +#include +#include + +#include "ha/oplog/oplog_applier.h" +#include "ha/oplog/oplog_test_failpoint.h" + +namespace mooncake { + +namespace { + +bool IsRetryableBackendError(ErrorCode error) { + return error == ErrorCode::ETCD_OPERATION_ERROR || + error == ErrorCode::ETCD_CTX_CANCELLED; +} + +void SetPollError(OpLogBatchStandbyPollResult& result, ErrorCode error, + bool retryable) { + result.error = error; + result.disposition = retryable ? OpLogBatchStandbyPollDisposition::RETRYABLE + : OpLogBatchStandbyPollDisposition::FATAL; +} + +} // namespace + +OpLogBatchStandbyReader::OpLogBatchStandbyReader(std::string cluster_id, + HaKvBackend& backend, + OpLogApplier& applier) + : storage_(std::move(cluster_id), backend), applier_(applier) {} + +OpLogBatchStandbyPollResult OpLogBatchStandbyReader::PollOnce( + size_t max_batches) { + OpLogBatchStandbyPollResult result; + DurablePrefix prefix; + ErrorCode err = storage_.ReadDurablePrefix(prefix); + if (err == ErrorCode::ETCD_KEY_NOT_EXIST) { + if (batch_format_seen_) { + SetPollError(result, ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, false); + } + return result; + } + if (err != ErrorCode::OK) { + SetPollError(result, err, IsRetryableBackendError(err)); + return result; + } + batch_format_seen_ = true; + result.durable_prefix_present = true; + result.durable_prefix = prefix; + + if (last_observed_prefix_ && + (prefix.batch_id < last_observed_prefix_->batch_id || + prefix.last_seq < last_observed_prefix_->last_seq || + (prefix.batch_id == last_observed_prefix_->batch_id && + prefix.last_seq != last_observed_prefix_->last_seq) || + (prefix.batch_id > last_observed_prefix_->batch_id && + prefix.last_seq == last_observed_prefix_->last_seq))) { + SetPollError(result, ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, false); + return result; + } + + if (prefix.batch_id == 0) { + last_observed_prefix_ = prefix; + if (prefix.last_seq != 0) { + SetPollError(result, ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, false); + } + return result; + } + + TestFailPoint::Wait("standby_prefix_read_before_batch"); + OpLogBatchRecord target_batch; + err = storage_.ReadBatch(prefix.batch_id, target_batch); + if (err != ErrorCode::OK) { + SetPollError(result, err, IsRetryableBackendError(err)); + return result; + } + if (target_batch.last_seq != prefix.last_seq) { + SetPollError(result, ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, false); + return result; + } + last_observed_prefix_ = prefix; + if (prefix.batch_id <= last_applied_batch_id_) { + return result; + } + + std::vector batches; + err = + storage_.ReadBatchesAfter(last_applied_batch_id_, max_batches, batches); + if (err != ErrorCode::OK) { + SetPollError(result, err, IsRetryableBackendError(err)); + return result; + } + for (const auto& batch : batches) { + if (last_scanned_batch_last_seq_ && + (last_applied_batch_id_ == UINT64_MAX || + batch.batch_id != last_applied_batch_id_ + 1)) { + SetPollError(result, ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, false); + return result; + } + if (last_scanned_batch_last_seq_ && + (*last_scanned_batch_last_seq_ == UINT64_MAX || + batch.first_seq != *last_scanned_batch_last_seq_ + 1)) { + SetPollError(result, ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, false); + return result; + } + for (const auto& entry : batch.entries) { + if (entry.sequence_id > prefix.last_seq) { + break; + } + const uint64_t expected = applier_.GetExpectedSequenceId(); + if (IsSequenceOlder(entry.sequence_id, expected)) { + continue; + } + if (IsSequenceNewer(entry.sequence_id, expected)) { + SetPollError(result, ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, + false); + return result; + } + if (!applier_.ApplyOpLogEntry(entry)) { + SetPollError(result, ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, + false); + return result; + } + ++result.applied_entries; + } + last_applied_batch_id_ = batch.batch_id; + last_scanned_batch_last_seq_ = batch.last_seq; + if (last_applied_batch_id_ >= prefix.batch_id) { + break; + } + } + const bool has_more_pages = max_batches != 0 && + batches.size() >= max_batches && + last_applied_batch_id_ < prefix.batch_id; + if (!has_more_pages && + IsSequenceOlderOrEqual(applier_.GetExpectedSequenceId(), + prefix.last_seq)) { + SetPollError(result, ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, false); + } + return result; +} + +} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_batch_storage.cpp b/mooncake-store/src/ha/oplog/oplog_batch_storage.cpp new file mode 100644 index 0000000000..8229b2262e --- /dev/null +++ b/mooncake-store/src/ha/oplog/oplog_batch_storage.cpp @@ -0,0 +1,360 @@ +#include "ha/oplog/oplog_batch_storage.h" + +#include +#include + +#include + +#include "ha/oplog/oplog_batch_codec.h" +#include "ha/oplog/oplog_types.h" +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS +#include "ha_metric_manager.h" +#endif + +namespace mooncake { +namespace { + +bool SameBatchRecord(const OpLogBatchRecord& lhs, const OpLogBatchRecord& rhs) { + if (lhs.schema_version != rhs.schema_version || + lhs.batch_id != rhs.batch_id || lhs.first_seq != rhs.first_seq || + lhs.last_seq != rhs.last_seq || + lhs.entries.size() != rhs.entries.size()) { + return false; + } + for (size_t i = 0; i < lhs.entries.size(); ++i) { + const auto& a = lhs.entries[i]; + const auto& b = rhs.entries[i]; + if (a.sequence_id != b.sequence_id || a.op_type != b.op_type || + a.tenant_id != b.tenant_id || a.object_key != b.object_key || + a.payload != b.payload) { + return false; + } + } + return true; +} + +bool TryParseBatchIdFromKey(const std::string& key, uint64_t& batch_id) { + const size_t slash = key.rfind('/'); + if (slash == std::string::npos || key.size() - slash - 1 != 20) { + return false; + } + const std::string_view suffix(key.data() + slash + 1, 20); + for (char c : suffix) { + if (c < '0' || c > '9') { + return false; + } + } + auto result = + std::from_chars(suffix.data(), suffix.data() + suffix.size(), batch_id); + return result.ec == std::errc() && + result.ptr == suffix.data() + suffix.size(); +} + +} // namespace + +OpLogBatchStorage::OpLogBatchStorage(std::string cluster_id, + HaKvBackend& backend) + : cluster_id_(std::move(cluster_id)), backend_(backend) { + cluster_id_valid_ = + NormalizeAndValidateClusterId(cluster_id_) && !cluster_id_.empty(); +} + +ErrorCode OpLogBatchStorage::InitDurablePrefix(DurablePrefix& prefix) { + if (!IsValidClusterId()) { + return ErrorCode::INVALID_PARAMS; + } + ErrorCode err = RejectLegacyLayout(); + if (err != ErrorCode::OK) { + return err; + } + err = ReadDurablePrefix(prefix); + if (err == ErrorCode::OK) { + return ValidateDurablePrefixAtStartup(prefix); + } + if (err != ErrorCode::ETCD_KEY_NOT_EXIST) { + return err; + } + if (!backend_.SupportsTxn()) { + return ErrorCode::INVALID_PARAMS; + } + + auto batch_range = BuildBatchRecordRange(cluster_id_, 0); + std::vector existing_batches; + err = backend_.Range(batch_range.begin_key, batch_range.end_key, + /*limit=*/1, existing_batches); + if (err != ErrorCode::OK) { + return err; + } + if (!existing_batches.empty()) { + LOG(ERROR) << "Durable prefix is missing but OpLog batch records exist"; + return ErrorCode::INTERNAL_ERROR; + } + + const std::string durable_key = BuildDurablePrefixKey(cluster_id_); + DurablePrefix initial{.batch_id = 0, .last_seq = 0}; + KvTxn txn; + txn.compares.push_back({.key = durable_key, + .kind = KvCompareKind::kKeyNotExists, + .expected_value = ""}); + txn.puts.push_back( + {.key = durable_key, .value = EncodeDurablePrefix(initial)}); + err = backend_.Txn(txn); + if (err == ErrorCode::OK) { + prefix = initial; + return ErrorCode::OK; + } + if (err == ErrorCode::ETCD_TRANSACTION_FAIL) { + err = ReadDurablePrefix(prefix); + if (err != ErrorCode::OK) { + return err; + } + return ValidateDurablePrefixAtStartup(prefix); + } + return err; +} + +ErrorCode OpLogBatchStorage::ReadDurablePrefix(DurablePrefix& prefix) { + if (!IsValidClusterId()) { + return ErrorCode::INVALID_PARAMS; + } + std::string value; + const std::string key = BuildDurablePrefixKey(cluster_id_); + ErrorCode err = backend_.Get(key, value); + if (err != ErrorCode::OK) { + return err; + } + std::string reason; + if (!DecodeDurablePrefix(value, &prefix, &reason)) { + LOG(ERROR) << "Failed to decode durable prefix: " << reason; + return ErrorCode::INTERNAL_ERROR; + } + return ErrorCode::OK; +} + +ErrorCode OpLogBatchStorage::ValidateDurablePrefixAtStartup( + const DurablePrefix& prefix) { + if ((prefix.batch_id == 0) != (prefix.last_seq == 0)) { + LOG(ERROR) << "Durable prefix has inconsistent zero fields: cluster=" + << cluster_id_ << ", batch_id=" << prefix.batch_id + << ", last_seq=" << prefix.last_seq; + return ErrorCode::INTERNAL_ERROR; + } + if (prefix.batch_id == 0) { + const auto range = BuildBatchRecordRange(cluster_id_, 0); + std::vector batches; + ErrorCode err = + backend_.Range(range.begin_key, range.end_key, 1, batches); + if (err != ErrorCode::OK) { + return err; + } + if (!batches.empty()) { + LOG(ERROR) << "Zero durable prefix has batch records: cluster=" + << cluster_id_; + return ErrorCode::INTERNAL_ERROR; + } + return ErrorCode::OK; + } + + OpLogBatchRecord batch; + ErrorCode err = ReadBatch(prefix.batch_id, batch); + if (err == ErrorCode::ETCD_KEY_NOT_EXIST) { + LOG(ERROR) << "Durable prefix terminal batch is missing: cluster=" + << cluster_id_ << ", batch_id=" << prefix.batch_id; + return ErrorCode::INTERNAL_ERROR; + } + if (err != ErrorCode::OK) { + return err; + } + if (batch.last_seq != prefix.last_seq) { + LOG(ERROR) << "Durable prefix last sequence does not match batch: " + << "cluster=" << cluster_id_ + << ", batch_id=" << prefix.batch_id + << ", prefix_last_seq=" << prefix.last_seq + << ", batch_last_seq=" << batch.last_seq; + return ErrorCode::INTERNAL_ERROR; + } + return ErrorCode::OK; +} + +ErrorCode OpLogBatchStorage::WriteBatchAndAdvancePrefix( + const OpLogBatchRecord& batch, const DurablePrefix& expected_prefix) { + if (!IsValidClusterId()) { + return ErrorCode::INVALID_PARAMS; + } + if (!backend_.SupportsTxn()) { + return ErrorCode::INVALID_PARAMS; + } + std::string reason; + if (!ValidateOpLogBatchRecordShape(batch, &reason)) { + LOG(ERROR) << "Invalid OpLog batch record: " << reason; + return ErrorCode::INVALID_PARAMS; + } + if (expected_prefix.batch_id == UINT64_MAX || + expected_prefix.last_seq == UINT64_MAX) { + return ErrorCode::INVALID_PARAMS; + } + const uint64_t expected_batch_id = expected_prefix.batch_id + 1; + const uint64_t expected_first_seq = expected_prefix.last_seq + 1; + if (batch.batch_id != expected_batch_id || + batch.first_seq != expected_first_seq) { + LOG(ERROR) << "OpLog batch does not advance durable prefix " + "contiguously: expected_batch_id=" + << expected_batch_id + << ", actual_batch_id=" << batch.batch_id + << ", expected_first_seq=" << expected_first_seq + << ", actual_first_seq=" << batch.first_seq; + return ErrorCode::INVALID_PARAMS; + } + + const std::string durable_key = BuildDurablePrefixKey(cluster_id_); + const std::string encoded_batch = EncodeOpLogBatchRecord(batch); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + HAMetricManager::instance().observe_batch_record_batch_bytes( + encoded_batch.size()); +#endif + KvTxn txn; + txn.compares.push_back( + {.key = durable_key, + .kind = KvCompareKind::kValueEquals, + .expected_value = EncodeDurablePrefix(expected_prefix)}); + txn.puts.push_back({.key = BuildBatchRecordKey(cluster_id_, batch.batch_id), + .value = encoded_batch}); + txn.puts.push_back( + {.key = durable_key, + .value = EncodeDurablePrefix( + {.batch_id = batch.batch_id, .last_seq = batch.last_seq})}); + ErrorCode err = backend_.Txn(txn); + if (err != ErrorCode::ETCD_TRANSACTION_FAIL) { + return err; + } + + DurablePrefix current_prefix; + if (ReadDurablePrefix(current_prefix) != ErrorCode::OK || + current_prefix.batch_id != batch.batch_id || + current_prefix.last_seq != batch.last_seq) { + return err; + } + OpLogBatchRecord current_batch; + if (ReadBatch(batch.batch_id, current_batch) != ErrorCode::OK || + !SameBatchRecord(batch, current_batch)) { + return err; + } + return ErrorCode::OK; +} + +ErrorCode OpLogBatchStorage::ReadBatch(uint64_t batch_id, + OpLogBatchRecord& batch) { + if (!IsValidClusterId()) { + return ErrorCode::INVALID_PARAMS; + } + std::string value; + ErrorCode err = + backend_.Get(BuildBatchRecordKey(cluster_id_, batch_id), value); + if (err != ErrorCode::OK) { + return err; + } + std::string reason; + if (!DecodeOpLogBatchRecord(value, &batch, &reason)) { + LOG(ERROR) << "Failed to decode OpLog batch record: " << reason; + return ErrorCode::INTERNAL_ERROR; + } + if (batch.batch_id != batch_id) { + LOG(ERROR) << "OpLog batch id does not match key: requested=" + << batch_id << ", payload=" << batch.batch_id; + return ErrorCode::INTERNAL_ERROR; + } + return ErrorCode::OK; +} + +ErrorCode OpLogBatchStorage::ReadBatchesAfter( + uint64_t after_batch_id, size_t limit, + std::vector& batches) { + batches.clear(); + if (!IsValidClusterId()) { + return ErrorCode::INVALID_PARAMS; + } + auto range = BuildBatchRecordRange(cluster_id_, after_batch_id); + std::string begin_key = range.begin_key; + do { + std::vector kvs; + const size_t remaining = limit == 0 ? 0 : limit - batches.size(); + ErrorCode err = + backend_.Range(begin_key, range.end_key, remaining, kvs); + if (err != ErrorCode::OK) { + return err; + } + for (const auto& kv : kvs) { + uint64_t key_batch_id = 0; + if (!TryParseBatchIdFromKey(kv.key, key_batch_id)) { + continue; + } + OpLogBatchRecord batch; + std::string reason; + if (!DecodeOpLogBatchRecord(kv.value, &batch, &reason)) { + LOG(ERROR) << "Failed to decode OpLog batch record at key=" + << kv.key << ": " << reason; + return ErrorCode::INTERNAL_ERROR; + } + if (batch.batch_id != key_batch_id) { + LOG(ERROR) << "OpLog batch id does not match key at key=" + << kv.key; + return ErrorCode::INTERNAL_ERROR; + } + batches.push_back(std::move(batch)); + } + if (limit == 0 || batches.size() >= limit || kvs.size() < remaining) { + break; + } + begin_key = kvs.back().key + '\0'; + } while (begin_key < range.end_key); + return ErrorCode::OK; +} + +bool OpLogBatchStorage::IsValidClusterId() const { return cluster_id_valid_; } + +ErrorCode OpLogBatchStorage::RejectLegacyLayout() const { + const std::string root = "/oplog/" + cluster_id_ + "/"; + std::string ignored; + ErrorCode err = backend_.Get(root + "latest", ignored); + if (err == ErrorCode::OK) { + LOG(ERROR) << "Legacy OpLog latest key exists for cluster=" + << cluster_id_ + << "; clear the legacy OpLog namespace before enabling " + "batch-record OpLog"; + return ErrorCode::INCOMPLETE_OPLOG_CATCH_UP; + } + if (err != ErrorCode::ETCD_KEY_NOT_EXIST) { + return err; + } + + std::vector entries; + err = backend_.Range(root + "00000000000000000000", root + ":", + /*limit=*/1, entries); + if (err != ErrorCode::OK) { + return err; + } + if (!entries.empty()) { + LOG(ERROR) << "Legacy per-entry OpLog key exists for cluster=" + << cluster_id_ + << "; clear the legacy OpLog namespace before enabling " + "batch-record OpLog"; + return ErrorCode::INCOMPLETE_OPLOG_CATCH_UP; + } + + entries.clear(); + err = backend_.Range(root + "snapshot/", root + "snapshot0", + /*limit=*/1, entries); + if (err != ErrorCode::OK) { + return err; + } + if (!entries.empty()) { + LOG(ERROR) << "Legacy OpLog snapshot sidecar exists for cluster=" + << cluster_id_ + << "; clear the legacy OpLog namespace before enabling " + "batch-record OpLog"; + return ErrorCode::INCOMPLETE_OPLOG_CATCH_UP; + } + return ErrorCode::OK; +} + +} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_batch_types.cpp b/mooncake-store/src/ha/oplog/oplog_batch_types.cpp new file mode 100644 index 0000000000..2e91730d31 --- /dev/null +++ b/mooncake-store/src/ha/oplog/oplog_batch_types.cpp @@ -0,0 +1,147 @@ +#include "ha/oplog/oplog_batch_types.h" + +#include +#include + +namespace mooncake { + +namespace { + +void SetReason(std::string* reason, const std::string& value) { + if (reason != nullptr) { + *reason = value; + } +} + +std::string PrefixEnd(std::string prefix) { + for (int i = static_cast(prefix.size()) - 1; i >= 0; --i) { + unsigned char c = static_cast(prefix[i]); + if (c < 0xFF) { + prefix[i] = static_cast(c + 1); + prefix.resize(i + 1); + return prefix; + } + } + return std::string(1, '\0'); +} + +std::string BatchPrefix(const std::string& cluster_id) { + return "/oplog/" + cluster_id + "/batches/"; +} + +} // namespace + +bool ValidateOpLogBatchRecordShape(const OpLogBatchRecord& batch, + std::string* reason) { + if (reason != nullptr) { + reason->clear(); + } + if (batch.schema_version != kOpLogBatchRecordSchemaVersion) { + SetReason(reason, "unsupported schema_version"); + return false; + } + if (batch.batch_id == 0) { + SetReason(reason, "batch_id must be non-zero"); + return false; + } + if (batch.first_seq == 0) { + SetReason(reason, "first_seq must be non-zero"); + return false; + } + if (batch.entries.empty()) { + SetReason(reason, "batch entries must not be empty"); + return false; + } + if (batch.first_seq != batch.entries.front().sequence_id) { + SetReason(reason, "first_seq does not match first entry sequence"); + return false; + } + if (batch.last_seq != batch.entries.back().sequence_id) { + SetReason(reason, "last_seq does not match last entry sequence"); + return false; + } + if (batch.last_seq < batch.first_seq || + batch.last_seq - batch.first_seq != batch.entries.size() - 1) { + SetReason(reason, "batch sequence range does not match entry count"); + return false; + } + for (size_t i = 0; i < batch.entries.size(); ++i) { + const uint64_t expected = batch.first_seq + i; + if (batch.entries[i].sequence_id != expected) { + SetReason(reason, "entry sequences must be contiguous"); + return false; + } + if (!ValidateOpLogBatchEntry(batch.entries[i], reason)) { + return false; + } + } + return true; +} + +bool ValidateOpLogBatchEntry(const OpLogEntry& entry, std::string* reason) { + if (reason != nullptr) { + reason->clear(); + } + const auto op_type = static_cast(entry.op_type); + if (op_type == 0 || op_type >= static_cast(OpType::OP_TYPE_MAX)) { + SetReason(reason, "op_type is outside the valid enum range"); + return false; + } + if (!TenantId(entry.tenant_id).IsValid()) { + SetReason(reason, "tenant_id is empty or invalid"); + return false; + } + return ValidateOpLogEntrySize(entry, reason); +} + +bool ValidateOpLogBatchClusterId(const std::string& cluster_id, + std::string* reason) { + std::string normalized = cluster_id; + if (!NormalizeAndValidateClusterId(normalized) || normalized.empty()) { + SetReason(reason, "invalid cluster_id"); + return false; + } + if (reason != nullptr) { + reason->clear(); + } + return true; +} + +std::string FormatOpLogBatchId(uint64_t batch_id) { + std::ostringstream oss; + oss << std::setw(kOpLogBatchIdWidth) << std::setfill('0') << batch_id; + return oss.str(); +} + +std::string BuildBatchRecordKey(const std::string& cluster_id, + uint64_t batch_id) { + std::string normalized = cluster_id; + if (!NormalizeAndValidateClusterId(normalized) || normalized.empty()) { + return {}; + } + return BatchPrefix(normalized) + FormatOpLogBatchId(batch_id); +} + +std::string BuildDurablePrefixKey(const std::string& cluster_id) { + std::string normalized = cluster_id; + if (!NormalizeAndValidateClusterId(normalized) || normalized.empty()) { + return {}; + } + return "/oplog/" + normalized + "/durable_prefix"; +} + +BatchRecordRange BuildBatchRecordRange(const std::string& cluster_id, + uint64_t after_batch_id) { + std::string normalized = cluster_id; + if (!NormalizeAndValidateClusterId(normalized) || normalized.empty()) { + return {}; + } + const std::string prefix = BatchPrefix(normalized); + if (after_batch_id == UINT64_MAX) { + return {.begin_key = PrefixEnd(prefix), .end_key = PrefixEnd(prefix)}; + } + return {.begin_key = prefix + FormatOpLogBatchId(after_batch_id + 1), + .end_key = PrefixEnd(prefix)}; +} + +} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_manager.cpp b/mooncake-store/src/ha/oplog/oplog_manager.cpp deleted file mode 100644 index b75ccecb56..0000000000 --- a/mooncake-store/src/ha/oplog/oplog_manager.cpp +++ /dev/null @@ -1,190 +0,0 @@ -#include "ha/oplog/oplog_manager.h" - -#include -#include -#include -#include - -#include "ha/oplog/oplog_store.h" - -namespace mooncake { - -OpLogManager::OpLogManager() = default; - -void OpLogManager::SetOpLogStore(std::shared_ptr oplog_store) { - std::unique_lock lock(mutex_); - oplog_store_ = oplog_store; -} - -uint64_t OpLogManager::Append(OpType type, const std::string& key, - const std::string& payload) { - OpLogEntry entry; - entry.op_type = type; - entry.object_key = key; - entry.payload = payload; - entry.timestamp_ms = NowMs(); - entry.checksum = ComputeChecksum(entry.payload); - entry.prefix_hash = ComputePrefixHash(entry.object_key); - - std::unique_lock lock(mutex_); - entry.sequence_id = ++last_seq_id_; - const uint64_t seq = entry.sequence_id; // save before potential unlock - - if (buffer_.size() >= kMaxBufferEntries_) { - buffer_.pop_front(); - ++first_seq_id_; - } - - buffer_.emplace_back(entry); // Copy entry to buffer - - // Write to etcd if EtcdOpLogStore is set. - // Strategy: PUT_END is async (sync=false) — only pushes to batch queue - // (microsecond-level), safe to hold mutex_. - // REMOVE / PUT_REVOKE are sync (sync=true) — blocks until etcd confirms - // persistence; must release mutex_ to avoid blocking other Append calls - // during the wait. The caller relies on sync semantics to know the - // entry is durable before freeing/reusing associated memory. - if (oplog_store_) { - bool sync = (type != OpType::PUT_END); - if (sync) { - // Release lock before the blocking wait to avoid holding - // mutex_ for the entire etcd round-trip. - lock.unlock(); - } - ErrorCode err = oplog_store_->WriteOpLog(entry, sync); - if (err != ErrorCode::OK) { - LOG(WARNING) << "Failed to write OpLog to store, sequence_id=" - << seq << ", but entry is in memory buffer"; - } - } - - return seq; -} - -OpLogEntry OpLogManager::AllocateEntry(OpType type, const std::string& key, - const std::string& payload) { - OpLogEntry entry; - entry.op_type = type; - entry.object_key = key; - entry.payload = payload; - entry.timestamp_ms = NowMs(); - entry.checksum = ComputeChecksum(entry.payload); - entry.prefix_hash = ComputePrefixHash(entry.object_key); - - std::unique_lock lock(mutex_); - entry.sequence_id = ++last_seq_id_; - - if (buffer_.size() >= kMaxBufferEntries_) { - buffer_.pop_front(); - ++first_seq_id_; - } - buffer_.emplace_back(entry); - return entry; -} - -ErrorCode OpLogManager::PersistEntry(const OpLogEntry& entry) const { - std::shared_lock lock(mutex_); - auto store = oplog_store_; - lock.unlock(); - if (!store) { - return ErrorCode::INTERNAL_ERROR; - } - // Strategy 2+: PUT_END is Async, REMOVE (and others) are Sync - bool sync = (entry.op_type != OpType::PUT_END); - return store->WriteOpLog(entry, sync); -} - -tl::expected OpLogManager::AppendAndPersist( - OpType type, const std::string& key, const std::string& payload) { - // Seq pre-allocation semantics: allocate first, then persist. - OpLogEntry entry = AllocateEntry(type, key, payload); - ErrorCode err = PersistEntry(entry); - if (err != ErrorCode::OK) { - return tl::make_unexpected(err); - } - return entry.sequence_id; -} - -uint64_t OpLogManager::GetLastSequenceId() const { - std::shared_lock lock(mutex_); - return last_seq_id_; -} - -void OpLogManager::SetInitialSequenceId(uint64_t sequence_id) { - std::unique_lock lock(mutex_); - if (last_seq_id_ == 0 && buffer_.empty()) { - // Only allow setting initial sequence_id if OpLogManager is empty - last_seq_id_ = sequence_id; - first_seq_id_ = sequence_id + - 1; // first_seq_id_ should be > last_seq_id_ when empty - LOG(INFO) << "OpLogManager initial sequence_id set to " << sequence_id; - } else { - LOG(WARNING) - << "Cannot set initial sequence_id: OpLogManager is not empty " - << "(last_seq_id_=" << last_seq_id_ - << ", buffer_size=" << buffer_.size() << ")"; - } -} - -size_t OpLogManager::GetEntryCount() const { - std::shared_lock lock(mutex_); - return buffer_.size(); -} - -ErrorCode OpLogManager::CleanupOpLogBefore(uint64_t before_sequence_id) { - std::shared_lock lock(mutex_); - if (!oplog_store_) { - return ErrorCode::OK; - } - return oplog_store_->CleanupOpLogBefore(before_sequence_id); -} - -uint64_t OpLogManager::NowMs() { - using namespace std::chrono; - return duration_cast(steady_clock::now().time_since_epoch()) - .count(); -} - -uint32_t OpLogManager::ComputeChecksum(const std::string& data) { - // Use xxHash XXH32 for a fast, deterministic 32-bit checksum. - // Requires linking against xxHash (e.g., libxxhash) and including - // . - return static_cast(XXH32(data.data(), data.size(), 0)); -} - -uint32_t OpLogManager::ComputePrefixHash(const std::string& key) { - if (key.empty()) { - return 0; - } - // Use XXH32 for consistency with ComputeChecksum and better performance. - // XXH32 provides faster hashing and lower collision rate than std::hash. - // Computing hash for the entire key ensures better distribution and fewer - // collisions. - return static_cast(XXH32(key.data(), key.size(), 0)); -} - -bool OpLogManager::VerifyChecksum(const OpLogEntry& entry) { - uint32_t computed = ComputeChecksum(entry.payload); - return computed == entry.checksum; -} - -bool OpLogManager::ValidateEntrySize(const OpLogEntry& entry, - std::string* reason) { - if (entry.object_key.size() > kMaxObjectKeySize) { - if (reason) { - *reason = "object_key too large: size=" + - std::to_string(entry.object_key.size()); - } - return false; - } - if (entry.payload.size() > kMaxPayloadSize) { - if (reason) { - *reason = "payload too large: size=" + - std::to_string(entry.payload.size()); - } - return false; - } - return true; -} - -} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_replicator.cpp b/mooncake-store/src/ha/oplog/oplog_replicator.cpp deleted file mode 100644 index b4b1d7891a..0000000000 --- a/mooncake-store/src/ha/oplog/oplog_replicator.cpp +++ /dev/null @@ -1,80 +0,0 @@ -#include "ha/oplog/oplog_replicator.h" - -#include - -#include "ha/oplog/oplog_applier.h" - -namespace mooncake { - -OpLogReplicator::OpLogReplicator(OpLogChangeNotifier* notifier, - OpLogApplier* applier) - : notifier_(notifier), applier_(applier) { - if (notifier_ == nullptr) { - LOG(FATAL) << "OpLogChangeNotifier cannot be null"; - } - if (applier_ == nullptr) { - LOG(FATAL) << "OpLogApplier cannot be null"; - } -} - -OpLogReplicator::~OpLogReplicator() { Stop(); } - -void OpLogReplicator::Start() { - // Backward-compatible: start from the last processed sequence id. - (void)StartFromSequenceId(last_processed_sequence_id_.load()); -} - -bool OpLogReplicator::StartFromSequenceId(uint64_t start_seq_id) { - if (running_.load()) { - LOG(WARNING) << "OpLogReplicator is already running"; - return true; - } - - auto on_entry = [this](const OpLogEntry& entry) { - if (applier_->ApplyOpLogEntry(entry)) { - uint64_t cur = last_processed_sequence_id_.load(); - while (IsSequenceNewer(entry.sequence_id, cur) && - !last_processed_sequence_id_.compare_exchange_weak( - cur, entry.sequence_id)) { - } - } - }; - - auto on_error = [this](ErrorCode err) { - LOG(ERROR) << "OpLogReplicator: notifier error=" - << static_cast(err); - NotifyStateEvent(StandbyEvent::WATCH_BROKEN); - }; - - ErrorCode err = notifier_->Start(start_seq_id, on_entry, on_error); - if (err != ErrorCode::OK) { - LOG(ERROR) << "Failed to start OpLogChangeNotifier, error=" - << static_cast(err); - return false; - } - - running_.store(true); - NotifyStateEvent(StandbyEvent::WATCH_HEALTHY); - LOG(INFO) << "OpLogReplicator started from sequence_id=" << start_seq_id; - return true; -} - -void OpLogReplicator::Stop() { - if (!running_.load()) { - return; - } - - running_.store(false); - notifier_->Stop(); - LOG(INFO) << "OpLogReplicator stopped"; -} - -uint64_t OpLogReplicator::GetLastProcessedSequenceId() const { - return last_processed_sequence_id_.load(); -} - -bool OpLogReplicator::IsHealthy() const { - return running_.load() && notifier_->IsHealthy(); -} - -} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_serializer.cpp b/mooncake-store/src/ha/oplog/oplog_serializer.cpp deleted file mode 100644 index b13c4b0fed..0000000000 --- a/mooncake-store/src/ha/oplog/oplog_serializer.cpp +++ /dev/null @@ -1,73 +0,0 @@ -#include "ha/oplog/oplog_serializer.h" - -#include -#include - -#if __has_include() -#include -#else -#include -#endif - -#include "utils/base64.h" - -namespace mooncake { - -std::string SerializeOpLogEntry(const OpLogEntry& entry) { - Json::Value root; - root["sequence_id"] = static_cast(entry.sequence_id); - root["timestamp_ms"] = static_cast(entry.timestamp_ms); - root["op_type"] = static_cast(entry.op_type); - root["object_key"] = entry.object_key; - // CRITICAL: Base64 encode binary payload to prevent UTF-8 corruption in - // JSON - root["payload"] = base64::Encode(entry.payload); - root["checksum"] = static_cast(entry.checksum); - root["prefix_hash"] = static_cast(entry.prefix_hash); - - Json::StreamWriterBuilder builder; - builder["indentation"] = ""; // Compact format - std::unique_ptr writer(builder.newStreamWriter()); - std::ostringstream oss; - writer->write(root, &oss); - return oss.str(); -} - -bool DeserializeOpLogEntry(const std::string& json_str, OpLogEntry& entry) { - Json::Value root; - Json::CharReaderBuilder builder; - std::unique_ptr reader(builder.newCharReader()); - std::string errors; - - if (!reader->parse(json_str.data(), json_str.data() + json_str.size(), - &root, &errors)) { - LOG(ERROR) << "Failed to parse JSON: " << errors; - return false; - } - - try { - entry.sequence_id = root["sequence_id"].asUInt64(); - entry.timestamp_ms = root["timestamp_ms"].asUInt64(); - entry.op_type = static_cast(root["op_type"].asInt()); - entry.object_key = root["object_key"].asString(); - // CRITICAL: Base64 decode payload to restore binary data - entry.payload = base64::Decode(root["payload"].asString()); - entry.checksum = root["checksum"].asUInt(); - entry.prefix_hash = root["prefix_hash"].asUInt(); - } catch (const std::exception& e) { - LOG(ERROR) << "Failed to deserialize OpLogEntry: " << e.what(); - return false; - } - - std::string size_reason; - if (!OpLogManager::ValidateEntrySize(entry, &size_reason)) { - LOG(ERROR) << "OpLogSerializer: entry size rejected, sequence_id=" - << entry.sequence_id << ", key=" << entry.object_key - << ", reason=" << size_reason; - return false; - } - - return true; -} - -} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_store_factory.cpp b/mooncake-store/src/ha/oplog/oplog_store_factory.cpp deleted file mode 100644 index a824ed73ab..0000000000 --- a/mooncake-store/src/ha/oplog/oplog_store_factory.cpp +++ /dev/null @@ -1,54 +0,0 @@ -#include "ha/oplog/oplog_store_factory.h" - -#include - -#ifdef STORE_USE_ETCD -#include "ha/oplog/etcd_oplog_store.h" -#endif - -#include "ha/oplog/localfs_oplog_store.h" - -namespace mooncake { - -std::unique_ptr OpLogStoreFactory::Create( - OpLogStoreType type, const std::string& cluster_id, OpLogStoreRole role, - const std::string& oplog_root_dir, int poll_interval_ms) { - switch (type) { - case OpLogStoreType::ETCD: { -#ifdef STORE_USE_ETCD - bool batch_update = (role == OpLogStoreRole::WRITER); - bool batch_write = (role == OpLogStoreRole::WRITER); - auto store = std::make_unique( - cluster_id, batch_update, batch_write); - if (store->Init() != ErrorCode::OK) { - LOG(ERROR) << "OpLogStoreFactory: failed to init EtcdOpLogStore" - << ", cluster_id=" << cluster_id; - return nullptr; - } - return store; -#else - LOG(ERROR) << "OpLogStoreFactory: ETCD support not compiled in"; - return nullptr; -#endif - } - case OpLogStoreType::LOCAL_FS: { - bool enable_batch_write = (role == OpLogStoreRole::WRITER); - auto store = std::make_unique( - cluster_id, oplog_root_dir, enable_batch_write, - poll_interval_ms); - if (store->Init() != ErrorCode::OK) { - LOG(ERROR) - << "OpLogStoreFactory: failed to init LocalFsOpLogStore" - << ", cluster_id=" << cluster_id - << ", root_dir=" << oplog_root_dir; - return nullptr; - } - return store; - } - default: - LOG(ERROR) << "OpLogStoreFactory: unknown store type"; - return nullptr; - } -} - -} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/oplog_test_failpoint.cpp b/mooncake-store/src/ha/oplog/oplog_test_failpoint.cpp new file mode 100644 index 0000000000..bfd5e67c51 --- /dev/null +++ b/mooncake-store/src/ha/oplog/oplog_test_failpoint.cpp @@ -0,0 +1,115 @@ +#include "ha/oplog/oplog_test_failpoint.h" + +#ifdef MOONCAKE_ENABLE_TEST_FAILPOINTS + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace mooncake { +namespace { + +bool IsValidName(std::string_view name) { + if (name.empty()) { + return false; + } + for (char c : name) { + if ((c < 'a' || c > 'z') && (c < '0' || c > '9') && c != '_') { + return false; + } + } + return true; +} + +std::chrono::seconds Timeout() { + const char* value = std::getenv("MOONCAKE_TEST_FAILPOINT_TIMEOUT_SEC"); + if (value == nullptr) { + return std::chrono::seconds(30); + } + char* end = nullptr; + const long seconds = std::strtol(value, &end, 10); + if (end == value || *end != '\0' || seconds <= 0 || seconds > 3600) { + return std::chrono::seconds(30); + } + return std::chrono::seconds(seconds); +} + +void Remove(const std::filesystem::path& path) { + std::error_code error; + std::filesystem::remove(path, error); +} + +} // namespace + +bool TestFailPoint::Wait(std::string_view name) { + const char* directory = std::getenv("MOONCAKE_TEST_FAILPOINT_DIR"); + if (directory == nullptr || !IsValidName(name)) { + return false; + } + const std::filesystem::path root(directory); + const std::string base(name); + const auto arm = root / (base + ".arm"); + const auto claim = root / (base + ".claimed." + std::to_string(::getpid())); + std::error_code error; + std::filesystem::rename(arm, claim, error); + if (error) { + return false; + } + + const auto hit = root / (base + ".hit"); + const auto release = root / (base + ".release"); + const auto temporary = + root / (base + ".hit.tmp." + std::to_string(::getpid())); + Remove(release); + { + std::ofstream output(temporary); + if (!output) { + Remove(claim); + LOG(ERROR) << "Failed to create failpoint hit file: " << temporary; + return false; + } + output << ::getpid() << '\n'; + } + error.clear(); + std::filesystem::rename(temporary, hit, error); + if (error) { + Remove(temporary); + Remove(claim); + LOG(ERROR) << "Failed to publish failpoint hit file: " + << error.message(); + return false; + } + + const auto deadline = std::chrono::steady_clock::now() + Timeout(); + while (std::chrono::steady_clock::now() < deadline) { + if (std::filesystem::exists(release)) { + Remove(release); + Remove(hit); + Remove(claim); + return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + Remove(hit); + Remove(claim); + LOG(ERROR) << "Timed out waiting for failpoint release: " << name; + return false; +} + +} // namespace mooncake + +#else + +namespace mooncake { + +bool TestFailPoint::Wait(std::string_view) { return false; } + +} // namespace mooncake + +#endif diff --git a/mooncake-store/src/ha/oplog/oplog_types.cpp b/mooncake-store/src/ha/oplog/oplog_types.cpp new file mode 100644 index 0000000000..9841da3ce3 --- /dev/null +++ b/mooncake-store/src/ha/oplog/oplog_types.cpp @@ -0,0 +1,53 @@ +#include "ha/oplog/oplog_types.h" + +#include + +namespace mooncake { + +namespace { + +uint32_t ComputePrefixHash(std::string_view key) { + if (key.empty()) { + return 0; + } + return static_cast(XXH32(key.data(), key.size(), 0)); +} + +} // namespace + +bool NormalizeAndValidateClusterId(std::string& cluster_id) { + while (!cluster_id.empty() && cluster_id.back() == '/') { + cluster_id.pop_back(); + } + return cluster_id.empty() || IsValidClusterIdComponent(cluster_id); +} + +uint32_t ComputeOpLogChecksum(std::string_view payload) { + return static_cast(XXH32(payload.data(), payload.size(), 0)); +} + +bool VerifyOpLogChecksum(const OpLogEntry& entry) { + return ComputeOpLogChecksum(entry.payload) == entry.checksum && + (entry.prefix_hash == 0 || + ComputePrefixHash(entry.object_key) == entry.prefix_hash); +} + +bool ValidateOpLogEntrySize(const OpLogEntry& entry, std::string* reason) { + if (entry.object_key.size() > kMaxOpLogObjectKeySize) { + if (reason != nullptr) { + *reason = "object_key too large: size=" + + std::to_string(entry.object_key.size()); + } + return false; + } + if (entry.payload.size() > kMaxOpLogPayloadSize) { + if (reason != nullptr) { + *reason = "payload too large: size=" + + std::to_string(entry.payload.size()); + } + return false; + } + return true; +} + +} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/ordered_oplog_writer.cpp b/mooncake-store/src/ha/oplog/ordered_oplog_writer.cpp new file mode 100644 index 0000000000..43807a4703 --- /dev/null +++ b/mooncake-store/src/ha/oplog/ordered_oplog_writer.cpp @@ -0,0 +1,399 @@ +#include "ha/oplog/ordered_oplog_writer.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "ha/oplog/oplog_test_failpoint.h" +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS +#include "ha_metric_manager.h" +#endif + +namespace mooncake { + +struct OrderedOpLogWriter::Impl { +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + using Clock = std::chrono::steady_clock; +#endif + + struct PendingEntry { + OpLogEntry entry; + DurableCallback callback; +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + Clock::time_point committed_at{}; + Clock::time_point durable_at{}; +#endif + }; + + explicit Impl(OrderedOpLogWriterConfig config, WriteBatchFn write_batch) + : config(std::move(config)), write_batch(std::move(write_batch)) { + if (this->config.max_entries_per_batch == 0) { + this->config.max_entries_per_batch = 1; + } + if (this->config.initial_durable_prefix.last_seq == UINT64_MAX || + this->config.initial_durable_prefix.batch_id == UINT64_MAX) { + accepting = false; + last_error = ErrorCode::INVALID_PARAMS; + return; + } + next_sequence_id = this->config.initial_durable_prefix.last_seq + 1; + } + + void SealCommittedEntriesIfIdle() { + if (batch_busy || committed_entries.empty()) { + return; + } + const size_t count = + std::min(committed_entries.size(), config.max_entries_per_batch); + ready_entries.reserve(count); + for (size_t i = 0; i < count; ++i) { + ready_entries.push_back(std::move(committed_entries.front())); + committed_entries.pop_front(); + } +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + HAMetricManager::instance().set_batch_record_committed_queue_depth( + committed_entries.size() + ready_entries.size()); +#endif + open_waiting_slots -= count; + batch_busy = true; + } + + OrderedOpLogWriterConfig config; + WriteBatchFn write_batch; + mutable std::mutex mutex; + std::condition_variable cv; + bool accepting{true}; + bool running{false}; + bool stop_requested{false}; + bool callback_stop_requested{false}; + ErrorCode last_error{ErrorCode::OK}; + uint64_t next_reservation_id{1}; + uint64_t next_sequence_id{1}; + DurablePrefix durable_prefix{config.initial_durable_prefix}; + size_t open_waiting_slots{0}; + std::unordered_set active_reservations; + std::deque committed_entries; + std::vector ready_entries; + bool batch_busy{false}; + std::deque callback_entries; + std::thread writer_thread; + std::thread callback_thread; +}; + +OrderedOpLogWriter::Reservation::Reservation() = default; + +OrderedOpLogWriter::Reservation::Reservation(OrderedOpLogWriter* writer, + uint64_t id) + : writer_(writer), id_(id) {} + +OrderedOpLogWriter::Reservation::Reservation(Reservation&& other) noexcept + : writer_(other.writer_), id_(other.id_) { + other.writer_ = nullptr; + other.id_ = 0; +} + +OrderedOpLogWriter::Reservation& OrderedOpLogWriter::Reservation::operator=( + Reservation&& other) noexcept { + if (this != &other) { + if (writer_ != nullptr) { + writer_->Abort(std::move(*this)); + } + writer_ = other.writer_; + id_ = other.id_; + other.writer_ = nullptr; + other.id_ = 0; + } + return *this; +} + +OrderedOpLogWriter::Reservation::~Reservation() { + if (writer_ != nullptr) { + writer_->Abort(std::move(*this)); + } +} + +OrderedOpLogWriter::PendingHandle::PendingHandle() = default; + +OrderedOpLogWriter::PendingHandle::PendingHandle(uint64_t sequence_id) + : sequence_id_(sequence_id) {} + +uint64_t OrderedOpLogWriter::PendingHandle::sequence_id() const { + return sequence_id_; +} + +OrderedOpLogWriter::OrderedOpLogWriter(OrderedOpLogWriterConfig config, + WriteBatchFn write_batch) + : impl_(std::make_unique(std::move(config), std::move(write_batch))) { +} + +OrderedOpLogWriter::~OrderedOpLogWriter() { Stop(); } + +tl::expected +OrderedOpLogWriter::Reserve() { + std::lock_guard lock(impl_->mutex); + if (!impl_->accepting) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } + if (impl_->open_waiting_slots >= impl_->config.max_entries_per_batch) { + return tl::make_unexpected(ErrorCode::TASK_PENDING_LIMIT_EXCEEDED); + } + const uint64_t id = impl_->next_reservation_id++; + impl_->active_reservations.insert(id); + ++impl_->open_waiting_slots; + return Reservation(this, id); +} + +tl::expected +OrderedOpLogWriter::Commit(Reservation&& reservation, OpLogEntry entry, + DurableCallback callback) { + std::lock_guard lock(impl_->mutex); + if (reservation.writer_ != this || + impl_->active_reservations.erase(reservation.id_) == 0) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + if (impl_->stop_requested) { + --impl_->open_waiting_slots; + reservation.writer_ = nullptr; + reservation.id_ = 0; + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } + std::string reason; + if (!ValidateOpLogBatchEntry(entry, &reason)) { + --impl_->open_waiting_slots; + reservation.writer_ = nullptr; + reservation.id_ = 0; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + reservation.writer_ = nullptr; + reservation.id_ = 0; + entry.timestamp_ms = 0; + entry.checksum = ComputeOpLogChecksum(entry.payload); + entry.prefix_hash = 0; + entry.sequence_id = impl_->next_sequence_id++; + const uint64_t sequence_id = entry.sequence_id; + Impl::PendingEntry pending{.entry = std::move(entry), + .callback = std::move(callback)}; +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + pending.committed_at = Impl::Clock::now(); +#endif + impl_->committed_entries.push_back(std::move(pending)); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + HAMetricManager::instance().set_batch_record_committed_queue_depth( + impl_->committed_entries.size() + impl_->ready_entries.size()); +#endif + impl_->SealCommittedEntriesIfIdle(); + impl_->cv.notify_all(); + return PendingHandle(sequence_id); +} + +void OrderedOpLogWriter::Abort(Reservation&& reservation) { + std::lock_guard lock(impl_->mutex); + if (reservation.writer_ == this && + impl_->active_reservations.erase(reservation.id_) != 0) { + --impl_->open_waiting_slots; + } + reservation.writer_ = nullptr; + reservation.id_ = 0; +} + +bool OrderedOpLogWriter::IsAccepting() const { + std::lock_guard lock(impl_->mutex); + return impl_->accepting; +} + +ErrorCode OrderedOpLogWriter::LastError() const { + std::lock_guard lock(impl_->mutex); + return impl_->last_error; +} + +void OrderedOpLogWriter::Start() { + std::lock_guard lock(impl_->mutex); + if (impl_->running || impl_->stop_requested) { + return; + } + impl_->callback_stop_requested = false; + impl_->running = true; + impl_->callback_thread = std::thread([this] { + while (true) { + Impl::PendingEntry callback_entry; + { + std::unique_lock lock(impl_->mutex); + impl_->cv.wait(lock, [&] { + return impl_->callback_stop_requested || + !impl_->callback_entries.empty(); + }); + if (impl_->callback_stop_requested && + impl_->callback_entries.empty()) { + return; + } + callback_entry = std::move(impl_->callback_entries.front()); + impl_->callback_entries.pop_front(); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + HAMetricManager::instance() + .set_batch_record_callback_queue_depth( + impl_->callback_entries.size()); +#endif + } + if (callback_entry.callback) { +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + const auto callback_started_at = Impl::Clock::now(); + const auto latency_us = + std::chrono::duration_cast( + callback_started_at - callback_entry.durable_at) + .count(); + HAMetricManager::instance() + .observe_batch_record_callback_latency_us(latency_us); +#endif + callback_entry.callback(callback_entry.entry); + } + } + }); + impl_->writer_thread = std::thread([this] { + while (true) { + std::vector entries; + DurablePrefix expected_prefix; + { + std::unique_lock lock(impl_->mutex); + impl_->cv.wait(lock, [&] { + return impl_->stop_requested || + !impl_->ready_entries.empty(); + }); + if (impl_->stop_requested && impl_->ready_entries.empty()) { + return; + } + entries = std::move(impl_->ready_entries); + impl_->ready_entries.clear(); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + HAMetricManager::instance() + .set_batch_record_committed_queue_depth( + impl_->committed_entries.size()); +#endif + expected_prefix = impl_->durable_prefix; + } + + OpLogBatchRecord batch; + batch.batch_id = expected_prefix.batch_id + 1; + batch.first_seq = entries.front().entry.sequence_id; + batch.last_seq = entries.back().entry.sequence_id; + batch.entries.reserve(entries.size()); + for (auto& entry : entries) { + batch.entries.push_back(std::move(entry.entry)); + } + + constexpr auto kInitialRetryDelay = std::chrono::milliseconds(1); + constexpr auto kMaxRetryDelay = std::chrono::milliseconds(1000); + auto retry_delay = kInitialRetryDelay; + + while (true) { + TestFailPoint::Wait("batch_before_txn"); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + const auto txn_started_at = Impl::Clock::now(); +#endif + ErrorCode err = impl_->write_batch(batch, expected_prefix); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + const auto txn_latency_us = + std::chrono::duration_cast( + Impl::Clock::now() - txn_started_at) + .count(); + HAMetricManager::instance().observe_batch_record_txn_latency_us( + txn_latency_us); +#endif + if (err == ErrorCode::OK) { + TestFailPoint::Wait("batch_txn_succeeded_before_callback"); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + const auto durable_at = Impl::Clock::now(); + for (const auto& entry : entries) { + HAMetricManager::instance() + .observe_batch_record_commit_to_durable_us( + std::chrono::duration_cast< + std::chrono::microseconds>( + durable_at - entry.committed_at) + .count()); + } +#endif + { + std::lock_guard lock(impl_->mutex); + impl_->durable_prefix = {.batch_id = batch.batch_id, + .last_seq = batch.last_seq}; + impl_->last_error = ErrorCode::OK; + impl_->accepting = !impl_->stop_requested; + for (size_t i = 0; i < entries.size(); ++i) { +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + entries[i].durable_at = durable_at; +#endif + entries[i].entry = batch.entries[i]; + impl_->callback_entries.push_back( + std::move(entries[i])); + } +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + auto& metrics = HAMetricManager::instance(); + metrics.inc_batch_record_durable_batches(); + metrics.inc_batch_record_durable_entries( + batch.entries.size()); + metrics.observe_batch_record_batch_entries( + batch.entries.size()); + metrics.set_batch_record_last_batch_id(batch.batch_id); + metrics.set_batch_record_durable_sequence( + batch.last_seq); + metrics.set_batch_record_callback_queue_depth( + impl_->callback_entries.size()); +#endif + impl_->batch_busy = false; + impl_->SealCommittedEntriesIfIdle(); + } + impl_->cv.notify_all(); + break; + } + + { + std::unique_lock lock(impl_->mutex); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + HAMetricManager::instance().inc_batch_record_retries(); +#endif + impl_->last_error = err; + impl_->accepting = false; + if (impl_->stop_requested || + impl_->cv.wait_for(lock, retry_delay, [this] { + return impl_->stop_requested; + })) { + return; + } + } + retry_delay = std::min(retry_delay * 2, kMaxRetryDelay); + } + } + }); +} + +void OrderedOpLogWriter::Stop() { + { + std::lock_guard lock(impl_->mutex); + impl_->accepting = false; + impl_->stop_requested = true; + if (!impl_->running) { + return; + } + } + impl_->cv.notify_all(); + if (impl_->writer_thread.joinable()) { + impl_->writer_thread.join(); + } + { + std::lock_guard lock(impl_->mutex); + impl_->callback_stop_requested = true; + } + impl_->cv.notify_all(); + if (impl_->callback_thread.joinable()) { + impl_->callback_thread.join(); + } + std::lock_guard lock(impl_->mutex); + impl_->running = false; +} + +} // namespace mooncake diff --git a/mooncake-store/src/ha/oplog/polling_oplog_change_notifier.cpp b/mooncake-store/src/ha/oplog/polling_oplog_change_notifier.cpp deleted file mode 100644 index b228e3f44c..0000000000 --- a/mooncake-store/src/ha/oplog/polling_oplog_change_notifier.cpp +++ /dev/null @@ -1,75 +0,0 @@ -#include "ha/oplog/polling_oplog_change_notifier.h" - -#include - -namespace mooncake { - -PollingOpLogChangeNotifier::PollingOpLogChangeNotifier(OpLogStore* store, - int poll_interval_ms) - : store_(store), poll_interval_ms_(poll_interval_ms) {} - -PollingOpLogChangeNotifier::~PollingOpLogChangeNotifier() { Stop(); } - -ErrorCode PollingOpLogChangeNotifier::Start(uint64_t start_sequence_id, - EntryCallback on_entry, - ErrorCallback on_error) { - if (running_.load()) { - return ErrorCode::INTERNAL_ERROR; - } - - on_entry_ = std::move(on_entry); - on_error_ = std::move(on_error); - last_sequence_id_.store(start_sequence_id); - running_.store(true); - healthy_.store(true); - poll_thread_ = std::thread(&PollingOpLogChangeNotifier::PollLoop, this); - - return ErrorCode::OK; -} - -void PollingOpLogChangeNotifier::Stop() { - running_.store(false); - stop_cv_.notify_all(); - if (poll_thread_.joinable()) { - poll_thread_.join(); - } - healthy_.store(false); -} - -bool PollingOpLogChangeNotifier::IsHealthy() const { - return running_.load() && healthy_.load(); -} - -void PollingOpLogChangeNotifier::PollLoop() { - while (running_.load()) { - uint64_t last_seq = last_sequence_id_.load(); - std::vector entries; - auto err = store_->ReadOpLogSince(last_seq, kPollBatchSize, entries); - - if (err == ErrorCode::OK && !entries.empty()) { - for (const auto& entry : entries) { - if (!running_.load()) break; - on_entry_(entry); - } - last_sequence_id_.store(entries.back().sequence_id); - healthy_.store(true); - // Data available — poll again immediately to drain backlog - continue; - } else if (err != ErrorCode::OK) { - if (on_error_) { - on_error_(err); - } - healthy_.store(false); - } - - // Interruptible sleep: wakes immediately on Stop() - { - std::unique_lock lock(stop_mutex_); - stop_cv_.wait_for(lock, - std::chrono::milliseconds(poll_interval_ms_), - [&] { return !running_.load(); }); - } - } -} - -} // namespace mooncake diff --git a/mooncake-store/src/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.cpp b/mooncake-store/src/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.cpp index a3f59e7d1a..1b60194a14 100644 --- a/mooncake-store/src/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.cpp +++ b/mooncake-store/src/ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.cpp @@ -6,6 +6,8 @@ #include +#include "ascii_string.h" + namespace mooncake { namespace ha { namespace backends { @@ -94,8 +96,7 @@ EmbeddedSnapshotCatalogStore::GetLatest() { return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL); } - latest_snapshot_id = snapshot_catalog_store_detail::TrimAsciiWhitespace( - std::move(latest_snapshot_id)); + latest_snapshot_id = std::string(TrimAsciiWhitespace(latest_snapshot_id)); if (latest_snapshot_id.empty()) { return std::optional(); } diff --git a/mooncake-store/src/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.cpp b/mooncake-store/src/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.cpp index 29e7371dc0..93dd5c9ae8 100644 --- a/mooncake-store/src/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.cpp +++ b/mooncake-store/src/ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.cpp @@ -9,6 +9,8 @@ #include #include "types.h" +#include "ascii_string.h" +#include "integer_parser.h" #ifdef STORE_USE_REDIS #include #endif @@ -43,11 +45,11 @@ tl::expected ParseSnapshotScore( } } - try { - return std::stoll(digits); - } catch (const std::exception&) { + const auto score = TryParseInteger(digits); + if (!score.has_value()) { return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } + return *score; } tl::expected LoadSnapshotDescriptor( @@ -211,9 +213,8 @@ RedisSnapshotCatalogStore::GetLatest() { return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL); } - auto latest_snapshot_id = - snapshot_catalog_store_detail::TrimAsciiWhitespace( - std::string(reply->str, reply->len)); + std::string latest_snapshot_id(reply->str, reply->len); + latest_snapshot_id = std::string(TrimAsciiWhitespace(latest_snapshot_id)); if (latest_snapshot_id.empty()) { return std::optional(); } diff --git a/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp b/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp index dc150e8ea9..54ee82056d 100644 --- a/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp +++ b/mooncake-store/src/ha/snapshot/catalog_backed_snapshot_provider.cpp @@ -16,7 +16,7 @@ #include "ha/snapshot/catalog/snapshot_catalog_store.h" #include "ha/snapshot/object/snapshot_object_store.h" #include "segment.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" #include "utils/zstd_util.h" namespace mooncake { @@ -141,11 +141,12 @@ DeserializeStandbyObjectMetadata( // v3: 9 + replica_count, data_type + hard_pinned or // hard_pinned + group_id // v4: 10 + replica_count, data_type + hard_pinned + group_id + // v5: 11 + replica_count, v4 + object_checksum (ignored here) // 64-bit arithmetic keeps an attacker-controlled near-UINT32_MAX // replica_count from wrapping the bounds and slipping an out-of-bounds // index through. constexpr uint64_t kBaseFieldCount = 7; - constexpr uint64_t kMaxOptionalFieldCount = 3; + constexpr uint64_t kMaxOptionalFieldCount = 4; const uint64_t total_elements = object.via.array.size; const uint64_t min_elements = kBaseFieldCount + replica_count; if (total_elements < min_elements || @@ -156,12 +157,12 @@ DeserializeStandbyObjectMetadata( return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); } - // Skip the optional data_type; the standby restore path does not use - // it. A leading positive integer is data_type, whereas a replica is - // serialized as an array. + // Read data_type if present (8+ or 10+ format) + ObjectDataType data_type = ObjectDataType::UNKNOWN; if (index < total_elements && array[index].type == msgpack::type::POSITIVE_INTEGER) { - ++index; // data_type + data_type = + static_cast(array[index++].as()); } const auto lease_timeout = std::chrono::system_clock::time_point( @@ -211,11 +212,25 @@ DeserializeStandbyObjectMetadata( return std::optional(); } + // Read hard_pinned and group_id if present. Standby does not need + // hard_pinned, but promotion needs group_id to rebuild group indexes. + if (index < total_elements && + array[index].type == msgpack::type::BOOLEAN) { + ++index; // hard_pinned + } + + std::string group_id; + if (index < total_elements && array[index].type == msgpack::type::STR) { + group_id = array[index++].as(); + } + StandbyObjectMetadata metadata; metadata.client_id = client_id; metadata.size = size; metadata.replicas = std::move(replicas); metadata.last_sequence_id = snapshot_sequence_id; + metadata.data_type = data_type; + metadata.group_id = std::move(group_id); return std::optional(std::move(metadata)); } catch (const std::exception& ex) { LOG(ERROR) << "Failed to parse snapshot metadata entry: " << ex.what(); @@ -223,8 +238,7 @@ DeserializeStandbyObjectMetadata( } } -tl::expected>, - ErrorCode> +tl::expected, ErrorCode> DeserializeStandbySnapshotMetadata(const std::vector& data, const SegmentView& segment_view, uint64_t snapshot_sequence_id) { @@ -245,7 +259,7 @@ DeserializeStandbySnapshotMetadata(const std::vector& data, return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); } - std::vector> snapshot; + std::vector snapshot; const auto now = std::chrono::system_clock::now(); for (uint32_t i = 0; i < shards->via.map.size; ++i) { const auto& shard_blob = shards->via.map.ptr[i].val; @@ -286,23 +300,46 @@ DeserializeStandbySnapshotMetadata(const std::vector& data, for (uint32_t j = 0; j < metadata_entries->via.array.size; ++j) { const auto& item = metadata_entries->via.array.ptr[j]; - if (item.type != msgpack::type::ARRAY || item.via.array.size != 2) { + if (item.type != msgpack::type::ARRAY) { LOG(ERROR) << "Snapshot metadata item has invalid shape"; return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); } try { - const std::string key = item.via.array.ptr[0].as(); + std::string tenant_id = "default"; + std::string key; + size_t metadata_index; + + if (item.via.array.size == 2) { + // Old format: [key, metadata] + key = item.via.array.ptr[0].as(); + metadata_index = 1; + } else if (item.via.array.size == 3) { + // New format: [tenant_id, key, metadata] + tenant_id = item.via.array.ptr[0].as(); + key = item.via.array.ptr[1].as(); + metadata_index = 2; + } else { + LOG(ERROR) + << "Snapshot metadata item has invalid array size: " + << item.via.array.size; + return tl::make_unexpected(ErrorCode::DESERIALIZE_FAIL); + } + + const auto normalized_tenant = NormalizeTenantId(tenant_id); + auto metadata_result = DeserializeStandbyObjectMetadata( - item.via.array.ptr[1], segment_view, snapshot_sequence_id, - now); + item.via.array.ptr[metadata_index], segment_view, + snapshot_sequence_id, now); if (!metadata_result) { return tl::make_unexpected(metadata_result.error()); } if (!metadata_result->has_value()) { continue; } - snapshot.emplace_back(key, std::move(metadata_result->value())); + snapshot.push_back( + StandbyObjectEntry{normalized_tenant, key, + std::move(metadata_result->value())}); } catch (const std::exception& ex) { LOG(ERROR) << "Failed to parse snapshot metadata item: " << ex.what(); @@ -458,6 +495,52 @@ class CatalogBackedSnapshotProvider final : public SnapshotProvider { snapshot.snapshot_id = descriptor.snapshot_id; snapshot.snapshot_sequence_id = descriptor.last_included_seq; snapshot.metadata = std::move(deserialize_metadata.value()); + + // Extract standby segment registry entries from the deserialized + // SegmentManager. The snapshot's SegmentSerializer::Serialize() + // currently carries enough data to rebuild only memory segments + // (segment_manager.mounted_segments_, where buf_allocator is non-null + // by construction — MountSegment is the only path that populates it). + // + // Local-disk segments are serialized as per-client offloading + // bookkeeping (client_local_disk_segment_'s offloading_objects map) + // without the transport_endpoint / file_path / capacity fields that + // StandbySegmentInfo needs, and NoF segments are not serialized at + // all in this snapshot path. Both have to be re-mounted explicitly + // via SEGMENT_MOUNT OpLog replay after standby promotion (see + // OpLogApplier::Apply / HotStandbyService::LoadSnapshotBaselineLocked). + // + // If a future change makes the serializer carry richer per-segment + // data, the predicate below should be replaced with explicit branches + // for each segment type. + ScopedSegmentAccess segment_access = segment_manager.getSegmentAccess(); + std::vector> all_segments; + segment_access.GetAllSegments(all_segments); + SegmentView view = segment_manager.getView(); + for (const auto& [seg, client_id] : all_segments) { + MountedSegment mounted; + if (view.GetMountedSegment(seg.id, mounted) != ErrorCode::OK) { + continue; + } + if (mounted.buf_allocator == nullptr) { + // Defensive: a MountedSegment without an allocator should + // not exist in the snapshot today. Log and skip rather + // than emitting a half-populated StandbySegmentInfo. + LOG(WARNING) + << "snapshot contains MountedSegment without allocator; " + << "skipping segment_name=" << seg.name + << " segment_id=" << seg.id; + continue; + } + StandbySegmentInfo info; + info.segment_name = seg.name; + info.transport_endpoint = seg.te_endpoint; + info.capacity = seg.size; + info.is_memory_segment = true; + // file_path stays empty for memory segments by contract. + snapshot.segments.push_back(std::move(info)); + } + return std::optional(std::move(snapshot)); } diff --git a/mooncake-store/src/ha/snapshot/master_snapshot_codec.cpp b/mooncake-store/src/ha/snapshot/master_snapshot_codec.cpp new file mode 100644 index 0000000000..96e7da5c02 --- /dev/null +++ b/mooncake-store/src/ha/snapshot/master_snapshot_codec.cpp @@ -0,0 +1,150 @@ +#include "ha/snapshot/master_snapshot_codec.h" + +#include +#include +#include + +#include "master_service.h" +#include "segment.h" +#include "serialize/serializer.h" +#include "task_manager.h" + +namespace mooncake::ha { + +std::vector MasterSnapshotCodec::EncodeManifest( + const std::string& type, const std::string& version, + const std::string& snapshot_id) { + std::string manifest = type + "|" + version + "|" + snapshot_id; + return std::vector(manifest.begin(), manifest.end()); +} + +tl::expected +MasterSnapshotCodec::Encode(MasterSnapshotStateView& state_view) const { + MasterSnapshotPayloads payloads; + + // 1. Encode metadata (shards, discarded replicas, replica_next_id) + auto metadata_result = EncodeMetadata(state_view.master_service); + if (!metadata_result) { + return tl::make_unexpected(metadata_result.error()); + } + payloads.metadata = std::move(metadata_result.value()); + + // 2. Encode segments (memory segments + NoF segments) + auto segments_result = EncodeSegments(state_view.segment_manager, + state_view.nof_segment_manager); + if (!segments_result) { + return tl::make_unexpected(segments_result.error()); + } + payloads.segments = std::move(segments_result.value()); + + // 3. Encode task manager + auto task_manager_result = EncodeTaskManager(state_view.task_manager); + if (!task_manager_result) { + return tl::make_unexpected(task_manager_result.error()); + } + payloads.task_manager = std::move(task_manager_result.value()); + + return payloads; +} + +tl::expected MasterSnapshotCodec::Decode( + MasterService* master_service, + const MasterSnapshotPayloads& payloads) const { + if (master_service == nullptr) { + return tl::make_unexpected(SerializationError( + ErrorCode::INVALID_PARAMS, "master_service is null")); + } + + // Decode() is the codec-level exception boundary for restore. Most + // serializer failures are already reported as SerializationError, but a + // few MessagePack conversions (e.g. TaskManagerSerializer::Deserialize() + // calling arr[0].as() on a structurally valid but + // wrongly-typed field) can still throw msgpack::type_error outside their + // local try blocks. Since the caller RestoreState() no longer wraps each + // candidate in a try/catch, any escaping exception would abort restore and + // prevent fallback to an older healthy snapshot. Converting all exceptions + // here into SerializationError preserves that per-candidate fallback. + try { + // 1. Decode segments first. A MEMORY replica's allocator is bound to + // its mounted segment, so the segment/allocator must be restored + // before metadata; otherwise GetMountedSegment() returns + // SEGMENT_NOT_FOUND while deserializing the replica. + auto segments_result = + DecodeSegments(master_service, payloads.segments); + if (!segments_result) { + return tl::make_unexpected(segments_result.error()); + } + + // 2. Decode metadata (shards, discarded replicas, replica_next_id) + auto metadata_result = + DecodeMetadata(master_service, payloads.metadata); + if (!metadata_result) { + return tl::make_unexpected(metadata_result.error()); + } + + // 3. Decode task manager + auto task_manager_result = + DecodeTaskManager(master_service, payloads.task_manager); + if (!task_manager_result) { + return tl::make_unexpected(task_manager_result.error()); + } + + return {}; + } catch (const std::exception& e) { + return tl::make_unexpected(SerializationError( + ErrorCode::DESERIALIZE_FAIL, + std::string("exception during snapshot decode: ") + e.what())); + } catch (...) { + return tl::make_unexpected( + SerializationError(ErrorCode::DESERIALIZE_FAIL, + "unknown exception during snapshot decode")); + } +} + +tl::expected, SerializationError> +MasterSnapshotCodec::EncodeMetadata(MasterService& master_service) const { + // Delegate to the existing MetadataSerializer for now. + // This maintains the exact same format as before. + MasterService::MetadataSerializer serializer(&master_service); + return serializer.Serialize(); +} + +tl::expected MasterSnapshotCodec::DecodeMetadata( + MasterService* master_service, const std::vector& data) const { + // Delegate to the existing MetadataSerializer for now. + MasterService::MetadataSerializer serializer(master_service); + return serializer.Deserialize(data); +} + +tl::expected, SerializationError> +MasterSnapshotCodec::EncodeSegments( + SegmentManager& segment_manager, + NoFSegmentManager& nof_segment_manager) const { + // Use the existing SegmentSerializer which only handles SegmentManager + // Note: NoFSegmentManager is not currently serialized in snapshots + SegmentSerializer serializer(&segment_manager); + return serializer.Serialize(); +} + +tl::expected MasterSnapshotCodec::DecodeSegments( + MasterService* master_service, const std::vector& data) const { + // Access the segment managers from MasterService + SegmentSerializer serializer(&master_service->segment_manager_); + return serializer.Deserialize(data); +} + +tl::expected, SerializationError> +MasterSnapshotCodec::EncodeTaskManager(ClientTaskManager& task_manager) const { + // Use the existing TaskManagerSerializer + TaskManagerSerializer serializer(&task_manager); + return serializer.Serialize(); +} + +tl::expected MasterSnapshotCodec::DecodeTaskManager( + MasterService* master_service, const std::vector& data) const { + // Access the task manager from MasterService + TaskManagerSerializer serializer(&master_service->task_manager_); + return serializer.Deserialize(data); +} + +} // namespace mooncake::ha diff --git a/mooncake-store/src/ha/standby_controller.cpp b/mooncake-store/src/ha/standby_controller.cpp index 0849dc660b..6576cd9a21 100644 --- a/mooncake-store/src/ha/standby_controller.cpp +++ b/mooncake-store/src/ha/standby_controller.cpp @@ -31,6 +31,8 @@ std::unique_ptr CreateStandbyService( .enable_verification = false, .enable_snapshot_bootstrap = config.enable_snapshot_restore, .enable_oplog_following = capabilities.has_oplog_following, + .oplog_poll_interval_ms = config.oplog_poll_interval_ms, + .batch_oplog_retry_timeout_sec = config.batch_oplog_retry_timeout_sec, }); } @@ -38,7 +40,8 @@ StandbyRuntimeCapabilities BuildStandbyRuntimeCapabilities( const HABackendSpec& spec, const MasterServiceSupervisorConfig& config) { StandbyRuntimeCapabilities capabilities; capabilities.has_snapshot_bootstrap = config.enable_snapshot_restore; - capabilities.has_oplog_following = spec.type == HABackendType::ETCD; + capabilities.has_oplog_following = + config.enable_oplog && spec.type == HABackendType::ETCD; return capabilities; } @@ -78,6 +81,11 @@ class NoopStandbyController final : public StandbyController { ErrorCode PromoteStandby() override { return ErrorCode::OK; } + tl::expected PromoteStandbyAndExport() + override { + return tl::unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + } + void UpdateObservedLeader(const std::optional&) override {} MasterRuntimeState GetStandbyRuntimeState() const override { @@ -125,11 +133,22 @@ class CapabilityDrivenStandbyController final : public StandbyController { } standby_service_->SetSyncStatusCallback( - [this](const StandbySyncStatus&) { + [this](const StandbySyncStatus& status) { + if (status.state == StandbyState::FAILED || + status.state == StandbyState::STOPPED) { + std::lock_guard lock(state_mutex_); + standby_running_ = false; + last_standby_error_ = status.last_error; + } NotifyRuntimeStateIfChanged(); }); } + ~CapabilityDrivenStandbyController() override { + standby_service_->SetSyncStatusCallback({}); + standby_service_->Stop(); + } + ErrorCode StartStandby( const std::optional& observed_leader) override { bool standby_running = false; @@ -155,8 +174,17 @@ class CapabilityDrivenStandbyController final : public StandbyController { { std::lock_guard lock(state_mutex_); - standby_running_ = err == ErrorCode::OK; - last_standby_error_ = err; + const StandbySyncStatus status = standby_service_->GetSyncStatus(); + standby_running_ = err == ErrorCode::OK && + status.state != StandbyState::FAILED && + status.state != StandbyState::STOPPED; + if (standby_running_) { + last_standby_error_ = ErrorCode::OK; + } else if (status.last_error != ErrorCode::OK) { + last_standby_error_ = status.last_error; + } else { + last_standby_error_ = err; + } } if (err == ErrorCode::OK) { NotifyRuntimeStateIfChanged(); @@ -165,13 +193,6 @@ class CapabilityDrivenStandbyController final : public StandbyController { } void StopStandby() override { - { - std::lock_guard lock(state_mutex_); - if (!standby_running_) { - return; - } - } - standby_service_->Stop(); { std::lock_guard lock(state_mutex_); @@ -209,6 +230,50 @@ class CapabilityDrivenStandbyController final : public StandbyController { return err; } + tl::expected PromoteStandbyAndExport() + override { + // Verify standby is running first (same check as PromoteStandby) + ErrorCode promote_error = ErrorCode::OK; + { + std::lock_guard lock(state_mutex_); + if (!standby_running_) { + promote_error = last_standby_error_ != ErrorCode::OK + ? last_standby_error_ + : ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS; + } + } + if (promote_error != ErrorCode::OK) { + return tl::unexpected(promote_error); + } + + // Atomic promote + export (final catch-up happens inside) + StandbySnapshot snapshot; + ErrorCode err = standby_service_->PromoteAndExportSnapshot(snapshot); + if (err != ErrorCode::OK) { + { + std::lock_guard lock(state_mutex_); + standby_running_ = false; + last_standby_error_ = err; + } + NotifyRuntimeStateIfChanged(); + return tl::unexpected(err); + } + + { + std::lock_guard lock(state_mutex_); + standby_running_ = false; + last_standby_error_ = ErrorCode::OK; + } + NotifyRuntimeStateIfChanged(); + + PromotionContext ctx; + ctx.applied_seq_id = snapshot.oplog_sequence_id; + ctx.objects = std::move(snapshot.objects); + ctx.segments = std::move(snapshot.segments); + + return ctx; + } + void UpdateObservedLeader( const std::optional& observed_leader) override { { diff --git a/mooncake-store/src/ha_metric_manager.cpp b/mooncake-store/src/ha_metric_manager.cpp index 2819232071..ee18a5d414 100644 --- a/mooncake-store/src/ha_metric_manager.cpp +++ b/mooncake-store/src/ha_metric_manager.cpp @@ -77,6 +77,45 @@ HAMetricManager::HAMetricManager() "Latency of OpLog entry application in microseconds", {10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000}), + batch_record_durable_batches_total_( + "ha_batch_record_durable_batches_total", + "Total durable batch-record OpLog batches"), + batch_record_durable_entries_total_( + "ha_batch_record_durable_entries_total", + "Total durable entries in batch-record OpLog batches"), + batch_record_retry_total_("ha_batch_record_retry_total", + "Total batch-record backend retries"), + batch_record_committed_queue_depth_( + "ha_batch_record_committed_queue_depth", + "Committed batch-record entries waiting for durability"), + batch_record_callback_queue_depth_( + "ha_batch_record_callback_queue_depth", + "Durable batch-record callbacks waiting to run"), + batch_record_last_batch_id_("ha_batch_record_last_batch_id", + "Latest durable batch-record batch ID"), + batch_record_durable_sequence_( + "ha_batch_record_durable_sequence", + "Latest durable batch-record OpLog sequence"), + batch_record_batch_entries_("ha_batch_record_batch_entries", + "Entries per durable batch-record batch", + {1, 8, 32, 64, 128, 256, 512, 1024, 4096}), + batch_record_batch_bytes_( + "ha_batch_record_batch_bytes", + "Encoded bytes per durable batch-record batch", + {256, 1024, 4096, 16384, 65536, 262144, 1048576, 4194304}), + batch_record_txn_latency_us_( + "ha_batch_record_txn_latency_us", + "Batch-record backend transaction latency in microseconds", + {100, 500, 1000, 5000, 10000, 50000, 100000, 500000, 1000000}), + batch_record_commit_to_durable_us_( + "ha_batch_record_commit_to_durable_us", + "Batch-record commit-to-durable latency in microseconds", + {100, 500, 1000, 5000, 10000, 50000, 100000, 500000, 1000000}), + batch_record_callback_latency_us_( + "ha_batch_record_callback_latency_us", + "Batch-record durable-to-callback queue latency in microseconds", + {10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000}), + // State Machine standby_state_( "ha_standby_state", @@ -92,6 +131,12 @@ HAMetricManager::HAMetricManager() oplog_standby_lag_.update(0); oplog_pending_entries_.update(0); pending_mutation_queue_size_.update(0); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + batch_record_committed_queue_depth_.update(0); + batch_record_callback_queue_depth_.update(0); + batch_record_last_batch_id_.update(0); + batch_record_durable_sequence_.update(0); +#endif standby_state_.update(0); } @@ -237,6 +282,84 @@ void HAMetricManager::observe_oplog_apply_latency_us(int64_t latency_us) { oplog_apply_latency_us_.observe(latency_us); } +void HAMetricManager::inc_batch_record_durable_batches(int64_t val) { + batch_record_durable_batches_total_.inc(val); +} + +int64_t HAMetricManager::get_batch_record_durable_batches_total() { + return static_cast(batch_record_durable_batches_total_.value()); +} + +void HAMetricManager::inc_batch_record_durable_entries(int64_t val) { + batch_record_durable_entries_total_.inc(val); +} + +int64_t HAMetricManager::get_batch_record_durable_entries_total() { + return static_cast(batch_record_durable_entries_total_.value()); +} + +void HAMetricManager::inc_batch_record_retries(int64_t val) { + batch_record_retry_total_.inc(val); +} + +int64_t HAMetricManager::get_batch_record_retries_total() { + return static_cast(batch_record_retry_total_.value()); +} + +void HAMetricManager::set_batch_record_committed_queue_depth(int64_t depth) { + batch_record_committed_queue_depth_.update(depth); +} + +int64_t HAMetricManager::get_batch_record_committed_queue_depth() { + return static_cast(batch_record_committed_queue_depth_.value()); +} + +void HAMetricManager::set_batch_record_callback_queue_depth(int64_t depth) { + batch_record_callback_queue_depth_.update(depth); +} + +int64_t HAMetricManager::get_batch_record_callback_queue_depth() { + return static_cast(batch_record_callback_queue_depth_.value()); +} + +void HAMetricManager::set_batch_record_last_batch_id(int64_t batch_id) { + batch_record_last_batch_id_.update(batch_id); +} + +int64_t HAMetricManager::get_batch_record_last_batch_id() { + return static_cast(batch_record_last_batch_id_.value()); +} + +void HAMetricManager::set_batch_record_durable_sequence(int64_t sequence_id) { + batch_record_durable_sequence_.update(sequence_id); +} + +int64_t HAMetricManager::get_batch_record_durable_sequence() { + return static_cast(batch_record_durable_sequence_.value()); +} + +void HAMetricManager::observe_batch_record_batch_entries(int64_t entries) { + batch_record_batch_entries_.observe(entries); +} + +void HAMetricManager::observe_batch_record_batch_bytes(int64_t bytes) { + batch_record_batch_bytes_.observe(bytes); +} + +void HAMetricManager::observe_batch_record_txn_latency_us(int64_t latency_us) { + batch_record_txn_latency_us_.observe(latency_us); +} + +void HAMetricManager::observe_batch_record_commit_to_durable_us( + int64_t latency_us) { + batch_record_commit_to_durable_us_.observe(latency_us); +} + +void HAMetricManager::observe_batch_record_callback_latency_us( + int64_t latency_us) { + batch_record_callback_latency_us_.observe(latency_us); +} + // ========== State Machine Metrics ========== void HAMetricManager::set_standby_state(int64_t state_value) { @@ -273,6 +396,12 @@ std::string HAMetricManager::serialize_metrics() { serialize_metric(oplog_standby_lag_); serialize_metric(oplog_pending_entries_); serialize_metric(pending_mutation_queue_size_); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + serialize_metric(batch_record_committed_queue_depth_); + serialize_metric(batch_record_callback_queue_depth_); + serialize_metric(batch_record_last_batch_id_); + serialize_metric(batch_record_durable_sequence_); +#endif serialize_metric(standby_state_); // Counters @@ -284,11 +413,26 @@ std::string HAMetricManager::serialize_metrics() { serialize_metric(oplog_etcd_write_retries_total_); serialize_metric(oplog_watch_disconnections_total_); serialize_metric(oplog_applied_entries_total_); + serialize_metric(oplog_dropped_put_end_total_); + serialize_metric(oplog_batch_commits_total_); + serialize_metric(oplog_sync_batch_commits_total_); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + serialize_metric(batch_record_durable_batches_total_); + serialize_metric(batch_record_durable_entries_total_); + serialize_metric(batch_record_retry_total_); +#endif serialize_metric(state_transitions_total_); // Histograms serialize_metric(oplog_etcd_write_latency_us_); serialize_metric(oplog_apply_latency_us_); +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS + serialize_metric(batch_record_batch_entries_); + serialize_metric(batch_record_batch_bytes_); + serialize_metric(batch_record_txn_latency_us_); + serialize_metric(batch_record_commit_to_durable_us_); + serialize_metric(batch_record_callback_latency_us_); +#endif return ss.str(); } diff --git a/mooncake-store/src/hot_standby_service.cpp b/mooncake-store/src/hot_standby_service.cpp index b733075ba0..c31ae1e2e4 100644 --- a/mooncake-store/src/hot_standby_service.cpp +++ b/mooncake-store/src/hot_standby_service.cpp @@ -2,25 +2,24 @@ #include +#include #include #include #include "etcd_helper.h" +#include "ha/kv/etcd_ha_kv_backend.h" #include "ha_metric_manager.h" #include "ha/oplog/oplog_applier.h" -#include "ha/oplog/oplog_manager.h" -#include "ha/oplog/oplog_replicator.h" -#include "ha/oplog/oplog_store_factory.h" +#include "ha/oplog/oplog_batch_standby_reader.h" +#include "ha/oplog/oplog_test_failpoint.h" +#include "ha/oplog/oplog_types.h" namespace mooncake { HotStandbyService::HotStandbyService(const HotStandbyConfig& config) : config_(config) { metadata_store_ = std::make_unique(); - // OpLogApplier will be re-created in Start() with the resolved cluster_id - // to enable etcd-based operations (e.g. requesting missing OpLog entries). - // Here we construct a minimal instance so that local metadata operations - // are available before etcd wiring is completed. + // OpLogApplier is re-created in Start() with the resolved cluster_id. oplog_applier_ = std::make_unique(metadata_store_.get()); // Register callback for state change logging and metrics. @@ -49,54 +48,77 @@ HotStandbyService::HotStandbyService(const HotStandbyConfig& config) // StandbyMetadataStore implementation bool HotStandbyService::StandbyMetadataStore::PutMetadata( - const std::string& key, const StandbyObjectMetadata& metadata) { + const std::string& tenant_id, const std::string& key, + const StandbyObjectMetadata& metadata) { + const auto normalized = NormalizeTenantId(tenant_id); std::lock_guard lock(mutex_); - store_[key] = metadata; - VLOG(2) << "StandbyMetadataStore: stored metadata for key=" << key - << ", replicas=" << metadata.replicas.size() + store_[normalized][key] = metadata; + VLOG(2) << "StandbyMetadataStore: stored metadata for tenant=" << normalized + << ", key=" << key << ", replicas=" << metadata.replicas.size() << ", size=" << metadata.size; return true; } bool HotStandbyService::StandbyMetadataStore::Put(const std::string& key, const std::string& payload) { - // Legacy interface - create empty metadata + // Legacy interface - create empty metadata for default tenant StandbyObjectMetadata metadata; std::lock_guard lock(mutex_); - store_[key] = metadata; + store_["default"][key] = metadata; return true; } std::optional HotStandbyService::StandbyMetadataStore::GetMetadata( - const std::string& key) const { + const std::string& tenant_id, const std::string& key) const { + const auto normalized = NormalizeTenantId(tenant_id); std::lock_guard lock(mutex_); - auto it = store_.find(key); - if (it != store_.end()) { - return it->second; - } - return std::nullopt; + auto tenant_it = store_.find(normalized); + if (tenant_it == store_.end()) return std::nullopt; + auto it = tenant_it->second.find(key); + if (it == tenant_it->second.end()) return std::nullopt; + return it->second; } -bool HotStandbyService::StandbyMetadataStore::Remove(const std::string& key) { +bool HotStandbyService::StandbyMetadataStore::Remove( + const std::string& tenant_id, const std::string& key) { + const auto normalized = NormalizeTenantId(tenant_id); std::lock_guard lock(mutex_); - auto it = store_.find(key); - if (it != store_.end()) { - store_.erase(it); - return true; + auto tenant_it = store_.find(normalized); + if (tenant_it == store_.end()) return false; + auto it = tenant_it->second.find(key); + if (it == tenant_it->second.end()) return false; + tenant_it->second.erase(it); + if (tenant_it->second.empty()) { + store_.erase(tenant_it); } - return false; + return true; } bool HotStandbyService::StandbyMetadataStore::Exists( - const std::string& key) const { + const std::string& tenant_id, const std::string& key) const { + const auto normalized = NormalizeTenantId(tenant_id); std::lock_guard lock(mutex_); - return store_.find(key) != store_.end(); + auto tenant_it = store_.find(normalized); + if (tenant_it == store_.end()) return false; + return tenant_it->second.find(key) != tenant_it->second.end(); } size_t HotStandbyService::StandbyMetadataStore::GetKeyCount() const { std::lock_guard lock(mutex_); - return store_.size(); + size_t total = 0; + for (const auto& [tid, tenant_map] : store_) { + total += tenant_map.size(); + } + return total; +} + +size_t HotStandbyService::StandbyMetadataStore::GetKeyCountForTenant( + const std::string& tenant_id) const { + const auto normalized = NormalizeTenantId(tenant_id); + std::lock_guard lock(mutex_); + auto it = store_.find(normalized); + return it == store_.end() ? 0 : it->second.size(); } void HotStandbyService::StandbyMetadataStore::Clear() { @@ -105,12 +127,14 @@ void HotStandbyService::StandbyMetadataStore::Clear() { } void HotStandbyService::StandbyMetadataStore::Snapshot( - std::vector>& out) const { + std::vector& out) const { std::lock_guard lock(mutex_); out.clear(); - out.reserve(store_.size()); - for (const auto& kv : store_) { - out.emplace_back(kv.first, kv.second); + for (const auto& [tenant_id, tenant_store] : store_) { + for (const auto& [key, metadata] : tenant_store) { + out.push_back(StandbyObjectEntry{NormalizeTenantId(tenant_id), key, + metadata}); + } } } @@ -140,6 +164,21 @@ ErrorCode HotStandbyService::Start(const std::string& primary_address, return ErrorCode::OK; } + // A failed asynchronous run leaves joinable worker threads behind. Reap + // them before constructing the next run's readers and workers. + replication_loop_running_.store(false, std::memory_order_release); + replication_loop_cv_.notify_all(); + if (replication_thread_.joinable()) { + replication_thread_.join(); + } + if (verification_thread_.joinable()) { + verification_thread_.join(); + } + batch_standby_reader_.reset(); + batch_standby_kv_backend_.reset(); + + last_error_.store(ErrorCode::OK, std::memory_order_release); + // Trigger START event auto result = state_machine_.ProcessEvent(StandbyEvent::START); if (!result.allowed) { @@ -152,9 +191,8 @@ ErrorCode HotStandbyService::Start(const std::string& primary_address, cluster_id_ = cluster_id; if (config_.enable_oplog_following) { - // Connect to etcd only when using ETCD backend - if (config_.oplog_store_type == OpLogStoreType::ETCD) { #ifdef STORE_USE_ETCD + if (!catch_up_batch_kv_backend_for_testing_) { ErrorCode err = EtcdHelper::ConnectToEtcdStoreClient(oplog_endpoints.c_str()); if (err != ErrorCode::OK) { @@ -162,13 +200,12 @@ ErrorCode HotStandbyService::Start(const std::string& primary_address, state_machine_.ProcessEvent(StandbyEvent::CONNECTION_FAILED); return err; } + } #else - state_machine_.ProcessEvent(StandbyEvent::FATAL_ERROR); - LOG(ERROR) << "ETCD backend requested but STORE_USE_ETCD is not " - "enabled at compile time"; - return ErrorCode::INTERNAL_ERROR; + state_machine_.ProcessEvent(StandbyEvent::FATAL_ERROR); + LOG(ERROR) << "Batch-record OpLog requires STORE_USE_ETCD"; + return ErrorCode::INTERNAL_ERROR; #endif - } } state_machine_.ProcessEvent(StandbyEvent::CONNECTED); @@ -286,8 +323,13 @@ ErrorCode HotStandbyService::LoadSnapshotBaselineLocked( << snapshot.snapshot_id << ", snapshot_seq_id=" << snapshot.snapshot_sequence_id << ", keys=" << snapshot.metadata.size(); - for (const auto& kv : snapshot.metadata) { - metadata_store_->PutMetadata(kv.first, kv.second); + for (const auto& entry : snapshot.metadata) { + metadata_store_->PutMetadata(entry.tenant_id, entry.key, + entry.metadata); + } + // Load segment registry from snapshot + if (oplog_applier_) { + oplog_applier_->LoadSegmentRegistry(snapshot.segments); } oplog_applier_->Recover(snapshot.snapshot_sequence_id); baseline_seq_id = snapshot.snapshot_sequence_id; @@ -296,52 +338,17 @@ ErrorCode HotStandbyService::LoadSnapshotBaselineLocked( ErrorCode HotStandbyService::StartOplogFollowingLocked( uint64_t baseline_seq_id) { - // Create OpLogStore, OpLogChangeNotifier, and OpLogReplicator via factory - watcher_oplog_store_ = OpLogStoreFactory::Create( - config_.oplog_store_type, cluster_id_, OpLogStoreRole::READER, - config_.oplog_store_root_dir, config_.oplog_poll_interval_ms); - if (watcher_oplog_store_) { - // Wire OpLogStore into OpLogApplier so gap resolution works - oplog_applier_->SetOpLogStore(watcher_oplog_store_.get()); - oplog_change_notifier_ = - watcher_oplog_store_->CreateChangeNotifier(cluster_id_); - } - if (oplog_change_notifier_) { - oplog_replicator_ = std::make_unique( - oplog_change_notifier_.get(), oplog_applier_.get()); - oplog_replicator_->SetStateCallback( - [this](StandbyEvent event) { OnWatcherEvent(event); }); + (void)baseline_seq_id; + if (catch_up_batch_kv_backend_for_testing_) { + batch_standby_kv_backend_ = catch_up_batch_kv_backend_for_testing_; } else { - LOG(ERROR) << "Failed to create OpLogChangeNotifier for replicator"; - state_machine_.ProcessEvent(StandbyEvent::FATAL_ERROR); - return ErrorCode::INTERNAL_ERROR; - } - - static constexpr int kMaxStartRetries = 3; - static constexpr int kStartRetryBaseMs = 500; - bool watcher_started = false; - for (int attempt = 0; attempt < kMaxStartRetries; ++attempt) { - if (oplog_replicator_->StartFromSequenceId(baseline_seq_id)) { - watcher_started = true; - break; - } - LOG(WARNING) << "Failed to start OpLogReplicator from sequence_id=" - << baseline_seq_id << " (attempt " << (attempt + 1) << "/" - << kMaxStartRetries << ")"; - if (attempt + 1 < kMaxStartRetries) { - std::this_thread::sleep_for( - std::chrono::milliseconds(kStartRetryBaseMs * (1 << attempt))); - } - } - - if (!watcher_started) { - LOG(ERROR) << "Failed to start OpLogReplicator after " - << kMaxStartRetries << " attempts, aborting Start()"; - state_machine_.ProcessEvent(StandbyEvent::FATAL_ERROR); - return ErrorCode::INTERNAL_ERROR; + batch_standby_kv_backend_ = std::make_shared(); } + batch_standby_reader_ = std::make_unique( + cluster_id_, *batch_standby_kv_backend_, *oplog_applier_); state_machine_.ProcessEvent(StandbyEvent::SYNC_COMPLETE); + replication_loop_running_.store(true, std::memory_order_release); replication_thread_ = std::thread(&HotStandbyService::ReplicationLoop, this); if (config_.enable_verification) { @@ -382,8 +389,17 @@ void HotStandbyService::NotifySyncStatus() { } } -void HotStandbyService::OnWatcherEvent(StandbyEvent event) { - state_machine_.ProcessEvent(event); +void HotStandbyService::SetCatchUpBatchKvBackendForTesting( + std::shared_ptr backend) { + catch_up_batch_kv_backend_for_testing_ = std::move(backend); +} + +void HotStandbyService::StopReplicationLoop() { + replication_loop_running_.store(false, std::memory_order_release); + replication_loop_cv_.notify_all(); + if (replication_thread_.joinable()) { + replication_thread_.join(); + } } void HotStandbyService::Stop() { @@ -396,17 +412,9 @@ void HotStandbyService::Stop() { } state_machine_.ProcessEvent(StandbyEvent::STOP); - - // Stop OpLogReplicator - if (oplog_replicator_) { - oplog_replicator_->Stop(); - oplog_replicator_.reset(); - } + StopReplicationLoop(); // Wait for threads to finish - if (replication_thread_.joinable()) { - replication_thread_.join(); - } if (verification_thread_.joinable()) { verification_thread_.join(); } @@ -436,6 +444,7 @@ StandbySyncStatus HotStandbyService::GetSyncStatus() const { // Use state machine for connection status status.is_connected = IsConnected(); status.state = GetState(); + status.last_error = last_error_.load(std::memory_order_acquire); status.time_in_state = state_machine_.GetTimeInCurrentState(); if (status.primary_seq_id > status.applied_seq_id) { @@ -471,102 +480,85 @@ bool HotStandbyService::IsReadyForPromotion() const { << "will be synced after promotion."; } + // NOTE: unresolved-gap check is deferred to PromoteLockedInternal, which + // runs gap resolution + final catch-up first and only rejects promotion + // when gaps remain after both attempts. Checking here would return a + // misleading UNAVAILABLE_IN_CURRENT_STATUS before gap resolution runs. return true; } -void HotStandbyService::ResolvePromotionGapsLocked() { - if (!config_.enable_oplog_following || !oplog_applier_) { - return; - } - - static constexpr int kMaxGapResolveRetries = 3; - for (int retry = 0; retry < kMaxGapResolveRetries; ++retry) { - auto res = oplog_applier_->TryResolveGapsOnceForPromotion( - /*max_ids=*/1024); - if (res.attempted == 0) { - return; - } - - LOG(INFO) << "Promotion gap resolve (attempt " << (retry + 1) << "/" - << kMaxGapResolveRetries << "): attempted=" << res.attempted - << ", fetched=" << res.fetched - << ", applied_deletes=" << res.applied_deletes; - if (res.fetched == res.attempted) { - return; - } - if (retry + 1 < kMaxGapResolveRetries) { - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - } - } -} - ErrorCode HotStandbyService::FinalCatchUpForPromotionLocked( uint64_t current_applied_seq_id) { + (void)current_applied_seq_id; if (!config_.enable_oplog_following) { LOG(INFO) << "Promotion does not require final OpLog catch-up"; return ErrorCode::OK; } - - LOG(INFO) << "Final catch-up sync before promotion..."; - auto catch_up_store = OpLogStoreFactory::Create( - config_.oplog_store_type, cluster_id_, OpLogStoreRole::READER, - config_.oplog_store_root_dir, config_.oplog_poll_interval_ms); - if (!catch_up_store) { - LOG(ERROR) << "Failed to create oplog_store for final catch-up"; + if (!oplog_applier_) { + LOG(ERROR) << "Final catch-up requires OpLogApplier"; return ErrorCode::INTERNAL_ERROR; } - static constexpr size_t kBatchSize = 1000; - static constexpr size_t kMaxCatchUpBatches = 100; - static constexpr auto kMaxCatchUpDuration = std::chrono::seconds(30); + if (catch_up_batch_kv_backend_for_testing_) { + return FinalCatchUpBatchRecordsLocked( + *catch_up_batch_kv_backend_for_testing_); + } - uint64_t read_from_seq = current_applied_seq_id; - auto catch_up_start = std::chrono::steady_clock::now(); - size_t total_applied = 0; - size_t batch_count = 0; + EtcdHaKvBackend batch_backend; + return FinalCatchUpBatchRecordsLocked(batch_backend); +} +ErrorCode HotStandbyService::FinalCatchUpBatchRecordsLocked( + HaKvBackend& backend) { + std::unique_ptr local_reader; + OpLogBatchStandbyReader* reader = batch_standby_reader_.get(); + if (reader == nullptr) { + local_reader = std::make_unique( + cluster_id_, backend, *oplog_applier_); + reader = local_reader.get(); + } + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(30); + const auto initial_retry_delay = + std::chrono::milliseconds(std::max(config_.oplog_poll_interval_ms, 1)); + const auto max_retry_delay = + std::max(initial_retry_delay, std::chrono::milliseconds(1000)); + auto retry_delay = initial_retry_delay; + auto wait_to_retry = [&] { + const auto now = std::chrono::steady_clock::now(); + if (now >= deadline) { + return false; + } + std::this_thread::sleep_for(std::min( + retry_delay, std::chrono::duration_cast( + deadline - now))); + retry_delay = std::min(retry_delay * 2, max_retry_delay); + return true; + }; for (;;) { - auto elapsed = std::chrono::steady_clock::now() - catch_up_start; - if (elapsed > kMaxCatchUpDuration) { - LOG(WARNING) << "Final catch-up: timeout after " - << std::chrono::duration_cast( - elapsed) - .count() - << "s. Proceeding with promotion. total_applied=" - << total_applied; - break; + if (std::chrono::steady_clock::now() >= deadline) { + return ErrorCode::INCOMPLETE_OPLOG_CATCH_UP; } - - if (batch_count >= kMaxCatchUpBatches) { - LOG(WARNING) << "Final catch-up: reached max batch limit (" - << kMaxCatchUpBatches - << "). Proceeding with promotion. total_applied=" - << total_applied; - break; + auto result = reader->PollOnce(); + if (result.error != ErrorCode::OK) { + if (result.disposition != + OpLogBatchStandbyPollDisposition::RETRYABLE || + !wait_to_retry()) { + return ErrorCode::INCOMPLETE_OPLOG_CATCH_UP; + } + continue; } - - std::vector batch; - ErrorCode read_err = - catch_up_store->ReadOpLogSince(read_from_seq, kBatchSize, batch); - if (read_err != ErrorCode::OK) { - LOG(WARNING) << "Final catch-up: failed to read OpLog since seq=" - << read_from_seq - << ", err=" << static_cast(read_err) - << ". Proceeding with promotion."; - break; + retry_delay = initial_retry_delay; + if (!result.durable_prefix_present) { + return GetLocalLastAppliedSequenceIdLocked() == 0 + ? ErrorCode::OK + : ErrorCode::INCOMPLETE_OPLOG_CATCH_UP; } - if (batch.empty()) { - break; + if (GetLocalLastAppliedSequenceIdLocked() >= + result.durable_prefix.last_seq) { + return ErrorCode::OK; } - - total_applied += oplog_applier_->ApplyOpLogEntries(batch); - read_from_seq = batch.back().sequence_id; - ++batch_count; } - - LOG(INFO) << "Final catch-up sync done. total_applied=" << total_applied - << ", batches=" << batch_count; - return ErrorCode::OK; } ErrorCode HotStandbyService::Promote() { @@ -592,18 +584,33 @@ ErrorCode HotStandbyService::Promote() { << current_applied_seq_id << ", lag: " << status.lag_entries << " entries" << ", state: " << StandbyStateToString(GetState()); - if (oplog_replicator_) { - oplog_replicator_->Stop(); + auto internal_err = PromoteLockedInternal(current_applied_seq_id); + if (internal_err != ErrorCode::OK) { + return internal_err; } - ResolvePromotionGapsLocked(); + lock.unlock(); + Stop(); + + if (config_.enable_oplog_following) { + LOG(INFO) << "Standby promoted to Primary successfully. " + << "All remaining OpLog entries have been synced."; + } else { + LOG(INFO) << "Standby promoted to Primary from snapshot baseline."; + } + return ErrorCode::OK; +} - auto catch_up_err = FinalCatchUpForPromotionLocked(current_applied_seq_id); +ErrorCode HotStandbyService::PromoteLockedInternal( + uint64_t current_applied_seq_id) { + StopReplicationLoop(); + ErrorCode catch_up_err = + FinalCatchUpForPromotionLocked(current_applied_seq_id); if (catch_up_err != ErrorCode::OK) { state_machine_.ProcessEvent(StandbyEvent::PROMOTION_FAILED); return catch_up_err; } - + TestFailPoint::Wait("promotion_final_catch_up_before_complete"); uint64_t latest_applied_seq_id = GetLocalLastAppliedSequenceIdLocked(); applied_seq_id_.store(latest_applied_seq_id, std::memory_order_release); primary_seq_id_.store(latest_applied_seq_id, std::memory_order_release); @@ -614,6 +621,51 @@ ErrorCode HotStandbyService::Promote() { LOG(ERROR) << "Cannot finish promotion: " << promotion_success.reason; return ErrorCode::INTERNAL_ERROR; } + return ErrorCode::OK; +} + +ErrorCode HotStandbyService::PromoteAndExportSnapshot(StandbySnapshot& out) { + std::unique_lock lock(mutex_); + + if (!IsReadyForPromotion()) { + LOG(ERROR) << "Standby is not ready for promotion, state=" + << StandbyStateToString(GetState()); + return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS; + } + + // Trigger PROMOTE event + auto result = state_machine_.ProcessEvent(StandbyEvent::PROMOTE); + if (!result.allowed) { + LOG(ERROR) << "Cannot promote: " << result.reason; + return ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS; + } + + StandbySyncStatus status = GetSyncStatus(); + uint64_t current_applied_seq_id = status.applied_seq_id; + + LOG(INFO) << "Promoting Standby to Primary. Applied seq_id: " + << current_applied_seq_id << ", lag: " << status.lag_entries + << " entries" + << ", state: " << StandbyStateToString(GetState()); + + auto internal_err = PromoteLockedInternal(current_applied_seq_id); + if (internal_err != ErrorCode::OK) { + return internal_err; + } + + // Export snapshot BEFORE unlocking mutex (atomic promotion + export) + uint64_t latest_applied_seq_id = GetLocalLastAppliedSequenceIdLocked(); + out.oplog_sequence_id = latest_applied_seq_id; + if (metadata_store_) { + metadata_store_->Snapshot(out.objects); + } else { + out.objects.clear(); + } + if (oplog_applier_) { + out.segments = oplog_applier_->GetSegmentRegistry().GetAllSegments(); + } else { + out.segments.clear(); + } lock.unlock(); Stop(); @@ -644,7 +696,7 @@ uint64_t HotStandbyService::GetLatestAppliedSequenceId() const { } bool HotStandbyService::ExportMetadataSnapshot( - std::vector>& out) const { + std::vector& out) const { std::lock_guard lock(mutex_); if (!metadata_store_) { out.clear(); @@ -654,6 +706,37 @@ bool HotStandbyService::ExportMetadataSnapshot( return true; } +bool HotStandbyService::ExportStandbySnapshot(StandbySnapshot& out) const { + std::lock_guard lock(mutex_); + if (!IsRunning()) { + return false; + } + + // Get applied sequence ID (inline to avoid recursive mutex lock) + if (oplog_applier_) { + uint64_t expected_seq = oplog_applier_->GetExpectedSequenceId(); + out.oplog_sequence_id = expected_seq > 0 ? expected_seq - 1 : 0; + } else { + out.oplog_sequence_id = applied_seq_id_.load(); + } + + // Export object metadata + if (metadata_store_) { + metadata_store_->Snapshot(out.objects); + } else { + out.objects.clear(); + } + + // Export segments from OpLogApplier's registry (Patch B) + if (oplog_applier_) { + out.segments = oplog_applier_->GetSegmentRegistry().GetAllSegments(); + } else { + out.segments.clear(); + } + + return true; +} + void HotStandbyService::SetSnapshotProvider( std::unique_ptr provider) { std::lock_guard lock(mutex_); @@ -667,31 +750,85 @@ void HotStandbyService::SetSnapshotProvider( void HotStandbyService::ReplicationLoop() { LOG(INFO) << "Replication loop started (OpLog sync)"; - // OpLogReplicator handles the actual watching in its own thread. - // This loop monitors the status and updates metrics. - - // Create OpLogStore once before the loop to query primary sequence_id. - // Reuse watcher_oplog_store_ if available, otherwise create a new one. - std::shared_ptr repl_oplog_store = watcher_oplog_store_; - if (!repl_oplog_store && !cluster_id_.empty()) { - repl_oplog_store = OpLogStoreFactory::Create( - config_.oplog_store_type, cluster_id_, OpLogStoreRole::READER, - config_.oplog_store_root_dir, config_.oplog_poll_interval_ms); - if (!repl_oplog_store) { - LOG(ERROR) << "Failed to create oplog_store in replication loop"; - } - } - uint64_t last_reported_applied_seq_id = applied_seq_id_.load(); uint64_t last_reported_primary_seq_id = primary_seq_id_.load(); - - while (IsRunning()) { + const auto retry_base = + std::chrono::milliseconds(std::max(config_.oplog_poll_interval_ms, 1)); + const auto retry_limit = + std::chrono::seconds(config_.batch_oplog_retry_timeout_sec); + const auto max_retry_delay = + std::max(retry_base, std::chrono::milliseconds(5000)); + auto retry_delay = retry_base; + std::optional retry_started; + + auto wait_for_next_poll = [this](std::chrono::milliseconds delay) { + std::unique_lock lock(replication_loop_mutex_); + replication_loop_cv_.wait_for(lock, delay, [this] { + return !replication_loop_running_.load(std::memory_order_acquire); + }); + }; + + while (replication_loop_running_.load(std::memory_order_acquire) && + IsRunning()) { if (!IsConnected()) { - // Not connected - wait a bit before checking again - std::this_thread::sleep_for(std::chrono::seconds(1)); + wait_for_next_poll(std::chrono::seconds(1)); continue; } + if (batch_standby_reader_) { + const uint64_t expected_before = + oplog_applier_->GetExpectedSequenceId(); + auto result = batch_standby_reader_->PollOnce(); + if (result.durable_prefix_present) { + const uint64_t current_primary = primary_seq_id_.load(); + if (result.durable_prefix.last_seq > current_primary) { + primary_seq_id_.store(result.durable_prefix.last_seq); + } + } + + const uint64_t expected_after = + oplog_applier_->GetExpectedSequenceId(); + if (expected_after > 0) { + applied_seq_id_.store(expected_after - 1); + } + if (result.error != ErrorCode::OK) { + last_error_.store(result.error, std::memory_order_release); + const bool made_progress = expected_after > expected_before; + if (result.disposition == + OpLogBatchStandbyPollDisposition::RETRYABLE) { + const auto now = std::chrono::steady_clock::now(); + if (!retry_started || made_progress) { + retry_started = now; + retry_delay = retry_base; + } + if (now - *retry_started < retry_limit) { + LOG(WARNING) + << "Transient batch-record standby poll failure, " + << "retrying in " << retry_delay.count() + << " ms, err=" << static_cast(result.error); + NotifySyncStatus(); + wait_for_next_poll(retry_delay); + retry_delay = + std::min(retry_delay * 2, max_retry_delay); + continue; + } + LOG(ERROR) + << "Batch-record standby retry timeout after " + << config_.batch_oplog_retry_timeout_sec + << " seconds, err=" << static_cast(result.error); + } else { + LOG(ERROR) << "Fatal batch-record standby poll failure, " + << "err=" << static_cast(result.error); + } + state_machine_.ProcessEvent(StandbyEvent::FATAL_ERROR); + replication_loop_cv_.notify_all(); + break; + } + retry_started.reset(); + retry_delay = retry_base; + last_error_.store(ErrorCode::OK, std::memory_order_release); + } + // Update applied_seq_id from OpLogApplier if (oplog_applier_) { uint64_t expected = oplog_applier_->GetExpectedSequenceId(); @@ -701,17 +838,6 @@ void HotStandbyService::ReplicationLoop() { } } - // Update primary_seq_id by querying etcd `/latest` (best-effort). - // Note: `/latest` is batch-updated on Primary, so this is for - // monitoring only. - if (repl_oplog_store) { - uint64_t latest_seq = 0; - ErrorCode err = repl_oplog_store->GetLatestSequenceId(latest_seq); - if (err == ErrorCode::OK) { - primary_seq_id_.store(latest_seq); - } - } - const uint64_t applied_seq_id = applied_seq_id_.load(); const uint64_t primary_seq_id = primary_seq_id_.load(); if (applied_seq_id != last_reported_applied_seq_id || @@ -721,8 +847,8 @@ void HotStandbyService::ReplicationLoop() { NotifySyncStatus(); } - // Sleep and check again - std::this_thread::sleep_for(std::chrono::milliseconds(1000)); + wait_for_next_poll( + std::chrono::milliseconds(config_.oplog_poll_interval_ms)); } LOG(INFO) << "Replication loop stopped"; @@ -732,8 +858,20 @@ void HotStandbyService::VerificationLoop() { LOG(INFO) << "Verification loop started"; while (IsRunning()) { - std::this_thread::sleep_for( - std::chrono::seconds(config_.verification_interval_sec)); + std::unique_lock lock(replication_loop_mutex_); + replication_loop_cv_.wait_for( + lock, std::chrono::seconds(config_.verification_interval_sec), + [this] { + return !replication_loop_running_.load( + std::memory_order_acquire) || + !IsRunning(); + }); + lock.unlock(); + + if (!IsRunning() || + !replication_loop_running_.load(std::memory_order_acquire)) { + break; + } if (!IsConnected()) { continue; @@ -753,45 +891,4 @@ void HotStandbyService::VerificationLoop() { LOG(INFO) << "Verification loop stopped"; } -void HotStandbyService::ApplyOpLogEntry(const OpLogEntry& entry) { - // NOTE: This method is deprecated. OpLog entries are now applied via - // OpLogApplier, which is called by OpLogReplicator. This method is kept - // for backward compatibility but should not be used in the new etcd-based - // implementation. - - // Update applied_seq_id for status tracking - applied_seq_id_.store(entry.sequence_id); - - // The actual application is handled by OpLogApplier via OpLogReplicator - VLOG(2) << "ApplyOpLogEntry called (deprecated), sequence_id=" - << entry.sequence_id - << ", op_type=" << static_cast(entry.op_type) - << ", key=" << entry.object_key; -} - -void HotStandbyService::ProcessOpLogBatch( - const std::vector& entries) { - for (const auto& entry : entries) { - ApplyOpLogEntry(entry); - } -} - -bool HotStandbyService::ConnectToPrimary() { - // With etcd-based OpLog sync, connection is handled by OpLogReplicator - // This method is kept for compatibility but is no longer used - LOG(INFO) << "ConnectToPrimary called (no-op with etcd-based sync)"; - return true; -} - -void HotStandbyService::DisconnectFromPrimary() { - // With etcd-based OpLog sync, disconnection is handled by OpLogReplicator - // This method is kept for compatibility - if (IsConnected()) { - state_machine_.ProcessEvent(StandbyEvent::DISCONNECTED); - replication_stream_.reset(); - LOG(INFO) << "Disconnected from Primary (etcd-based sync), state=" - << StandbyStateToString(GetState()); - } -} - } // namespace mooncake diff --git a/mooncake-store/src/http_metadata_server.cpp b/mooncake-store/src/http_metadata_server.cpp index 72e152fcc3..446e7c9c4f 100644 --- a/mooncake-store/src/http_metadata_server.cpp +++ b/mooncake-store/src/http_metadata_server.cpp @@ -112,7 +112,16 @@ bool HttpMetadataServer::start() { return true; } - server_->async_start(); + // async_start() binds synchronously and hands back a future that is already + // resolved (hasResult()) when the bind failed; otherwise the server keeps + // running. Mirror MasterAdminServer::Start() so a failed bind is surfaced + // instead of reporting a healthy server that never came up. + auto ec = server_->async_start(); + if (ec.hasResult()) { + LOG(ERROR) << "Failed to start HTTP metadata server on " << host_ << ":" + << port_; + return false; + } running_ = true; LOG(INFO) << "HTTP metadata server started on " << host_ << ":" << port_; return true; @@ -135,4 +144,25 @@ KVPoll HttpMetadataServer::poll() const { return KVPoll::Success; } +bool HttpMetadataServer::removeKey(const std::string& key) { + std::lock_guard lock(store_mutex_); + if (store_.erase(key) > 0) { + LOG(INFO) << "HttpMetadataServer: removed key=" << key; + return true; + } + return false; +} + +size_t HttpMetadataServer::removeKeys(const std::vector& keys) { + std::lock_guard lock(store_mutex_); + size_t removed = 0; + for (const auto& key : keys) { + if (store_.erase(key) > 0) { + LOG(INFO) << "HttpMetadataServer: removed key=" << key; + ++removed; + } + } + return removed; +} + } // namespace mooncake diff --git a/mooncake-store/src/k8s_lease_helper.cpp b/mooncake-store/src/k8s_lease_helper.cpp index 15323bbc81..301fc31de6 100644 --- a/mooncake-store/src/k8s_lease_helper.cpp +++ b/mooncake-store/src/k8s_lease_helper.cpp @@ -149,6 +149,38 @@ ErrorCode K8sLeaseHelper::CancelWatch(const std::string& ns, return ErrorCode::OK; } +ErrorCode K8sLeaseHelper::SetPodLabel(const std::string& ns, + const std::string& pod, + const std::string& key, + const std::string& value) { + char* err_msg = nullptr; + int ret = K8sPatchPodLabel(const_cast(ns.c_str()), + const_cast(pod.c_str()), + const_cast(key.c_str()), + const_cast(value.c_str()), &err_msg); + if (ret != 0) { + LOG(ERROR) << "SetPodLabel failed: " << err_msg; + free(err_msg); + return ErrorCode::K8S_LEASE_OPERATION_ERROR; + } + return ErrorCode::OK; +} + +ErrorCode K8sLeaseHelper::ClearPodLabel(const std::string& ns, + const std::string& pod, + const std::string& key) { + char* err_msg = nullptr; + int ret = K8sRemovePodLabel(const_cast(ns.c_str()), + const_cast(pod.c_str()), + const_cast(key.c_str()), &err_msg); + if (ret != 0) { + LOG(ERROR) << "ClearPodLabel failed: " << err_msg; + free(err_msg); + return ErrorCode::K8S_LEASE_OPERATION_ERROR; + } + return ErrorCode::OK; +} + #else // !STORE_USE_K8S_LEASE ErrorCode K8sLeaseHelper::Init() { @@ -198,6 +230,18 @@ ErrorCode K8sLeaseHelper::CancelWatch(const std::string&, const std::string&) { return ErrorCode::K8S_LEASE_OPERATION_ERROR; } +ErrorCode K8sLeaseHelper::SetPodLabel(const std::string&, const std::string&, + const std::string&, const std::string&) { + LOG(ERROR) << "K8s Lease is not enabled in compilation"; + return ErrorCode::K8S_LEASE_OPERATION_ERROR; +} + +ErrorCode K8sLeaseHelper::ClearPodLabel(const std::string&, const std::string&, + const std::string&) { + LOG(ERROR) << "K8s Lease is not enabled in compilation"; + return ErrorCode::K8S_LEASE_OPERATION_ERROR; +} + #endif // STORE_USE_K8S_LEASE } // namespace mooncake diff --git a/mooncake-store/src/kv_event/kv_event_publisher.cpp b/mooncake-store/src/kv_event/kv_event_publisher.cpp new file mode 100644 index 0000000000..8349bfc7a1 --- /dev/null +++ b/mooncake-store/src/kv_event/kv_event_publisher.cpp @@ -0,0 +1,378 @@ +#include "kv_event/kv_event_publisher.h" + +#if defined(MOONCAKE_ENABLE_KV_EVENTS) && MOONCAKE_ENABLE_KV_EVENTS + +#include +#include +#include + +#include +#include +#include +#include + +namespace mooncake { +namespace { + +constexpr int kZmqSendHwm = 10000; +constexpr size_t kMaxBatchSize = 64; + +int64_t CurrentUnixTimeMs() { + return std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); +} + +void PackOptionalString(msgpack::packer& packer, + const std::string& value) { + if (value.empty()) { + packer.pack_nil(); + } else { + packer.pack(value); + } +} + +void PackOptionalU32(msgpack::packer& packer, uint32_t value, + bool has_value) { + if (!has_value) { + packer.pack_nil(); + } else { + packer.pack(value); + } +} + +size_t ComputeEventMapSize(bool is_stored, bool emit_legacy, + bool emit_object_key) { + // Base envelope: event_id, timestamp, event_type, model_name, block_size, + // additional_salt, lora_name, tenant_id, backend_id, medium, dp_rank, + // seq_hashes, group_id. + constexpr size_t kBaseFields = 13; + size_t map_size = kBaseFields; + if (emit_legacy) { + map_size += 2; // type, block_hashes + } + if (emit_object_key) { + map_size += 1; // object_key + } + if (is_stored) { + map_size += 3; // base_block_idx, parent_hash, token_ids + if (emit_legacy) { + map_size += 1; // parent_block_hash + } + } else { + map_size += 1; // base_block_idx + } + return map_size; +} + +} // namespace + +KvEventPublisher::KvEventPublisher(KvEventConfig config) + : config_(std::move(config)) { + if (!config_.enabled) { + return; + } + if (config_.bind_endpoint.empty()) { + LOG(ERROR) << "kv_events enabled but bind_endpoint is empty"; + config_.enabled = false; + return; + } + if (config_.backend_id.empty()) { + LOG(ERROR) << "kv_events enabled but backend_id is empty"; + config_.enabled = false; + return; + } + + zmq_context_ = zmq_ctx_new(); + if (!zmq_context_) { + LOG(ERROR) << "kv_events: failed to create ZMQ context"; + config_.enabled = false; + return; + } + zmq_socket_ = zmq_socket(zmq_context_, ZMQ_PUB); + if (!zmq_socket_) { + LOG(ERROR) << "kv_events: failed to create ZMQ PUB socket: " + << zmq_strerror(zmq_errno()); + zmq_ctx_destroy(zmq_context_); + zmq_context_ = nullptr; + config_.enabled = false; + return; + } + int hwm = kZmqSendHwm; + zmq_setsockopt(zmq_socket_, ZMQ_SNDHWM, &hwm, sizeof(hwm)); + int linger_ms = 0; + zmq_setsockopt(zmq_socket_, ZMQ_LINGER, &linger_ms, sizeof(linger_ms)); + + if (zmq_bind(zmq_socket_, config_.bind_endpoint.c_str()) != 0) { + LOG(ERROR) << "kv_events: zmq_bind failed for " << config_.bind_endpoint + << ": " << zmq_strerror(zmq_errno()); + zmq_close(zmq_socket_); + zmq_ctx_destroy(zmq_context_); + zmq_socket_ = nullptr; + zmq_context_ = nullptr; + config_.enabled = false; + return; + } + + worker_ = std::thread(&KvEventPublisher::WorkerLoop, this); + LOG(INFO) << "kv_events publisher enabled on " << config_.bind_endpoint + << " backend_id=" << config_.backend_id; +} + +KvEventPublisher::~KvEventPublisher() { + if (!config_.enabled) { + return; + } + stop_.store(true); + queue_cv_.notify_all(); + if (worker_.joinable()) { + worker_.join(); + } + if (zmq_socket_) { + zmq_close(zmq_socket_); + zmq_socket_ = nullptr; + } + if (zmq_context_) { + zmq_ctx_destroy(zmq_context_); + zmq_context_ = nullptr; + } +} + +void KvEventPublisher::PublishStored(const std::string& object_key, + const std::string& medium, + const TenantId& tenant_id, + const std::string& group_id) { + if (!config_.enabled) { + return; + } + Enqueue(PendingEvent{EventKind::kStored, object_key, medium, tenant_id, + group_id}); +} + +void KvEventPublisher::PublishRemoved(const std::string& object_key, + const std::string& medium, + const TenantId& tenant_id, + const std::string& group_id) { + if (!config_.enabled) { + return; + } + Enqueue(PendingEvent{EventKind::kRemoved, object_key, medium, tenant_id, + group_id}); +} + +KvEventPublisher::Stats KvEventPublisher::GetStats() const { + Stats stats; + stats.published_batches = published_batches_.load(); + stats.published_events = published_events_.load(); + stats.dropped_events = dropped_events_.load(); + stats.skipped_unparsed_keys = skipped_unparsed_keys_.load(); + return stats; +} + +void KvEventPublisher::Enqueue(PendingEvent event) { + { + std::lock_guard lock(queue_mutex_); + if (config_.queue_capacity > 0 && + queue_.size() >= config_.queue_capacity) { + queue_.pop_front(); + dropped_events_.fetch_add(1, std::memory_order_relaxed); + // Reserve a ZMQ sequence gap so consumers can detect loss. + next_zmq_sequence_.fetch_add(1, std::memory_order_relaxed); + } + queue_.push_back(std::move(event)); + } + queue_cv_.notify_one(); +} + +void KvEventPublisher::DrainRemainingQueue(std::vector& batch) { + while (true) { + batch.clear(); + { + std::lock_guard lock(queue_mutex_); + if (queue_.empty()) { + break; + } + while (!queue_.empty() && batch.size() < kMaxBatchSize) { + batch.push_back(std::move(queue_.front())); + queue_.pop_front(); + } + } + PublishBatch(batch); + } +} + +void KvEventPublisher::WorkerLoop() { + std::vector batch; + batch.reserve(kMaxBatchSize); + while (!stop_.load()) { + { + std::unique_lock lock(queue_mutex_); + queue_cv_.wait(lock, + [this] { return stop_.load() || !queue_.empty(); }); + while (!queue_.empty() && batch.size() < kMaxBatchSize) { + batch.push_back(std::move(queue_.front())); + queue_.pop_front(); + } + } + if (!batch.empty()) { + PublishBatch(batch); + batch.clear(); + } + } + DrainRemainingQueue(batch); +} + +void KvEventPublisher::PublishBatch(const std::vector& batch) { + struct EncodedEvent { + PendingEvent pending; + std::optional seq_hash; + uint64_t event_id{0}; + }; + std::vector encoded; + encoded.reserve(batch.size()); + for (const auto& pending : batch) { + const auto seq_hash = ParseSeqHashFromObjectKey(pending.object_key); + if (!seq_hash.has_value()) { + if (!config_.emit_object_key || pending.object_key.empty()) { + skipped_unparsed_keys_.fetch_add(1, std::memory_order_relaxed); + continue; + } + skipped_unparsed_keys_.fetch_add(1, std::memory_order_relaxed); + } + encoded.push_back(EncodedEvent{ + pending, seq_hash, + next_event_id_.fetch_add(1, std::memory_order_relaxed)}); + } + if (encoded.empty()) { + return; + } + + msgpack::sbuffer payload_buffer; + msgpack::packer packer(&payload_buffer); + + const int64_t timestamp_ms = CurrentUnixTimeMs(); + + packer.pack_array(3); + packer.pack(timestamp_ms); + + packer.pack_array(encoded.size()); + for (const auto& item : encoded) { + const bool is_stored = item.pending.kind == EventKind::kStored; + const char* rfc_type = is_stored ? "stored" : "removed"; + const char* legacy_type = is_stored ? "BlockStored" : "BlockRemoved"; + const std::string& tenant_id = item.pending.tenant_id.value(); + + const size_t map_size = + ComputeEventMapSize(is_stored, config_.emit_legacy_compat_fields, + config_.emit_object_key); + + packer.pack_map(map_size); + packer.pack("event_id"); + packer.pack(item.event_id); + packer.pack("timestamp"); + packer.pack(timestamp_ms); + packer.pack("event_type"); + packer.pack(rfc_type); + if (config_.emit_legacy_compat_fields) { + packer.pack("type"); + packer.pack(legacy_type); + } + // Per-block envelope fields unknown to the storage pool are omitted + // (nil). Indexer registration supplies model/block_size/dp_rank. + packer.pack("model_name"); + packer.pack_nil(); + packer.pack("block_size"); + packer.pack_nil(); + packer.pack("additional_salt"); + packer.pack_nil(); + packer.pack("lora_name"); + packer.pack_nil(); + packer.pack("tenant_id"); + packer.pack(tenant_id); + packer.pack("backend_id"); + packer.pack(config_.backend_id); + packer.pack("group_id"); + PackOptionalString(packer, item.pending.group_id); + packer.pack("medium"); + PackOptionalString(packer, item.pending.medium); + packer.pack("dp_rank"); + packer.pack_nil(); + + if (config_.emit_object_key) { + packer.pack("object_key"); + packer.pack(item.pending.object_key); + } + + packer.pack("seq_hashes"); + if (item.seq_hash.has_value()) { + packer.pack_array(1); + packer.pack(item.seq_hash.value()); + } else { + packer.pack_array(0); + } + + if (config_.emit_legacy_compat_fields && item.seq_hash.has_value()) { + packer.pack("block_hashes"); + packer.pack_array(1); + packer.pack(static_cast(item.seq_hash.value())); + } else if (config_.emit_legacy_compat_fields) { + packer.pack("block_hashes"); + packer.pack_array(0); + } + + if (is_stored) { + // Master keys are standalone pool blocks; depth 0 satisfies RFC + // #1527 requirement that base_block_idx or parent_hash be present. + packer.pack("base_block_idx"); + packer.pack(static_cast(0)); + packer.pack("parent_hash"); + packer.pack_nil(); + packer.pack("token_ids"); + packer.pack_nil(); + if (config_.emit_legacy_compat_fields) { + packer.pack("parent_block_hash"); + packer.pack_nil(); + } + } else { + packer.pack("base_block_idx"); + packer.pack_nil(); + } + } + + // Batch-level dp_rank; storage pool has no DP context (0). + packer.pack(static_cast(0)); + + const uint64_t seq = next_zmq_sequence_.fetch_add(1); + const uint64_t seq_be = htobe64(seq); + + zmq_msg_t topic_msg; + zmq_msg_t seq_msg; + zmq_msg_t payload_msg; + zmq_msg_init_size(&topic_msg, 0); + zmq_msg_init_size(&seq_msg, sizeof(seq_be)); + std::memcpy(zmq_msg_data(&seq_msg), &seq_be, sizeof(seq_be)); + zmq_msg_init_size(&payload_msg, payload_buffer.size()); + std::memcpy(zmq_msg_data(&payload_msg), payload_buffer.data(), + payload_buffer.size()); + + const int rc_topic = zmq_sendmsg(zmq_socket_, &topic_msg, ZMQ_SNDMORE); + const int rc_seq = + (rc_topic >= 0) ? zmq_sendmsg(zmq_socket_, &seq_msg, ZMQ_SNDMORE) : -1; + const int rc_payload = + (rc_seq >= 0) ? zmq_sendmsg(zmq_socket_, &payload_msg, 0) : -1; + + zmq_msg_close(&topic_msg); + zmq_msg_close(&seq_msg); + zmq_msg_close(&payload_msg); + + if (rc_topic >= 0 && rc_seq >= 0 && rc_payload >= 0) { + published_batches_.fetch_add(1, std::memory_order_relaxed); + published_events_.fetch_add(encoded.size(), std::memory_order_relaxed); + } else { + dropped_events_.fetch_add(encoded.size(), std::memory_order_relaxed); + } +} + +} // namespace mooncake + +#endif // MOONCAKE_ENABLE_KV_EVENTS diff --git a/mooncake-store/src/master.cpp b/mooncake-store/src/master.cpp index 700775c680..e4e9dec07c 100644 --- a/mooncake-store/src/master.cpp +++ b/mooncake-store/src/master.cpp @@ -4,8 +4,12 @@ #include // For std::atomic #include // For std::chrono #include -#include // For std::unique_ptr +#include // For std::getenv +#include // For std::ifstream +#include // For std::unique_ptr +#include #include // For std::thread +#include #include #include @@ -26,14 +30,14 @@ using namespace coro_rpc; using namespace async_simple; using namespace async_simple::coro; -static_assert(mooncake::DEFAULT_DEFAULT_KV_LEASE_TTL == 5000, +static_assert(mooncake::DEFAULT_DEFAULT_KV_LEASE_TTL == 10000, "Update kDefaultKvLeaseTtlFlagValue when " "DEFAULT_DEFAULT_KV_LEASE_TTL changes"); static_assert(mooncake::DEFAULT_KV_SOFT_PIN_TTL_MS == 30 * 60 * 1000, "Update kDefaultKvSoftPinTtlFlagValue when " "DEFAULT_KV_SOFT_PIN_TTL_MS changes"); -constexpr char kDefaultKvLeaseTtlFlagValue[] = "5000"; +constexpr char kDefaultKvLeaseTtlFlagValue[] = "10000"; constexpr char kDefaultKvSoftPinTtlFlagValue[] = "1800000"; namespace { @@ -60,13 +64,57 @@ uint64_t ParseDurationFlagOrDie(const char* flag_name, return parsed_value; } +// Derive the metadata server address for cleanup when it is deployed +// separately. Priority: MOONCAKE_TE_META_DATA_SERVER, then the +// "metadata_server" field of MOONCAKE_CONFIG_PATH. Returns "" if none found; +// the caller validates the scheme (only http(s) is supported). +std::string ResolveMetadataServerForCleanup() { + if (const char* env = std::getenv("MOONCAKE_TE_META_DATA_SERVER")) { + std::string value(env); + // P2PHANDSHAKE has no central metadata server, nothing to clean up. + if (!value.empty() && value != "P2PHANDSHAKE") { + return value; + } + } + + if (const char* cfg = std::getenv("MOONCAKE_CONFIG_PATH")) { + if (cfg[0] != '\0') { + try { + std::ifstream fin(cfg); + if (fin) { + Json::CharReaderBuilder builder; + Json::Value root; + std::string errs; + if (Json::parseFromStream(builder, fin, &root, &errs) && + root.isMember("metadata_server") && + root["metadata_server"].isString()) { + return root["metadata_server"].asString(); + } + if (!errs.empty()) { + LOG(WARNING) << "Failed to parse MOONCAKE_CONFIG_PATH (" + << cfg << "): " << errs; + } + } else { + LOG(WARNING) << "Cannot open MOONCAKE_CONFIG_PATH (" << cfg + << "): file does not exist or is not readable"; + } + } catch (const std::exception& e) { + LOG(WARNING) << "Error reading MOONCAKE_CONFIG_PATH (" << cfg + << "): " << e.what(); + } + } + } + + return {}; +} + } // namespace DEFINE_string(config_path, "", "master service config file path"); DEFINE_int32(port, 50051, "Port for master service to listen on (deprecated, use rpc_port)"); DEFINE_int32( - max_threads, 4, + max_threads, 16, "Maximum number of threads to use (deprecated, use rpc_thread_num)"); DEFINE_bool(enable_metric_reporting, true, "Enable periodic metric reporting"); DEFINE_int32(metrics_port, 9003, "Port for HTTP metrics server to listen on"); @@ -129,6 +177,40 @@ DEFINE_bool(offload_on_evict, false, "Defer LOCAL_DISK offload to eviction time instead of PutEnd"); DEFINE_bool(offload_force_evict, false, "Force-evict objects exceeding offload cap without disk offload"); +DEFINE_uint64(offloading_queue_limit, 50000, + "Maximum number of objects allowed in the offloading queue per " + "local disk segment. Increase to allow more objects to be " + "offloaded to SSD before force-eviction kicks in"); +DEFINE_validator(offloading_queue_limit, [](const char* flagname, + uint64_t value) { + // Zero would cause PushOffloadingQueue to always return + // KEYS_ULTRA_LIMIT, disabling offload entirely. The upper + // bound (1e8) keeps `offloading_queue_limit_ * + // offload_cap_ratio_` well within signed long range to + // avoid overflow when computing offload_cap in + // BatchEvict / EvictTenantMemoryForQuota. + if (value == 0) { + LOG(FATAL) << "offloading_queue_limit must be greater than 0"; + return false; + } + if (value > 100'000'000ULL) { + LOG(FATAL) << "offloading_queue_limit must be <= " + "100000000 to avoid overflow"; + return false; + } + return true; +}); +DEFINE_double(offload_cap_ratio, 0.5, + "Per-cycle offload cap as a fraction of offloading_queue_limit. " + "Controls how many objects can be queued for offload in a single " + "eviction cycle before falling back to force-evict"); +DEFINE_validator(offload_cap_ratio, [](const char* flagname, double value) { + if (value < 0.0 || value > 1.0) { + LOG(FATAL) << "offload_cap_ratio must be between 0.0 and 1.0"; + return false; + } + return true; +}); DEFINE_bool(promotion_on_hit, false, "Promote LOCAL_DISK-only keys to MEMORY on read access (mirror of " "offload_on_evict)"); @@ -143,6 +225,31 @@ DEFINE_uint32(promotion_max_per_heartbeat, 1, "SSD-read + RDMA-write on the client; serializing them avoids " "blocking past the client-liveness window. Default 1 is " "conservative."); +DEFINE_bool(enable_kv_events, false, + "Enable RFC #1527 KV cache event publisher over ZMQ"); +DEFINE_string(kv_events_bind_endpoint, "", + "ZMQ PUB bind endpoint for KV events, e.g. tcp://0.0.0.0:5557"); +DEFINE_string(kv_events_model_name, "", + "Deprecated: not emitted on events; use indexer POST /register"); +DEFINE_string(kv_events_backend_id, "", + "backend_id for published KV events (cache owner identity)"); +DEFINE_string(kv_events_tenant_id, "default", + "Deprecated: tenant_id comes from each object on events"); +DEFINE_string(kv_events_additional_salt, "", + "Deprecated: not emitted on events; use indexer POST /register"); +DEFINE_string(kv_events_lora_name, "", + "Deprecated: not emitted on events (no LoRA context in master)"); +DEFINE_uint32(kv_events_block_size, 0, + "Deprecated: not emitted on events; use indexer POST /register"); +DEFINE_uint32(kv_events_dp_rank, 0, + "Deprecated: not emitted on events; use indexer POST /register"); +DEFINE_bool(kv_events_emit_legacy_compat, true, + "Include vLLM/SGLang-compatible type/block_hashes fields"); +DEFINE_bool(kv_events_emit_object_key, true, + "Include Mooncake object_key in published KV events"); +DEFINE_uint32(kv_events_queue_capacity, 65536, + "Maximum pending KV events; oldest events are dropped when " + "the queue is full (0 = unbounded)"); DEFINE_string(ha_backend_type, "etcd", "HA backend type, e.g. etcd | redis | k8s"); DEFINE_string(ha_backend_connstring, "", @@ -177,17 +284,41 @@ DEFINE_string(cluster_id, mooncake::DEFAULT_CLUSTER_ID, "Cluster ID for the master service, used for kvcache persistence " "in HA mode"); +// OpLog store configuration +DEFINE_bool(enable_oplog, false, + "Enable HA metadata replication through batch-record OpLog"); +DEFINE_int32(oplog_poll_interval_ms, 1000, + "Batch-record standby poll interval."); +DEFINE_uint32(oplog_batch_max_entries, 1024, + "Maximum number of committed/reserved entries in the open " + "batch-record OpLog waiting batch."); +DEFINE_uint32(batch_oplog_retry_timeout_sec, 180, + "Maximum time to retry transient batch OpLog standby errors."); + DEFINE_string(memory_allocator, "offset", "Memory allocator for global segments, cachelib | offset"); DEFINE_string( allocation_strategy, "random", - "Allocation strategy for segments, random | free_ratio_first | cxl"); + "Allocation strategy for segments, random | free_ratio_first | cxl | " + "ssd_free_ratio_first | local_first"); DEFINE_bool(enable_http_metadata_server, false, "Enable HTTP metadata server instead of etcd"); DEFINE_int32(http_metadata_server_port, 8080, "Port for HTTP metadata server to listen on"); DEFINE_string(http_metadata_server_host, "0.0.0.0", "Host for HTTP metadata server to bind to"); +DEFINE_bool( + enable_metadata_cleanup_on_timeout, false, + "Enable cleanup of HTTP metadata (mooncake/ram/*, mooncake/rpc_meta/*) " + "when client heartbeat times out. Works in two modes: (1) co-located " + "(enable_http_metadata_server=true) via in-process removal, or " + "(2) separately-deployed metadata server via async HTTP DELETE."); + +DEFINE_string(pod_name, "", + "Pod name for K8s label-based routing (default: $POD_NAME)"); +DEFINE_string(pod_namespace, "", + "Pod namespace for K8s label-based routing " + "(default: $POD_NAMESPACE)"); DEFINE_uint64(put_start_discard_timeout_sec, mooncake::DEFAULT_PUT_START_DISCARD_TIMEOUT, @@ -201,13 +332,12 @@ DEFINE_bool(enable_disk_eviction, true, DEFINE_uint64( quota_bytes, 0, "Quota for storage backend in bytes (0 = use default 90% of capacity)"); -DEFINE_bool(enable_tenant_quota, false, - "Enable per-tenant memory quota admission"); -DEFINE_uint64(default_tenant_quota_bytes, 0, - "Default per-tenant memory quota in bytes (0 = unlimited)"); -DEFINE_uint64(tenant_quota_pool_capacity_bytes, 0, - "Capacity used to compute effective tenant quotas " - "(0 = mounted memory capacity)"); +DEFINE_bool(enable_multi_tenants, false, + "Enable strict multi-tenant namespace and quota admission"); +DEFINE_string(tenant_quota_connector_type, "file", + "Tenant quota policy connector type"); +DEFINE_string(tenant_quota_connector_uri, "", + "Tenant quota policy connector URI"); // Snapshot related configuration flags (migrated from global_flags) DEFINE_string(snapshot_backup_dir, "", @@ -366,6 +496,17 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config, default_config.GetBool("offload_force_evict", &master_config.offload_force_evict, FLAGS_offload_force_evict); + { + uint64_t tmp_offloading_queue_limit = FLAGS_offloading_queue_limit; + default_config.GetUInt64("offloading_queue_limit", + &tmp_offloading_queue_limit, + FLAGS_offloading_queue_limit); + master_config.offloading_queue_limit = + static_cast(tmp_offloading_queue_limit); + } + default_config.GetDouble("offload_cap_ratio", + &master_config.offload_cap_ratio, + FLAGS_offload_cap_ratio); default_config.GetBool("promotion_on_hit", &master_config.promotion_on_hit, FLAGS_promotion_on_hit); default_config.GetUInt32("promotion_admission_threshold", @@ -377,6 +518,41 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config, default_config.GetUInt32("promotion_max_per_heartbeat", &master_config.promotion_max_per_heartbeat, FLAGS_promotion_max_per_heartbeat); + default_config.GetBool("enable_kv_events", &master_config.enable_kv_events, + FLAGS_enable_kv_events); + default_config.GetString("kv_events_bind_endpoint", + &master_config.kv_events_bind_endpoint, + FLAGS_kv_events_bind_endpoint); + default_config.GetString("kv_events_model_name", + &master_config.kv_events_model_name, + FLAGS_kv_events_model_name); + default_config.GetString("kv_events_backend_id", + &master_config.kv_events_backend_id, + FLAGS_kv_events_backend_id); + default_config.GetString("kv_events_tenant_id", + &master_config.kv_events_tenant_id, + FLAGS_kv_events_tenant_id); + default_config.GetString("kv_events_additional_salt", + &master_config.kv_events_additional_salt, + FLAGS_kv_events_additional_salt); + default_config.GetString("kv_events_lora_name", + &master_config.kv_events_lora_name, + FLAGS_kv_events_lora_name); + default_config.GetUInt32("kv_events_block_size", + &master_config.kv_events_block_size, + FLAGS_kv_events_block_size); + default_config.GetUInt32("kv_events_dp_rank", + &master_config.kv_events_dp_rank, + FLAGS_kv_events_dp_rank); + default_config.GetBool("kv_events_emit_legacy_compat", + &master_config.kv_events_emit_legacy_compat, + FLAGS_kv_events_emit_legacy_compat); + default_config.GetBool("kv_events_emit_object_key", + &master_config.kv_events_emit_object_key, + FLAGS_kv_events_emit_object_key); + default_config.GetUInt32("kv_events_queue_capacity", + &master_config.kv_events_queue_capacity, + FLAGS_kv_events_queue_capacity); default_config.GetString("ha_backend_type", &master_config.ha_backend_type, FLAGS_ha_backend_type); default_config.GetString("ha_backend_connstring", @@ -386,6 +562,17 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config, FLAGS_etcd_endpoints); default_config.GetString("cluster_id", &master_config.cluster_id, FLAGS_cluster_id); + default_config.GetBool("enable_oplog", &master_config.enable_oplog, + FLAGS_enable_oplog); + default_config.GetInt32("oplog_poll_interval_ms", + &master_config.oplog_poll_interval_ms, + FLAGS_oplog_poll_interval_ms); + default_config.GetUInt32("oplog_batch_max_entries", + &master_config.oplog_batch_max_entries, + FLAGS_oplog_batch_max_entries); + default_config.GetUInt32("batch_oplog_retry_timeout_sec", + &master_config.batch_oplog_retry_timeout_sec, + FLAGS_batch_oplog_retry_timeout_sec); default_config.GetString("root_fs_dir", &master_config.root_fs_dir, FLAGS_root_fs_dir); default_config.GetInt64("global_file_segment_size", @@ -406,6 +593,13 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config, default_config.GetString("http_metadata_server_host", &master_config.http_metadata_server_host, FLAGS_http_metadata_server_host); + default_config.GetString("pod_name", &master_config.pod_name, + FLAGS_pod_name); + default_config.GetString("pod_namespace", &master_config.pod_namespace, + FLAGS_pod_namespace); + default_config.GetBool("enable_metadata_cleanup_on_timeout", + &master_config.enable_metadata_cleanup_on_timeout, + FLAGS_enable_metadata_cleanup_on_timeout); default_config.GetUInt64("put_start_discard_timeout_sec", &master_config.put_start_discard_timeout_sec, FLAGS_put_start_discard_timeout_sec); @@ -417,15 +611,15 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config, FLAGS_enable_disk_eviction); default_config.GetUInt64("quota_bytes", &master_config.quota_bytes, FLAGS_quota_bytes); - default_config.GetBool("enable_tenant_quota", - &master_config.enable_tenant_quota, - FLAGS_enable_tenant_quota); - default_config.GetUInt64("default_tenant_quota_bytes", - &master_config.default_tenant_quota_bytes, - FLAGS_default_tenant_quota_bytes); - default_config.GetUInt64("tenant_quota_pool_capacity_bytes", - &master_config.tenant_quota_pool_capacity_bytes, - FLAGS_tenant_quota_pool_capacity_bytes); + default_config.GetBool("enable_multi_tenants", + &master_config.enable_multi_tenants, + FLAGS_enable_multi_tenants); + default_config.GetString("tenant_quota_connector_type", + &master_config.tenant_quota_connector_type, + FLAGS_tenant_quota_connector_type); + default_config.GetString("tenant_quota_connector_uri", + &master_config.tenant_quota_connector_uri, + FLAGS_tenant_quota_connector_uri); default_config.GetString("snapshot_backup_dir", &master_config.snapshot_backup_dir, @@ -496,7 +690,7 @@ void InitMasterConf(const mooncake::DefaultConfig& default_config, void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, bool conf_set) { - if (FLAGS_max_threads != 4) { // 4 is the default value + if (FLAGS_max_threads != 16) { // 16 is the default value LOG(WARNING) << "max_threads is deprecated, use rpc_thread_num instead"; } if (FLAGS_port != 50051) { // 50051 is the default value @@ -512,7 +706,7 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, size_t rpc_thread_num; if (FLAGS_rpc_thread_num > 0) { rpc_thread_num = static_cast(FLAGS_rpc_thread_num); - if (FLAGS_max_threads != 4) { // 4 is the default value + if (FLAGS_max_threads != 16) { // 16 is the default value LOG(WARNING) << "Both rpc_thread_num and max_threads are set. " << "Using rpc_thread_num=" << FLAGS_rpc_thread_num << ". Please migrate to use rpc_thread_num only."; @@ -649,6 +843,17 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, !conf_set) { master_config.offload_force_evict = FLAGS_offload_force_evict; } + if ((google::GetCommandLineFlagInfo("offloading_queue_limit", &info) && + !info.is_default) || + !conf_set) { + master_config.offloading_queue_limit = + static_cast(FLAGS_offloading_queue_limit); + } + if ((google::GetCommandLineFlagInfo("offload_cap_ratio", &info) && + !info.is_default) || + !conf_set) { + master_config.offload_cap_ratio = FLAGS_offload_cap_ratio; + } if ((google::GetCommandLineFlagInfo("promotion_on_hit", &info) && !info.is_default) || !conf_set) { @@ -672,6 +877,70 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, master_config.promotion_max_per_heartbeat = FLAGS_promotion_max_per_heartbeat; } + if ((google::GetCommandLineFlagInfo("enable_kv_events", &info) && + !info.is_default) || + !conf_set) { + master_config.enable_kv_events = FLAGS_enable_kv_events; + } + if ((google::GetCommandLineFlagInfo("kv_events_bind_endpoint", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_bind_endpoint = FLAGS_kv_events_bind_endpoint; + } + if ((google::GetCommandLineFlagInfo("kv_events_model_name", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_model_name = FLAGS_kv_events_model_name; + } + if ((google::GetCommandLineFlagInfo("kv_events_backend_id", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_backend_id = FLAGS_kv_events_backend_id; + } + if ((google::GetCommandLineFlagInfo("kv_events_tenant_id", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_tenant_id = FLAGS_kv_events_tenant_id; + } + if ((google::GetCommandLineFlagInfo("kv_events_additional_salt", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_additional_salt = + FLAGS_kv_events_additional_salt; + } + if ((google::GetCommandLineFlagInfo("kv_events_lora_name", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_lora_name = FLAGS_kv_events_lora_name; + } + if ((google::GetCommandLineFlagInfo("kv_events_block_size", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_block_size = FLAGS_kv_events_block_size; + } + if ((google::GetCommandLineFlagInfo("kv_events_dp_rank", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_dp_rank = FLAGS_kv_events_dp_rank; + } + if ((google::GetCommandLineFlagInfo("kv_events_emit_legacy_compat", + &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_emit_legacy_compat = + FLAGS_kv_events_emit_legacy_compat; + } + if ((google::GetCommandLineFlagInfo("kv_events_emit_object_key", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_emit_object_key = + FLAGS_kv_events_emit_object_key; + } + if ((google::GetCommandLineFlagInfo("kv_events_queue_capacity", &info) && + !info.is_default) || + !conf_set) { + master_config.kv_events_queue_capacity = FLAGS_kv_events_queue_capacity; + } // Clamp promotion_admission_threshold into the sketch counter's // representable range. The CountMinSketch uses 8-bit saturating // counters (max 255) so any threshold beyond that would silently @@ -730,6 +999,28 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, !conf_set) { master_config.cluster_id = FLAGS_cluster_id; } + if ((google::GetCommandLineFlagInfo("enable_oplog", &info) && + !info.is_default) || + !conf_set) { + master_config.enable_oplog = FLAGS_enable_oplog; + } + if ((google::GetCommandLineFlagInfo("oplog_poll_interval_ms", &info) && + !info.is_default) || + !conf_set) { + master_config.oplog_poll_interval_ms = FLAGS_oplog_poll_interval_ms; + } + if ((google::GetCommandLineFlagInfo("oplog_batch_max_entries", &info) && + !info.is_default) || + !conf_set) { + master_config.oplog_batch_max_entries = FLAGS_oplog_batch_max_entries; + } + if ((google::GetCommandLineFlagInfo("batch_oplog_retry_timeout_sec", + &info) && + !info.is_default) || + !conf_set) { + master_config.batch_oplog_retry_timeout_sec = + FLAGS_batch_oplog_retry_timeout_sec; + } if ((google::GetCommandLineFlagInfo("root_fs_dir", &info) && !info.is_default) || !conf_set) { @@ -768,6 +1059,23 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, master_config.http_metadata_server_host = FLAGS_http_metadata_server_host; } + if ((google::GetCommandLineFlagInfo("pod_name", &info) && + !info.is_default) || + !conf_set) { + master_config.pod_name = FLAGS_pod_name; + } + if ((google::GetCommandLineFlagInfo("pod_namespace", &info) && + !info.is_default) || + !conf_set) { + master_config.pod_namespace = FLAGS_pod_namespace; + } + if ((google::GetCommandLineFlagInfo("enable_metadata_cleanup_on_timeout", + &info) && + !info.is_default) || + !conf_set) { + master_config.enable_metadata_cleanup_on_timeout = + FLAGS_enable_metadata_cleanup_on_timeout; + } if ((google::GetCommandLineFlagInfo("put_start_discard_timeout_sec", &info) && !info.is_default) || @@ -792,23 +1100,22 @@ void LoadConfigFromCmdline(mooncake::MasterConfig& master_config, !conf_set) { master_config.quota_bytes = FLAGS_quota_bytes; } - if ((google::GetCommandLineFlagInfo("enable_tenant_quota", &info) && + if ((google::GetCommandLineFlagInfo("enable_multi_tenants", &info) && !info.is_default) || !conf_set) { - master_config.enable_tenant_quota = FLAGS_enable_tenant_quota; + master_config.enable_multi_tenants = FLAGS_enable_multi_tenants; } - if ((google::GetCommandLineFlagInfo("default_tenant_quota_bytes", &info) && + if ((google::GetCommandLineFlagInfo("tenant_quota_connector_type", &info) && !info.is_default) || !conf_set) { - master_config.default_tenant_quota_bytes = - FLAGS_default_tenant_quota_bytes; + master_config.tenant_quota_connector_type = + FLAGS_tenant_quota_connector_type; } - if ((google::GetCommandLineFlagInfo("tenant_quota_pool_capacity_bytes", - &info) && + if ((google::GetCommandLineFlagInfo("tenant_quota_connector_uri", &info) && !info.is_default) || !conf_set) { - master_config.tenant_quota_pool_capacity_bytes = - FLAGS_tenant_quota_pool_capacity_bytes; + master_config.tenant_quota_connector_uri = + FLAGS_tenant_quota_connector_uri; } if ((google::GetCommandLineFlagInfo("max_total_finished_tasks", &info) && !info.is_default) || @@ -1019,6 +1326,16 @@ int main(int argc, char* argv[]) { LoadConfigFromCmdline(master_config, !conf_path.empty()); ResolveRpcAddressFromInterfaceOrDie(master_config); + // Fall back to environment variables for pod identity (K8s Downward API) + if (master_config.pod_name.empty()) { + const char* env = std::getenv("POD_NAME"); + if (env) master_config.pod_name = env; + } + if (master_config.pod_namespace.empty()) { + const char* env = std::getenv("POD_NAMESPACE"); + if (env) master_config.pod_namespace = env; + } + const std::string ha_backend_connstring = ResolveHABackendConnstring(master_config); if (master_config.enable_ha && ha_backend_connstring.empty()) { @@ -1029,6 +1346,14 @@ int main(int argc, char* argv[]) { << "etcd_endpoints"; return 1; } + if (master_config.enable_oplog && !master_config.enable_ha) { + LOG(FATAL) << "enable_oplog requires enable_ha=true"; + return 1; + } + if (master_config.enable_oplog && master_config.ha_backend_type != "etcd") { + LOG(FATAL) << "enable_oplog currently requires ha_backend_type=etcd"; + return 1; + } if (!master_config.enable_ha && (!ha_backend_connstring.empty() || !master_config.etcd_endpoints.empty())) { LOG(WARNING) @@ -1056,6 +1381,41 @@ int main(int argc, char* argv[]) { if (value && std::string_view(value) == "rdma") { protocol = "rdma"; } + + // enable_metadata_cleanup_on_timeout requires a reachable HTTP metadata + // server. Two topologies are supported: + // 1) Co-located: enable_http_metadata_server=true -> the master cleans + // up via the in-process server (no network overhead). + // 2) Separately deployed: the master derives the metadata server address + // from the cluster's existing configuration + // (MOONCAKE_TE_META_DATA_SERVER, or MOONCAKE_CONFIG_PATH json's + // "metadata_server") and cleans up via HTTP DELETE. Only http(s) + // endpoints are supported for now (etcd/redis left for future work). + // If neither is available, cleanup is disabled with a warning so the main + // process is never affected. + std::string http_metadata_remote_url; + if (master_config.enable_metadata_cleanup_on_timeout && + !master_config.enable_http_metadata_server) { + std::string derived = ResolveMetadataServerForCleanup(); + if (derived.rfind("http://", 0) == 0 || + derived.rfind("https://", 0) == 0) { + http_metadata_remote_url = std::move(derived); + LOG(INFO) << "enable_metadata_cleanup_on_timeout: HTTP metadata " + "server is deployed separately; cleanup will target " + << http_metadata_remote_url; + } else { + LOG(WARNING) + << "enable_metadata_cleanup_on_timeout is set to true but " + "enable_http_metadata_server is false and no HTTP metadata " + "server address could be derived from the cluster config " + "(set " + "MOONCAKE_TE_META_DATA_SERVER=http://host:port/metadata " + "or MOONCAKE_CONFIG_PATH). Disabling metadata cleanup on " + "timeout."; + master_config.enable_metadata_cleanup_on_timeout = false; + } + } + LOG(INFO) << "Master service started on port " << master_config.rpc_port << ", max_threads=" << master_config.rpc_thread_num @@ -1069,9 +1429,15 @@ int main(int argc, char* argv[]) { << ", eviction_high_watermark_ratio=" << master_config.eviction_high_watermark_ratio << ", enable_ha=" << master_config.enable_ha + << ", enable_oplog=" << master_config.enable_oplog << ", enable_offload=" << master_config.enable_offload + << ", enable_kv_events=" << master_config.enable_kv_events + << ", kv_events_bind_endpoint=" << master_config.kv_events_bind_endpoint + << ", kv_events_backend_id=" << master_config.kv_events_backend_id << ", offload_on_evict=" << master_config.offload_on_evict << ", offload_force_evict=" << master_config.offload_force_evict + << ", offloading_queue_limit=" << master_config.offloading_queue_limit + << ", offload_cap_ratio=" << master_config.offload_cap_ratio << ", ha_backend_type=" << master_config.ha_backend_type << ", ha_backend_connstring=" << ha_backend_connstring << ", etcd_endpoints=" << master_config.etcd_endpoints @@ -1095,6 +1461,8 @@ int main(int argc, char* argv[]) { << master_config.http_metadata_server_port << ", http_metadata_server_host=" << master_config.http_metadata_server_host + << ", enable_metadata_cleanup_on_timeout=" + << master_config.enable_metadata_cleanup_on_timeout << ", put_start_discard_timeout_sec=" << master_config.put_start_discard_timeout_sec << ", put_start_release_timeout_sec=" @@ -1140,9 +1508,20 @@ int main(int argc, char* argv[]) { std::this_thread::sleep_for(std::chrono::seconds(1)); } + // Metadata cleanup on client timeout (used by both the HA and non-HA + // paths): prefer the co-located in-process server, else the separate URL. + mooncake::HttpMetadataServer* metadata_server_ptr = nullptr; + if (master_config.enable_metadata_cleanup_on_timeout && + master_config.enable_http_metadata_server) { + metadata_server_ptr = http_metadata_server.get(); + } + if (master_config.enable_ha) { - mooncake::ha::MasterServiceSupervisor supervisor( - mooncake::MasterServiceSupervisorConfig{master_config}); + mooncake::MasterServiceSupervisorConfig supervisor_config{ + master_config}; + supervisor_config.http_metadata_server = metadata_server_ptr; + supervisor_config.http_metadata_remote_url = http_metadata_remote_url; + mooncake::ha::MasterServiceSupervisor supervisor(supervisor_config); return supervisor.Start(); } else { // version is not used in non-HA mode, just pass a dummy value @@ -1158,7 +1537,8 @@ int main(int argc, char* argv[]) { } auto wrapped_master_service = std::make_shared( - mooncake::WrappedMasterServiceConfig(master_config, version)); + mooncake::WrappedMasterServiceConfig(master_config, version), + metadata_server_ptr, http_metadata_remote_url); mooncake::MasterAdminServer admin_server( static_cast(master_config.metrics_port), master_config.enable_metric_reporting); diff --git a/mooncake-store/src/master_admin_service.cpp b/mooncake-store/src/master_admin_service.cpp index 834d06cf12..950f02fd60 100644 --- a/mooncake-store/src/master_admin_service.cpp +++ b/mooncake-store/src/master_admin_service.cpp @@ -71,7 +71,13 @@ coro_http::status_type ErrorCodeToHttpStatus(ErrorCode error) { return coro_http::status_type::bad_request; case ErrorCode::JOB_NOT_FOUND: case ErrorCode::SEGMENT_NOT_FOUND: + case ErrorCode::OBJECT_NOT_FOUND: + case ErrorCode::TENANT_NOT_REGISTERED: return coro_http::status_type::not_found; + case ErrorCode::UNAVAILABLE_IN_CURRENT_MODE: + case ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS: + case ErrorCode::TENANT_NOT_EMPTY: + return coro_http::status_type::conflict; default: return coro_http::status_type::internal_server_error; } @@ -143,6 +149,106 @@ std::string EscapeJson(std::string_view input) { return escaped; } +std::string EscapePrometheusLabel(std::string_view input) { + std::string escaped; + escaped.reserve(input.size()); + for (char ch : input) { + switch (ch) { + case '\\': + escaped += "\\\\"; + break; + case '"': + escaped += "\\\""; + break; + case '\n': + escaped += "\\n"; + break; + default: + escaped.push_back(ch); + break; + } + } + return escaped; +} + +struct HttpTenantQuotaSnapshot { + std::string tenant_id; + uint64_t requested_quota_bytes{0}; + uint64_t effective_quota_bytes{0}; + uint64_t used_bytes{0}; + uint64_t reserved_bytes{0}; + uint64_t committed_count{0}; + uint64_t metadata_object_count{0}; + bool over_quota{false}; + bool has_explicit_policy{false}; +}; +YLT_REFL(HttpTenantQuotaSnapshot, tenant_id, requested_quota_bytes, + effective_quota_bytes, used_bytes, reserved_bytes, committed_count, + metadata_object_count, over_quota, has_explicit_policy); + +HttpTenantQuotaSnapshot ToHttpTenantQuotaSnapshot( + const TenantQuotaSnapshot& snapshot) { + return HttpTenantQuotaSnapshot{ + .tenant_id = snapshot.tenant_id.value(), + .requested_quota_bytes = snapshot.requested_quota_bytes, + .effective_quota_bytes = snapshot.effective_quota_bytes, + .used_bytes = snapshot.used_bytes, + .reserved_bytes = snapshot.reserved_bytes, + .committed_count = snapshot.committed_count, + .metadata_object_count = snapshot.metadata_object_count, + .over_quota = snapshot.over_quota, + .has_explicit_policy = snapshot.has_explicit_policy, + }; +} + +struct HttpTenantQuotaListResponse { + bool success{true}; + std::vector data; +}; +YLT_REFL(HttpTenantQuotaListResponse, success, data); + +struct HttpTenantQuotaResponse { + bool success{true}; + HttpTenantQuotaSnapshot data; +}; +YLT_REFL(HttpTenantQuotaResponse, success, data); + +struct HttpTenantQuotaDeleteResponse { + bool success{true}; + std::optional data; +}; +YLT_REFL(HttpTenantQuotaDeleteResponse, success, data); + +struct HttpTenantQuotaPolicyRequest { + uint64_t requested_quota_bytes{0}; +}; +YLT_REFL(HttpTenantQuotaPolicyRequest, requested_quota_bytes); + +tl::expected ParseAdminTenantId( + coro_http::coro_http_request& req) { + auto tenant_id_view = req.get_decode_query_value("tenant_id"); + if (tenant_id_view.empty()) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + TenantId tenant_id{std::string(tenant_id_view)}; + if (!tenant_id.IsValid()) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + return tenant_id.value(); +} + +tl::expected ParseQuotaPolicyBody( + coro_http::coro_http_request& req) { + HttpTenantQuotaPolicyRequest request; + try { + struct_json::from_json(request, req.get_body()); + } catch (const std::exception& e) { + return tl::make_unexpected(std::string("Invalid JSON body: ") + + e.what()); + } + return request; +} + } // namespace MasterAdminServer::MasterAdminServer(uint16_t http_port, @@ -253,9 +359,98 @@ MasterAdminServer::RuntimeSnapshot MasterAdminServer::SnapshotState() const { } std::string MasterAdminServer::BuildMetricsText() const { - return AppendMetricSections( + std::string metrics = AppendMetricSections( MasterMetricManager::instance().serialize_metrics(), HAMetricManager::instance().serialize_metrics()); + auto tenant_metrics = BuildTenantQuotaMetricsText(); + if (tenant_metrics.empty()) { + return metrics; + } + return AppendMetricSections(std::move(metrics), std::move(tenant_metrics)); +} + +std::string MasterAdminServer::BuildTenantQuotaMetricsText() const { + auto service = GetActiveService(); + if (!service) { + return ""; + } + auto snapshots_result = service->ListTenantQuotaSnapshots(); + auto capacity_result = service->GetTenantQuotaAllocatableCapacityBytes(); + if (!snapshots_result || !capacity_result) { + return ""; + } + + const auto& snapshots = snapshots_result.value(); + uint64_t requested_sum = 0; + uint64_t effective_sum = 0; + std::ostringstream tenant_metrics; + tenant_metrics + << "# HELP mooncake_tenant_quota_requested_bytes Requested tenant " + "quota policy in bytes\n" + << "# TYPE mooncake_tenant_quota_requested_bytes gauge\n" + << "# HELP mooncake_tenant_quota_effective_bytes Effective tenant " + "quota in bytes\n" + << "# TYPE mooncake_tenant_quota_effective_bytes gauge\n" + << "# HELP mooncake_tenant_quota_used_bytes Tenant committed quota " + "usage in bytes\n" + << "# TYPE mooncake_tenant_quota_used_bytes gauge\n" + << "# HELP mooncake_tenant_quota_reserved_bytes Tenant reserved quota " + "usage in bytes\n" + << "# TYPE mooncake_tenant_quota_reserved_bytes gauge\n" + << "# HELP mooncake_tenant_quota_committed_count Tenant committed " + "object count\n" + << "# TYPE mooncake_tenant_quota_committed_count gauge\n" + << "# HELP mooncake_tenant_quota_metadata_object_count Tenant " + "metadata object count\n" + << "# TYPE mooncake_tenant_quota_metadata_object_count gauge\n" + << "# HELP mooncake_tenant_quota_over_quota Tenant over-quota flag\n" + << "# TYPE mooncake_tenant_quota_over_quota gauge\n" + << "# HELP mooncake_tenant_quota_explicit_policy Tenant explicit " + "policy flag\n" + << "# TYPE mooncake_tenant_quota_explicit_policy gauge\n"; + for (const auto& snapshot : snapshots) { + requested_sum += snapshot.requested_quota_bytes; + effective_sum += snapshot.effective_quota_bytes; + const auto tenant = EscapePrometheusLabel(snapshot.tenant_id.value()); + tenant_metrics << "mooncake_tenant_quota_requested_bytes{tenant_id=\"" + << tenant << "\"} " << snapshot.requested_quota_bytes + << "\n"; + tenant_metrics << "mooncake_tenant_quota_effective_bytes{tenant_id=\"" + << tenant << "\"} " << snapshot.effective_quota_bytes + << "\n"; + tenant_metrics << "mooncake_tenant_quota_used_bytes{tenant_id=\"" + << tenant << "\"} " << snapshot.used_bytes << "\n"; + tenant_metrics << "mooncake_tenant_quota_reserved_bytes{tenant_id=\"" + << tenant << "\"} " << snapshot.reserved_bytes << "\n"; + tenant_metrics << "mooncake_tenant_quota_committed_count{tenant_id=\"" + << tenant << "\"} " << snapshot.committed_count << "\n"; + tenant_metrics + << "mooncake_tenant_quota_metadata_object_count{tenant_id=\"" + << tenant << "\"} " << snapshot.metadata_object_count << "\n"; + tenant_metrics << "mooncake_tenant_quota_over_quota{tenant_id=\"" + << tenant << "\"} " << (snapshot.over_quota ? 1 : 0) + << "\n"; + tenant_metrics << "mooncake_tenant_quota_explicit_policy{tenant_id=\"" + << tenant << "\"} " + << (snapshot.has_explicit_policy ? 1 : 0) << "\n"; + } + tenant_metrics + << "# HELP mooncake_tenant_quota_allocatable_capacity_bytes Global " + "tenant quota allocatable capacity in bytes\n" + << "# TYPE mooncake_tenant_quota_allocatable_capacity_bytes gauge\n" + << "mooncake_tenant_quota_allocatable_capacity_bytes " + << capacity_result.value() << "\n" + << "# HELP mooncake_tenant_quota_requested_bytes_sum Global requested " + "tenant quota sum in bytes\n" + << "# TYPE mooncake_tenant_quota_requested_bytes_sum gauge\n" + << "mooncake_tenant_quota_requested_bytes_sum " << requested_sum << "\n" + << "# HELP mooncake_tenant_quota_effective_bytes_sum Global effective " + "tenant quota sum in bytes\n" + << "# TYPE mooncake_tenant_quota_effective_bytes_sum gauge\n" + << "mooncake_tenant_quota_effective_bytes_sum " << effective_sum + << "\n"; + + return tenant_metrics.str(); } std::string MasterAdminServer::BuildMetricsSummaryText() const { @@ -367,6 +562,31 @@ void MasterAdminServer::HandleHaStatus(coro_http::coro_http_request&, ha::MasterRuntimeStateToString(snapshot.state)); } +struct HttpKvEventsStatusResponse { + bool enabled{false}; + uint64_t published_batches{0}; + uint64_t published_events{0}; + uint64_t dropped_events{0}; + uint64_t skipped_unparsed_keys{0}; +}; +YLT_REFL(HttpKvEventsStatusResponse, enabled, published_batches, + published_events, dropped_events, skipped_unparsed_keys); + +void MasterAdminServer::HandleKvEventsStatus( + coro_http::coro_http_request&, coro_http::coro_http_response& resp) { + WithActiveService( + resp, [&](const std::shared_ptr& service) { + const auto stats = service->GetKvEventStats(); + HttpKvEventsStatusResponse payload; + payload.enabled = service->KvEventsEnabled(); + payload.published_batches = stats.published_batches; + payload.published_events = stats.published_events; + payload.dropped_events = stats.dropped_events; + payload.skipped_unparsed_keys = stats.skipped_unparsed_keys; + WriteJsonResponse(resp, coro_http::status_type::ok, payload); + }); +} + void MasterAdminServer::HandleQueryKey(coro_http::coro_http_request& req, coro_http::coro_http_response& resp) { WithActiveService(resp, [&](auto service) { @@ -725,12 +945,30 @@ void MasterAdminServer::HandleSegmentStatus( }); } +struct HttpDiskReplicaInfo { + std::string file_path; + uint64_t object_size = 0; + YLT_REFL(HttpDiskReplicaInfo, file_path, object_size); +}; + +struct HttpLocalDiskReplicaInfo { + std::string client_id; + uint64_t object_size = 0; + std::string transport_endpoint; + YLT_REFL(HttpLocalDiskReplicaInfo, client_id, object_size, + transport_endpoint); +}; + struct HttpBatchQueryKeyResult { bool ok{false}; std::optional error; std::optional> values; + std::optional> disk_values; + std::optional> local_disk_values; + std::optional> nof_values; }; -YLT_REFL(HttpBatchQueryKeyResult, ok, error, values); +YLT_REFL(HttpBatchQueryKeyResult, ok, error, values, disk_values, + local_disk_values, nof_values); struct HttpBatchQueryKeysResponse { bool success{false}; @@ -740,63 +978,210 @@ YLT_REFL(HttpBatchQueryKeysResponse, success, data); void MasterAdminServer::HandleBatchQueryKeys( coro_http::coro_http_request& req, coro_http::coro_http_response& resp) { - auto service = GetActiveService(); - if (!service) { - WriteSimpleErrorResponse(resp, - coro_http::status_type::service_unavailable, - "service plane is not active"); - return; - } + WithActiveService(resp, [&](auto service) { + auto keys_str = req.get_decode_query_value("keys"); + std::vector keys; + if (!keys_str.empty()) { + std::string_view sv(keys_str); + size_t pos = 0; + while ((pos = sv.find(',')) != std::string_view::npos) { + keys.emplace_back(sv.substr(0, pos)); + sv.remove_prefix(pos + 1); + } + keys.emplace_back(sv); + } - auto keys_str = req.get_query_value("keys"); - std::vector keys; - if (!keys_str.empty()) { - std::string_view sv(keys_str); - size_t pos = 0; - while ((pos = sv.find(',')) != std::string_view::npos) { - keys.emplace_back(sv.substr(0, pos)); - sv.remove_prefix(pos + 1); + if (keys.empty()) { + WriteSimpleErrorResponse( + resp, coro_http::status_type::bad_request, + "No keys provided. Use ?keys=key1,key2,..."); + return; } - keys.emplace_back(sv); - } - if (keys.empty()) { - WriteSimpleErrorResponse(resp, coro_http::status_type::bad_request, - "No keys provided. Use ?keys=key1,key2,..."); - return; - } + auto results = service->BatchGetReplicaListForAdmin(keys, "default"); + const size_t n = std::min(keys.size(), results.size()); + HttpBatchQueryKeysResponse payload; + payload.success = true; - auto results = service->BatchGetReplicaList(keys, "default"); - const size_t n = std::min(keys.size(), results.size()); - HttpBatchQueryKeysResponse payload; - payload.success = true; + for (size_t i = 0; i < n; ++i) { + const auto& result = results[i]; + HttpBatchQueryKeyResult item; + if (!result.has_value()) { + item.error = toString(result.error()); + payload.data.emplace(keys[i], std::move(item)); + continue; + } - for (size_t i = 0; i < n; ++i) { - const auto& result = results[i]; - HttpBatchQueryKeyResult item; - if (!result.has_value()) { - item.error = toString(result.error()); + item.ok = true; + item.values = std::vector{}; + for (const auto& replica : result.value().replicas) { + if (replica.is_memory_replica()) { + item.values->emplace_back( + replica.get_memory_descriptor().buffer_descriptor); + } else if (replica.is_disk_replica()) { + if (!item.disk_values.has_value()) { + item.disk_values = std::vector{}; + } + auto& d = replica.get_disk_descriptor(); + item.disk_values->emplace_back( + HttpDiskReplicaInfo{d.file_path, d.object_size}); + } else if (replica.is_local_disk_replica()) { + if (!item.local_disk_values.has_value()) { + item.local_disk_values = + std::vector{}; + } + auto& d = replica.get_local_disk_descriptor(); + item.local_disk_values->emplace_back( + HttpLocalDiskReplicaInfo{UuidToString(d.client_id), + d.object_size, + d.transport_endpoint}); + } else if (replica.is_nof_replica()) { + if (!item.nof_values.has_value()) { + item.nof_values = + std::vector{}; + } + item.nof_values->emplace_back( + replica.get_nof_descriptor().buffer_descriptor); + } + } payload.data.emplace(keys[i], std::move(item)); - continue; } - item.ok = true; - item.values = std::vector{}; - for (const auto& replica : result.value().replicas) { - if (!replica.is_memory_replica()) { - continue; + if (results.size() != keys.size()) { + LOG(WARNING) << "BatchGetReplicaListForAdmin size mismatch: keys=" + << keys.size() << " results=" << results.size(); + } + WriteJsonResponse(resp, coro_http::status_type::ok, payload); + }); +} + +void MasterAdminServer::HandleGetTenantQuotas( + coro_http::coro_http_request& req, coro_http::coro_http_response& resp) { + auto tenant_id_view = req.get_decode_query_value("tenant_id"); + WithActiveService(resp, [&](auto service) { + if (tenant_id_view.empty()) { + auto result = service->ListTenantQuotaSnapshots(); + if (!result.has_value()) { + WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()), + result.error()); + return; + } + HttpTenantQuotaListResponse payload; + payload.data.reserve(result->size()); + for (const auto& snapshot : result.value()) { + payload.data.push_back(ToHttpTenantQuotaSnapshot(snapshot)); } - item.values->emplace_back( - replica.get_memory_descriptor().buffer_descriptor); + WriteJsonResponse(resp, coro_http::status_type::ok, payload); + return; + } + + auto tenant_id_result = ParseAdminTenantId(req); + if (!tenant_id_result.has_value()) { + WriteErrorResponse(resp, coro_http::status_type::bad_request, + ErrorCode::INVALID_PARAMS, "Invalid tenant_id"); + return; + } + auto result = service->GetTenantQuotaSnapshot(tenant_id_result.value()); + if (!result.has_value()) { + WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()), + result.error()); + return; + } + WriteJsonResponse( + resp, coro_http::status_type::ok, + HttpTenantQuotaResponse{ + .data = ToHttpTenantQuotaSnapshot(result.value())}); + }); +} + +void MasterAdminServer::HandleUpsertTenantQuota( + coro_http::coro_http_request& req, coro_http::coro_http_response& resp) { + auto tenant_id_result = ParseAdminTenantId(req); + if (!tenant_id_result.has_value()) { + WriteErrorResponse(resp, coro_http::status_type::bad_request, + tenant_id_result.error(), + "Missing or invalid tenant_id"); + return; + } + auto body_result = ParseQuotaPolicyBody(req); + if (!body_result.has_value()) { + WriteErrorResponse(resp, coro_http::status_type::bad_request, + ErrorCode::INVALID_PARAMS, body_result.error()); + return; + } + if (body_result->requested_quota_bytes == 0) { + WriteErrorResponse(resp, coro_http::status_type::bad_request, + ErrorCode::INVALID_PARAMS, + "Tenant quota must be positive"); + return; + } + + WithActiveService(resp, [&](auto service) { + auto result = service->UpsertTenantQuotaPolicy( + tenant_id_result.value(), body_result->requested_quota_bytes); + if (!result.has_value()) { + WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()), + result.error()); + return; } - payload.data.emplace(keys[i], std::move(item)); + WriteJsonResponse( + resp, coro_http::status_type::ok, + HttpTenantQuotaResponse{ + .data = ToHttpTenantQuotaSnapshot(result.value())}); + }); +} + +void MasterAdminServer::HandleDeleteTenantQuota( + coro_http::coro_http_request& req, coro_http::coro_http_response& resp) { + auto tenant_id_result = ParseAdminTenantId(req); + if (!tenant_id_result.has_value()) { + WriteErrorResponse(resp, coro_http::status_type::bad_request, + tenant_id_result.error(), + "Missing or invalid tenant_id"); + return; } - if (results.size() != keys.size()) { - LOG(WARNING) << "BatchGetReplicaList size mismatch: keys=" - << keys.size() << " results=" << results.size(); + WithActiveService(resp, [&](auto service) { + auto result = + service->DeleteTenantQuotaPolicy(tenant_id_result.value()); + if (!result.has_value()) { + WriteErrorResponse(resp, ErrorCodeToHttpStatus(result.error()), + result.error()); + return; + } + HttpTenantQuotaDeleteResponse payload; + if (result.value().has_value()) { + payload.data = ToHttpTenantQuotaSnapshot(result.value().value()); + } + WriteJsonResponse(resp, coro_http::status_type::ok, payload); + }); +} + +struct HttpRemoveAllResponse { + bool success{true}; + long removed_count{0}; +}; +YLT_REFL(HttpRemoveAllResponse, success, removed_count); + +void MasterAdminServer::HandleRemoveAll(coro_http::coro_http_request& req, + coro_http::coro_http_response& resp) { + bool force = false; + if (auto it = req.get_query_value("force"); !it.empty()) { + force = (it == "true" || it == "1"); } - WriteJsonResponse(resp, coro_http::status_type::ok, payload); + std::string tenant_id; + if (auto it = req.get_query_value("tenant_id"); !it.empty()) { + tenant_id = std::string(it); + } + + WithActiveService(resp, [&](auto service) { + // Empty tenant_id => clear all tenants; pass "" so WrappedMasterService + // dispatches to the global (broadcast) RemoveAll, not the "default" + // tenant-scoped one. + long count = service->RemoveAll(force, tenant_id); + WriteJsonResponse(resp, coro_http::status_type::ok, + HttpRemoveAllResponse{.removed_count = count}); + }); } void MasterAdminServer::RegisterHandler() { @@ -823,6 +1208,11 @@ void MasterAdminServer::RegisterHandler() { "/ha_status", [this](coro_http_request& req, coro_http_response& resp) { HandleHaStatus(req, resp); }); + http_server_.set_http_handler( + "/kv_events/status", + [this](coro_http_request& req, coro_http_response& resp) { + HandleKvEventsStatus(req, resp); + }); http_server_.set_http_handler( "/leader", [this](coro_http_request& req, coro_http_response& resp) { HandleLeader(req, resp); @@ -872,10 +1262,30 @@ void MasterAdminServer::RegisterHandler() { [this](coro_http_request& req, coro_http_response& resp) { HandleSegmentStatus(req, resp); }); + http_server_.set_http_handler( + "/api/v1/tenant_quotas", + [this](coro_http_request& req, coro_http_response& resp) { + HandleGetTenantQuotas(req, resp); + }); + http_server_.set_http_handler( + "/api/v1/tenant_quotas", + [this](coro_http_request& req, coro_http_response& resp) { + HandleUpsertTenantQuota(req, resp); + }); + http_server_.set_http_handler( + "/api/v1/tenant_quotas", + [this](coro_http_request& req, coro_http_response& resp) { + HandleDeleteTenantQuota(req, resp); + }); http_server_.set_http_handler( "/batch_query_keys", [this](coro_http_request& req, coro_http_response& resp) { HandleBatchQueryKeys(req, resp); }); + http_server_.set_http_handler( + "/api/v1/remove_all", + [this](coro_http_request& req, coro_http_response& resp) { + HandleRemoveAll(req, resp); + }); } } // namespace mooncake diff --git a/mooncake-store/src/master_client.cpp b/mooncake-store/src/master_client.cpp index e2d9db347f..42fa13974b 100644 --- a/mooncake-store/src/master_client.cpp +++ b/mooncake-store/src/master_client.cpp @@ -317,6 +317,11 @@ struct RpcNameTraits<&WrappedMasterService::BatchEvictDiskReplica> { static constexpr const char* value = "BatchEvictDiskReplica"; }; +template <> +struct RpcNameTraits<&WrappedMasterService::PollRemoveAll> { + static constexpr const char* value = "PollRemoveAll"; +}; + template tl::expected MasterClient::invoke_rpc(Args&&... args) { auto pool = client_accessor_.GetClientPool(); @@ -419,13 +424,9 @@ ErrorCode MasterClient::Connect(const std::string& master_addr) { MutexLocker lock(&connect_mutex_); if (client_addr_param_ != master_addr) { - // WARNING: The existing client pool cannot be erased. So if there are a - // lot of different addresses, there will be resource leak problems. - auto client_pool = client_pools_->at(master_addr); - client_accessor_.SetClientPool(client_pool); + client_accessor_.GetOrCreateClientPool(master_addr); client_addr_param_ = master_addr; } - auto pool = client_accessor_.GetClientPool(); // The client pool does not have native connection check method, so we need // to use custom ServiceReady API. auto result = @@ -452,8 +453,8 @@ tl::expected MasterClient::ExistKey( ScopedVLogTimer timer(1, "MasterClient::ExistKey"); timer.LogRequest("object_key=", object_key); - auto result = invoke_rpc<&WrappedMasterService::ExistKey, bool>(object_key, - tenant_id_); + auto result = invoke_rpc<&WrappedMasterService::ExistKey, bool>( + object_key, tenant_id_.value()); timer.LogResponseExpected(result); return result; } @@ -464,7 +465,7 @@ std::vector> MasterClient::BatchExistKey( timer.LogRequest("keys_count=", object_keys.size()); auto result = invoke_batch_rpc<&WrappedMasterService::BatchExistKey, bool>( - object_keys.size(), object_keys, tenant_id_); + object_keys.size(), object_keys, tenant_id_.value()); timer.LogResponse("result=", result.size(), " keys"); return result; } @@ -515,7 +516,7 @@ MasterClient::GetReplicaListByRegex(const std::string& str) { auto result = invoke_rpc< &WrappedMasterService::GetReplicaListByRegex, std::unordered_map>>( - str, tenant_id_); + str, tenant_id_.value()); timer.LogResponseExpected(result); return result; @@ -523,7 +524,7 @@ MasterClient::GetReplicaListByRegex(const std::string& str) { tl::expected MasterClient::GetReplicaList( const std::string& object_key) { - return GetReplicaList(object_key, tenant_id_); + return GetReplicaList(object_key, tenant_id_.value()); } tl::expected MasterClient::GetReplicaList( @@ -539,7 +540,7 @@ tl::expected MasterClient::GetReplicaList( std::vector> MasterClient::BatchGetReplicaList(const std::vector& object_keys) { - return BatchGetReplicaList(object_keys, tenant_id_); + return BatchGetReplicaList(object_keys, tenant_id_.value()); } std::vector> @@ -570,7 +571,7 @@ MasterClient::PutStart(const std::string& key, auto result = invoke_rpc<&WrappedMasterService::PutStart, std::vector>( - client_id_, key, total_slice_length, config, tenant_id_); + client_id_, key, total_slice_length, config, tenant_id_.value()); timer.LogResponseExpected(result); return result; } @@ -595,29 +596,31 @@ MasterClient::BatchPutStart( auto result = invoke_batch_rpc<&WrappedMasterService::BatchPutStart, std::vector>( - keys.size(), client_id_, keys, total_slice_lengths, config, tenant_id_); + keys.size(), client_id_, keys, total_slice_lengths, config, + tenant_id_.value()); timer.LogResponse("result=", result.size(), " operations"); return result; } -tl::expected MasterClient::PutEnd(const std::string& key, - ReplicaType replica_type) { +tl::expected MasterClient::PutEnd( + const ObjectMeta& object_meta, ReplicaType replica_type) { ScopedVLogTimer timer(1, "MasterClient::PutEnd"); - timer.LogRequest("key=", key); + timer.LogRequest("key=", object_meta.key); auto result = invoke_rpc<&WrappedMasterService::PutEnd, void>( - client_id_, key, replica_type, tenant_id_); + client_id_, object_meta, replica_type, tenant_id_.value()); timer.LogResponseExpected(result); return result; } std::vector> MasterClient::BatchPutEnd( - const std::vector& keys, ReplicaType replica_type) { + const std::vector& object_metas, ReplicaType replica_type) { ScopedVLogTimer timer(1, "MasterClient::BatchPutEnd"); - timer.LogRequest("keys_count=", keys.size()); + timer.LogRequest("keys_count=", object_metas.size()); auto result = invoke_batch_rpc<&WrappedMasterService::BatchPutEnd, void>( - keys.size(), client_id_, keys, replica_type, tenant_id_); + object_metas.size(), client_id_, object_metas, replica_type, + tenant_id_.value()); timer.LogResponse("result=", result.size(), " operations"); return result; } @@ -628,7 +631,7 @@ tl::expected MasterClient::PutRevoke( timer.LogRequest("key=", key); auto result = invoke_rpc<&WrappedMasterService::PutRevoke, void>( - client_id_, key, replica_type, tenant_id_); + client_id_, key, replica_type, tenant_id_.value()); timer.LogResponseExpected(result); return result; } @@ -639,7 +642,7 @@ std::vector> MasterClient::BatchPutRevoke( timer.LogRequest("keys_count=", keys.size()); auto result = invoke_batch_rpc<&WrappedMasterService::BatchPutRevoke, void>( - keys.size(), client_id_, keys, replica_type, tenant_id_); + keys.size(), client_id_, keys, replica_type, tenant_id_.value()); timer.LogResponse("result=", result.size(), " operations"); return result; } @@ -658,7 +661,7 @@ MasterClient::UpsertStart(const std::string& key, auto result = invoke_rpc<&WrappedMasterService::UpsertStart, std::vector>( - client_id_, key, total_slice_length, config, tenant_id_); + client_id_, key, total_slice_length, config, tenant_id_.value()); timer.LogResponseExpected(result); return result; } @@ -683,29 +686,30 @@ MasterClient::BatchUpsertStart( auto result = invoke_batch_rpc<&WrappedMasterService::BatchUpsertStart, std::vector>( - keys.size(), client_id_, keys, total_slice_lengths, config, tenant_id_); + keys.size(), client_id_, keys, total_slice_lengths, config, + tenant_id_.value()); timer.LogResponse("result=", result.size(), " operations"); return result; } tl::expected MasterClient::UpsertEnd( - const std::string& key, ReplicaType replica_type) { + const ObjectMeta& object_meta, ReplicaType replica_type) { ScopedVLogTimer timer(1, "MasterClient::UpsertEnd"); - timer.LogRequest("key=", key); + timer.LogRequest("key=", object_meta.key); auto result = invoke_rpc<&WrappedMasterService::UpsertEnd, void>( - client_id_, key, replica_type, tenant_id_); + client_id_, object_meta, replica_type, tenant_id_.value()); timer.LogResponseExpected(result); return result; } std::vector> MasterClient::BatchUpsertEnd( - const std::vector& keys) { + const std::vector& object_metas) { ScopedVLogTimer timer(1, "MasterClient::BatchUpsertEnd"); - timer.LogRequest("keys_count=", keys.size()); + timer.LogRequest("keys_count=", object_metas.size()); auto result = invoke_batch_rpc<&WrappedMasterService::BatchUpsertEnd, void>( - keys.size(), client_id_, keys, tenant_id_); + object_metas.size(), client_id_, object_metas, tenant_id_.value()); timer.LogResponse("result=", result.size(), " operations"); return result; } @@ -716,7 +720,7 @@ tl::expected MasterClient::UpsertRevoke( timer.LogRequest("key=", key); auto result = invoke_rpc<&WrappedMasterService::UpsertRevoke, void>( - client_id_, key, replica_type, tenant_id_); + client_id_, key, replica_type, tenant_id_.value()); timer.LogResponseExpected(result); return result; } @@ -728,7 +732,7 @@ std::vector> MasterClient::BatchUpsertRevoke( auto result = invoke_batch_rpc<&WrappedMasterService::BatchUpsertRevoke, void>( - keys.size(), client_id_, keys, tenant_id_); + keys.size(), client_id_, keys, tenant_id_.value()); timer.LogResponse("result=", result.size(), " operations"); return result; } @@ -738,8 +742,8 @@ tl::expected MasterClient::Remove(const std::string& key, ScopedVLogTimer timer(1, "MasterClient::Remove"); timer.LogRequest("key=", key, ", force=", force); - auto result = - invoke_rpc<&WrappedMasterService::Remove, void>(key, force, tenant_id_); + auto result = invoke_rpc<&WrappedMasterService::Remove, void>( + key, force, tenant_id_.value()); timer.LogResponseExpected(result); return result; } @@ -750,7 +754,7 @@ tl::expected MasterClient::RemoveByRegex( timer.LogRequest("key=", str, ", force=", force); auto result = invoke_rpc<&WrappedMasterService::RemoveByRegex, long>( - str, force, tenant_id_); + str, force, tenant_id_.value()); timer.LogResponseExpected(result); return result; } @@ -759,8 +763,8 @@ tl::expected MasterClient::RemoveAll(bool force) { ScopedVLogTimer timer(1, "MasterClient::RemoveAll"); timer.LogRequest("action=remove_all_objects, force=", force); - auto result = - invoke_rpc<&WrappedMasterService::RemoveAll, long>(force, tenant_id_); + auto result = invoke_rpc<&WrappedMasterService::RemoveAll, long>( + force, tenant_id_.value()); timer.LogResponseExpected(result); return result; } @@ -771,7 +775,7 @@ std::vector> MasterClient::BatchRemove( timer.LogRequest("keys_count=", keys.size(), ", force=", force); auto result = invoke_batch_rpc<&WrappedMasterService::BatchRemove, void>( - keys.size(), keys, force, tenant_id_); + keys.size(), keys, force, tenant_id_.value()); timer.LogResponse("result=", result.size(), " operations"); return result; } @@ -940,7 +944,7 @@ tl::expected MasterClient::MountLocalDiskSegment( tl::expected MasterClient::CreateCopyTask( const std::string& key, const std::vector& targets) { - return CreateCopyTask(key, tenant_id_, targets); + return CreateCopyTask(key, tenant_id_.value(), targets); } tl::expected MasterClient::CreateCopyTask( @@ -959,7 +963,7 @@ tl::expected MasterClient::CreateCopyTask( tl::expected MasterClient::CreateMoveTask( const std::string& key, const std::string& source, const std::string& target) { - return CreateMoveTask(key, tenant_id_, source, target); + return CreateMoveTask(key, tenant_id_.value(), source, target); } tl::expected MasterClient::CreateMoveTask( @@ -988,6 +992,17 @@ MasterClient::OffloadObjectHeartbeat(const UUID& client_id, return result; } +tl::expected MasterClient::PollRemoveAll() { + ScopedVLogTimer timer(1, "MasterClient::PollRemoveAll"); + timer.LogRequest("client_id=", client_id_); + + auto result = + invoke_rpc<&WrappedMasterService::PollRemoveAll, bool>(client_id_); + timer.LogResponse("should_remove_all=", + result.has_value() ? result.value() : false); + return result; +} + tl::expected MasterClient::ReportSsdCapacity( const UUID& client_id, int64_t ssd_total_capacity_bytes) { ScopedVLogTimer timer(1, "MasterClient::ReportSsdCapacity"); @@ -1003,8 +1018,8 @@ tl::expected MasterClient::NotifyOffloadSuccess( std::vector tasks; tasks.reserve(keys.size()); for (const auto& key : keys) { - tasks.push_back( - OffloadTaskItem{.tenant_id = tenant_id_, .key = key, .size = 0}); + tasks.push_back(OffloadTaskItem{ + .tenant_id = tenant_id_.value(), .key = key, .size = 0}); } return NotifyOffloadSuccess(client_id, tasks, metadatas); } @@ -1034,7 +1049,7 @@ tl::expected MasterClient::PromotionAllocStart( const UUID& client_id, const std::string& key, uint64_t size, const std::vector& preferred_segments) { - return PromotionAllocStart(client_id, key, tenant_id_, size, + return PromotionAllocStart(client_id, key, tenant_id_.value(), size, preferred_segments); } @@ -1055,7 +1070,7 @@ MasterClient::PromotionAllocStart( tl::expected MasterClient::NotifyPromotionSuccess( const UUID& client_id, const std::string& key) { - return NotifyPromotionSuccess(client_id, key, tenant_id_); + return NotifyPromotionSuccess(client_id, key, tenant_id_.value()); } tl::expected MasterClient::NotifyPromotionSuccess( @@ -1073,7 +1088,7 @@ tl::expected MasterClient::NotifyPromotionSuccess( tl::expected MasterClient::NotifyPromotionFailure( const UUID& client_id, const std::string& key) { - return NotifyPromotionFailure(client_id, key, tenant_id_); + return NotifyPromotionFailure(client_id, key, tenant_id_.value()); } tl::expected MasterClient::NotifyPromotionFailure( @@ -1092,7 +1107,7 @@ tl::expected MasterClient::NotifyPromotionFailure( tl::expected MasterClient::CopyStart( const std::string& key, const std::string& src_segment, const std::vector& tgt_segments) { - return CopyStart(key, tenant_id_, src_segment, tgt_segments); + return CopyStart(key, tenant_id_.value(), src_segment, tgt_segments); } tl::expected MasterClient::CopyStart( @@ -1124,7 +1139,7 @@ tl::expected MasterClient::QueryTask( } tl::expected MasterClient::CopyEnd(const std::string& key) { - return CopyEnd(key, tenant_id_); + return CopyEnd(key, tenant_id_.value()); } tl::expected MasterClient::CopyEnd( @@ -1150,7 +1165,7 @@ tl::expected, ErrorCode> MasterClient::FetchTasks( } tl::expected MasterClient::CopyRevoke(const std::string& key) { - return CopyRevoke(key, tenant_id_); + return CopyRevoke(key, tenant_id_.value()); } tl::expected MasterClient::CopyRevoke( @@ -1167,7 +1182,7 @@ tl::expected MasterClient::CopyRevoke( tl::expected MasterClient::MoveStart( const std::string& key, const std::string& src_segment, const std::string& tgt_segment) { - return MoveStart(key, tenant_id_, src_segment, tgt_segment); + return MoveStart(key, tenant_id_.value(), src_segment, tgt_segment); } tl::expected MasterClient::MoveStart( @@ -1186,7 +1201,7 @@ tl::expected MasterClient::MoveStart( } tl::expected MasterClient::MoveEnd(const std::string& key) { - return MoveEnd(key, tenant_id_); + return MoveEnd(key, tenant_id_.value()); } tl::expected MasterClient::MoveEnd( @@ -1201,7 +1216,7 @@ tl::expected MasterClient::MoveEnd( } tl::expected MasterClient::MoveRevoke(const std::string& key) { - return MoveRevoke(key, tenant_id_); + return MoveRevoke(key, tenant_id_.value()); } tl::expected MasterClient::MoveRevoke( @@ -1227,7 +1242,7 @@ tl::expected MasterClient::MarkTaskToComplete( tl::expected MasterClient::EvictDiskReplica( const std::string& key, ReplicaType replica_type) { - return EvictDiskReplica(key, tenant_id_, replica_type); + return EvictDiskReplica(key, tenant_id_.value(), replica_type); } tl::expected MasterClient::EvictDiskReplica( @@ -1245,7 +1260,7 @@ tl::expected MasterClient::EvictDiskReplica( std::vector> MasterClient::BatchEvictDiskReplica( const std::vector& keys, ReplicaType replica_type) { - return BatchEvictDiskReplica(keys, tenant_id_, replica_type); + return BatchEvictDiskReplica(keys, tenant_id_.value(), replica_type); } std::vector> MasterClient::BatchEvictDiskReplica( diff --git a/mooncake-store/src/master_metric_manager.cpp b/mooncake-store/src/master_metric_manager.cpp index baa4ec53a0..8af2cb8597 100644 --- a/mooncake-store/src/master_metric_manager.cpp +++ b/mooncake-store/src/master_metric_manager.cpp @@ -71,6 +71,10 @@ MasterMetricManager::MasterMetricManager() "master_put_start_alloc_failures_total", "Total number of PutStart failures caused by replica allocation " "failure"), + put_start_partial_allocations_( + "master_put_start_partial_allocations_total", + "Total number of PutStart requests that succeeded with fewer " + "replicas than requested (best-effort degradation)"), put_end_requests_("master_put_end_requests_total", "Total number of PutEnd requests received"), put_end_failures_("master_put_end_failures_total", @@ -277,10 +281,16 @@ MasterMetricManager::MasterMetricManager() file_cache_hit_nums_("file_cache_hit_nums_", "Total number of GetReplicaList results served from " "the SSD cache"), + mem_cache_hit_bytes_("mem_cache_hit_bytes_total", + "Total bytes of GetReplicaList results served from " + "the memory pool"), + file_cache_hit_bytes_("file_cache_hit_bytes_total", + "Total bytes of GetReplicaList results served from " + "the SSD cache"), mem_cache_nums_("mem_cache_nums_", "Current number of cached values in the memory pool"), file_cache_nums_("file_cache_nums_", - "Current number of cached values in the SSD cache"), + "Current number of cached values in SSD cache"), valid_get_nums_("valid_get_nums_", "Total number of GetReplicaList operations that returned " "at least one completed replica"), @@ -368,6 +378,35 @@ MasterMetricManager::MasterMetricManager() "master_promotion_rejected_cap_total", "Promotion attempts rejected because promotion_in_flight was at " "promotion_queue_limit"), + promotion_candidate_recorded_( + "master_promotion_candidate_recorded_total", + "New promotion retry candidate entries created"), + promotion_candidate_admitted_( + "master_promotion_candidate_admitted_total", + "Promotion retry candidates successfully queued on background retry"), + promotion_candidate_admission_rejected_( + "master_promotion_candidate_admission_rejected_total", + "Promotion retry candidates that hit a gate again during retry scan"), + promotion_candidate_expired_evaluated_( + "master_promotion_candidate_expired_evaluated_total", + "Promotion retry candidates expired or exhausted during a retry " + "scan"), + promotion_candidate_expired_unevaluated_( + "master_promotion_candidate_expired_unevaluated_total", + "Promotion retry candidates that aged out before the background " + "scheduler ever evaluated them (scan budget too small to reach " + "their shard within the TTL window)"), + promotion_candidate_dropped_limit_( + "master_promotion_candidate_dropped_limit_total", + "Promotion retry candidates dropped at record time because the " + "global candidate count limit was reached"), + tenant_quota_reject_total_( + "mooncake_tenant_quota_reject_total", + "Total number of tenant quota admission rejects", + {"tenant_id", "reason"}), + tenant_evict_bytes_total_( + "mooncake_tenant_evict_bytes_total", + "Total bytes evicted by tenant-scoped quota eviction", {"tenant_id"}), // Snapshot Metrics snapshot_duration_ms_( @@ -469,9 +508,16 @@ void MasterMetricManager::update_metrics_for_zero_output() { promotion_rejected_frequency_.inc(0); promotion_rejected_watermark_.inc(0); promotion_rejected_cap_.inc(0); + promotion_candidate_recorded_.inc(0); + promotion_candidate_admitted_.inc(0); + promotion_candidate_admission_rejected_.inc(0); + promotion_candidate_expired_evaluated_.inc(0); + promotion_candidate_expired_unevaluated_.inc(0); + promotion_candidate_dropped_limit_.inc(0); put_start_requests_.inc(0); put_start_failures_.inc(0); put_start_alloc_failures_.inc(0); + put_start_partial_allocations_.inc(0); put_end_requests_.inc(0); put_end_failures_.inc(0); put_revoke_requests_.inc(0); @@ -564,6 +610,8 @@ void MasterMetricManager::update_metrics_for_zero_output() { // Update Store-observed cache reuse metrics mem_cache_hit_nums_.inc(0); file_cache_hit_nums_.inc(0); + mem_cache_hit_bytes_.inc(0); + file_cache_hit_bytes_.inc(0); valid_get_nums_.inc(0); total_get_nums_.inc(0); @@ -657,6 +705,11 @@ int64_t MasterMetricManager::get_segment_total_mem_capacity( return mem_total_capacity_per_segment_.value({segment}); } +void MasterMetricManager::remove_segment_metrics(const std::string& segment) { + mem_allocated_size_per_segment_.remove_label_value({{"segment", segment}}); + mem_total_capacity_per_segment_.remove_label_value({{"segment", segment}}); +} + double MasterMetricManager::get_segment_mem_used_ratio( const std::string& segment) { double allocated = get_segment_allocated_mem_size(segment); @@ -727,6 +780,12 @@ int64_t MasterMetricManager::get_segment_total_nof_capacity( return nof_total_capacity_per_segment_.value({segment}); } +void MasterMetricManager::remove_nof_segment_metrics( + const std::string& segment) { + nof_allocated_size_per_segment_.remove_label_value({{"segment", segment}}); + nof_total_capacity_per_segment_.remove_label_value({{"segment", segment}}); +} + double MasterMetricManager::get_segment_nof_used_ratio( const std::string& segment) { double allocated = get_segment_allocated_nof_size(segment); @@ -824,12 +883,30 @@ void MasterMetricManager::inc_mem_cache_hit_nums(int64_t val) { void MasterMetricManager::inc_file_cache_hit_nums(int64_t val) { file_cache_hit_nums_.inc(val); } +void MasterMetricManager::inc_mem_cache_hit_bytes(int64_t val) { + mem_cache_hit_bytes_.inc(val); +} +void MasterMetricManager::inc_file_cache_hit_bytes(int64_t val) { + file_cache_hit_bytes_.inc(val); +} +int64_t MasterMetricManager::get_mem_cache_hit_bytes() { + return mem_cache_hit_bytes_.value(); +} +int64_t MasterMetricManager::get_file_cache_hit_bytes() { + return file_cache_hit_bytes_.value(); +} void MasterMetricManager::inc_mem_cache_nums(int64_t val) { mem_cache_nums_.inc(val); } void MasterMetricManager::inc_file_cache_nums(int64_t val) { file_cache_nums_.inc(val); } +int64_t MasterMetricManager::get_mem_cache_nums() { + return mem_cache_nums_.value(); +} +int64_t MasterMetricManager::get_file_cache_nums() { + return file_cache_nums_.value(); +} void MasterMetricManager::dec_mem_cache_nums(int64_t val) { mem_cache_nums_.dec(val); } @@ -863,6 +940,9 @@ void MasterMetricManager::inc_put_start_failures(int64_t val) { void MasterMetricManager::inc_put_start_alloc_failures(int64_t val) { put_start_alloc_failures_.inc(val); } +void MasterMetricManager::inc_put_start_partial_allocations(int64_t val) { + put_start_partial_allocations_.inc(val); +} void MasterMetricManager::inc_put_end_requests(int64_t val) { put_end_requests_.inc(val); } @@ -1112,6 +1192,38 @@ void MasterMetricManager::inc_promotion_rejected_watermark(int64_t val) { void MasterMetricManager::inc_promotion_rejected_cap(int64_t val) { promotion_rejected_cap_.inc(val); } +void MasterMetricManager::inc_promotion_candidate_recorded(int64_t val) { + promotion_candidate_recorded_.inc(val); +} +void MasterMetricManager::inc_promotion_candidate_admitted(int64_t val) { + promotion_candidate_admitted_.inc(val); +} +void MasterMetricManager::inc_promotion_candidate_admission_rejected( + int64_t val) { + promotion_candidate_admission_rejected_.inc(val); +} +void MasterMetricManager::inc_promotion_candidate_expired_evaluated( + int64_t val) { + promotion_candidate_expired_evaluated_.inc(val); +} +void MasterMetricManager::inc_promotion_candidate_expired_unevaluated( + int64_t val) { + promotion_candidate_expired_unevaluated_.inc(val); +} +void MasterMetricManager::inc_promotion_candidate_dropped_limit(int64_t val) { + promotion_candidate_dropped_limit_.inc(val); +} + +void MasterMetricManager::inc_tenant_quota_reject(const std::string& tenant_id, + const std::string& reason, + int64_t val) { + tenant_quota_reject_total_.inc({tenant_id, reason}, val); +} + +void MasterMetricManager::inc_tenant_evict_bytes(const std::string& tenant_id, + int64_t bytes) { + tenant_evict_bytes_total_.inc({tenant_id}, bytes); +} void MasterMetricManager::set_snapshot_duration_ms(int64_t size) { snapshot_duration_ms_.observe(size); @@ -1133,6 +1245,10 @@ int64_t MasterMetricManager::get_put_start_alloc_failures() { return put_start_alloc_failures_.value(); } +int64_t MasterMetricManager::get_put_start_partial_allocations() { + return put_start_partial_allocations_.value(); +} + int64_t MasterMetricManager::get_put_end_requests() { return put_end_requests_.value(); } @@ -1496,6 +1612,24 @@ int64_t MasterMetricManager::get_promotion_rejected_watermark() { int64_t MasterMetricManager::get_promotion_rejected_cap() { return promotion_rejected_cap_.value(); } +int64_t MasterMetricManager::get_promotion_candidate_recorded() { + return promotion_candidate_recorded_.value(); +} +int64_t MasterMetricManager::get_promotion_candidate_admitted() { + return promotion_candidate_admitted_.value(); +} +int64_t MasterMetricManager::get_promotion_candidate_admission_rejected() { + return promotion_candidate_admission_rejected_.value(); +} +int64_t MasterMetricManager::get_promotion_candidate_expired_evaluated() { + return promotion_candidate_expired_evaluated_.value(); +} +int64_t MasterMetricManager::get_promotion_candidate_expired_unevaluated() { + return promotion_candidate_expired_unevaluated_.value(); +} +int64_t MasterMetricManager::get_promotion_candidate_dropped_limit() { + return promotion_candidate_dropped_limit_.value(); +} // CopyStart, CopyEnd, CopyRevoke, MoveStart, MoveEnd, MoveRevoke Metrics void MasterMetricManager::inc_copy_start_requests(int64_t val) { @@ -1669,6 +1803,10 @@ std::string MasterMetricManager::serialize_metrics() { serialize_metric(mem_total_capacity_); serialize_metric(mem_allocated_size_per_segment_); serialize_metric(mem_total_capacity_per_segment_); + serialize_metric(nof_allocated_size_); + serialize_metric(nof_total_capacity_); + serialize_metric(nof_allocated_size_per_segment_); + serialize_metric(nof_total_capacity_per_segment_); serialize_metric(file_allocated_size_); serialize_metric(file_total_capacity_); serialize_metric(key_count_); @@ -1684,6 +1822,7 @@ std::string MasterMetricManager::serialize_metrics() { serialize_metric(put_start_requests_); serialize_metric(put_start_failures_); serialize_metric(put_start_alloc_failures_); + serialize_metric(put_start_partial_allocations_); serialize_metric(put_end_requests_); serialize_metric(put_end_failures_); serialize_metric(put_revoke_requests_); @@ -1704,6 +1843,12 @@ std::string MasterMetricManager::serialize_metrics() { serialize_metric(unmount_segment_failures_); serialize_metric(remount_segment_requests_); serialize_metric(remount_segment_failures_); + serialize_metric(mount_nof_segment_requests_); + serialize_metric(mount_nof_segment_failures_); + serialize_metric(unmount_nof_segment_requests_); + serialize_metric(unmount_nof_segment_failures_); + serialize_metric(remount_nof_segment_requests_); + serialize_metric(remount_nof_segment_failures_); serialize_metric(ping_requests_); serialize_metric(ping_failures_); serialize_metric(nof_heartbeat_success_total_); @@ -1745,24 +1890,63 @@ std::string MasterMetricManager::serialize_metrics() { // Serialize Batch Request Counters serialize_metric(batch_exist_key_requests_); serialize_metric(batch_exist_key_failures_); + serialize_metric(batch_exist_key_partial_successes_); + serialize_metric(batch_exist_key_items_); + serialize_metric(batch_exist_key_failed_items_); serialize_metric(batch_query_ip_requests_); serialize_metric(batch_query_ip_failures_); + serialize_metric(batch_query_ip_partial_successes_); + serialize_metric(batch_query_ip_items_); + serialize_metric(batch_query_ip_failed_items_); serialize_metric(batch_replica_clear_requests_); serialize_metric(batch_replica_clear_failures_); + serialize_metric(batch_replica_clear_partial_successes_); + serialize_metric(batch_replica_clear_items_); + serialize_metric(batch_replica_clear_failed_items_); serialize_metric(batch_get_replica_list_requests_); serialize_metric(batch_get_replica_list_failures_); + serialize_metric(batch_get_replica_list_partial_successes_); + serialize_metric(batch_get_replica_list_items_); + serialize_metric(batch_get_replica_list_failed_items_); serialize_metric(batch_put_start_requests_); serialize_metric(batch_put_start_failures_); + serialize_metric(batch_put_start_partial_successes_); + serialize_metric(batch_put_start_items_); + serialize_metric(batch_put_start_failed_items_); serialize_metric(batch_put_end_requests_); serialize_metric(batch_put_end_failures_); + serialize_metric(batch_put_end_partial_successes_); + serialize_metric(batch_put_end_items_); + serialize_metric(batch_put_end_failed_items_); serialize_metric(batch_put_revoke_requests_); serialize_metric(batch_put_revoke_failures_); + serialize_metric(batch_put_revoke_partial_successes_); + serialize_metric(batch_put_revoke_items_); + serialize_metric(batch_put_revoke_failed_items_); + + // Serialize Store-observed cache reuse metrics + serialize_metric(mem_cache_hit_nums_); + serialize_metric(file_cache_hit_nums_); + serialize_metric(mem_cache_hit_bytes_); + serialize_metric(file_cache_hit_bytes_); + serialize_metric(mem_cache_nums_); + serialize_metric(file_cache_nums_); + serialize_metric(valid_get_nums_); + serialize_metric(total_get_nums_); // Serialize Eviction Counters serialize_metric(eviction_success_); serialize_metric(eviction_attempts_); serialize_metric(evicted_key_count_); serialize_metric(evicted_size_); + serialize_metric(mem_eviction_success_); + serialize_metric(mem_eviction_attempts_); + serialize_metric(mem_evicted_key_count_); + serialize_metric(mem_evicted_size_); + serialize_metric(nof_eviction_success_); + serialize_metric(nof_eviction_attempts_); + serialize_metric(nof_evicted_key_count_); + serialize_metric(nof_evicted_size_); // Serialize PutStart Discard Metrics serialize_metric(put_start_discard_cnt_); @@ -1780,6 +1964,14 @@ std::string MasterMetricManager::serialize_metrics() { serialize_metric(promotion_rejected_frequency_); serialize_metric(promotion_rejected_watermark_); serialize_metric(promotion_rejected_cap_); + serialize_metric(promotion_candidate_recorded_); + serialize_metric(promotion_candidate_admitted_); + serialize_metric(promotion_candidate_admission_rejected_); + serialize_metric(promotion_candidate_expired_evaluated_); + serialize_metric(promotion_candidate_expired_unevaluated_); + serialize_metric(promotion_candidate_dropped_limit_); + serialize_metric(tenant_quota_reject_total_); + serialize_metric(tenant_evict_bytes_total_); // Serialize Snapshot Metrics serialize_metric(snapshot_duration_ms_); @@ -1881,7 +2073,7 @@ std::string MasterMetricManager::get_summary_string( int64_t nof_allocated = nof_allocated_size_.value(); int64_t nof_capacity = nof_total_capacity_.value(); int64_t file_allocated = file_allocated_size_.value(); - int64_t file_capacity = file_total_capacity_.value(); + [[maybe_unused]] int64_t file_capacity = file_total_capacity_.value(); int64_t keys = key_count_.value(); int64_t soft_pin_keys = soft_pin_key_count_.value(); int64_t active_clients = active_clients_.value(); @@ -1892,6 +2084,7 @@ std::string MasterMetricManager::get_summary_string( int64_t put_starts = put_start_requests_.value(); int64_t put_start_fails = put_start_failures_.value(); int64_t put_start_alloc_fails = put_start_alloc_failures_.value(); + int64_t put_start_partial_allocs = put_start_partial_allocations_.value(); int64_t put_ends = put_end_requests_.value(); int64_t put_end_fails = put_end_failures_.value(); int64_t put_revoke_requests = put_revoke_requests_.value(); @@ -2013,6 +2206,7 @@ std::string MasterMetricManager::get_summary_string( current_counters.put_starts = put_starts; current_counters.put_start_fails = put_start_fails; current_counters.put_start_alloc_fails = put_start_alloc_fails; + current_counters.put_start_partial_allocs = put_start_partial_allocs; current_counters.put_ends = put_ends; current_counters.put_end_fails = put_end_fails; current_counters.put_revoke_requests = put_revoke_requests; @@ -2391,6 +2585,8 @@ std::string MasterMetricManager::get_summary_string( << "Success/Attempts=" << eviction_success << "/" << eviction_attempts << ", " << "AllocFail=" << delta(&SummaryCounters::put_start_alloc_fails) << ", " + << "PartialAlloc=" << delta(&SummaryCounters::put_start_partial_allocs) + << ", " << "keys=" << evicted_key_count << ", " << "size=" << byte_size_to_string(evicted_size); // mem eviction diff --git a/mooncake-store/src/master_service.cpp b/mooncake-store/src/master_service.cpp index 9669883d49..a201e8cd05 100644 --- a/mooncake-store/src/master_service.cpp +++ b/mooncake-store/src/master_service.cpp @@ -1,12 +1,20 @@ #include "master_service.h" +#include +#include +#include #include +#include #include +#include #include +#include +#include #include -#include #include #include +#include +#include #include #include #include @@ -15,51 +23,51 @@ #include #include +#include "http_metadata_server.h" #include "master_metric_manager.h" #include "common.h" #include "segment.h" +#ifdef USE_HTTP +#include "transfer_metadata_plugin.h" +#endif #ifdef USE_NOF #include "spdk/spdk_wrapper.h" #endif #ifdef STORE_USE_ETCD #include "etcd_helper.h" -#include "ha/oplog/etcd_oplog_store.h" +#include "ha/kv/etcd_ha_kv_backend.h" #endif +#include "ha/oplog/oplog_batch_storage.h" +#include "ha/oplog/ordered_oplog_writer.h" #include "ha/snapshot/catalog/backends/embedded/embedded_snapshot_catalog_store.h" #include "ha/snapshot/catalog/backends/redis/redis_snapshot_catalog_store.h" #include "ha/snapshot/object/snapshot_object_store.h" +#include "ha/snapshot/snapshot_constants.h" #include "types.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" #include "ha/snapshot/snapshot_logger.h" #include "utils/zstd_util.h" #include "utils/file_util.h" +#include "random.h" #include "utils.h" +#include "kv_event/kv_event_config.h" +#include "master_snapshot_manager.h" +#include "master_snapshot_repository.h" +#include "ha_metric_manager.h" +#include "metadata_store.h" namespace mooncake { -// Snapshot file names -static const std::string SNAPSHOT_METADATA_FILE = "metadata"; -static const std::string SNAPSHOT_SEGMENTS_FILE = "segments"; -static const std::string SNAPSHOT_TASK_MANAGER_FILE = "task_manager"; -static const std::string SNAPSHOT_MANIFEST_FILE = "manifest.txt"; -static const std::string SNAPSHOT_LATEST_FILE = "latest.txt"; -static const std::string SNAPSHOT_BACKUP_SAVE_DIR = - "mooncake_snapshot_save_backup"; -static const std::string SNAPSHOT_BACKUP_RESTORE_DIR = - "mooncake_snapshot_restore_backup"; -static const std::string SNAPSHOT_SERIALIZER_VERSION = "1.0.0"; -static const std::string SNAPSHOT_SERIALIZER_TYPE = "messagepack"; - namespace { -constexpr size_t kUnlimitedSnapshotList = 0; +constexpr int kMaxTenantQuotaEvictionRetries = 2; // Per-cycle offload cap as a fraction of `offloading_queue_limit_`. Used only // when offload-on-evict mode is active. Defers memory eviction for at most // this fraction of the queue limit per BatchEvict cycle; beyond that, eviction -// falls back according to `offload_force_evict_`. A future change may expose -// this as a configurable parameter if workloads demand tuning. -constexpr double kOffloadCapRatio = 0.5; +// falls back according to `offload_force_evict_`. +// NOTE: Both offloading_queue_limit_ and offload_cap_ratio_ are now +// configurable via --offloading_queue_limit and --offload_cap_ratio flags. enum class SnapshotCatalogBackendKind { kEmbedded, @@ -79,18 +87,38 @@ tl::expected ParseSnapshotCatalogKind( std::string(store_type)); } -int64_t CurrentTimeMs() { - return std::chrono::duration_cast( - std::chrono::system_clock::now().time_since_epoch()) - .count(); +uint64_t SaturatingAdd(uint64_t lhs, uint64_t rhs) { + if (lhs > std::numeric_limits::max() - rhs) { + return std::numeric_limits::max(); + } + return lhs + rhs; } -size_t RandomIndex(size_t upper_bound) { - static thread_local std::mt19937 generator(std::random_device{}()); - std::uniform_int_distribution dist(0, upper_bound - 1); - return dist(generator); +uint64_t SaturatingMultiply(uint64_t lhs, uint64_t rhs) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + return std::numeric_limits::max(); + } + return lhs * rhs; } +// Decides whether PutStart may proceed with the replicas that were +// actually allocated. Three deliberately different policies apply: +// +// - Memory-only (nof_replica_num == 0): best-effort. Fewer than +// config.replica_num replicas (but at least one) still succeed, even +// though DetermineReplicaWriteMode() classifies such configs as +// RELIABLE_MULTI_REPLICA. The shortfall is surfaced via a WARNING log +// (action=put_start_partial_allocation) and the +// master_put_start_partial_allocations_total metric. +// - FLEXIBLE_DUAL_REPLICA (1 memory + 1 NoF): allocating either side +// alone is sufficient. +// - Any other config with nof_replica_num > 0: strict. Both replica +// types must match the requested counts exactly, otherwise PutStart +// fails with NO_AVAILABLE_HANDLE. +// +// The "reliable" guarantee of RELIABLE_MULTI_REPLICA is enforced at the +// transfer stage (all allocated replicas must complete or the put is +// revoked), not at the allocation stage for memory-only configs. bool HasExpectedReplicaAllocation(const ReplicateConfig& config, size_t allocated_memory_replicas, size_t allocated_nof_replicas) { @@ -155,20 +183,24 @@ MasterService::MasterService(const MasterServiceConfig& config) enable_offload_(config.enable_offload), ha_backend_type_(config.ha_backend_type), ha_backend_connstring_(config.ha_backend_connstring), + enable_oplog_(config.enable_ha && config.enable_oplog && + config.ha_backend_type == "etcd"), + oplog_batch_max_entries_(config.oplog_batch_max_entries), cluster_id_(config.cluster_id), root_fs_dir_(config.root_fs_dir), global_file_segment_size_(config.global_file_segment_size), enable_disk_eviction_(config.enable_disk_eviction), quota_bytes_(config.quota_bytes), - enable_tenant_quota_(config.enable_tenant_quota), - default_tenant_quota_bytes_(config.default_tenant_quota_bytes), - tenant_quota_pool_capacity_bytes_( - config.tenant_quota_pool_capacity_bytes), + enable_multi_tenants_(config.enable_multi_tenants), + tenant_quota_connector_type_(config.tenant_quota_connector_type), + tenant_quota_connector_uri_(config.tenant_quota_connector_uri), segment_manager_(config.memory_allocator, config.enable_cxl), nof_segment_manager_(config.memory_allocator), memory_allocator_type_(config.memory_allocator), - allocation_strategy_( - CreateAllocationStrategy(config.allocation_strategy_type)), + allocation_strategy_type_(config.enable_cxl + ? AllocationStrategyType::CXL + : config.allocation_strategy_type), + allocation_strategy_(CreateAllocationStrategy(allocation_strategy_type_)), enable_snapshot_restore_(config.enable_snapshot_restore), enable_snapshot_(config.enable_snapshot), snapshot_backup_dir_(config.snapshot_backup_dir), @@ -180,10 +212,26 @@ MasterService::MasterService(const MasterServiceConfig& config) config.snapshot_catalog_store_connstring), put_start_discard_timeout_sec_(config.put_start_discard_timeout_sec), put_start_release_timeout_sec_(config.put_start_release_timeout_sec), - task_manager_(config.task_manager_config), cxl_path_(config.cxl_path), cxl_size_(config.cxl_size), - enable_cxl_(config.enable_cxl) { + enable_cxl_(config.enable_cxl), + offloading_queue_limit_(config.offloading_queue_limit), + offload_cap_ratio_(config.offload_cap_ratio), + task_manager_(config.task_manager_config) { + // Initialize HTTP metadata key prefix (read env var once at startup) + const char* custom_prefix = std::getenv("MC_METADATA_CLUSTER_ID"); + if (custom_prefix && std::strlen(custom_prefix) > 0) { + http_metadata_prefix_ = "mooncake/" + std::string(custom_prefix); + if (http_metadata_prefix_.back() != '/') { + http_metadata_prefix_ += '/'; + } + } else { + http_metadata_prefix_ = "mooncake/"; + } + if (allocation_strategy_type_ == AllocationStrategyType::LOCAL_FIRST) { + LOG(INFO) << "Local-first allocation strategy enabled"; + } + if (enable_snapshot_ || enable_snapshot_restore_) { try { auto object_store_type = @@ -199,11 +247,31 @@ MasterService::MasterService(const MasterServiceConfig& config) if (!snapshot_backup_dir_.empty()) { use_snapshot_backup_dir_ = true; } + + // Initialize repository and codec for both save and restore + snapshot_repository_ = std::make_unique( + snapshot_object_store_.get(), snapshot_catalog_store_.get(), + snapshot_backup_dir_, use_snapshot_backup_dir_); + snapshot_codec_ = std::make_unique(); + } + + if (enable_multi_tenants_) { + auto store = CreateTenantQuotaPolicyStore(tenant_quota_connector_type_, + tenant_quota_connector_uri_, + cluster_id_); + if (!store) { + throw std::invalid_argument(store.error()); + } + tenant_quota_policy_store_ = std::move(store.value()); } if (enable_snapshot_restore_) { RestoreState(); } + if (enable_multi_tenants_) { + LoadTenantQuotaPoliciesFromStoreOrThrow(); + RebuildTenantQuotaUsageFromMetadata(); + } if (enable_snapshot_ && snapshot_retention_count_ == 0) { LOG(ERROR) << "snapshot_retention_count must be greater than 0"; throw std::invalid_argument("snapshot_retention_count must be > 0"); @@ -221,6 +289,25 @@ MasterService::MasterService(const MasterServiceConfig& config) throw std::invalid_argument("Invalid eviction high watermark ratio"); } + // Validate offload tuning knobs here (not only via gflags validator), + // because values loaded from a configuration file bypass the gflags + // validator chain. + if (offload_cap_ratio_ < 0.0 || offload_cap_ratio_ > 1.0) { + LOG(ERROR) << "offload_cap_ratio must be between 0.0 and 1.0, " + << "current value: " << offload_cap_ratio_; + throw std::invalid_argument("Invalid offload_cap_ratio"); + } + if (offloading_queue_limit_ == 0) { + LOG(ERROR) << "offloading_queue_limit must be greater than 0"; + throw std::invalid_argument("Invalid offloading_queue_limit"); + } + if (offloading_queue_limit_ > 100'000'000ULL) { + LOG(ERROR) << "offloading_queue_limit must be <= 100000000 to avoid " + << "overflow when computing offload_cap, current value: " + << offloading_queue_limit_; + throw std::invalid_argument("Invalid offloading_queue_limit"); + } + if (put_start_release_timeout_sec_ <= put_start_discard_timeout_sec_) { LOG(ERROR) << "put_start_release_timeout=" << put_start_release_timeout_sec_.count() @@ -306,6 +393,43 @@ MasterService::MasterService(const MasterServiceConfig& config) << ")"; } + kv_event_publisher_ = + std::make_unique(BuildKvEventConfig(config)); + + if (enable_oplog_ && !cluster_id_.empty()) { +#ifdef STORE_USE_ETCD + if (ha_backend_connstring_.empty()) { + LOG(INFO) << "Skipping automatic batch-record OpLog writer " + "initialization; no HA backend connstring configured"; + } else { + ErrorCode connect_err = EtcdHelper::ConnectToEtcdStoreClient( + ha_backend_connstring_.c_str()); + if (connect_err != ErrorCode::OK) { + throw std::runtime_error(fmt::format( + "failed to connect HA batch-record OpLog writer to etcd: " + "{}", + toString(connect_err))); + } + auto backend = std::make_shared(); + ErrorCode err = InitializeBatchOpLogWriter(std::move(backend)); + if (err != ErrorCode::OK) { + throw std::runtime_error(fmt::format( + "failed to create HA batch-record OpLog writer: {}", + toString(err))); + } + } +#else + if (ha_backend_connstring_.empty()) { + LOG(INFO) << "Skipping automatic batch-record OpLog writer " + "initialization; no HA backend connstring configured"; + } else { + throw std::runtime_error( + "failed to create HA batch-record OpLog writer: ETCD support " + "not compiled in"); + } +#endif + } + eviction_running_ = true; eviction_thread_ = std::thread(&MasterService::EvictionThreadFunc, this); VLOG(1) << "action=start_eviction_thread"; @@ -329,6 +453,11 @@ MasterService::MasterService(const MasterServiceConfig& config) std::thread(&MasterService::TaskCleanupThreadFunc, this); VLOG(1) << "action=start_task_cleanup_thread"; + // NOTE: The async HTTP metadata cleanup worker is started lazily in + // setHttpMetadataRemoteUrl() once http_metadata_remote_ is initialized, + // since that happens after this constructor returns (in + // WrappedMasterService). + job_dispatch_running_ = true; job_dispatch_thread_ = std::thread(&MasterService::JobDispatchThreadFunc, this); @@ -344,12 +473,36 @@ MasterService::MasterService(const MasterServiceConfig& config) } } - if (enable_snapshot_) { + if (enable_snapshot_ && !enable_oplog_) { if (memory_allocator_type_ == BufferAllocatorType::OFFSET) { - snapshot_running_ = true; - snapshot_thread_ = - std::thread(&MasterService::SnapshotThreadFunc, this); - } + // Initialize and start snapshot manager + MasterSnapshotManagerOptions snapshot_options; + snapshot_options.enable_snapshot = enable_snapshot_; + snapshot_options.snapshot_interval_seconds = + snapshot_interval_seconds_; + snapshot_options.snapshot_child_timeout_seconds = + snapshot_child_timeout_seconds_; + snapshot_options.snapshot_retention_count = + snapshot_retention_count_; + snapshot_options.snapshot_backup_dir = snapshot_backup_dir_; + snapshot_options.use_snapshot_backup_dir = use_snapshot_backup_dir_; + snapshot_options.snapshot_catalog_store_type = + snapshot_catalog_store_type_; + snapshot_options.snapshot_catalog_store_connstring = + snapshot_catalog_store_connstring_; + snapshot_options.ha_backend_type = ha_backend_type_; + snapshot_options.ha_backend_connstring = ha_backend_connstring_; + snapshot_options.cluster_id = cluster_id_; + snapshot_options.enable_ha = enable_ha_; + + snapshot_manager_ = std::make_unique( + this, snapshot_options, snapshot_mutex_, + snapshot_object_store_.get(), snapshot_catalog_store_.get()); + snapshot_manager_->Start(); + } + } else if (enable_snapshot_ && enable_oplog_) { + LOG(INFO) << "Skipping primary snapshot generation in batch-record " + "OpLog mode; snapshots are owned by standby"; } if (enable_cxl_) { @@ -396,12 +549,22 @@ MasterService::CreateSnapshotCatalogStore() { } MasterService::~MasterService() { + if (ordered_oplog_writer_) { + ordered_oplog_writer_->Stop(); + } + // Stop and join the threads eviction_running_ = false; client_monitor_running_ = false; - snapshot_running_ = false; + + // Stop snapshot manager (non-blocking) + if (snapshot_manager_) { + snapshot_manager_->Stop(); + } + task_cleanup_running_ = false; job_dispatch_running_ = false; + http_metadata_cleanup_running_ = false; graceful_unmount_scheduler_.Stop(); #ifdef USE_NOF nof_heartbeat_running_ = false; @@ -409,6 +572,7 @@ MasterService::~MasterService() { // Wake sleepers so join() doesn't block for long sleep intervals. task_cleanup_cv_.notify_all(); + http_metadata_cleanup_cv_.notify_all(); if (eviction_thread_.joinable()) { eviction_thread_.join(); @@ -421,15 +585,47 @@ MasterService::~MasterService() { nof_heartbeat_thread_.join(); } #endif - if (snapshot_thread_.joinable()) { - snapshot_thread_.join(); - } if (task_cleanup_thread_.joinable()) { task_cleanup_thread_.join(); } + if (http_metadata_cleanup_thread_.joinable()) { + http_metadata_cleanup_thread_.join(); + } if (job_dispatch_thread_.joinable()) { job_dispatch_thread_.join(); } + + // Reset snapshot manager after all other threads have joined + // This triggers the destructor which joins the snapshot thread + if (snapshot_manager_) { + snapshot_manager_.reset(); + } + for (const auto& [segment, bytes] : standby_accounted_memory_bytes_) { + MasterMetricManager::instance().dec_allocated_mem_size( + segment, static_cast(bytes)); + MasterMetricManager::instance().remove_segment_metrics(segment); + } + + // Segments still mounted here never went through CommitUnmountSegment; + // release their capacity contribution so the process-lifetime + // MasterMetricManager stays consistent when the next leadership term + // constructs a fresh MasterService and the clients remount. + segment_manager_.releaseCapacityMetrics(); +} + +ErrorCode MasterService::SetBatchOpLogBackendForTesting( + std::shared_ptr backend) { + return InitializeBatchOpLogWriter(std::move(backend)); +} + +void MasterService::RunBatchEvictForTesting(double evict_ratio_target, + double evict_ratio_lowerbound) { + BatchEvict(evict_ratio_target, evict_ratio_lowerbound); +} + +void MasterService::RunNoFBatchEvictForTesting(double evict_ratio_target, + double evict_ratio_lowerbound) { + NoFBatchEvict(evict_ratio_target, evict_ratio_lowerbound); } void MasterService::SetNoFProbeFnForTesting(NoFProbeFn fn) { @@ -476,31 +672,111 @@ std::optional MasterService::GetNoFHeartbeatFailureCountForTesting( return it->second.consecutive_failures; } -std::optional -MasterService::GetTenantQuotaSnapshotForTesting( - const std::string& tenant_id) const { - const auto normalized_tenant = NormalizeTenantId(tenant_id); - const auto shard_idx = getTenantQuotaShardIndex(normalized_tenant); - const auto& shard = tenant_quota_shards_[shard_idx]; - std::lock_guard lock(shard.mutex); - auto it = shard.tenants.find(normalized_tenant); - if (it == shard.tenants.end()) { - return std::nullopt; +bool MasterService::IsTenantQuotaEnabled() const { + return enable_multi_tenants_; +} + +std::vector MasterService::ListTenantQuotaSnapshots() + const { + return tenant_quota_table_.ListTenantSnapshots(); +} + +std::optional MasterService::GetTenantQuotaSnapshot( + const TenantId& tenant_id) const { + assert(tenant_id.IsValid()); + return tenant_quota_table_.GetTenantSnapshot(tenant_id); +} + +tl::expected +MasterService::UpsertTenantQuotaPolicy(const TenantId& tenant_id, + uint64_t requested_quota_bytes) { + assert(tenant_id.IsValid()); + if (!enable_multi_tenants_) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + if (requested_quota_bytes == 0) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } - const auto& state = it->second; - return TenantQuotaSnapshot{ - .tenant_id = normalized_tenant, - .requested_quota_bytes = state.requested_quota_bytes, - .effective_quota_bytes = state.effective_quota_bytes, - .used_bytes = state.used_bytes, - .reserved_bytes = state.reserved_bytes, - .committed_count = state.committed_count, - .has_explicit_policy = state.has_explicit_policy, - .over_quota = state.over_quota}; + + std::lock_guard policy_lock(tenant_quota_policy_mutex_); + auto policy = BuildTenantQuotaPolicySnapshot(); + policy.tenant_quotas[tenant_id.value()] = requested_quota_bytes; + auto save_result = tenant_quota_policy_store_->Save(policy); + if (!save_result) { + LOG(ERROR) << "failed to save tenant quota policy: " + << save_result.error(); + return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL); + } + ApplyTenantQuotaPolicies(policy); + auto result_snapshot = GetTenantQuotaSnapshot(tenant_id); + if (!result_snapshot.has_value()) { + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + return result_snapshot.value(); +} + +tl::expected, ErrorCode> +MasterService::DeleteTenantQuotaPolicy(const TenantId& tenant_id) { + assert(tenant_id.IsValid()); + if (!enable_multi_tenants_) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + + std::lock_guard policy_lock(tenant_quota_policy_mutex_); + auto policy = BuildTenantQuotaPolicySnapshot(); + auto policy_it = policy.tenant_quotas.find(tenant_id.value()); + if (policy_it == policy.tenant_quotas.end()) { + return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + } + const uint64_t requested_quota_bytes = policy_it->second; + + auto restore_policy = [&] { + std::lock_guard recompute_lock( + tenant_quota_recompute_mutex_); + const uint64_t capacity = GetTenantQuotaAllocatableCapacityBytes(); + auto result = tenant_quota_table_.UpsertTenantPolicy( + tenant_id, requested_quota_bytes, capacity); + if (!result) { + LOG(ERROR) << "failed to restore tenant quota policy tenant=" + << tenant_id.value(); + } + }; + + auto disable_result = + tenant_quota_table_.DisableTenantPolicyIfEmpty(tenant_id); + if (!disable_result) { + return tl::make_unexpected(disable_result.error() == + TenantQuotaError::kTenantNotEmpty + ? ErrorCode::TENANT_NOT_EMPTY + : ErrorCode::OBJECT_NOT_FOUND); + } + + auto post_mark_snapshot = GetTenantQuotaSnapshot(tenant_id); + if (TenantHasObjects(tenant_id) || + (post_mark_snapshot.has_value() && + (post_mark_snapshot->used_bytes != 0 || + post_mark_snapshot->reserved_bytes != 0 || + post_mark_snapshot->committed_count != 0 || + post_mark_snapshot->metadata_object_count != 0))) { + restore_policy(); + return tl::make_unexpected(ErrorCode::TENANT_NOT_EMPTY); + } + + policy.tenant_quotas.erase(policy_it); + auto save_result = tenant_quota_policy_store_->Save(policy); + if (!save_result) { + restore_policy(); + LOG(ERROR) << "failed to save tenant quota policy: " + << save_result.error(); + return tl::make_unexpected(ErrorCode::PERSISTENT_FAIL); + } + ApplyTenantQuotaPolicies(policy); + return GetTenantQuotaSnapshot(tenant_id); } auto MasterService::MountSegment(const Segment& segment, const UUID& client_id) -> tl::expected { + ErrorCode mount_result = ErrorCode::OK; std::shared_lock shared_lock(snapshot_mutex_); { ScopedSegmentAccess segment_access = @@ -534,12 +810,28 @@ auto MasterService::MountSegment(const Segment& segment, const UUID& client_id) auto err = segment_access.MountSegment(segment, client_id); if (err == ErrorCode::SEGMENT_ALREADY_EXISTS) { // Return OK because this is an idempotent operation - return {}; + mount_result = err; } else if (err != ErrorCode::OK) { return tl::make_unexpected(err); } } - RecomputeTenantEffectiveQuotas(); + + if (enable_oplog_ && ordered_oplog_writer_) { + SegmentMountOp op; + op.segment_name = segment.name; + op.transport_endpoint = segment.te_endpoint; + op.capacity = segment.size; + op.is_memory_segment = true; + op.file_path.clear(); + auto bytes = struct_pack::serialize(op); + PersistSegmentOpForHAOrEnqueue("MountSegment", OpType::SEGMENT_MOUNT, + segment.te_endpoint, + std::string(bytes.begin(), bytes.end())); + } + UpdateClientHostId(client_id, segment.host_id); + if (mount_result == ErrorCode::OK) { + RecomputeTenantEffectiveQuotas(); + } return {}; } @@ -569,12 +861,56 @@ auto MasterService::MountNoFSegment(const NoFSegment& segment, #endif } +ErrorCode MasterService::ValidateStandbyRemountSegment( + const Segment& segment) const { + const StandbySegmentInfo* match = nullptr; + for (const auto& standby : standby_memory_segments_) { + if (standby.transport_endpoint == segment.te_endpoint || + standby.segment_name == segment.name) { + if (match != nullptr && match != &standby) { + return ErrorCode::INVALID_PARAMS; + } + match = &standby; + } + } + if (match != nullptr && segment.protocol == "cxl") { + return ErrorCode::UNAVAILABLE_IN_CURRENT_MODE; + } + if (match != nullptr && (match->segment_name != segment.name || + match->transport_endpoint != segment.te_endpoint || + match->capacity != segment.size)) { + return ErrorCode::INVALID_PARAMS; + } + return ErrorCode::OK; +} + auto MasterService::ReMountSegment(const std::vector& segments, const UUID& client_id) -> tl::expected { - std::shared_lock shared_lock(snapshot_mutex_); + std::unique_lock snapshot_lock(snapshot_mutex_); { std::unique_lock lock(client_mutex_); + for (const auto& segment : segments) { + if (!segment.host_id.empty()) { + client_host_id_[client_id] = segment.host_id; + break; + } + } + { + auto segment_access = segment_manager_.getSegmentAccess(); + for (const auto& segment : segments) { + auto standby_validation = + ValidateStandbyRemountSegment(segment); + if (standby_validation != ErrorCode::OK) { + return tl::make_unexpected(standby_validation); + } + auto validation = + segment_access.ValidateRemountSegment(segment, client_id); + if (validation != ErrorCode::OK) { + return tl::make_unexpected(validation); + } + } + } if (ok_client_.contains(client_id)) { LOG(WARNING) << "client_id=" << client_id << ", warn=client_already_remounted"; @@ -585,6 +921,37 @@ auto MasterService::ReMountSegment(const std::vector& segments, { ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); + std::vector segment_existed(segments.size()); + for (size_t i = 0; i < segments.size(); ++i) { + segment_existed[i] = + segment_access.GetAllocator(segments[i].id) != nullptr; + } + auto rollback_new_segments = [&] { + for (size_t i = 0; i < segments.size(); ++i) { + if (segment_existed[i] || + !segment_access.GetAllocator(segments[i].id)) { + continue; + } + size_t capacity = 0; + if (segment_access.PrepareUnmountSegment( + segments[i].id, capacity) != ErrorCode::OK) { + LOG(ERROR) << "segment_name=" << segments[i].name + << ", error=remount_rollback_prepare_failed"; + continue; + } + if (segment_access.CommitUnmountSegment( + segments[i].id, client_id, capacity) != + ErrorCode::OK) { + LOG(ERROR) << "segment_name=" << segments[i].name + << ", error=remount_rollback_commit_failed"; + } + } + }; + auto fail_remount = + [&](ErrorCode error) -> tl::expected { + rollback_new_segments(); + return tl::make_unexpected(error); + }; // Tell the client monitor thread to start timing for this client. // To avoid the following undesired situations, this message must be @@ -609,7 +976,171 @@ auto MasterService::ReMountSegment(const std::vector& segments, ErrorCode err = segment_access.ReMountSegment(segments, client_id); if (err != ErrorCode::OK) { - return tl::make_unexpected(err); + return fail_remount(err); + } + + struct SegmentRestore { + Segment segment; + std::shared_ptr old_allocator; + std::shared_ptr restored_allocator; + std::vector replicas; + std::vector descriptors; + std::vector> buffers; + uint64_t imported_size{0}; + }; + std::vector restores; + restores.reserve(segments.size()); + for (const auto& segment : segments) { + auto allocator = segment_access.GetAllocator(segment.id); + Segment authoritative; + if (!allocator || + !segment_access.GetSegment(segment.id, authoritative)) { + return fail_remount(ErrorCode::INTERNAL_ERROR); + } + restores.push_back({std::move(authoritative), + std::move(allocator), + nullptr, + {}, + {}, + {}, + 0}); + } + + bool ambiguous_endpoint = false; + bool unsupported_cxl = false; + for (size_t shard_index = 0; shard_index < kNumShards; + ++shard_index) { + MetadataShardAccessorRW shard(this, shard_index); + for (auto& [tenant_id, tenant] : shard->tenants) { + (void)tenant_id; + for (auto& [key, metadata] : tenant.metadata) { + (void)key; + metadata.VisitReplicas( + [](const Replica& replica) { + return replica.is_memory_replica() && + replica.status() != + ReplicaStatus::REMOVED && + replica.status() != + ReplicaStatus::FAILED; + }, + [&](Replica& replica) { + auto descriptor = replica.get_descriptor() + .get_memory_descriptor() + .buffer_descriptor; + SegmentRestore* match = nullptr; + for (auto& restore : restores) { + if (descriptor.transport_endpoint_ == + restore.segment.te_endpoint || + descriptor.transport_endpoint_ == + restore.segment.name) { + if (match != nullptr) { + ambiguous_endpoint = true; + return; + } + match = &restore; + } + } + if (match != nullptr) { + if (descriptor.protocol_ == "cxl") { + unsupported_cxl = true; + return; + } + descriptor.transport_endpoint_ = + match->segment.te_endpoint; + match->replicas.push_back(&replica); + match->descriptors.push_back(descriptor); + } + }); + } + } + } + if (ambiguous_endpoint) { + return fail_remount(ErrorCode::INVALID_PARAMS); + } + if (unsupported_cxl) { + return fail_remount(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + + for (auto& restore : restores) { + if (restore.descriptors.empty()) { + continue; + } + if (std::dynamic_pointer_cast( + restore.old_allocator)) { + auto restored = RestoreOffsetBufferAllocator( + restore.segment.name, restore.segment.base, + restore.segment.size, restore.segment.te_endpoint, + restore.descriptors); + if (!restored) { + return fail_remount(ErrorCode::INVALID_PARAMS); + } + restore.restored_allocator = std::move(restored->allocator); + restore.buffers = std::move(restored->buffers); + } else if (std::dynamic_pointer_cast( + restore.old_allocator)) { + auto restored = RestoreCachelibBufferAllocator( + restore.segment.name, restore.segment.base, + restore.segment.size, restore.segment.te_endpoint, + restore.descriptors); + if (!restored) { + return fail_remount(ErrorCode::INVALID_PARAMS); + } + restore.restored_allocator = std::move(restored->allocator); + restore.buffers = std::move(restored->buffers); + } else { + return fail_remount(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + } + + std::vector + allocator_replacements; + for (auto& restore : restores) { + if (restore.restored_allocator) { + if (restore.buffers.size() != restore.replicas.size() || + std::any_of( + restore.buffers.begin(), restore.buffers.end(), + [](const auto& buffer) { return !buffer; })) { + return fail_remount(ErrorCode::INTERNAL_ERROR); + } + restore.imported_size = std::accumulate( + restore.descriptors.begin(), restore.descriptors.end(), + uint64_t{0}, [](uint64_t sum, const auto& descriptor) { + return sum + descriptor.size_; + }); + auto accounted = standby_accounted_memory_bytes_.find( + restore.segment.name); + if (accounted == standby_accounted_memory_bytes_.end() || + accounted->second < restore.imported_size) { + return fail_remount(ErrorCode::INTERNAL_ERROR); + } + allocator_replacements.push_back( + {restore.segment.id, restore.old_allocator, + restore.restored_allocator}); + } + } + if (!segment_access.ReplaceAllocators(allocator_replacements)) { + return fail_remount(ErrorCode::INTERNAL_ERROR); + } + for (auto& restore : restores) { + if (restore.imported_size != 0) { + MasterMetricManager::instance().dec_allocated_mem_size( + restore.segment.name, + static_cast(restore.imported_size)); + auto accounted = standby_accounted_memory_bytes_.find( + restore.segment.name); + accounted->second -= restore.imported_size; + if (accounted->second == 0) { + standby_accounted_memory_bytes_.erase(accounted); + } + } + for (size_t i = 0; i < restore.replicas.size(); ++i) { + (void)restore.replicas[i]->replace_memory_buffer( + std::move(restore.buffers[i])); + } + invalid_replica_endpoints_.erase(restore.segment.te_endpoint); + invalid_replica_endpoints_.erase(restore.segment.name); + standby_allocator_keepalive_.erase(restore.segment.te_endpoint); + standby_allocator_keepalive_.erase(restore.segment.name); } } @@ -617,6 +1148,21 @@ auto MasterService::ReMountSegment(const std::vector& segments, ok_client_.insert(client_id); MasterMetricManager::instance().inc_active_clients(); } + + if (enable_oplog_ && ordered_oplog_writer_) { + for (const auto& seg : segments) { + SegmentMountOp op; + op.segment_name = seg.name; + op.transport_endpoint = seg.te_endpoint; + op.capacity = seg.size; + op.is_memory_segment = true; + op.file_path.clear(); + auto bytes = struct_pack::serialize(op); + PersistSegmentOpForHAOrEnqueue( + "ReMountSegment", OpType::SEGMENT_MOUNT, seg.name, + std::string(bytes.begin(), bytes.end())); + } + } RecomputeTenantEffectiveQuotas(); return {}; @@ -647,22 +1193,135 @@ MasterService::getAliveClientsSnapshot() const { return ok_client_; } -size_t MasterService::getMetadataShardIndex(const std::string& tenant_id, +void MasterService::UpdateClientHostId(const UUID& client_id, + const std::string& host_id) { + if (host_id.empty()) { + return; + } + { + std::shared_lock lock(client_mutex_); + auto it = client_host_id_.find(client_id); + if (it != client_host_id_.end() && it->second == host_id) { + return; + } + } + + std::unique_lock lock(client_mutex_); + auto it = client_host_id_.find(client_id); + if (it == client_host_id_.end() || it->second != host_id) { + client_host_id_[client_id] = host_id; + } +} + +std::string MasterService::GetClientHostId(const UUID& client_id) const { + std::shared_lock lock(client_mutex_); + auto it = client_host_id_.find(client_id); + return it == client_host_id_.end() ? std::string() : it->second; +} + +size_t MasterService::getMetadataShardIndex(const TenantId& tenant_id, const std::string& key) const { - const auto normalized_tenant = NormalizeTenantId(tenant_id); std::shared_lock lock(group_routing_mutex_); - auto it = - object_group_ids_.find(MakeTenantScopedKey(normalized_tenant, key)); + auto it = object_group_ids_.find(tenant_id.MakeScopedKey(key)); if (it == object_group_ids_.end()) { - return getShardIndex(normalized_tenant, key); + return getShardIndex(tenant_id, key); } return getShardIndex(it->second); } -size_t MasterService::getTenantQuotaShardIndex( - const std::string& tenant_id) const { - return std::hash{}(NormalizeTenantId(tenant_id)) % - kNumTenantQuotaShards; +const TenantId& MasterService::ResolveRequestTenantId( + const TenantId& tenant_id) const { + assert(tenant_id.IsValid()); + if (!enable_multi_tenants_) { + return TenantId::Default(); + } + return tenant_id; +} + +MasterService::ObjectIdentity MasterService::MakeObjectIdentityForRequest( + const std::string& user_key, const TenantId& tenant_id) const { + return {ResolveRequestTenantId(tenant_id), user_key}; +} + +bool MasterService::IsTenantRegistered(const TenantId& tenant_id) const { + if (!enable_multi_tenants_) { + return true; + } + return tenant_quota_table_.IsTenantRegistered(tenant_id); +} + +tl::expected MasterService::ResolveTenantIdForWrite( + const TenantId& tenant_id) const { + assert(tenant_id.IsValid()); + if (!enable_multi_tenants_) { + return TenantId::Default(); + } + std::lock_guard policy_lock(tenant_quota_policy_mutex_); + return ResolveTenantIdForWriteLocked(tenant_id); +} + +tl::expected MasterService::ResolveTenantIdForWriteLocked( + const TenantId& tenant_id) const { + assert(tenant_id.IsValid()); + if (!enable_multi_tenants_) { + return TenantId::Default(); + } + if (!IsTenantRegistered(tenant_id)) { + return tl::make_unexpected(ErrorCode::TENANT_NOT_REGISTERED); + } + return tenant_id; +} + +bool MasterService::TenantHasObjects(const TenantId& tenant_id) const { + for (size_t i = 0; i < kNumShards; ++i) { + MetadataShardAccessorRO shard(this, i); + auto tenant_it = shard->tenants.find(tenant_id); + if (tenant_it != shard->tenants.end() && + !tenant_it->second.metadata.empty()) { + return true; + } + } + return false; +} + +TenantQuotaPolicySnapshot MasterService::BuildTenantQuotaPolicySnapshot() + const { + TenantQuotaPolicySnapshot snapshot; + for (const auto& [tenant_id, requested_quota_bytes] : + tenant_quota_table_.GetTenantPolicies()) { + snapshot.tenant_quotas.emplace(tenant_id.value(), + requested_quota_bytes); + } + return snapshot; +} + +void MasterService::ApplyTenantQuotaPolicies( + const TenantQuotaPolicySnapshot& snapshot) { + TenantQuotaPolicyMap policies; + for (const auto& [tenant_id, requested_quota_bytes] : + snapshot.tenant_quotas) { + policies.emplace(TenantId(tenant_id), requested_quota_bytes); + } + std::lock_guard recompute_lock(tenant_quota_recompute_mutex_); + const uint64_t capacity = GetTenantQuotaAllocatableCapacityBytes(); + tenant_quota_table_.ApplyTenantPolicies(policies, capacity); +} + +void MasterService::LoadTenantQuotaPoliciesFromStoreOrThrow() { + if (!enable_multi_tenants_) { + return; + } + if (!tenant_quota_policy_store_) { + throw std::runtime_error( + "tenant quota policy store is not initialized"); + } + std::lock_guard policy_lock(tenant_quota_policy_mutex_); + auto snapshot = tenant_quota_policy_store_->Load(); + if (!snapshot) { + throw std::runtime_error("failed to load tenant quota policy: " + + snapshot.error()); + } + ApplyTenantQuotaPolicies(snapshot.value()); } uint64_t MasterService::CompletedMemoryQuotaCharge( @@ -683,10 +1342,12 @@ uint64_t MasterService::RequestedMemoryQuotaCharge( return static_cast(charge); } -uint64_t MasterService::GetTenantQuotaCapacityBytes() { - if (tenant_quota_pool_capacity_bytes_ != 0) { - return tenant_quota_pool_capacity_bytes_; - } +bool MasterService::ShouldProtectZeroChargeMetadataCreate( + uint64_t requested_quota_charge) const { + return enable_multi_tenants_ && requested_quota_charge == 0; +} + +uint64_t MasterService::GetTenantQuotaAllocatableCapacityBytes() { uint64_t capacity = 0; ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); std::vector> segments; @@ -703,325 +1364,108 @@ uint64_t MasterService::GetTenantQuotaCapacityBytes() { } void MasterService::RecomputeTenantEffectiveQuotas() { - if (!enable_tenant_quota_) { - return; - } - const uint64_t capacity = GetTenantQuotaCapacityBytes(); - - std::vector> active_tenants; - for (size_t i = 0; i < kNumTenantQuotaShards; ++i) { - auto& shard = tenant_quota_shards_[i]; - std::lock_guard lock(shard.mutex); - for (auto it = shard.tenants.begin(); it != shard.tenants.end();) { - const auto& tenant_id = it->first; - auto& state = it->second; - if (!state.has_explicit_policy) { - state.requested_quota_bytes = default_tenant_quota_bytes_; - } - if (!state.has_explicit_policy && state.used_bytes == 0 && - state.reserved_bytes == 0 && state.committed_count == 0) { - it = shard.tenants.erase(it); - continue; - } - active_tenants.emplace_back(i, tenant_id); - ++it; - } - } - - if (default_tenant_quota_bytes_ == 0) { - for (const auto& [shard_idx, tenant_id] : active_tenants) { - auto& shard = tenant_quota_shards_[shard_idx]; - std::lock_guard lock(shard.mutex); - auto it = shard.tenants.find(tenant_id); - if (it == shard.tenants.end()) { - continue; - } - auto& state = it->second; - state.effective_quota_bytes = std::numeric_limits::max(); - state.over_quota = false; - } + if (!enable_multi_tenants_) { return; } - - const uint64_t active_count = active_tenants.size(); - for (size_t ordinal = 0; ordinal < active_tenants.size(); ++ordinal) { - const auto& [shard_idx, tenant_id] = active_tenants[ordinal]; - auto& shard = tenant_quota_shards_[shard_idx]; - std::lock_guard lock(shard.mutex); - auto it = shard.tenants.find(tenant_id); - if (it == shard.tenants.end()) { - continue; - } - auto& state = it->second; - uint64_t effective = 0; - if (active_count > 0 && capacity > 0) { - effective = capacity / active_count; - if (ordinal < capacity % active_count) { - ++effective; - } - } - state.effective_quota_bytes = effective; - state.over_quota = static_cast(state.used_bytes) + - state.reserved_bytes > - state.effective_quota_bytes; - } + std::lock_guard recompute_lock(tenant_quota_recompute_mutex_); + const uint64_t capacity = GetTenantQuotaAllocatableCapacityBytes(); + tenant_quota_table_.RecomputeEffectiveQuotas(capacity); } tl::expected MasterService::ReserveTenantQuota( - const std::string& tenant_id, uint64_t bytes) { - if (!enable_tenant_quota_ || bytes == 0) { + const TenantId& tenant_id, uint64_t bytes) { + if (!enable_multi_tenants_) { return {}; } - const auto normalized_tenant = NormalizeTenantId(tenant_id); - auto exceeds_effective_quota = [](const TenantQuotaState& state, - uint64_t additional_bytes) { - return static_cast(state.used_bytes) + - state.reserved_bytes + additional_bytes > - state.effective_quota_bytes; - }; - auto refresh_over_quota = [](TenantQuotaState& state) { - state.over_quota = state.effective_quota_bytes != - std::numeric_limits::max() && - static_cast(state.used_bytes) + - state.reserved_bytes > - state.effective_quota_bytes; - }; - - auto& shard = - tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)]; - { - std::lock_guard lock(shard.mutex); - auto it = shard.tenants.find(normalized_tenant); - if (it != shard.tenants.end()) { - auto& state = it->second; - if (exceeds_effective_quota(state, bytes)) { - return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); - } - - state.reserved_bytes += bytes; - refresh_over_quota(state); - return {}; - } - - auto [insert_it, _] = shard.tenants.try_emplace(normalized_tenant); - auto& state = insert_it->second; - state.requested_quota_bytes = default_tenant_quota_bytes_; - state.effective_quota_bytes = default_tenant_quota_bytes_ == 0 - ? std::numeric_limits::max() - : 0; - state.over_quota = false; - - if (default_tenant_quota_bytes_ == 0) { - if (exceeds_effective_quota(state, bytes)) { - shard.tenants.erase(insert_it); - return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); - } - state.reserved_bytes += bytes; - return {}; - } - - state.reserved_bytes = bytes; - } - - RecomputeTenantEffectiveQuotas(); - - bool rollback_needs_recompute = false; - { - std::lock_guard lock(shard.mutex); - auto it = shard.tenants.find(normalized_tenant); - if (it == shard.tenants.end()) { - return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); - } - - auto& state = it->second; - if (!exceeds_effective_quota(state, 0)) { - refresh_over_quota(state); - return {}; - } - - if (state.reserved_bytes < bytes) { - LOG(ERROR) << "tenant quota reserve rollback mismatch tenant=" - << normalized_tenant << ", bytes=" << bytes - << ", reserved=" << state.reserved_bytes; - return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); - } - state.reserved_bytes -= bytes; - if (!state.has_explicit_policy && state.used_bytes == 0 && - state.reserved_bytes == 0 && state.committed_count == 0) { - shard.tenants.erase(it); - rollback_needs_recompute = true; - } else { - refresh_over_quota(state); - } - } - - if (rollback_needs_recompute) { - RecomputeTenantEffectiveQuotas(); + auto result = tenant_quota_table_.Reserve(tenant_id, bytes); + if (result) { + return {}; } - return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); + return tl::make_unexpected( + result.error() == TenantQuotaError::kTenantNotRegistered + ? ErrorCode::TENANT_NOT_REGISTERED + : result.error() == TenantQuotaError::kQuotaExceeded + ? ErrorCode::TENANT_QUOTA_EXCEEDED + : ErrorCode::INTERNAL_ERROR); } -void MasterService::CommitTenantQuota(const std::string& tenant_id, +void MasterService::CommitTenantQuota(const TenantId& tenant_id, uint64_t bytes) { - if (!enable_tenant_quota_ || bytes == 0) { + if (!enable_multi_tenants_ || bytes == 0) { return; } - const auto normalized_tenant = NormalizeTenantId(tenant_id); - auto& shard = - tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)]; - std::lock_guard lock(shard.mutex); - auto it = shard.tenants.find(normalized_tenant); - if (it == shard.tenants.end()) { + if (!tenant_quota_table_.Commit(tenant_id, bytes)) { LOG(ERROR) << "tenant quota commit mismatch tenant=" - << normalized_tenant << ", bytes=" << bytes - << ", reserved=0"; + << tenant_id.value() << ", bytes=" << bytes; + } +} + +void MasterService::AbortTenantQuota(const TenantId& tenant_id, + uint64_t bytes) { + if (!enable_multi_tenants_ || bytes == 0) { return; } - auto& state = it->second; - if (state.reserved_bytes < bytes) { - LOG(ERROR) << "tenant quota commit mismatch tenant=" - << normalized_tenant << ", bytes=" << bytes - << ", reserved=" << state.reserved_bytes; + if (!tenant_quota_table_.Abort(tenant_id, bytes)) { + LOG(ERROR) << "tenant quota abort mismatch tenant=" << tenant_id.value() + << ", bytes=" << bytes; + } +} + +void MasterService::ReleaseTenantQuota(const TenantId& tenant_id, + uint64_t bytes) { + if (!enable_multi_tenants_ || bytes == 0) { return; } - state.reserved_bytes -= bytes; - if (state.used_bytes > std::numeric_limits::max() - bytes) { - state.used_bytes = std::numeric_limits::max(); - } else { - state.used_bytes += bytes; + if (!tenant_quota_table_.Release(tenant_id, bytes)) { + LOG(ERROR) << "tenant quota release mismatch tenant=" + << tenant_id.value() << ", bytes=" << bytes; } - ++state.committed_count; - state.over_quota = - state.effective_quota_bytes != std::numeric_limits::max() && - static_cast(state.used_bytes) + - state.reserved_bytes > - state.effective_quota_bytes; } -void MasterService::AbortTenantQuota(const std::string& tenant_id, - uint64_t bytes) { - if (!enable_tenant_quota_ || bytes == 0) { +void MasterService::ReleaseTenantQuotaPartial(const TenantId& tenant_id, + uint64_t bytes) { + if (!enable_multi_tenants_ || bytes == 0) { return; } - const auto normalized_tenant = NormalizeTenantId(tenant_id); - auto& shard = - tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)]; - bool recompute_needed = false; - { - std::lock_guard lock(shard.mutex); - auto it = shard.tenants.find(normalized_tenant); - if (it == shard.tenants.end()) { - LOG(ERROR) << "tenant quota abort mismatch tenant=" - << normalized_tenant << ", bytes=" << bytes - << ", reserved=0"; - return; - } - auto& state = it->second; - if (state.reserved_bytes < bytes) { - LOG(ERROR) << "tenant quota abort mismatch tenant=" - << normalized_tenant << ", bytes=" << bytes - << ", reserved=" << state.reserved_bytes; - return; - } - state.reserved_bytes -= bytes; - if (!state.has_explicit_policy && state.used_bytes == 0 && - state.reserved_bytes == 0 && state.committed_count == 0) { - shard.tenants.erase(it); - recompute_needed = true; - } else { - state.over_quota = - state.effective_quota_bytes != - std::numeric_limits::max() && - static_cast(state.used_bytes) + - state.reserved_bytes > - state.effective_quota_bytes; - } - } - if (recompute_needed) { - RecomputeTenantEffectiveQuotas(); + if (!tenant_quota_table_.ReleasePartial(tenant_id, bytes)) { + LOG(ERROR) << "tenant quota partial release mismatch tenant=" + << tenant_id.value() << ", bytes=" << bytes; } } -void MasterService::ReleaseTenantQuota(const std::string& tenant_id, - uint64_t bytes) { - if (!enable_tenant_quota_ || bytes == 0) { +void MasterService::CommitAdditionalTenantQuota(const TenantId& tenant_id, + uint64_t bytes) { + if (!enable_multi_tenants_ || bytes == 0) { return; } - const auto normalized_tenant = NormalizeTenantId(tenant_id); - auto& shard = - tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)]; - bool recompute_needed = false; - { - std::lock_guard lock(shard.mutex); - auto it = shard.tenants.find(normalized_tenant); - if (it == shard.tenants.end()) { - LOG(ERROR) << "tenant quota release mismatch tenant=" - << normalized_tenant << ", bytes=" << bytes - << ", used=0"; - return; - } - auto& state = it->second; - if (state.used_bytes < bytes) { - LOG(ERROR) << "tenant quota release mismatch tenant=" - << normalized_tenant << ", bytes=" << bytes - << ", used=" << state.used_bytes; - return; - } - state.used_bytes -= bytes; - if (state.committed_count > 0) { - --state.committed_count; - } - if (!state.has_explicit_policy && state.used_bytes == 0 && - state.reserved_bytes == 0 && state.committed_count == 0) { - shard.tenants.erase(it); - recompute_needed = true; - } else { - state.over_quota = - state.effective_quota_bytes != - std::numeric_limits::max() && - static_cast(state.used_bytes) + - state.reserved_bytes > - state.effective_quota_bytes; - } - } - if (recompute_needed) { - RecomputeTenantEffectiveQuotas(); + if (!tenant_quota_table_.CommitAdditional(tenant_id, bytes)) { + LOG(ERROR) << "tenant quota additional commit mismatch tenant=" + << tenant_id.value() << ", bytes=" << bytes; } } -void MasterService::ReleaseTenantQuotaPartial(const std::string& tenant_id, - uint64_t bytes) { - if (!enable_tenant_quota_ || bytes == 0) { +void MasterService::IncrementTenantMetadataObjectCount( + const TenantId& tenant_id) { + if (!enable_multi_tenants_) { return; } - const auto normalized_tenant = NormalizeTenantId(tenant_id); - auto& shard = - tenant_quota_shards_[getTenantQuotaShardIndex(normalized_tenant)]; - std::lock_guard lock(shard.mutex); - auto it = shard.tenants.find(normalized_tenant); - if (it == shard.tenants.end()) { - LOG(ERROR) << "tenant quota partial release mismatch tenant=" - << normalized_tenant << ", bytes=" << bytes << ", used=0"; + tenant_quota_table_.IncrementMetadataObjectCount(tenant_id); +} + +void MasterService::DecrementTenantMetadataObjectCount( + const TenantId& tenant_id) { + if (!enable_multi_tenants_) { return; } - auto& state = it->second; - if (state.used_bytes < bytes) { - LOG(ERROR) << "tenant quota partial release mismatch tenant=" - << normalized_tenant << ", bytes=" << bytes - << ", used=" << state.used_bytes; - return; + if (!tenant_quota_table_.DecrementMetadataObjectCount(tenant_id)) { + LOG(WARNING) << "tenant metadata object count decrement mismatch " + << "tenant=" << tenant_id.value(); } - state.used_bytes -= bytes; - state.over_quota = - state.effective_quota_bytes != std::numeric_limits::max() && - static_cast(state.used_bytes) + - state.reserved_bytes > - state.effective_quota_bytes; } void MasterService::ReleaseCommittedQuotaCharge(ObjectMetadata& metadata, uint64_t bytes) { - if (!enable_tenant_quota_ || bytes == 0) { + if (!enable_multi_tenants_ || bytes == 0) { return; } const uint64_t release_bytes = @@ -1035,16 +1479,17 @@ void MasterService::ReleaseCommittedQuotaCharge(ObjectMetadata& metadata, } void MasterService::RebuildTenantQuotaUsageFromMetadata() { - if (!enable_tenant_quota_) { + if (!enable_multi_tenants_) { return; } - std::unordered_map used_by_tenant; - std::unordered_map committed_count_by_tenant; + TenantQuotaUsageMap usage; for (size_t i = 0; i < kNumShards; ++i) { MetadataShardAccessorRW shard(this, i); for (auto& [tenant_id, tenant_state] : shard->tenants) { for (auto& [_, metadata] : tenant_state.metadata) { + auto& tenant_usage = usage[tenant_id]; + ++tenant_usage.metadata_object_count; const uint64_t charge = CompletedMemoryQuotaCharge(metadata); metadata.reserved_quota_charge_bytes = 0; metadata.committed_quota_charge_bytes = charge; @@ -1052,38 +1497,29 @@ void MasterService::RebuildTenantQuotaUsageFromMetadata() { if (charge == 0) { continue; } - used_by_tenant[tenant_id] += charge; - committed_count_by_tenant[tenant_id]++; + tenant_usage.used_bytes += charge; + ++tenant_usage.committed_count; } } } - for (size_t i = 0; i < kNumTenantQuotaShards; ++i) { - auto& shard = tenant_quota_shards_[i]; - std::lock_guard lock(shard.mutex); - for (auto& [_, state] : shard.tenants) { - state.used_bytes = 0; - state.reserved_bytes = 0; - state.committed_count = 0; + for (const auto& [tenant_id, _] : usage) { + if (!tenant_quota_table_.IsTenantRegistered(tenant_id)) { + LOG(WARNING) + << "tenant " << tenant_id.value() + << " exists in metadata but has no connector quota policy; " + "creating orphan quota state"; } } - for (const auto& [tenant_id, used_bytes] : used_by_tenant) { - auto& shard = tenant_quota_shards_[getTenantQuotaShardIndex(tenant_id)]; - std::lock_guard lock(shard.mutex); - auto& state = shard.tenants[tenant_id]; - state.requested_quota_bytes = default_tenant_quota_bytes_; - state.used_bytes = used_bytes; - state.committed_count = committed_count_by_tenant[tenant_id]; - } - RecomputeTenantEffectiveQuotas(); + std::lock_guard recompute_lock(tenant_quota_recompute_mutex_); + const uint64_t capacity = GetTenantQuotaAllocatableCapacityBytes(); + tenant_quota_table_.RebuildUsage(usage, capacity); } std::optional MasterService::GetGroupRoute( - const std::string& tenant_id, const std::string& key) const { - const auto normalized_tenant = NormalizeTenantId(tenant_id); + const TenantId& tenant_id, const std::string& key) const { std::shared_lock lock(group_routing_mutex_); - auto it = - object_group_ids_.find(MakeTenantScopedKey(normalized_tenant, key)); + auto it = object_group_ids_.find(tenant_id.MakeScopedKey(key)); if (it == object_group_ids_.end()) { return std::nullopt; } @@ -1091,36 +1527,33 @@ std::optional MasterService::GetGroupRoute( } MasterService::ObjectOperationLock MasterService::AcquireObjectOperationLock( - const std::string& tenant_id, const std::string& key) { - const auto scoped_key = MakeTenantScopedKey(tenant_id, key); + const TenantId& tenant_id, const std::string& key) { + const auto scoped_key = tenant_id.MakeScopedKey(key); const auto stripe_idx = std::hash{}(scoped_key) % kObjectOperationLockStripes; return {std::unique_lock(object_operation_locks_[stripe_idx])}; } void MasterService::RegisterGroupMember(TenantState& tenant_state, - const std::string& tenant_id, + const TenantId& tenant_id, const std::string& key, const std::string& group_id) { if (group_id.empty()) { return; } - const auto normalized_tenant = NormalizeTenantId(tenant_id); std::unique_lock lock(group_routing_mutex_); - object_group_ids_[MakeTenantScopedKey(normalized_tenant, key)] = group_id; - groups_needing_lease_refresh_.insert( - MakeTenantScopedKey(normalized_tenant, group_id)); + object_group_ids_[tenant_id.MakeScopedKey(key)] = group_id; + groups_needing_lease_refresh_.insert(tenant_id.MakeScopedKey(group_id)); tenant_state.group_members[group_id].insert(key); } void MasterService::UnregisterGroupMember(TenantState& tenant_state, - const std::string& tenant_id, + const TenantId& tenant_id, const std::string& key, const std::string& group_id) { if (group_id.empty()) { return; } - const auto normalized_tenant = NormalizeTenantId(tenant_id); bool group_empty = false; auto group_it = tenant_state.group_members.find(group_id); if (group_it != tenant_state.group_members.end()) { @@ -1131,14 +1564,12 @@ void MasterService::UnregisterGroupMember(TenantState& tenant_state, } } std::unique_lock lock(group_routing_mutex_); - auto route_it = - object_group_ids_.find(MakeTenantScopedKey(normalized_tenant, key)); + auto route_it = object_group_ids_.find(tenant_id.MakeScopedKey(key)); if (route_it != object_group_ids_.end() && route_it->second == group_id) { object_group_ids_.erase(route_it); } if (group_empty) { - groups_needing_lease_refresh_.erase( - MakeTenantScopedKey(normalized_tenant, group_id)); + groups_needing_lease_refresh_.erase(tenant_id.MakeScopedKey(group_id)); } } @@ -1152,7 +1583,8 @@ bool MasterService::HasCompletedMemoryCacheReplica( bool MasterService::HasCompletedDiskCacheReplica( const ObjectMetadata& metadata) { return metadata.HasReplica([](const Replica& replica) { - return replica.is_disk_replica() && replica.is_completed(); + return (replica.is_disk_replica() || replica.is_local_disk_replica()) && + replica.is_completed(); }); } @@ -1218,17 +1650,245 @@ std::vector MasterService::PopReplicasWithCacheTotalAccounting( size_t MasterService::EraseReplicasWithCacheTotalAccounting( ObjectMetadata& metadata, - const std::function& pred_fn) { + const std::function& pred_fn, + std::vector* erased_replica_ids) { auto erased_replicas = PopReplicasWithCacheTotalAccounting(metadata, pred_fn); + if (erased_replica_ids != nullptr) { + erased_replica_ids->reserve(erased_replica_ids->size() + + erased_replicas.size()); + for (const auto& replica : erased_replicas) { + erased_replica_ids->push_back(replica.id()); + } + } + // Release SSD/local-disk usage for any local-disk replicas being removed. + // No-op for memory/noF replicas, so it is safe to call unconditionally. + ReleaseLocalDiskUsage(erased_replicas); return erased_replicas.size(); } +void MasterService::FinalizeRemovedReplicasAfterDurable( + const OpLogEntry& durable_entry, const std::vector& replica_ids, + QuotaEraseMode quota_mode) { + if (replica_ids.empty()) { + return; + } + + std::shared_lock shared_lock(snapshot_mutex_); + const TenantId tenant_id(durable_entry.tenant_id); + const size_t shard_idx = + getMetadataShardIndex(tenant_id, durable_entry.object_key); + MetadataShardAccessorRW shard(this, shard_idx); + auto tenant_it = shard->tenants.find(tenant_id); + if (tenant_it == shard->tenants.end()) { + return; + } + auto& tenant_state = tenant_it->second; + auto metadata_it = tenant_state.metadata.find(durable_entry.object_key); + if (metadata_it == tenant_state.metadata.end()) { + return; + } + + std::unordered_set ids(replica_ids.begin(), replica_ids.end()); + auto& metadata = metadata_it->second; + auto erased_replicas = PopReplicasWithCacheTotalAccounting( + metadata, [&ids](const Replica& replica) { + return replica.status() == ReplicaStatus::REMOVED && + ids.contains(replica.id()); + }); + if (erased_replicas.empty()) { + return; + } + std::vector erased_replica_ids; + erased_replica_ids.reserve(erased_replicas.size()); + for (const auto& replica : erased_replicas) { + erased_replica_ids.push_back(replica.id()); + } + const uint64_t erased_memory_replicas = static_cast(std::count_if( + erased_replicas.begin(), erased_replicas.end(), + [](const Replica& replica) { return replica.is_memory_replica(); })); + if (erased_memory_replicas > 0) { + ReleaseCommittedQuotaCharge( + metadata, SaturatingMultiply(static_cast(metadata.size), + erased_memory_replicas)); + } + const bool erased_local_disk = std::any_of( + erased_replicas.begin(), erased_replicas.end(), + [](const Replica& replica) { return replica.is_local_disk_replica(); }); + ReleaseLocalDiskUsage(erased_replicas); + if (erased_local_disk) { + shard.OnDiskReplicaRemoved(erased_local_disk, metadata); + } + CancelPromotionTaskForRemovedReplicas(tenant_state, metadata, + erased_replica_ids); + if (!metadata.IsValid()) { + EraseMetadata(tenant_state, metadata_it, tenant_id, quota_mode, &shard); + if (tenant_state.Empty()) { + shard->tenants.erase(tenant_it); + } + } +} + +void MasterService::FinalizeMetadataEraseAfterDurable( + const OpLogEntry& durable_entry, QuotaEraseMode quota_mode) { + std::shared_lock shared_lock(snapshot_mutex_); + const TenantId tenant_id(durable_entry.tenant_id); + const size_t shard_idx = + getMetadataShardIndex(tenant_id, durable_entry.object_key); + MetadataShardAccessorRW shard(this, shard_idx); + auto tenant_it = shard->tenants.find(tenant_id); + if (tenant_it == shard->tenants.end()) { + return; + } + auto& tenant_state = tenant_it->second; + auto metadata_it = tenant_state.metadata.find(durable_entry.object_key); + if (metadata_it == tenant_state.metadata.end()) { + return; + } + EraseMetadata(tenant_state, metadata_it, tenant_id, quota_mode, &shard); + if (tenant_state.Empty()) { + shard->tenants.erase(tenant_it); + } +} + +void MasterService::FinalizeExpiredProcessingReplicasAfterDurable( + const OpLogEntry& durable_entry, + const std::chrono::system_clock::time_point& ttl) { + std::shared_lock shared_lock(snapshot_mutex_); + const TenantId tenant_id(durable_entry.tenant_id); + MetadataAccessorRW accessor(this, MakeObjectIdentityForRequest( + durable_entry.object_key, tenant_id)); + if (!accessor.Exists()) { + return; + } + + auto& metadata = accessor.Get(); + auto replicas = PopReplicasWithCacheTotalAccounting( + metadata, &Replica::fn_is_processing); + if (!replicas.empty()) { + std::lock_guard lock(discarded_replicas_mutex_); + discarded_replicas_.emplace_back(std::move(replicas), ttl); + } + if (!metadata.IsValid()) { + accessor.Erase(); + } else if (accessor.InProcessing()) { + accessor.EraseFromProcessing(); + } +} + +void MasterService::FinalizeExpiredReplicationTaskAfterDurable( + const OpLogEntry& durable_entry, ReplicaID source_id, + const std::vector& target_ids, + const std::chrono::system_clock::time_point& ttl) { + if (target_ids.empty()) { + return; + } + + std::shared_lock shared_lock(snapshot_mutex_); + const TenantId tenant_id(durable_entry.tenant_id); + MetadataAccessorRW accessor(this, MakeObjectIdentityForRequest( + durable_entry.object_key, tenant_id)); + if (!accessor.Exists()) { + return; + } + + auto& metadata = accessor.Get(); + if (auto source = metadata.GetReplicaByID(source_id); source != nullptr) { + source->dec_refcnt(); + } + + std::unordered_set ids(target_ids.begin(), target_ids.end()); + auto replicas = PopReplicasWithCacheTotalAccounting( + metadata, + [&ids](const Replica& replica) { return ids.contains(replica.id()); }); + if (!replicas.empty()) { + std::lock_guard lock(discarded_replicas_mutex_); + discarded_replicas_.emplace_back(std::move(replicas), ttl); + } + if (!metadata.IsValid()) { + accessor.Erase(); + } else if (accessor.HasReplicationTask()) { + AbortTenantQuota( + tenant_id, + accessor.GetReplicationTask().reserved_quota_charge_bytes); + accessor.EraseReplicationTask(); + } +} + +MasterService::StaleHandleCleanupPlan +MasterService::BuildStaleHandleCleanupPlan( + const ObjectMetadata& metadata, + const std::unordered_set>& alive_clients) const { + StaleHandleCleanupPlan plan; + bool has_valid_after_cleanup = false; + for (const auto& replica : metadata.GetAllReplicas()) { + const bool stale = + (replica.has_invalid_mem_handle() || + replica.has_invalid_nof_handle() || + replica.has_stale_local_disk_client(alive_clients)) && + replica.is_completed(); + if (stale) { + plan.removed_ids.push_back(replica.id()); + continue; + } + if (replica.status() == ReplicaStatus::COMPLETE) { + plan.remaining.push_back(replica.get_descriptor()); + } + if (!replica.is_memory_replica() || !replica.has_invalid_mem_handle()) { + has_valid_after_cleanup = true; + } + } + plan.would_invalidate = metadata.size == 0 || !has_valid_after_cleanup; + return plan; +} + +tl::expected MasterService::PersistStaleHandleCleanupForHA( + const std::string& why, const TenantId& tenant_id, const std::string& key, + ObjectMetadata& metadata, const StaleHandleCleanupPlan& plan) { + if (plan.removed_ids.empty() || !enable_oplog_) { + return {}; + } + + const auto op_type = + plan.would_invalidate ? OpType::REMOVE : OpType::PUT_END; + const std::string payload = + plan.would_invalidate + ? std::string{} + : SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, plan.remaining, + metadata.group_id, metadata.data_type); + + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + return tl::make_unexpected(reservation.error()); + } + const std::unordered_set ids(plan.removed_ids.begin(), + plan.removed_ids.end()); + metadata.VisitReplicas( + [&ids](const Replica& replica) { return ids.contains(replica.id()); }, + [](Replica& replica) { replica.mark_removed(); }); + auto result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), op_type, tenant_id.value(), key, + payload, + [this, + removed_ids = plan.removed_ids](const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable(durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + if (!result) { + LOG(WARNING) << why + << ": stale cleanup OpLog queue failed for key=" << key + << ", err=" << static_cast(result.error()); + return tl::make_unexpected(result.error()); + } + return {}; +} + std::unordered_map::iterator MasterService::EraseMetadata( TenantState& tenant_state, std::unordered_map::iterator it, - const std::string& tenant_id) { + const TenantId& tenant_id) { return EraseMetadata(tenant_state, it, tenant_id, QuotaEraseMode::kFull); } @@ -1236,10 +1896,55 @@ std::unordered_map::iterator MasterService::EraseMetadata( TenantState& tenant_state, std::unordered_map::iterator it, - const std::string& tenant_id, QuotaEraseMode quota_mode) { + const TenantId& tenant_id, QuotaEraseMode quota_mode) { + return EraseMetadata(tenant_state, it, tenant_id, quota_mode, nullptr); +} + +// EraseMetadata deletes the object metadata and also cleans up all +// associated per-key state: offloading_tasks (with dec_refcnt), +// processing_keys, replication_tasks, and promotion tasks. +// Callers no longer need to clean these up manually before calling. +std::unordered_map::iterator +MasterService::EraseMetadata( + TenantState& tenant_state, + std::unordered_map::iterator it, + const TenantId& tenant_id, QuotaEraseMode quota_mode, + MetadataShardAccessorRW* shard) { + bool had_completed_disk = it->second.HasReplica([](const Replica& r) { + return r.is_local_disk_replica() && r.is_completed(); + }); const std::string key = it->first; const std::string group_id = it->second.group_id; auto& metadata = it->second; + + // Clean up offloading_task + dec_refcnt before erasing metadata. + // When BatchEvict deletes metadata, Store Worker may still have an + // in-flight offload for this key. Without this cleanup the task + // becomes an orphan that only expires after 600s. + auto offload_it = tenant_state.offloading_tasks.find(key); + if (offload_it != tenant_state.offloading_tasks.end()) { + auto source = metadata.GetReplicaByID(offload_it->second.source_id); + if (source != nullptr) { + source->dec_refcnt(); + } + tenant_state.offloading_tasks.erase(offload_it); + + // Mirror entry in local_disk_segment.offloading_objects must be + // dropped too, otherwise the next OffloadObjectHeartbeat drains a + // task-less key back to the client and produces an orphan bucket. + const std::string scoped_key = tenant_id.MakeScopedKey(key); + ScopedLocalDiskSegmentAccess ssd_access = + segment_manager_.getLocalDiskSegmentAccess(); + for (auto& [_, segment] : ssd_access.getClientLocalDiskSegment()) { + MutexLocker locker(&segment->offloading_mutex_); + segment->offloading_objects.erase(scoped_key); + } + } + tenant_state.processing_keys.erase(key); + tenant_state.replication_tasks.erase(key); + ErasePromotionTaskIfPresent(tenant_state, key, tenant_id); + + ReleaseLocalDiskUsage(metadata.GetAllReplicas()); AccountCacheTotalRemoval(metadata); switch (quota_mode) { case QuotaEraseMode::kFull: @@ -1257,10 +1962,43 @@ MasterService::EraseMetadata( break; } auto next = tenant_state.metadata.erase(it); + DecrementTenantMetadataObjectCount(tenant_id); + if (had_completed_disk && shard) { + shard->OnDiskReplicaRemoved(had_completed_disk); + } UnregisterGroupMember(tenant_state, tenant_id, key, group_id); return next; } +void MasterService::ReleaseLocalDiskUsage( + const std::vector& replicas) { + std::unordered_map> bytes_by_client; + for (const auto& replica : replicas) { + if (!replica.is_local_disk_replica()) { + continue; + } + const auto descriptor = + replica.get_descriptor().get_local_disk_descriptor(); + if (descriptor.object_size > 0) { + bytes_by_client[descriptor.client_id] += descriptor.object_size; + } + } + if (bytes_by_client.empty()) { + return; + } + + ScopedLocalDiskSegmentAccess ssd_access = + segment_manager_.getLocalDiskSegmentAccess(); + auto& client_segments = ssd_access.getClientLocalDiskSegment(); + for (const auto& [client_id, bytes] : bytes_by_client) { + auto disk_it = client_segments.find(client_id); + if (disk_it != client_segments.end()) { + disk_it->second->ssd_used_bytes.fetch_sub( + bytes, std::memory_order_relaxed); + } + } +} + void MasterService::RebuildGroupRoutingIndex() { std::unordered_map rebuilt_group_ids; std::unordered_set groups_needing_refresh; @@ -1273,10 +2011,10 @@ void MasterService::RebuildGroupRoutingIndex() { continue; } tenant_state.group_members[metadata.group_id].insert(key); - rebuilt_group_ids[MakeTenantScopedKey(tenant_id, key)] = + rebuilt_group_ids[tenant_id.MakeScopedKey(key)] = metadata.group_id; groups_needing_refresh.insert( - MakeTenantScopedKey(tenant_id, metadata.group_id)); + tenant_id.MakeScopedKey(metadata.group_id)); } } } @@ -1299,9 +2037,9 @@ void MasterService::GrantLeaseForGroup(const TenantState& tenant_state, default_kv_soft_pin_ttl_); if (!needs_refresh) { std::shared_lock lock(group_routing_mutex_); - needs_refresh = groups_needing_lease_refresh_.find(MakeTenantScopedKey( - metadata.tenant_id, metadata.group_id)) != - groups_needing_lease_refresh_.end(); + needs_refresh = + groups_needing_lease_refresh_.find(metadata.tenant_id.MakeScopedKey( + metadata.group_id)) != groups_needing_lease_refresh_.end(); } if (!needs_refresh) { return; @@ -1326,7 +2064,7 @@ void MasterService::GrantLeaseForGroup(const TenantState& tenant_state, { std::unique_lock lock(group_routing_mutex_); groups_needing_lease_refresh_.erase( - MakeTenantScopedKey(metadata.tenant_id, metadata.group_id)); + metadata.tenant_id.MakeScopedKey(metadata.group_id)); } } @@ -1343,13 +2081,57 @@ void MasterService::ClearInvalidHandles( auto& tenant_state = tenant_it->second; auto it = tenant_state.metadata.begin(); while (it != tenant_state.metadata.end()) { - if (CleanupStaleHandles(it->second, alive_clients)) { - tenant_state.processing_keys.erase(it->first); - tenant_state.replication_tasks.erase(it->first); - tenant_state.offloading_tasks.erase(it->first); - ErasePromotionTaskIfPresent(tenant_state, it->first, - tenant_it->first); - it = EraseMetadata(tenant_state, it, tenant_it->first); + const auto cleanup_plan = + BuildStaleHandleCleanupPlan(it->second, alive_clients); + if (!cleanup_plan.removed_ids.empty()) { + if (enable_ha_) { + if (enable_oplog_) { + auto persist_result = + PersistStaleHandleCleanupForHA( + "ClearInvalidHandles", tenant_it->first, + it->first, it->second, cleanup_plan); + if (!persist_result) { + ++it; + continue; + } + ++it; + continue; + } + } + if (CleanupStaleHandles(tenant_state, it->second, + alive_clients, &shard)) { + it = EraseMetadata(tenant_state, it, tenant_it->first, + QuotaEraseMode::kFull, &shard); + } else { + ++it; + } + } else if (!it->second.IsValid()) { + if (enable_ha_) { + if (enable_oplog_) { + auto persist_result = + AppendOpLogWithDurableFinalize( + OpType::REMOVE, tenant_it->first.value(), + it->first, {}, + [this](const OpLogEntry& durable_entry) { + FinalizeMetadataEraseAfterDurable( + durable_entry, + QuotaEraseMode::kFull); + }); + if (!persist_result) { + LOG(WARNING) + << "ClearInvalidHandles(last replica)" + << ": REMOVE persist failed for key=" + << it->first << ", err=" + << static_cast(persist_result.error()); + ++it; + continue; + } + ++it; + continue; + } + } + it = EraseMetadata(tenant_state, it, tenant_it->first, + QuotaEraseMode::kFull, &shard); } else { ++it; } @@ -1411,6 +2193,14 @@ auto MasterService::UnmountSegment(const UUID& segment_id, // 2. Remove the metadata of the related objects ClearInvalidHandles(); + // Cache endpoint before commit removes segment from registry. + std::string segment_name; + std::string te_endpoint; + if (!segment_manager_.GetSegmentBasicInfo(segment_id, segment_name, + te_endpoint)) { + return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND); + } + // 3. Commit the unmount operation { ScopedSegmentAccess segment_access = @@ -1421,6 +2211,14 @@ auto MasterService::UnmountSegment(const UUID& segment_id, return tl::make_unexpected(err); } } + + if (enable_oplog_ && ordered_oplog_writer_ && !te_endpoint.empty()) { + SegmentUnmountOp op{te_endpoint}; + auto bytes = struct_pack::serialize(op); + PersistSegmentOpForHAOrEnqueue("UnmountSegment", + OpType::SEGMENT_UNMOUNT, te_endpoint, + std::string(bytes.begin(), bytes.end())); + } RecomputeTenantEffectiveQuotas(); return {}; } @@ -1508,36 +2306,34 @@ auto MasterService::UnmountNoFSegment(const UUID& segment_id, #endif } -auto MasterService::ExistKey(const std::string& key, - const std::string& tenant_id) +auto MasterService::ExistKey(const std::string& key, const TenantId& tenant_id) -> tl::expected { std::shared_lock shared_lock(snapshot_mutex_); - MetadataAccessorRO accessor(this, MakeObjectIdentity(key, tenant_id)); + MetadataAccessorRO accessor(this, + MakeObjectIdentityForRequest(key, tenant_id)); if (!accessor.Exists()) { VLOG(1) << "key=" << key << ", info=object_not_found"; return false; } const auto& metadata = accessor.Get(); - if (metadata.HasReplica(&Replica::fn_is_completed)) { - // Grant a lease to the object as it may be further used by the - // client. - auto* ts = accessor.GetTenantState(); - if (ts) { - GrantLeaseForGroup(*ts, key, metadata); - } else { - metadata.GrantLease(default_kv_lease_ttl_, - default_kv_soft_pin_ttl_); - } - return true; + if (!metadata.HasReplica(&Replica::fn_is_completed)) { + return false; } - return false; // If no complete replica is found, return false + // Grant a lease to the object as it may be further used by the client. + auto* ts = accessor.GetTenantState(); + if (ts) { + GrantLeaseForGroup(*ts, key, metadata); + } else { + metadata.GrantLease(default_kv_lease_ttl_, default_kv_soft_pin_ttl_); + } + return true; } std::vector> MasterService::BatchExistKey( - const std::vector& keys, const std::string& tenant_id) { - const std::string normalized_tenant = NormalizeTenantId(tenant_id); + const std::vector& keys, const TenantId& tenant_id) { + const TenantId& normalized_tenant = ResolveRequestTenantId(tenant_id); std::vector> results(keys.size()); if (keys.empty()) { return results; @@ -1549,7 +2345,7 @@ std::vector> MasterService::BatchExistKey( group_routing_mutex_); for (size_t i = 0; i < keys.size(); ++i) { auto route_it = object_group_ids_.find( - MakeTenantScopedKey(normalized_tenant, keys[i])); + normalized_tenant.MakeScopedKey(keys[i])); const size_t shard_idx = route_it == object_group_ids_.end() ? getShardIndex(normalized_tenant, keys[i]) @@ -1558,7 +2354,7 @@ std::vector> MasterService::BatchExistKey( } } - const size_t start_shard = RandomIndex(kNumShards); + const size_t start_shard = randomIndex(kNumShards); for (size_t scanned = 0; scanned < kNumShards; ++scanned) { const size_t shard_idx = (start_shard + kNumShards - scanned) % kNumShards; @@ -1592,21 +2388,21 @@ std::vector> MasterService::BatchExistKey( } const auto& metadata = it->second; - if (metadata.HasReplica(&Replica::fn_is_completed)) { - GrantLeaseForGroup(tenant_state, key, metadata); - results[i] = true; - } else { + if (!metadata.HasReplica(&Replica::fn_is_completed)) { results[i] = false; + continue; } + GrantLeaseForGroup(tenant_state, key, metadata); + results[i] = true; } } return results; } -auto MasterService::GetAllKeys(const std::string& tenant_id) +auto MasterService::GetAllKeys(const TenantId& tenant_id) -> tl::expected, ErrorCode> { std::vector all_keys; - const auto normalized_tenant = NormalizeTenantId(tenant_id); + const TenantId& normalized_tenant = ResolveRequestTenantId(tenant_id); for (size_t i = 0; i < kNumShards; i++) { MetadataShardAccessorRO shard(this, i); auto tenant_it = shard->tenants.find(normalized_tenant); @@ -1723,6 +2519,146 @@ auto MasterService::QuerySegmentStatusById(const UUID& segment_id) return status; } +void MasterService::RestoreFromStandbySnapshot( + const std::vector& objects, + uint64_t initial_oplog_sequence_id, + const std::vector& segments) { + // The ordered writer initializes its sequence from durable_prefix. + (void)initial_oplog_sequence_id; + + // 2. Build allocator keepalive map for standby segments. + for (const auto& [segment, bytes] : standby_accounted_memory_bytes_) { + MasterMetricManager::instance().dec_allocated_mem_size( + segment, static_cast(bytes)); + } + standby_accounted_memory_bytes_.clear(); + standby_memory_segments_.clear(); + standby_allocator_keepalive_.clear(); + invalid_replica_endpoints_.clear(); + for (const auto& seg : segments) { + if (seg.is_memory_segment) { + standby_memory_segments_.push_back(seg); + auto allocator = std::make_shared( + seg.segment_name, seg.transport_endpoint); + standby_allocator_keepalive_[seg.transport_endpoint] = allocator; + if (seg.segment_name != seg.transport_endpoint) { + standby_allocator_keepalive_[seg.segment_name] = allocator; + } + } + if (!segment_manager_.HasSegmentByEndpoint(seg.transport_endpoint)) { + invalid_replica_endpoints_.insert(seg.transport_endpoint); + if (seg.segment_name != seg.transport_endpoint) { + invalid_replica_endpoints_.insert(seg.segment_name); + } + } + } + + // 3. Restore object metadata. + const auto resolve_standby_object = [](const StandbyObjectEntry& entry) { + auto [scoped_tenant_id, user_key] = TenantId::ParseScopedKey(entry.key); + TenantId tenant_id(entry.tenant_id); + if (tenant_id.IsDefault() && !scoped_tenant_id.IsDefault()) { + tenant_id = std::move(scoped_tenant_id); + } + return std::make_pair(std::move(tenant_id), std::move(user_key)); + }; + + std::unordered_map> + objects_by_shard; + for (const auto& entry : objects) { + auto [tenant_id, user_key] = resolve_standby_object(entry); + if (!tenant_id.IsValid()) { + LOG(WARNING) << "RestoreFromStandbySnapshot: invalid tenant_id=" + << entry.tenant_id << ", key=" << entry.key + << ", skipping"; + continue; + } + const auto shard_idx = entry.metadata.group_id.empty() + ? getShardIndex(tenant_id, user_key) + : getShardIndex(entry.metadata.group_id); + objects_by_shard[shard_idx].push_back(&entry); + } + + for (const auto& [shard_idx, shard_objects] : objects_by_shard) { + MetadataShardAccessorRW shard(this, shard_idx); + auto now = std::chrono::system_clock::now(); + for (const auto* entry_ptr : shard_objects) { + const auto& entry = *entry_ptr; + auto [tenant_id, user_key] = resolve_standby_object(entry); + const auto& standby_meta = entry.metadata; + std::vector replicas; + replicas.reserve(standby_meta.replicas.size()); + + for (const auto& desc : standby_meta.replicas) { + if (desc.is_memory_replica()) { + const auto& mem_desc = desc.get_memory_descriptor(); + const std::string& endpoint = + mem_desc.buffer_descriptor.transport_endpoint_; + auto it = standby_allocator_keepalive_.find(endpoint); + if (it != standby_allocator_keepalive_.end()) { + auto alloc = it->second; + replicas.emplace_back( + std::make_unique( + alloc, mem_desc.buffer_descriptor), + desc.status); + MasterMetricManager::instance().inc_allocated_mem_size( + alloc->getSegmentName(), + static_cast( + mem_desc.buffer_descriptor.size_)); + standby_accounted_memory_bytes_ + [alloc->getSegmentName()] += + mem_desc.buffer_descriptor.size_; + } else { + invalid_replica_endpoints_.insert(endpoint); + } + } else if (desc.is_nof_replica()) { + const auto& nof_desc = desc.get_nof_descriptor(); + const std::string& endpoint = + nof_desc.buffer_descriptor.transport_endpoint_; + auto& alloc = standby_allocator_keepalive_[endpoint]; + if (!alloc) { + alloc = std::make_shared( + endpoint, endpoint); + } + replicas.emplace_back( + std::make_unique( + alloc, nof_desc.buffer_descriptor), + desc.status, ReplicaType::NOF_SSD); + } else if (desc.is_disk_replica()) { + const auto& disk_desc = desc.get_disk_descriptor(); + replicas.emplace_back(disk_desc.file_path, + disk_desc.object_size, desc.status); + } else if (desc.is_local_disk_replica()) { + const auto& local_disk_desc = + desc.get_local_disk_descriptor(); + replicas.emplace_back( + local_disk_desc.client_id, local_disk_desc.object_size, + local_disk_desc.transport_endpoint, desc.status); + } + } + + auto& tenant_state = shard->tenants[tenant_id]; + tenant_state.metadata.emplace( + std::piecewise_construct, std::forward_as_tuple(user_key), + std::forward_as_tuple( + standby_meta.client_id, now, standby_meta.size, + std::move(replicas), false, false, standby_meta.data_type, + standby_meta.group_id, tenant_id, user_key)); + if (!standby_meta.group_id.empty()) { + RegisterGroupMember(tenant_state, tenant_id, user_key, + standby_meta.group_id); + } + tenant_state.processing_keys.erase(user_key); + } + } + + // 4. Log the result. + LOG(INFO) << "Restored from standby: " << objects.size() << " objects, " + << segments.size() + << " segments, initial_seq_id=" << initial_oplog_sequence_id + << ", invalid_endpoints=" << invalid_replica_endpoints_.size(); +} + auto MasterService::QueryIp(const UUID& client_id) -> tl::expected, ErrorCode> { ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); @@ -1778,21 +2714,35 @@ auto MasterService::BatchReplicaClear( const std::vector& object_keys, const UUID& client_id, const std::string& segment_name) -> tl::expected, ErrorCode> { + return BatchReplicaClear(object_keys, client_id, segment_name, "default"); +} + +auto MasterService::BatchReplicaClear( + const std::vector& object_keys, const UUID& client_id, + const std::string& segment_name, const std::string& tenant_id) + -> tl::expected, ErrorCode> { std::shared_lock shared_lock(snapshot_mutex_); std::vector cleared_keys; cleared_keys.reserve(object_keys.size()); const bool clear_all_segments = segment_name.empty(); + const TenantId requested_tenant(tenant_id); + if (!requested_tenant.IsValid()) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + const TenantId& normalized_tenant = + ResolveRequestTenantId(requested_tenant); for (const auto& key : object_keys) { if (key.empty()) { - LOG(WARNING) << "BatchReplicaClear: empty key, skipping"; + LOG(WARNING) << "BatchReplicaClear: tenant=" << normalized_tenant + << " empty key, skipping"; continue; } - // BatchReplicaClear is a default-tenant compatibility/admin helper. - MetadataAccessorRW accessor(this, MakeObjectIdentity(key, "default")); + MetadataAccessorRW accessor(this, + MakeObjectIdentity(key, normalized_tenant)); if (!accessor.Exists()) { - LOG(WARNING) << "BatchReplicaClear: key=" << key - << " not found, skipping"; + LOG(WARNING) << "BatchReplicaClear: tenant=" << normalized_tenant + << " key=" << key << " not found, skipping"; continue; } @@ -1800,8 +2750,8 @@ auto MasterService::BatchReplicaClear( // Security check: Ensure the requesting client owns the object. if (metadata.client_id != client_id) { - LOG(WARNING) << "BatchReplicaClear: key=" << key - << " belongs to different client_id=" + LOG(WARNING) << "BatchReplicaClear: tenant=" << normalized_tenant + << " key=" << key << " belongs to different client_id=" << metadata.client_id << ", expected=" << client_id << ", skipping"; continue; @@ -1809,8 +2759,8 @@ auto MasterService::BatchReplicaClear( // Safety check: Do not clear an object that has an active lease. if (!metadata.IsLeaseExpired()) { - LOG(WARNING) << "BatchReplicaClear: key=" << key - << " has active lease, skipping"; + LOG(WARNING) << "BatchReplicaClear: tenant=" << normalized_tenant + << " key=" << key << " has active lease, skipping"; continue; } @@ -1819,20 +2769,57 @@ auto MasterService::BatchReplicaClear( // indicate an ongoing Put operation, and clearing during this time // could lead to an inconsistent state or interfere with the write. if (!metadata.AllReplicas(&Replica::fn_is_completed)) { - LOG(WARNING) << "BatchReplicaClear: key=" << key - << " has incomplete replicas, skipping"; + LOG(WARNING) + << "BatchReplicaClear: tenant=" << normalized_tenant + << " key=" << key << " has incomplete replicas, skipping"; continue; } + if (enable_ha_) { + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + continue; + } + std::vector removed_ids; + metadata.VisitReplicas( + &Replica::fn_is_completed, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + auto persist_result = + AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::REMOVE, + normalized_tenant.value(), key, {}, + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + if (!persist_result) { + continue; + } + cleared_keys.emplace_back(key); + VLOG(1) + << "BatchReplicaClear: tenant=" << normalized_tenant + << " successfully cleared all replicas for key=" << key + << " for client_id=" << client_id; + continue; + } + } + // Erase the entire metadata (all replicas will be deallocated) + // accessor.Erase() internally calls EraseMetadata which already + // decrements disk_object_count via OnDiskReplicaRemoved. accessor.Erase(); cleared_keys.emplace_back(key); - VLOG(1) << "BatchReplicaClear: successfully cleared all replicas " - "for key=" - << key << " for client_id=" << client_id; + VLOG(1) << "BatchReplicaClear: tenant=" << normalized_tenant + << " successfully cleared all replicas for key=" << key + << " for client_id=" << client_id; } else { // Clear only replicas on the specified segment_name - bool has_replica_on_segment = false; const auto match_replica_on_segment = [&](const Replica& replica) -> bool { if (!replica.is_completed()) { @@ -1848,28 +2835,104 @@ auto MasterService::BatchReplicaClear( return false; }; - has_replica_on_segment = - metadata.HasReplica(match_replica_on_segment); - - if (!has_replica_on_segment) { + if (!metadata.HasReplica(match_replica_on_segment)) { LOG(WARNING) - << "BatchReplicaClear: key=" << key + << "BatchReplicaClear: tenant=" << normalized_tenant + << " key=" << key << " has no replica on segment_name=" << segment_name << ", skipping"; continue; } + bool had_completed_disk_on_segment = + metadata.HasReplica([&segment_name](const Replica& r) { + if (!r.is_local_disk_replica() || !r.is_completed()) + return false; + for (const auto& name : r.get_segment_names()) { + if (name.has_value() && name.value() == segment_name) + return true; + } + return false; + }); + + if (enable_ha_) { + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + continue; + } + auto remaining = BuildRemainingReplicaDescriptors( + metadata, + [&match_replica_on_segment](const Replica& r) { + return match_replica_on_segment(r); + }); + std::vector removed_ids; + metadata.VisitReplicas( + match_replica_on_segment, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + + tl::expected persist_result; + if (remaining.empty()) { + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::REMOVE, + normalized_tenant.value(), key, {}, + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + } else { + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::PUT_END, + normalized_tenant.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, remaining, + metadata.group_id, metadata.data_type), + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + } + if (!persist_result) { + continue; + } + cleared_keys.emplace_back(key); + VLOG(1) << "BatchReplicaClear: tenant=" << normalized_tenant + << " successfully cleared replicas on segment_name=" + << segment_name << " for key=" << key + << " for client_id=" << client_id; + continue; + } + } + EraseReplicasWithCacheTotalAccounting(metadata, match_replica_on_segment); + if (had_completed_disk_on_segment && + !metadata.HasReplica([](const Replica& r) { + return r.is_local_disk_replica() && r.is_completed(); + })) { + auto& shard = accessor.GetShard(); + shard.OnDiskReplicaRemoved(had_completed_disk_on_segment, + metadata); + } + // If no valid replicas remain, erase the entire metadata + // accessor.Erase() internally calls EraseMetadata which already + // decrements disk_object_count via OnDiskReplicaRemoved. if (!metadata.IsValid()) { accessor.Erase(); } cleared_keys.emplace_back(key); - VLOG(1) << "BatchReplicaClear: successfully cleared replicas on " - "segment_name=" + VLOG(1) << "BatchReplicaClear: tenant=" << normalized_tenant + << " successfully cleared replicas on segment_name=" << segment_name << " for key=" << key << " for client_id=" << client_id; } @@ -1878,8 +2941,27 @@ auto MasterService::BatchReplicaClear( return cleared_keys; } +bool MasterService::IsReplicaReadable(const Replica& replica) const { + if (!replica.is_completed() || replica.has_invalid_mem_handle() || + replica.has_invalid_nof_handle()) { + return false; + } + const auto descriptor = replica.get_descriptor(); + std::optional endpoint; + if (descriptor.is_memory_replica()) { + endpoint = descriptor.get_memory_descriptor() + .buffer_descriptor.transport_endpoint_; + } else if (descriptor.is_nof_replica()) { + endpoint = descriptor.get_nof_descriptor() + .buffer_descriptor.transport_endpoint_; + } else if (descriptor.is_local_disk_replica()) { + endpoint = descriptor.get_local_disk_descriptor().transport_endpoint; + } + return !endpoint || !invalid_replica_endpoints_.contains(*endpoint); +} + auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern, - const std::string& tenant_id) + const TenantId& tenant_id) -> tl::expected< std::unordered_map>, ErrorCode> { @@ -1895,7 +2977,7 @@ auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern, } std::shared_lock shared_lock(snapshot_mutex_); - const auto normalized_tenant = NormalizeTenantId(tenant_id); + const TenantId& normalized_tenant = ResolveRequestTenantId(tenant_id); for (size_t i = 0; i < kNumShards; ++i) { MetadataShardAccessorRO shard(this, i); auto tenant_it = shard->tenants.find(normalized_tenant); @@ -1906,7 +2988,9 @@ auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern, if (std::regex_search(key, pattern)) { std::vector replica_list; metadata.VisitReplicas( - &Replica::fn_is_completed, + [this](const Replica& replica) { + return IsReplicaReadable(replica); + }, [&replica_list](const Replica& replica) { replica_list.emplace_back(replica.get_descriptor()); }); @@ -1928,10 +3012,10 @@ auto MasterService::GetReplicaListByRegex(const std::string& regex_pattern, } auto MasterService::GetReplicaList(const std::string& key, - const std::string& tenant_id) + const TenantId& tenant_id) -> tl::expected { std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); + const auto object_id = MakeObjectIdentityForRequest(key, tenant_id); GetReplicaListResponse resp({}, default_kv_lease_ttl_); bool promotion_eligible = false; @@ -1948,11 +3032,19 @@ auto MasterService::GetReplicaList(const std::string& key, std::vector replica_list; metadata.VisitReplicas( - &Replica::fn_is_completed, [&replica_list](const Replica& replica) { + [this](const Replica& replica) { + return IsReplicaReadable(replica); + }, + [&replica_list](const Replica& replica) { replica_list.emplace_back(replica.get_descriptor()); }); if (replica_list.empty()) { + if (metadata.AllReplicas([](const Replica& replica) { + return replica.status() == ReplicaStatus::REMOVED; + })) { + return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + } LOG(WARNING) << "key=" << key << ", error=replica_not_ready"; return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY); } @@ -1960,8 +3052,13 @@ auto MasterService::GetReplicaList(const std::string& key, // TODO: NoF SSD support (ranhaojia) if (replica_list[0].is_memory_replica()) { MasterMetricManager::instance().inc_mem_cache_hit_nums(); - } else if (replica_list[0].is_disk_replica()) { + MasterMetricManager::instance().inc_mem_cache_hit_bytes( + static_cast(metadata.size)); + } else if (replica_list[0].is_local_disk_replica() || + replica_list[0].is_disk_replica()) { MasterMetricManager::instance().inc_file_cache_hit_nums(); + MasterMetricManager::instance().inc_file_cache_hit_bytes( + static_cast(metadata.size)); } MasterMetricManager::instance().inc_valid_get_nums(); // Grant a lease to the object so it will not be removed @@ -1987,7 +3084,8 @@ auto MasterService::GetReplicaList(const std::string& key, } resp = GetReplicaListResponse(std::move(replica_list), - default_kv_lease_ttl_); + default_kv_lease_ttl_, + metadata.object_checksum); } // RO accessor released. Safe to take a fresh RW accessor now. if (promotion_eligible) { @@ -1996,18 +3094,51 @@ auto MasterService::GetReplicaList(const std::string& key, return resp; } +auto MasterService::GetReplicaListForAdmin(const std::string& key, + const TenantId& tenant_id) + -> tl::expected { + assert(tenant_id.IsValid()); + const auto object_id = MakeObjectIdentity(key, tenant_id); + + std::shared_lock shared_lock(snapshot_mutex_); + MetadataAccessorRO accessor(this, object_id); + + if (!accessor.Exists()) { + VLOG(1) << "key=" << key << ", info=object_not_found"; + return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + } + const auto& metadata = accessor.Get(); + + std::vector replica_list; + metadata.VisitReplicas( + &Replica::fn_is_completed, [&replica_list](const Replica& replica) { + replica_list.emplace_back(replica.get_descriptor()); + }); + + if (replica_list.empty()) { + LOG(WARNING) << "key=" << key << ", error=replica_not_ready"; + return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY); + } + + return GetReplicaListResponse(std::move(replica_list), + default_kv_lease_ttl_, + metadata.object_checksum); +} + std::vector> MasterService::BatchGetReplicaList(const std::vector& keys, - const std::string& tenant_id) { + const TenantId& tenant_id) { using GetResult = tl::expected; + assert(tenant_id.IsValid()); + std::vector results( keys.size(), tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND)); if (keys.empty()) { return results; } - const auto normalized_tenant = NormalizeTenantId(tenant_id); + const TenantId& normalized_tenant = ResolveRequestTenantId(tenant_id); constexpr size_t kInvalidKeyIndex = std::numeric_limits::max(); std::array key_list_heads; key_list_heads.fill(kInvalidKeyIndex); @@ -2017,7 +3148,7 @@ MasterService::BatchGetReplicaList(const std::vector& keys, for (size_t i = keys.size(); i > 0; --i) { const size_t original_idx = i - 1; const auto scoped_key = - MakeTenantScopedKey(normalized_tenant, keys[original_idx]); + normalized_tenant.MakeScopedKey(keys[original_idx]); const auto route_it = object_group_ids_.find(scoped_key); const size_t shard_idx = route_it == object_group_ids_.end() @@ -2028,7 +3159,7 @@ MasterService::BatchGetReplicaList(const std::vector& keys, } } - const size_t start_shard = RandomIndex(kNumShards); + const size_t start_shard = randomIndex(kNumShards); for (size_t scanned = 0; scanned < kNumShards; ++scanned) { const size_t shard_idx = (start_shard + kNumShards - scanned) % kNumShards; @@ -2067,12 +3198,21 @@ MasterService::BatchGetReplicaList(const std::vector& keys, const auto& metadata = metadata_it->second; std::vector replica_list; metadata.VisitReplicas( - &Replica::fn_is_completed, + [this](const Replica& replica) { + return IsReplicaReadable(replica); + }, [&replica_list](const Replica& replica) { replica_list.emplace_back(replica.get_descriptor()); }); if (replica_list.empty()) { + if (metadata.AllReplicas([](const Replica& replica) { + return replica.status() == ReplicaStatus::REMOVED; + })) { + results[original_idx] = + tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + continue; + } LOG(WARNING) << "key=" << key << ", error=replica_not_ready"; results[original_idx] = @@ -2082,8 +3222,13 @@ MasterService::BatchGetReplicaList(const std::vector& keys, if (replica_list[0].is_memory_replica()) { MasterMetricManager::instance().inc_mem_cache_hit_nums(); - } else if (replica_list[0].is_disk_replica()) { + MasterMetricManager::instance().inc_mem_cache_hit_bytes( + static_cast(metadata.size)); + } else if (replica_list[0].is_local_disk_replica() || + replica_list[0].is_disk_replica()) { MasterMetricManager::instance().inc_file_cache_hit_nums(); + MasterMetricManager::instance().inc_file_cache_hit_bytes( + static_cast(metadata.size)); } MasterMetricManager::instance().inc_valid_get_nums(); GrantLeaseForGroup(tenant_state, key, metadata); @@ -2100,7 +3245,8 @@ MasterService::BatchGetReplicaList(const std::vector& keys, } results[original_idx] = GetReplicaListResponse( - std::move(replica_list), default_kv_lease_ttl_); + std::move(replica_list), default_kv_lease_ttl_, + metadata.object_checksum); } } @@ -2112,12 +3258,101 @@ MasterService::BatchGetReplicaList(const std::vector& keys, return results; } +std::vector> +MasterService::BatchGetReplicaListForAdmin(const std::vector& keys, + const TenantId& tenant_id) { + using GetResult = tl::expected; + + assert(tenant_id.IsValid()); + + std::vector results( + keys.size(), tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND)); + if (keys.empty()) { + return results; + } + + const TenantId& normalized_tenant = tenant_id; + constexpr size_t kInvalidKeyIndex = std::numeric_limits::max(); + std::array key_list_heads; + key_list_heads.fill(kInvalidKeyIndex); + std::vector next_key_indexes(keys.size(), kInvalidKeyIndex); + { + std::shared_lock lock(group_routing_mutex_); + for (size_t i = keys.size(); i > 0; --i) { + const size_t original_idx = i - 1; + const auto scoped_key = + normalized_tenant.MakeScopedKey(keys[original_idx]); + const auto route_it = object_group_ids_.find(scoped_key); + const size_t shard_idx = + route_it == object_group_ids_.end() + ? getShardIndex(normalized_tenant, keys[original_idx]) + : getShardIndex(route_it->second); + next_key_indexes[original_idx] = key_list_heads[shard_idx]; + key_list_heads[shard_idx] = original_idx; + } + } + + const size_t start_shard = randomIndex(kNumShards); + for (size_t scanned = 0; scanned < kNumShards; ++scanned) { + const size_t shard_idx = + (start_shard + kNumShards - scanned) % kNumShards; + if (key_list_heads[shard_idx] == kInvalidKeyIndex) { + continue; + } + + std::shared_lock shared_lock(snapshot_mutex_); + { + MetadataShardAccessorRO shard(this, shard_idx); + const auto tenant_it = shard->tenants.find(normalized_tenant); + for (size_t original_idx = key_list_heads[shard_idx]; + original_idx != kInvalidKeyIndex; + original_idx = next_key_indexes[original_idx]) { + const std::string& key = keys[original_idx]; + + if (tenant_it == shard->tenants.end()) { + results[original_idx] = + tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + continue; + } + + const auto& tenant_state = tenant_it->second; + const auto metadata_it = tenant_state.metadata.find(key); + if (metadata_it == tenant_state.metadata.end() || + !metadata_it->second.IsValid()) { + results[original_idx] = + tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + continue; + } + + const auto& metadata = metadata_it->second; + std::vector replica_list; + metadata.VisitReplicas( + &Replica::fn_is_completed, + [&replica_list](const Replica& replica) { + replica_list.emplace_back(replica.get_descriptor()); + }); + + if (replica_list.empty()) { + results[original_idx] = + tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY); + continue; + } + + results[original_idx] = GetReplicaListResponse( + std::move(replica_list), default_kv_lease_ttl_, + metadata.object_checksum); + } + } + } + + return results; +} + auto MasterService::AllocateAndInsertMetadata( MetadataShardAccessorRW& shard, const UUID& client_id, const std::string& key, uint64_t value_length, const ReplicateConfig& config, const std::string& group_id, - const std::string& tenant_id, - const std::chrono::system_clock::time_point& now) + const TenantId& tenant_id, const std::chrono::system_clock::time_point& now) -> tl::expected, ErrorCode> { auto& tenant_state = shard->tenants[tenant_id]; if (tenant_state.metadata.contains(key)) { @@ -2144,20 +3379,61 @@ auto MasterService::AllocateAndInsertMetadata( size_t allocated_memory_replicas = 0; size_t allocated_nof_replicas = 0; if (config.replica_num > 0) { + const bool use_local_first = + allocation_strategy_type_ == AllocationStrategyType::LOCAL_FIRST && + config.replica_num == 1; + std::string writer_host_id; + if (use_local_first) { + writer_host_id = config.host_id.empty() ? GetClientHostId(client_id) + : config.host_id; + } + ScopedAllocatorAccess allocator_access = segment_manager_.getAllocatorAccess(); const auto& allocator_manager = allocator_access.getAllocatorManager(); std::vector preferred_segments; + auto append_preferred_segment = [&preferred_segments]( + const std::string& segment_name) { + if (!segment_name.empty() && + std::find(preferred_segments.begin(), preferred_segments.end(), + segment_name) == preferred_segments.end()) { + preferred_segments.push_back(segment_name); + } + }; if (!config.preferred_segment.empty()) { - preferred_segments.push_back(config.preferred_segment); - } else if (!config.preferred_segments.empty()) { - preferred_segments = config.preferred_segments; + append_preferred_segment(config.preferred_segment); + } else { + for (const auto& preferred_segment : config.preferred_segments) { + append_preferred_segment(preferred_segment); + } + } + if (!writer_host_id.empty()) { + auto host_ordered_segments = + allocator_access.GetHostOrderedSegments(writer_host_id, key); + for (const auto& segment_name : host_ordered_segments) { + append_preferred_segment(segment_name); + } + if (!host_ordered_segments.empty()) { + VLOG(1) << "key=" << key + << ", writer_host_id=" << writer_host_id + << ", local_first_preferred_segments=" + << host_ordered_segments.size(); + } + } + + const SsdMetricsProvider* ssd_provider = nullptr; + std::optional ssd_access; + if (allocation_strategy_type_ == + AllocationStrategyType::SSD_FREE_RATIO_FIRST) { + ssd_access.emplace(segment_manager_.getLocalDiskSegmentAccess()); + ssd_provider = &*ssd_access; } auto allocation_result = allocation_strategy_->Allocate( allocator_manager, value_length, config.replica_num, - preferred_segments); + preferred_segments, std::set(), ReplicaType::MEMORY, + ssd_provider); if (!allocation_result.has_value()) { VLOG(1) << "Failed to allocate replicas for key=" << key @@ -2239,6 +3515,21 @@ auto MasterService::AllocateAndInsertMetadata( return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); } + // Best-effort / flexible modes may pass the check above with fewer + // replicas than requested (see HasExpectedReplicaAllocation). Surface + // the degradation so callers and operators can detect the reduced + // redundancy instead of failing silently. + if (allocated_memory_replicas < config.replica_num || + allocated_nof_replicas < config.nof_replica_num) { + MasterMetricManager::instance().inc_put_start_partial_allocations(); + LOG(WARNING) << "key=" << key << ", action=put_start_partial_allocation" + << ", requested_memory_replicas=" << config.replica_num + << ", allocated_memory_replicas=" + << allocated_memory_replicas + << ", requested_nof_replicas=" << config.nof_replica_num + << ", allocated_nof_replicas=" << allocated_nof_replicas; + } + if (use_disk_replica_) { std::string file_path = ResolvePathFromKey(key, root_fs_dir_, cluster_id_); @@ -2280,6 +3571,7 @@ auto MasterService::AllocateAndInsertMetadata( abort_reserved_quota(); return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS); } + IncrementTenantMetadataObjectCount(tenant_id); it->second.reserved_quota_charge_bytes = reserved_quota_charge; RegisterGroupMember(tenant_state, tenant_id, key, group_id); tenant_state.processing_keys.insert(key); @@ -2288,11 +3580,16 @@ auto MasterService::AllocateAndInsertMetadata( } auto MasterService::PutStart(const UUID& client_id, const std::string& key, - const std::string& tenant_id, + const TenantId& tenant_id, const uint64_t slice_length, const ReplicateConfig& config) -> tl::expected, ErrorCode> { - const auto object_id = MakeObjectIdentity(key, tenant_id); + auto normalized_tenant_result = ResolveTenantIdForWrite(tenant_id); + if (!normalized_tenant_result) { + return tl::make_unexpected(normalized_tenant_result.error()); + } + const ObjectIdentity object_id{std::move(normalized_tenant_result.value()), + key}; if ((config.replica_num == 0 && config.nof_replica_num == 0) || key.empty() || slice_length == 0) { LOG(ERROR) << "key=" << key << ", replica_num=" << config.replica_num @@ -2318,6 +3615,8 @@ auto MasterService::PutStart(const UUID& client_id, const std::string& key, } #endif + UpdateClientHostId(client_id, config.host_id); + if ((memory_allocator_type_ == BufferAllocatorType::CACHELIB) && (slice_length > kMaxSliceSize)) { LOG(ERROR) << "key=" << key << ", slice_length=" << slice_length @@ -2337,87 +3636,145 @@ auto MasterService::PutStart(const UUID& client_id, const std::string& key, [[maybe_unused]] auto object_operation_lock = AcquireObjectOperationLock(object_id.tenant_id, object_id.user_key); - const auto now = std::chrono::system_clock::now(); - std::optional retry_shard_idx; - { - auto alive_clients = getAliveClientsSnapshot(); - std::shared_lock shared_lock(snapshot_mutex_); - const size_t lookup_shard_idx = - getMetadataShardIndex(object_id.tenant_id, object_id.user_key); - MetadataShardAccessorRW shard(this, lookup_shard_idx); - auto& tenant_state = shard->tenants[object_id.tenant_id]; - - auto it = tenant_state.metadata.find(key); - if (it != tenant_state.metadata.end()) { - if (CleanupStaleHandles(it->second, alive_clients)) { - tenant_state.processing_keys.erase(key); - tenant_state.replication_tasks.erase(key); - tenant_state.offloading_tasks.erase(key); - ErasePromotionTaskIfPresent(tenant_state, key, - object_id.tenant_id); - EraseMetadata(tenant_state, it, object_id.tenant_id); - it = tenant_state.metadata.end(); - } else { - auto& metadata = it->second; - if (metadata.HasReplica(&Replica::fn_is_completed) || - metadata.put_start_time + put_start_discard_timeout_sec_ >= - now) { - LOG(INFO) - << "key=" << key << ", info=object_already_exists"; - return tl::make_unexpected( - ErrorCode::OBJECT_ALREADY_EXISTS); + const uint64_t requested_quota_charge = + RequestedMemoryQuotaCharge(slice_length, config); + + auto attempt_once = + [&]() -> tl::expected, ErrorCode> { + std::unique_lock zero_charge_policy_lock( + tenant_quota_policy_mutex_, std::defer_lock); + if (ShouldProtectZeroChargeMetadataCreate(requested_quota_charge)) { + zero_charge_policy_lock.lock(); + auto latest_tenant_result = + ResolveTenantIdForWriteLocked(tenant_id); + if (!latest_tenant_result) { + return tl::make_unexpected(latest_tenant_result.error()); + } + } + + auto now = std::chrono::system_clock::now(); + std::optional retry_shard_idx; + { + auto alive_clients = getAliveClientsSnapshot(); + std::shared_lock shared_lock(snapshot_mutex_); + const size_t lookup_shard_idx = + getMetadataShardIndex(object_id.tenant_id, object_id.user_key); + MetadataShardAccessorRW shard(this, lookup_shard_idx); + auto& tenant_state = shard->tenants[object_id.tenant_id]; + + auto it = tenant_state.metadata.find(key); + if (it != tenant_state.metadata.end()) { + auto cleanup_plan = + BuildStaleHandleCleanupPlan(it->second, alive_clients); + if (!cleanup_plan.removed_ids.empty()) { + auto persist_result = PersistStaleHandleCleanupForHA( + "PutStart(stale cleanup)", object_id.tenant_id, key, + it->second, cleanup_plan); + if (!persist_result) { + return tl::make_unexpected(persist_result.error()); + } + if (enable_oplog_) { + return tl::make_unexpected( + ErrorCode::OBJECT_ALREADY_EXISTS); + } else if (CleanupStaleHandles(tenant_state, it->second, + alive_clients, &shard)) { + EraseMetadata(tenant_state, it, object_id.tenant_id, + QuotaEraseMode::kFull, &shard); + it = tenant_state.metadata.end(); + } } - auto replicas = - metadata.PopReplicas(&Replica::fn_is_processing); - if (!replicas.empty()) { - std::lock_guard lock(discarded_replicas_mutex_); - discarded_replicas_.emplace_back( - std::move(replicas), + if (it != tenant_state.metadata.end()) { + auto& metadata = it->second; + if (metadata.HasReplica(&Replica::fn_is_completed) || metadata.put_start_time + - put_start_release_timeout_sec_); + put_start_discard_timeout_sec_ >= + now) { + LOG(INFO) + << "key=" << key << ", info=object_already_exists"; + return tl::make_unexpected( + ErrorCode::OBJECT_ALREADY_EXISTS); + } + if (enable_oplog_ && ordered_oplog_writer_) { + auto err = + PersistRemoveForHA("PutStart(stale cleanup REMOVE)", + object_id.tenant_id, key); + if (!err) { + return tl::make_unexpected(err.error()); + } + } + auto replicas = PopReplicasWithCacheTotalAccounting( + metadata, &Replica::fn_is_processing); + if (!replicas.empty()) { + std::lock_guard lock(discarded_replicas_mutex_); + discarded_replicas_.emplace_back( + std::move(replicas), + metadata.put_start_time + + put_start_release_timeout_sec_); + } + EraseMetadata(tenant_state, it, object_id.tenant_id, + QuotaEraseMode::kFull, &shard); + it = tenant_state.metadata.end(); } - tenant_state.processing_keys.erase(key); - EraseMetadata(tenant_state, it, object_id.tenant_id); - it = tenant_state.metadata.end(); } - } - if (it == tenant_state.metadata.end()) { - const size_t target_shard_idx = - group_id.empty() - ? getShardIndex(object_id.tenant_id, object_id.user_key) - : getShardIndex(group_id); - if (target_shard_idx != lookup_shard_idx) { - retry_shard_idx = target_shard_idx; - if (tenant_state.Empty()) { - shard->tenants.erase(object_id.tenant_id); + if (it == tenant_state.metadata.end()) { + const size_t target_shard_idx = + group_id.empty() + ? getShardIndex(object_id.tenant_id, object_id.user_key) + : getShardIndex(group_id); + if (target_shard_idx != lookup_shard_idx) { + retry_shard_idx = target_shard_idx; + if (tenant_state.Empty()) { + shard->tenants.erase(object_id.tenant_id); + } + } else { + return AllocateAndInsertMetadata( + shard, client_id, key, slice_length, config, group_id, + object_id.tenant_id, now); } - } else { - return AllocateAndInsertMetadata(shard, client_id, key, - slice_length, config, group_id, - object_id.tenant_id, now); } } + + std::shared_lock shared_lock(snapshot_mutex_); + MetadataShardAccessorRW shard(this, retry_shard_idx.value()); + auto& retry_tenant_state = shard->tenants[object_id.tenant_id]; + if (GetGroupRoute(object_id.tenant_id, object_id.user_key) + .has_value() || + retry_tenant_state.metadata.contains(key)) { + LOG(INFO) << "key=" << key << ", info=object_already_exists"; + return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS); + } + return AllocateAndInsertMetadata(shard, client_id, key, slice_length, + config, group_id, object_id.tenant_id, + now); + }; + + for (int attempt = 0; attempt <= kMaxTenantQuotaEvictionRetries; + ++attempt) { + auto result = attempt_once(); + if (result.has_value() || + result.error() != ErrorCode::TENANT_QUOTA_EXCEEDED) { + return result; + } + if (attempt == kMaxTenantQuotaEvictionRetries) { + MasterMetricManager::instance().inc_tenant_quota_reject( + object_id.tenant_id.value(), "quota_exceeded"); + return result; + } + EvictTenantMemoryForQuota( + object_id.tenant_id, + tenant_quota_table_.ComputeDeficit(object_id.tenant_id, + requested_quota_charge)); } - std::shared_lock shared_lock(snapshot_mutex_); - MetadataShardAccessorRW shard(this, retry_shard_idx.value()); - auto& retry_tenant_state = shard->tenants[object_id.tenant_id]; - if (GetGroupRoute(object_id.tenant_id, object_id.user_key).has_value() || - retry_tenant_state.metadata.contains(key)) { - LOG(INFO) << "key=" << key << ", info=object_already_exists"; - return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS); - } - return AllocateAndInsertMetadata(shard, client_id, key, slice_length, - config, group_id, object_id.tenant_id, - now); + return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); } -auto MasterService::PutEnd(const UUID& client_id, const std::string& key, - const std::string& tenant_id, - ReplicaType replica_type) +auto MasterService::PutEnd(const UUID& client_id, const ObjectMeta& object_meta, + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected { + const auto& key = object_meta.key; std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); + const auto object_id = MakeObjectIdentityForRequest(key, tenant_id); MetadataAccessorRW accessor(this, object_id); if (!accessor.Exists()) { LOG(ERROR) << "key=" << key << ", error=object_not_found"; @@ -2431,26 +3788,61 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key, return tl::make_unexpected(ErrorCode::ILLEGAL_CLIENT); } + auto is_target_replica = [replica_type](const Replica& replica) { + if (replica_type == ReplicaType::ALL) { + return (replica.is_memory_replica() && + !replica.has_invalid_mem_handle()) || + (replica.is_nof_replica() && + !replica.has_invalid_nof_handle()); + } + if (replica_type == ReplicaType::MEMORY) { + return replica.is_memory_replica() && + !replica.has_invalid_mem_handle(); + } + if (replica_type == ReplicaType::NOF_SSD) { + return replica.is_nof_replica() && + !replica.has_invalid_nof_handle(); + } + return replica.type() == replica_type; + }; + + // A successful End removes the processing marker. Treat a retry as a + // no-op only when every replica targeted by that End is already COMPLETE. + // In particular, a promotion-owned PROCESSING replica keeps this check + // from accepting a MEMORY/ALL End and is never modified here. + if (!accessor.InProcessing()) { + bool has_target_replica = false; + bool all_target_replicas_complete = true; + for (const auto& replica : metadata.GetAllReplicas()) { + if (!is_target_replica(replica)) { + continue; + } + has_target_replica = true; + if (!replica.is_completed()) { + all_target_replicas_complete = false; + break; + } + } + if (has_target_replica && all_target_replicas_complete) { + return {}; + } + LOG(ERROR) << "key=" << key << ", error=no_primary_write_in_progress"; + return tl::make_unexpected(ErrorCode::INVALID_WRITE); + } + metadata.VisitReplicas( - [replica_type](const Replica& replica) { - if (replica_type == ReplicaType::ALL) { - return (replica.is_memory_replica() && - !replica.has_invalid_mem_handle()) || - (replica.is_nof_replica() && - !replica.has_invalid_nof_handle()); - } - if (replica_type == ReplicaType::MEMORY) { - return replica.is_memory_replica() && - !replica.has_invalid_mem_handle(); - } - if (replica_type == ReplicaType::NOF_SSD) { - return replica.is_nof_replica() && - !replica.has_invalid_nof_handle(); - } - return replica.type() == replica_type; + [&is_target_replica](const Replica& replica) { + return replica.is_processing() && is_target_replica(replica); }, [](Replica& replica) { replica.mark_complete(); }); + if (object_meta.object_checksum.has_value() || + replica_type == ReplicaType::ALL || + replica_type == ReplicaType::MEMORY || + replica_type == ReplicaType::NOF_SSD) { + metadata.object_checksum = object_meta.object_checksum; + } + const bool has_memory_replica = metadata.HasMemReplica(); const bool should_settle_quota = replica_type == ReplicaType::MEMORY || @@ -2477,18 +3869,22 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key, if (enable_offload_ && !offload_on_evict_) { auto& tenant_state = accessor.GetTenantState(); + bool task_created = false; metadata.VisitReplicas( [](const Replica& replica) { return replica.is_completed() && replica.is_memory_replica(); }, - [this, &object_id, &tenant_state](Replica& replica) { + [this, &object_id, &tenant_state, &task_created](Replica& replica) { auto result = PushOffloadingQueue(object_id, replica); if (result) { - replica.inc_refcnt(); - tenant_state.offloading_tasks.emplace( - object_id.user_key, - OffloadingTask{replica.id(), - std::chrono::system_clock::now()}); + if (!task_created) { + replica.inc_refcnt(); + tenant_state.offloading_tasks.emplace( + object_id.user_key, + OffloadingTask{replica.id(), + std::chrono::system_clock::now()}); + task_created = true; + } } }); } @@ -2505,14 +3901,38 @@ auto MasterService::PutEnd(const UUID& client_id, const std::string& key, // at beginning. 2. If this object has soft pin enabled, set it to be soft // pinned. metadata.GrantLease(0, default_kv_soft_pin_ttl_); + PublishKvStored(key, replica_type, metadata, object_id.tenant_id); + + if (enable_oplog_ && ordered_oplog_writer_) { + std::string payload = SerializeMetadataForOpLog(metadata); + auto result = AppendOpLogVisibleBeforeDurable( + OpType::PUT_END, object_id.tenant_id.value(), key, payload); + if (!result) { + LOG(WARNING) << "PutEnd: OpLog queue failed for key=" << key + << ", err=" << static_cast(result.error()); + } + } return {}; } auto MasterService::AddReplica(const UUID& client_id, const std::string& key, - const std::string& tenant_id, Replica& replica) - -> tl::expected { + const TenantId& tenant_id, Replica& replica) + -> tl::expected { + assert(tenant_id.IsValid()); + TenantId normalized_tenant; + std::unique_lock policy_lock(tenant_quota_policy_mutex_, + std::defer_lock); + if (enable_multi_tenants_) { + policy_lock.lock(); + auto normalized_tenant_result = + ResolveTenantIdForWriteLocked(tenant_id); + if (!normalized_tenant_result) { + return tl::make_unexpected(normalized_tenant_result.error()); + } + normalized_tenant = std::move(normalized_tenant_result.value()); + } std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); + const ObjectIdentity object_id{std::move(normalized_tenant), key}; MetadataAccessorRW accessor(this, object_id); if (!accessor.Exists()) { accessor.Create( @@ -2527,11 +3947,56 @@ auto MasterService::AddReplica(const UUID& client_id, const std::string& key, return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } - if (!metadata.HasReplica(&Replica::fn_is_local_disk_replica)) { + const bool replacing_existing = + metadata.HasReplica(&Replica::fn_is_local_disk_replica); + + if (enable_oplog_ && ordered_oplog_writer_) { + std::vector post; + for (const auto& existing : metadata.GetAllReplicas()) { + if (existing.status() != ReplicaStatus::COMPLETE) continue; + if (replacing_existing && + existing.type() == ReplicaType::LOCAL_DISK && + existing.get_descriptor() + .get_local_disk_descriptor() + .client_id == client_id) { + // Substitute with the updated descriptor. + Replica::Descriptor updated = existing.get_descriptor(); + updated.get_local_disk_descriptor().transport_endpoint = + replica.get_descriptor() + .get_local_disk_descriptor() + .transport_endpoint; + updated.get_local_disk_descriptor().object_size = + replica.get_descriptor() + .get_local_disk_descriptor() + .object_size; + post.push_back(std::move(updated)); + } else { + post.push_back(existing.get_descriptor()); + } + } + if (!replacing_existing) { + // The new LOCAL_DISK replica is COMPLETE upon AddReplica. + post.push_back(replica.get_descriptor()); + } + + auto persist_result = AppendOpLogVisibleBeforeDurable( + OpType::PUT_END, object_id.tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, post, metadata.group_id, + metadata.data_type)); + if (!persist_result) { + return tl::make_unexpected(persist_result.error()); + } + } + + if (!replacing_existing) { std::vector replicas; replicas.emplace_back(std::move(replica)); metadata.AddReplicas(std::move(replicas)); - return {}; + auto& shard = accessor.GetShard(); + shard.OnDiskReplicaAdded(metadata); + SyncCacheTotalAccounting(metadata); + return true; } metadata.VisitReplicas( @@ -2551,15 +4016,15 @@ auto MasterService::AddReplica(const UUID& client_id, const std::string& key, .get_local_disk_descriptor() .object_size; }); - return {}; + return false; } auto MasterService::PutRevoke(const UUID& client_id, const std::string& key, - const std::string& tenant_id, + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected { std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); + const auto object_id = MakeObjectIdentityForRequest(key, tenant_id); MetadataAccessorRW accessor(this, object_id); if (!accessor.Exists()) { LOG(INFO) << "key=" << key << ", info=object_not_found"; @@ -2573,6 +4038,11 @@ auto MasterService::PutRevoke(const UUID& client_id, const std::string& key, return tl::make_unexpected(ErrorCode::ILLEGAL_CLIENT); } + if (!accessor.InProcessing()) { + LOG(ERROR) << "key=" << key << ", error=no_primary_write_in_progress"; + return tl::make_unexpected(ErrorCode::INVALID_WRITE); + } + auto processing_rep = metadata.GetFirstReplica([replica_type]( const Replica& replica) { if (replica_type == ReplicaType::ALL) { @@ -2587,14 +4057,60 @@ auto MasterService::PutRevoke(const UUID& client_id, const std::string& key, return tl::make_unexpected(ErrorCode::INVALID_WRITE); } - const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata); - EraseReplicasWithCacheTotalAccounting( - metadata, [replica_type](const Replica& replica) { - if (replica_type == ReplicaType::ALL) { - return replica.is_memory_replica() || replica.is_nof_replica(); - } - return replica.type() == replica_type; + auto target_pred = [replica_type](const Replica& r) { + if (!r.is_processing()) { + return false; + } + if (replica_type == ReplicaType::ALL) { + return r.is_memory_replica() || r.is_nof_replica(); + } + return r.type() == replica_type; + }; + + if (enable_oplog_ && ordered_oplog_writer_) { + auto remaining = + BuildRemainingReplicaDescriptors(metadata, target_pred); + std::vector removed_ids; + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + return tl::make_unexpected(reservation.error()); + } + metadata.VisitReplicas(target_pred, [&removed_ids](Replica& r) { + removed_ids.push_back(r.id()); + r.mark_removed(); }); + + tl::expected persist_result; + if (remaining.empty()) { + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::REMOVE, + tenant_id.value(), key, {}, + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, QuotaEraseMode::kFull); + }); + } else { + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::PUT_END, + tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, remaining, + metadata.group_id, metadata.data_type), + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, QuotaEraseMode::kFull); + }); + } + if (!persist_result) { + return tl::make_unexpected(persist_result.error()); + } + return {}; + } + + const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata); + EraseReplicasWithCacheTotalAccounting(metadata, target_pred); const uint64_t after_charge = CompletedMemoryQuotaCharge(metadata); if (before_charge > after_charge) { ReleaseCommittedQuotaCharge(metadata, before_charge - after_charge); @@ -2617,20 +4133,30 @@ auto MasterService::PutRevoke(const UUID& client_id, const std::string& key, return {}; } +auto MasterService::PutEnd(const UUID& client_id, const std::string& key, + const TenantId& tenant_id, ReplicaType replica_type) + -> tl::expected { + return PutEnd(client_id, ObjectMeta{key, std::nullopt}, tenant_id, + replica_type); +} + std::vector> MasterService::BatchPutEnd( - const UUID& client_id, const std::vector& keys, - const std::string& tenant_id, ReplicaType replica_type) { + const UUID& client_id, const std::vector& object_metas, + const TenantId& tenant_id, ReplicaType replica_type) { + assert(tenant_id.IsValid()); std::vector> results; - results.reserve(keys.size()); - for (const auto& key : keys) { - results.emplace_back(PutEnd(client_id, key, tenant_id, replica_type)); + results.reserve(object_metas.size()); + for (const auto& object_meta : object_metas) { + results.emplace_back( + PutEnd(client_id, object_meta, tenant_id, replica_type)); } return results; } std::vector> MasterService::BatchPutRevoke( const UUID& client_id, const std::vector& keys, - const std::string& tenant_id, ReplicaType replica_type) { + const TenantId& tenant_id, ReplicaType replica_type) { + assert(tenant_id.IsValid()); std::vector> results; results.reserve(keys.size()); for (const auto& key : keys) { @@ -2655,11 +4181,16 @@ std::vector> MasterService::BatchPutRevoke( // Note: during Case B the key is temporarily unreadable (all replicas are // PROCESSING). Readers will get REPLICA_IS_NOT_READY until UpsertEnd. auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, - const std::string& tenant_id, + const TenantId& tenant_id, const uint64_t slice_length, const ReplicateConfig& config) -> tl::expected, ErrorCode> { - const auto object_id = MakeObjectIdentity(key, tenant_id); + auto normalized_tenant_result = ResolveTenantIdForWrite(tenant_id); + if (!normalized_tenant_result) { + return tl::make_unexpected(normalized_tenant_result.error()); + } + const ObjectIdentity object_id{std::move(normalized_tenant_result.value()), + key}; // --- Parameter validation (same as PutStart) --- if ((config.replica_num == 0 && config.nof_replica_num == 0) || key.empty() || slice_length == 0) { @@ -2686,6 +4217,8 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, } #endif + UpdateClientHostId(client_id, config.host_id); + if ((memory_allocator_type_ == BufferAllocatorType::CACHELIB) && (slice_length > kMaxSliceSize)) { LOG(ERROR) << "key=" << key << ", slice_length=" << slice_length @@ -2705,236 +4238,314 @@ auto MasterService::UpsertStart(const UUID& client_id, const std::string& key, [[maybe_unused]] auto object_operation_lock = AcquireObjectOperationLock(object_id.tenant_id, object_id.user_key); - const auto now = std::chrono::system_clock::now(); - std::optional case_a_retry_shard_idx; - { - // --- Lock acquisition --- - auto alive_clients = getAliveClientsSnapshot(); - std::shared_lock shared_lock(snapshot_mutex_); - // Use getMetadataShardIndex to find the object at its current shard - // (handles both grouped and ungrouped routing). - const size_t lookup_shard_idx = - getMetadataShardIndex(object_id.tenant_id, object_id.user_key); - MetadataShardAccessorRW shard(this, lookup_shard_idx); - auto& tenant_state = shard->tenants[object_id.tenant_id]; - - auto it = tenant_state.metadata.find(key); - - // --- Step 0: stale handle cleanup --- - if (it != tenant_state.metadata.end() && - CleanupStaleHandles(it->second, alive_clients)) { - tenant_state.processing_keys.erase(key); - ErasePromotionTaskIfPresent(tenant_state, key, object_id.tenant_id); - EraseMetadata(tenant_state, it, object_id.tenant_id); - it = tenant_state.metadata.end(); - } - - // --- Step 1: safety checks and preemption (only if key exists) --- - if (it != tenant_state.metadata.end()) { - auto& metadata = it->second; + const uint64_t requested_quota_charge = + RequestedMemoryQuotaCharge(slice_length, config); - // Reject if the caller tries to change group membership. - // Group membership is immutable while an object exists. - if (config.group_ids.has_value() && metadata.group_id != group_id) { - LOG(ERROR) << "key=" << key - << ", error=group_membership_is_immutable"; - return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + auto attempt_once = + [&]() -> tl::expected, ErrorCode> { + std::unique_lock zero_charge_policy_lock( + tenant_quota_policy_mutex_, std::defer_lock); + if (ShouldProtectZeroChargeMetadataCreate(requested_quota_charge)) { + zero_charge_policy_lock.lock(); + auto latest_tenant_result = + ResolveTenantIdForWriteLocked(tenant_id); + if (!latest_tenant_result) { + return tl::make_unexpected(latest_tenant_result.error()); } + } - // Reject if a Copy/Move task is actively reading this key's - // replicas. - if (tenant_state.replication_tasks.count(key) > 0) { - LOG(INFO) << "key=" << key - << ", error=object_has_replication_task"; - return tl::make_unexpected( - ErrorCode::OBJECT_HAS_REPLICATION_TASK); - } + auto now = std::chrono::system_clock::now(); + std::optional case_a_retry_shard_idx; + { + // --- Lock acquisition --- + auto alive_clients = getAliveClientsSnapshot(); + std::shared_lock shared_lock(snapshot_mutex_); + // Use getMetadataShardIndex to find the object at its current shard + // (handles both grouped and ungrouped routing). + const size_t lookup_shard_idx = + getMetadataShardIndex(object_id.tenant_id, object_id.user_key); + MetadataShardAccessorRW shard(this, lookup_shard_idx); + auto& tenant_state = shard->tenants[object_id.tenant_id]; - // Reject if an offload-to-disk task is in progress (same reason). - if (tenant_state.offloading_tasks.count(key) > 0) { - LOG(INFO) << "key=" << key - << ", error=object_has_offloading_task"; - return tl::make_unexpected( - ErrorCode::OBJECT_HAS_REPLICATION_TASK); + auto it = tenant_state.metadata.find(key); + + // --- Step 0: stale handle cleanup --- + if (it != tenant_state.metadata.end()) { + auto cleanup_plan = + BuildStaleHandleCleanupPlan(it->second, alive_clients); + if (!cleanup_plan.removed_ids.empty()) { + auto persist_result = PersistStaleHandleCleanupForHA( + "UpsertStart(stale cleanup)", object_id.tenant_id, key, + it->second, cleanup_plan); + if (!persist_result) { + return tl::make_unexpected(persist_result.error()); + } + if (enable_oplog_) { + return tl::make_unexpected( + ErrorCode::OBJECT_ALREADY_EXISTS); + } else if (CleanupStaleHandles(tenant_state, it->second, + alive_clients, &shard)) { + // EraseMetadata handles processing_keys, + // replication_tasks, offloading_tasks (with + // dec_refcnt), and promotion task cleanup. + EraseMetadata(tenant_state, it, object_id.tenant_id, + QuotaEraseMode::kFull, &shard); + it = tenant_state.metadata.end(); + } + } } - // Preempt an in-progress Put/Upsert on the same key. The previous - // writer's PROCESSING replicas are moved to discarded_replicas_ - // with a TTL so they are not freed while the old writer may still - // be doing RDMA writes. Unlike PutStart (which only preempts after - // a timeout), UpsertStart preempts immediately. - if (tenant_state.processing_keys.count(key) > 0) { - auto processing_replicas = - metadata.PopReplicas(&Replica::fn_is_processing); - if (!processing_replicas.empty()) { - std::lock_guard lock(discarded_replicas_mutex_); - discarded_replicas_.emplace_back( - std::move(processing_replicas), - now + put_start_release_timeout_sec_); + // --- Step 1: safety checks and preemption (only if key exists) --- + if (it != tenant_state.metadata.end()) { + auto& metadata = it->second; + + // Reject if the caller tries to change group membership. + // Group membership is immutable while an object exists. + if (config.group_ids.has_value() && + metadata.group_id != group_id) { + LOG(ERROR) << "key=" << key + << ", error=group_membership_is_immutable"; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } - tenant_state.processing_keys.erase(key); - - // If no COMPLETE replicas survive the preemption, this key - // effectively does not exist — fall through to Case A. - if (!metadata.HasReplica(&Replica::fn_is_completed)) { - ErasePromotionTaskIfPresent(tenant_state, key, - object_id.tenant_id); - EraseMetadata(tenant_state, it, object_id.tenant_id); - it = tenant_state.metadata.end(); + + // Reject if a Copy/Move task is actively reading this key's + // replicas. + if (tenant_state.replication_tasks.count(key) > 0) { + LOG(INFO) << "key=" << key + << ", error=object_has_replication_task"; + return tl::make_unexpected( + ErrorCode::OBJECT_HAS_REPLICATION_TASK); } - } - } - // --- Case A: key does not exist (or was erased above) --- - // Allocate fresh buffers, identical to PutStart. - if (it == tenant_state.metadata.end()) { - VLOG(1) << "key=" << key << ", action=upsert_start_case_a"; - const size_t case_a_shard_idx = - group_id.empty() - ? getShardIndex(object_id.tenant_id, object_id.user_key) - : getShardIndex(group_id); - if (case_a_shard_idx != lookup_shard_idx) { - case_a_retry_shard_idx = case_a_shard_idx; - if (tenant_state.Empty()) { - shard->tenants.erase(object_id.tenant_id); + if (tenant_state.promotion_tasks.count(key) > 0) { + LOG(INFO) + << "key=" << key << ", error=object_has_promotion_task"; + return tl::make_unexpected( + ErrorCode::OBJECT_HAS_REPLICATION_TASK); } - } else { - return AllocateAndInsertMetadata(shard, client_id, key, - slice_length, config, group_id, - object_id.tenant_id, now); - } - } else { - // --- Step 2: key exists with COMPLETE replicas → Case B or C --- - auto& metadata = it->second; - // Reject if any reader holds a reference (refcnt > 0). Overwriting - // a buffer that an RDMA read is streaming from would cause data - // corruption. The client should retry after readers finish. - if (metadata.HasReplica(&Replica::fn_is_busy)) { - LOG(INFO) << "key=" << key << ", error=object_replica_busy"; - return tl::make_unexpected(ErrorCode::OBJECT_REPLICA_BUSY); - } + // Reject if an offload-to-disk task is in progress (same + // reason). + if (tenant_state.offloading_tasks.count(key) > 0) { + LOG(INFO) << "key=" << key + << ", error=object_has_offloading_task"; + return tl::make_unexpected( + ErrorCode::OBJECT_HAS_REPLICATION_TASK); + } - if (metadata.size == slice_length) { - // --- Case B: same size — in-place update --- - // Reuse existing buffer addresses. No allocation or - // deallocation. The client will RDMA-write new data to the same - // addresses. - // - // hard_pinned is const and preserved automatically — upsert - // does not change the eviction protection level of an existing - // object. - metadata.client_id = client_id; - metadata.put_start_time = now; + // Preempt an in-progress Put/Upsert on the same key. The + // previous writer's PROCESSING replicas are moved to + // discarded_replicas_ with a TTL so they are not freed while + // the old writer may still be doing RDMA writes. Unlike + // PutStart (which only preempts after a timeout), UpsertStart + // preempts immediately. + if (tenant_state.processing_keys.count(key) > 0) { + auto processing_replicas = + metadata.PopReplicas(&Replica::fn_is_processing); + if (!processing_replicas.empty()) { + std::lock_guard lock(discarded_replicas_mutex_); + discarded_replicas_.emplace_back( + std::move(processing_replicas), + now + put_start_release_timeout_sec_); + } + tenant_state.processing_keys.erase(key); + + // If no COMPLETE replicas survive the preemption, this key + // effectively does not exist — fall through to Case A. + if (!metadata.HasReplica(&Replica::fn_is_completed)) { + EraseMetadata(tenant_state, it, object_id.tenant_id, + QuotaEraseMode::kFull, &shard); + it = tenant_state.metadata.end(); + } + } + } - // Reconcile soft_pin state with the incoming config. - { - SpinLocker locker(&metadata.lock); - if (config.with_soft_pin && !metadata.soft_pin_timeout) { - metadata.soft_pin_timeout.emplace(); - MasterMetricManager::instance().inc_soft_pin_key_count( - 1); - } else if (!config.with_soft_pin && - metadata.soft_pin_timeout) { - metadata.soft_pin_timeout.reset(); - MasterMetricManager::instance().dec_soft_pin_key_count( - 1); + // --- Case A: key does not exist (or was erased above) --- + // Allocate fresh buffers, identical to PutStart. + if (it == tenant_state.metadata.end()) { + VLOG(1) << "key=" << key << ", action=upsert_start_case_a"; + const size_t case_a_shard_idx = + group_id.empty() + ? getShardIndex(object_id.tenant_id, object_id.user_key) + : getShardIndex(group_id); + if (case_a_shard_idx != lookup_shard_idx) { + case_a_retry_shard_idx = case_a_shard_idx; + if (tenant_state.Empty()) { + shard->tenants.erase(object_id.tenant_id); } + } else { + return AllocateAndInsertMetadata( + shard, client_id, key, slice_length, config, group_id, + object_id.tenant_id, now); } + } else { + // --- Step 2: key exists with COMPLETE replicas → Case B or C + // --- + auto& metadata = it->second; - // Mark COMPLETE → PROCESSING so readers won't see stale data - // mid-transfer. The key becomes unreadable until UpsertEnd. - metadata.VisitReplicas( - &Replica::fn_is_completed, - [](Replica& replica) { replica.mark_processing(); }); - SyncCacheTotalAccounting(metadata); + // Reject if any reader holds a reference (refcnt > 0). + // Overwriting a buffer that an RDMA read is streaming from + // would cause data corruption. The client should retry after + // readers finish. + if (metadata.HasReplica(&Replica::fn_is_busy)) { + LOG(INFO) << "key=" << key << ", error=object_replica_busy"; + return tl::make_unexpected(ErrorCode::OBJECT_REPLICA_BUSY); + } + + if (metadata.size == slice_length) { + // --- Case B: same size — in-place update --- + // Reuse existing buffer addresses. No allocation or + // deallocation. The client will RDMA-write new data to the + // same addresses. + // + // hard_pinned is const and preserved automatically — upsert + // does not change the eviction protection level of an + // existing object. + metadata.client_id = client_id; + metadata.put_start_time = now; + + // Reconcile soft_pin state with the incoming config. + { + SpinLocker locker(&metadata.lock); + if (config.with_soft_pin && + !metadata.soft_pin_timeout) { + metadata.soft_pin_timeout.emplace(); + MasterMetricManager::instance() + .inc_soft_pin_key_count(1); + } else if (!config.with_soft_pin && + metadata.soft_pin_timeout) { + metadata.soft_pin_timeout.reset(); + MasterMetricManager::instance() + .dec_soft_pin_key_count(1); + } + } + + // Mark COMPLETE → PROCESSING so readers won't see stale + // data mid-transfer. The key becomes unreadable until + // UpsertEnd. + metadata.VisitReplicas( + &Replica::fn_is_completed, + [](Replica& replica) { replica.mark_processing(); }); + SyncCacheTotalAccounting(metadata); + + tenant_state.processing_keys.insert(key); + + // Return the existing descriptors — same buffer addresses + // as before. + std::vector replica_list; + const auto& all_replicas = metadata.GetAllReplicas(); + replica_list.reserve(all_replicas.size()); + for (const auto& replica : all_replicas) { + replica_list.emplace_back(replica.get_descriptor()); + } - tenant_state.processing_keys.insert(key); + VLOG(1) << "key=" << key + << ", action=upsert_start_case_b_inplace"; + return replica_list; + } - // Return the existing descriptors — same buffer addresses as - // before. - std::vector replica_list; - const auto& all_replicas = metadata.GetAllReplicas(); - replica_list.reserve(all_replicas.size()); - for (const auto& replica : all_replicas) { - replica_list.emplace_back(replica.get_descriptor()); + // --- Case C: different size — discard old replicas and + // reallocate + // --- Old buffers cannot be reused. Move them to + // discarded_replicas_ for delayed release (readers may still + // hold descriptors without refcnt), then allocate fresh buffers + // at the new size. + // + // Preserve hard_pin and soft_pin from the old metadata so that + // eviction protection survives a size-changing upsert (RFC + // §2.2.2). + ReplicateConfig merged_config = config; + merged_config.with_hard_pin = + merged_config.with_hard_pin || metadata.IsHardPinned(); + merged_config.with_soft_pin = + merged_config.with_soft_pin || metadata.IsSoftPinned(); + + const std::string existing_group_id = metadata.group_id; + const uint64_t old_quota_charge = + metadata.committed_quota_charge_bytes != 0 + ? metadata.committed_quota_charge_bytes + : CompletedMemoryQuotaCharge(metadata); + auto old_replicas = + PopReplicasWithCacheTotalAccounting(metadata); + if (!old_replicas.empty()) { + std::lock_guard lock(discarded_replicas_mutex_); + discarded_replicas_.emplace_back( + std::move(old_replicas), + now + put_start_release_timeout_sec_); } + EraseMetadata(tenant_state, it, object_id.tenant_id, + QuotaEraseMode::kPreserveOld, &shard); VLOG(1) << "key=" << key - << ", action=upsert_start_case_b_inplace"; - return replica_list; - } - - // --- Case C: different size — discard old replicas and reallocate - // --- Old buffers cannot be reused. Move them to - // discarded_replicas_ for delayed release (readers may still hold - // descriptors without refcnt), then allocate fresh buffers at the - // new size. - // - // Preserve hard_pin and soft_pin from the old metadata so that - // eviction protection survives a size-changing upsert (RFC §2.2.2). - ReplicateConfig merged_config = config; - merged_config.with_hard_pin = - merged_config.with_hard_pin || metadata.IsHardPinned(); - merged_config.with_soft_pin = - merged_config.with_soft_pin || metadata.IsSoftPinned(); - - const std::string existing_group_id = metadata.group_id; - const uint64_t old_quota_charge = - metadata.committed_quota_charge_bytes != 0 - ? metadata.committed_quota_charge_bytes - : CompletedMemoryQuotaCharge(metadata); - auto old_replicas = PopReplicasWithCacheTotalAccounting(metadata); - if (!old_replicas.empty()) { - std::lock_guard lock(discarded_replicas_mutex_); - discarded_replicas_.emplace_back( - std::move(old_replicas), - now + put_start_release_timeout_sec_); - } - EraseMetadata(tenant_state, it, object_id.tenant_id, - QuotaEraseMode::kPreserveOld); - - VLOG(1) << "key=" << key - << ", action=upsert_start_case_c_reallocate"; - auto allocate_result = AllocateAndInsertMetadata( - shard, client_id, key, slice_length, merged_config, - existing_group_id, object_id.tenant_id, now); - if (!allocate_result) { - ReleaseTenantQuota(object_id.tenant_id, old_quota_charge); + << ", action=upsert_start_case_c_reallocate"; + auto allocate_result = AllocateAndInsertMetadata( + shard, client_id, key, slice_length, merged_config, + existing_group_id, object_id.tenant_id, now); + if (!allocate_result) { + ReleaseTenantQuota(object_id.tenant_id, old_quota_charge); + return allocate_result; + } + auto new_it = tenant_state.metadata.find(key); + if (new_it != tenant_state.metadata.end()) { + new_it->second.pending_replaced_quota_charge_bytes = + old_quota_charge; + } return allocate_result; } - auto new_it = tenant_state.metadata.find(key); - if (new_it != tenant_state.metadata.end()) { - new_it->second.pending_replaced_quota_charge_bytes = - old_quota_charge; - } - return allocate_result; } + std::shared_lock shared_lock(snapshot_mutex_); + MetadataShardAccessorRW shard(this, case_a_retry_shard_idx.value()); + auto& retry_tenant_state = shard->tenants[object_id.tenant_id]; + const auto current_route = + GetGroupRoute(object_id.tenant_id, object_id.user_key); + if (current_route.has_value() || + retry_tenant_state.metadata.contains(key)) { + LOG(INFO) << "key=" << key << ", info=object_already_exists"; + return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS); + } + return AllocateAndInsertMetadata(shard, client_id, key, slice_length, + config, group_id, object_id.tenant_id, + now); + }; + + for (int attempt = 0; attempt <= kMaxTenantQuotaEvictionRetries; + ++attempt) { + auto result = attempt_once(); + if (result.has_value() || + result.error() != ErrorCode::TENANT_QUOTA_EXCEEDED) { + return result; + } + if (attempt == kMaxTenantQuotaEvictionRetries) { + MasterMetricManager::instance().inc_tenant_quota_reject( + object_id.tenant_id.value(), "quota_exceeded"); + return result; + } + EvictTenantMemoryForQuota( + object_id.tenant_id, + tenant_quota_table_.ComputeDeficit(object_id.tenant_id, + requested_quota_charge)); } - std::shared_lock shared_lock(snapshot_mutex_); - MetadataShardAccessorRW shard(this, case_a_retry_shard_idx.value()); - auto& retry_tenant_state = shard->tenants[object_id.tenant_id]; - const auto current_route = - GetGroupRoute(object_id.tenant_id, object_id.user_key); - if (current_route.has_value() || - retry_tenant_state.metadata.contains(key)) { - LOG(INFO) << "key=" << key << ", info=object_already_exists"; - return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS); - } - return AllocateAndInsertMetadata(shard, client_id, key, slice_length, - config, group_id, object_id.tenant_id, - now); + return tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); +} + +auto MasterService::UpsertEnd(const UUID& client_id, + const ObjectMeta& object_meta, + const TenantId& tenant_id, + ReplicaType replica_type) + -> tl::expected { + return PutEnd(client_id, object_meta, tenant_id, replica_type); } auto MasterService::UpsertEnd(const UUID& client_id, const std::string& key, - const std::string& tenant_id, + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected { - return PutEnd(client_id, key, tenant_id, replica_type); + return UpsertEnd(client_id, ObjectMeta{key, std::nullopt}, tenant_id, + replica_type); } auto MasterService::UpsertRevoke(const UUID& client_id, const std::string& key, - const std::string& tenant_id, + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected { return PutRevoke(client_id, key, tenant_id, replica_type); @@ -2943,9 +4554,10 @@ auto MasterService::UpsertRevoke(const UUID& client_id, const std::string& key, std::vector, ErrorCode>> MasterService::BatchUpsertStart(const UUID& client_id, const std::vector& keys, - const std::string& tenant_id, + const TenantId& tenant_id, const std::vector& slice_lengths, const ReplicateConfig& config) { + assert(tenant_id.IsValid()); if (keys.size() != slice_lengths.size()) { LOG(ERROR) << "BatchUpsertStart: keys.size()=" << keys.size() << " != slice_lengths.size()=" << slice_lengths.size(); @@ -2974,50 +4586,127 @@ MasterService::BatchUpsertStart(const UUID& client_id, } std::vector> MasterService::BatchUpsertEnd( - const UUID& client_id, const std::vector& keys, - const std::string& tenant_id) { - return BatchPutEnd(client_id, keys, tenant_id); + const UUID& client_id, const std::vector& object_metas, + const TenantId& tenant_id) { + return BatchPutEnd(client_id, object_metas, tenant_id, ReplicaType::ALL); } std::vector> MasterService::BatchUpsertRevoke( const UUID& client_id, const std::vector& keys, - const std::string& tenant_id) { + const TenantId& tenant_id) { return BatchPutRevoke(client_id, keys, tenant_id); } auto MasterService::EvictDiskReplica(const UUID& client_id, const std::string& key, - const std::string& tenant_id, + const TenantId& tenant_id, ReplicaType replica_type) -> tl::expected { - const auto object_id = MakeObjectIdentity(key, tenant_id); + const auto object_id = MakeObjectIdentityForRequest(key, tenant_id); MetadataAccessorRW accessor(this, object_id); if (!accessor.Exists()) { - LOG(INFO) << "key=" << key << ", tenant_id=" << object_id.tenant_id + LOG(INFO) << "key=" << key + << ", tenant_id=" << object_id.tenant_id.value() << ", info=object_not_found_for_eviction"; return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); } auto& metadata = accessor.Get(); - if (replica_type == ReplicaType::DISK) { - EraseReplicasWithCacheTotalAccounting( - metadata, - [](const Replica& replica) { return replica.is_disk_replica(); }); - } else if (replica_type == ReplicaType::LOCAL_DISK) { - metadata.EraseReplicas([&client_id](const Replica& replica) { - return replica.is_local_disk_replica() && - replica.get_descriptor() - .get_local_disk_descriptor() - .client_id == client_id; - }); - } else { + if (replica_type != ReplicaType::DISK && + replica_type != ReplicaType::LOCAL_DISK) { LOG(ERROR) << "key=" << key << ", error=invalid_replica_type_for_eviction"; return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } + auto target_pred = [replica_type, &client_id](const Replica& r) { + if (replica_type == ReplicaType::DISK) { + return r.is_disk_replica(); + } else if (replica_type == ReplicaType::LOCAL_DISK) { + return r.is_local_disk_replica() && + r.get_descriptor().get_local_disk_descriptor().client_id == + client_id; + } + return false; + }; + + if (enable_oplog_ && ordered_oplog_writer_) { + auto remaining = + BuildRemainingReplicaDescriptors(metadata, target_pred); + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + return tl::make_unexpected(reservation.error()); + } + std::vector removed_ids; + metadata.VisitReplicas(target_pred, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + + tl::expected persist_result; + if (remaining.empty()) { + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::REMOVE, + metadata.tenant_id.value(), key, {}, + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, QuotaEraseMode::kFull); + }); + } else { + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::PUT_END, + metadata.tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, remaining, + metadata.group_id, metadata.data_type), + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, QuotaEraseMode::kFull); + }); + } + if (!persist_result) { + return tl::make_unexpected(persist_result.error()); + } + return {}; + } + + tl::expected persist_result; + if (remaining.empty()) { + persist_result = AppendOpLogWithDurableFinalize( + OpType::REMOVE, metadata.tenant_id.value(), key, {}, nullptr); + } else { + persist_result = AppendOpLogWithDurableFinalize( + OpType::PUT_END, metadata.tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, remaining, + metadata.group_id, metadata.data_type), + nullptr); + } + if (!persist_result) { + return tl::make_unexpected(persist_result.error()); + } + } + + if (replica_type == ReplicaType::DISK) { + EraseReplicasWithCacheTotalAccounting(metadata, target_pred); + } else if (replica_type == ReplicaType::LOCAL_DISK) { + bool had_completed_disk = metadata.HasReplica([](const Replica& r) { + return r.is_local_disk_replica() && r.is_completed(); + }); + EraseReplicasWithCacheTotalAccounting(metadata, target_pred); + if (had_completed_disk) { + auto& shard = accessor.GetShard(); + shard.OnDiskReplicaRemoved(had_completed_disk, metadata); + } + } + if (!metadata.IsValid()) { + PublishKvRemoved(key, metadata, object_id.tenant_id); accessor.Erase(); } return {}; @@ -3025,7 +4714,8 @@ auto MasterService::EvictDiskReplica(const UUID& client_id, std::vector> MasterService::BatchEvictDiskReplica( const UUID& client_id, const std::vector& keys, - const std::string& tenant_id, ReplicaType replica_type) { + const TenantId& tenant_id, ReplicaType replica_type) { + assert(tenant_id.IsValid()); std::vector> results; results.reserve(keys.size()); for (const auto& key : keys) { @@ -3036,11 +4726,16 @@ std::vector> MasterService::BatchEvictDiskReplica( } tl::expected MasterService::CopyStart( - const UUID& client_id, const std::string& key, const std::string& tenant_id, + const UUID& client_id, const std::string& key, const TenantId& tenant_id, const std::string& src_segment, const std::vector& tgt_segments) { + auto normalized_tenant_result = ResolveTenantIdForWrite(tenant_id); + if (!normalized_tenant_result) { + return tl::make_unexpected(normalized_tenant_result.error()); + } + const ObjectIdentity object_id{std::move(normalized_tenant_result.value()), + key}; std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); { ScopedSegmentAccess segment_access = segment_manager_.getSegmentAccess(); @@ -3079,12 +4774,32 @@ tl::expected MasterService::CopyStart( return tl::make_unexpected(ErrorCode::REPLICA_NOT_FOUND); } + size_t new_replica_count = 0; + for (const auto& tgt_segment : tgt_segments) { + if (metadata.GetReplicaBySegmentName(tgt_segment) == nullptr) { + ++new_replica_count; + } + } + + const uint64_t reserved_quota_charge = + SaturatingMultiply(static_cast(metadata.size), + static_cast(new_replica_count)); + auto quota_result = + ReserveTenantQuota(object_id.tenant_id, reserved_quota_charge); + if (!quota_result) { + if (quota_result.error() == ErrorCode::TENANT_QUOTA_EXCEEDED) { + MasterMetricManager::instance().inc_tenant_quota_reject( + object_id.tenant_id.value(), "quota_exceeded"); + } + return tl::make_unexpected(quota_result.error()); + } + auto abort_reserved_quota = [&] { + AbortTenantQuota(object_id.tenant_id, reserved_quota_charge); + }; + std::vector replicas; - replicas.reserve(tgt_segments.size()); + replicas.reserve(new_replica_count); { - // PR2 limitation: Copy can allocate extra physical MEMORY replicas - // without tenant quota admission. It does not change the logical - // object set; full quota-aware Copy admission is deferred. ScopedAllocatorAccess allocator_access = segment_manager_.getAllocatorAccess(); const auto& allocator_manager = allocator_access.getAllocatorManager(); @@ -3100,6 +4815,7 @@ tl::expected MasterService::CopyStart( if (!replica.has_value()) { LOG(ERROR) << "key=" << key << ", tgt_segment=" << tgt_segment << ", failed to allocate replica"; + abort_reserved_quota(); return tl::make_unexpected(replica.error()); } replicas.push_back(std::move(*replica)); @@ -3119,11 +4835,15 @@ tl::expected MasterService::CopyStart( // Create replication task for tracking. auto& tenant_state = accessor.GetTenantState(); - tenant_state.replication_tasks.emplace( + auto task_insert = tenant_state.replication_tasks.emplace( std::piecewise_construct, std::forward_as_tuple(key), std::forward_as_tuple(client_id, std::chrono::system_clock::now(), ReplicationTask::Type::COPY, source->id(), - std::move(replica_ids))); + std::move(replica_ids), reserved_quota_charge)); + if (!task_insert.second) { + abort_reserved_quota(); + return tl::make_unexpected(ErrorCode::OBJECT_HAS_REPLICATION_TASK); + } // Increase source refcnt to protect it from eviction. source->inc_refcnt(); @@ -3136,10 +4856,10 @@ tl::expected MasterService::CopyStart( } tl::expected MasterService::CopyEnd( - const UUID& client_id, const std::string& key, - const std::string& tenant_id) { + const UUID& client_id, const std::string& key, const TenantId& tenant_id) { std::shared_lock shared_lock(snapshot_mutex_); - MetadataAccessorRW accessor(this, MakeObjectIdentity(key, tenant_id)); + MetadataAccessorRW accessor(this, + MakeObjectIdentityForRequest(key, tenant_id)); if (!accessor.Exists()) { LOG(ERROR) << "key=" << key << ", error=object_not_found"; return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); @@ -3171,6 +4891,12 @@ tl::expected MasterService::CopyEnd( LOG(ERROR) << "key=" << key << ", source_id=" << source_id << ", status=" << (source == nullptr ? "nullptr" : "invalid") << ", copy source becomes invalid during data transfer"; + // Release the refcnt taken in CopyStart. The success path below does + // this once the copy completes; this error path must do it too, or the + // source replica stays pinned and can never be evicted. + if (source != nullptr) { + source->dec_refcnt(); + } // Discard target replicas and clear the replication task. EraseReplicasWithCacheTotalAccounting( metadata, [&task](const Replica& replica) { @@ -3178,6 +4904,7 @@ tl::expected MasterService::CopyEnd( task.replica_ids.end(), replica.id()) != task.replica_ids.end(); }); + AbortTenantQuota(metadata.tenant_id, task.reserved_quota_charge_bytes); accessor.EraseReplicationTask(); if (!metadata.IsValid()) { // Remove the object if it does not have any replicas. @@ -3186,11 +4913,13 @@ tl::expected MasterService::CopyEnd( return tl::make_unexpected(ErrorCode::REPLICA_IS_GONE); } - // Decrement source reference count - source->dec_refcnt(); - - // Mark all replica_ids as complete + // First validate that all target replicas are still healthy. If any + // replica is invalid we won't be able to mark it complete; this affects + // the post-mutation descriptor list. bool all_complete = true; + uint64_t completed_quota_charge = 0; + std::vector commit_target_ids; + commit_target_ids.reserve(task.replica_ids.size()); for (const auto& replica_id : task.replica_ids) { auto replica = metadata.GetReplicaByID(replica_id); if (replica == nullptr || replica->has_invalid_mem_handle()) { @@ -3199,11 +4928,69 @@ tl::expected MasterService::CopyEnd( << ", copy target becomes invalid during data transfer"; all_complete = false; } else { + commit_target_ids.push_back(replica_id); + } + } + + std::optional batch_reservation; + if (enable_ha_ && enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + return tl::make_unexpected(reservation.error()); + } + batch_reservation = std::move(reservation.value()); + } + + source->dec_refcnt(); + for (const auto& replica_id : commit_target_ids) { + auto replica = metadata.GetReplicaByID(replica_id); + if (replica != nullptr) { replica->mark_complete(); + completed_quota_charge = SaturatingAdd( + completed_quota_charge, static_cast(metadata.size)); + } + } + + if (enable_oplog_ && ordered_oplog_writer_) { + std::vector post; + metadata.VisitReplicas(&Replica::fn_is_completed, + [&post](const Replica& replica) { + post.push_back(replica.get_descriptor()); + }); + auto payload = SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, post, metadata.group_id, + metadata.data_type); + if (batch_reservation) { + auto persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(*batch_reservation), OpType::PUT_END, + tenant_id.value(), key, payload, nullptr); + if (!persist_result) { + LOG(WARNING) + << "CopyEnd: PUT_END persist failed for key=" << key + << ", err=" << static_cast(persist_result.error()); + } + } else { + auto persist_result = AppendOpLogVisibleBeforeDurable( + OpType::PUT_END, tenant_id.value(), key, payload); + if (!persist_result) { + LOG(WARNING) + << "CopyEnd: PUT_END persist failed for key=" << key + << ", err=" << static_cast(persist_result.error()); + } } } + SyncCacheTotalAccounting(metadata); + const uint64_t commit_charge = + std::min(completed_quota_charge, task.reserved_quota_charge_bytes); + const uint64_t abort_charge = + task.reserved_quota_charge_bytes - commit_charge; + CommitAdditionalTenantQuota(metadata.tenant_id, commit_charge); + AbortTenantQuota(metadata.tenant_id, abort_charge); + metadata.committed_quota_charge_bytes = + SaturatingAdd(metadata.committed_quota_charge_bytes, commit_charge); + accessor.EraseReplicationTask(); return all_complete ? tl::expected() @@ -3211,10 +4998,10 @@ tl::expected MasterService::CopyEnd( } tl::expected MasterService::CopyRevoke( - const UUID& client_id, const std::string& key, - const std::string& tenant_id) { + const UUID& client_id, const std::string& key, const TenantId& tenant_id) { std::shared_lock shared_lock(snapshot_mutex_); - MetadataAccessorRW accessor(this, MakeObjectIdentity(key, tenant_id)); + MetadataAccessorRW accessor(this, + MakeObjectIdentityForRequest(key, tenant_id)); if (!accessor.Exists()) { LOG(ERROR) << "key=" << key << ", error=object_not_found"; return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); @@ -3257,6 +5044,7 @@ tl::expected MasterService::CopyRevoke( }); } + AbortTenantQuota(metadata.tenant_id, task.reserved_quota_charge_bytes); accessor.EraseReplicationTask(); if (!metadata.IsValid()) { @@ -3268,10 +5056,15 @@ tl::expected MasterService::CopyRevoke( } tl::expected MasterService::MoveStart( - const UUID& client_id, const std::string& key, const std::string& tenant_id, + const UUID& client_id, const std::string& key, const TenantId& tenant_id, const std::string& src_segment, const std::string& tgt_segment) { + auto normalized_tenant_result = ResolveTenantIdForWrite(tenant_id); + if (!normalized_tenant_result) { + return tl::make_unexpected(normalized_tenant_result.error()); + } + const ObjectIdentity object_id{std::move(normalized_tenant_result.value()), + key}; std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); if (src_segment == tgt_segment) { LOG(ERROR) << "key=" << key << ", move_tgt=" << tgt_segment << " cannot be the same as move_src=" << src_segment; @@ -3316,9 +5109,21 @@ tl::expected MasterService::MoveStart( std::vector replicas; if (metadata.GetReplicaBySegmentName(tgt_segment) == nullptr) { - // PR2 limitation: Move can allocate a replacement physical MEMORY - // replica without tenant quota admission. Logical object accounting is - // unchanged; full quota-aware Move admission is deferred. + const uint64_t reserved_quota_charge = + SaturatingMultiply(static_cast(metadata.size), 1); + auto quota_result = + ReserveTenantQuota(object_id.tenant_id, reserved_quota_charge); + if (!quota_result) { + if (quota_result.error() == ErrorCode::TENANT_QUOTA_EXCEEDED) { + MasterMetricManager::instance().inc_tenant_quota_reject( + object_id.tenant_id.value(), "quota_exceeded"); + } + return tl::make_unexpected(quota_result.error()); + } + auto abort_reserved_quota = [&] { + AbortTenantQuota(object_id.tenant_id, reserved_quota_charge); + }; + ScopedAllocatorAccess allocator_access = segment_manager_.getAllocatorAccess(); const auto& allocator_manager = allocator_access.getAllocatorManager(); @@ -3328,11 +5133,22 @@ tl::expected MasterService::MoveStart( if (!replica.has_value()) { LOG(ERROR) << "key=" << key << ", tgt_segment=" << tgt_segment << ", failed to allocate replica"; + abort_reserved_quota(); return tl::make_unexpected(replica.error()); } replicas.push_back(std::move(*replica)); + } else { + auto quota_result = ReserveTenantQuota(object_id.tenant_id, 0); + if (!quota_result) { + return tl::make_unexpected(quota_result.error()); + } } + const uint64_t reserved_quota_charge = + replicas.empty() + ? 0 + : SaturatingMultiply(static_cast(metadata.size), 1); + MoveStartResponse response; std::vector replica_ids; @@ -3346,11 +5162,15 @@ tl::expected MasterService::MoveStart( // Create replication task for tracking. auto& tenant_state = accessor.GetTenantState(); - tenant_state.replication_tasks.emplace( + auto task_insert = tenant_state.replication_tasks.emplace( std::piecewise_construct, std::forward_as_tuple(key), std::forward_as_tuple(client_id, std::chrono::system_clock::now(), ReplicationTask::Type::MOVE, source->id(), - std::move(replica_ids))); + std::move(replica_ids), reserved_quota_charge)); + if (!task_insert.second) { + AbortTenantQuota(object_id.tenant_id, reserved_quota_charge); + return tl::make_unexpected(ErrorCode::OBJECT_HAS_REPLICATION_TASK); + } // Increase source refcnt to protect it from eviction. source->inc_refcnt(); @@ -3363,10 +5183,10 @@ tl::expected MasterService::MoveStart( } tl::expected MasterService::MoveEnd( - const UUID& client_id, const std::string& key, - const std::string& tenant_id) { + const UUID& client_id, const std::string& key, const TenantId& tenant_id) { std::shared_lock shared_lock(snapshot_mutex_); - MetadataAccessorRW accessor(this, MakeObjectIdentity(key, tenant_id)); + MetadataAccessorRW accessor(this, + MakeObjectIdentityForRequest(key, tenant_id)); if (!accessor.Exists()) { LOG(ERROR) << "key=" << key << ", error=object_not_found"; return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); @@ -3398,6 +5218,12 @@ tl::expected MasterService::MoveEnd( LOG(ERROR) << "key=" << key << ", source_id=" << source_id << ", status=" << (source == nullptr ? "nullptr" : "invalid") << ", move source becomes invalid during data transfer"; + // Release the refcnt taken in MoveStart. The success path below does + // this once the move completes; this error path must do it too, or the + // source replica stays pinned and can never be evicted. + if (source != nullptr) { + source->dec_refcnt(); + } // Discard target replica and clear the replication task. EraseReplicasWithCacheTotalAccounting( metadata, [&task](const Replica& replica) { @@ -3405,6 +5231,7 @@ tl::expected MasterService::MoveEnd( task.replica_ids.end(), replica.id()) != task.replica_ids.end(); }); + AbortTenantQuota(metadata.tenant_id, task.reserved_quota_charge_bytes); accessor.EraseReplicationTask(); if (!metadata.IsValid()) { // Remove the object if it does not have any replicas. @@ -3413,50 +5240,108 @@ tl::expected MasterService::MoveEnd( return tl::make_unexpected(ErrorCode::REPLICA_IS_GONE); } - // Decrement source reference count - source->dec_refcnt(); - - // If the move target has already existed on MoveStart, task.replica_ids - // will be empty. Thus we need to check whether we have replica_ids to - // process. - if (!task.replica_ids.empty()) { - auto replica_id = task.replica_ids[0]; - auto replica = metadata.GetReplicaByID(replica_id); + // Validate the target replica before any mutation. Source dec_refcnt + // and target mark_complete are deferred until after persist. + bool has_target = !task.replica_ids.empty(); + ReplicaID target_id = has_target ? task.replica_ids[0] : ReplicaID{}; + if (has_target) { + auto replica = metadata.GetReplicaByID(target_id); if (replica == nullptr || replica->has_invalid_mem_handle()) { LOG(WARNING) - << "key=" << key << ", replica_id=" << replica_id + << "key=" << key << ", replica_id=" << target_id << ", move target becomes invalid during data transfer"; + AbortTenantQuota(metadata.tenant_id, + task.reserved_quota_charge_bytes); + // Source untouched; safe to drop the broken task. accessor.EraseReplicationTask(); return tl::make_unexpected(ErrorCode::REPLICA_IS_GONE); } + } - // Mark replica as complete - replica->mark_complete(); - SyncCacheTotalAccounting(metadata); + if (enable_oplog_ && ordered_oplog_writer_) { + // Build post-mutation descriptors: + // - existing COMPLETE replicas, except the source (about to be + // popped) + // - target (if any) flipped to COMPLETE + std::vector post; + for (const auto& rep : metadata.GetAllReplicas()) { + if (rep.id() == source_id) continue; + if (rep.status() == ReplicaStatus::COMPLETE) { + post.push_back(rep.get_descriptor()); + continue; + } + if (has_target && rep.id() == target_id) { + Replica::Descriptor desc = rep.get_descriptor(); + desc.status = ReplicaStatus::COMPLETE; + post.push_back(std::move(desc)); + } + } + + tl::expected persist_result; + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + return tl::make_unexpected(reservation.error()); + } + source->mark_removed(); + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::PUT_END, + metadata.tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, post, metadata.group_id, + metadata.data_type), + [this, removed_ids = std::vector{source_id}]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, QuotaEraseMode::kFull); + }); + } else { + persist_result = AppendOpLogWithDurableFinalize( + OpType::PUT_END, metadata.tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, post, metadata.group_id, + metadata.data_type), + nullptr); + } + if (!persist_result) { + return tl::make_unexpected(persist_result.error()); + } } - // Remove the source replica and release its space later. - auto source_replica = PopReplicasWithCacheTotalAccounting( - metadata, [&source_id](const Replica& replica) { - return replica.id() == source_id; - }); - if (!source_replica.empty()) { - std::lock_guard lock(discarded_replicas_mutex_); - discarded_replicas_.emplace_back( - std::move(source_replica), - std::chrono::system_clock::now() + put_start_release_timeout_sec_); + // Persist OK — apply local commit. + source->dec_refcnt(); + if (has_target) { + auto replica = metadata.GetReplicaByID(target_id); + if (replica != nullptr) { + replica->mark_complete(); + } + } + + if (!(enable_ha_ && enable_oplog_)) { + // Remove the source replica and release its space later. + auto source_replica = PopReplicasWithCacheTotalAccounting( + metadata, [&source_id](const Replica& replica) { + return replica.id() == source_id; + }); + if (!source_replica.empty()) { + std::lock_guard lock(discarded_replicas_mutex_); + discarded_replicas_.emplace_back( + std::move(source_replica), std::chrono::system_clock::now() + + put_start_release_timeout_sec_); + } } + AbortTenantQuota(metadata.tenant_id, task.reserved_quota_charge_bytes); accessor.EraseReplicationTask(); return {}; } tl::expected MasterService::MoveRevoke( - const UUID& client_id, const std::string& key, - const std::string& tenant_id) { + const UUID& client_id, const std::string& key, const TenantId& tenant_id) { std::shared_lock shared_lock(snapshot_mutex_); - MetadataAccessorRW accessor(this, MakeObjectIdentity(key, tenant_id)); + MetadataAccessorRW accessor(this, + MakeObjectIdentityForRequest(key, tenant_id)); if (!accessor.Exists()) { LOG(ERROR) << "key=" << key << ", error=object_not_found"; return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); @@ -3499,6 +5384,7 @@ tl::expected MasterService::MoveRevoke( }); } + AbortTenantQuota(metadata.tenant_id, task.reserved_quota_charge_bytes); accessor.EraseReplicationTask(); if (!metadata.IsValid()) { @@ -3509,10 +5395,10 @@ tl::expected MasterService::MoveRevoke( return {}; } -auto MasterService::Remove(const std::string& key, const std::string& tenant_id, +auto MasterService::Remove(const std::string& key, const TenantId& tenant_id, bool force) -> tl::expected { std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); + const auto object_id = MakeObjectIdentityForRequest(key, tenant_id); MetadataAccessorRW accessor(this, object_id); if (!accessor.Exists()) { VLOG(1) << "key=" << key << ", error=object_not_found"; @@ -3542,15 +5428,41 @@ auto MasterService::Remove(const std::string& key, const std::string& tenant_id, return tl::make_unexpected(ErrorCode::OBJECT_HAS_REPLICATION_TASK); } - auto& tenant_state = accessor.GetTenantState(); - ErasePromotionTaskIfPresent(tenant_state, key, object_id.tenant_id); + if (enable_ha_) { + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + return tl::make_unexpected(reservation.error()); + } + std::vector removed_ids; + metadata.VisitReplicas(&Replica::fn_is_completed, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + auto persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::REMOVE, + object_id.tenant_id.value(), key, {}, + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, QuotaEraseMode::kFull); + }); + if (!persist_result) { + return tl::make_unexpected(persist_result.error()); + } + return {}; + } + } + PublishKvRemoved(key, metadata, object_id.tenant_id); accessor.Erase(); return {}; } auto MasterService::RemoveByRegex(const std::string& regex_pattern, - const std::string& tenant_id, bool force) + const TenantId& tenant_id, bool force) -> tl::expected { + assert(tenant_id.IsValid()); long removed_count = 0; std::regex pattern; @@ -3563,7 +5475,7 @@ auto MasterService::RemoveByRegex(const std::string& regex_pattern, } std::shared_lock shared_lock(snapshot_mutex_); - const auto normalized_tenant = NormalizeTenantId(tenant_id); + const TenantId& normalized_tenant = ResolveRequestTenantId(tenant_id); for (size_t i = 0; i < kNumShards; ++i) { MetadataShardAccessorRW shard(this, i); auto tenant_it = shard->tenants.find(normalized_tenant); @@ -3606,9 +5518,41 @@ auto MasterService::RemoveByRegex(const std::string& regex_pattern, VLOG(1) << "key=" << it->first << " matched by regex. Removing."; - ErasePromotionTaskIfPresent(tenant_state, it->first, - normalized_tenant); - it = EraseMetadata(tenant_state, it, normalized_tenant); + if (enable_ha_) { + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + ++it; + continue; + } + std::vector removed_ids; + it->second.VisitReplicas( + &Replica::fn_is_completed, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + auto persist_result = + AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::REMOVE, + normalized_tenant.value(), it->first, {}, + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + if (!persist_result) { + ++it; + continue; + } + ++it; + removed_count++; + continue; + } + } + it = EraseMetadata(tenant_state, it, normalized_tenant, + QuotaEraseMode::kFull, &shard); removed_count++; } else { ++it; @@ -3626,10 +5570,25 @@ auto MasterService::RemoveByRegex(const std::string& regex_pattern, long MasterService::RemoveAll(bool force) { long removed_count = 0; - uint64_t total_freed_size = 0; + int64_t total_freed_size = 0; std::shared_lock shared_lock(snapshot_mutex_); auto now = std::chrono::system_clock::now(); + // Since RemoveAll clears everything, signal ALL clients with a + // LocalDiskSegment to physically clear their SSD immediately. + // This lets client cleanup overlap with master metadata deletion. + { + ScopedLocalDiskSegmentAccess local_disk_segment_access = + segment_manager_.getLocalDiskSegmentAccess(); + auto& client_local_disk_segment = + local_disk_segment_access.getClientLocalDiskSegment(); + for (auto& [client_id, segment] : client_local_disk_segment) { + MutexLocker locker(&segment->offloading_mutex_); + segment->pending_remove_all = true; + } + } + + // Delete metadata — runs concurrently with client SSD cleanup. for (size_t i = 0; i < kNumShards; i++) { MetadataShardAccessorRW shard(this, i); for (auto tenant_it = shard->tenants.begin(); @@ -3642,10 +5601,49 @@ long MasterService::RemoveAll(bool force) { !tenant_state.replication_tasks.contains(it->first)) { auto mem_rep_count = it->second.CountReplicas( &Replica::fn_is_memory_replica); + + if (enable_ha_) { + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + ++it; + continue; + } + std::vector removed_ids; + it->second.VisitReplicas( + &Replica::fn_is_completed, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + auto persist_result = + AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), + OpType::REMOVE, tenant_it->first.value(), + it->first, {}, + [this, + removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + if (!persist_result) { + ++it; + continue; + } + total_freed_size += it->second.size * mem_rep_count; + ++it; + removed_count++; + continue; + } + } + total_freed_size += it->second.size * mem_rep_count; ErasePromotionTaskIfPresent(tenant_state, it->first, tenant_it->first); - it = EraseMetadata(tenant_state, it, tenant_it->first); + it = EraseMetadata(tenant_state, it, tenant_it->first, + QuotaEraseMode::kFull, &shard); removed_count++; } else { ++it; @@ -3665,14 +5663,17 @@ long MasterService::RemoveAll(bool force) { return removed_count; } -long MasterService::RemoveAll(const std::string& tenant_id, bool force) { +long MasterService::RemoveAll(const TenantId& tenant_id, bool force) { long removed_count = 0; - uint64_t total_freed_size = 0; - // Store the current time to avoid repeatedly - // calling std::chrono::steady_clock::now() + int64_t total_freed_size = 0; std::shared_lock shared_lock(snapshot_mutex_); auto now = std::chrono::system_clock::now(); - const auto normalized_tenant = NormalizeTenantId(tenant_id); + const TenantId& normalized_tenant = ResolveRequestTenantId(tenant_id); + + // For the tenant-scoped overload, only signal clients that own LOCAL_DISK + // replicas of THIS tenant — clearing all clients would cross-delete other + // tenants' SSD data. + std::unordered_set> clients_with_disk_replicas; for (size_t i = 0; i < kNumShards; i++) { MetadataShardAccessorRW shard(this, i); @@ -3686,12 +5687,55 @@ long MasterService::RemoveAll(const std::string& tenant_id, bool force) { if ((force || it->second.IsLeaseExpired(now)) && it->second.AllReplicas(&Replica::fn_is_completed) && !tenant_state.replication_tasks.contains(it->first)) { + it->second.VisitReplicas( + &Replica::fn_is_local_disk_replica, + [&clients_with_disk_replicas](const Replica& replica) { + auto cid = replica.get_local_disk_client_id(); + if (cid) { + clients_with_disk_replicas.insert(*cid); + } + }); auto mem_rep_count = it->second.CountReplicas(&Replica::fn_is_memory_replica); + if (enable_ha_) { + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + ++it; + continue; + } + std::vector removed_ids; + it->second.VisitReplicas( + &Replica::fn_is_completed, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + auto persist_result = + AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::REMOVE, + normalized_tenant.value(), it->first, {}, + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + if (!persist_result) { + ++it; + continue; + } + total_freed_size += it->second.size * mem_rep_count; + ++it; + removed_count++; + continue; + } + } total_freed_size += it->second.size * mem_rep_count; ErasePromotionTaskIfPresent(tenant_state, it->first, normalized_tenant); - it = EraseMetadata(tenant_state, it, normalized_tenant); + it = EraseMetadata(tenant_state, it, normalized_tenant, + QuotaEraseMode::kFull, &shard); removed_count++; } else { ++it; @@ -3702,18 +5746,33 @@ long MasterService::RemoveAll(const std::string& tenant_id, bool force) { } } + if (!clients_with_disk_replicas.empty()) { + ScopedLocalDiskSegmentAccess local_disk_segment_access = + segment_manager_.getLocalDiskSegmentAccess(); + auto& client_local_disk_segment = + local_disk_segment_access.getClientLocalDiskSegment(); + for (const auto& client_id : clients_with_disk_replicas) { + auto seg_it = client_local_disk_segment.find(client_id); + if (seg_it != client_local_disk_segment.end()) { + MutexLocker locker(&seg_it->second->offloading_mutex_); + seg_it->second->pending_remove_all = true; + } + } + } + VLOG(1) << "action=remove_all_objects" - << ", tenant_id=" << normalized_tenant + << ", tenant_id=" << normalized_tenant.value() << ", removed_count=" << removed_count - << ", total_freed_size=" << total_freed_size; + << ", total_freed_size=" << total_freed_size + << ", signaled_clients=" << clients_with_disk_replicas.size(); return removed_count; } auto MasterService::BatchRemove(const std::vector& keys, - const std::string& tenant_id, bool force) + const TenantId& tenant_id, bool force) -> std::vector> { std::vector> results(keys.size()); - const auto normalized_tenant = NormalizeTenantId(tenant_id); + const TenantId& normalized_tenant = ResolveRequestTenantId(tenant_id); // Group keys by shard to reduce lock contention std::unordered_map& keys, continue; } - // Clean up stale replica handles (consistent with single Remove) - if (CleanupStaleHandles(it->second, alive_clients)) { - tenant_state.processing_keys.erase(key); - tenant_state.replication_tasks.erase(key); - tenant_state.offloading_tasks.erase(key); - ErasePromotionTaskIfPresent(tenant_state, key, - normalized_tenant); - EraseMetadata(tenant_state, it, normalized_tenant); - if (tenant_state.Empty()) { - shard->tenants.erase(tenant_it); + // Clean up stale replica handles (consistent with single Remove). + auto cleanup_plan = + BuildStaleHandleCleanupPlan(it->second, alive_clients); + if (!cleanup_plan.removed_ids.empty()) { + auto persist_result = PersistStaleHandleCleanupForHA( + "BatchRemove(stale cleanup)", normalized_tenant, key, + it->second, cleanup_plan); + if (!persist_result) { + results[original_idx] = + tl::make_unexpected(persist_result.error()); + continue; + } + if (enable_oplog_) { + results[original_idx] = tl::make_unexpected( + cleanup_plan.would_invalidate + ? ErrorCode::OBJECT_NOT_FOUND + : ErrorCode::OBJECT_ALREADY_EXISTS); + continue; + } else if (CleanupStaleHandles(tenant_state, it->second, + alive_clients, &shard)) { + EraseMetadata(tenant_state, it, normalized_tenant, + QuotaEraseMode::kFull, &shard); + if (tenant_state.Empty()) { + shard->tenants.erase(tenant_it); + } + results[original_idx] = + tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + continue; } - results[original_idx] = - tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); - continue; } if (!it->second.IsValid()) { results[original_idx] = @@ -3801,8 +5875,42 @@ auto MasterService::BatchRemove(const std::vector& keys, } // Remove object metadata - ErasePromotionTaskIfPresent(tenant_state, key, normalized_tenant); - EraseMetadata(tenant_state, it, normalized_tenant); + if (enable_ha_) { + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + results[original_idx] = + tl::make_unexpected(reservation.error()); + continue; + } + std::vector removed_ids; + metadata.VisitReplicas( + &Replica::fn_is_completed, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + auto persist_result = + AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::REMOVE, + normalized_tenant.value(), key, {}, + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + if (!persist_result) { + results[original_idx] = + tl::make_unexpected(persist_result.error()); + continue; + } + results[original_idx] = {}; + continue; + } + } + EraseMetadata(tenant_state, it, normalized_tenant, + QuotaEraseMode::kFull, &shard); if (tenant_state.Empty()) { shard->tenants.erase(tenant_it); } @@ -3813,23 +5921,74 @@ auto MasterService::BatchRemove(const std::vector& keys, return results; } +void MasterService::CancelPromotionTaskForRemovedReplicas( + TenantState& tenant_state, ObjectMetadata& metadata, + const std::vector& removed_replica_ids) { + if (removed_replica_ids.empty()) { + return; + } + + auto task_it = tenant_state.promotion_tasks.find(metadata.user_key); + if (task_it == tenant_state.promotion_tasks.end() || + task_it->second.alloc_id == 0 || + std::find(removed_replica_ids.begin(), removed_replica_ids.end(), + task_it->second.alloc_id) == removed_replica_ids.end()) { + return; + } + + if (auto* source = metadata.GetReplicaByID(task_it->second.source_id); + source != nullptr) { + source->dec_refcnt(); + } + const UUID holder_id = task_it->second.holder_id; + ErasePromotionTaskIfPresent(tenant_state, metadata.user_key, + metadata.tenant_id); + + // Best-effort cleanup of a task that may still be queued on the holder. + ScopedLocalDiskSegmentAccess local_disk_segment_access = + segment_manager_.getLocalDiskSegmentAccess(); + auto& client_local_disk_segment = + local_disk_segment_access.getClientLocalDiskSegment(); + auto segment_it = client_local_disk_segment.find(holder_id); + if (segment_it != client_local_disk_segment.end()) { + MutexLocker locker(&segment_it->second->offloading_mutex_); + segment_it->second->promotion_objects.erase( + metadata.tenant_id.MakeScopedKey(metadata.user_key)); + } +} + bool MasterService::CleanupStaleHandles( - ObjectMetadata& metadata, - const std::unordered_set>& alive_clients) { + TenantState& tenant_state, ObjectMetadata& metadata, + const std::unordered_set>& alive_clients, + MetadataShardAccessorRW* shard) { + bool had_completed_disk = metadata.HasReplica([](const Replica& r) { + return r.is_local_disk_replica() && r.is_completed(); + }); // Remove those with invalid allocators (memory replicas on unmounted // segments) and local_disk replicas whose owner client is no longer alive. const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata); + std::vector removed_replica_ids; EraseReplicasWithCacheTotalAccounting( - metadata, [&alive_clients](const Replica& replica) { + metadata, + [&alive_clients](const Replica& replica) { return (replica.has_invalid_mem_handle() || replica.has_invalid_nof_handle() || replica.has_stale_local_disk_client(alive_clients)) && replica.is_completed(); - }); + }, + &removed_replica_ids); + CancelPromotionTaskForRemovedReplicas(tenant_state, metadata, + removed_replica_ids); const uint64_t after_charge = CompletedMemoryQuotaCharge(metadata); if (before_charge > after_charge) { ReleaseCommittedQuotaCharge(metadata, before_charge - after_charge); } + if (had_completed_disk && shard && + !metadata.HasReplica([](const Replica& r) { + return r.is_local_disk_replica() && r.is_completed(); + })) { + shard->OnDiskReplicaRemoved(had_completed_disk, metadata); + } // Return true if no valid replicas remain after cleanup return !metadata.IsValid(); @@ -3848,13 +6007,15 @@ size_t MasterService::GetKeyCount() const { auto MasterService::Ping(const UUID& client_id) -> tl::expected { - std::shared_lock lock(client_mutex_); ClientStatus client_status; - auto it = ok_client_.find(client_id); - if (it != ok_client_.end()) { - client_status = ClientStatus::OK; - } else { - client_status = ClientStatus::NEED_REMOUNT; + { + std::shared_lock lock(client_mutex_); + auto it = ok_client_.find(client_id); + if (it != ok_client_.end()) { + client_status = ClientStatus::OK; + } else { + client_status = ClientStatus::NEED_REMOUNT; + } } PodUUID pod_client_id = {client_id.first, client_id.second}; if (!client_ping_queue_.push(pod_client_id)) { @@ -3938,12 +6099,12 @@ auto MasterService::OffloadObjectHeartbeat(const UUID& client_id, << client_id; return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND); } + std::vector result; std::unordered_map offloading_objects_copy; { MutexLocker locker(&local_disk_segment_it->second->offloading_mutex_); local_disk_segment_it->second->enable_offloading = enable_offloading; if (enable_offloading) { - std::vector result; result.reserve( local_disk_segment_it->second->offloading_objects.size()); for (const auto& [_, task] : @@ -3967,7 +6128,8 @@ auto MasterService::OffloadObjectHeartbeat(const UUID& client_id, } for (auto& [_, task] : offloading_objects_copy) { - const auto object_id = MakeObjectIdentity(task.key, task.tenant_id); + const auto object_id = + MakeObjectIdentity(task.key, TenantId(task.tenant_id)); MetadataAccessorRW accessor(this, object_id); if (accessor.Exists()) { auto& tenant_state = accessor.GetTenantState(); @@ -3983,7 +6145,27 @@ auto MasterService::OffloadObjectHeartbeat(const UUID& client_id, } } } - return {}; + return result; +} + +auto MasterService::PollRemoveAll(const UUID& client_id) + -> tl::expected { + std::shared_lock shared_lock(snapshot_mutex_); + ScopedLocalDiskSegmentAccess local_disk_segment_access = + segment_manager_.getLocalDiskSegmentAccess(); + auto& client_local_disk_segment = + local_disk_segment_access.getClientLocalDiskSegment(); + auto local_disk_segment_it = client_local_disk_segment.find(client_id); + if (local_disk_segment_it == client_local_disk_segment.end()) { + return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND); + } + bool result; + { + MutexLocker locker(&local_disk_segment_it->second->offloading_mutex_); + result = local_disk_segment_it->second->pending_remove_all; + local_disk_segment_it->second->pending_remove_all = false; + } + return result; } auto MasterService::ReportSsdCapacity(const UUID& client_id, @@ -4028,19 +6210,69 @@ auto MasterService::NotifyOffloadSuccess( if (tasks.size() != metadatas.size()) { return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } + std::shared_ptr local_disk_segment; + { + ScopedLocalDiskSegmentAccess ssd_access = + segment_manager_.getLocalDiskSegmentAccess(); + auto& client_segments = ssd_access.getClientLocalDiskSegment(); + auto disk_it = client_segments.find(client_id); + if (disk_it != client_segments.end()) { + local_disk_segment = disk_it->second; + } + } + for (size_t i = 0; i < tasks.size(); ++i) { const auto& task = tasks[i]; const auto& metadata = metadatas[i]; - const auto object_id = MakeObjectIdentity(task.key, task.tenant_id); + const TenantId task_tenant = enable_multi_tenants_ + ? TenantId(task.tenant_id) + : TenantId::Default(); + const auto request_object_id = + MakeObjectIdentityForRequest(task.key, task_tenant); + + // NACK sentinel: offload failed on worker. Clean up the + // offloading_task + dec_refcnt but skip AddReplica. + if (metadata.data_size < 0) { + std::shared_lock shared_lock(snapshot_mutex_); + MetadataAccessorRW accessor(this, request_object_id); + if (accessor.Exists()) { + auto& tenant_state = accessor.GetTenantState(); + auto task_it = tenant_state.offloading_tasks.find( + request_object_id.user_key); + if (task_it != tenant_state.offloading_tasks.end()) { + auto source = accessor.Get().GetReplicaByID( + task_it->second.source_id); + if (source != nullptr) { + source->dec_refcnt(); + } + tenant_state.offloading_tasks.erase(task_it); + } + } + continue; + } - // Release refcnt and clear offloading task. + Replica replica(client_id, metadata.data_size, + metadata.transport_endpoint, ReplicaStatus::COMPLETE); + bool handled_existing_object = false; + bool added_new_local_disk_replica = false; { - MetadataAccessorRW accessor(this, object_id); + std::shared_lock shared_lock(snapshot_mutex_); + MetadataAccessorRW accessor(this, request_object_id); if (accessor.Exists()) { auto& obj_metadata = accessor.Get(); auto& tenant_state = accessor.GetTenantState(); - auto task_it = - tenant_state.offloading_tasks.find(object_id.user_key); + auto task_it = tenant_state.offloading_tasks.find( + request_object_id.user_key); + if (task_it != tenant_state.offloading_tasks.end() && + replica.type() != ReplicaType::LOCAL_DISK) { + LOG(ERROR) << "Invalid replica type: " << replica.type() + << ". Expected ReplicaType::LOCAL_DISK."; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + // Existing orphan objects can only bypass tenant registration + // for a master-admitted offload completion. Without this task + // marker, fall through to the regular registration check. if (task_it != tenant_state.offloading_tasks.end()) { auto source = obj_metadata.GetReplicaByID(task_it->second.source_id); @@ -4048,23 +6280,75 @@ auto MasterService::NotifyOffloadSuccess( source->dec_refcnt(); } tenant_state.offloading_tasks.erase(task_it); + + if (!obj_metadata.HasReplica( + &Replica::fn_is_local_disk_replica)) { + std::vector replicas; + replicas.emplace_back(std::move(replica)); + obj_metadata.AddReplicas(std::move(replicas)); + auto& shard = accessor.GetShard(); + shard.OnDiskReplicaAdded(obj_metadata); + SyncCacheTotalAccounting(obj_metadata); + added_new_local_disk_replica = true; + } else { + obj_metadata.VisitReplicas( + [client_id](const Replica& rep) { + return rep.type() == ReplicaType::LOCAL_DISK && + rep.get_descriptor() + .get_local_disk_descriptor() + .client_id == client_id; + }, + [&replica](Replica& rep) { + rep.get_descriptor() + .get_local_disk_descriptor() + .transport_endpoint = + replica.get_descriptor() + .get_local_disk_descriptor() + .transport_endpoint; + rep.get_descriptor() + .get_local_disk_descriptor() + .object_size = + replica.get_descriptor() + .get_local_disk_descriptor() + .object_size; + }); + } + handled_existing_object = true; } } } - // Add LOCAL_DISK replica. - Replica replica(client_id, metadata.data_size, - metadata.transport_endpoint, ReplicaStatus::COMPLETE); - auto res = AddReplica(client_id, object_id.user_key, - object_id.tenant_id, replica); - if (!res && res.error() != ErrorCode::OBJECT_NOT_FOUND) { - LOG(ERROR) << "Failed to add replica: error=" << res.error() - << ", client_id=" << client_id - << ", tenant_id=" << object_id.tenant_id - << ", key=" << object_id.user_key; - return tl::make_unexpected(res.error()); + if (!handled_existing_object) { + auto normalized_tenant_result = + ResolveTenantIdForWrite(request_object_id.tenant_id); + if (!normalized_tenant_result) { + return tl::make_unexpected(normalized_tenant_result.error()); + } + const ObjectIdentity object_id{ + std::move(normalized_tenant_result.value()), + request_object_id.user_key}; + + auto res = AddReplica(client_id, object_id.user_key, + object_id.tenant_id, replica); + if (!res) { + if (res.error() == ErrorCode::OBJECT_NOT_FOUND) { + continue; + } + LOG(ERROR) << "Failed to add replica: error=" << res.error() + << ", client_id=" << client_id + << ", tenant_id=" << object_id.tenant_id.value() + << ", key=" << object_id.user_key; + return tl::make_unexpected(res.error()); + } + added_new_local_disk_replica = res.value(); + } + if (local_disk_segment && metadata.data_size > 0 && + added_new_local_disk_replica) { + local_disk_segment->ssd_used_bytes.fetch_add( + metadata.data_size, std::memory_order_relaxed); } } + return {}; } @@ -4084,7 +6368,6 @@ tl::expected MasterService::PushOffloadingQueue( local_disk_segment_access.getClientByName(); auto client_id_it = client_by_name.find(segment_name_it.value()); if (client_id_it == client_by_name.end()) { - LOG(ERROR) << "Segment " << segment_name_it.value() << " not found"; return tl::make_unexpected(ErrorCode::SEGMENT_NOT_FOUND); } auto& client_local_disk_segment = @@ -4106,8 +6389,8 @@ tl::expected MasterService::PushOffloadingQueue( .get_memory_descriptor() .buffer_descriptor.size_; auto res = local_disk_segment_it->second->offloading_objects.emplace( - MakeTenantScopedStorageKey(object_id.tenant_id, object_id.user_key), - OffloadTaskItem{.tenant_id = object_id.tenant_id, + object_id.tenant_id.MakeScopedKey(object_id.user_key), + OffloadTaskItem{.tenant_id = object_id.tenant_id.value(), .key = object_id.user_key, .size = size}); if (!res.second) { @@ -4142,9 +6425,9 @@ tl::expected MasterService::PushPromotionQueue( } MutexLocker locker(&local_disk_segment_it->second->offloading_mutex_); auto res = local_disk_segment_it->second->promotion_objects.emplace( - MakeTenantScopedStorageKey(object_id.tenant_id, object_id.user_key), + object_id.tenant_id.MakeScopedKey(object_id.user_key), PromotionTaskItem{ - .tenant_id = object_id.tenant_id, + .tenant_id = object_id.tenant_id.value(), .key = object_id.user_key, .size = static_cast(source_replica.get_descriptor() .get_local_disk_descriptor() @@ -4155,13 +6438,315 @@ tl::expected MasterService::PushPromotionQueue( return {}; } -void MasterService::TryPushPromotionQueue(const ObjectIdentity& object_id) { - if (!promotion_on_hit_ || !promotion_sketch_) { +// --- Promotion retry candidate helpers --- + +void MasterService::DecrementCandidateCount() { + uint64_t count = promotion_candidate_count_.load(std::memory_order_relaxed); + while (count > 0) { + if (promotion_candidate_count_.compare_exchange_weak( + count, count - 1, std::memory_order_relaxed)) { + return; + } + } +} + +void MasterService::EraseCandidate(TenantState& tenant_state, + const std::string& key) { + if (tenant_state.promotion_candidates.erase(key) > 0) { + DecrementCandidateCount(); + } +} + +void MasterService::EraseCandidate(const ObjectIdentity& object_id) { + MetadataShardAccessorRW shard( + this, getMetadataShardIndex(object_id.tenant_id, object_id.user_key)); + auto tenant_it = shard->tenants.find(object_id.tenant_id); + if (tenant_it == shard->tenants.end()) return; + EraseCandidate(tenant_it->second, object_id.user_key); + if (tenant_it->second.Empty()) { + shard->tenants.erase(tenant_it); + } +} + +void MasterService::RecordOrUpdateCandidate(TenantState& tenant_state, + const std::string& key, + uint8_t sketch_score, + PromotionCandidateReason reason, + ErrorCode last_error) { + const auto now = std::chrono::steady_clock::now(); + auto it = tenant_state.promotion_candidates.find(key); + if (it != tenant_state.promotion_candidates.end()) { + // Update existing entry: refresh last_seen, reset + // retry_after/retry_count. + it->second.last_seen = now; + it->second.last_reason = reason; + it->second.last_error = last_error; + if (sketch_score > it->second.sketch_score) { + it->second.sketch_score = sketch_score; + } + it->second.retry_after = now; + it->second.retry_count = 0; + return; + } + + // Reserve a slot in the global candidate limit. + uint64_t count = promotion_candidate_count_.load(std::memory_order_relaxed); + while (count < kPromotionCandidateLimit) { + if (promotion_candidate_count_.compare_exchange_weak( + count, count + 1, std::memory_order_relaxed)) { + break; + } + } + if (count >= kPromotionCandidateLimit) { + VLOG(1) << "promotion_candidate_dropped key=" << key + << " reason=global_limit"; + MasterMetricManager::instance().inc_promotion_candidate_dropped_limit(); return; } + + auto [emplace_it, inserted] = tenant_state.promotion_candidates.emplace( + key, PromotionCandidate{.sketch_score = sketch_score, + .first_seen = now, + .last_seen = now, + .retry_after = now, + .last_reason = reason, + .last_error = last_error, + .retry_count = 0}); + if (inserted) { + MasterMetricManager::instance().inc_promotion_candidate_recorded(); + VLOG(1) << "promotion_candidate_recorded key=" << key; + } else { + DecrementCandidateCount(); + } +} + +std::chrono::milliseconds MasterService::CandidateBackoff( + uint32_t retry_count) const { + uint64_t backoff_ms = + static_cast(kPromotionCandidateInitialBackoff.count()); + for (uint32_t i = 1; i < retry_count; ++i) { + backoff_ms = std::min( + backoff_ms * 2, + static_cast(kPromotionCandidateMaxBackoff.count())); + } + return std::chrono::milliseconds(backoff_ms); +} + +bool MasterService::IsTransientResult(PromotionQueueResult result) const { + return result == PromotionQueueResult::kWatermarkRejected || + result == PromotionQueueResult::kQueueCapRejected || + result == PromotionQueueResult::kPushFailed; +} + +void MasterService::BackoffCandidate(const ObjectIdentity& object_id, + PromotionQueueResult result) { + const auto now = std::chrono::steady_clock::now(); + MetadataShardAccessorRW shard( + this, getMetadataShardIndex(object_id.tenant_id, object_id.user_key)); + auto tenant_it = shard->tenants.find(object_id.tenant_id); + if (tenant_it == shard->tenants.end()) return; + auto& tenant_state = tenant_it->second; + auto candidate_it = + tenant_state.promotion_candidates.find(object_id.user_key); + if (candidate_it == tenant_state.promotion_candidates.end()) return; + + auto& c = candidate_it->second; + c.retry_count++; + if (result == PromotionQueueResult::kWatermarkRejected) { + c.last_reason = PromotionCandidateReason::kWatermark; + c.last_error = ErrorCode::OK; + } else if (result == PromotionQueueResult::kQueueCapRejected) { + c.last_reason = PromotionCandidateReason::kQueueCap; + c.last_error = ErrorCode::OK; + } else { + c.last_reason = PromotionCandidateReason::kPushFailed; + } + + const bool ttl_expired = now - c.last_seen >= kPromotionCandidateTtl; + if (ttl_expired || c.retry_count >= kPromotionCandidateMaxRetries) { + VLOG(1) << "promotion_candidate_gave_up key=" << object_id.user_key + << " retries=" << c.retry_count; + EraseCandidate(tenant_state, object_id.user_key); + MasterMetricManager::instance() + .inc_promotion_candidate_expired_evaluated(); + } else { + c.retry_after = now + CandidateBackoff(c.retry_count); + } + + if (tenant_state.Empty()) { + shard->tenants.erase(tenant_it); + } +} + +void MasterService::ClearCandidatesForReload() { + for (size_t i = 0; i < kNumShards; ++i) { + MetadataShardAccessorRW shard(this, i); + for (auto& [tenant_id, tenant_state] : shard->tenants) { + (void)tenant_id; + tenant_state.promotion_candidates.clear(); + } + } + promotion_candidate_count_.store(0, std::memory_order_relaxed); + promotion_retry_cursor_.store(0, std::memory_order_relaxed); + promotion_in_flight_.store(0, std::memory_order_relaxed); +} + +size_t MasterService::RunPromotionCandidateRetry() { + return RunPromotionCandidateRetry(kPromotionRetryShardBatch); +} + +size_t MasterService::RunPromotionCandidateRetryForTesting() { + return RunPromotionCandidateRetry(kNumShards); +} + +size_t MasterService::CountCandidatesForTesting(const TenantId& tenant_id) { + size_t count = 0; + std::shared_lock lock(snapshot_mutex_); + for (size_t i = 0; i < kNumShards; i++) { + MetadataShardAccessorRO shard(this, i); + auto it = shard->tenants.find(tenant_id); + if (it != shard->tenants.end()) { + count += it->second.promotion_candidates.size(); + } + } + return count; +} + +void MasterService::ResetCandidateBackoffsForTesting() { + const auto epoch = std::chrono::steady_clock::time_point{}; + for (size_t i = 0; i < kNumShards; i++) { + MetadataShardAccessorRW shard(this, i); + for (auto& [tenant_id, tenant_state] : shard->tenants) { + (void)tenant_id; + for (auto& [key, candidate] : tenant_state.promotion_candidates) { + (void)key; + candidate.retry_after = epoch; + } + } + } +} + +size_t MasterService::RunPromotionCandidateRetry(size_t max_shards_to_scan) { + if (!promotion_on_hit_ || + promotion_candidate_count_.load(std::memory_order_relaxed) == 0) { + return 0; + } + + const auto now = std::chrono::steady_clock::now(); + std::vector due_candidates; + due_candidates.reserve(kPromotionRetryBatchSize); + + const size_t shards_to_scan = std::min(max_shards_to_scan, kNumShards); + if (shards_to_scan == 0) return 0; + const size_t start_shard = promotion_retry_cursor_.fetch_add( + shards_to_scan, std::memory_order_relaxed) % + kNumShards; + + { + std::shared_lock snap_lock(snapshot_mutex_); + for (size_t scanned = 0; + scanned < shards_to_scan && + due_candidates.size() < kPromotionRetryBatchSize; + ++scanned) { + const size_t i = (start_shard + scanned) % kNumShards; + MetadataShardAccessorRW shard(this, i); + for (auto tenant_it = shard->tenants.begin(); + tenant_it != shard->tenants.end() && + due_candidates.size() < kPromotionRetryBatchSize;) { + auto& tenant_state = tenant_it->second; + for (auto cit = tenant_state.promotion_candidates.begin(); + cit != tenant_state.promotion_candidates.end() && + due_candidates.size() < kPromotionRetryBatchSize;) { + const auto& key = cit->first; + auto& c = cit->second; + + const bool ttl_expired = + now - c.last_seen >= kPromotionCandidateTtl; + if (ttl_expired || + c.retry_count >= kPromotionCandidateMaxRetries) { + VLOG(1) << "promotion_candidate_expired key=" << key + << " retry_count=" << c.retry_count; + const uint32_t saved_retry_count = c.retry_count; + cit = tenant_state.promotion_candidates.erase(cit); + DecrementCandidateCount(); + // retry_count == 0: scheduler never reached this + // candidate before TTL elapsed — scan budget was + // too small. retry_count > 0: scheduler evaluated + // it but gave up after retries or TTL. + if (saved_retry_count == 0) { + MasterMetricManager::instance() + .inc_promotion_candidate_expired_unevaluated(); + } else { + MasterMetricManager::instance() + .inc_promotion_candidate_expired_evaluated(); + } + continue; + } + if (c.retry_after > now) { + ++cit; + continue; + } + + // Quick pre-filter under shard lock to avoid adding + // candidates that are obviously ineligible. + auto meta_it = tenant_state.metadata.find(key); + if (meta_it == tenant_state.metadata.end() || + !meta_it->second.IsValid() || + tenant_state.processing_keys.count(key) > 0 || + tenant_state.promotion_tasks.count(key) > 0 || + meta_it->second.HasReplica( + &Replica::fn_is_memory_replica) || + !meta_it->second.HasReplica( + &Replica::fn_is_local_disk_replica)) { + cit = tenant_state.promotion_candidates.erase(cit); + DecrementCandidateCount(); + continue; + } + + due_candidates.push_back(ObjectIdentity{ + .tenant_id = tenant_it->first, .user_key = key}); + ++cit; + } + + if (tenant_state.Empty()) { + tenant_it = shard->tenants.erase(tenant_it); + } else { + ++tenant_it; + } + } + } + } + + size_t queued = 0; + { + std::shared_lock snap_lock(snapshot_mutex_); + for (const auto& object_id : due_candidates) { + const auto result = + TryPushPromotionQueue(object_id, /*record_candidate=*/false); + if (result == PromotionQueueResult::kQueued) { + queued++; + MasterMetricManager::instance() + .inc_promotion_candidate_admitted(); + } else if (IsTransientResult(result)) { + MasterMetricManager::instance() + .inc_promotion_candidate_admission_rejected(); + BackoffCandidate(object_id, result); + } else { + EraseCandidate(object_id); + } + } + } + + return queued; +} + +MasterService::PromotionQueueResult MasterService::TryPushPromotionQueue( + const ObjectIdentity& object_id, bool record_candidate) { + if (!promotion_on_hit_ || !promotion_sketch_) { + return PromotionQueueResult::kDisabled; + } const auto& key = object_id.user_key; - const auto admission_key = - MakeTenantScopedStorageKey(object_id.tenant_id, key); + const auto admission_key = object_id.tenant_id.MakeScopedKey(key); // Frequency gate: bump and compare against the threshold. The sketch // returns uint8_t (saturating at 255); promotion_admission_threshold_ @@ -4171,7 +6756,7 @@ void MasterService::TryPushPromotionQueue(const ObjectIdentity& object_id) { const uint8_t freq = promotion_sketch_->increment(admission_key); if (freq < promotion_admission_threshold_) { MasterMetricManager::instance().inc_promotion_rejected_frequency(); - return; + return PromotionQueueResult::kFrequencyRejected; } // Watermark gate: don't promote if DRAM is already under eviction @@ -4181,7 +6766,15 @@ void MasterService::TryPushPromotionQueue(const ObjectIdentity& object_id) { MasterMetricManager::instance().get_global_mem_used_ratio(); if (used_ratio >= eviction_high_watermark_ratio_) { MasterMetricManager::instance().inc_promotion_rejected_watermark(); - return; + if (record_candidate) { + MetadataAccessorRW accessor(this, object_id); + if (accessor.Exists()) { + RecordOrUpdateCandidate(accessor.GetTenantState(), key, freq, + PromotionCandidateReason::kWatermark, + ErrorCode::OK); + } + } + return PromotionQueueResult::kWatermarkRejected; } // Acquire a fresh RW shard accessor for dedup, refcnt-pin, and task @@ -4189,18 +6782,38 @@ void MasterService::TryPushPromotionQueue(const ObjectIdentity& object_id) { // its RO accessor. MetadataAccessorRW accessor(this, object_id); if (!accessor.Exists()) { - return; + return PromotionQueueResult::kNotFound; } auto& metadata = accessor.Get(); auto& tenant_state = accessor.GetTenantState(); + // A primary Put/Upsert owns all PROCESSING replicas while the key is in + // processing_keys. Promotion must not establish a second owner. + if (accessor.InProcessing()) { + EraseCandidate(tenant_state, key); + return PromotionQueueResult::kAlreadyInFlight; + } + // Dedup: don't queue twice if a promotion is already in flight or if a // MEMORY replica has appeared since GetReplicaList observed only-disk. if (tenant_state.promotion_tasks.count(key) > 0) { - return; + EraseCandidate(tenant_state, key); + return PromotionQueueResult::kAlreadyInFlight; } if (metadata.HasReplica(&Replica::fn_is_memory_replica)) { - return; + EraseCandidate(tenant_state, key); + return PromotionQueueResult::kMemoryReplicaPresent; + } + + // Find the LOCAL_DISK source replica. + Replica* source = nullptr; + metadata.VisitReplicas(&Replica::fn_is_local_disk_replica, + [&source](Replica& r) { + if (source == nullptr) source = &r; + }); + if (source == nullptr) { + EraseCandidate(tenant_state, key); + return PromotionQueueResult::kNoLocalDiskSource; } // Cap gate: read the cluster-wide in-flight count. Soft cap — a @@ -4213,17 +6826,12 @@ void MasterService::TryPushPromotionQueue(const ObjectIdentity& object_id) { if (promotion_in_flight_.load(std::memory_order_relaxed) >= promotion_queue_limit_) { MasterMetricManager::instance().inc_promotion_rejected_cap(); - return; - } - - // Find the LOCAL_DISK source replica. - Replica* source = nullptr; - metadata.VisitReplicas(&Replica::fn_is_local_disk_replica, - [&source](Replica& r) { - if (source == nullptr) source = &r; - }); - if (source == nullptr) { - return; + if (record_candidate) { + RecordOrUpdateCandidate(tenant_state, key, freq, + PromotionCandidateReason::kQueueCap, + ErrorCode::OK); + } + return PromotionQueueResult::kQueueCapRejected; } // Pin the source replica. @@ -4237,7 +6845,21 @@ void MasterService::TryPushPromotionQueue(const ObjectIdentity& object_id) { source->dec_refcnt(); VLOG(1) << "promotion_push_failed key=" << key << " error=" << push_result.error(); - return; + if (push_result.error() == ErrorCode::OBJECT_ALREADY_EXISTS) { + EraseCandidate(tenant_state, key); + return PromotionQueueResult::kAlreadyInFlight; + } + if (push_result.error() == ErrorCode::SEGMENT_NOT_FOUND || + push_result.error() == ErrorCode::INVALID_PARAMS) { + EraseCandidate(tenant_state, key); + return PromotionQueueResult::kNoLocalDiskSource; + } + if (record_candidate) { + RecordOrUpdateCandidate(tenant_state, key, freq, + PromotionCandidateReason::kPushFailed, + push_result.error()); + } + return PromotionQueueResult::kPushFailed; } // Capture the holder client_id so NotifyPromotionSuccess can reject @@ -4247,6 +6869,7 @@ void MasterService::TryPushPromotionQueue(const ObjectIdentity& object_id) { // Record the in-flight task. alloc_id is filled in by // PromotionAllocStart once the new MEMORY replica is staged. + EraseCandidate(tenant_state, key); tenant_state.promotion_tasks.emplace( key, PromotionTask{.source_id = source->id(), .alloc_id = 0, @@ -4257,6 +6880,7 @@ void MasterService::TryPushPromotionQueue(const ObjectIdentity& object_id) { MasterMetricManager::instance().inc_promotion_in_flight(); MasterMetricManager::instance().inc_promotion_admitted(); VLOG(1) << "promotion_queued key=" << key << " size=" << object_size; + return PromotionQueueResult::kQueued; } auto MasterService::PromotionObjectHeartbeat(const UUID& client_id) @@ -4288,17 +6912,26 @@ auto MasterService::PromotionObjectHeartbeat(const UUID& client_id) } auto MasterService::PromotionAllocStart( - const UUID& client_id, const std::string& key, const std::string& tenant_id, + const UUID& client_id, const std::string& key, const TenantId& tenant_id, uint64_t size, const std::vector& preferred_segments) -> tl::expected { + auto normalized_tenant_result = ResolveTenantIdForWrite(tenant_id); + if (!normalized_tenant_result) { + return tl::make_unexpected(normalized_tenant_result.error()); + } + const ObjectIdentity object_id{std::move(normalized_tenant_result.value()), + key}; std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); MetadataAccessorRW accessor(this, object_id); if (!accessor.Exists()) { return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); } auto& metadata = accessor.Get(); + if (accessor.InProcessing()) { + return tl::make_unexpected(ErrorCode::REPLICA_IS_NOT_READY); + } + // Verify the in-flight task still exists before allocating. The // reaper can sweep it between the holder's heartbeat and this // AllocStart call (a hung client, GC pause, or HA failover can @@ -4403,10 +7036,10 @@ auto MasterService::PromotionAllocStart( auto MasterService::NotifyPromotionSuccess(const UUID& client_id, const std::string& key, - const std::string& tenant_id) + const TenantId& tenant_id) -> tl::expected { std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); + const auto object_id = MakeObjectIdentityForRequest(key, tenant_id); MetadataAccessorRW accessor(this, object_id); if (!accessor.Exists()) { return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); @@ -4434,8 +7067,48 @@ auto MasterService::NotifyPromotionSuccess(const UUID& client_id, Replica* staged = metadata.GetReplicaByID(task_it->second.alloc_id); if (staged != nullptr && staged->is_memory_replica() && staged->is_processing()) { + std::optional batch_reservation; + if (enable_ha_ && enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + return tl::make_unexpected(reservation.error()); + } + batch_reservation = std::move(reservation.value()); + } staged->mark_complete(); committed = true; + if (enable_oplog_ && ordered_oplog_writer_) { + std::vector post; + metadata.VisitReplicas(&Replica::fn_is_completed, + [&post](const Replica& replica) { + post.push_back(replica.get_descriptor()); + }); + + const auto payload = + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, post, metadata.group_id, + metadata.data_type); + if (batch_reservation) { + auto persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(*batch_reservation), OpType::PUT_END, + tenant_id.value(), key, payload, nullptr); + if (!persist_result) { + LOG(WARNING) + << "NotifyPromotionSuccess: PUT_END persist failed " + << "for key=" << key + << ", err=" << static_cast(persist_result.error()); + } + } else { + auto persist_result = AppendOpLogVisibleBeforeDurable( + OpType::PUT_END, tenant_id.value(), key, payload); + if (!persist_result) { + LOG(WARNING) + << "NotifyPromotionSuccess: PUT_END persist failed " + << "for key=" << key + << ", err=" << static_cast(persist_result.error()); + } + } + } } // Drop the source LOCAL_DISK replica's refcnt and erase the task. @@ -4485,8 +7158,8 @@ auto MasterService::NotifyPromotionSuccess(const UUID& client_id, auto it = client_local_disk_segment.find(client_id); if (it != client_local_disk_segment.end()) { MutexLocker locker(&it->second->offloading_mutex_); - it->second->promotion_objects.erase(MakeTenantScopedStorageKey( - object_id.tenant_id, object_id.user_key)); + it->second->promotion_objects.erase( + object_id.tenant_id.MakeScopedKey(object_id.user_key)); } } @@ -4498,10 +7171,10 @@ auto MasterService::NotifyPromotionSuccess(const UUID& client_id, auto MasterService::NotifyPromotionFailure(const UUID& client_id, const std::string& key, - const std::string& tenant_id) + const TenantId& tenant_id) -> tl::expected { std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); + const auto object_id = MakeObjectIdentityForRequest(key, tenant_id); MetadataAccessorRW accessor(this, object_id); if (!accessor.Exists()) { return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); @@ -4555,8 +7228,8 @@ auto MasterService::NotifyPromotionFailure(const UUID& client_id, auto it = client_local_disk_segment.find(client_id); if (it != client_local_disk_segment.end()) { MutexLocker locker(&it->second->offloading_mutex_); - it->second->promotion_objects.erase(MakeTenantScopedStorageKey( - object_id.tenant_id, object_id.user_key)); + it->second->promotion_objects.erase( + object_id.tenant_id.MakeScopedKey(object_id.user_key)); } } @@ -4617,6 +7290,10 @@ void MasterService::EvictionThreadFunc() { } #endif + if (promotion_candidate_count_.load(std::memory_order_relaxed) > 0) { + RunPromotionCandidateRetry(); + } + std::this_thread::sleep_for( std::chrono::milliseconds(kEvictionThreadSleepMs)); } @@ -4647,24 +7324,83 @@ void MasterService::DiscardExpiredProcessingReplicas( if (!metadata.IsValid() || metadata.AllReplicas(&Replica::fn_is_completed)) { if (!metadata.IsValid()) { - EraseMetadata(tenant_state, it, tenant_it->first); + auto next_key_it = std::next(key_it); + EraseMetadata(tenant_state, it, tenant_it->first, + QuotaEraseMode::kFull, &shard); + key_it = next_key_it; + } else { + key_it = tenant_state.processing_keys.erase(key_it); } - key_it = tenant_state.processing_keys.erase(key_it); continue; } const auto ttl = metadata.put_start_time + put_start_release_timeout_sec_; if (ttl < now) { - auto replicas = - metadata.PopReplicas(&Replica::fn_is_processing); - if (!replicas.empty()) { + const bool had_complete_replica = + metadata.HasReplica(&Replica::fn_is_completed); + // Predict post-discard descriptors WITHOUT mutating: drop + // PROCESSING replicas; keep COMPLETE replicas. + auto post_descriptors = BuildRemainingReplicaDescriptors( + metadata, &Replica::fn_is_processing); + const bool would_invalidate = post_descriptors.empty(); + + if (had_complete_replica && enable_oplog_ && + ordered_oplog_writer_) { + tl::expected persist_result; + if (would_invalidate) { + persist_result = AppendOpLogWithDurableFinalize( + OpType::REMOVE, tenant_it->first.value(), *key_it, + {}, + enable_oplog_ + ? [this, ttl](const OpLogEntry& durable_entry) { + FinalizeExpiredProcessingReplicasAfterDurable( + durable_entry, ttl); + } + : DurableFinalizeCallback{}); + } else { + persist_result = AppendOpLogWithDurableFinalize( + OpType::PUT_END, tenant_it->first.value(), *key_it, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, + post_descriptors, metadata.group_id, + metadata.data_type), + enable_oplog_ + ? [this, ttl](const OpLogEntry& durable_entry) { + FinalizeExpiredProcessingReplicasAfterDurable( + durable_entry, ttl); + } + : DurableFinalizeCallback{}); + } + if (!persist_result) { + LOG(WARNING) << "DiscardExpiredProcessingReplicas: " + "OpLog persist failed for key=" + << *key_it << ", err=" + << static_cast(persist_result.error()) + << ", deferring discard"; + ++key_it; + continue; + } + if (enable_oplog_) { + ++key_it; + continue; + } + } + + // Persist OK (or HA disabled / never published) — apply. + auto replicas = + metadata.PopReplicas(&Replica::fn_is_processing); + if (!replicas.empty()) { discarded_replicas.emplace_back(std::move(replicas), ttl); } if (!metadata.IsValid()) { - EraseMetadata(tenant_state, it, tenant_it->first); + auto next_key_it = std::next(key_it); + EraseMetadata(tenant_state, it, tenant_it->first, + QuotaEraseMode::kFull, &shard); + key_it = next_key_it; + } else { + key_it = tenant_state.processing_keys.erase(key_it); } - key_it = tenant_state.processing_keys.erase(key_it); continue; } key_it++; @@ -4676,6 +7412,8 @@ void MasterService::DiscardExpiredProcessingReplicas( if (metadata_it == tenant_state.metadata.end()) { LOG(ERROR) << "Key " << task_it->first << " was removed with ongoing replication task"; + AbortTenantQuota(tenant_it->first, + task_it->second.reserved_quota_charge_bytes); task_it = tenant_state.replication_tasks.erase(task_it); continue; } @@ -4688,25 +7426,90 @@ void MasterService::DiscardExpiredProcessingReplicas( } auto& metadata = metadata_it->second; + + const bool had_complete_replica = + metadata.HasReplica(&Replica::fn_is_completed); + auto& replica_ids = task_it->second.replica_ids; + + const auto target_pred = [&replica_ids](const Replica& r) { + return std::find(replica_ids.begin(), replica_ids.end(), + r.id()) != replica_ids.end(); + }; + // Predict post-discard descriptor list WITHOUT mutating: drop + // task target replicas; keep the rest of the COMPLETE replicas. + auto post_descriptors = + BuildRemainingReplicaDescriptors(metadata, target_pred); + const bool would_invalidate = post_descriptors.empty(); + + if (had_complete_replica && enable_oplog_ && + ordered_oplog_writer_) { + tl::expected persist_result; + auto source_id = task_it->second.source_id; + auto target_ids = replica_ids; + if (would_invalidate) { + persist_result = AppendOpLogWithDurableFinalize( + OpType::REMOVE, tenant_it->first.value(), + task_it->first, {}, + enable_oplog_ + ? [this, source_id, + target_ids = std::move(target_ids), + ttl](const OpLogEntry& durable_entry) { + FinalizeExpiredReplicationTaskAfterDurable( + durable_entry, source_id, target_ids, + ttl); + } + : DurableFinalizeCallback{}); + } else { + persist_result = AppendOpLogWithDurableFinalize( + OpType::PUT_END, tenant_it->first.value(), + task_it->first, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, post_descriptors, + metadata.group_id, metadata.data_type), + enable_oplog_ + ? [this, source_id, + target_ids = std::move(target_ids), + ttl](const OpLogEntry& durable_entry) { + FinalizeExpiredReplicationTaskAfterDurable( + durable_entry, source_id, target_ids, + ttl); + } + : DurableFinalizeCallback{}); + } + if (!persist_result) { + LOG(WARNING) + << "DiscardExpiredProcessingReplicas: OpLog persist " + "failed for replication task key=" + << task_it->first + << ", err=" << static_cast(persist_result.error()) + << ", deferring discard"; + ++task_it; + continue; + } + if (enable_oplog_) { + ++task_it; + continue; + } + } + auto source = metadata.GetReplicaByID(task_it->second.source_id); if (source != nullptr) { source->dec_refcnt(); } - auto& replica_ids = task_it->second.replica_ids; - auto replicas = PopReplicasWithCacheTotalAccounting( - metadata, [&replica_ids](const Replica& replica) { - auto it = std::find(replica_ids.begin(), replica_ids.end(), - replica.id()); - return it != replica_ids.end(); - }); + auto replicas = + PopReplicasWithCacheTotalAccounting(metadata, target_pred); if (!replicas.empty()) { discarded_replicas.emplace_back(std::move(replicas), ttl); } if (!metadata.IsValid()) { - EraseMetadata(tenant_state, metadata_it, tenant_it->first); + auto next_task_it = std::next(task_it); + EraseMetadata(tenant_state, metadata_it, tenant_it->first, + QuotaEraseMode::kFull, &shard); + task_it = next_task_it; + } else { + task_it = tenant_state.replication_tasks.erase(task_it); } - task_it = tenant_state.replication_tasks.erase(task_it); } for (auto task_it = tenant_state.offloading_tasks.begin(); @@ -4726,7 +7529,7 @@ void MasterService::DiscardExpiredProcessingReplicas( } } LOG(WARNING) << "Offloading task expired for key: " - << task_it->first; + << task_it->first << " tenant=" << tenant_it->first; task_it = tenant_state.offloading_tasks.erase(task_it); } @@ -4794,1068 +7597,481 @@ uint64_t MasterService::ReleaseExpiredDiscardedReplicas( return released_cnt; } -void MasterService::SnapshotThreadFunc() { - LOG(INFO) << "[Snapshot] snapshot_thread started"; - while (snapshot_running_) { - std::this_thread::sleep_for( - std::chrono::seconds(snapshot_interval_seconds_)); - if (!enable_snapshot_) { - // Snapshot is disabled - LOG(INFO) - << "[Snapshot] Snapshot is disabled, waiting for next cycle"; - continue; - } - // Fork a child process to save current state - - std::string snapshot_id = - FormatTimestamp(std::chrono::system_clock::now()); - LOG(INFO) << "[Snapshot] Preparing to fork child process, snapshot_id=" - << snapshot_id; - - // Create pipe for child process logging - int log_pipe[2]; - if (pipe(log_pipe) == -1) { - LOG(ERROR) << "[Snapshot] Failed to create log pipe: " - << strerror(errno) << ", snapshot_id=" << snapshot_id; - continue; - } +/** + * @brief Restore master state from snapshot using three-phase architecture. + * + * Phase 1 (Repository): Load candidate snapshots from catalog + * Phase 2 (Repository + Codec): Download payloads and decode to memory + * Phase 3 (Service): Apply decoded state and rebuild metrics + * + * Attempts restore from candidates in chronological order until one succeeds. + * If all candidates fail, starts with a fresh state. + */ +void MasterService::RestoreState() { + auto* snapshot_catalog_store = snapshot_catalog_store_.get(); + if (!snapshot_catalog_store) { + LOG(ERROR) << "[Restore] Snapshot catalog store is not initialized, " + "starting fresh"; + return; + } - const std::string& snapshot_root = - snapshot_catalog_store_->GetSnapshotRoot(); - const std::string path_prefix = snapshot_root + snapshot_id + "/"; - const std::string manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE; - auto descriptor = - BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix); - if (!descriptor) { - LOG(ERROR) << "[Snapshot] Failed to build descriptor before fork, " - "snapshot_id=" - << snapshot_id - << ", code=" << toString(descriptor.error().code) - << ", msg=" << descriptor.error().message; - close(log_pipe[0]); - close(log_pipe[1]); - continue; - } + LOG(INFO) << "[Restore] Backend info: " + << snapshot_object_store_->GetConnectionInfo(); - pid_t pid; - { - std::unique_lock lock(snapshot_mutex_); - LOG(INFO) << "[Snapshot] Locking snapshot mutex, snapshot_id=" - << snapshot_id; - pid = fork(); - } - if (pid == -1) { - // Fork failed - LOG(ERROR) << "[Snapshot] Failed to fork child process for state " - "persistence: " - << strerror(errno) << ", snapshot_id=" << snapshot_id; - close(log_pipe[0]); - close(log_pipe[1]); - } else if (pid == 0) { - // Child process - // Close read end, set write end for logging - close(log_pipe[0]); - g_snapshot_log_pipe_fd = log_pipe[1]; - - // Save current state using the configured persistence mechanism - SNAP_LOG_INFO("[Snapshot] Child process started, snapshot_id={}", - snapshot_id); - auto result = PersistState(descriptor.value()); - if (!result) { - SNAP_LOG_ERROR( - "[Snapshot] Child process failed to persist state, " - "snapshot_id={},code={},msg={}", - snapshot_id, toString(result.error().code), - result.error().message); - close(log_pipe[1]); - _exit(1); // Exit child process with error - } - SNAP_LOG_INFO( - "[Snapshot] Child process successfully persisted state, " - "snapshot_id={}", - snapshot_id); - - close(log_pipe[1]); - _exit(0); // Exit child process successfully - } else { - // Parent process - // Close write end, pass read end to wait function - close(log_pipe[1]); - WaitForSnapshotChild(pid, snapshot_id, log_pipe[0]); - close(log_pipe[0]); - } - } - LOG(INFO) << "[Snapshot] snapshot_thread stopped"; -} - -void MasterService::WaitForSnapshotChild(pid_t pid, - const std::string& snapshot_id, - int log_pipe_fd) { - // Default 5 minute timeout - const int64_t timeout_seconds = snapshot_child_timeout_seconds_; - - LOG(INFO) - << "[Snapshot] waiting for child process to complete, snapshot_id=" - << snapshot_id << ", child_pid=" << pid - << ", timeout=" << timeout_seconds << "s"; - - // Set pipe to non-blocking mode - int flags = fcntl(log_pipe_fd, F_GETFL, 0); - if (flags == -1 || fcntl(log_pipe_fd, F_SETFL, flags | O_NONBLOCK) == -1) { - LOG(WARNING) << "[Snapshot] Failed to set pipe non-blocking: " - << strerror(errno); - } - - // Buffer for reading child logs - char buf[4096]; - std::string log_buffer; - - // Helper lambda to read and output child logs - auto flush_child_logs = [&]() { - while (true) { - ssize_t n = read(log_pipe_fd, buf, sizeof(buf) - 1); - if (n > 0) { - buf[n] = '\0'; - log_buffer += buf; - // Output complete lines - size_t pos; - while ((pos = log_buffer.find('\n')) != std::string::npos) { - std::string line = log_buffer.substr(0, pos); - log_buffer.erase(0, pos + 1); - if (!line.empty()) { - LOG(INFO) << "[Snapshot:Child] " << line; - } - } - } else { - break; - } - } - }; + // Phase 1: Find snapshot candidates (repository responsibility) + auto latest_result = snapshot_repository_->LoadLatestSnapshot(); + std::optional latest_snapshot; + if (!latest_result) { + LOG(WARNING) << "[Restore] Failed to load latest snapshot marker: " + << toString(latest_result.error()) + << ", falling back to published snapshot listing"; + } else { + latest_snapshot = latest_result.value(); + } - // Record start time - auto start_time = std::chrono::steady_clock::now(); + auto candidates_result = + snapshot_repository_->LoadRestoreCandidates(latest_snapshot); + if (!candidates_result || candidates_result->empty()) { + LOG(ERROR) << "[Restore] No previous snapshot found, starting fresh"; + return; + } - // Use non-blocking polling to wait - while (true) { - // Read child logs first - flush_child_logs(); + // Phase 2 & 3: Try each candidate + const auto now = std::chrono::system_clock::now(); + for (const auto& snapshot : candidates_result.value()) { + ResetStateAfterFailedRestoreAttempt(); - int status; - pid_t result = waitpid(pid, &status, WNOHANG); + try { + // Phase 2a: Download payloads (repository responsibility) + auto payloads_result = + snapshot_repository_->DownloadSnapshotPayloads(snapshot); + if (!payloads_result) { + LOG(WARNING) + << "[Restore] Snapshot candidate " << snapshot.snapshot_id + << " is unusable: failed to download payloads: " + << payloads_result.error().message; + continue; + } - if (result == -1) { - LOG(ERROR) << "[Snapshot] Failed to wait for child process: " - << strerror(errno) << ", snapshot_id=" << snapshot_id - << ", child_pid=" << pid; - MasterMetricManager::instance().inc_snapshot_fail(); - return; - } else if (result == 0) { - // Child process is still running - auto elapsed = std::chrono::duration_cast( - std::chrono::steady_clock::now() - start_time) - .count(); - - if (elapsed >= timeout_seconds) { - // Timeout handling - flush remaining logs before killing - flush_child_logs(); - if (!log_buffer.empty()) { - LOG(INFO) << "[Snapshot:Child] " << log_buffer; - } - HandleChildTimeout(pid, snapshot_id); - MasterMetricManager::instance().inc_snapshot_fail(); - return; + // Phase 2b: Decode payloads (codec responsibility) + auto decode_result = + snapshot_codec_->Decode(this, payloads_result.value()); + if (!decode_result) { + LOG(WARNING) + << "[Restore] Snapshot candidate " << snapshot.snapshot_id + << " is unusable: " << decode_result.error().message; + continue; } - // Brief sleep before checking again - std::this_thread::sleep_for(std::chrono::seconds(2)); - } else { - // Child process has exited - // Flush remaining logs from child - flush_child_logs(); - // Output any remaining incomplete line - if (!log_buffer.empty()) { - LOG(INFO) << "[Snapshot:Child] " << log_buffer; + // Phase 3: Apply state (master service responsibility) + auto apply_result = ApplySnapshotState(now); + if (!apply_result) { + LOG(WARNING) + << "[Restore] Snapshot candidate " << snapshot.snapshot_id + << " is unusable: failed to apply state: " + << apply_result.error().message; + continue; } - HandleChildExit(pid, status, snapshot_id); - auto elapsed = - std::chrono::duration_cast( - std::chrono::steady_clock::now() - start_time) - .count(); - MasterMetricManager::instance().set_snapshot_duration_ms(elapsed); + LOG(INFO) << "[Restore] Successfully restored state from snapshot: " + << snapshot.snapshot_id; return; + } catch (const std::exception& e) { + LOG(WARNING) << "[Restore] Snapshot candidate " + << snapshot.snapshot_id + << " is unusable: exception during restore: " + << e.what(); + // State reset already happened at loop start; continue to next + continue; + } catch (...) { + LOG(WARNING) << "[Restore] Snapshot candidate " + << snapshot.snapshot_id + << " is unusable: unknown exception during restore"; + continue; } } -} - -void MasterService::HandleChildTimeout(pid_t pid, - const std::string& snapshot_id) { - LOG(WARNING) << "[Snapshot] Child process timeout, snapshot_id=" - << snapshot_id << ", child_pid=" << pid - << ", killing child process"; - // Try to gracefully terminate the child process - if (kill(pid, SIGTERM) == 0) { - // Wait a few seconds to see if it exits gracefully - std::this_thread::sleep_for(std::chrono::seconds(5)); + ResetStateAfterFailedRestoreAttempt(); + LOG(ERROR) << "[Restore] Failed to restore from all candidate snapshots " + << "(count=" << candidates_result->size() << "), starting fresh"; +} - // Check if it has exited - int status; - if (waitpid(pid, &status, WNOHANG) == 0) { - // Child process still not exited, force kill - LOG(WARNING) << "[Snapshot] Child process still running, force " - "killing, snapshot_id=" - << snapshot_id << ", child_pid=" << pid; - kill(pid, SIGKILL); +void MasterService::ResetStateAfterFailedRestoreAttempt() { + SegmentSerializer segment_serializer(&segment_manager_); + MetadataSerializer metadata_serializer(this); + TaskManagerSerializer task_manager_serializer(&task_manager_); - // Wait for force termination to complete - waitpid(pid, &status, 0); - LOG(WARNING) - << "[Snapshot] Child process force killed, snapshot_id=" - << snapshot_id << ", child_pid=" << pid; - } else { - LOG(INFO) << "[Snapshot] Child process terminated gracefully after " - "SIGTERM, snapshot_id=" - << snapshot_id << ", child_pid=" << pid; - } - } else { - LOG(ERROR) << "[Snapshot] Failed to send SIGTERM to child process, " - "snapshot_id=" - << snapshot_id << ", child_pid=" << pid - << ", error=" << strerror(errno); - } -} - -void MasterService::HandleChildExit(pid_t pid, int status, - const std::string& snapshot_id) { - if (WIFEXITED(status)) { - int exit_code = WEXITSTATUS(status); - if (exit_code != 0) { - LOG(ERROR) << "[Snapshot] Child process exited with error code: " - << exit_code << ", snapshot_id=" << snapshot_id - << ", child_pid=" << pid; - MasterMetricManager::instance().inc_snapshot_fail(); - } else { - LOG(INFO) << "[Snapshot] Child process successfully persisted " - "state, snapshot_id=" - << snapshot_id << ", child_pid=" << pid; - MasterMetricManager::instance().inc_snapshot_success(); - } - } else if (WIFSIGNALED(status)) { - int signal = WTERMSIG(status); - LOG(ERROR) << "[Snapshot] Child process terminated by signal: " - << signal << ", snapshot_id=" << snapshot_id - << ", child_pid=" << pid; - MasterMetricManager::instance().inc_snapshot_fail(); - } -} - -tl::expected -MasterService::ResolveSnapshotSequenceId() const { - if (!enable_ha_ || ha_backend_type_ != "etcd") { - // OpLog sequence ids start at 1. Returning 0 here is a sentinel that - // means "no persisted OpLog boundary", so a standby that later calls - // Recover(0) will replay from the first entry when oplog following is - // enabled. - return ha::OpLogSequenceId{0}; - } - -#ifndef STORE_USE_ETCD - return tl::make_unexpected(SerializationError( - ErrorCode::UNAVAILABLE_IN_CURRENT_MODE, - "etcd snapshot sequence resolution is unavailable in this build")); -#else - auto oplog_store = GetSnapshotBoundaryOpLogStore(); - if (!oplog_store) { - return tl::make_unexpected(oplog_store.error()); - } + task_manager_serializer.Reset(); + metadata_serializer.Reset(); + segment_serializer.Reset(); - uint64_t sequence_id = 0; - auto err = oplog_store.value()->GetLatestSequenceId(sequence_id); - if (err == ErrorCode::OPLOG_ENTRY_NOT_FOUND) { - return ha::OpLogSequenceId{0}; + { + std::unique_lock lock(client_mutex_); + ok_client_.clear(); } - if (err != ErrorCode::OK) { - return tl::make_unexpected(SerializationError( - err, fmt::format("failed to resolve snapshot sequence boundary: {}", - toString(err)))); + PodUUID pod_uuid; + while (client_ping_queue_.pop(pod_uuid)) { } - return static_cast(sequence_id); -#endif + MasterMetricManager::instance().reset_allocated_mem_size(); + MasterMetricManager::instance().reset_total_mem_capacity(); + MasterMetricManager::instance().reset_cache_total_nums(); } -#ifdef STORE_USE_ETCD -tl::expected -MasterService::GetSnapshotBoundaryOpLogStore() const { - if (ha_backend_connstring_.empty()) { - return tl::make_unexpected(SerializationError( - ErrorCode::INVALID_PARAMS, - "etcd snapshot sequence resolution requires a backend connstring")); - } - - std::lock_guard lock(snapshot_boundary_oplog_store_mutex_); - if (snapshot_boundary_oplog_store_ != nullptr) { - return snapshot_boundary_oplog_store_.get(); - } - - auto err = - EtcdHelper::ConnectToEtcdStoreClient(ha_backend_connstring_.c_str()); - if (err != ErrorCode::OK) { - return tl::make_unexpected(SerializationError( - err, fmt::format("failed to connect to etcd for snapshot boundary: " - "{}", - toString(err)))); - } +tl::expected MasterService::ApplySnapshotState( + const std::chrono::system_clock::time_point& now) { + // Note: Codec has already called Deserialize() on all payloads, + // so the internal state is already restored. This method handles + // post-restore cleanup and metrics rebuilding. - auto oplog_store = std::make_unique(cluster_id_); - err = oplog_store->Init(); - if (err != ErrorCode::OK) { - return tl::make_unexpected(SerializationError( - err, fmt::format("failed to initialize etcd oplog store: {}", - toString(err)))); + std::vector segment_names; + { + ScopedSegmentAccess segment_access = + segment_manager_.getSegmentAccess(); + segment_access.GetAllSegmentNames(segment_names); } - snapshot_boundary_oplog_store_ = std::move(oplog_store); - return snapshot_boundary_oplog_store_.get(); -} -#endif - -tl::expected -MasterService::BuildSnapshotDescriptor(const std::string& snapshot_id, - const std::string& manifest_path, - const std::string& object_prefix) const { - auto sequence_id = ResolveSnapshotSequenceId(); - if (!sequence_id) { - return tl::make_unexpected(sequence_id.error()); - } - - const std::string& snapshot_root = - snapshot_catalog_store_->GetSnapshotRoot(); - auto descriptor = ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor( - snapshot_root, snapshot_id); - descriptor.last_included_seq = sequence_id.value(); - descriptor.producer_view_version = view_version_; - descriptor.manifest_key = manifest_path; - descriptor.object_prefix = object_prefix; - descriptor.created_at_ms = CurrentTimeMs(); - return descriptor; -} - -tl::expected MasterService::PersistState( - const std::string& snapshot_id) { - const std::string& snapshot_root = - snapshot_catalog_store_->GetSnapshotRoot(); - const std::string path_prefix = snapshot_root + snapshot_id + "/"; - const std::string manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE; - auto descriptor = - BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix); - if (!descriptor) { - return tl::make_unexpected(descriptor.error()); - } - return PersistState(descriptor.value()); -} - -tl::expected MasterService::PersistState( - const ha::SnapshotDescriptor& descriptor) { - const std::string& snapshot_id = descriptor.snapshot_id; - const std::string& path_prefix = descriptor.object_prefix; - const std::string& manifest_path = descriptor.manifest_key; - - try { - auto* snapshot_catalog_store = GetSnapshotCatalogStore(); - if (!snapshot_catalog_store) { - return tl::make_unexpected(SerializationError( - ErrorCode::PERSISTENT_FAIL, - "snapshot catalog store is not initialized")); + // Cleanup expired metadata (unless test environment disables it) + { + const bool skip_cleanup = + std::getenv("MOONCAKE_MASTER_SERVICE_SNAPSHOT_TEST_SKIP_CLEANUP"); + if (!skip_cleanup) { + auto cleanup_now = now; + for (auto& shard : metadata_shards_) { + for (auto tenant_it = shard.tenants.begin(); + tenant_it != shard.tenants.end();) { + auto& tenant_state = tenant_it->second; + for (auto it = tenant_state.metadata.begin(); + it != tenant_state.metadata.end();) { + if (it->second.HasDiffRepStatus( + ReplicaStatus::COMPLETE) || + (it->second.IsLeaseExpired(cleanup_now) && + !it->second.IsSoftPinned(cleanup_now))) { + VLOG(1) << "clear metadata key=" << it->first; + it = EraseMetadata(tenant_state, it, + tenant_it->first); + } else { + ++it; + } + } + if (tenant_state.Empty()) { + tenant_it = shard.tenants.erase(tenant_it); + } else { + ++tenant_it; + } + } + } } - SNAP_LOG_INFO( - "[Snapshot] action=persisting_state start, snapshot_id={}, " - "serializer_type={}, version={}", - snapshot_id, SNAPSHOT_SERIALIZER_TYPE, SNAPSHOT_SERIALIZER_VERSION); - MetadataSerializer metadata_serializer(this); - SegmentSerializer segment_serializer(&segment_manager_); - TaskManagerSerializer task_manager_serializer(&task_manager_); + // Rebuild allocated memory metrics + MasterMetricManager::instance().reset_allocated_mem_size(); + RebuildCacheTotalAccounting(); + for (auto& segment_name : segment_names) { + MasterMetricManager::instance().reset_segment_allocated_mem_size( + segment_name); + } - auto metadata_result = metadata_serializer.Serialize(); - if (!metadata_result) { - SNAP_LOG_ERROR( - "[Snapshot] metadata serialization failed, snapshot_id={}, " - "code={}, msg={}", - snapshot_id, toString(metadata_result.error().code), - metadata_result.error().message); - - return tl::make_unexpected(metadata_result.error()); - } - SNAP_LOG_INFO( - "[Snapshot] metadata serialization_successful, snapshot_id={}", - snapshot_id); - - auto segment_result = segment_serializer.Serialize(); - if (!segment_result) { - SNAP_LOG_ERROR( - "[Snapshot] segment serialization failed, snapshot_id={}, " - "code={}, msg={}", - snapshot_id, toString(segment_result.error().code), - segment_result.error().message); - return tl::make_unexpected(segment_result.error()); - } - SNAP_LOG_INFO( - "[Snapshot] segment serialization_successful, snapshot_id={}", - snapshot_id); - - auto task_manager_result = task_manager_serializer.Serialize(); - if (!task_manager_result) { - SNAP_LOG_ERROR( - "[Snapshot] task manager serialization failed, snapshot_id={}, " - "code={}, msg={}", - snapshot_id, toString(task_manager_result.error().code), - task_manager_result.error().message); - return tl::make_unexpected(task_manager_result.error()); - } - SNAP_LOG_INFO( - "[Snapshot] task manager serialization_successful, snapshot_id={}", - snapshot_id); - - const auto& serialized_metadata = metadata_result.value(); - const auto& serialized_segment = segment_result.value(); - const auto& serialized_task_manager = task_manager_result.value(); - - // When backup_dir is enabled, try all uploads to ensure complete backup - // When backup_dir is disabled, use fail-fast mode - bool upload_success = true; - std::string error_msg; - SNAP_LOG_INFO("[Snapshot] Backend info: {}", - snapshot_object_store_->GetConnectionInfo()); - - // Upload metadata - std::string metadata_path = path_prefix + SNAPSHOT_METADATA_FILE; - auto upload_result = - UploadSnapshotPayloadFile(serialized_metadata, metadata_path, - SNAPSHOT_METADATA_FILE, snapshot_id); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] metadata upload failed, snapshot_id={}, " - "path={}, code={}, msg={}", - snapshot_id, metadata_path, - toString(upload_result.error().code), - upload_result.error().message); - if (!use_snapshot_backup_dir_) { - return tl::make_unexpected(upload_result.error()); - } - error_msg.append(upload_result.error().message + "\n"); - upload_success = false; - } - - // Upload segment - std::string segment_path = path_prefix + SNAPSHOT_SEGMENTS_FILE; - upload_result = - UploadSnapshotPayloadFile(serialized_segment, segment_path, - SNAPSHOT_SEGMENTS_FILE, snapshot_id); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] segment upload failed, snapshot_id={}, " - "path={}, code={}, msg={}", - snapshot_id, segment_path, toString(upload_result.error().code), - upload_result.error().message); - if (!use_snapshot_backup_dir_) { - return tl::make_unexpected(upload_result.error()); - } - error_msg.append(upload_result.error().message + "\n"); - upload_success = false; - } - - // Upload task manager - std::string task_manager_path = - path_prefix + SNAPSHOT_TASK_MANAGER_FILE; - upload_result = UploadSnapshotPayloadFile( - serialized_task_manager, task_manager_path, - SNAPSHOT_TASK_MANAGER_FILE, snapshot_id); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] task_manager upload failed, snapshot_id={}, " - "path={}, code={}, msg={}", - snapshot_id, task_manager_path, - toString(upload_result.error().code), - upload_result.error().message); - if (!use_snapshot_backup_dir_) { - return tl::make_unexpected(upload_result.error()); - } - error_msg.append(upload_result.error().message + "\n"); - upload_success = false; - } - - // Upload manifest - std::string manifest_content = - fmt::format("{}|{}|{}", SNAPSHOT_SERIALIZER_TYPE, - SNAPSHOT_SERIALIZER_VERSION, snapshot_id); - std::vector manifest_bytes(manifest_content.begin(), - manifest_content.end()); - upload_result = UploadSnapshotPayloadFile( - manifest_bytes, manifest_path, SNAPSHOT_MANIFEST_FILE, snapshot_id); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] manifest upload failed, snapshot_id={}, " - "path={}, code={}, msg={}", - snapshot_id, manifest_path, - toString(upload_result.error().code), - upload_result.error().message); - if (!use_snapshot_backup_dir_) { - return tl::make_unexpected(upload_result.error()); - } - error_msg.append(upload_result.error().message + "\n"); - upload_success = false; - } - - if (!upload_success) { - return tl::make_unexpected( - SerializationError(ErrorCode::PERSISTENT_FAIL, error_msg)); - } - - // Publish snapshot catalog entry and advance the latest marker. - std::string latest_path = - snapshot_catalog_store->GetSnapshotRoot() + SNAPSHOT_LATEST_FILE; - std::string latest_content = snapshot_id; - - auto publish_result = snapshot_catalog_store->Publish(descriptor); - if (publish_result != ErrorCode::OK) { - SNAP_LOG_ERROR( - "[Snapshot] latest update failed, snapshot_id={}, file={}, " - "code={}", - snapshot_id, latest_path, toString(publish_result)); - if (use_snapshot_backup_dir_) { - auto save_path = fs::path(snapshot_backup_dir_) / - SNAPSHOT_BACKUP_SAVE_DIR / - SNAPSHOT_LATEST_FILE; - auto save_result = - FileUtil::SaveStringToFile(latest_content, save_path); - if (!save_result) { - SNAP_LOG_ERROR( - "[Snapshot] save latest to disk failed, " - "snapshot_id={}, " - "content={}, file={}", - snapshot_id, latest_content, save_path.string()); + for (auto& shard : metadata_shards_) { + for (auto& [tenant_id, tenant_state] : shard.tenants) { + for (auto it = tenant_state.metadata.begin(); + it != tenant_state.metadata.end();) { + for (auto& replica : it->second.GetAllReplicas()) { + if (!replica.get_descriptor().is_memory_replica()) { + continue; + } + auto temp_segment_names = replica.get_segment_names(); + if (temp_segment_names.empty()) { + continue; + } + if (!temp_segment_names[0].has_value()) { + continue; + } + auto buffer_descriptor = replica.get_descriptor() + .get_memory_descriptor() + .buffer_descriptor; + MasterMetricManager::instance().inc_allocated_mem_size( + temp_segment_names[0].value(), + static_cast(buffer_descriptor.size_)); + } + ++it; } } - - return tl::make_unexpected(SerializationError( - ErrorCode::PERSISTENT_FAIL, - fmt::format("latest update {} failed", latest_path))); } - SNAP_LOG_INFO( - "[Snapshot] Upload latest success: {}, snapshot_id={}, " - "content={}", - latest_path, snapshot_id, latest_content); - - CleanupOldSnapshot(snapshot_retention_count_, snapshot_id); - SNAP_LOG_INFO("[Snapshot] action=persisting_state end, snapshot_id={}", - snapshot_id); - } catch (const std::exception& e) { - SNAP_LOG_ERROR( - "[Snapshot] Exception during state persistent, snapshot_id={}, " - "error={}", - snapshot_id, e.what()); - return tl::make_unexpected(SerializationError( - ErrorCode::PERSISTENT_FAIL, - fmt::format("Exception during state persistent: {}", e.what()))); - } catch (...) { - SNAP_LOG_ERROR( - "[Snapshot] Unknown exception during state persistent, " - "snapshot_id={}", - snapshot_id); - return tl::make_unexpected( - SerializationError(ErrorCode::PERSISTENT_FAIL, - "Unknown exception during state persistent")); - } - return {}; -} -tl::expected MasterService::UploadSnapshotPayloadFile( - const std::vector& data, const std::string& path, - const std::string& local_filename, const std::string& snapshot_id) { - SNAP_LOG_INFO("[Snapshot] Uploading {} to: {}, snapshot_id={}", - local_filename, path, snapshot_id); - - std::string error_msg; - auto upload_result = snapshot_object_store_->UploadBuffer(path, data); - if (!upload_result) { - SNAP_LOG_ERROR( - "[Snapshot] {} upload failed, snapshot_id={}, file={}, error={}", - local_filename, snapshot_id, path, upload_result.error()); - - // Upload failed, save locally for manual recovery in exception - // scenarios - if (use_snapshot_backup_dir_) { - auto save_path = fs::path(snapshot_backup_dir_) / - SNAPSHOT_BACKUP_SAVE_DIR / local_filename; - auto save_result = FileUtil::SaveBinaryToFile(data, save_path); - if (!save_result) { - SNAP_LOG_ERROR( - "[Snapshot] save {} to disk failed, snapshot_id={}, " - "file={}", - local_filename, snapshot_id, save_path.string()); - } - } - - error_msg.append(local_filename) - .append(" upload ") - .append(path) - .append(" failed; "); - return tl::make_unexpected( - SerializationError(ErrorCode::PERSISTENT_FAIL, error_msg)); - } else { - SNAP_LOG_INFO("[Snapshot] Upload {} success: {}, snapshot_id={}", - local_filename, path, snapshot_id); + LOG(INFO) << "[Restore] Total allocated size after restore: " + << MasterMetricManager::instance().get_allocated_mem_size(); } - return {}; -} - -void MasterService::CleanupOldSnapshot(int keep_count, - const std::string& snapshot_id) { - auto* snapshot_catalog_store = GetSnapshotCatalogStore(); - if (!snapshot_catalog_store) { - SNAP_LOG_ERROR( - "[Snapshot] snapshot catalog store is not initialized, " - "snapshot_id={}", - snapshot_id); - return; - } - - // List() loads one descriptor per published snapshot. This remains cheap - // because CleanupOldSnapshot() itself enforces snapshot_retention_count_ - // and keeps the catalog single-digit in normal deployments. - auto list_result = snapshot_catalog_store->List(kUnlimitedSnapshotList); - if (!list_result) { - SNAP_LOG_ERROR("[Snapshot] error=list failed, snapshot_id={}, code={}", - snapshot_id, toString(list_result.error())); - return; - } - - const auto& snapshots = list_result.value(); - - if (static_cast(snapshots.size()) > keep_count) { - for (int i = keep_count; i < static_cast(snapshots.size()); i++) { - const std::string& old_state_dir = snapshots[i].snapshot_id; + // Rebuild total capacity metrics + { + MasterMetricManager::instance().reset_total_mem_capacity(); + for (auto& segment_name : segment_names) { + MasterMetricManager::instance().reset_segment_total_mem_capacity( + segment_name); + } - if (old_state_dir == snapshot_id) { - SNAP_LOG_WARN( - "[Snapshot] Skipping deletion of current snapshot " - "directory {}, " - "snapshot_id={}", - old_state_dir, snapshot_id); - continue; + ScopedSegmentAccess segment_access = + segment_manager_.getSegmentAccess(); + std::vector> unready_segments; + if (segment_access.GetUnreadySegments(unready_segments) == + ErrorCode::OK) { + for (const auto& [segment, client_id] : unready_segments) { + UnmountSegment(segment.id, client_id); } + } - auto delete_result = snapshot_catalog_store->Delete(old_state_dir); - if (delete_result != ErrorCode::OK) { - SNAP_LOG_ERROR( - "[Snapshot] Failed to delete old snapshot {}, " - "snapshot_id={}, code={}", - old_state_dir, snapshot_id, toString(delete_result)); - } else { - SNAP_LOG_INFO( - "[Snapshot] Successfully deleted old snapshot {}, " - "snapshot_id={}", - old_state_dir, snapshot_id); + std::vector> all_segments; + auto err = segment_access.GetAllSegments(all_segments); + + if (err == ErrorCode::OK) { + int64_t total_size = 0; + for (const auto& [segment, client_id] : all_segments) { + Ping(client_id); + total_size += static_cast(segment.size); + MasterMetricManager::instance().inc_total_mem_capacity( + segment.name, segment.size); } + LOG(INFO) << "[Restore] Total capacity size after restore: " + << total_size; + } else { + LOG(ERROR) << "[Restore] Failed to get all segments, error: " + << err; } } + + return {}; } -void MasterService::RestoreState() { - auto* snapshot_catalog_store = GetSnapshotCatalogStore(); - if (!snapshot_catalog_store) { - LOG(ERROR) << "[Restore] Snapshot catalog store is not initialized, " - "starting fresh"; - return; +MasterService::TenantQuotaEvictionResult +MasterService::EvictTenantMemoryForQuota(const TenantId& tenant_id, + uint64_t target_bytes) { + TenantQuotaEvictionResult total; + if (!enable_multi_tenants_ || target_bytes == 0) { + return total; } - LOG(INFO) << "[Restore] Backend info: " - << snapshot_object_store_->GetConnectionInfo(); - - std::vector restore_candidates; - std::unordered_set candidate_ids; - std::optional latest_snapshot_id; + const TenantId normalized_tenant(tenant_id); + auto now = std::chrono::system_clock::now(); + std::shared_lock shared_lock(snapshot_mutex_); - auto latest_result = snapshot_catalog_store->GetLatest(); - if (!latest_result) { - LOG(WARNING) << "[Restore] Failed to load latest snapshot marker: " - << toString(latest_result.error()) - << ", falling back to published snapshot listing"; - } else if (latest_result->has_value()) { - const auto& latest_snapshot = latest_result->value(); - latest_snapshot_id = latest_snapshot.snapshot_id; - restore_candidates.push_back(latest_snapshot); - candidate_ids.emplace(latest_snapshot.snapshot_id); - } - - // Snapshot ids use YYYYMMDD_HHMMSS_mmm, so lexicographic order matches - // creation order. List() may perform one descriptor read per published - // snapshot; retention cleanup keeps that set bounded in practice. - auto snapshots_result = - snapshot_catalog_store->List(kUnlimitedSnapshotList); - if (!snapshots_result) { - if (restore_candidates.empty()) { - LOG(ERROR) << "[Restore] Failed to list restorable snapshots: " - << toString(snapshots_result.error()) - << ", starting fresh"; - return; - } - LOG(WARNING) << "[Restore] Failed to list fallback snapshots: " - << toString(snapshots_result.error()) - << ", attempting latest marker only"; - } else { - for (const auto& snapshot : snapshots_result.value()) { - // Snapshot ids are timestamp-derived, so string comparison keeps - // only candidates at or before the latest marker chronologically. - if (latest_snapshot_id.has_value() && - snapshot.snapshot_id > latest_snapshot_id.value()) { - continue; + auto is_evictable_memory_replica = [](const Replica& replica) { + return replica.is_memory_replica() && replica.is_completed() && + replica.get_refcnt() == 0; + }; + auto can_evict_replicas = [&](const ObjectMetadata& metadata) { + return metadata.HasReplica(is_evictable_memory_replica); + }; + auto has_local_disk_replica = [](const ObjectMetadata& metadata) { + return metadata.HasReplica(&Replica::fn_is_local_disk_replica); + }; + auto evict_replicas = + [&, this](ObjectMetadata& metadata, + std::vector>& deferred_replicas) { + const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata); + auto replicas = PopReplicasWithCacheTotalAccounting( + metadata, is_evictable_memory_replica); + const uint64_t replica_count = replicas.size(); + if (!replicas.empty()) { + deferred_replicas.emplace_back(std::move(replicas)); } - if (!candidate_ids.emplace(snapshot.snapshot_id).second) { - continue; + const uint64_t after_charge = CompletedMemoryQuotaCharge(metadata); + if (before_charge > after_charge) { + ReleaseCommittedQuotaCharge(metadata, + before_charge - after_charge); } - restore_candidates.push_back(snapshot); - } - } + return metadata.size * replica_count; + }; + long offload_queued_this_call = 0; + long offload_deferred_count = 0; + long offload_cap_forced_count = 0; + long offload_push_failed_forced = 0; + const long offload_cap = + offload_on_evict_ + ? static_cast(offloading_queue_limit_ * offload_cap_ratio_) + : 0; - if (restore_candidates.empty()) { - LOG(ERROR) << "[Restore] No previous snapshot found, starting fresh"; - return; - } + auto try_evict_or_offload = + [&, this](const std::string& key, ObjectMetadata& metadata, + TenantState& tenant_state, + std::vector>& deferred_replicas) { + if (!offload_on_evict_) { + return evict_replicas(metadata, deferred_replicas); + } - const auto now = std::chrono::system_clock::now(); - for (const auto& snapshot : restore_candidates) { - ResetStateAfterFailedRestoreAttempt(); - if (TryRestoreStateFromSnapshot(snapshot, now)) { - return; - } - } + if (has_local_disk_replica(metadata)) { + return evict_replicas(metadata, deferred_replicas); + } - ResetStateAfterFailedRestoreAttempt(); - LOG(ERROR) << "[Restore] Failed to restore from all candidate snapshots " - << "(count=" << restore_candidates.size() << "), starting fresh"; -} + if (offload_force_evict_ && + offload_queued_this_call >= offload_cap) { + ++offload_cap_forced_count; + return evict_replicas(metadata, deferred_replicas); + } -bool MasterService::TryRestoreStateFromSnapshot( - const ha::SnapshotDescriptor& snapshot, - const std::chrono::system_clock::time_point& now) { - const std::string& state_id = snapshot.snapshot_id; - std::string path_prefix = snapshot.object_prefix; - if (path_prefix.empty()) { - path_prefix = - snapshot_catalog_store_->GetSnapshotRoot() + state_id + "/"; - } + bool queued = false; + metadata.VisitReplicas( + is_evictable_memory_replica, + [this, &key, &normalized_tenant, &tenant_state, &queued, + &now](Replica& replica) { + if (queued) { + return; + } + auto result = PushOffloadingQueue( + MakeObjectIdentity(key, normalized_tenant), replica); + if (result) { + replica.inc_refcnt(); + tenant_state.offloading_tasks.emplace( + key, OffloadingTask{replica.id(), now}); + queued = true; + } + }); - std::string manifest_path = snapshot.manifest_key; - if (manifest_path.empty()) { - manifest_path = path_prefix + SNAPSHOT_MANIFEST_FILE; - } + if (queued) { + ++offload_queued_this_call; + ++offload_deferred_count; + return evict_replicas(metadata, deferred_replicas); + } - auto fail_restore = [&](const std::string& message) { - LOG(WARNING) << "[Restore] Snapshot candidate " << state_id - << " is unusable: " << message; - ResetStateAfterFailedRestoreAttempt(); - return false; - }; + if (offload_force_evict_) { + ++offload_push_failed_forced; + return evict_replicas(metadata, deferred_replicas); + } + return uint64_t{0}; + }; - try { - std::string manifest_content; - auto manifest_result = snapshot_object_store_->DownloadString( - manifest_path, manifest_content); - if (!manifest_result) { - return fail_restore("failed to download manifest '" + - manifest_path + - "': " + manifest_result.error()); - } - - if (use_snapshot_backup_dir_) { - auto save_result = FileUtil::SaveStringToFile( - manifest_content, fs::path(snapshot_backup_dir_) / - SNAPSHOT_BACKUP_RESTORE_DIR / - SNAPSHOT_MANIFEST_FILE); - if (!save_result) { - LOG(ERROR) << "[Restore] Failed to save manifest to file: " - << save_result.error(); - } - } - - std::vector parts; - boost::split(parts, manifest_content, boost::is_any_of("|")); - if (parts.size() < 3) { - return fail_restore("invalid snapshot manifest format"); - } - - const std::string& protocol_type = parts[0]; - const std::string& version = parts[1]; - - LOG(INFO) << "[Restore] Trying snapshot: " << state_id - << " version: " << version << " protocol: " << protocol_type; - - if (protocol_type != SNAPSHOT_SERIALIZER_TYPE) { - return fail_restore("unsupported protocol type '" + protocol_type + - "', expected '" + SNAPSHOT_SERIALIZER_TYPE + - "'"); - } - if (version != SNAPSHOT_SERIALIZER_VERSION) { - return fail_restore("incompatible snapshot version '" + version + - "', expected '" + SNAPSHOT_SERIALIZER_VERSION + - "'"); - } - - std::string metadata_path = path_prefix + SNAPSHOT_METADATA_FILE; - std::vector metadata_content; - auto download_result = snapshot_object_store_->DownloadBuffer( - metadata_path, metadata_content); - if (!download_result) { - return fail_restore("failed to download metadata '" + - metadata_path + - "': " + download_result.error()); - } - - if (use_snapshot_backup_dir_) { - auto save_result = FileUtil::SaveBinaryToFile( - metadata_content, fs::path(snapshot_backup_dir_) / - SNAPSHOT_BACKUP_RESTORE_DIR / - SNAPSHOT_METADATA_FILE); - if (!save_result) { - LOG(ERROR) << "[Restore] Failed to save metadata to file: " - << save_result.error(); - } - } - LOG(INFO) << "[Restore] Download metadata file success"; - - std::string segments_path = path_prefix + SNAPSHOT_SEGMENTS_FILE; - std::vector segments_content; - download_result = snapshot_object_store_->DownloadBuffer( - segments_path, segments_content); - if (!download_result) { - return fail_restore("failed to download segments '" + - segments_path + - "': " + download_result.error()); - } - if (use_snapshot_backup_dir_) { - auto save_result = FileUtil::SaveBinaryToFile( - segments_content, fs::path(snapshot_backup_dir_) / - SNAPSHOT_BACKUP_RESTORE_DIR / - SNAPSHOT_SEGMENTS_FILE); - if (!save_result) { - LOG(ERROR) << "[Restore] Failed to save segments to file: " - << save_result.error(); - } - } - LOG(INFO) << "[Restore] Download segments file success"; - - std::string task_manager_path = - path_prefix + SNAPSHOT_TASK_MANAGER_FILE; - std::vector task_manager_content; - download_result = snapshot_object_store_->DownloadBuffer( - task_manager_path, task_manager_content); - if (!download_result) { - return fail_restore("failed to download task_manager '" + - task_manager_path + - "': " + download_result.error()); - } - if (use_snapshot_backup_dir_) { - auto save_result = FileUtil::SaveBinaryToFile( - task_manager_content, fs::path(snapshot_backup_dir_) / - SNAPSHOT_BACKUP_RESTORE_DIR / - SNAPSHOT_TASK_MANAGER_FILE); - if (!save_result) { - LOG(ERROR) << "[Restore] Failed to save task manager to file: " - << save_result.error(); - } - } - LOG(INFO) << "[Restore] Download task manager file success"; - - SegmentSerializer segment_serializer(&segment_manager_); - MetadataSerializer metadata_serializer(this); - TaskManagerSerializer task_manager_serializer(&task_manager_); - - auto segments_result = segment_serializer.Deserialize(segments_content); - if (!segments_result) { - return fail_restore( - fmt::format("failed to deserialize segments: {} - {}", - static_cast(segments_result.error().code), - segments_result.error().message)); - } - LOG(INFO) << "[Restore] Deserialize segments success"; - - auto metadata_result = - metadata_serializer.Deserialize(metadata_content); - if (!metadata_result) { - return fail_restore( - fmt::format("failed to deserialize metadata: {} - {}", - static_cast(metadata_result.error().code), - metadata_result.error().message)); + auto try_evict_group_or_object = + [&, this](const std::string& key, ObjectMetadata& metadata, + TenantState& tenant_state, + std::vector>& deferred_replicas, + bool allow_soft_pinned) -> TenantQuotaEvictionResult { + if (!metadata.IsGrouped()) { + uint64_t freed = try_evict_or_offload(key, metadata, tenant_state, + deferred_replicas); + return {.freed_bytes = freed, + .evicted_objects = freed > 0 ? 1U : 0U}; } - LOG(INFO) << "[Restore] Deserialize metadata success"; - auto task_manager_result = - task_manager_serializer.Deserialize(task_manager_content); - if (!task_manager_result) { - return fail_restore( - fmt::format("failed to deserialize task manager: {} - {}", - static_cast(task_manager_result.error().code), - task_manager_result.error().message)); + auto group_it = tenant_state.group_members.find(metadata.group_id); + if (group_it == tenant_state.group_members.end()) { + uint64_t freed = try_evict_or_offload(key, metadata, tenant_state, + deferred_replicas); + return {.freed_bytes = freed, + .evicted_objects = freed > 0 ? 1U : 0U}; } - LOG(INFO) << "[Restore] Deserialize task manager success"; - std::vector segment_names; - { - ScopedSegmentAccess segment_access = - segment_manager_.getSegmentAccess(); - segment_access.GetAllSegmentNames(segment_names); + for (const auto& member_key : group_it->second) { + auto member_it = tenant_state.metadata.find(member_key); + if (member_it != tenant_state.metadata.end() && + !member_it->second.IsLeaseExpired(now)) { + return {}; + } } - { - const bool skip_cleanup = std::getenv( - "MOONCAKE_MASTER_SERVICE_SNAPSHOT_TEST_SKIP_CLEANUP"); - if (!skip_cleanup) { - auto cleanup_now = now; - for (auto& shard : metadata_shards_) { - for (auto tenant_it = shard.tenants.begin(); - tenant_it != shard.tenants.end();) { - auto& tenant_state = tenant_it->second; - for (auto it = tenant_state.metadata.begin(); - it != tenant_state.metadata.end();) { - if (it->second.HasDiffRepStatus( - ReplicaStatus::COMPLETE) || - (it->second.IsLeaseExpired(cleanup_now) && - !it->second.IsSoftPinned(cleanup_now))) { - VLOG(1) << "clear metadata key=" << it->first; - it = EraseMetadata(tenant_state, it, - tenant_it->first); - } else { - ++it; - } - } - if (tenant_state.Empty()) { - tenant_it = shard.tenants.erase(tenant_it); - } else { - ++tenant_it; - } - } - } + TenantQuotaEvictionResult result; + std::vector member_keys(group_it->second.begin(), + group_it->second.end()); + for (const auto& member_key : member_keys) { + auto member_it = tenant_state.metadata.find(member_key); + if (member_it == tenant_state.metadata.end()) { + continue; } - - MasterMetricManager::instance().reset_allocated_mem_size(); - RebuildCacheTotalAccounting(); - for (auto& segment_name : segment_names) { - MasterMetricManager::instance() - .reset_segment_allocated_mem_size(segment_name); + auto& member_metadata = member_it->second; + if (member_metadata.IsHardPinned() || + !member_metadata.IsLeaseExpired(now) || + (!allow_soft_pinned && member_metadata.IsSoftPinned(now)) || + !can_evict_replicas(member_metadata)) { + continue; } - for (auto& shard : metadata_shards_) { - for (auto& [tenant_id, tenant_state] : shard.tenants) { - for (auto it = tenant_state.metadata.begin(); - it != tenant_state.metadata.end();) { - for (auto& replica : it->second.GetAllReplicas()) { - if (!replica.get_descriptor().is_memory_replica()) { - continue; - } - auto temp_segment_names = - replica.get_segment_names(); - if (temp_segment_names.empty()) { - continue; - } - if (!temp_segment_names[0].has_value()) { - continue; - } - auto buffer_descriptor = - replica.get_descriptor() - .get_memory_descriptor() - .buffer_descriptor; - MasterMetricManager::instance() - .inc_allocated_mem_size( - temp_segment_names[0].value(), - static_cast( - buffer_descriptor.size_)); - } - ++it; - } - } + const uint64_t freed = try_evict_or_offload( + member_key, member_metadata, tenant_state, deferred_replicas); + result.freed_bytes += freed; + if (freed > 0) { + ++result.evicted_objects; } - - LOG(INFO) - << "[Restore] Total allocated size after restore: " - << MasterMetricManager::instance().get_allocated_mem_size(); - } - - { - MasterMetricManager::instance().reset_total_mem_capacity(); - for (auto& segment_name : segment_names) { - MasterMetricManager::instance() - .reset_segment_total_mem_capacity(segment_name); + if (member_key != key && !member_metadata.IsValid()) { + EraseMetadata(tenant_state, member_it, normalized_tenant); } + } + return result; + }; - ScopedSegmentAccess segment_access = - segment_manager_.getSegmentAccess(); - std::vector> unready_segments; - if (segment_access.GetUnreadySegments(unready_segments) == - ErrorCode::OK) { - for (const auto& [segment, client_id] : unready_segments) { - UnmountSegment(segment.id, client_id); + auto pass = [&](bool allow_soft_pinned) { + const size_t start_shard = randomIndex(kNumShards); + for (size_t scanned = 0; + scanned < kNumShards && total.freed_bytes < target_bytes; + ++scanned) { + const size_t shard_idx = (start_shard + scanned) % kNumShards; + std::vector> deferred_replicas; + { + MetadataShardAccessorRW shard(this, shard_idx); + auto tenant_it = shard->tenants.find(normalized_tenant); + if (tenant_it == shard->tenants.end()) { + continue; } - } - - std::vector> all_segments; - auto err = segment_access.GetAllSegments(all_segments); + auto& tenant_state = tenant_it->second; + for (auto it = tenant_state.metadata.begin(); + it != tenant_state.metadata.end() && + total.freed_bytes < target_bytes;) { + auto& metadata = it->second; + if (metadata.IsHardPinned() || + !metadata.IsLeaseExpired(now) || + (!allow_soft_pinned && metadata.IsSoftPinned(now)) || + !can_evict_replicas(metadata)) { + ++it; + continue; + } - if (err == ErrorCode::OK) { - int64_t total_size = 0; - for (const auto& [segment, client_id] : all_segments) { - Ping(client_id); - total_size += static_cast(segment.size); - MasterMetricManager::instance().inc_total_mem_capacity( - segment.name, segment.size); + auto evict_result = try_evict_group_or_object( + it->first, metadata, tenant_state, deferred_replicas, + allow_soft_pinned); + total.freed_bytes += evict_result.freed_bytes; + total.evicted_objects += evict_result.evicted_objects; + if (!metadata.IsValid()) { + it = EraseMetadata(tenant_state, it, normalized_tenant); + } else { + ++it; + } + } + if (tenant_state.Empty()) { + shard->tenants.erase(tenant_it); } - LOG(INFO) << "[Restore] Total capacity size after restore: " - << total_size; - } else { - LOG(ERROR) << "[Restore] Failed to get all segments, error: " - << err; } } + }; - LOG(INFO) << "[Restore] Successfully restored state from snapshot: " - << state_id; - RebuildTenantQuotaUsageFromMetadata(); - return true; - } catch (const std::exception& e) { - return fail_restore("exception during state restoration: " + - std::string(e.what())); - } catch (...) { - return fail_restore("unknown exception during state restoration"); + pass(/*allow_soft_pinned=*/false); + if (allow_evict_soft_pinned_objects_ && total.freed_bytes < target_bytes) { + pass(/*allow_soft_pinned=*/true); } -} -void MasterService::ResetStateAfterFailedRestoreAttempt() { - SegmentSerializer segment_serializer(&segment_manager_); - MetadataSerializer metadata_serializer(this); - TaskManagerSerializer task_manager_serializer(&task_manager_); - - task_manager_serializer.Reset(); - metadata_serializer.Reset(); - segment_serializer.Reset(); - - { - std::unique_lock lock(client_mutex_); - ok_client_.clear(); + if (total.freed_bytes > 0) { + MasterMetricManager::instance().inc_tenant_evict_bytes( + normalized_tenant.value(), + static_cast(std::min( + total.freed_bytes, + static_cast(std::numeric_limits::max())))); } - PodUUID pod_uuid; - while (client_ping_queue_.pop(pod_uuid)) { + if (offload_on_evict_ && total.freed_bytes == 0 && + offload_deferred_count > 0) { + LOG(WARNING) << "[TENANT-EVICT] No memory freed for tenant " + << normalized_tenant << "; " << offload_deferred_count + << " object(s) deferred for disk offload."; } - - MasterMetricManager::instance().reset_allocated_mem_size(); - MasterMetricManager::instance().reset_total_mem_capacity(); - MasterMetricManager::instance().reset_cache_total_nums(); -} - -ha::SnapshotCatalogStore* MasterService::GetSnapshotCatalogStore() { - return snapshot_catalog_store_.get(); + if (offload_cap_forced_count > 0) { + LOG(WARNING) << "[TENANT-EVICT] Offload cap (" << offload_cap + << ") reached for tenant " << normalized_tenant + << "; force-evicted " << offload_cap_forced_count + << " object(s) without disk offload."; + } + if (offload_push_failed_forced > 0) { + LOG(WARNING) << "[TENANT-EVICT] PushOffloadingQueue failed for tenant " + << normalized_tenant << " on " + << offload_push_failed_forced + << " object(s); force-evicted without disk offload " + "(offload_force_evict=true)."; + } + return total; } void MasterService::BatchEvict(double evict_ratio_target, @@ -5868,13 +8084,6 @@ void MasterService::BatchEvict(double evict_ratio_target, } auto now = std::chrono::system_clock::now(); - long evicted_count = 0; - long object_count = 0; - uint64_t total_freed_size = 0; - - // Candidates for second pass eviction - std::vector no_pin_objects; - std::vector soft_pin_objects; auto is_evictable_memory_replica = [](const Replica& replica) { return replica.is_memory_replica() && replica.is_completed() && @@ -5888,6 +8097,13 @@ void MasterService::BatchEvict(double evict_ratio_target, auto evict_replicas = [&, this](ObjectMetadata& metadata, std::vector>& deferred_replicas) { + if (enable_oplog_) { + return metadata.size * + metadata.CountReplicas([](const Replica& replica) { + return replica.is_memory_replica() && + replica.status() == ReplicaStatus::REMOVED; + }); + } const uint64_t before_charge = CompletedMemoryQuotaCharge(metadata); auto replicas = PopReplicasWithCacheTotalAccounting( metadata, is_evictable_memory_replica); @@ -5910,7 +8126,7 @@ void MasterService::BatchEvict(double evict_ratio_target, long offload_push_failed_forced = 0; // #keys force-evicted on push fail const long offload_cap = offload_on_evict_ - ? static_cast(offloading_queue_limit_ * kOffloadCapRatio) + ? static_cast(offloading_queue_limit_ * offload_cap_ratio_) : 0; auto has_local_disk_replica = [](const ObjectMetadata& metadata) { @@ -5921,9 +8137,12 @@ void MasterService::BatchEvict(double evict_ratio_target, // replicas were evicted (all MEMORY replicas of the key are now pinned). auto try_evict_or_offload = [&, this]( - const std::string& tenant_id, const std::string& key, + const TenantId& tenant_id, const std::string& key, ObjectMetadata& metadata, TenantState& tenant_state, std::vector>& deferred_replicas) -> uint64_t { + if (enable_oplog_) { + return evict_replicas(metadata, deferred_replicas); + } if (!offload_on_evict_) { // Original behavior return evict_replicas(metadata, deferred_replicas); @@ -5982,18 +8201,111 @@ void MasterService::BatchEvict(double evict_ratio_target, return 0; }; + // HA strong-consistency: persist the post-eviction state BEFORE the + // helper mutates `metadata`. Returns true on success (or when HA is + // disabled). On false, the caller must NOT call try_evict_or_offload + // and must NOT erase the metadata entry — local state must stay in + // sync with what was published. + auto persist_evict_oplog_or_skip = + [&, this](const TenantId& tenant_id, const std::string& key, + ObjectMetadata& metadata) -> bool { + if (!enable_oplog_ || !ordered_oplog_writer_) { + return true; + } + + // Predict the descriptor list after evict_replicas() runs: + // drop COMPLETE memory replicas with refcnt==0; keep everything else + // that is COMPLETE. + auto remaining = + BuildRemainingReplicaDescriptors(metadata, [](const Replica& r) { + return r.is_memory_replica() && r.is_completed() && + r.get_refcnt() == 0; + }); + + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + LOG(WARNING) + << "BatchEvict: OpLog reservation failed for key=" << key + << ", err=" << static_cast(reservation.error()) + << ", skipping eviction"; + return false; + } + std::vector removed_ids; + metadata.VisitReplicas(is_evictable_memory_replica, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + tl::expected persist_result; + if (remaining.empty()) { + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::REMOVE, + tenant_id.value(), key, {}, + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, QuotaEraseMode::kFull); + }); + } else { + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::PUT_END, + tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, remaining, + metadata.group_id, metadata.data_type), + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, QuotaEraseMode::kFull); + }); + } + if (!persist_result) { + LOG(WARNING) + << "BatchEvict: OpLog persist failed for key=" << key + << ", err=" << static_cast(persist_result.error()) + << ", skipping eviction"; + return false; + } + return true; + } + + tl::expected persist_result; + if (remaining.empty()) { + persist_result = AppendOpLogWithDurableFinalize( + OpType::REMOVE, tenant_id.value(), key, {}, nullptr); + } else { + persist_result = AppendOpLogWithDurableFinalize( + OpType::PUT_END, tenant_id.value(), key, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, remaining, + metadata.group_id, metadata.data_type), + nullptr); + } + if (!persist_result) { + LOG(WARNING) << "BatchEvict: OpLog persist failed for key=" << key + << ", err=" << static_cast(persist_result.error()) + << ", skipping eviction"; + return false; + } + return true; + }; + struct EvictionResult { uint64_t freed_bytes{0}; long evicted_objects{0}; }; auto try_evict_group_or_object = - [&, this](const std::string& tenant_id, const std::string& key, + [&, this](const TenantId& tenant_id, const std::string& key, ObjectMetadata& metadata, MetadataShardAccessorRW& shard, TenantState& tenant_state, std::vector>& deferred_replicas, bool allow_soft_pinned) -> EvictionResult { if (!metadata.IsGrouped()) { + if (!persist_evict_oplog_or_skip(tenant_id, key, metadata)) { + return {}; + } uint64_t freed = try_evict_or_offload( tenant_id, key, metadata, tenant_state, deferred_replicas); return {.freed_bytes = freed, .evicted_objects = freed > 0 ? 1 : 0}; @@ -6001,6 +8313,9 @@ void MasterService::BatchEvict(double evict_ratio_target, auto group_it = tenant_state.group_members.find(metadata.group_id); if (group_it == tenant_state.group_members.end()) { + if (!persist_evict_oplog_or_skip(tenant_id, key, metadata)) { + return {}; + } uint64_t freed = try_evict_or_offload( tenant_id, key, metadata, tenant_state, deferred_replicas); return {.freed_bytes = freed, .evicted_objects = freed > 0 ? 1 : 0}; @@ -6030,119 +8345,354 @@ void MasterService::BatchEvict(double evict_ratio_target, continue; } + if (!persist_evict_oplog_or_skip(tenant_id, member_key, + member_metadata)) { + continue; + } uint64_t freed = try_evict_or_offload(tenant_id, member_key, member_metadata, tenant_state, deferred_replicas); result.freed_bytes += freed; if (freed > 0) { result.evicted_objects++; + if (!enable_oplog_) { + PublishKvRemovedAfterEvict(member_key, freed, "cpu", + member_metadata, tenant_id); + } } - if (member_key != key && !member_metadata.IsValid()) { - EraseMetadata(tenant_state, member_it, tenant_id); + if (member_key != key && !enable_oplog_ && + !member_metadata.IsValid()) { + EraseMetadata(tenant_state, member_it, tenant_id, + QuotaEraseMode::kFull, &shard); + } + } + return result; + }; + + // Candidate carries key for safe lookup after releasing shard lock. + // Iterators would be invalid if the shard is modified between phases. + struct Candidate { + size_t shard_idx; + TenantId tenant_id; + std::string key; + std::chrono::system_clock::time_point lease_timeout; + }; + + // Randomly select a starting shard to avoid imbalance eviction between + // shards. + size_t start_idx = randomIndex(kNumShards); + std::shared_lock shared_lock(snapshot_mutex_); + + // ===== Phase 1: Parallel candidate census ===== + // N threads each scan a batch of shards. For selective ratios only the + // lease timestamps are collected here; full tenant/key identities are + // materialized afterwards for a bounded frontier around the eviction + // cutoff. High ratios collect full Candidates directly, because a census + // followed by a second scan would cost more than the identities it saves. + int num_threads = std::min((int)kNumShards, 16); + size_t shards_per_thread = (kNumShards + num_threads - 1) / num_threads; + + constexpr size_t kMinReserveSlack = 1024; + constexpr size_t kMinFrontierLimit = 64 * 1024; + constexpr size_t kReserveSlackDivisor = 10; + constexpr size_t kFrontierDivisor = 4; + // Above this target ratio the reserve frontier would already cover a + // large share of the population, so selective materialization stops + // paying for the extra scan it costs. + constexpr double kCompactPrebypassTargetRatio = + static_cast(kReserveSlackDivisor) / + static_cast(kFrontierDivisor * (kReserveSlackDivisor + 1)); + const bool compact_frontier_prebypass = + evict_ratio_target >= kCompactPrebypassTargetRatio; + + std::vector> local_candidates(num_threads); + std::vector> + local_no_pin(num_threads); + std::vector local_eviction_base(num_threads, 0); + std::vector local_object_count(num_threads, 0); + std::vector> + local_soft_pin(num_threads); + + std::vector threads; + for (int t = 0; t < num_threads; t++) { + threads.emplace_back([&, t] { + size_t s_start = t * shards_per_thread; + size_t s_end = std::min(s_start + shards_per_thread, kNumShards); + for (size_t s = s_start; s < s_end; s++) { + MetadataShardAccessorRW shard(this, s); + DiscardExpiredProcessingReplicas(shard, now); + + size_t shard_metadata_count = 0; + size_t shard_evictable_count = 0; + for (const auto& [tenant_id, tenant_state] : shard->tenants) { + shard_metadata_count += tenant_state.metadata.size(); + for (auto it = tenant_state.metadata.begin(); + it != tenant_state.metadata.end(); ++it) { + if (it->second.IsHardPinned()) continue; + bool has_evictable = can_evict_replicas(it->second); + if (has_evictable) shard_evictable_count++; + if (!it->second.IsLeaseExpired(now) || !has_evictable) + continue; + if (!it->second.IsSoftPinned(now)) { + if (compact_frontier_prebypass) { + local_candidates[t].push_back( + {s, tenant_id, it->first, + it->second.lease_timeout}); + } else { + local_no_pin[t].push_back( + it->second.lease_timeout); + } + } else if (allow_evict_soft_pinned_objects_) { + local_soft_pin[t].push_back( + it->second.lease_timeout); + } + } + } + local_object_count[t] += shard_metadata_count; + local_eviction_base[t] += shard_evictable_count; } + }); + } + for (auto& t : threads) t.join(); + + // Merge per-thread results + long total_eviction_base = 0; + for (auto v : local_eviction_base) total_eviction_base += v; + + long object_count = 0; + for (auto v : local_object_count) object_count += v; + + std::vector candidates; + if (compact_frontier_prebypass) { + size_t total = 0; + for (auto& v : local_candidates) total += v.size(); + candidates.reserve(total); + for (auto& v : local_candidates) { + candidates.insert(candidates.end(), + std::make_move_iterator(v.begin()), + std::make_move_iterator(v.end())); } - return result; - }; + } - // Randomly select a starting shard to avoid imbalance eviction between - // shards. - size_t start_idx = RandomIndex(kNumShards); - std::shared_lock shared_lock(snapshot_mutex_); + std::vector no_pin_timeouts; + { + size_t total = 0; + for (auto& v : local_no_pin) total += v.size(); + no_pin_timeouts.reserve(total); + } + for (auto& v : local_no_pin) { + no_pin_timeouts.insert(no_pin_timeouts.end(), + std::make_move_iterator(v.begin()), + std::make_move_iterator(v.end())); + } - // First pass: evict objects without soft pin and lease expired - std::vector> deferred_replicas; - for (size_t i = 0; i < kNumShards; i++) { - { - MetadataShardAccessorRW shard(this, (start_idx + i) % kNumShards); + std::vector soft_pin_objects; + { + size_t total = 0; + for (auto& v : local_soft_pin) total += v.size(); + soft_pin_objects.reserve(total); + } + for (auto& v : local_soft_pin) { + soft_pin_objects.insert(soft_pin_objects.end(), + std::make_move_iterator(v.begin()), + std::make_move_iterator(v.end())); + } - // Discard expired processing keys first so that they won't be - // counted in later evictions. - DiscardExpiredProcessingReplicas(shard, now); + if (total_eviction_base == 0) { + need_mem_eviction_ = false; + VLOG(1) << "[EVICT-DIAG] object_count=" << object_count + << " eviction_base=0 (no evictable memory objects)"; + return; + } - size_t shard_object_count = 0; - for (const auto& [tenant_id, tenant_state] : shard->tenants) { - shard_object_count += tenant_state.metadata.size(); + const long ideal_evict_num = + std::ceil(total_eviction_base * evict_ratio_target); + const size_t no_pin_count = + compact_frontier_prebypass ? candidates.size() : no_pin_timeouts.size(); + const long primary_no_pin_num = + std::min(ideal_evict_num, static_cast(no_pin_count)); + + // Re-scan metadata and copy full identities only for objects inside the + // requested timestamp range. The eligibility conditions are identical to + // the census above, so the selected set matches what the census counted. + auto collect_candidates = [&](bool use_cutoff, + std::chrono::system_clock::time_point cutoff, + bool collect_older_or_equal) { + std::vector> local_frontier(num_threads); + std::vector collectors; + collectors.reserve(num_threads); + + for (int t = 0; t < num_threads; t++) { + collectors.emplace_back([&, t] { + size_t s_start = t * shards_per_thread; + size_t s_end = + std::min(s_start + shards_per_thread, kNumShards); + for (size_t s = s_start; s < s_end; s++) { + MetadataShardAccessorRW shard(this, s); + for (const auto& [tenant_id, tenant_state] : + shard->tenants) { + for (const auto& [key, metadata] : + tenant_state.metadata) { + if (metadata.IsHardPinned() || + !metadata.IsLeaseExpired(now) || + metadata.IsSoftPinned(now) || + !can_evict_replicas(metadata)) { + continue; + } + if (use_cutoff) { + const bool in_range = + collect_older_or_equal + ? metadata.lease_timeout <= cutoff + : metadata.lease_timeout > cutoff; + if (!in_range) continue; + } + local_frontier[t].push_back( + {s, tenant_id, key, metadata.lease_timeout}); + } + } + } + }); + } + for (auto& collector : collectors) collector.join(); + + size_t total = 0; + for (const auto& v : local_frontier) total += v.size(); + std::vector merged; + merged.reserve(total); + for (auto& v : local_frontier) { + merged.insert(merged.end(), std::make_move_iterator(v.begin()), + std::make_move_iterator(v.end())); + } + return merged; + }; + + bool compact_frontier_used = false; + std::chrono::system_clock::time_point reserve_cutoff{}; + + if (primary_no_pin_num > 0 && !compact_frontier_prebypass) { + const size_t primary_count = static_cast(primary_no_pin_num); + // The reserve absorbs objects that stop being evictable between the + // census and the eviction pass, so ordinary churn does not require a + // second materialization pass. + const size_t reserve_slack = std::max( + kMinReserveSlack, + (primary_count + kReserveSlackDivisor - 1) / kReserveSlackDivisor); + const size_t reserve_count = + std::min(no_pin_count, primary_count + reserve_slack); + const size_t frontier_limit = + std::max(kMinFrontierLimit, + (no_pin_count + kFrontierDivisor - 1) / kFrontierDivisor); + + if (reserve_count <= frontier_limit) { + std::nth_element(no_pin_timeouts.begin(), + no_pin_timeouts.begin() + (reserve_count - 1), + no_pin_timeouts.end()); + reserve_cutoff = no_pin_timeouts[reserve_count - 1]; + candidates = collect_candidates(/*use_cutoff=*/true, reserve_cutoff, + /*collect_older_or_equal=*/true); + + // Shortfall guard: if churn left the frontier holding fewer + // objects than the target needs, fall back to the full candidate + // set so evict_num below still derives from the requested target + // rather than from a shrunken frontier. + if (candidates.size() >= primary_count) { + compact_frontier_used = true; + } else { + candidates = + collect_candidates(/*use_cutoff=*/false, {}, + /*collect_older_or_equal=*/true); } - object_count += shard_object_count; + } else { + candidates = collect_candidates(/*use_cutoff=*/false, {}, + /*collect_older_or_equal=*/true); + } + } + // ===== Phase 2: Serial eviction via key lookup ===== + long evicted_count = 0; + uint64_t total_freed_size = 0; + std::vector no_pin_objects; + std::vector> deferred_replicas; - // To achieve evicted_count / object_count = evict_ratio_target, - // ideally how many object should be evicted in this shard - const long ideal_evict_num = - std::ceil(object_count * evict_ratio_target) - evicted_count; + // First pass: evict candidates with no soft pin + if (!candidates.empty()) { + long evict_num = std::min(ideal_evict_num, (long)candidates.size()); + + std::nth_element(candidates.begin(), + candidates.begin() + (evict_num - 1), candidates.end(), + [](const Candidate& a, const Candidate& b) { + return a.lease_timeout < b.lease_timeout; + }); + auto target_timeout = candidates[evict_num - 1].lease_timeout; + + // Treat evict_num as a minimum: if re-validation skips a candidate, + // continue trying the next one so actual evicted count reaches + // evict_num. This matches the old per-shard over-eviction behavior. + long evicted_this_pass = 0; + auto evict_candidate_batch = [&](std::vector& batch) { + for (auto& c : batch) { + if (evicted_this_pass >= evict_num && + c.lease_timeout > target_timeout) { + no_pin_objects.push_back(c.lease_timeout); + continue; + } + { + MetadataShardAccessorRW shard(this, c.shard_idx); + auto tenant_it = shard->tenants.find(c.tenant_id); + if (tenant_it == shard->tenants.end()) continue; + auto& tenant_state = tenant_it->second; + auto it = tenant_state.metadata.find(c.key); + if (it == tenant_state.metadata.end()) continue; - std::vector - candidates; // can be removed - for (const auto& [tenant_id, tenant_state] : shard->tenants) { - for (auto it = tenant_state.metadata.begin(); - it != tenant_state.metadata.end(); it++) { - if (it->second.IsHardPinned()) { - continue; - } + // Re-validate: state may have changed since Phase 1 if (!it->second.IsLeaseExpired(now) || + it->second.IsSoftPinned(now) || !can_evict_replicas(it->second)) { + no_pin_objects.push_back(c.lease_timeout); continue; } - if (!it->second.IsSoftPinned(now)) { - if (ideal_evict_num > 0) { - candidates.push_back(it->second.lease_timeout); - } else { - no_pin_objects.push_back(it->second.lease_timeout); - } - } else if (allow_evict_soft_pinned_objects_) { - soft_pin_objects.push_back(it->second.lease_timeout); + + auto evict_result = try_evict_group_or_object( + c.tenant_id, c.key, it->second, shard, tenant_state, + deferred_replicas, + /*allow_soft_pinned=*/false); + + total_freed_size += evict_result.freed_bytes; + + if (!enable_oplog_ && !it->second.IsGrouped()) { + PublishKvRemovedAfterEvict( + c.key, evict_result.freed_bytes, "cpu", it->second, + c.tenant_id); } - } - } - if (ideal_evict_num > 0 && !candidates.empty()) { - long evict_num = - std::min(ideal_evict_num, (long)candidates.size()); - long shard_evicted_count = - 0; // number of objects evicted from this shard - std::nth_element(candidates.begin(), - candidates.begin() + (evict_num - 1), - candidates.end()); - auto target_timeout = candidates[evict_num - 1]; - for (auto tenant_it = shard->tenants.begin(); - tenant_it != shard->tenants.end();) { - auto& tenant_state = tenant_it->second; - auto it = tenant_state.metadata.begin(); - while (it != tenant_state.metadata.end()) { - if (it->second.IsHardPinned() || - !it->second.IsLeaseExpired(now) || - it->second.IsSoftPinned(now) || - !can_evict_replicas(it->second)) { - ++it; - continue; - } - if (it->second.lease_timeout <= target_timeout) { - auto evict_result = try_evict_group_or_object( - tenant_it->first, it->first, it->second, shard, - tenant_state, deferred_replicas, - /*allow_soft_pinned=*/false); - total_freed_size += evict_result.freed_bytes; - if (it->second.IsValid() == false) { - it = EraseMetadata(tenant_state, it, - tenant_it->first); - } else { - ++it; - } - shard_evicted_count += evict_result.evicted_objects; - } else { - no_pin_objects.push_back(it->second.lease_timeout); - ++it; - } + if (!enable_oplog_ && !it->second.IsValid()) { + EraseMetadata(tenant_state, it, c.tenant_id, + QuotaEraseMode::kFull, &shard); } + if (tenant_state.Empty()) { - tenant_it = shard->tenants.erase(tenant_it); - } else { - ++tenant_it; + shard->tenants.erase(tenant_it); } + + evicted_count += evict_result.evicted_objects; + evicted_this_pass += evict_result.evicted_objects; } - evicted_count += shard_evicted_count; + deferred_replicas.clear(); } + }; + + evict_candidate_batch(candidates); + + // Metadata may change after the frontier is materialized. If the + // reserve is exhausted before the target is met, refill from the + // remainder of the current no-soft-pin population. This recovery + // scan is paid only on churn and preserves the behavior of + // continuing past the cutoff until evict_num is reached. + if (compact_frontier_used && evicted_this_pass < evict_num) { + auto refill_candidates = collect_candidates( + /*use_cutoff=*/true, reserve_cutoff, + /*collect_older_or_equal=*/false); + evict_candidate_batch(refill_candidates); } - deferred_replicas.clear(); } // Try releasing discarded replicas before we decide whether to do the @@ -6150,8 +8700,9 @@ void MasterService::BatchEvict(double evict_ratio_target, uint64_t released_discarded_cnt = ReleaseExpiredDiscardedReplicas(now); // The ideal number of objects to evict in the second pass - long target_evict_num = std::ceil(object_count * evict_ratio_lowerbound) - - evicted_count - released_discarded_cnt; + long target_evict_num = + std::ceil(total_eviction_base * evict_ratio_lowerbound) - + evicted_count - released_discarded_cnt; // The actual number of objects we can evict in the second pass target_evict_num = std::min(target_evict_num, @@ -6161,21 +8712,14 @@ void MasterService::BatchEvict(double evict_ratio_target, // evicted AND 2). The evicted number in the first pass is less than // evict_ratio_lowerbound. if (target_evict_num > 0) { - // If 1). there are enough candidates without soft pin OR 2). soft pin - // candidates are empty, then do second pass A. Otherwise, do second - // pass B. Note that the second condition is ensured implicitly by the - // calculation of target_evict_num. if (target_evict_num <= static_cast(no_pin_objects.size())) { - // Second pass A: only evict objects without soft pin. The following - // code is error-prone if target_evict_num > no_pin_objects.size(). - + // Second pass A: only evict objects without soft pin. std::nth_element(no_pin_objects.begin(), no_pin_objects.begin() + (target_evict_num - 1), no_pin_objects.end()); auto target_timeout = no_pin_objects[target_evict_num - 1]; - // Evict objects with lease timeout less than or equal to target. - // Stop when the target is reached. + // Evict via key lookup — avoid full metadata traversal for (size_t i = 0; i < kNumShards && target_evict_num > 0; i++) { { MetadataShardAccessorRW shard(this, @@ -6197,9 +8741,15 @@ void MasterService::BatchEvict(double evict_ratio_target, shard, tenant_state, deferred_replicas, /*allow_soft_pinned=*/false); total_freed_size += evict_result.freed_bytes; - if (!it->second.IsValid()) { - it = EraseMetadata(tenant_state, it, - tenant_it->first); + if (!enable_oplog_ && !it->second.IsGrouped()) { + PublishKvRemovedAfterEvict( + it->first, evict_result.freed_bytes, + "cpu", it->second, tenant_it->first); + } + if (!enable_oplog_ && !it->second.IsValid()) { + it = EraseMetadata( + tenant_state, it, tenant_it->first, + QuotaEraseMode::kFull, &shard); } else { ++it; } @@ -6220,23 +8770,16 @@ void MasterService::BatchEvict(double evict_ratio_target, deferred_replicas.clear(); } } else if (!soft_pin_objects.empty()) { - // allow_evict_soft_pinned_objects_ is implicitly true if - // soft_pin_objects is not empty Second pass B: Prioritize evicting - // objects without soft pin, but also allow to evict soft pinned - // objects. The following code is error-prone if the soft pin - // objects are empty. - + // Second pass B: Prioritize evicting objects without soft pin, + // but also allow evicting soft pinned objects. const long soft_pin_evict_num = target_evict_num - static_cast(no_pin_objects.size()); - // For soft pin objects, prioritize to evict the ones with smaller - // lease timeout. std::nth_element( soft_pin_objects.begin(), soft_pin_objects.begin() + (soft_pin_evict_num - 1), soft_pin_objects.end()); auto soft_target_timeout = soft_pin_objects[soft_pin_evict_num - 1]; - // Stop when the target is reached. for (size_t i = 0; i < kNumShards && target_evict_num > 0; i++) { { MetadataShardAccessorRW shard(this, @@ -6263,9 +8806,15 @@ void MasterService::BatchEvict(double evict_ratio_target, shard, tenant_state, deferred_replicas, /*allow_soft_pinned=*/true); total_freed_size += evict_result.freed_bytes; - if (!it->second.IsValid()) { - it = EraseMetadata(tenant_state, it, - tenant_it->first); + if (!enable_oplog_ && !it->second.IsGrouped()) { + PublishKvRemovedAfterEvict( + it->first, evict_result.freed_bytes, + "cpu", it->second, tenant_it->first); + } + if (!enable_oplog_ && !it->second.IsValid()) { + it = EraseMetadata( + tenant_state, it, tenant_it->first, + QuotaEraseMode::kFull, &shard); } else { ++it; } @@ -6286,32 +8835,30 @@ void MasterService::BatchEvict(double evict_ratio_target, deferred_replicas.clear(); } } else { - // This should not happen. LOG(ERROR) << "Error in second pass eviction: target_evict_num=" << target_evict_num << ", no_pin_objects.size()=" << no_pin_objects.size() << ", soft_pin_objects.size()=" << soft_pin_objects.size() << ", evicted_count=" << evicted_count - << ", object_count=" << object_count + << ", eviction_base=" << total_eviction_base << ", evict_ratio_target=" << evict_ratio_target << ", evict_ratio_lowerbound=" << evict_ratio_lowerbound; } } - if (evicted_count > 0 || released_discarded_cnt > 0 || - offload_deferred_count > 0) { - // Offload-deferred counts as partial success: work was done (objects - // queued for disk offload), so suppress re-triggering until the next - // watermark breach or explicit need_mem_eviction_ signal. + if (evicted_count > 0 || released_discarded_cnt > 0) { need_mem_eviction_ = false; MasterMetricManager::instance().inc_eviction_success(evicted_count, total_freed_size); MasterMetricManager::instance().inc_mem_eviction_success( evicted_count, total_freed_size); + } else if (offload_deferred_count > 0) { + need_mem_eviction_ = false; + MasterMetricManager::instance().inc_eviction_success(0, 0); + MasterMetricManager::instance().inc_mem_eviction_success(0, 0); } else { - if (object_count == 0) { - // No objects to evict, no need to check again + if (total_eviction_base == 0) { need_mem_eviction_ = false; } MasterMetricManager::instance().inc_eviction_fail(); @@ -6322,7 +8869,31 @@ void MasterService::BatchEvict(double evict_ratio_target, << ", offload_deferred=" << offload_deferred_count << ", offload_cap_forced=" << offload_cap_forced_count << ", offload_push_failed_forced=" << offload_push_failed_forced - << ", total_freed_size=" << total_freed_size; + << ", total_freed_size=" << total_freed_size + << ", eviction_base=" << total_eviction_base + << ", actual_evict_ratio=" + << (total_eviction_base > 0 + ? (double)evicted_count / total_eviction_base + : 0.0) + << ", target_evict_ratio=" << evict_ratio_target; + VLOG(1) << "[EVICT-DIAG] object_count=" << object_count + << " disk_object_count=" << (object_count - total_eviction_base) + << " eviction_base=" << total_eviction_base << " disk_ratio=" + << (object_count > 0 + ? (double)(object_count - total_eviction_base) / + object_count + : 0.0) + << " ideal_evict_num_inflated=" + << (long)std::ceil(object_count * evict_ratio_target) + << " ideal_evict_num_correct=" + << (long)std::ceil(total_eviction_base * evict_ratio_target); + LOG(INFO) << "[EVICT-RESULT] evicted_count=" << evicted_count + << ", eviction_base=" << total_eviction_base + << ", actual_evict_ratio=" + << (total_eviction_base > 0 + ? (double)evicted_count / total_eviction_base + : 0.0) + << ", target_evict_ratio=" << evict_ratio_target; if (offload_on_evict_ && evicted_count == 0 && offload_deferred_count > 0) { LOG(WARNING) << "[EVICT] No memory freed this cycle; " << offload_deferred_count @@ -6356,7 +8927,12 @@ void MasterService::NoFBatchEvict(double evict_ratio_target, long object_count = 0; uint64_t total_freed_size = 0; - size_t start_idx = RandomIndex(metadata_shards_.size()); + auto is_evictable_nof_replica = [](const Replica& replica) { + return replica.is_nof_replica() && replica.is_completed() && + replica.get_refcnt() == 0; + }; + + size_t start_idx = randomIndex(metadata_shards_.size()); for (size_t i = 0; i < metadata_shards_.size(); i++) { MetadataShardAccessorRW shard( this, (start_idx + i) % metadata_shards_.size()); @@ -6386,12 +8962,111 @@ void MasterService::NoFBatchEvict(double evict_ratio_target, continue; } + // Probe: any NoF replicas eligible for eviction? + const bool has_evictable_nof = + metadata.HasReplica(is_evictable_nof_replica); + if (!has_evictable_nof) { + ++it; + continue; + } + + // HA strong consistency: persist BEFORE erasing NoF replicas. + // Skip the key on persist failure. + if (enable_oplog_ && ordered_oplog_writer_) { + auto remaining = BuildRemainingReplicaDescriptors( + metadata, is_evictable_nof_replica); + if (enable_oplog_) { + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + LOG(WARNING) + << "NoFBatchEvict: OpLog reservation failed " + "for key=" + << it->first << ", err=" + << static_cast(reservation.error()) + << ", skipping eviction"; + ++it; + continue; + } + std::vector removed_ids; + metadata.VisitReplicas( + is_evictable_nof_replica, + [&removed_ids](Replica& replica) { + removed_ids.push_back(replica.id()); + replica.mark_removed(); + }); + const size_t removed_count = removed_ids.size(); + tl::expected persist_result; + if (remaining.empty()) { + persist_result = + AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), + OpType::REMOVE, tenant_it->first.value(), + it->first, {}, + [this, + removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + } else { + persist_result = AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), OpType::PUT_END, + tenant_it->first.value(), it->first, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, + remaining, metadata.group_id, + metadata.data_type), + [this, removed_ids = std::move(removed_ids)]( + const OpLogEntry& durable_entry) { + FinalizeRemovedReplicasAfterDurable( + durable_entry, removed_ids, + QuotaEraseMode::kFull); + }); + } + if (!persist_result) { + LOG(WARNING) + << "NoFBatchEvict: OpLog persist failed for " + "key=" + << it->first << ", err=" + << static_cast(persist_result.error()) + << ", skipping eviction"; + ++it; + continue; + } + total_freed_size += metadata.size * removed_count; + shard_evicted_count++; + ++it; + continue; + } + + tl::expected persist_result; + if (remaining.empty()) { + persist_result = AppendOpLogWithDurableFinalize( + OpType::REMOVE, tenant_it->first.value(), it->first, + {}, nullptr); + } else { + persist_result = AppendOpLogWithDurableFinalize( + OpType::PUT_END, tenant_it->first.value(), + it->first, + SerializeMetadataForOpLogFromReplicaDescriptors( + metadata.client_id, metadata.size, remaining, + metadata.group_id, metadata.data_type), + nullptr); + } + if (!persist_result) { + LOG(WARNING) + << "NoFBatchEvict: OpLog persist failed for key=" + << it->first << ", err=" + << static_cast(persist_result.error()) + << ", skipping eviction"; + ++it; + continue; + } + } + const size_t erased = - metadata.EraseReplicas([](const Replica& replica) { - return replica.is_nof_replica() && - replica.is_completed() && - replica.get_refcnt() == 0; - }); + metadata.EraseReplicas(is_evictable_nof_replica); if (erased == 0) { ++it; continue; @@ -6399,8 +9074,11 @@ void MasterService::NoFBatchEvict(double evict_ratio_target, total_freed_size += metadata.size * erased; shard_evicted_count++; + PublishKvRemovedAfterEvict(it->first, metadata.size * erased, + "disk", metadata, tenant_it->first); if (!metadata.IsValid()) { - it = EraseMetadata(tenant_state, it, tenant_it->first); + it = EraseMetadata(tenant_state, it, tenant_it->first, + QuotaEraseMode::kFull, &shard); } else { ++it; } @@ -6488,6 +9166,7 @@ void MasterService::ClientMonitorFunc() { ok_client_.erase(it); MasterMetricManager::instance().dec_active_clients(); } + client_host_id_.erase(client_id); } ScopedSegmentAccess segment_access = @@ -6534,6 +9213,8 @@ void MasterService::ClientMonitorFunc() { LOG(INFO) << "client_id=" << client_ids[i] << ", segment_name=" << segment_names[i] << ", action=unmount_expired_mem_segment"; + // Clean up HTTP metadata if enabled + cleanupHttpMetadata(segment_names[i]); } for (auto& client_id : expired_clients) { segment_access.UnmountLocalDiskSegment(client_id); @@ -6795,10 +9476,17 @@ MasterService::MetadataSerializer::Serialize() { // 1. Serialize metadata shards packer.pack("shards"); - // First count non-empty shards + // First count shards that have actual metadata entries. + // A shard may have empty tenants left after eviction erased all + // metadata but didn't clean up the tenant map; using metadata_count + // (not tenants.empty()) ensures the count matches the skip logic below. size_t valid_shards = 0; for (size_t i = 0; i < kNumShards; ++i) { - if (!service_->metadata_shards_[i].tenants.empty()) { + size_t metadata_count = 0; + for (const auto& [tid, ts] : service_->metadata_shards_[i].tenants) { + metadata_count += ts.metadata.size(); + } + if (metadata_count > 0) { valid_shards++; } } @@ -6810,8 +9498,16 @@ MasterService::MetadataSerializer::Serialize() { for (size_t shard_idx = 0; shard_idx < kNumShards; ++shard_idx) { const auto& shard = service_->metadata_shards_[shard_idx]; - // Skip if shard is empty - if (shard.tenants.empty()) { + // Skip shards with no actual metadata entries. + // A shard may have empty tenants left after eviction erased all + // metadata but didn't clean up the tenant map; serializing those + // would produce an entry that deserialization never recreates, + // breaking the snapshot round-trip comparison. + size_t metadata_count = 0; + for (const auto& [tid, ts] : shard.tenants) { + metadata_count += ts.metadata.size(); + } + if (metadata_count == 0) { continue; } @@ -6983,6 +9679,7 @@ MasterService::MetadataSerializer::Deserialize( Replica::next_id_.store(next_id); LOG(INFO) << "Restored Replica::next_id_ to " << next_id; service_->RebuildGroupRoutingIndex(); + service_->ClearCandidatesForReload(); return {}; } @@ -7001,6 +9698,7 @@ void MasterService::MetadataSerializer::Reset() { service_->discarded_replicas_.clear(); } Replica::next_id_.store(1); + service_->ClearCandidatesForReload(); } tl::expected @@ -7028,7 +9726,7 @@ MasterService::MetadataSerializer::SerializeShard(const MetadataShard& shard, sorted_entries.reserve(metadata_count); for (const auto& [tenant_id, tenant_state] : shard.tenants) { for (const auto& [key, metadata] : tenant_state.metadata) { - sorted_entries.push_back({tenant_id, key, &metadata}); + sorted_entries.push_back({tenant_id.value(), key, &metadata}); } } std::sort(sorted_entries.begin(), sorted_entries.end(), @@ -7102,15 +9800,14 @@ MasterService::MetadataSerializer::DeserializeShard(const msgpack::object& obj, "[tenant_id, key, metadata]")); } - std::string tenant_id = "default"; + TenantId tenant_id; std::string key; const msgpack::object* value_obj = nullptr; if (item.via.array.size == 2) { key = item.via.array.ptr[0].as(); value_obj = &item.via.array.ptr[1]; } else { - tenant_id = - NormalizeTenantId(item.via.array.ptr[0].as()); + tenant_id = TenantId(item.via.array.ptr[0].as()); key = item.via.array.ptr[1].as(); value_obj = &item.via.array.ptr[2]; } @@ -7136,6 +9833,14 @@ MasterService::MetadataSerializer::DeserializeShard(const msgpack::object& obj, it->second.lease_timeout = metadata_ptr->lease_timeout; it->second.soft_pin_timeout = metadata_ptr->soft_pin_timeout; + it->second.object_checksum = metadata_ptr->object_checksum; + + // Recompute disk_object_count for restored metadata + if (it->second.HasReplica([](const Replica& r) { + return r.is_local_disk_replica() && r.is_completed(); + })) { + shard.disk_object_count++; + } } return {}; @@ -7148,12 +9853,15 @@ MasterService::MetadataSerializer::SerializeMetadata( // Pack ObjectMetadata using array structure for efficiency // Format: [client_id, put_start_time, size, lease_timeout, // has_soft_pin_timeout, soft_pin_timeout, replicas_count, data_type, - // replicas..., hard_pinned, group_id] + // replicas..., hard_pinned, group_id, object_checksum?] size_t array_size = 10; // client_id, put_start_time, size, lease_timeout, // has_soft_pin_timeout, soft_pin_timeout, // replicas_count, data_type, hard_pinned, group_id array_size += metadata.CountReplicas(); // One element per replica + if (metadata.object_checksum.has_value()) { + ++array_size; + } packer.pack_array(array_size); // Serialize client_id @@ -7206,6 +9914,9 @@ MasterService::MetadataSerializer::SerializeMetadata( packer.pack(metadata.IsHardPinned()); packer.pack(metadata.group_id); + if (metadata.object_checksum.has_value()) { + packer.pack(*metadata.object_checksum); + } return {}; } @@ -7234,7 +9945,12 @@ MasterService::MetadataSerializer::DeserializeMetadata( // Deserialize client_id string std::string client_id_str = array[index++].as(); UUID client_id; - StringToUuid(client_id_str, client_id); + if (!StringToUuid(client_id_str, client_id)) { + return tl::unexpected(SerializationError( + ErrorCode::DESERIALIZE_FAIL, + fmt::format("deserialize ObjectMetadata invalid client_id UUID: {}", + client_id_str))); + } // Deserialize put_start_time uint64_t put_start_time_timestamp = array[index++].as(); @@ -7259,11 +9975,12 @@ MasterService::MetadataSerializer::DeserializeMetadata( // v2: 8 + replicas_count, either data_type or hard_pinned // v3: 9 + replicas_count, data_type + hard_pinned or hard_pinned + // group_id v4: 10 + replicas_count, data_type + hard_pinned + group_id + // v5: 11 + replicas_count, v4 + object_checksum // 64-bit arithmetic keeps an attacker-controlled near-UINT32_MAX // replicas_count from wrapping the bounds and slipping an out-of-bounds // index past the size check. constexpr uint64_t kBaseFieldCount = 7; - constexpr uint64_t kMaxOptionalFieldCount = 3; + constexpr uint64_t kMaxOptionalFieldCount = 4; const uint64_t total_elements = obj.via.array.size; const uint64_t min_elements = kBaseFieldCount + replicas_count; if (total_elements < min_elements || @@ -7313,6 +10030,17 @@ MasterService::MetadataSerializer::DeserializeMetadata( group_id = array[index++].as(); } + std::optional object_checksum; + if (index < total_elements && + array[index].type == msgpack::type::POSITIVE_INTEGER) { + object_checksum = array[index++].as(); + } + if (index != total_elements) { + return tl::unexpected(SerializationError( + ErrorCode::DESERIALIZE_FAIL, + "deserialize ObjectMetadata optional field type mismatch")); + } + // Create ObjectMetadata instance bool enable_soft_pin = has_soft_pin_timeout; auto metadata = std::make_unique( @@ -7321,6 +10049,7 @@ MasterService::MetadataSerializer::DeserializeMetadata( std::chrono::milliseconds(put_start_time_timestamp)), size, std::move(replicas), enable_soft_pin, is_hard_pinned, data_type, group_id); + metadata->object_checksum = object_checksum; metadata->lease_timeout = std::chrono::system_clock::time_point( std::chrono::milliseconds(lease_timestamp)); @@ -7334,28 +10063,16 @@ MasterService::MetadataSerializer::DeserializeMetadata( return metadata; } -std::string MasterService::FormatTimestamp( - const std::chrono::system_clock::time_point& tp) { - auto time_t = std::chrono::system_clock::to_time_t(tp); - - std::stringstream ss; - ss << std::put_time(std::localtime(&time_t), "%Y%m%d_%H%M%S"); - - // Add milliseconds to ensure uniqueness - auto ms = std::chrono::duration_cast( - tp.time_since_epoch()) % - 1000; - - ss << "_" << std::setfill('0') << std::setw(3) << ms.count(); - - return ss.str(); -} - tl::expected MasterService::CreateCopyTask( - const std::string& key, const std::string& tenant_id, + const std::string& key, const TenantId& tenant_id, const std::vector& targets) { + auto normalized_tenant_result = ResolveTenantIdForWrite(tenant_id); + if (!normalized_tenant_result) { + return tl::make_unexpected(normalized_tenant_result.error()); + } + const ObjectIdentity object_id{std::move(normalized_tenant_result.value()), + key}; std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); if (targets.empty()) { LOG(ERROR) << "key=" << key << ", error=empty_targets"; return tl::make_unexpected(ErrorCode::INVALID_PARAMS); @@ -7389,9 +10106,8 @@ tl::expected MasterService::CreateCopyTask( } // Randomly pick a segment from the source replicas - static thread_local std::mt19937 gen(std::random_device{}()); - std::uniform_int_distribution dis(0, segment_names.size() - 1); - std::string selected_source_segment = segment_names[dis(gen)]; + std::string selected_source_segment = + segment_names[randomIndex(segment_names.size())]; UUID select_client; ErrorCode error = segment_accessor.GetClientIdBySegmentName( selected_source_segment, select_client); @@ -7403,17 +10119,22 @@ tl::expected MasterService::CreateCopyTask( } return task_manager_.get_write_access() .submit_task_typed( - select_client, {.tenant_id = object_id.tenant_id, + select_client, {.tenant_id = object_id.tenant_id.value(), .key = object_id.user_key, .source = selected_source_segment, .targets = targets}); } tl::expected MasterService::CreateMoveTask( - const std::string& key, const std::string& tenant_id, + const std::string& key, const TenantId& tenant_id, const std::string& source, const std::string& target) { + auto normalized_tenant_result = ResolveTenantIdForWrite(tenant_id); + if (!normalized_tenant_result) { + return tl::make_unexpected(normalized_tenant_result.error()); + } + const ObjectIdentity object_id{std::move(normalized_tenant_result.value()), + key}; std::shared_lock shared_lock(snapshot_mutex_); - const auto object_id = MakeObjectIdentity(key, tenant_id); MetadataAccessorRO accessor(this, object_id); if (!accessor.Exists()) { VLOG(1) << "key=" << key << ", info=object_not_found"; @@ -7460,7 +10181,7 @@ tl::expected MasterService::CreateMoveTask( return task_manager_.get_write_access() .submit_task_typed( - select_client, {.tenant_id = object_id.tenant_id, + select_client, {.tenant_id = object_id.tenant_id.value(), .key = object_id.user_key, .source = source, .target = target}); @@ -7670,10 +10391,9 @@ tl::expected MasterService::CancelDrainJob( } std::string MasterService::MakeDrainUnitKey( - const std::string& tenant_id, const std::string& key, + const TenantId& tenant_id, const std::string& key, const std::string& source_segment) const { - const auto normalized_tenant = NormalizeTenantId(tenant_id); - return std::to_string(normalized_tenant.size()) + ":" + normalized_tenant + + return std::to_string(tenant_id.value().size()) + ":" + tenant_id.value() + ":" + std::to_string(key.size()) + ":" + key + ":" + source_segment; } @@ -7769,7 +10489,7 @@ void MasterService::ScheduleDrainJobTasks(DrainJob& job) { } struct DrainPlan { - std::string tenant_id; + TenantId tenant_id; std::string key; std::string source_segment; std::string target_segment; @@ -8094,4 +10814,474 @@ MasterService::MetadataSerializer::DeserializeDiscardedReplicas( return {}; } +KvEventConfig MasterService::BuildKvEventConfig( + const MasterServiceConfig& config) { + KvEventConfig kv_config; + kv_config.enabled = config.enable_kv_events; + kv_config.bind_endpoint = config.kv_events_bind_endpoint; + kv_config.model_name = config.kv_events_model_name; + kv_config.backend_id = config.kv_events_backend_id; + kv_config.tenant_id = config.kv_events_tenant_id; + kv_config.additional_salt = config.kv_events_additional_salt; + kv_config.lora_name = config.kv_events_lora_name; + kv_config.block_size = config.kv_events_block_size; + kv_config.dp_rank = config.kv_events_dp_rank; + kv_config.emit_legacy_compat_fields = config.kv_events_emit_legacy_compat; + kv_config.emit_object_key = config.kv_events_emit_object_key; + kv_config.queue_capacity = config.kv_events_queue_capacity; + return kv_config; +} + +std::string MasterService::MediumForReplicaType(ReplicaType replica_type) { + switch (replica_type) { + case ReplicaType::MEMORY: + return "cpu"; + case ReplicaType::DISK: + case ReplicaType::LOCAL_DISK: + case ReplicaType::NOF_SSD: + return "disk"; + case ReplicaType::ALL: + default: + return "cpu"; + } +} + +std::string MasterService::MediumForMetadata(const ObjectMetadata& metadata) { + if (metadata.HasMemReplica()) { + return "cpu"; + } + if (metadata.HasReplica(&Replica::fn_is_nof_replica) || + metadata.HasReplica(&Replica::fn_is_disk_replica) || + metadata.HasReplica(&Replica::fn_is_local_disk_replica)) { + return "disk"; + } + return "cpu"; +} + +void MasterService::PublishKvStored(const std::string& key, + ReplicaType replica_type, + const ObjectMetadata& metadata, + const TenantId& tenant_id) { + if (!kv_event_publisher_ || !kv_event_publisher_->enabled()) { + return; + } + std::string medium = MediumForReplicaType(replica_type); + if (replica_type == ReplicaType::ALL) { + medium = MediumForMetadata(metadata); + } + kv_event_publisher_->PublishStored(key, medium, tenant_id, + metadata.group_id); +} + +void MasterService::PublishKvRemoved(const std::string& key, + const std::string& medium, + const TenantId& tenant_id, + const std::string& group_id) { + if (!kv_event_publisher_ || !kv_event_publisher_->enabled()) { + return; + } + kv_event_publisher_->PublishRemoved(key, medium, tenant_id, group_id); +} + +void MasterService::PublishKvRemoved(const std::string& key, + const ObjectMetadata& metadata, + const TenantId& tenant_id) { + PublishKvRemoved(key, MediumForMetadata(metadata), tenant_id, + metadata.group_id); +} + +void MasterService::PublishKvRemovedAfterEvict(const std::string& key, + uint64_t freed_bytes, + const std::string& medium, + const ObjectMetadata& metadata, + const TenantId& tenant_id) { + (void)freed_bytes; + (void)medium; + if (!kv_event_publisher_ || !kv_event_publisher_->enabled()) { + return; + } + if (!metadata.IsValid()) { + PublishKvRemoved(key, metadata, tenant_id); + } +} + +bool MasterService::KvEventsEnabled() const { + return kv_event_publisher_ && kv_event_publisher_->enabled(); +} + +KvEventPublisher::Stats MasterService::GetKvEventStats() const { + if (!kv_event_publisher_) { + return {}; + } + return kv_event_publisher_->GetStats(); +} + +void MasterService::setHttpMetadataServer(HttpMetadataServer* server) { + http_metadata_server_ = server; + if (server) { + LOG(INFO) << "HTTP metadata cleanup on client timeout: enabled " + "(co-located metadata server)"; + } +} + +void MasterService::setHttpMetadataRemoteUrl( + const std::string& metadata_connstring) { +#ifdef USE_HTTP + // Only http(s) is supported; guard the scheme to avoid + // MetadataStoragePlugin::Create()'s LOG(FATAL) on other backends. + if (metadata_connstring.rfind("http://", 0) == 0 || + metadata_connstring.rfind("https://", 0) == 0) { + try { + http_metadata_remote_ = + MetadataStoragePlugin::Create(metadata_connstring); + LOG(INFO) << "HTTP metadata cleanup on client timeout: enabled " + "(remote metadata server " + << metadata_connstring << ")"; + // Start async cleanup worker now that http_metadata_remote_ is + // ready + http_metadata_cleanup_running_ = true; + http_metadata_cleanup_thread_ = std::thread( + &MasterService::HttpMetadataCleanupThreadFunc, this); + LOG(INFO) << "HTTP metadata cleanup worker thread started"; + } catch (const std::exception& e) { + LOG(WARNING) << "Failed to initialize remote HTTP metadata client " + "for " + << metadata_connstring << ": " << e.what() + << ". Metadata cleanup on timeout disabled."; + http_metadata_remote_.reset(); + } + return; + } + LOG(WARNING) << "enable_metadata_cleanup_on_timeout is set but the " + "configured metadata server '" + << metadata_connstring + << "' is not an HTTP endpoint; remote cleanup currently " + "supports only http(s). Metadata cleanup on timeout " + "disabled."; +#else + (void)metadata_connstring; + LOG(WARNING) << "enable_metadata_cleanup_on_timeout is set but this build " + "has no HTTP metadata support (USE_HTTP=OFF); metadata " + "cleanup on timeout disabled."; +#endif +} + +void MasterService::cleanupHttpMetadata(const std::string& segment_name) { + // Co-located: remove in-process, safe to run inline (no network I/O). + if (http_metadata_server_) { + const std::string ram_key = + http_metadata_prefix_ + "ram/" + segment_name; + const std::string rpc_key = + http_metadata_prefix_ + "rpc_meta/" + segment_name; + bool ram_removed = http_metadata_server_->removeKey(ram_key); + bool rpc_removed = http_metadata_server_->removeKey(rpc_key); + LOG(INFO) << "Cleaned up HTTP metadata for segment: " << segment_name + << ", ram_key_removed=" << ram_removed + << ", rpc_key_removed=" << rpc_removed; + return; + } + + // Separately-deployed: enqueue for async cleanup so a slow/unreachable + // server never blocks the client monitor thread. + if (http_metadata_remote_) { + { + std::lock_guard lk(http_metadata_cleanup_mutex_); + http_metadata_cleanup_queue_.push_back(segment_name); + } + http_metadata_cleanup_cv_.notify_one(); + return; + } + + // Neither configured: cleanup is disabled, nothing to do. +} + +void MasterService::HttpMetadataCleanupThreadFunc() { + LOG(INFO) << "HTTP metadata cleanup worker started"; + while (http_metadata_cleanup_running_) { + std::vector batch; + { + std::unique_lock lk(http_metadata_cleanup_mutex_); + http_metadata_cleanup_cv_.wait(lk, [&] { + return !http_metadata_cleanup_queue_.empty() || + !http_metadata_cleanup_running_.load(); + }); + if (!http_metadata_cleanup_running_ && + http_metadata_cleanup_queue_.empty()) { + break; + } + batch.swap(http_metadata_cleanup_queue_); + } + + for (const auto& segment_name : batch) { + const std::string ram_key = + http_metadata_prefix_ + "ram/" + segment_name; + const std::string rpc_key = + http_metadata_prefix_ + "rpc_meta/" + segment_name; + + // Each key attempted independently so one failure does not + // prevent cleanup of the other. + bool ram_removed = false; + bool rpc_removed = false; + try { + ram_removed = http_metadata_remote_->remove(ram_key); + } catch (const std::exception& e) { + LOG(WARNING) + << "Remote HTTP metadata cleanup failed for ram_key: " + << ram_key << ": " << e.what(); + } + try { + rpc_removed = http_metadata_remote_->remove(rpc_key); + } catch (const std::exception& e) { + LOG(WARNING) + << "Remote HTTP metadata cleanup failed for rpc_key: " + << rpc_key << ": " << e.what(); + } + LOG(INFO) << "Cleaned up remote HTTP metadata for segment: " + << segment_name << ", ram_key_removed=" << ram_removed + << ", rpc_key_removed=" << rpc_removed; + } + } + LOG(INFO) << "HTTP metadata cleanup worker stopped"; +} + +std::string MasterService::SerializeMetadataForOpLog( + const ObjectMetadata& metadata) const { + MetadataPayload payload; + payload.client_id = metadata.client_id; + payload.size = metadata.size; + payload.group_id = metadata.group_id; + payload.data_type = metadata.data_type; + + // Extract replica descriptors - get them all at once + const auto& replicas = metadata.GetAllReplicas(); + payload.replicas.reserve(replicas.size()); + for (const auto& replica : replicas) { + payload.replicas.push_back(replica.get_descriptor()); + } + + // NOTE: Lease information is NOT serialized because: + // 1. Standby does not perform eviction, so lease info is not used + // 2. After promotion, new Primary should grant fresh leases, not restore + // old ones + + // Serialize using struct_pack (msgpack binary format) + auto result = struct_pack::serialize(payload); + return std::string(result.begin(), result.end()); +} + +std::string MasterService::SerializeMetadataForOpLogWithoutMemReplicas( + const ObjectMetadata& metadata) const { + MetadataPayload payload; + payload.client_id = metadata.client_id; + payload.size = metadata.size; + payload.group_id = metadata.group_id; + payload.data_type = metadata.data_type; + + const auto& replicas = metadata.GetAllReplicas(); + payload.replicas.reserve(replicas.size()); + for (const auto& replica : replicas) { + if (replica.type() == ReplicaType::MEMORY) { + continue; + } + payload.replicas.push_back(replica.get_descriptor()); + } + + auto result = struct_pack::serialize(payload); + return std::string(result.begin(), result.end()); +} + +std::string MasterService::SerializeMetadataForOpLogFromReplicaDescriptors( + const UUID& client_id, uint64_t size, + const std::vector& replicas, + const std::string& group_id, ObjectDataType data_type) const { + MetadataPayload payload; + payload.client_id = client_id; + payload.size = size; + payload.replicas = replicas; + payload.group_id = group_id; + payload.data_type = data_type; + auto result = struct_pack::serialize(payload); + return std::string(result.begin(), result.end()); +} + +ErrorCode MasterService::InitializeBatchOpLogWriter( + std::shared_ptr backend) { + if (!backend || !backend->SupportsTxn()) { + return ErrorCode::INVALID_PARAMS; + } + + auto storage = std::make_unique(cluster_id_, *backend); + DurablePrefix durable_prefix; + ErrorCode err = storage->InitDurablePrefix(durable_prefix); + if (err != ErrorCode::OK) { + return err; + } + + OrderedOpLogWriterConfig writer_config; + writer_config.max_entries_per_batch = oplog_batch_max_entries_; + writer_config.initial_durable_prefix = durable_prefix; + OpLogBatchStorage* storage_ptr = storage.get(); + auto writer = std::make_unique( + writer_config, [storage_ptr](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage_ptr->WriteBatchAndAdvancePrefix(batch, + expected_prefix); + }); + if (!writer->IsAccepting()) { + return writer->LastError(); + } + writer->Start(); + + if (ordered_oplog_writer_) { + ordered_oplog_writer_->Stop(); + } + batch_oplog_kv_backend_ = std::move(backend); + batch_oplog_storage_ = std::move(storage); + ordered_oplog_writer_ = std::move(writer); + return ErrorCode::OK; +} + +tl::expected +MasterService::AppendOpLogVisibleBeforeDurable(OpType type, + const std::string& tenant_id, + const std::string& key, + const std::string& payload) { + if (!enable_oplog_) { + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + if (!ordered_oplog_writer_) { + return tl::unexpected(ErrorCode::INTERNAL_ERROR); + } + + const TenantId resolved_tenant(enable_multi_tenants_ + ? tenant_id + : std::string(TenantId::kDefaultValue)); + if (!resolved_tenant.IsValid()) { + return tl::unexpected(ErrorCode::TENANT_NOT_REGISTERED); + } + + auto reservation = ordered_oplog_writer_->Reserve(); + if (!reservation) { + return tl::unexpected(reservation.error()); + } + OpLogEntry entry; + entry.op_type = type; + entry.tenant_id = resolved_tenant.value(); + entry.object_key = key; + entry.payload = payload; + auto pending = ordered_oplog_writer_->Commit(std::move(reservation.value()), + std::move(entry), nullptr); + if (!pending) { + return tl::unexpected(pending.error()); + } + return pending.value().sequence_id(); +} + +tl::expected +MasterService::AppendOpLogWithDurableFinalize( + OpType type, const std::string& tenant_id, const std::string& key, + const std::string& payload, DurableFinalizeCallback callback) { + if (!enable_oplog_) { + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + + auto reservation = ReserveBatchOpLogSlot(); + if (!reservation) { + return tl::unexpected(reservation.error()); + } + return AppendReservedOpLogWithDurableFinalize( + std::move(reservation.value()), type, tenant_id, key, payload, + std::move(callback)); +} + +tl::expected +MasterService::ReserveBatchOpLogSlot() { + if (!enable_oplog_) { + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + if (!ordered_oplog_writer_) { + return tl::unexpected(ErrorCode::INTERNAL_ERROR); + } + return ordered_oplog_writer_->Reserve(); +} + +tl::expected +MasterService::AppendReservedOpLogWithDurableFinalize( + OrderedOpLogWriter::Reservation&& reservation, OpType type, + const std::string& tenant_id, const std::string& key, + const std::string& payload, DurableFinalizeCallback callback) { + if (!enable_oplog_) { + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + const TenantId resolved_tenant(enable_multi_tenants_ + ? tenant_id + : std::string(TenantId::kDefaultValue)); + if (!resolved_tenant.IsValid()) { + return tl::unexpected(ErrorCode::TENANT_NOT_REGISTERED); + } + OpLogEntry entry; + entry.op_type = type; + entry.tenant_id = resolved_tenant.value(); + entry.object_key = key; + entry.payload = payload; + auto pending = ordered_oplog_writer_->Commit(std::move(reservation), entry, + std::move(callback)); + if (!pending) { + return tl::unexpected(pending.error()); + } + entry.sequence_id = pending.value().sequence_id(); + return entry; +} + +tl::expected MasterService::PersistRemoveForHA( + const char* why, const std::string& key) { + return PersistRemoveForHA(why, TenantId::Default(), key); +} + +tl::expected MasterService::PersistRemoveForHA( + const char* why, const TenantId& tenant_id, const std::string& key) { + auto result = AppendOpLogWithDurableFinalize( + OpType::REMOVE, tenant_id.value(), key, {}, nullptr); + if (!result) { + LOG(WARNING) << why << ": REMOVE persist failed for key=" << key + << ", err=" << static_cast(result.error()); + return tl::unexpected(result.error()); + } + return {}; +} + +void MasterService::PersistSegmentOpForHAOrEnqueue(const char* why, OpType type, + const std::string& key, + const std::string& payload) { + PersistSegmentOpForHAOrEnqueue(why, type, TenantId::Default(), key, + payload); +} + +void MasterService::PersistSegmentOpForHAOrEnqueue(const char* why, OpType type, + const TenantId& tenant_id, + const std::string& key, + const std::string& payload) { + auto result = + AppendOpLogVisibleBeforeDurable(type, tenant_id.value(), key, payload); + if (!result) { + LOG(WARNING) << why << ": segment OpLog queue failed for key=" << key + << ", type=" << static_cast(type) + << ", err=" << static_cast(result.error()); + } +} + +std::vector +MasterService::BuildRemainingReplicaDescriptors( + const ObjectMetadata& metadata, + const std::function& should_remove) const { + std::vector remaining; + for (const auto& replica : metadata.GetAllReplicas()) { + if (!should_remove(replica) && + replica.status() == ReplicaStatus::COMPLETE) { + remaining.push_back(replica.get_descriptor()); + } + } + return remaining; +} + } // namespace mooncake diff --git a/mooncake-store/src/master_snapshot_manager.cpp b/mooncake-store/src/master_snapshot_manager.cpp new file mode 100644 index 0000000000..91f2e38609 --- /dev/null +++ b/mooncake-store/src/master_snapshot_manager.cpp @@ -0,0 +1,624 @@ +#include "master_snapshot_manager.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "master_service.h" +#include "master_metric_manager.h" +#include "master_snapshot_repository.h" +#include "ha/snapshot/catalog/snapshot_catalog_store.h" +#include "ha/snapshot/object/snapshot_object_store.h" +#include "ha/snapshot/snapshot_constants.h" +#include "ha/snapshot/snapshot_logger.h" +#include "ha/oplog/oplog_batch_storage.h" +#include "serialize/serializer.h" +#include "segment.h" +#include "task_manager.h" +#include "utils/file_util.h" +#include "utils/zstd_util.h" + +namespace mooncake { + +namespace { +int64_t CurrentTimeMs() { + return std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count(); +} +} // namespace + +MasterSnapshotManager::MasterSnapshotManager( + MasterService* master_service, MasterSnapshotManagerOptions options, + std::shared_mutex& snapshot_mutex, + SnapshotObjectStore* snapshot_object_store, + ha::SnapshotCatalogStore* snapshot_catalog_store) + : master_service_(master_service), + options_(std::move(options)), + snapshot_mutex_(snapshot_mutex), + snapshot_object_store_(snapshot_object_store), + snapshot_catalog_store_(snapshot_catalog_store), + repository_(std::make_unique( + snapshot_object_store, snapshot_catalog_store, + options_.snapshot_backup_dir, options_.use_snapshot_backup_dir)) {} + +MasterSnapshotManager::~MasterSnapshotManager() { + Stop(); + if (snapshot_thread_.joinable()) { + snapshot_thread_.join(); + } +} + +void MasterSnapshotManager::Start() { + if (snapshot_running_.load()) { + return; + } + snapshot_running_ = true; + snapshot_thread_ = + std::thread(&MasterSnapshotManager::SnapshotThreadFunc, this); + LOG(INFO) << "[MasterSnapshotManager] Started"; +} + +void MasterSnapshotManager::Stop() { + { + std::lock_guard lk(snapshot_thread_mutex_); + snapshot_running_ = false; + } + snapshot_thread_cv_.notify_all(); + LOG(INFO) << "[MasterSnapshotManager] Stop signaled"; +} + +std::string MasterSnapshotManager::FormatTimestamp( + const std::chrono::system_clock::time_point& tp) { + auto time_t = std::chrono::system_clock::to_time_t(tp); + + std::stringstream ss; + std::tm tm_now; + localtime_r(&time_t, &tm_now); + ss << std::put_time(&tm_now, "%Y%m%d_%H%M%S"); + + // Add milliseconds to ensure uniqueness + auto ms = std::chrono::duration_cast( + tp.time_since_epoch()) % + 1000; + + ss << "_" << std::setfill('0') << std::setw(3) << ms.count(); + + return ss.str(); +} + +void MasterSnapshotManager::SnapshotThreadFunc() { + LOG(INFO) << "[Snapshot] snapshot_thread started"; + while (snapshot_running_) { + // Wait for the next snapshot cycle, but allow fast shutdown. + { + std::unique_lock lk(snapshot_thread_mutex_); + snapshot_thread_cv_.wait_for( + lk, std::chrono::seconds(options_.snapshot_interval_seconds), + [&] { return !snapshot_running_.load(); }); + } + + if (!snapshot_running_) { + break; + } + + if (!options_.enable_snapshot) { + // Snapshot is disabled + LOG(INFO) + << "[Snapshot] Snapshot is disabled, waiting for next cycle"; + continue; + } + // Fork a child process to save current state + + std::string snapshot_id = + FormatTimestamp(std::chrono::system_clock::now()); + LOG(INFO) << "[Snapshot] Preparing to fork child process, snapshot_id=" + << snapshot_id; + + // Create pipe for child process logging + int log_pipe[2]; + if (pipe(log_pipe) == -1) { + LOG(ERROR) << "[Snapshot] Failed to create log pipe: " + << strerror(errno) << ", snapshot_id=" << snapshot_id; + continue; + } + + const std::string& snapshot_root = + snapshot_catalog_store_->GetSnapshotRoot(); + const std::string path_prefix = snapshot_root + snapshot_id + "/"; + const std::string manifest_path = + path_prefix + ha::kSnapshotManifestFile; + auto descriptor = + BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix); + if (!descriptor) { + LOG(ERROR) << "[Snapshot] Failed to build descriptor before fork, " + "snapshot_id=" + << snapshot_id + << ", code=" << toString(descriptor.error().code) + << ", msg=" << descriptor.error().message; + close(log_pipe[0]); + close(log_pipe[1]); + continue; + } + + pid_t pid; + { + std::unique_lock lock(snapshot_mutex_); + LOG(INFO) << "[Snapshot] Locking snapshot mutex, snapshot_id=" + << snapshot_id; + pid = fork(); + } + if (pid == -1) { + // Fork failed + LOG(ERROR) << "[Snapshot] Failed to fork child process for state " + "persistence: " + << strerror(errno) << ", snapshot_id=" << snapshot_id; + close(log_pipe[0]); + close(log_pipe[1]); + } else if (pid == 0) { + // Child process + // Close read end, set write end for logging + close(log_pipe[0]); + g_snapshot_log_pipe_fd = log_pipe[1]; + + // Save current state using the configured persistence mechanism + SNAP_LOG_INFO("[Snapshot] Child process started, snapshot_id={}", + snapshot_id); + auto result = PersistState(descriptor.value()); + if (!result) { + SNAP_LOG_ERROR( + "[Snapshot] Child process failed to persist state, " + "snapshot_id={},code={},msg={}", + snapshot_id, toString(result.error().code), + result.error().message); + close(log_pipe[1]); + _exit(1); // Exit child process with error + } + SNAP_LOG_INFO( + "[Snapshot] Child process successfully persisted state, " + "snapshot_id={}", + snapshot_id); + + close(log_pipe[1]); + _exit(0); // Exit child process successfully + } else { + // Parent process + // Close write end, pass read end to wait function + close(log_pipe[1]); + WaitForSnapshotChild(pid, snapshot_id, log_pipe[0]); + close(log_pipe[0]); + } + } + LOG(INFO) << "[Snapshot] snapshot_thread stopped"; +} + +void MasterSnapshotManager::WaitForSnapshotChild(pid_t pid, + const std::string& snapshot_id, + int log_pipe_fd) { + // Default 5 minute timeout + const int64_t timeout_seconds = options_.snapshot_child_timeout_seconds; + + LOG(INFO) + << "[Snapshot] waiting for child process to complete, snapshot_id=" + << snapshot_id << ", child_pid=" << pid + << ", timeout=" << timeout_seconds << "s"; + + // Set pipe to non-blocking mode + int flags = fcntl(log_pipe_fd, F_GETFL, 0); + if (flags == -1 || fcntl(log_pipe_fd, F_SETFL, flags | O_NONBLOCK) == -1) { + LOG(WARNING) << "[Snapshot] Failed to set pipe non-blocking: " + << strerror(errno); + } + + // Buffer for reading child logs + char buf[4096]; + std::string log_buffer; + + // Helper lambda to read and output child logs + auto flush_child_logs = [&]() { + while (true) { + ssize_t n = read(log_pipe_fd, buf, sizeof(buf) - 1); + if (n > 0) { + buf[n] = '\0'; + log_buffer += buf; + // Output complete lines + size_t pos; + while ((pos = log_buffer.find('\n')) != std::string::npos) { + std::string line = log_buffer.substr(0, pos); + log_buffer.erase(0, pos + 1); + if (!line.empty()) { + LOG(INFO) << "[Snapshot:Child] " << line; + } + } + } else { + break; + } + } + }; + + // Record start time + auto start_time = std::chrono::steady_clock::now(); + + // Use non-blocking polling to wait + while (true) { + // Read child logs first + flush_child_logs(); + + int status; + pid_t result = waitpid(pid, &status, WNOHANG); + + if (result == -1) { + LOG(ERROR) << "[Snapshot] Failed to wait for child process: " + << strerror(errno) << ", snapshot_id=" << snapshot_id + << ", child_pid=" << pid; + MasterMetricManager::instance().inc_snapshot_fail(); + return; + } else if (result == 0) { + // Child process is still running + auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start_time) + .count(); + + if (elapsed >= timeout_seconds) { + // Timeout handling - flush remaining logs before killing + flush_child_logs(); + if (!log_buffer.empty()) { + LOG(INFO) << "[Snapshot:Child] " << log_buffer; + } + HandleChildTimeout(pid, snapshot_id); + MasterMetricManager::instance().inc_snapshot_fail(); + return; + } + + // Brief sleep before checking again + std::this_thread::sleep_for(std::chrono::seconds(2)); + } else { + // Child process has exited + // Flush remaining logs from child + flush_child_logs(); + // Output any remaining incomplete line + if (!log_buffer.empty()) { + LOG(INFO) << "[Snapshot:Child] " << log_buffer; + } + + HandleChildExit(pid, status, snapshot_id); + auto elapsed = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - start_time) + .count(); + MasterMetricManager::instance().set_snapshot_duration_ms(elapsed); + return; + } + } +} + +void MasterSnapshotManager::HandleChildTimeout(pid_t pid, + const std::string& snapshot_id) { + LOG(WARNING) << "[Snapshot] Child process timeout, snapshot_id=" + << snapshot_id << ", child_pid=" << pid + << ", killing child process"; + + // Try to gracefully terminate the child process + if (kill(pid, SIGTERM) == 0) { + // Wait a few seconds to see if it exits gracefully + std::this_thread::sleep_for(std::chrono::seconds(5)); + + // Check if it has exited + int status; + if (waitpid(pid, &status, WNOHANG) == 0) { + // Child process still not exited, force kill + LOG(WARNING) << "[Snapshot] Child process still running, force " + "killing, snapshot_id=" + << snapshot_id << ", child_pid=" << pid; + kill(pid, SIGKILL); + + // Wait for force termination to complete + waitpid(pid, &status, 0); + LOG(WARNING) + << "[Snapshot] Child process force killed, snapshot_id=" + << snapshot_id << ", child_pid=" << pid; + } else { + LOG(INFO) << "[Snapshot] Child process terminated gracefully after " + "SIGTERM, snapshot_id=" + << snapshot_id << ", child_pid=" << pid; + } + } else { + LOG(ERROR) << "[Snapshot] Failed to send SIGTERM to child process, " + "snapshot_id=" + << snapshot_id << ", child_pid=" << pid + << ", error=" << strerror(errno); + } +} + +void MasterSnapshotManager::HandleChildExit(pid_t pid, int status, + const std::string& snapshot_id) { + if (WIFEXITED(status)) { + int exit_code = WEXITSTATUS(status); + if (exit_code != 0) { + LOG(ERROR) << "[Snapshot] Child process exited with error code: " + << exit_code << ", snapshot_id=" << snapshot_id + << ", child_pid=" << pid; + MasterMetricManager::instance().inc_snapshot_fail(); + } else { + LOG(INFO) << "[Snapshot] Child process successfully persisted " + "state, snapshot_id=" + << snapshot_id << ", child_pid=" << pid; + MasterMetricManager::instance().inc_snapshot_success(); + } + } else if (WIFSIGNALED(status)) { + int signal = WTERMSIG(status); + LOG(ERROR) << "[Snapshot] Child process terminated by signal: " + << signal << ", snapshot_id=" << snapshot_id + << ", child_pid=" << pid; + MasterMetricManager::instance().inc_snapshot_fail(); + } +} + +tl::expected +MasterSnapshotManager::ResolveSnapshotSequenceId() const { + if (!options_.enable_ha || !master_service_->enable_oplog_) { + // OpLog sequence ids start at 1. Returning 0 here is a sentinel that + // means "no persisted OpLog boundary", so a standby that later calls + // Recover(0) will replay from the first entry when oplog following is + // enabled. + return ha::OpLogSequenceId{0}; + } + + if (!master_service_->batch_oplog_storage_) { + return tl::make_unexpected(SerializationError( + ErrorCode::INTERNAL_ERROR, + "snapshot sequence resolution requires batch OpLog storage")); + } + DurablePrefix prefix; + ErrorCode err = + master_service_->batch_oplog_storage_->ReadDurablePrefix(prefix); + if (err == ErrorCode::ETCD_KEY_NOT_EXIST || + err == ErrorCode::OPLOG_ENTRY_NOT_FOUND) { + return ha::OpLogSequenceId{0}; + } + if (err != ErrorCode::OK) { + return tl::make_unexpected(SerializationError( + err, + fmt::format("failed to resolve batch snapshot sequence boundary: " + "{}", + toString(err)))); + } + return static_cast(prefix.last_seq); +} + +tl::expected +MasterSnapshotManager::BuildSnapshotDescriptor( + const std::string& snapshot_id, const std::string& manifest_path, + const std::string& object_prefix) const { + auto sequence_id = ResolveSnapshotSequenceId(); + if (!sequence_id) { + return tl::make_unexpected(sequence_id.error()); + } + + const std::string& snapshot_root = + snapshot_catalog_store_->GetSnapshotRoot(); + auto descriptor = ha::snapshot_catalog_store_detail::MakeSnapshotDescriptor( + snapshot_root, snapshot_id); + descriptor.last_included_seq = sequence_id.value(); + descriptor.producer_view_version = master_service_->view_version_; + descriptor.manifest_key = manifest_path; + descriptor.object_prefix = object_prefix; + descriptor.created_at_ms = CurrentTimeMs(); + return descriptor; +} + +tl::expected MasterSnapshotManager::PersistState( + const std::string& snapshot_id) { + const std::string& snapshot_root = + snapshot_catalog_store_->GetSnapshotRoot(); + const std::string path_prefix = snapshot_root + snapshot_id + "/"; + const std::string manifest_path = path_prefix + ha::kSnapshotManifestFile; + auto descriptor = + BuildSnapshotDescriptor(snapshot_id, manifest_path, path_prefix); + if (!descriptor) { + return tl::make_unexpected(descriptor.error()); + } + return PersistState(descriptor.value()); +} + +tl::expected MasterSnapshotManager::PersistState( + const ha::SnapshotDescriptor& descriptor) { + const std::string& snapshot_id = descriptor.snapshot_id; + const std::string& path_prefix = descriptor.object_prefix; + const std::string& manifest_path = descriptor.manifest_key; + + try { + if (!snapshot_catalog_store_) { + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, + "snapshot catalog store is not initialized")); + } + + SNAP_LOG_INFO( + "[Snapshot] action=persisting_state start, snapshot_id={}, " + "serializer_type={}, version={}", + snapshot_id, ha::kSnapshotSerializerType, + ha::kSnapshotSerializerVersion); + + // Use the new MasterSnapshotCodec to encode all state + ha::MasterSnapshotCodec codec; + ha::MasterSnapshotStateView state_view( + *master_service_, master_service_->segment_manager_, + master_service_->nof_segment_manager_, + master_service_->task_manager_); + + auto encode_result = codec.Encode(state_view); + if (!encode_result) { + SNAP_LOG_ERROR( + "[Snapshot] state encoding failed, snapshot_id={}, " + "code={}, msg={}", + snapshot_id, toString(encode_result.error().code), + encode_result.error().message); + return tl::make_unexpected(encode_result.error()); + } + + SNAP_LOG_INFO("[Snapshot] state encoding successful, snapshot_id={}", + snapshot_id); + + const auto& payloads = encode_result.value(); + const auto& serialized_metadata = payloads.metadata; + const auto& serialized_segment = payloads.segments; + const auto& serialized_task_manager = payloads.task_manager; + + // When backup_dir is enabled, try all uploads to ensure complete backup + // When backup_dir is disabled, use fail-fast mode + bool upload_success = true; + std::string error_msg; + SNAP_LOG_INFO("[Snapshot] Backend info: {}", + repository_->GetObjectStoreConnectionInfo()); + + // Upload metadata + std::string metadata_path = path_prefix + ha::kSnapshotMetadataFile; + auto upload_result = repository_->UploadPayloadFile( + serialized_metadata, metadata_path, ha::kSnapshotMetadataFile, + snapshot_id); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] metadata upload failed, snapshot_id={}, " + "path={}, code={}, msg={}", + snapshot_id, metadata_path, + toString(upload_result.error().code), + upload_result.error().message); + if (!options_.use_snapshot_backup_dir) { + return tl::make_unexpected(upload_result.error()); + } + error_msg.append(upload_result.error().message + "\n"); + upload_success = false; + } + + // Upload segment + std::string segment_path = path_prefix + ha::kSnapshotSegmentsFile; + upload_result = repository_->UploadPayloadFile( + serialized_segment, segment_path, ha::kSnapshotSegmentsFile, + snapshot_id); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] segment upload failed, snapshot_id={}, " + "path={}, code={}, msg={}", + snapshot_id, segment_path, toString(upload_result.error().code), + upload_result.error().message); + if (!options_.use_snapshot_backup_dir) { + return tl::make_unexpected(upload_result.error()); + } + error_msg.append(upload_result.error().message + "\n"); + upload_success = false; + } + + // Upload task manager + std::string task_manager_path = + path_prefix + ha::kSnapshotTaskManagerFile; + upload_result = repository_->UploadPayloadFile( + serialized_task_manager, task_manager_path, + ha::kSnapshotTaskManagerFile, snapshot_id); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] task_manager upload failed, snapshot_id={}, " + "path={}, code={}, msg={}", + snapshot_id, task_manager_path, + toString(upload_result.error().code), + upload_result.error().message); + if (!options_.use_snapshot_backup_dir) { + return tl::make_unexpected(upload_result.error()); + } + error_msg.append(upload_result.error().message + "\n"); + upload_success = false; + } + + // Upload manifest + std::vector manifest_bytes = + ha::MasterSnapshotCodec::EncodeManifest( + ha::kSnapshotSerializerType, ha::kSnapshotSerializerVersion, + snapshot_id); + upload_result = repository_->UploadPayloadFile( + manifest_bytes, manifest_path, ha::kSnapshotManifestFile, + snapshot_id); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] manifest upload failed, snapshot_id={}, " + "path={}, code={}, msg={}", + snapshot_id, manifest_path, + toString(upload_result.error().code), + upload_result.error().message); + if (!options_.use_snapshot_backup_dir) { + return tl::make_unexpected(upload_result.error()); + } + error_msg.append(upload_result.error().message + "\n"); + upload_success = false; + } + + if (!upload_success) { + return tl::make_unexpected( + SerializationError(ErrorCode::PERSISTENT_FAIL, error_msg)); + } + + // Publish snapshot catalog entry and advance the latest marker. + std::string latest_path = snapshot_catalog_store_->GetSnapshotRoot() + + ha::kSnapshotLatestFile; + std::string latest_content = snapshot_id; + + auto publish_result = repository_->PublishSnapshot(descriptor); + if (publish_result != ErrorCode::OK) { + SNAP_LOG_ERROR( + "[Snapshot] latest update failed, snapshot_id={}, file={}, " + "code={}", + snapshot_id, latest_path, toString(publish_result)); + if (options_.use_snapshot_backup_dir) { + auto save_path = fs::path(options_.snapshot_backup_dir) / + ha::kSnapshotBackupSaveDir / + ha::kSnapshotLatestFile; + auto save_result = + FileUtil::SaveStringToFile(latest_content, save_path); + if (!save_result) { + SNAP_LOG_ERROR( + "[Snapshot] save latest to disk failed, " + "snapshot_id={}, " + "content={}, file={}", + snapshot_id, latest_content, save_path.string()); + } + } + + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, + fmt::format("latest update {} failed", latest_path))); + } + SNAP_LOG_INFO( + "[Snapshot] Upload latest success: {}, snapshot_id={}, " + "content={}", + latest_path, snapshot_id, latest_content); + + repository_->CleanupOldSnapshots(options_.snapshot_retention_count, + snapshot_id); + SNAP_LOG_INFO("[Snapshot] action=persisting_state end, snapshot_id={}", + snapshot_id); + } catch (const std::exception& e) { + SNAP_LOG_ERROR( + "[Snapshot] Exception during state persistent, snapshot_id={}, " + "error={}", + snapshot_id, e.what()); + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, + fmt::format("Exception during state persistent: {}", e.what()))); + } catch (...) { + SNAP_LOG_ERROR( + "[Snapshot] Unknown exception during state persistent, " + "snapshot_id={}", + snapshot_id); + return tl::make_unexpected( + SerializationError(ErrorCode::PERSISTENT_FAIL, + "Unknown exception during state persistent")); + } + return {}; +} + +} // namespace mooncake diff --git a/mooncake-store/src/master_snapshot_repository.cpp b/mooncake-store/src/master_snapshot_repository.cpp new file mode 100644 index 0000000000..b441196bdb --- /dev/null +++ b/mooncake-store/src/master_snapshot_repository.cpp @@ -0,0 +1,344 @@ +#include "master_snapshot_repository.h" + +#include +#include + +#include + +#include "ha/snapshot/catalog/snapshot_catalog_store.h" +#include "ha/snapshot/object/snapshot_object_store.h" +#include "ha/snapshot/snapshot_constants.h" +#include "ha/snapshot/snapshot_logger.h" +#include "utils/file_util.h" + +namespace mooncake { + +namespace fs = std::filesystem; + +MasterSnapshotRepository::MasterSnapshotRepository( + SnapshotObjectStore* object_store, ha::SnapshotCatalogStore* catalog_store, + const std::string& backup_dir, bool use_backup_dir) + : object_store_(object_store), + catalog_store_(catalog_store), + backup_dir_(backup_dir), + use_backup_dir_(use_backup_dir) {} + +tl::expected +MasterSnapshotRepository::UploadPayloadFile(const std::vector& data, + const std::string& path, + const std::string& local_filename, + const std::string& snapshot_id) { + SNAP_LOG_INFO("[Snapshot] Uploading {} to: {}, snapshot_id={}", + local_filename, path, snapshot_id); + + std::string error_msg; + auto upload_result = object_store_->UploadBuffer(path, data); + if (!upload_result) { + SNAP_LOG_ERROR( + "[Snapshot] {} upload failed, snapshot_id={}, file={}, error={}", + local_filename, snapshot_id, path, upload_result.error()); + + // Upload failed, save locally for manual recovery in exception + // scenarios + if (use_backup_dir_) { + auto save_path = fs::path(backup_dir_) / + ha::kSnapshotBackupSaveDir / local_filename; + auto save_result = FileUtil::SaveBinaryToFile(data, save_path); + if (!save_result) { + SNAP_LOG_ERROR( + "[Snapshot] save {} to disk failed, snapshot_id={}, " + "file={}", + local_filename, snapshot_id, save_path.string()); + } + } + + error_msg.append(local_filename) + .append(" upload ") + .append(path) + .append(" failed; "); + return tl::make_unexpected( + SerializationError(ErrorCode::PERSISTENT_FAIL, error_msg)); + } else { + SNAP_LOG_INFO("[Snapshot] Upload {} success: {}, snapshot_id={}", + local_filename, path, snapshot_id); + } + + return {}; +} + +ErrorCode MasterSnapshotRepository::PublishSnapshot( + const ha::SnapshotDescriptor& descriptor) { + return catalog_store_->Publish(descriptor); +} + +void MasterSnapshotRepository::CleanupOldSnapshots( + size_t keep_count, const std::string& current_snapshot_id) { + if (!catalog_store_) { + SNAP_LOG_ERROR( + "[Snapshot] snapshot catalog store is not initialized, " + "snapshot_id={}", + current_snapshot_id); + return; + } + + // List() loads one descriptor per published snapshot. This remains cheap + // because CleanupOldSnapshots() itself enforces retention count + // and keeps the catalog single-digit in normal deployments. + auto list_result = catalog_store_->List(ha::kUnlimitedSnapshotList); + if (!list_result) { + SNAP_LOG_ERROR("[Snapshot] error=list failed, snapshot_id={}, code={}", + current_snapshot_id, toString(list_result.error())); + return; + } + + const auto& snapshots = list_result.value(); + + if (snapshots.size() > keep_count) { + for (size_t i = keep_count; i < snapshots.size(); i++) { + const std::string& old_state_dir = snapshots[i].snapshot_id; + + if (old_state_dir == current_snapshot_id) { + SNAP_LOG_WARN( + "[Snapshot] Skipping deletion of current snapshot " + "directory {}, " + "snapshot_id={}", + old_state_dir, current_snapshot_id); + continue; + } + + auto delete_result = catalog_store_->Delete(old_state_dir); + if (delete_result != ErrorCode::OK) { + SNAP_LOG_ERROR( + "[Snapshot] Failed to delete old snapshot {}, " + "snapshot_id={}, code={}", + old_state_dir, current_snapshot_id, + toString(delete_result)); + } else { + SNAP_LOG_INFO( + "[Snapshot] Successfully deleted old snapshot {}, " + "snapshot_id={}", + old_state_dir, current_snapshot_id); + } + } + } +} + +tl::expected, ErrorCode> +MasterSnapshotRepository::ListSnapshots(size_t limit) { + return catalog_store_->List(limit); +} + +ErrorCode MasterSnapshotRepository::DeleteSnapshot( + const ha::SnapshotId& snapshot_id) { + return catalog_store_->Delete(snapshot_id); +} + +std::string MasterSnapshotRepository::GetObjectStoreConnectionInfo() const { + return object_store_->GetConnectionInfo(); +} + +tl::expected +MasterSnapshotRepository::LoadLatestSnapshot() { + if (!catalog_store_) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + auto latest_result = catalog_store_->GetLatest(); + if (!latest_result) { + return tl::make_unexpected(latest_result.error()); + } + + if (!latest_result->has_value()) { + return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + } + + return latest_result->value(); +} + +tl::expected, ErrorCode> +MasterSnapshotRepository::LoadRestoreCandidates( + const std::optional& latest_snapshot) { + if (!catalog_store_) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + std::vector candidates; + std::unordered_set candidate_ids; + + // Add latest snapshot if provided + if (latest_snapshot.has_value()) { + candidates.push_back(latest_snapshot.value()); + candidate_ids.emplace(latest_snapshot->snapshot_id); + } + + // List all snapshots as fallback + auto list_result = ListSnapshots(ha::kUnlimitedSnapshotList); + if (!list_result) { + if (candidates.empty()) { + return tl::make_unexpected(list_result.error()); + } + // Return latest only if list failed + return candidates; + } + + // Filter by latest_id chronologically (snapshot IDs use timestamp format) + for (const auto& snapshot : list_result.value()) { + if (latest_snapshot.has_value() && + snapshot.snapshot_id > latest_snapshot->snapshot_id) { + continue; + } + if (candidate_ids.emplace(snapshot.snapshot_id).second) { + candidates.push_back(snapshot); + } + } + + return candidates; +} + +tl::expected +MasterSnapshotRepository::DownloadSnapshotPayloads( + const ha::SnapshotDescriptor& descriptor) { + if (!object_store_) { + return tl::make_unexpected(SerializationError(ErrorCode::INVALID_PARAMS, + "object_store is null")); + } + + const std::string& snapshot_id = descriptor.snapshot_id; + std::string path_prefix = descriptor.object_prefix; + if (path_prefix.empty()) { + if (!catalog_store_) { + return tl::make_unexpected(SerializationError( + ErrorCode::INVALID_PARAMS, "catalog_store is null")); + } + path_prefix = catalog_store_->GetSnapshotRoot() + snapshot_id + "/"; + } + + std::string manifest_path = descriptor.manifest_key; + if (manifest_path.empty()) { + manifest_path = path_prefix + ha::kSnapshotManifestFile; + } + + // Download and validate manifest + std::string manifest_content; + auto manifest_result = + object_store_->DownloadString(manifest_path, manifest_content); + if (!manifest_result) { + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, "failed to download manifest '" + + manifest_path + + "': " + manifest_result.error())); + } + + if (use_backup_dir_) { + auto save_result = FileUtil::SaveStringToFile( + manifest_content, fs::path(backup_dir_) / + ha::kSnapshotBackupRestoreDir / + ha::kSnapshotManifestFile); + if (!save_result) { + SNAP_LOG_ERROR("[Restore] Failed to save manifest to file: {}", + save_result.error()); + } + } + + // Parse and validate manifest + std::vector parts; + boost::split(parts, manifest_content, boost::is_any_of("|")); + if (parts.size() < 3) { + return tl::make_unexpected(SerializationError( + ErrorCode::INVALID_PARAMS, "invalid snapshot manifest format")); + } + + const std::string& protocol_type = parts[0]; + const std::string& version = parts[1]; + + SNAP_LOG_INFO("[Restore] Loading snapshot: {} version: {} protocol: {}", + snapshot_id, version, protocol_type); + + if (protocol_type != ha::kSnapshotSerializerType) { + return tl::make_unexpected(SerializationError( + ErrorCode::INVALID_PARAMS, "unsupported protocol type '" + + protocol_type + "', expected '" + + ha::kSnapshotSerializerType + "'")); + } + if (version != ha::kSnapshotSerializerVersion) { + return tl::make_unexpected(SerializationError( + ErrorCode::INVALID_PARAMS, + "incompatible snapshot version '" + version + "', expected '" + + ha::kSnapshotSerializerVersion + "'")); + } + + ha::MasterSnapshotPayloads payloads; + + // Download metadata + std::string metadata_path = path_prefix + ha::kSnapshotMetadataFile; + auto metadata_result = + object_store_->DownloadBuffer(metadata_path, payloads.metadata); + if (!metadata_result) { + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, "failed to download metadata '" + + metadata_path + + "': " + metadata_result.error())); + } + + if (use_backup_dir_) { + auto save_result = FileUtil::SaveBinaryToFile( + payloads.metadata, fs::path(backup_dir_) / + ha::kSnapshotBackupRestoreDir / + ha::kSnapshotMetadataFile); + if (!save_result) { + SNAP_LOG_ERROR("[Restore] Failed to save metadata to file: {}", + save_result.error()); + } + } + SNAP_LOG_INFO("[Restore] Downloaded metadata file successfully"); + + // Download segments + std::string segments_path = path_prefix + ha::kSnapshotSegmentsFile; + auto segments_result = + object_store_->DownloadBuffer(segments_path, payloads.segments); + if (!segments_result) { + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, "failed to download segments '" + + segments_path + + "': " + segments_result.error())); + } + + if (use_backup_dir_) { + auto save_result = FileUtil::SaveBinaryToFile( + payloads.segments, fs::path(backup_dir_) / + ha::kSnapshotBackupRestoreDir / + ha::kSnapshotSegmentsFile); + if (!save_result) { + SNAP_LOG_ERROR("[Restore] Failed to save segments to file: {}", + save_result.error()); + } + } + SNAP_LOG_INFO("[Restore] Downloaded segments file successfully"); + + // Download task_manager + std::string task_manager_path = path_prefix + ha::kSnapshotTaskManagerFile; + auto task_manager_result = + object_store_->DownloadBuffer(task_manager_path, payloads.task_manager); + if (!task_manager_result) { + return tl::make_unexpected(SerializationError( + ErrorCode::PERSISTENT_FAIL, + "failed to download task_manager '" + task_manager_path + + "': " + task_manager_result.error())); + } + + if (use_backup_dir_) { + auto save_result = FileUtil::SaveBinaryToFile( + payloads.task_manager, fs::path(backup_dir_) / + ha::kSnapshotBackupRestoreDir / + ha::kSnapshotTaskManagerFile); + if (!save_result) { + SNAP_LOG_ERROR("[Restore] Failed to save task manager to file: {}", + save_result.error()); + } + } + SNAP_LOG_INFO("[Restore] Downloaded task manager file successfully"); + + return payloads; +} + +} // namespace mooncake diff --git a/mooncake-store/src/metadata_store.cpp b/mooncake-store/src/metadata_store.cpp new file mode 100644 index 0000000000..2c653c6aca --- /dev/null +++ b/mooncake-store/src/metadata_store.cpp @@ -0,0 +1,53 @@ +#include "metadata_store.h" + +namespace mooncake { + +void StandbySegmentRegistry::OnSegmentMount(const StandbySegmentInfo& info) { + std::lock_guard lock(mutex_); + segments_by_endpoint_[info.transport_endpoint] = info; +} + +void StandbySegmentRegistry::OnSegmentUnmount( + const std::string& transport_endpoint) { + std::lock_guard lock(mutex_); + segments_by_endpoint_.erase(transport_endpoint); +} + +void StandbySegmentRegistry::OnSegmentUpdate(const StandbySegmentInfo& info) { + std::lock_guard lock(mutex_); + segments_by_endpoint_[info.transport_endpoint] = info; +} + +bool StandbySegmentRegistry::HasSegment( + const std::string& transport_endpoint) const { + std::shared_lock lock(mutex_); + return segments_by_endpoint_.find(transport_endpoint) != + segments_by_endpoint_.end(); +} + +std::optional StandbySegmentRegistry::GetSegment( + const std::string& transport_endpoint) const { + std::shared_lock lock(mutex_); + auto it = segments_by_endpoint_.find(transport_endpoint); + if (it == segments_by_endpoint_.end()) { + return std::nullopt; + } + return it->second; +} + +std::vector StandbySegmentRegistry::GetAllSegments() const { + std::shared_lock lock(mutex_); + std::vector result; + result.reserve(segments_by_endpoint_.size()); + for (const auto& [endpoint, info] : segments_by_endpoint_) { + result.push_back(info); + } + return result; +} + +void StandbySegmentRegistry::Clear() { + std::lock_guard lock(mutex_); + segments_by_endpoint_.clear(); +} + +} // namespace mooncake \ No newline at end of file diff --git a/mooncake-store/src/offset_allocator.cpp b/mooncake-store/src/offset_allocator.cpp index 68b44025f6..17f2369e2e 100644 --- a/mooncake-store/src/offset_allocator.cpp +++ b/mooncake-store/src/offset_allocator.cpp @@ -1,7 +1,7 @@ // (C) Sebastian Aaltonen 2023 // MIT License (see file: LICENSE) -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include #include @@ -284,9 +284,9 @@ OffsetAllocation __Allocator::allocate(uint32 size) { return OffsetAllocation(node.dataOffset, nodeIndex); } -void __Allocator::free(OffsetAllocation allocation) { +uint32 __Allocator::free(OffsetAllocation allocation) { ASSERT(allocation.metadata != OffsetAllocation::NO_SPACE); - if (m_nodes.empty()) return; + if (m_nodes.empty()) return 0; uint32 nodeIndex = allocation.metadata; Node& node = m_nodes[nodeIndex]; @@ -348,6 +348,8 @@ void __Allocator::free(OffsetAllocation allocation) { m_nodes[combinedNodeIndex].neighborPrev = neighborPrev; m_nodes[neighborPrev].neighborNext = combinedNodeIndex; } + + return m_freeOffset < m_max_capacity ? size : 0; } uint32 __Allocator::insertNodeIntoBin(uint32 size, uint32 dataOffset) { @@ -555,6 +557,9 @@ OffsetAllocator::OffsetAllocator(uint64_t base, size_t size, m_capacity(size) { m_allocator = std::make_unique<__Allocator>(size >> m_multiplier_bits, init_capacity, max_capacity); + m_largest_free_region.store( + m_allocator->storageReport().largestFreeRegion << m_multiplier_bits, + std::memory_order_relaxed); } OffsetAllocator::OffsetAllocator(uint64_t base, size_t size, @@ -563,19 +568,29 @@ OffsetAllocator::OffsetAllocator(uint64_t base, size_t size, : m_allocator(std::move(allocator)), m_base(base), m_multiplier_bits(multiplier_bits), - m_capacity(size) {} + m_capacity(size) { + const uint64_t largest_free_region = + m_allocator->storageReport().largestFreeRegion << m_multiplier_bits; + m_largest_free_region.store(largest_free_region, std::memory_order_relaxed); + const uint64_t allocator_capacity = + static_cast(m_allocator->m_size) << m_multiplier_bits; + m_largest_free_region_tightened = largest_free_region < allocator_capacity; +} std::optional OffsetAllocator::allocate(size_t size) { if (size == 0) { return std::nullopt; } - MutexLocker guard(&m_mutex); - if (!m_allocator) { + // Free regions are grouped into size bins. A request larger than the + // highest free-bin boundary rounds up to a higher bin and cannot fit. The + // cached value is only a fast-fail hint: a stale larger value merely falls + // through to the mutex-protected allocator. + if (size > getLargestFreeRegion()) { return std::nullopt; } - size_t fake_size = + const size_t fake_size = m_multiplier_bits > 0 ? ((size + (static_cast(1) << m_multiplier_bits) - 1u) >> m_multiplier_bits) @@ -585,13 +600,23 @@ std::optional OffsetAllocator::allocate(size_t size) { return std::nullopt; } + MutexLocker guard(&m_mutex); + if (!m_allocator) { + return std::nullopt; + } + OffsetAllocation allocation = m_allocator->allocate(fake_size); if (allocation.isNoSpace()) { - // Log metrics to help understand why allocation failed - // Note: We're already holding m_mutex, so use internal method - OffsetAllocatorMetrics metrics = get_metrics_internal(); - VLOG(1) << "OffsetAllocator allocation failed: size=" << size - << ", fake_size=" << fake_size << ", " << metrics; + // A request can pass the conservative hint but still fail because + // the allocator state changed concurrently. Tighten the hint after + // observing the authoritative state under the mutex. + refreshLargestFreeRegion(); + if (VLOG_IS_ON(1)) { + // We're already holding m_mutex, so use the internal method. + const OffsetAllocatorMetrics metrics = get_metrics_internal(); + VLOG(1) << "OffsetAllocator allocation failed: size=" << size + << ", fake_size=" << fake_size << ", " << metrics; + } return std::nullopt; } @@ -599,10 +624,31 @@ std::optional OffsetAllocator::allocate(size_t size) { m_allocated_size += size; m_allocated_num++; - // Use shared_from_this to get a shared_ptr to this OffsetAllocator - return OffsetAllocationHandle( - shared_from_this(), allocation, - m_base + (allocation.getOffset() << m_multiplier_bits), size); + const uint64_t real_base = + m_base + (allocation.getOffset() << m_multiplier_bits); + guard.unlock(); + + // Handle construction and shared_ptr reference-counting do not access + // allocator state and should not extend the serialized critical section. + return OffsetAllocationHandle(shared_from_this(), allocation, real_base, + size); +} + +uint64_t OffsetAllocator::normalizedAllocationSize(size_t size) const { + if (size == 0) { + return 0; + } + const uint64_t quantum = uint64_t{1} << m_multiplier_bits; + if (size > std::numeric_limits::max() - (quantum - 1)) { + return 0; + } + const uint64_t fake_size = (size + quantum - 1) >> m_multiplier_bits; + if (fake_size > SmallFloat::MAX_BIN_SIZE) { + return 0; + } + return static_cast(SmallFloat::floatToUint( + SmallFloat::uintToFloatRoundUp(static_cast(fake_size)))) + << m_multiplier_bits; } OffsetAllocStorageReport OffsetAllocator::storageReport() const { @@ -652,14 +698,55 @@ OffsetAllocatorMetrics OffsetAllocator::get_metrics() const { return get_metrics_internal(); } +void OffsetAllocator::refreshLargestFreeRegion() { + const uint64_t largest_free_region = + m_allocator ? m_allocator->storageReport().largestFreeRegion + << m_multiplier_bits + : 0; + m_largest_free_region.store(largest_free_region, std::memory_order_relaxed); + const uint64_t allocator_capacity = + m_allocator + ? static_cast(m_allocator->m_size) << m_multiplier_bits + : 0; + m_largest_free_region_tightened = + m_allocator && largest_free_region < allocator_capacity; +} + void OffsetAllocator::freeAllocation(const OffsetAllocation& allocation, uint64_t size) { MutexLocker lock(&m_mutex); if (m_allocator) { - m_allocator->free(allocation); - // Update lightweight metrics - m_allocated_size -= size; - m_allocated_num--; + const uint64_t freed_region = + static_cast(m_allocator->free(allocation)) + << m_multiplier_bits; + if (m_largest_free_region_tightened) { + // Before free, the hint is an upper bound for every existing free + // region. Only the newly merged region can raise that bound. + const uint64_t current_hint = + m_largest_free_region.load(std::memory_order_relaxed); + if (freed_region > current_hint) { + m_largest_free_region.store(freed_region, + std::memory_order_relaxed); + } + const uint64_t allocator_capacity = + static_cast(m_allocator->m_size) << m_multiplier_bits; + if (freed_region >= allocator_capacity) { + m_largest_free_region_tightened = false; + } + } + // Update lightweight metrics. Saturate instead of wrapping: + // recovery may free nodes whose exact requested size is unknown + // (corrupt record), and an unsigned underflow would poison the + // metric permanently. + if (size > m_allocated_size) { + LOG(WARNING) << "freeAllocation: size " << size + << " exceeds allocated_size " << m_allocated_size + << " -- clamping to 0"; + m_allocated_size = 0; + } else { + m_allocated_size -= size; + } + if (m_allocated_num > 0) m_allocated_num--; } } @@ -683,4 +770,31 @@ std::ostream& operator<<(std::ostream& os, return os; } -} // namespace mooncake::offset_allocator \ No newline at end of file +// ============================================================================ +// Recovery helpers +// ============================================================================ + +std::optional OffsetAllocator::createHandleAtNode( + uint32_t node_index, uint64_t real_offset, uint64_t requested_size) { + MutexLocker guard(&m_mutex); + if (!m_allocator || node_index >= m_allocator->m_current_capacity) + return std::nullopt; + const auto& node = m_allocator->m_nodes[node_index]; + if (!node.used) return std::nullopt; + + // Cross-validate: real_offset must match the node's stored offset. + uint64_t expected_offset = + m_base + (static_cast(node.dataOffset) << m_multiplier_bits); + if (expected_offset != real_offset) { + LOG(ERROR) << "node/offset mismatch: node_index=" << node_index + << " expected_offset=" << expected_offset + << " real_offset=" << real_offset; + return std::nullopt; + } + + OffsetAllocation allocation(node.dataOffset, node_index); + return OffsetAllocationHandle(shared_from_this(), allocation, real_offset, + requested_size); +} + +} // namespace mooncake::offset_allocator diff --git a/mooncake-store/src/posix_file.cpp b/mooncake-store/src/posix_file.cpp index e5f0158276..9f16b4fa0d 100644 --- a/mooncake-store/src/posix_file.cpp +++ b/mooncake-store/src/posix_file.cpp @@ -14,6 +14,14 @@ PosixFile::PosixFile(const std::string &filename, int fd) } } +tl::expected PosixFile::datasync() { + if (fdatasync(fd_) != 0) { + LOG(ERROR) << "fdatasync failed: " << strerror(errno); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + return {}; +} + PosixFile::~PosixFile() { if (fd_ >= 0) { if (close(fd_) != 0) { @@ -21,7 +29,8 @@ PosixFile::~PosixFile() { } // If the file was opened with an error code indicating a write failure, // attempt to delete the file to prevent corruption. - if (error_code_ == ErrorCode::FILE_WRITE_FAIL) { + if (delete_on_write_fail_ && + error_code_ == ErrorCode::FILE_WRITE_FAIL) { if (::unlink(filename_.c_str()) == -1) { LOG(ERROR) << "Failed to delete corrupted file: " << filename_; } else { diff --git a/mooncake-store/src/real_client.cpp b/mooncake-store/src/real_client.cpp index 81fd32d8e6..3f5eb35fbd 100644 --- a/mooncake-store/src/real_client.cpp +++ b/mooncake-store/src/real_client.cpp @@ -1,6 +1,3 @@ -#include -#include -#include #include #include #include @@ -14,23 +11,30 @@ #include // for dlsym (Python detection) #include // for atexit #include -#include #include #include #include #include #include "real_client.h" -#include "client_buffer.hpp" +#include "registered_pinned_memory.h" +#include "client_buffer.h" +#include "replica_selection.h" #include "common.h" #include "config.h" +#include "store_rpc_client_io_context.h" +#include "bool_parser.h" +#include "environ.h" +#include "integer_parser.h" #include "mutex.h" #include "types.h" #include "utils.h" #include "rpc_types.h" #include "file_storage.h" -#include "gpu_staging_utils.h" +#include "device/accelerator_registry.h" #include "default_config.h" +#include "uds_transport.h" +#include "device/cuda_ipc_buffer.h" #include "shm_helper.h" #include "memory_location.h" #ifdef USE_NOF @@ -40,6 +44,15 @@ #include "acl/acl_rt.h" #include "transport/ascend_transport/ascend_direct_transport/context_manager.h" #endif +#ifdef USE_CUDA +#include +#endif +#ifdef USE_INTRA_NVLINK +#include "gpu_vendor/intra_nvlink.h" +#endif +#if defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) +#include "ascend_allocator.h" +#endif DEFINE_bool(enable_http_server, false, "Enable embedded HTTP server for health check and metrics."); @@ -49,6 +62,48 @@ DEFINE_int32(http_port, 9300, namespace mooncake { namespace { +constexpr std::chrono::seconds kIpcRequestRecvTimeout{5}; + +bool IsHostStoreSegmentProtocol(const std::string &protocol) { + return protocol.empty() || protocol == "tcp" || protocol == "rdma" || + protocol == "efa" || protocol == "cxi" || protocol == "rpc_only"; +} + +std::shared_ptr TryPinStoreSegment( + void *ptr, size_t size, const std::string &protocol, + const char *segment_owner) { + if (!IsHostStoreSegmentProtocol(protocol)) return nullptr; + return RegisteredPinnedMemoryManager::instance().try_pin( + ptr, size, + std::string("Store segment ") + segment_owner + + " protocol=" + protocol); +} + +bool ReleasePinnedRegionForFree( + std::shared_ptr &pinned_region, + const char *segment_owner) { + if (!pinned_region) return true; + const bool safe_to_free = pinned_region->release(); + pinned_region.reset(); + if (!safe_to_free) { + LOG(ERROR) << "Leaking " << segment_owner + << " backing memory because cudaHostUnregister failed"; + } + return safe_to_free; +} + +bool ReleasePinnedRegionsForFree( + std::vector> &pinned_regions, + const char *segment_owner) { + bool safe_to_free = true; + for (auto &pinned_region : pinned_regions) { + safe_to_free &= + ReleasePinnedRegionForFree(pinned_region, segment_owner); + } + pinned_regions.clear(); + return safe_to_free; +} + #ifdef USE_ASCEND_DIRECT bool checkAcl(aclError result, const char *message) { if (result != ACL_ERROR_NONE) { @@ -80,15 +135,6 @@ tl::expected set_context_if_needed(const std::string &protocol, } #endif -struct PreparedRangedReadRequest { - std::vector>>> - results; - std::vector>> valid_fragments; - std::vector required_buffer_sizes; - bool top_level_valid = true; - bool has_any_valid_fragment = false; -}; - size_t sum_value_sizes(const std::vector> &values) { size_t total = 0; for (const auto &value : values) { @@ -170,97 +216,9 @@ size_t sum_buffer_handle_sizes( return total; } -PreparedRangedReadRequest prepare_ranged_read_request( - size_t buffer_count, const std::vector> &all_keys, - const std::vector>> &all_dst_offsets, - const std::vector>> &all_src_offsets, - const std::vector>> &all_sizes, - const char *log_prefix) { - PreparedRangedReadRequest prepared; - prepared.results.resize(buffer_count); - prepared.valid_fragments.resize(buffer_count); - prepared.required_buffer_sizes.resize(buffer_count, 0); - - if (buffer_count != all_keys.size() || - buffer_count != all_dst_offsets.size() || - buffer_count != all_src_offsets.size() || - buffer_count != all_sizes.size()) { - LOG(ERROR) << log_prefix << ": top-level size mismatch"; - prepared.results = build_ranged_read_internal_error_results( - buffer_count, all_keys, all_dst_offsets, ErrorCode::INVALID_PARAMS); - prepared.top_level_valid = false; - return prepared; - } - - for (size_t i = 0; i < buffer_count; ++i) { - const size_t key_count = all_keys[i].size(); - prepared.results[i].resize(key_count); - prepared.valid_fragments[i].resize(key_count); - - if (key_count != all_dst_offsets[i].size() || - key_count != all_src_offsets[i].size() || - key_count != all_sizes[i].size()) { - LOG(ERROR) << log_prefix - << ": key-group size mismatch for buffer index " << i; - for (size_t j = 0; j < key_count; ++j) { - prepared.results[i][j] = - std::vector>( - 1, tl::unexpected(ErrorCode::INVALID_PARAMS)); - prepared.valid_fragments[i][j] = std::vector(1, false); - } - continue; - } - - size_t max_required = 0; - for (size_t j = 0; j < key_count; ++j) { - const size_t fragment_count = all_dst_offsets[i][j].size(); - prepared.results[i][j] = - std::vector>( - fragment_count, tl::unexpected(ErrorCode::INVALID_PARAMS)); - prepared.valid_fragments[i][j] = - std::vector(fragment_count, false); - - if (fragment_count != all_src_offsets[i][j].size() || - fragment_count != all_sizes[i][j].size()) { - LOG(ERROR) << log_prefix << ": fragment size mismatch, " - << "buffer_index=" << i << " key_index=" << j; - continue; - } - - for (size_t k = 0; k < fragment_count; ++k) { - const size_t dst_offset = all_dst_offsets[i][j][k]; - const size_t fragment_size = all_sizes[i][j][k]; - if (dst_offset > - std::numeric_limits::max() - fragment_size) { - LOG(ERROR) - << log_prefix - << ": destination range overflow, buffer_index=" << i - << " key_index=" << j << " fragment_index=" << k; - continue; - } - prepared.valid_fragments[i][j][k] = true; - prepared.has_any_valid_fragment = true; - max_required = - std::max(max_required, dst_offset + fragment_size); - } - } - prepared.required_buffer_sizes[i] = max_required; - } - - return prepared; -} - -void fill_ranged_read_results_with_error( - std::vector>>> - &results, - ErrorCode error) { - for (auto &key_rows : results) { - for (auto &row : key_rows) { - for (auto &fragment : row) { - fragment = tl::unexpected(error); - } - } - } +ErrorCode scatter_transfer_error(const Status &status) { + return status.IsInvalidArgument() ? ErrorCode::INVALID_PARAMS + : ErrorCode::TRANSFER_FAIL; } // Scatter host (CPU) memory to a destination that may be GPU or host. @@ -268,69 +226,40 @@ void fill_ranged_read_results_with_error( // tl::expected. inline tl::expected scatter_host_to_maybe_device( void *dst, const void *src, size_t size, const std::string &context) { - int device_id = -1; - if (gpu_staging::IsDevicePointer(dst, &device_id)) { - gpu_staging::SetDevice(device_id); - if (!gpu_staging::CopyHostToDevice(dst, src, size)) { - LOG(ERROR) << "H2D copy failed: " << context; - return tl::unexpected(ErrorCode::TRANSFER_FAIL); - } - } else if (gpu_staging::IsHostPointer(dst)) { - memcpy(dst, src, size); - } else { - LOG(ERROR) << "Unknown memory type for dst buffer: " << context; - return tl::unexpected(ErrorCode::INVALID_PARAMS); + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); + if (!runtime_accelerator.CopyFromHost(dst, src, size)) { + LOG(ERROR) << "H2D copy failed: " << context; + return tl::unexpected(ErrorCode::TRANSFER_FAIL); } return {}; } -// Select the best replica from a list: prefer local MEMORY, then any -// MEMORY, then LOCAL_DISK, then DISK. Master may return replicas in any -// order, so we always scan. -inline const Replica::Descriptor *SelectBestReplica( - const std::vector &replicas, - const std::unordered_set &local_endpoints) { - const Replica::Descriptor *first_memory = nullptr; - const Replica::Descriptor *first_nof = nullptr; - for (const auto &r : replicas) { - if (r.status != ReplicaStatus::COMPLETE) continue; - if (r.is_memory_replica()) { - if (local_endpoints.count( - r.get_memory_descriptor() - .buffer_descriptor.transport_endpoint_)) { - return &r; // local MEMORY — best case - } - if (!first_memory) first_memory = &r; - } else if (r.is_nof_replica()) { - if (local_endpoints.count( - r.get_nof_descriptor() - .buffer_descriptor.transport_endpoint_)) { - return &r; // local NOF_SSD — also good - } - if (!first_nof) first_nof = &r; - } - } - if (first_memory) return first_memory; - if (first_nof) return first_nof; - - const Replica::Descriptor *best = nullptr; - for (const auto &r : replicas) { - if (r.status != ReplicaStatus::COMPLETE) continue; - if (r.is_local_disk_replica()) { - best = &r; // LOCAL_DISK always overrides DISK - } else if (r.is_disk_replica() && !best) { - best = &r; - } +// Gather memory that may be GPU or host into a host destination. +inline tl::expected gather_maybe_device_to_host( + void *dst, const void *src, size_t size, const std::string &context) { + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); + if (!runtime_accelerator.CopyToHost(dst, src, size)) { + LOG(ERROR) << "D2H copy failed: " << context; + return tl::unexpected(ErrorCode::TRANSFER_FAIL); } - return best; + return {}; } +// SelectBestReplica and the replica-scoring helpers live in +// replica_selection.h (included above) so they can be unit-tested directly. +using mooncake::SelectBestReplica; + // Build a QueryResult containing only the chosen replica so that // Client::Get / Client::BatchGet (which internally call // FindFirstCompleteReplica) cannot pick a different replica type. inline QueryResult FilterQueryResult(const QueryResult &qr, - const Replica::Descriptor &replica) { - return QueryResult({replica}, qr.lease_timeout); + const Replica::Descriptor &replica, + bool include_object_checksum = true) { + return QueryResult( + {replica}, qr.lease_timeout, + include_object_checksum ? qr.object_checksum : std::nullopt); } } // namespace @@ -359,14 +288,23 @@ bool RealClient::map_dummy_range_in_shm(const MappedShm &shm, bool RealClient::map_dummy_buffer_to_real(const ShmContext &shm_ctx, uint64_t dummy_addr, size_t buf_size, const MappedShm *&last_hit_shm, - void *&out_real) const { + void *&out_real, + size_t *out_capacity) const { if (last_hit_shm && map_dummy_range_in_shm(*last_hit_shm, dummy_addr, 0, buf_size, out_real)) { + if (out_capacity) { + *out_capacity = last_hit_shm->shm_size - + (dummy_addr - last_hit_shm->dummy_base_addr); + } return true; } for (const auto &shm : shm_ctx.mapped_shms) { if (map_dummy_range_in_shm(shm, dummy_addr, 0, buf_size, out_real)) { last_hit_shm = &shm; + if (out_capacity) { + *out_capacity = + shm.shm_size - (dummy_addr - shm.dummy_base_addr); + } return true; } } @@ -389,7 +327,8 @@ bool RealClient::map_dummy_buffer_range_to_real(const ShmContext &shm_ctx, tl::expected, ErrorCode> RealClient::map_dummy_addrs_to_real_ptrs( const ShmContext &context, const std::vector &dummy_addrs, - const std::vector &sizes, const UUID &client_id) const { + const std::vector &sizes, const UUID &client_id, + std::vector *capacities) const { if (dummy_addrs.size() != sizes.size()) { LOG(ERROR) << "Mismatched dummy_addrs and sizes, client_id=" << client_id; @@ -397,11 +336,17 @@ RealClient::map_dummy_addrs_to_real_ptrs( } std::vector buffers; buffers.reserve(dummy_addrs.size()); + if (capacities) { + capacities->clear(); + capacities->reserve(dummy_addrs.size()); + } const MappedShm *last_hit_shm = nullptr; for (size_t i = 0; i < dummy_addrs.size(); ++i) { void *real_ptr = nullptr; + size_t capacity = 0; if (!map_dummy_buffer_to_real(context, dummy_addrs[i], sizes[i], - last_hit_shm, real_ptr)) { + last_hit_shm, real_ptr, + capacities ? &capacity : nullptr)) { LOG(ERROR) << "Dummy buffer at " << dummy_addrs[i] << " (size " << sizes[i] << ") not found in any mapped shared memory, client_id=" @@ -409,6 +354,7 @@ RealClient::map_dummy_addrs_to_real_ptrs( return tl::unexpected(ErrorCode::INVALID_PARAMS); } buffers.push_back(real_ptr); + if (capacities) capacities->push_back(capacity); } return buffers; } @@ -633,7 +579,8 @@ tl::expected RealClient::setup_internal( const std::shared_ptr &transfer_engine, const std::string &ipc_socket_path, int local_rpc_port, bool enable_ssd_offload, bool start_offload_rpc_server, - const std::string &ssd_offload_path, const std::string &tenant_id) { + const std::string &ssd_offload_path, const std::string &tenant_id, + bool enable_client_http_server, int client_http_port) { this->protocol = protocol; this->ipc_socket_path_ = ipc_socket_path; const bool should_use_hugepage = @@ -657,8 +604,9 @@ tl::expected RealClient::setup_internal( #endif std::optional device_name = - (rdma_devices.empty() ? std::nullopt - : std::make_optional(rdma_devices)); + ((rdma_devices.empty() || rdma_devices == "auto-discovery") + ? std::nullopt + : std::make_optional(rdma_devices)); // Validate required parameters if (local_hostname.empty()) { @@ -687,9 +635,11 @@ tl::expected RealClient::setup_internal( } else { // Auto port binding with retry on metadata registration failure const int kMaxRetries = - GetEnvOr("MC_STORE_CLIENT_SETUP_RETRIES", 20); - const int rawMinPort = GetEnvOr("MC_STORE_CLIENT_MIN_PORT", 12300); - const int rawMaxPort = GetEnvOr("MC_STORE_CLIENT_MAX_PORT", 14300); + Environ::GetInt("MC_STORE_CLIENT_SETUP_RETRIES", 20); + const int rawMinPort = + Environ::GetInt("MC_STORE_CLIENT_MIN_PORT", 12300); + const int rawMaxPort = + Environ::GetInt("MC_STORE_CLIENT_MAX_PORT", 14300); constexpr int kDefaultMinPort = 12300; constexpr int kDefaultMaxPort = 14300; auto [minPort, maxPort] = ValidatePortRange( @@ -780,10 +730,10 @@ tl::expected RealClient::setup_internal( size_t cxl_dev_size = 0; const char *env = std::getenv("MC_CXL_DEV_SIZE"); if (env) { - char *end = nullptr; - unsigned long long val = strtoull(env, &end, 10); - if (end != env && *end == '\0') - cxl_dev_size = static_cast(val); + cxl_dev_size = + TryParseInteger(env, {.trim_ascii_whitespace = true, + .allow_leading_plus = true}) + .value_or(0); } else { LOG(FATAL) << "MC_CXL_DEV_SIZE not set"; return tl::unexpected(ErrorCode::INVALID_PARAMS); @@ -822,12 +772,12 @@ tl::expected RealClient::setup_internal( } } + const bool parallel_hugetlb_population = + protocol == "rdma" && should_use_hugepage; + while (global_segment_size > 0) { size_t segment_size = std::min(global_segment_size, max_mr_size); global_segment_size -= segment_size; - current_glbseg_size += segment_size; - LOG(INFO) << "Mounting segment: " << segment_size << " bytes, " - << current_glbseg_size << " of " << total_glbseg_size; size_t mapped_size = segment_size; void *ptr = nullptr; @@ -847,7 +797,18 @@ tl::expected RealClient::setup_internal( mapped_size = align_up(segment_size, get_hugepage_size_from_env()); ptr = allocate_buffer_mmap_memory(mapped_size, - get_hugepage_size_from_env()); + get_hugepage_size_from_env(), + parallel_hugetlb_population); +#if defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) + } else if ((protocol == "ascend" || protocol == "ubshmem") && + globalConfig().ascend_use_fabric_mem) { + size_t actual_size = 0; + ptr = ascend_allocate_memory_best_effort( + segment_size, this->protocol, &actual_size); + if (ptr) { + mapped_size = actual_size; + } +#endif } else { ptr = allocate_buffer_allocator_memory(segment_size, this->protocol); @@ -857,9 +818,22 @@ tl::expected RealClient::setup_internal( LOG(ERROR) << "Failed to allocate segment memory"; return tl::unexpected(ErrorCode::INVALID_PARAMS); } + current_glbseg_size += mapped_size; + LOG(INFO) << "Mounting segment: " << mapped_size << " bytes, " + << current_glbseg_size << " of " << total_glbseg_size; + if (this->protocol == "ascend" || this->protocol == "ubshmem") { ascend_segment_ptrs_.emplace_back( ptr, AscendSegmentDeleter{this->protocol}); + } else if (this->protocol == "sunrise_link") { +#if defined(USE_SUNRISE) + sunrise_segment_ptrs_.emplace_back(ptr, + SunriseSegmentDeleter{}); +#else + LOG(ERROR) + << "sunrise_link protocol requires USE_SUNRISE build"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); +#endif } else if (this->protocol == "ub") { ub_segment_ptrs_.emplace_back(ptr, UbSegmentDeleter{mapped_size}); @@ -869,15 +843,42 @@ tl::expected RealClient::setup_internal( hugepage_segment_ptrs_.emplace_back( ptr, HugepageSegmentDeleter{mapped_size}); } else { +#ifdef USE_VRAM_SEGMENT + vram_segment_ptrs_.emplace_back(ptr); +#else segment_ptrs_.emplace_back(ptr); +#endif } + + // Populate HugeTLB pages in parallel immediately before transfer- + // engine registration. NUMA mappings use node-local workers for + // each mbind region; direct mappings use the generic worker pool. + if (parallel_hugetlb_population) { + if (!seg_numa_nodes.empty()) { + populate_hugetlb_numa_mapping(ptr, mapped_size, + seg_numa_nodes); + } else if (!is_mmap_arena_allocation(ptr)) { + populate_hugetlb_mapping(ptr, mapped_size); + } + } + + auto pinned_region = + TryPinStoreSegment(ptr, mapped_size, this->protocol, "setup"); auto mount_result = client_->MountSegment(ptr, mapped_size, protocol, seg_location); if (!mount_result.has_value()) { + if (!ReleasePinnedRegionForFree(pinned_region, + "Store setup segment")) { + setup_segment_memory_must_leak_ = true; + } LOG(ERROR) << "Failed to mount segment: " << toString(mount_result.error()); return tl::unexpected(mount_result.error()); } + if (pinned_region) { + setup_segment_pinned_regions_.push_back( + std::move(pinned_region)); + } } if (total_glbseg_size == 0) { LOG(INFO) << "Global segment size is 0, skip mounting segment"; @@ -932,10 +933,20 @@ tl::expected RealClient::setup_internal( } } client_requester_ = std::make_shared(); - if (FLAGS_enable_http_server) { - if (start_http_server() != 0) { - LOG(ERROR) << "Failed to start HTTP server on port " - << FLAGS_http_port; + const bool should_start_http_server = + enable_client_http_server || FLAGS_enable_http_server; + const int selected_http_port = + enable_client_http_server ? client_http_port : FLAGS_http_port; + if (should_start_http_server) { + if (selected_http_port <= 0 || selected_http_port > 65535) { + LOG(ERROR) << "Invalid client HTTP server port: " + << selected_http_port; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + if (start_http_server(selected_http_port) != 0) { + LOG(WARNING) << "Failed to start client HTTP server on port " + << selected_http_port + << "; continuing without HTTP endpoints"; } } @@ -949,12 +960,13 @@ int RealClient::setup_real( const std::string &master_server_addr, const std::shared_ptr &transfer_engine, const std::string &ipc_socket_path, bool enable_ssd_offload, - const std::string &ssd_offload_path, const std::string &tenant_id) { + const std::string &ssd_offload_path, const std::string &tenant_id, + bool enable_client_http_server, int client_http_port) { return to_py_ret(setup_internal( local_hostname, metadata_server, global_segment_size, local_buffer_size, protocol, rdma_devices, master_server_addr, transfer_engine, ipc_socket_path, 50052, enable_ssd_offload, true, ssd_offload_path, - tenant_id)); + tenant_id, enable_client_http_server, client_http_port)); } namespace { @@ -981,6 +993,41 @@ inline std::optional get_config_size(const ConfigDict &config, } return static_cast(parsed_size_opt.value()); } + +inline bool get_config_bool(const ConfigDict &config, const std::string &key, + bool default_value) { + auto it = config.find(key); + if (it == config.end()) { + return default_value; + } + + const auto parsed = TryParseBool(it->second); + if (parsed.has_value()) { + return *parsed; + } + + LOG(WARNING) << "Invalid boolean value for config key '" << key + << "': " << it->second << ", using default: " << default_value; + return default_value; +} + +inline std::optional get_config_int(const ConfigDict &config, + const std::string &key, + int default_value) { + auto it = config.find(key); + if (it == config.end()) { + return default_value; + } + + const auto parsed_value = + TryParseInteger(it->second, {.trim_ascii_whitespace = true}); + if (!parsed_value.has_value()) { + LOG(ERROR) << "Invalid integer value for config key '" << key + << "': " << it->second; + return std::nullopt; + } + return *parsed_value; +} } // namespace tl::expected RealClient::setup_internal( @@ -1022,7 +1069,17 @@ tl::expected RealClient::setup_internal( get_config(config, CONFIG_KEY_IPC_SOCKET_PATH); // A size of 0 keeps the pure client/server setup semantics. - auto validate_size = [](const char *key, size_t value) { + // global_segment_size is a total capacity and may exceed max_mr_size; the + // setup path splits it into mountable chunks below. + auto validate_min_size = [](const char *key, size_t value) { + if (value != 0 && value < MIN_SEGMENT_SIZE) { + LOG(ERROR) << "Invalid " << key << ": " << value + << ", must be 0 or at least " << MIN_SEGMENT_SIZE; + return false; + } + return true; + }; + auto validate_single_segment_size = [](const char *key, size_t value) { if ((value != 0 && value < MIN_SEGMENT_SIZE) || value > MAX_SEGMENT_SIZE) { LOG(ERROR) << "Invalid " << key << ": " << value @@ -1032,33 +1089,31 @@ tl::expected RealClient::setup_internal( } return true; }; - if (!validate_size(CONFIG_KEY_GLOBAL_SEGMENT_SIZE, global_segment_size) || - !validate_size(CONFIG_KEY_LOCAL_BUFFER_SIZE, local_buffer_size)) { - return tl::unexpected(ErrorCode::INVALID_PARAMS); - } - - // Validate protocol is supported - if (protocol != "tcp" && protocol != "rdma") { - LOG(ERROR) << "Invalid " << CONFIG_KEY_PROTOCOL << ": " << protocol - << ", must be 'tcp' or 'rdma'"; + if (!validate_min_size(CONFIG_KEY_GLOBAL_SEGMENT_SIZE, + global_segment_size) || + !validate_single_segment_size(CONFIG_KEY_LOCAL_BUFFER_SIZE, + local_buffer_size)) { return tl::unexpected(ErrorCode::INVALID_PARAMS); } std::string ssd_offload_path = get_config(config, "ssd_offload_path"); std::string tenant_id = get_config(config, CONFIG_KEY_TENANT_ID, "default"); - - std::string enable_ssd_offload_str = - get_config(config, "enable_ssd_offload", "false"); - std::transform(enable_ssd_offload_str.begin(), enable_ssd_offload_str.end(), - enable_ssd_offload_str.begin(), - [](unsigned char c) { return std::tolower(c); }); bool enable_ssd_offload = - (enable_ssd_offload_str == "true" || enable_ssd_offload_str == "1"); + get_config_bool(config, "enable_ssd_offload", false); + bool enable_client_http_server = + get_config_bool(config, CONFIG_KEY_ENABLE_CLIENT_HTTP_SERVER, false); + auto client_http_port_opt = get_config_int( + config, CONFIG_KEY_CLIENT_HTTP_PORT, DEFAULT_CLIENT_HTTP_PORT); + if (!client_http_port_opt.has_value()) { + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + int client_http_port = client_http_port_opt.value(); - return setup_internal( - local_hostname, metadata_server, global_segment_size, local_buffer_size, - protocol, rdma_devices, master_server_addr, nullptr, ipc_socket_path, - 50052, enable_ssd_offload, true, ssd_offload_path, tenant_id); + return setup_internal(local_hostname, metadata_server, global_segment_size, + local_buffer_size, protocol, rdma_devices, + master_server_addr, nullptr, ipc_socket_path, 50052, + enable_ssd_offload, true, ssd_offload_path, tenant_id, + enable_client_http_server, client_http_port); } tl::expected RealClient::initAll_internal( @@ -1116,9 +1171,22 @@ tl::expected RealClient::tearDownAll_internal() { ReleaseAllAllocatedSegmentRecords(); client_buffer_allocator_.reset(); port_binder_.reset(); + const bool setup_segments_safe_to_free = ReleasePinnedRegionsForFree( + setup_segment_pinned_regions_, "Store setup segment"); + if (!setup_segments_safe_to_free || setup_segment_memory_must_leak_) { + for (auto &ptr : hugepage_segment_ptrs_) ptr.release(); + for (auto &ptr : segment_ptrs_) ptr.release(); + setup_segment_memory_must_leak_ = false; + } hugepage_segment_ptrs_.clear(); - ub_segment_ptrs_.clear(); segment_ptrs_.clear(); + ub_segment_ptrs_.clear(); +#ifdef USE_VRAM_SEGMENT + vram_segment_ptrs_.clear(); +#endif +#if defined(USE_SUNRISE) + sunrise_segment_ptrs_.clear(); +#endif local_hostname = ""; device_name = ""; protocol = ""; @@ -1293,6 +1361,13 @@ void RealClient::ReleaseAllMountedSegmentRecords() { } } +void RealClient::FreeAllocatedStoreSegment(AllocatedSegmentRecord &record) { + if (record.base && ReleasePinnedRegionForFree(record.pinned_region, + "allocated Store segment")) { + free_memory(record.protocol, record.base); + } +} + void RealClient::ReleaseAllocatedSegmentRecord(const std::string &segment_id) { AllocatedSegmentRecord record; bool found = false; @@ -1306,7 +1381,7 @@ void RealClient::ReleaseAllocatedSegmentRecord(const std::string &segment_id) { } } if (found && record.base) { - free_memory(record.protocol, record.base); + FreeAllocatedStoreSegment(record); } } @@ -1317,9 +1392,7 @@ void RealClient::ReleaseAllAllocatedSegmentRecords() { records.swap(allocated_segment_records_); } for (auto &entry : records) { - if (entry.second.base) { - free_memory(entry.second.protocol, entry.second.base); - } + FreeAllocatedStoreSegment(entry.second); } } @@ -1459,17 +1532,23 @@ int RealClient::allocateAndMountSegment( break; } + auto pinned_region = + TryPinStoreSegment(ptr, chunk_size, protocol, "allocated"); auto result = client_->MountSegmentAndGetId(ptr, chunk_size, protocol, location); if (!result.has_value()) { LOG(ERROR) << "MountSegmentAndGetId failed"; - free_memory(protocol, ptr); + if (ReleasePinnedRegionForFree(pinned_region, + "allocated Store segment")) { + free_memory(protocol, ptr); + } break; } std::string segment_id = UuidToString(result.value()); mounted_ids.push_back(segment_id); - allocated_records.push_back({ptr, chunk_size, protocol}); + allocated_records.push_back( + {ptr, chunk_size, protocol, std::move(pinned_region)}); remaining -= chunk_size; } @@ -1481,8 +1560,7 @@ int RealClient::allocateAndMountSegment( client_->UnmountSegmentById(id); } if (allocated_records[i].base) { - free_memory(allocated_records[i].protocol, - allocated_records[i].base); + FreeAllocatedStoreSegment(allocated_records[i]); } } out_segment_ids.clear(); @@ -1568,9 +1646,7 @@ int RealClient::unmountAndFreeSegment( } for (auto &p : to_cleanup) { - if (p.second.base) { - free_memory(p.second.protocol, p.second.base); - } + FreeAllocatedStoreSegment(p.second); } return first_error; @@ -1583,11 +1659,15 @@ int RealClient::health_check() { return HC_HEALTHY; } -int RealClient::start_http_server() { +int RealClient::start_http_server(int port) { using namespace coro_http; - http_server_ = - std::make_unique(/*thread_num=*/1, FLAGS_http_port); + if (http_server_) { + LOG(WARNING) << "Client HTTP server is already running"; + return 0; + } + + http_server_ = std::make_unique(/*thread_num=*/1, port); http_server_->set_http_handler( "/health", [this](coro_http_request &req, coro_http_response &resp) { @@ -1653,11 +1733,11 @@ int RealClient::start_http_server() { auto ec = http_server_->async_start(); if (ec.hasResult()) { - LOG(ERROR) << "Failed to start HTTP server on port " << FLAGS_http_port; + LOG(WARNING) << "Failed to start HTTP server on port " << port; http_server_.reset(); return -1; } - LOG(INFO) << "Client HTTP server started on port " << FLAGS_http_port; + LOG(INFO) << "Client HTTP server started on port " << port; return 0; } @@ -1692,7 +1772,11 @@ tl::expected RealClient::put_internal( return tl::unexpected(ErrorCode::INVALID_PARAMS); } auto &buffer_handle = *alloc_result; - memcpy(buffer_handle.ptr(), value.data(), value.size_bytes()); + auto scatter_result = gather_maybe_device_to_host( + buffer_handle.ptr(), value.data(), value.size_bytes(), "put:" + key); + if (!scatter_result) { + return tl::unexpected(scatter_result.error()); + } std::vector slices = split_into_slices(buffer_handle); @@ -1770,7 +1854,12 @@ tl::expected RealClient::put_batch_internal( return tl::unexpected(ErrorCode::INVALID_PARAMS); } auto &buffer_handle = *alloc_result; - memcpy(buffer_handle.ptr(), value.data(), value.size_bytes()); + auto scatter_result = + gather_maybe_device_to_host(buffer_handle.ptr(), value.data(), + value.size_bytes(), "put_batch:" + key); + if (!scatter_result) { + return tl::unexpected(scatter_result.error()); + } auto slices = split_into_slices(buffer_handle); buffer_handles.emplace_back(std::move(*alloc_result)); batched_slices.emplace(key, std::move(slices)); @@ -1874,8 +1963,12 @@ tl::expected RealClient::put_parts_internal( // Copy all parts into the contiguous buffer size_t offset = 0; for (const auto &value : values) { - memcpy(static_cast(buffer_handle.ptr()) + offset, value.data(), - value.size_bytes()); + auto scatter_result = gather_maybe_device_to_host( + static_cast(buffer_handle.ptr()) + offset, value.data(), + value.size_bytes(), "put_multi_value"); + if (!scatter_result) { + return tl::unexpected(scatter_result.error()); + } offset += value.size_bytes(); } @@ -2597,14 +2690,24 @@ std::shared_ptr RealClient::get_buffer_internal( << "': " << toString(read_result.error()); return nullptr; } + auto checksum_result = + client_->VerifyObjectChecksum(key, objects.at(key), total_length, + query_result.value().object_checksum); + if (!checksum_result) { + LOG(ERROR) << "SSD checksum verification failed for key '" << key + << "': " << toString(checksum_result.error()); + return nullptr; + } return buffer_handle; } // MEMORY / DISK: use client_->Get. FilterQueryResult ensures - // Client::Get's internal FindFirstCompleteReplica can only see + // Client::Get internal FindFirstCompleteReplica can only see // the replica we selected, preventing accidental LOCAL_DISK picks. + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); if (replica.is_disk_replica() && - gpu_staging::IsDevicePointer(buffer_handle->ptr(), nullptr)) { + runtime_accelerator.FindDeviceForPointer(buffer_handle->ptr())) { LOG(WARNING) << "DISK replica for key '" << key << "' received a device pointer from the allocator; " << "file I/O cannot write to GPU memory — read will fail. " @@ -2711,6 +2814,40 @@ tl::expected RealClient::release_buffer_dummy( return {}; } +tl::expected, ErrorCode> +RealClient::allocate_buffer_dummy(size_t size, const UUID &client_id) { + std::unique_lock lock(dummy_client_mutex_); + auto it = shm_contexts_.find(client_id); + if (it == shm_contexts_.end()) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + auto &context = it->second; + if (!context.client_buffer_allocator) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + auto alloc_result = context.client_buffer_allocator->allocate(size); + if (!alloc_result) { + return tl::make_unexpected(ErrorCode::NO_AVAILABLE_HANDLE); + } + + auto handle = std::make_shared(std::move(*alloc_result)); + const uint64_t real_addr = reinterpret_cast(handle->ptr()); + const size_t allocated_size = handle->size(); + for (const auto &shm : context.mapped_shms) { + const uint64_t shm_start = reinterpret_cast(shm.shm_buffer); + const uint64_t shm_end = shm_start + shm.shm_size; + if (real_addr >= shm_start && allocated_size <= shm_end - real_addr) { + const uint64_t dummy_addr = real_addr - shm.shm_addr_offset; + context.active_handles[dummy_addr] = std::move(handle); + return std::make_tuple(dummy_addr, allocated_size); + } + } + + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); +} + std::vector, ErrorCode>> RealClient::batch_acquire_hot_cache(const std::vector &keys) { std::vector, ErrorCode>> results; @@ -2909,8 +3046,10 @@ RealClient::batch_get_buffer_internal( // DISK replicas use storage_backend::vector_read (file I/O) which // can only write to CPU-addressable memory. If the allocator ever // returns device memory for DISK, the read will silently fail. + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); if (replica.is_disk_replica() && - gpu_staging::IsDevicePointer(buffer_handle->ptr(), nullptr)) { + runtime_accelerator.FindDeviceForPointer(buffer_handle->ptr())) { LOG(WARNING) << "DISK replica for key '" << key << "' received a device pointer from the allocator; " @@ -3000,6 +3139,16 @@ RealClient::batch_get_buffer_internal( if (idx_it == disk_key_to_idx.end()) continue; auto &op = disk_ops[idx_it->second]; if (read_result) { + auto checksum_result = client_->VerifyObjectChecksum( + key, slices, op.total_size, + op.query_result.object_checksum); + if (!checksum_result) { + LOG(ERROR) + << "SSD checksum verification failed for key '" + << key + << "': " << toString(checksum_result.error()); + continue; + } final_results[op.original_index] = std::make_shared( std::move(*op.buffer_handle)); @@ -3107,12 +3256,21 @@ RealClient::resolve_writable_buffer_region(void *buffer) const { } tl::expected -RealClient::resolve_ranged_read_metadata(const std::string &key) { +RealClient::resolve_ranged_read_metadata( + const std::string &key, const QueryResultCache *query_result_cache) { if (!client_) { LOG(ERROR) << "Client is not initialized"; return tl::unexpected(ErrorCode::INVALID_PARAMS); } + if (query_result_cache) { + auto cached = query_result_cache->find(key); + if (cached != query_result_cache->end() && + (!cached->second || !cached->second->IsLeaseExpired())) { + return build_ranged_read_metadata_from_query_result(key, + cached->second); + } + } return build_ranged_read_metadata_from_query_result(key, client_->Query(key)); } @@ -3120,7 +3278,7 @@ RealClient::resolve_ranged_read_metadata(const std::string &key) { tl::expected RealClient::execute_ranged_read( const std::string &key, void *buffer, size_t dst_offset, size_t src_offset, size_t size, const RangedReadMetadata &metadata, - bool size_is_buffer_capacity) { + bool size_is_buffer_capacity, bool verify_checksum) { const auto &query_result = metadata.query_result; const auto &replica = metadata.replica; const uint64_t total_size = metadata.total_size; @@ -3152,6 +3310,14 @@ tl::expected RealClient::execute_ranged_read( auto result = batch_get_into_offload_object_internal(endpoint, objects); if (!result) return tl::unexpected(result.error()); + if (verify_checksum) { + auto checksum_result = client_->VerifyObjectChecksum( + key, objects.at(key), total_size, + query_result.object_checksum); + if (!checksum_result) { + return tl::unexpected(checksum_result.error()); + } + } return static_cast(total_size); } @@ -3167,7 +3333,8 @@ tl::expected RealClient::execute_ranged_read( BufferHandle tmp_handle(std::move(*alloc_result)); std::vector tmp_slices; allocateSlices(tmp_slices, replica, tmp_handle.ptr()); - auto filtered_qr = FilterQueryResult(query_result, replica); + auto filtered_qr = + FilterQueryResult(query_result, replica, verify_checksum); auto get_result = client_->Get(key, filtered_qr, tmp_slices); if (!get_result) { LOG(ERROR) << "DISK Get failed for key: " << key @@ -3184,11 +3351,46 @@ tl::expected RealClient::execute_ranged_read( return static_cast(total_size); } + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); + void *dst = static_cast(buffer) + dst_offset; + if (runtime_accelerator.FindDeviceForPointer(dst)) { + if (!client_buffer_allocator_) { + LOG(ERROR) << "Client buffer allocator is not provided"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + auto alloc_result = client_buffer_allocator_->allocate(total_size); + if (!alloc_result) { + LOG(ERROR) << "Failed to allocate temp buffer for GPU memory " + << "read, key: " << key << ", size: " << total_size; + return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE); + } + BufferHandle tmp_handle(std::move(*alloc_result)); + std::vector tmp_slices; + allocateSlices(tmp_slices, replica, tmp_handle.ptr()); + + auto filtered_qr = + FilterQueryResult(query_result, replica, verify_checksum); + auto get_result = client_->Get(key, filtered_qr, tmp_slices); + if (!get_result) { + LOG(ERROR) << "Get failed for key: " << key + << " with error: " << toString(get_result.error()); + return tl::unexpected(get_result.error()); + } + if (auto r = scatter_host_to_maybe_device( + dst, tmp_handle.ptr(), total_size, + "MEMORY full read, key: " + key); + !r) { + return tl::unexpected(r.error()); + } + return static_cast(total_size); + } + std::vector slices; - allocateSlices(slices, replica, - static_cast(buffer) + dst_offset); + allocateSlices(slices, replica, dst); - auto filtered_qr = FilterQueryResult(query_result, replica); + auto filtered_qr = + FilterQueryResult(query_result, replica, verify_checksum); auto get_result = client_->Get(key, filtered_qr, slices); if (!get_result) { LOG(ERROR) << "Get failed for key: " << key @@ -3252,7 +3454,8 @@ tl::expected RealClient::execute_ranged_read( [&](void *tmp_buf) -> tl::expected { std::vector tmp_slices; allocateSlices(tmp_slices, replica, tmp_buf); - auto filtered_qr = FilterQueryResult(query_result, replica); + auto filtered_qr = + FilterQueryResult(query_result, replica, false); auto get_result = client_->Get(key, filtered_qr, tmp_slices); if (!get_result) { LOG(ERROR) @@ -3270,8 +3473,39 @@ tl::expected RealClient::execute_ranged_read( return tl::unexpected(ErrorCode::INVALID_REPLICA); } + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); + void *dst = static_cast(buffer) + dst_offset; + if (runtime_accelerator.FindDeviceForPointer(dst)) { + if (!client_buffer_allocator_) { + LOG(ERROR) << "Client buffer allocator is not provided"; + return tl::unexpected(ErrorCode::INVALID_PARAMS); + } + auto alloc_result = client_buffer_allocator_->allocate(size); + if (!alloc_result) { + LOG(ERROR) << "Failed to allocate temp buffer for GPU ranged " + << "read, key: " << key << ", size: " << size; + return tl::unexpected(ErrorCode::NO_AVAILABLE_HANDLE); + } + BufferHandle tmp_handle(std::move(*alloc_result)); + std::vector tmp_slices; + tmp_slices.emplace_back(Slice{tmp_handle.ptr(), size}); + + auto get_result = + client_->Get(key, query_result, tmp_slices, src_offset); + if (!get_result) { + return tl::unexpected(get_result.error()); + } + if (auto r = scatter_host_to_maybe_device( + dst, tmp_handle.ptr(), size, "MEMORY ranged read, key: " + key); + !r) { + return tl::unexpected(r.error()); + } + return static_cast(size); + } + std::vector slices; - slices.emplace_back(Slice{static_cast(buffer) + dst_offset, size}); + slices.emplace_back(Slice{dst, size}); auto get_result = client_->Get(key, query_result, slices, src_offset); if (!get_result) { @@ -3282,7 +3516,7 @@ tl::expected RealClient::execute_ranged_read( tl::expected RealClient::get_into_range_internal( const std::string &key, void *buffer, size_t dst_offset, size_t src_offset, - size_t size, bool size_is_buffer_capacity) { + size_t size, bool size_is_buffer_capacity, bool verify_checksum) { auto metadata_result = resolve_ranged_read_metadata(key); if (!metadata_result) { if ((metadata_result.error() == ErrorCode::OBJECT_NOT_FOUND || @@ -3294,15 +3528,15 @@ tl::expected RealClient::get_into_range_internal( } return execute_ranged_read(key, buffer, dst_offset, src_offset, size, - metadata_result.value(), - size_is_buffer_capacity); + metadata_result.value(), size_is_buffer_capacity, + verify_checksum); } int64_t RealClient::get_into(const std::string &key, void *buffer, size_t size) { auto result = execute_timed_operation>( [&]() { - return get_into_range_internal(key, buffer, 0, 0, size, true); + return get_into_range_internal(key, buffer, 0, 0, size, true, true); }, [](const auto &ret) { return ret.has_value(); }, [&](uint64_t latency_us, const auto &ret) { @@ -3321,143 +3555,184 @@ RealClient::get_into_ranges_internal( const std::vector>> &all_src_offsets, const std::vector>> &all_sizes, const std::vector *buffer_capacities, - std::vector>>> - *prepared_results, - const std::vector>> *valid_fragments, const QueryResultCache *query_result_cache) { + auto results = build_ranged_read_internal_error_results( + buffers.size(), all_keys, all_dst_offsets, ErrorCode::INVALID_PARAMS); if (!client_) { LOG(ERROR) << "Client is not initialized"; - return build_ranged_read_internal_error_results( - buffers.size(), all_keys, all_dst_offsets, - ErrorCode::INVALID_PARAMS); + return results; } const size_t buffer_count = buffers.size(); - PreparedRangedReadRequest prepared; - if (prepared_results != nullptr && valid_fragments != nullptr) { - prepared.results = std::move(*prepared_results); - prepared.valid_fragments = *valid_fragments; - prepared.required_buffer_sizes.resize(buffer_count, 0); - prepared.top_level_valid = true; - } else { - prepared = prepare_ranged_read_request(buffer_count, all_keys, - all_dst_offsets, all_src_offsets, - all_sizes, "get_into_ranges"); - } - if (!prepared.top_level_valid) { - return std::move(prepared.results); + if (buffer_count != all_keys.size() || + buffer_count != all_dst_offsets.size() || + buffer_count != all_src_offsets.size() || + buffer_count != all_sizes.size() || + (buffer_capacities && buffer_capacities->size() != buffer_count)) { + LOG(ERROR) << "get_into_ranges: top-level size mismatch"; + return results; } - std::vector resolved_buffer_capacities; - if (buffer_capacities != nullptr) { - if (buffer_capacities->size() != buffer_count) { - LOG(ERROR) << "get_into_ranges: buffer capacities size mismatch"; - return build_ranged_read_internal_error_results( - buffer_count, all_keys, all_dst_offsets, - ErrorCode::INVALID_PARAMS); - } - resolved_buffer_capacities = *buffer_capacities; - } else { - resolved_buffer_capacities.resize(buffer_count, 0); + std::vector capacities = buffer_capacities + ? *buffer_capacities + : std::vector(buffer_count); + if (!buffer_capacities) { for (size_t i = 0; i < buffer_count; ++i) { auto region = resolve_writable_buffer_region(buffers[i]); - if (!region.has_value()) { - LOG(ERROR) - << "get_into_ranges: buffer is not Store-managed writable " - "memory at index " - << i; - continue; - } - resolved_buffer_capacities[i] = region->size - region->offset; + if (region) capacities[i] = region->size - region->offset; } } std::unordered_map> metadata_cache; - size_t key_count_hint = 0; - for (const auto &keys : all_keys) { - key_count_hint += keys.size(); - } - metadata_cache.reserve(key_count_hint); - auto now = std::chrono::steady_clock::now(); - if (query_result_cache != nullptr) { - for (const auto &[key, query_result] : *query_result_cache) { - if (!query_result) { - metadata_cache.emplace(key, - tl::unexpected(query_result.error())); - continue; - } - if (query_result->IsLeaseExpired(now)) { - continue; - } - metadata_cache.emplace(key, - build_ranged_read_metadata_from_query_result( - key, query_result)); - } - } + auto metadata_for = [&](const std::string &key) -> auto & { + auto found = metadata_cache.find(key); + if (found != metadata_cache.end()) return found->second; + return metadata_cache + .emplace(key, resolve_ranged_read_metadata(key, query_result_cache)) + .first->second; + }; + struct ScatterLease { + std::chrono::steady_clock::time_point expires_at; + std::optional error; + }; + std::unordered_map scatter_leases; + std::vector memory_transfers; for (size_t i = 0; i < buffer_count; ++i) { - const size_t key_count = prepared.results[i].size(); - const size_t buffer_size = resolved_buffer_capacities[i]; - if (buffer_size == 0 && key_count > 0 && buffer_capacities == nullptr) { + if (!buffers[i] || (!buffer_capacities && capacities[i] == 0)) { + continue; + } + const auto &keys = all_keys[i]; + const auto &dst_groups = all_dst_offsets[i]; + const auto &src_groups = all_src_offsets[i]; + const auto &size_groups = all_sizes[i]; + if (keys.size() != dst_groups.size() || + keys.size() != src_groups.size() || + keys.size() != size_groups.size()) { continue; } - for (size_t j = 0; j < key_count; ++j) { - auto metadata_it = metadata_cache.find(all_keys[i][j]); - if (metadata_it == metadata_cache.end()) { - metadata_it = - metadata_cache - .emplace( - all_keys[i][j], - build_ranged_read_metadata_from_query_result( - all_keys[i][j], client_->Query(all_keys[i][j]))) - .first; + for (size_t j = 0; j < keys.size(); ++j) { + const auto &dst_offsets = dst_groups[j]; + const auto &src_offsets = src_groups[j]; + const auto &sizes = size_groups[j]; + auto &range_results = results[i][j]; + if (dst_offsets.size() != src_offsets.size() || + dst_offsets.size() != sizes.size()) { + continue; + } + auto &metadata_result = metadata_for(keys[j]); + if (!metadata_result) { + std::fill(range_results.begin(), range_results.end(), + tl::unexpected(metadata_result.error())); + continue; } - auto &metadata_result = metadata_it->second; - for (size_t k = 0; k < prepared.results[i][j].size(); ++k) { - if (!prepared.valid_fragments[i][j][k]) { - continue; - } + const auto &metadata = metadata_result.value(); + if (metadata.replica.is_memory_replica()) { + const auto &handle = + metadata.replica.get_memory_descriptor().buffer_descriptor; + auto [lease_it, inserted] = scatter_leases.try_emplace(keys[j]); + if (inserted) + lease_it->second.expires_at = + metadata.query_result.lease_timeout; + memory_transfers.push_back(TransferEngine::ScatterTransferRange{ + .opcode = TransferRequest::READ, + .remote_segment = handle.transport_endpoint_, + .remote_base_offset = handle.buffer_address_, + .remote_size = handle.size_, + .local_buffer = buffers[i], + .local_capacity = capacities[i], + .local_offsets = dst_offsets, + .remote_offsets = src_offsets, + .lengths = sizes, + .on_fragment_complete = + [results = &range_results, sizes = &sizes, + lease = &lease_it->second](size_t k, + const Status &status) { + if (status.ok() && !lease->error.has_value() && + std::chrono::steady_clock::now() < + lease->expires_at) { + (*results)[k] = + static_cast((*sizes)[k]); + return; + } + const auto error = lease->error.value_or( + status.ok() ? ErrorCode::LEASE_EXPIRED + : scatter_transfer_error(status)); + (*results)[k] = tl::unexpected(error); + }, + }); + continue; + } - if (all_sizes[i][j][k] > 0 && - (all_dst_offsets[i][j][k] > buffer_size || - all_sizes[i][j][k] > - buffer_size - all_dst_offsets[i][j][k])) { - LOG(ERROR) - << "get_into_ranges: destination overflow, " - "buffer_index=" - << i << " key_index=" << j << " fragment_index=" << k - << " dst_offset=" << all_dst_offsets[i][j][k] - << " size=" << all_sizes[i][j][k] - << " buffer_size=" << buffer_size; + for (size_t k = 0; k < range_results.size(); ++k) { + const size_t dst_offset = dst_offsets[k]; + if (dst_offset > capacities[i] || + sizes[k] > capacities[i] - dst_offset) { continue; } - if (!metadata_result) { - if ((metadata_result.error() == - ErrorCode::OBJECT_NOT_FOUND || - metadata_result.error() == - ErrorCode::REPLICA_IS_NOT_READY) && - all_src_offsets[i][j][k] == 0) { - VLOG(1) - << "Object not found for key: " << all_keys[i][j]; - } - prepared.results[i][j][k] = - tl::unexpected(metadata_result.error()); - continue; - } + range_results[k] = execute_ranged_read( + keys[j], buffers[i], dst_offset, src_offsets[k], sizes[k], + metadata, false, false); + } + } + } - prepared.results[i][j][k] = execute_ranged_read( - all_keys[i][j], buffers[i], all_dst_offsets[i][j][k], - all_src_offsets[i][j][k], all_sizes[i][j][k], - metadata_result.value()); + auto next_refresh_delay = [&]() { + const auto now = std::chrono::steady_clock::now(); + auto delay = std::chrono::nanoseconds::max(); + for (const auto &[key, lease] : scatter_leases) { + (void)key; + if (lease.error.has_value()) continue; + const auto remaining = + std::chrono::duration_cast( + lease.expires_at - now); + delay = std::min(delay, std::max(remaining / 2, + std::chrono::nanoseconds::zero())); + } + return delay; + }; + auto refresh_leases = [&]() { + std::vector keys; + keys.reserve(scatter_leases.size()); + for (const auto &[key, lease] : scatter_leases) + if (!lease.error.has_value()) keys.push_back(key); + auto refreshed = client_->BatchQuery(keys); + for (size_t i = 0; i < refreshed.size(); ++i) { + auto &lease = scatter_leases.at(keys[i]); + if (!refreshed[i]) { + lease.error = refreshed[i].error(); + continue; } + lease.expires_at = refreshed[i]->lease_timeout; + } + }; + + // Planning may consume most of a short lease; renew before submission. + if (!scatter_leases.empty()) refresh_leases(); + auto operation = client_->SubmitScatter(memory_transfers); + if (!operation.has_value()) { + const auto failure = + Status::InvalidArgument("TransferSubmitter not initialized"); + for (const auto &transfer : memory_transfers) { + for (size_t i = 0; i < transfer.lengths.size(); ++i) + transfer.on_fragment_complete(i, failure); } + return results; } - return prepared.results; + while (true) { + const auto delay = next_refresh_delay(); + const auto status = delay == std::chrono::nanoseconds::max() + ? operation->wait() + : operation->waitFor(delay); + if (!status.IsClock()) break; + refresh_leases(); + } + return results; } std::vector>> RealClient::get_into_ranges( @@ -3472,7 +3747,7 @@ std::vector>> RealClient::get_into_ranges( [&]() { return convert_ranged_read_results(get_into_ranges_internal( buffers, all_keys, all_dst_offsets, all_src_offsets, - all_sizes, nullptr, nullptr, nullptr, query_result_cache)); + all_sizes, nullptr, query_result_cache)); }, [](const auto &) { return true; }, [&](uint64_t latency_us, const auto &ret) { @@ -3619,17 +3894,7 @@ std::vector> RealClient::batch_put_from_internal( void *buffer = buffers[i]; size_t size = sizes[i]; - std::vector slices; - uint64_t offset = 0; - - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - - all_slices[key] = std::move(slices); + all_slices[key] = split_into_slices(buffer, size); } std::vector> ordered_batched_slices; @@ -3669,15 +3934,7 @@ tl::expected RealClient::put_from_internal( } // Create slices directly from the user buffer - std::vector slices; - uint64_t offset = 0; - - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } + std::vector slices = split_into_slices(buffer, size); auto put_result = client_->Put(key, slices, config); if (!put_result) { @@ -3724,7 +3981,11 @@ tl::expected RealClient::upsert_internal( return tl::unexpected(ErrorCode::INVALID_PARAMS); } auto &buffer_handle = *alloc_result; - memcpy(buffer_handle.ptr(), value.data(), value.size_bytes()); + auto scatter_result = gather_maybe_device_to_host( + buffer_handle.ptr(), value.data(), value.size_bytes(), "upsert:" + key); + if (!scatter_result) { + return tl::unexpected(scatter_result.error()); + } std::vector slices = split_into_slices(buffer_handle); @@ -3780,14 +4041,7 @@ tl::expected RealClient::upsert_from_internal( return {}; } - std::vector slices; - uint64_t offset = 0; - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } + std::vector slices = split_into_slices(buffer, size); auto result = client_->Upsert(key, slices, config); if (!result) { @@ -3833,15 +4087,8 @@ RealClient::batch_upsert_from_internal(const std::vector &keys, ordered_batched_slices.reserve(keys.size()); for (size_t i = 0; i < keys.size(); ++i) { - std::vector slices; - uint64_t offset = 0; - while (offset < sizes[i]) { - auto chunk_size = std::min(sizes[i] - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffers[i]) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } - ordered_batched_slices.emplace_back(std::move(slices)); + ordered_batched_slices.emplace_back( + split_into_slices(buffers[i], sizes[i])); } return client_->BatchUpsert(keys, ordered_batched_slices, config); @@ -3970,8 +4217,12 @@ tl::expected RealClient::upsert_parts_internal( auto &buffer_handle = *alloc_result; size_t offset = 0; for (const auto &value : values) { - memcpy(static_cast(buffer_handle.ptr()) + offset, value.data(), - value.size_bytes()); + auto scatter_result = gather_maybe_device_to_host( + static_cast(buffer_handle.ptr()) + offset, value.data(), + value.size_bytes(), "upsert_parts:" + key); + if (!scatter_result) { + return tl::unexpected(scatter_result.error()); + } offset += value.size_bytes(); } @@ -4054,7 +4305,12 @@ tl::expected RealClient::upsert_batch_internal( return tl::unexpected(ErrorCode::INVALID_PARAMS); } auto &buffer_handle = *alloc_result; - memcpy(buffer_handle.ptr(), value.data(), value.size_bytes()); + auto scatter_result = gather_maybe_device_to_host( + buffer_handle.ptr(), value.data(), value.size_bytes(), + "upsert_batch:" + key); + if (!scatter_result) { + return tl::unexpected(scatter_result.error()); + } auto slices = split_into_slices(buffer_handle); buffer_handles.emplace_back(std::move(*alloc_result)); batched_slices.emplace(key, std::move(slices)); @@ -4152,15 +4408,6 @@ RealClient::batch_get_into_dummy_helper( const std::vector &dummy_buffers, const std::vector &sizes, int32_t device_id, const UUID &client_id) { -#ifdef USE_ASCEND_DIRECT - if (!ContextManager::getInstance().setCurrentContextByPhysicalId( - device_id)) { - LOG(ERROR) << "Failed to set context for physical device " << device_id; - co_return std::vector>( - keys.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); - } -#endif - // Hold shared_lock for the entire operation to prevent SHM from being // unmapped while batch_get_into_internal is using the translated buffers. auto lock = std::make_shared>( @@ -4193,15 +4440,25 @@ RealClient::batch_get_into_dummy_helper( std::vector sizes; std::vector buffers; std::shared_ptr> lock; + int32_t device_id; }; auto state = std::make_unique(); state->keys = keys; state->sizes = sizes; state->buffers = std::move(buffers_result.value()); state->lock = std::move(lock); + state->device_id = device_id; auto *s = state.get(); auto try_result = co_await coro_io::post([this, s]() { +#ifdef USE_ASCEND_DIRECT + auto context_result = + set_context_if_needed(protocol, s->device_id, "batch_get worker"); + if (!context_result) { + return std::vector>( + s->keys.size(), tl::unexpected(context_result.error())); + } +#endif return batch_get_into_internal(s->keys, s->buffers, s->sizes); }); co_return try_result.value(); @@ -4242,6 +4499,145 @@ RealClient::batch_put_from_multi_buffers_dummy_helper( keys, real_buffers_result.value(), all_sizes, config); } +std::vector> +RealClient::batch_put_from_cuda_ipc_dummy_helper( + const std::vector &requests, + const ReplicateConfig &config, const UUID &client_id) { + std::shared_lock lock(dummy_client_mutex_); + auto it = shm_contexts_.find(client_id); + if (it == shm_contexts_.end()) { + LOG(ERROR) << "client_id=" << client_id << ", error=shm_not_mapped"; + return std::vector>( + requests.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + + const ShmContext &context = it->second; + const bool has_metadata = + !requests.empty() && requests[0].metadata.size != 0; + for (const auto &request : requests) { + if ((request.metadata.size != 0) != has_metadata) { + LOG(ERROR) << "Mixed cuda ipc metadata batches are not supported"; + return std::vector>( + requests.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + } + + std::vector payload_mappings; + std::vector keys; + std::vector> all_buffers; + std::vector> all_sizes; + payload_mappings.reserve(requests.size()); + keys.reserve(requests.size()); + all_buffers.reserve(requests.size()); + all_sizes.reserve(requests.size()); + + const MappedShm *last_hit_shm = nullptr; + for (const auto &request : requests) { + auto mapping = device::CudaIpcBufferMapping::Open(request.payload); + if (!mapping) { + return std::vector>( + requests.size(), tl::unexpected(mapping.error())); + } + payload_mappings.push_back(std::move(*mapping)); + void *payload_ptr = payload_mappings.back().ptr(); + keys.push_back(request.key); + if (has_metadata) { + void *metadata_ptr = nullptr; + if (!map_dummy_buffer_to_real( + context, request.metadata.ptr, + static_cast(request.metadata.size), last_hit_shm, + metadata_ptr)) { + LOG(ERROR) << "Dummy metadata buffer at " + << request.metadata.ptr << " (size " + << request.metadata.size + << ") not found in any mapped shared memory, " + << "client_id=" << client_id; + return std::vector>( + requests.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + all_buffers.push_back({metadata_ptr, payload_ptr}); + all_sizes.push_back({static_cast(request.metadata.size), + static_cast(request.payload.size)}); + } else { + all_buffers.push_back({payload_ptr}); + all_sizes.push_back({static_cast(request.payload.size)}); + } + } + + return batch_put_from_multi_buffers_internal(keys, all_buffers, all_sizes, + config); +} + +std::vector> +RealClient::batch_get_into_cuda_ipc_dummy_helper( + const std::vector &requests, const UUID &client_id) { + { + std::shared_lock lock(dummy_client_mutex_); + if (shm_contexts_.find(client_id) == shm_contexts_.end()) { + LOG(ERROR) << "client_id=" << client_id << ", error=shm_not_mapped"; + return std::vector>( + requests.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + } + } + + std::vector> results( + requests.size(), tl::unexpected(ErrorCode::INVALID_PARAMS)); + std::vector mappings; + std::vector buffers; + std::vector> all_keys; + std::vector>> all_dst_offsets; + std::vector>> all_src_offsets; + std::vector>> all_sizes; + std::vector buffer_capacities; + std::vector original_indices; + mappings.reserve(requests.size()); + buffers.reserve(requests.size()); + all_keys.reserve(requests.size()); + all_dst_offsets.reserve(requests.size()); + all_src_offsets.reserve(requests.size()); + all_sizes.reserve(requests.size()); + buffer_capacities.reserve(requests.size()); + original_indices.reserve(requests.size()); + + for (size_t i = 0; i < requests.size(); ++i) { + const auto &request = requests[i]; + auto mapping = device::CudaIpcBufferMapping::Open(request.destination); + if (!mapping) { + results[i] = tl::unexpected(mapping.error()); + continue; + } + mappings.push_back(std::move(*mapping)); + buffers.push_back(mappings.back().ptr()); + all_keys.push_back({request.key}); + all_dst_offsets.push_back({{0}}); + all_src_offsets.push_back( + {{static_cast(request.source_offset)}}); + all_sizes.push_back({{static_cast(request.size)}}); + buffer_capacities.push_back(static_cast(request.size)); + original_indices.push_back(i); + } + + if (buffers.empty()) { + return results; + } + + auto range_results = get_into_ranges_internal( + buffers, all_keys, all_dst_offsets, all_src_offsets, all_sizes, + &buffer_capacities, nullptr); + for (size_t i = 0; i < original_indices.size(); ++i) { + if (i < range_results.size() && range_results[i].size() == 1 && + range_results[i][0].size() == 1) { + results[original_indices[i]] = range_results[i][0][0]; + } else { + LOG(ERROR) << "Invalid cuda ipc tensor read result shape for key " + << requests[original_indices[i]].key; + results[original_indices[i]] = + tl::unexpected(ErrorCode::INTERNAL_ERROR); + } + } + return results; +} + std::vector> RealClient::batch_get_into_multi_buffers_dummy_helper( const std::vector &keys, @@ -4280,7 +4676,8 @@ RealClient::batch_get_into_multi_buffers_dummy_helper( tl::expected RealClient::get_into_range_shm_helper( const std::string &key, uint64_t dummy_buffer, size_t dst_offset, - size_t src_offset, size_t size, const UUID &client_id) { + size_t src_offset, size_t size, bool size_is_buffer_capacity, + bool verify_checksum, const UUID &client_id) { std::shared_lock lock(dummy_client_mutex_); auto it = shm_contexts_.find(client_id); if (it == shm_contexts_.end()) { @@ -4299,7 +4696,8 @@ tl::expected RealClient::get_into_range_shm_helper( return tl::unexpected(ErrorCode::INVALID_PARAMS); } - return get_into_range_internal(key, real_buffer, 0, src_offset, size); + return get_into_range_internal(key, real_buffer, 0, src_offset, size, + size_is_buffer_capacity, verify_checksum); } std::vector>>> @@ -4325,41 +4723,30 @@ RealClient::get_into_ranges_shm_helper( } #endif - const size_t buffer_count = dummy_buffers.size(); - auto prepared = prepare_ranged_read_request( - buffer_count, all_keys, all_dst_offsets, all_src_offsets, all_sizes, - "get_into_ranges_shm_helper"); - if (!prepared.top_level_valid) { - return std::move(prepared.results); - } - std::shared_lock lock(dummy_client_mutex_); auto it = shm_contexts_.find(client_id); if (it == shm_contexts_.end()) { LOG(ERROR) << "client_id=" << client_id << ", error=shm_not_mapped"; - fill_ranged_read_results_with_error(prepared.results, - ErrorCode::INVALID_PARAMS); - return prepared.results; - } - - if (!prepared.has_any_valid_fragment) { - return prepared.results; + return build_ranged_read_internal_error_results( + dummy_buffers.size(), all_keys, all_dst_offsets, + ErrorCode::INVALID_PARAMS); } + std::vector capacities; auto real_buffers_result = map_dummy_addrs_to_real_ptrs( - it->second, dummy_buffers, prepared.required_buffer_sizes, client_id); + it->second, dummy_buffers, std::vector(dummy_buffers.size()), + client_id, &capacities); if (!real_buffers_result) { - fill_ranged_read_results_with_error(prepared.results, - real_buffers_result.error()); - return prepared.results; + return build_ranged_read_internal_error_results( + dummy_buffers.size(), all_keys, all_dst_offsets, + real_buffers_result.error()); } auto query_result_cache = build_query_result_cache_from_cached_results(cached_query_results); return get_into_ranges_internal( real_buffers_result.value(), all_keys, all_dst_offsets, all_src_offsets, - all_sizes, &prepared.required_buffer_sizes, &prepared.results, - &prepared.valid_fragments, + all_sizes, &capacities, query_result_cache.empty() ? nullptr : &query_result_cache); } @@ -4367,7 +4754,7 @@ std::vector> RealClient::batch_get_into_internal(const std::vector &keys, const std::vector &buffers, const std::vector &sizes) { - auto start_time = std::chrono::steady_clock::now(); + [[maybe_unused]] auto start_time = std::chrono::steady_clock::now(); // Validate preconditions if (!client_) { LOG(ERROR) << "Client is not initialized"; @@ -4636,8 +5023,9 @@ RealClient::batch_get_into_internal(const std::vector &keys, store_segment_it->second.emplace(op_it.first, op_it.second.slices); } - size_t offload_object_count = 0; - auto start_read_store_time = std::chrono::steady_clock::now(); + [[maybe_unused]] size_t offload_object_count = 0; + [[maybe_unused]] auto start_read_store_time = + std::chrono::steady_clock::now(); for (auto &offload_objects_it : offload_objects) { offload_object_count += offload_objects_it.second.size(); auto batch_get_offload_result = batch_get_into_offload_object_internal( @@ -4650,14 +5038,27 @@ RealClient::batch_get_into_internal(const std::vector &keys, .original_index] = tl::make_unexpected(batch_get_offload_result.error()); } + continue; + } + for (const auto &offload_object_it : offload_objects_it.second) { + const auto &op = + valid_local_disk_operations.at(offload_object_it.first); + auto checksum_result = client_->VerifyObjectChecksum( + offload_object_it.first, offload_object_it.second, + op.total_size, op.query_result.object_checksum); + if (!checksum_result) { + results[op.original_index] = + tl::make_unexpected(checksum_result.error()); + } } } auto end_time = std::chrono::steady_clock::now(); - auto elapsed_time = std::chrono::duration_cast( - end_time - start_time) - .count(); - auto read_store_time = + [[maybe_unused]] auto elapsed_time = + std::chrono::duration_cast(end_time - + start_time) + .count(); + [[maybe_unused]] auto read_store_time = std::chrono::duration_cast( end_time - start_read_store_time) .count(); @@ -4708,25 +5109,10 @@ int RealClient::put_from_with_metadata(const std::string &key, void *buffer, } // Create slices directly from the user buffer - std::vector slices; - // Add metadata slice - uint64_t metadata_offset = 0; - while (metadata_offset < metadata_size) { - auto metadata_chunk_size = - std::min(metadata_size - metadata_offset, kMaxSliceSize); - void *metadata_chunk_ptr = - static_cast(metadata_buffer) + metadata_offset; - slices.emplace_back(Slice{metadata_chunk_ptr, metadata_chunk_size}); - metadata_offset += metadata_chunk_size; - } - - uint64_t offset = 0; - while (offset < size) { - auto chunk_size = std::min(size - offset, kMaxSliceSize); - void *chunk_ptr = static_cast(buffer) + offset; - slices.emplace_back(Slice{chunk_ptr, chunk_size}); - offset += chunk_size; - } + std::vector slices = + split_into_slices(metadata_buffer, metadata_size); + auto data_slices = split_into_slices(buffer, size); + slices.insert(slices.end(), data_slices.begin(), data_slices.end()); auto put_result = client_->Put(key, slices, config); if (!put_result) { LOG(ERROR) << "Put operation failed with error: " @@ -5079,6 +5465,19 @@ RealClient::batch_get_into_multi_buffers_internal( results[disk_it->second.original_index] = tl::make_unexpected(read_result.error()); } + continue; + } + for (auto &[key, slices] : objects) { + auto disk_it = valid_local_disk_ops.find(key); + if (disk_it == valid_local_disk_ops.end()) continue; + auto &op = disk_it->second; + auto checksum_result = client_->VerifyObjectChecksum( + key, slices, op.total_size, + op.query_result.object_checksum); + if (!checksum_result) { + results[op.original_index] = + tl::make_unexpected(checksum_result.error()); + } } } } @@ -5274,107 +5673,48 @@ void RealClient::stop_dummy_client_monitor() { } } int RealClient::start_ipc_server() { - ipc_running_ = true; - ipc_thread_ = std::jthread(&RealClient::ipc_server_func, this); - return 0; -} - -int RealClient::stop_ipc_server() { - ipc_running_ = false; - // Connect to self to unblock accept if blocked, or unlink - if (!ipc_socket_path_.empty()) { - // Create a dummy socket and connect - int sock = socket(AF_UNIX, SOCK_STREAM, 0); - if (sock >= 0) { - struct sockaddr_un addr; - memset(&addr, 0, sizeof(addr)); - addr.sun_family = AF_UNIX; - strncpy(&addr.sun_path[1], ipc_socket_path_.c_str(), - sizeof(addr.sun_path) - 2); - connect(sock, (struct sockaddr *)&addr, - sizeof(sa_family_t) + strlen(&addr.sun_path[1]) + 1); - close(sock); - } - } - // jthread will join on destruction or we can explicitly join if needed - // But recvmsg/accept might block. The logic above attempts to unblock. - return 0; -} - -void RealClient::ipc_server_func() { - int server_sock = socket(AF_UNIX, SOCK_STREAM, 0); - if (server_sock < 0) { - LOG(ERROR) << "Failed to create IPC socket"; - return; - } - - struct sockaddr_un addr; - memset(&addr, 0, sizeof(addr)); - addr.sun_family = AF_UNIX; - // Using abstract namespace, so we don't need to unlink - strncpy(&addr.sun_path[1], ipc_socket_path_.c_str(), - sizeof(addr.sun_path) - 2); - - if (bind(server_sock, (struct sockaddr *)&addr, - sizeof(sa_family_t) + strlen(&addr.sun_path[1]) + 1) < 0) { - LOG(ERROR) << "Failed to bind IPC socket: " << strerror(errno); - close(server_sock); - return; - } - - if (listen(server_sock, 5) < 0) { - LOG(ERROR) << "Failed to listen on IPC socket: " << strerror(errno); - close(server_sock); - return; - } - - LOG(INFO) << "IPC server is listening"; - - while (ipc_running_) { - int client_sock = accept(server_sock, nullptr, nullptr); - if (client_sock < 0) { - if (ipc_running_) { - LOG(ERROR) << "Accept failed: " << strerror(errno); - } - continue; - } - - if (!ipc_running_) { - close(client_sock); - break; + uds_acceptor_ = std::make_unique(ipc_socket_path_); + uds_acceptor_->registerHandler([this](UdsConnection &connection) { + auto timeout_result = connection.setRecvTimeout(kIpcRequestRecvTimeout); + if (!timeout_result) { + LOG(ERROR) << timeout_result.error(); + return; } - // Set recv timeout to prevent slow/malicious clients from blocking - struct timeval tv = {.tv_sec = 5, .tv_usec = 0}; - setsockopt(client_sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)); - - // Read request type discriminator IpcRequestType req_type; - if (recv(client_sock, &req_type, sizeof(req_type), MSG_WAITALL) != - sizeof(req_type)) { + if (connection.recvRaw(&req_type, sizeof(req_type)) != 0) { LOG(ERROR) << "Failed to read IPC request type"; - close(client_sock); - continue; + return; } if (req_type == IPC_SHM_REGISTER) { - handle_ipc_shm_register(client_sock); + handle_ipc_shm_register(connection); } else if (req_type == IPC_SHM_FD_REQUEST) { - handle_ipc_shm_fd_request(client_sock); + handle_ipc_shm_fd_request(connection); } else { LOG(ERROR) << "Unknown IPC request type: " << req_type; } - - close(client_sock); + }); + auto start_result = uds_acceptor_->start(); + if (!start_result) { + LOG(ERROR) << start_result.error(); + uds_acceptor_.reset(); + return -1; } + return 0; +} - close(server_sock); - LOG(INFO) << "IPC server stopped"; +int RealClient::stop_ipc_server() { + if (uds_acceptor_) { + uds_acceptor_->stop(); + uds_acceptor_.reset(); + } + return 0; } -void RealClient::handle_ipc_shm_register(int client_sock) { +void RealClient::handle_ipc_shm_register(UdsConnection &connection) { ShmRegisterRequest req; - int fd = ipc_recv_fd(client_sock, &req, sizeof(req)); + int fd = connection.recvFd(&req, sizeof(req)); if (fd < 0) { LOG(ERROR) << "Failed to receive fd for SHM_REGISTER"; return; @@ -5386,12 +5726,12 @@ void RealClient::handle_ipc_shm_register(int client_sock) { req.device_id, client_id); int status = result.has_value() ? 0 : -1; - ::send(client_sock, &status, sizeof(status), 0); + connection.sendRaw(&status, sizeof(status)); } -void RealClient::handle_ipc_shm_fd_request(int client_sock) { +void RealClient::handle_ipc_shm_fd_request(UdsConnection &connection) { ShmFdRequest req; - if (recv(client_sock, &req, sizeof(req), MSG_WAITALL) != sizeof(req)) { + if (connection.recvRaw(&req, sizeof(req)) != 0) { LOG(ERROR) << "Failed to read ShmFdRequest payload"; return; } @@ -5407,7 +5747,7 @@ void RealClient::handle_ipc_shm_fd_request(int client_sock) { if (shm_contexts_.find(client_id) == shm_contexts_.end()) { LOG(ERROR) << "Unregistered client_id in fd request: " << client_id.first << ":" << client_id.second; - ::send(client_sock, &resp, sizeof(resp), 0); + connection.sendRaw(&resp, sizeof(resp)); return; } } @@ -5415,13 +5755,13 @@ void RealClient::handle_ipc_shm_fd_request(int client_sock) { // Currently only SHM_SEG_HOT_CACHE is supported if (req.segment_type != SHM_SEG_HOT_CACHE) { LOG(ERROR) << "Unknown segment_type: " << req.segment_type; - ::send(client_sock, &resp, sizeof(resp), 0); + connection.sendRaw(&resp, sizeof(resp)); return; } if (!client_ || !client_->IsHotCacheEnabled()) { LOG(ERROR) << "Hot cache not available for fd request"; - ::send(client_sock, &resp, sizeof(resp), 0); + connection.sendRaw(&resp, sizeof(resp)); return; } @@ -5429,14 +5769,14 @@ void RealClient::handle_ipc_shm_fd_request(int client_sock) { auto seg = hot_cache->GetShmSegment(); if (!seg || seg->fd < 0) { LOG(ERROR) << "Hot cache shm segment not available"; - ::send(client_sock, &resp, sizeof(resp), 0); + connection.sendRaw(&resp, sizeof(resp)); return; } resp.status = 0; resp.shm_size = seg->size; - if (ipc_send_fd(client_sock, seg->fd, &resp, sizeof(resp)) < 0) { + if (connection.sendFd(seg->fd, &resp, sizeof(resp)) < 0) { LOG(ERROR) << "Failed to send hot cache fd to dummy client"; } } @@ -5589,10 +5929,10 @@ RealClient::batch_get_into_offload_object_internal( std::vector keys; std::vector storage_keys; std::vector sizes; + const TenantId tenant_id(client_->tenant_id()); for (const auto &object_it : objects) { keys.emplace_back(object_it.first); - storage_keys.emplace_back( - MakeTenantScopedStorageKey(client_->tenant_id(), object_it.first)); + storage_keys.emplace_back(tenant_id.MakeScopedKey(object_it.first)); int64_t total = 0; for (const auto &s : object_it.second) total += s.size; sizes.emplace_back(total); @@ -5661,7 +6001,7 @@ ClientRequester::ClientRequester() { client_pools_ = std::make_shared>( - pool_conf); + pool_conf, GetStoreRpcClientIoContextPool()); } tl::expected diff --git a/mooncake-store/src/real_client_main.cpp b/mooncake-store/src/real_client_main.cpp index 15c7202f33..a20029a8ee 100644 --- a/mooncake-store/src/real_client_main.cpp +++ b/mooncake-store/src/real_client_main.cpp @@ -3,6 +3,7 @@ #include #include "client_service.h" +#include "common.h" #include "config.h" #include "real_client.h" @@ -17,7 +18,9 @@ DEFINE_string(master_server_address, "127.0.0.1:50051", DEFINE_string(protocol, "tcp", "Protocol"); DEFINE_int32(port, 50052, "Real Client service port"); DEFINE_string(global_segment_size, "4 GB", "Size of global segment"); +DEFINE_string(local_buffer_size, "0", "Size of local buffer (e.g., 16MB, 1GB)"); DEFINE_int32(threads, 1, "Number of threads for client service"); +DEFINE_string(tenant_id, "default", "Tenant identifier"); DEFINE_bool(enable_offload, false, "Enable offload availability"); DEFINE_bool(start_offload_rpc_server, true, "Expose TCP RPC for disk-tier reads " @@ -44,6 +47,8 @@ void RegisterClientRpcService(coro_rpc::coro_rpc_server &server, &real_client); server.register_handler< &RealClient::batch_put_from_multi_buffers_dummy_helper>(&real_client); + server.register_handler<&RealClient::batch_put_from_cuda_ipc_dummy_helper>( + &real_client); server.register_handler<&RealClient::upsert_dummy_helper>(&real_client); server.register_handler<&RealClient::upsert_from_dummy_helper>( &real_client); @@ -57,6 +62,8 @@ void RegisterClientRpcService(coro_rpc::coro_rpc_server &server, &real_client); server.register_handler< &RealClient::batch_get_into_multi_buffers_dummy_helper>(&real_client); + server.register_handler<&RealClient::batch_get_into_cuda_ipc_dummy_helper>( + &real_client); server.register_handler<&RealClient::get_into_range_shm_helper>( &real_client); server.register_handler<&RealClient::get_into_ranges_shm_helper>( @@ -80,6 +87,7 @@ void RegisterClientRpcService(coro_rpc::coro_rpc_server &server, server.register_handler<&RealClient::release_buffer_dummy>(&real_client); server.register_handler<&RealClient::batch_acquire_buffer_dummy>( &real_client); + server.register_handler<&RealClient::allocate_buffer_dummy>(&real_client); server.register_handler<&RealClient::create_copy_task>(&real_client); server.register_handler<&RealClient::create_move_task>(&real_client); server.register_handler<&RealClient::query_task>(&real_client); @@ -102,6 +110,7 @@ int main(int argc, char *argv[]) { } size_t global_segment_size = string_to_byte_size(FLAGS_global_segment_size); + size_t local_buffer_size = string_to_byte_size(FLAGS_local_buffer_size); #ifdef USE_ASCEND_DIRECT // just set to true, does not affect GPU process. globalConfig().ascend_agent_mode = true; @@ -109,10 +118,12 @@ int main(int argc, char *argv[]) { auto client_inst = RealClient::create(); auto res = client_inst->setup_internal( - FLAGS_host, FLAGS_metadata_server, global_segment_size, 0, - FLAGS_protocol, FLAGS_device_names, FLAGS_master_server_address, - nullptr, "@mooncake_client_" + std::to_string(FLAGS_port) + ".sock", - FLAGS_port, FLAGS_enable_offload, FLAGS_start_offload_rpc_server); + FLAGS_host, FLAGS_metadata_server, global_segment_size, + local_buffer_size, FLAGS_protocol, FLAGS_device_names, + FLAGS_master_server_address, nullptr, + "@mooncake_client_" + std::to_string(FLAGS_port) + ".sock", FLAGS_port, + FLAGS_enable_offload, FLAGS_start_offload_rpc_server, "", + FLAGS_tenant_id, FLAGS_enable_http_server, FLAGS_http_port); if (!res) { LOG(FATAL) << "Failed to setup client: " << toString(res.error()); return -1; @@ -123,10 +134,11 @@ int main(int argc, char *argv[]) { return -1; } - coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, FLAGS_host); + auto rpc_bind_host = getHostNameWithoutPort(FLAGS_host); + coro_rpc::coro_rpc_server server(FLAGS_threads, FLAGS_port, rpc_bind_host); RegisterClientRpcService(server, *client_inst); - LOG(INFO) << "Starting real client service on " << FLAGS_host << ":" + LOG(INFO) << "Starting real client service on " << rpc_bind_host << ":" << FLAGS_port; return server.start(); diff --git a/mooncake-store/src/registered_pinned_memory.cpp b/mooncake-store/src/registered_pinned_memory.cpp new file mode 100644 index 0000000000..92aecb6917 --- /dev/null +++ b/mooncake-store/src/registered_pinned_memory.cpp @@ -0,0 +1,241 @@ +#include "registered_pinned_memory.h" + +#include +#include +#include + +#include + +#include "ascii_string.h" +#include "integer_parser.h" + +#if defined(USE_CUDA) +#include +#endif + +namespace mooncake { +namespace { + +std::pair ParsePinnedMemoryConfig() { + const char* raw_value = std::getenv("MC_STORE_PIN_MEMORY_MAX_BYTES"); + if (!raw_value || raw_value[0] == '\0') return {false, 0}; + + const auto limit = + TryParseInteger(TrimAsciiWhitespace(raw_value)); + if (!limit.has_value()) { + LOG(WARNING) << "Invalid MC_STORE_PIN_MEMORY_MAX_BYTES='" << raw_value + << "', disabling Store segment pinning"; + return {false, 0}; + } + return {*limit != 0, *limit}; +} + +void LogPinSkip(const std::string& owner, const char* reason, size_t size) { + LOG(WARNING) << "Skip cudaHostRegister for " << owner << ": " << reason + << ", size=" << size; +} + +#if defined(USE_CUDA) +bool RegisterPinnedRegionWithCuda(void* addr, size_t size, + std::string* error_message) { + cudaError_t err = cudaHostRegister(addr, size, cudaHostRegisterPortable); + if (err == cudaSuccess) return true; + if (error_message) *error_message = cudaGetErrorString(err); + cudaGetLastError(); + return false; +} + +RegisteredPinnedMemoryManager::UnregisterResult UnregisterPinnedRegionWithCuda( + void* addr, std::string* error_message) { + cudaError_t err = cudaHostUnregister(addr); + if (err == cudaSuccess) { + return RegisteredPinnedMemoryManager::UnregisterResult::kSuccess; + } + if (error_message) *error_message = cudaGetErrorString(err); + if (err == cudaErrorCudartUnloading) { + return RegisteredPinnedMemoryManager::UnregisterResult:: + kRuntimeUnloading; + } + return RegisteredPinnedMemoryManager::UnregisterResult::kError; +} + +#endif + +RegisteredPinnedMemoryManager::PinOps DefaultPinOps() { +#if defined(USE_CUDA) + return {RegisterPinnedRegionWithCuda, UnregisterPinnedRegionWithCuda}; +#else + return {}; +#endif +} + +} // namespace + +RegisteredPinnedRegion::~RegisteredPinnedRegion() { release(); } + +bool RegisteredPinnedRegion::release() { + if (!manager_) return release_succeeded_; + auto* manager = manager_; + manager_ = nullptr; + release_succeeded_ = manager->release(this); + return release_succeeded_; +} + +RegisteredPinnedMemoryManager& RegisteredPinnedMemoryManager::instance() { + static RegisteredPinnedMemoryManager* manager = + new RegisteredPinnedMemoryManager(); + return *manager; +} + +RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager() + : RegisteredPinnedMemoryManager(ParsePinnedMemoryConfig(), + DefaultPinOps()) {} + +RegisteredPinnedMemoryManager::RegisteredPinnedMemoryManager( + std::pair config, PinOps pin_ops) + : enabled_(config.first), limit_bytes_(config.second), pin_ops_(pin_ops) { +#if defined(USE_CUDA) + LOG(INFO) << "Store segment pinned memory is " + << (enabled_ ? "enabled" : "disabled") + << ", max_bytes=" << limit_bytes_; +#else + if (enabled_) { + LOG(INFO) << "Store segment pinning requested but this build has no " + "CUDA runtime support"; + } +#endif +} + +std::shared_ptr RegisteredPinnedMemoryManager::try_pin( + void* addr, size_t size, const std::string& owner) { + if (!addr || size == 0 || !enabled_) return nullptr; + if (!pin_ops_.register_region || !pin_ops_.unregister_region) { + return nullptr; + } + + const auto start = reinterpret_cast(addr); + const auto end = start + size; + if (end < start) { + LogPinSkip(owner, "address range overflow", size); + return nullptr; + } + + std::shared_ptr region; + try { + region.reset(new RegisteredPinnedRegion(this, addr, size)); + } catch (...) { + LogPinSkip(owner, "failed to allocate pin tracking", size); + return nullptr; + } + + { + std::lock_guard lock(mutex_); + for (const auto& entry : regions_) { + const auto region_start = reinterpret_cast(entry.addr); + const auto region_end = region_start + entry.size; + const bool overlaps = start < region_end && end > region_start; + if (overlaps) { + LogPinSkip(owner, "overlaps an active pinned region", size); + return nullptr; + } + } + + if (size > limit_bytes_ || pinned_bytes_ > limit_bytes_ - size) { + LOG(WARNING) << "Skip cudaHostRegister for " << owner + << ": quota exceeded, requested=" << size + << ", pinned=" << pinned_bytes_ + << ", limit=" << limit_bytes_; + return nullptr; + } + + try { + regions_.push_back({addr, size, nullptr}); + } catch (...) { + LogPinSkip(owner, "failed to allocate pin tracking", size); + return nullptr; + } + pinned_bytes_ += size; + } + + std::string error_message; + const bool registered = + pin_ops_.register_region(addr, size, &error_message); + if (!registered) { + { + std::lock_guard lock(mutex_); + remove_inactive_region_locked(addr, size); + } + LOG(WARNING) << "cudaHostRegister failed for " << owner + << ", size=" << size << ", error=" << error_message + << ". Continue with pageable host memory."; + return nullptr; + } + + uint64_t pinned_bytes = 0; + { + std::lock_guard lock(mutex_); + for (auto& entry : regions_) { + if (entry.addr == addr && entry.size == size && !entry.region) { + entry.region = region.get(); + pinned_bytes = pinned_bytes_; + break; + } + } + } + + LOG(INFO) << "cudaHostRegister succeeded for " << owner << ", size=" << size + << ", pinned=" << pinned_bytes << ", limit=" << limit_bytes_; + return region; +} + +bool RegisteredPinnedMemoryManager::release(RegisteredPinnedRegion* region) { + if (!region || !region->addr_ || region->size_ == 0) return true; + + bool should_unregister = false; + { + std::lock_guard lock(mutex_); + for (auto& entry : regions_) { + if (entry.addr == region->addr_ && entry.size == region->size_ && + entry.region == region) { + entry.region = nullptr; + should_unregister = true; + break; + } + } + } + if (!should_unregister) return true; + + std::string error_message; + auto unregister_result = + pin_ops_.unregister_region(region->addr_, &error_message); + if (unregister_result != UnregisterResult::kSuccess) { + if (unregister_result == UnregisterResult::kRuntimeUnloading) { + LOG(WARNING) << "Skip cudaHostUnregister because CUDA runtime " + "is unloading, size=" + << region->size_; + } else { + LOG(ERROR) << "cudaHostUnregister failed, size=" << region->size_ + << ", error=" << error_message + << ". Keep the range reserved; backing memory must not " + "be freed."; + return false; + } + } + + std::lock_guard lock(mutex_); + remove_inactive_region_locked(region->addr_, region->size_); + return true; +} + +void RegisteredPinnedMemoryManager::remove_inactive_region_locked(void* addr, + size_t size) { + for (auto it = regions_.begin(); it != regions_.end(); ++it) { + if (it->addr == addr && it->size == size && !it->region) { + regions_.erase(it); + pinned_bytes_ = pinned_bytes_ >= size ? pinned_bytes_ - size : 0; + return; + } + } +} + +} // namespace mooncake diff --git a/mooncake-store/src/registered_pinned_memory.h b/mooncake-store/src/registered_pinned_memory.h new file mode 100644 index 0000000000..cbc43c1388 --- /dev/null +++ b/mooncake-store/src/registered_pinned_memory.h @@ -0,0 +1,83 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace mooncake { + +class RegisteredPinnedMemoryManager; + +class RegisteredPinnedRegion { + public: + RegisteredPinnedRegion(const RegisteredPinnedRegion&) = delete; + RegisteredPinnedRegion& operator=(const RegisteredPinnedRegion&) = delete; + ~RegisteredPinnedRegion(); + + bool release(); + + private: + friend class RegisteredPinnedMemoryManager; + + RegisteredPinnedRegion(RegisteredPinnedMemoryManager* manager, void* addr, + size_t size) + : manager_(manager), addr_(addr), size_(size) {} + + RegisteredPinnedMemoryManager* manager_ = nullptr; + void* addr_ = nullptr; + size_t size_ = 0; + bool release_succeeded_ = true; +}; + +class RegisteredPinnedMemoryManager { + public: + enum class UnregisterResult { kSuccess, kRuntimeUnloading, kError }; + + struct PinOps { + bool (*register_region)(void* addr, size_t size, + std::string* error_message) = nullptr; + UnregisterResult (*unregister_region)( + void* addr, std::string* error_message) = nullptr; + }; + + static RegisteredPinnedMemoryManager& instance(); + + std::shared_ptr try_pin(void* addr, size_t size, + const std::string& owner); + + private: + friend class RegisteredPinnedRegion; + + struct ActiveRegion { + void* addr; + size_t size; + RegisteredPinnedRegion* region; + }; + + RegisteredPinnedMemoryManager(); +#if defined(MOONCAKE_STORE_TEST) + public: +#endif + RegisteredPinnedMemoryManager(std::pair config, + PinOps pin_ops); +#if defined(MOONCAKE_STORE_TEST) + private: +#endif + + bool release(RegisteredPinnedRegion* region); + void remove_inactive_region_locked(void* addr, size_t size); + + const bool enabled_; + const uint64_t limit_bytes_; + const PinOps pin_ops_; + + mutable std::mutex mutex_; + uint64_t pinned_bytes_ = 0; + std::vector regions_; +}; + +} // namespace mooncake diff --git a/mooncake-store/src/rpc_service.cpp b/mooncake-store/src/rpc_service.cpp index 386b3cdef0..c6f3a98604 100644 --- a/mooncake-store/src/rpc_service.cpp +++ b/mooncake-store/src/rpc_service.cpp @@ -1,5 +1,10 @@ #include "rpc_service.h" +#include +#include +#include +#include + #include #include @@ -11,10 +16,80 @@ #include "version.h" namespace mooncake { +namespace { + +tl::expected ResolveRequestTenantId(std::string_view raw) { + TenantId tenant_id{std::string(raw)}; + if (!tenant_id.IsValid()) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + return tenant_id; +} + +tl::expected ResolveTenantIdForWrite( + std::string_view raw, bool enable_multi_tenants) { + if (!enable_multi_tenants) { + return TenantId::Default(); + } + if (raw.empty()) { + return tl::make_unexpected(ErrorCode::TENANT_NOT_REGISTERED); + } + TenantId tenant_id{std::string(raw)}; + if (!tenant_id.IsValid()) { + return tl::make_unexpected(ErrorCode::TENANT_NOT_REGISTERED); + } + return tenant_id; +} + +template +auto WithRequestTenant(std::string_view raw, Fn&& fn) { + using Result = std::invoke_result_t; + auto tenant_id = ResolveRequestTenantId(raw); + if (!tenant_id) { + return Result(tl::make_unexpected(tenant_id.error())); + } + return std::invoke(std::forward(fn), tenant_id.value()); +} + +template +auto WithRequestTenantBatch(std::string_view raw, size_t result_count, + Fn&& fn) { + using Results = std::invoke_result_t; + auto tenant_id = ResolveRequestTenantId(raw); + if (!tenant_id) { + using Result = typename Results::value_type; + return Results(result_count, + Result(tl::make_unexpected(tenant_id.error()))); + } + return std::invoke(std::forward(fn), tenant_id.value()); +} + +template +auto WithWriteTenant(std::string_view raw, bool enable_multi_tenants, Fn&& fn) { + using Result = std::invoke_result_t; + auto tenant_id = ResolveTenantIdForWrite(raw, enable_multi_tenants); + if (!tenant_id) { + return Result(tl::make_unexpected(tenant_id.error())); + } + return std::invoke(std::forward(fn), tenant_id.value()); +} + +} // namespace WrappedMasterService::WrappedMasterService( - const WrappedMasterServiceConfig& config) - : master_service_(MasterServiceConfig(config)) {} + const WrappedMasterServiceConfig& config, + HttpMetadataServer* http_metadata_server, + const std::string& http_metadata_remote_url) + : master_service_(MasterServiceConfig(config)) { + // Configure metadata cleanup on client timeout. Prefer the co-located + // in-process server; otherwise fall back to a separately-deployed HTTP + // metadata server derived from the cluster configuration. + if (http_metadata_server) { + master_service_.setHttpMetadataServer(http_metadata_server); + } else if (!http_metadata_remote_url.empty()) { + master_service_.setHttpMetadataRemoteUrl(http_metadata_remote_url); + } +} WrappedMasterService::~WrappedMasterService() = default; @@ -26,7 +101,16 @@ WrappedMasterService::CalcCacheStats() { tl::expected WrappedMasterService::ExistKey( const std::string& key, const std::string& tenant_id) { return execute_rpc( - "ExistKey", [&] { return master_service_.ExistKey(key, tenant_id); }, + "ExistKey", + [&] { + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.ExistKey( + key, resolved_tenant_id); + }); + }, [&](auto& timer) { timer.LogRequest("key=", key); }, [] { MasterMetricManager::instance().inc_exist_key_requests(); }, [] { MasterMetricManager::instance().inc_exist_key_failures(); }); @@ -39,7 +123,12 @@ std::vector> WrappedMasterService::BatchExistKey( timer.LogRequest("keys_count=", total_keys); MasterMetricManager::instance().inc_batch_exist_key_requests(total_keys); - auto result = master_service_.BatchExistKey(keys, tenant_id); + auto result = WithRequestTenantBatch( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + keys.size(), [&](const TenantId& resolved_tenant_id) { + return master_service_.BatchExistKey(keys, resolved_tenant_id); + }); size_t failure_count = 0; for (size_t i = 0; i < result.size(); ++i) { @@ -150,7 +239,16 @@ WrappedMasterService::GetReplicaListByRegex(const std::string& str, const std::string& tenant_id) { return execute_rpc( "GetReplicaListByRegex", - [&] { return master_service_.GetReplicaListByRegex(str, tenant_id); }, + [&] { + return WithRequestTenant( + master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.GetReplicaListByRegex( + str, resolved_tenant_id); + }); + }, [&](auto& timer) { timer.LogRequest("Regex=", str); }, [] { MasterMetricManager::instance() @@ -167,7 +265,15 @@ WrappedMasterService::GetReplicaList(const std::string& key, const std::string& tenant_id) { return execute_rpc( "GetReplicaList", - [&] { return master_service_.GetReplicaList(key, tenant_id); }, + [&] { + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.GetReplicaList( + key, resolved_tenant_id); + }); + }, [&](auto& timer) { timer.LogRequest("key=", key); }, [] { MasterMetricManager::instance().inc_get_replica_list_requests(); }, [] { @@ -187,7 +293,13 @@ WrappedMasterService::BatchGetReplicaList(const std::vector& keys, std::vector> results; results.reserve(keys.size()); - results = master_service_.BatchGetReplicaList(keys, tenant_id); + results = WithRequestTenantBatch( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + keys.size(), [&](const TenantId& resolved_tenant_id) { + return master_service_.BatchGetReplicaList(keys, + resolved_tenant_id); + }); size_t failure_count = 0; for (size_t i = 0; i < results.size(); ++i) { @@ -219,6 +331,36 @@ WrappedMasterService::BatchGetReplicaList(const std::vector& keys, return results; } +std::vector> +WrappedMasterService::BatchGetReplicaListForAdmin( + const std::vector& keys, const std::string& tenant_id) { + return WithRequestTenantBatch( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + keys.size(), [&](const TenantId& resolved_tenant_id) { + return master_service_.BatchGetReplicaListForAdmin( + keys, resolved_tenant_id); + }); +} + +tl::expected +WrappedMasterService::GetReplicaListForAdmin(const std::string& key, + const std::string& tenant_id) { + return execute_rpc( + "GetReplicaListForAdmin", + [&] { + return WithRequestTenant( + master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.GetReplicaListForAdmin( + key, resolved_tenant_id); + }); + }, + [&](auto& timer) { timer.LogRequest("key=", key); }, [] {}, [] {}); +} + tl::expected, ErrorCode> WrappedMasterService::PutStart(const UUID& client_id, const std::string& key, const uint64_t slice_length, @@ -227,8 +369,13 @@ WrappedMasterService::PutStart(const UUID& client_id, const std::string& key, return execute_rpc( "PutStart", [&] { - return master_service_.PutStart(client_id, key, tenant_id, - slice_length, config); + return WithWriteTenant(tenant_id, + master_service_.IsTenantQuotaEnabled(), + [&](const TenantId& resolved_tenant_id) { + return master_service_.PutStart( + client_id, key, resolved_tenant_id, + slice_length, config); + }); }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, @@ -239,16 +386,22 @@ WrappedMasterService::PutStart(const UUID& client_id, const std::string& key, } tl::expected WrappedMasterService::PutEnd( - const UUID& client_id, const std::string& key, ReplicaType replica_type, - const std::string& tenant_id) { + const UUID& client_id, const ObjectMeta& object_meta, + ReplicaType replica_type, const std::string& tenant_id) { return execute_rpc( "PutEnd", [&] { - return master_service_.PutEnd(client_id, key, tenant_id, - replica_type); + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.PutEnd( + client_id, object_meta, + resolved_tenant_id, replica_type); + }); }, [&](auto& timer) { - timer.LogRequest("client_id=", client_id, ", key=", key, + timer.LogRequest("client_id=", client_id, ", key=", object_meta.key, ", replica_type=", replica_type); }, [] { MasterMetricManager::instance().inc_put_end_requests(); }, @@ -261,8 +414,14 @@ tl::expected WrappedMasterService::PutRevoke( return execute_rpc( "PutRevoke", [&] { - return master_service_.PutRevoke(client_id, key, tenant_id, + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.PutRevoke( + client_id, key, resolved_tenant_id, replica_type); + }); }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, @@ -286,6 +445,8 @@ WrappedMasterService::BatchPutStart(const UUID& client_id, std::vector, ErrorCode>> results; results.reserve(keys.size()); + auto resolved_tenant_id = ResolveTenantIdForWrite( + tenant_id, master_service_.IsTenantQuotaEnabled()); if (keys.size() != slice_lengths.size()) { LOG(ERROR) << "BatchPutStart: keys.size()=" << keys.size() @@ -299,12 +460,16 @@ WrappedMasterService::BatchPutStart(const UUID& client_id, << " != keys.size()=" << keys.size(); results.assign(keys.size(), tl::make_unexpected(ErrorCode::INVALID_PARAMS)); + } else if (!resolved_tenant_id) { + results.assign(keys.size(), + tl::make_unexpected(resolved_tenant_id.error())); } else if (config.prefer_alloc_in_same_node) { ReplicateConfig new_config = config; for (size_t i = 0; i < keys.size(); ++i) { auto key_config = new_config.ForSingleKey(i); auto result = master_service_.PutStart( - client_id, keys[i], tenant_id, slice_lengths[i], key_config); + client_id, keys[i], resolved_tenant_id.value(), + slice_lengths[i], key_config); results.emplace_back(result); if ((i == 0) && result.has_value()) { std::string preferred_segment; @@ -326,7 +491,8 @@ WrappedMasterService::BatchPutStart(const UUID& client_id, for (size_t i = 0; i < keys.size(); ++i) { auto key_config = config.ForSingleKey(i); results.emplace_back(master_service_.PutStart( - client_id, keys[i], tenant_id, slice_lengths[i], key_config)); + client_id, keys[i], resolved_tenant_id.value(), + slice_lengths[i], key_config)); } } @@ -368,28 +534,28 @@ WrappedMasterService::BatchPutStart(const UUID& client_id, } std::vector> WrappedMasterService::BatchPutEnd( - const UUID& client_id, const std::vector& keys, + const UUID& client_id, const std::vector& object_metas, ReplicaType replica_type, const std::string& tenant_id) { ScopedVLogTimer timer(1, "BatchPutEnd"); - const size_t total_keys = keys.size(); + const size_t total_keys = object_metas.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); MasterMetricManager::instance().inc_batch_put_end_requests(total_keys); - std::vector> results; - results.reserve(keys.size()); - - for (const auto& key : keys) { - results.emplace_back( - master_service_.PutEnd(client_id, key, tenant_id, replica_type)); - } + auto results = WithRequestTenantBatch( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + object_metas.size(), [&](const TenantId& resolved_tenant_id) { + return master_service_.BatchPutEnd( + client_id, object_metas, resolved_tenant_id, replica_type); + }); size_t failure_count = 0; for (size_t i = 0; i < results.size(); ++i) { if (!results[i].has_value()) { failure_count++; auto error = results[i].error(); - LOG(ERROR) << "BatchPutEnd failed for key[" << i << "] '" << keys[i] - << "': " << toString(error); + LOG(ERROR) << "BatchPutEnd failed for key[" << i << "] '" + << object_metas[i].key << "': " << toString(error); } } @@ -415,13 +581,18 @@ std::vector> WrappedMasterService::BatchPutRevoke( timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); MasterMetricManager::instance().inc_batch_put_revoke_requests(total_keys); - std::vector> results; - results.reserve(keys.size()); - - for (const auto& key : keys) { - results.emplace_back( - master_service_.PutRevoke(client_id, key, tenant_id, replica_type)); - } + auto results = WithRequestTenantBatch( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + keys.size(), [&](const TenantId& resolved_tenant_id) { + std::vector> batch_results; + batch_results.reserve(keys.size()); + for (const auto& key : keys) { + batch_results.emplace_back(master_service_.PutRevoke( + client_id, key, resolved_tenant_id, replica_type)); + } + return batch_results; + }); size_t failure_count = 0; for (size_t i = 0; i < results.size(); ++i) { @@ -455,8 +626,13 @@ WrappedMasterService::UpsertStart(const UUID& client_id, const std::string& key, return execute_rpc( "UpsertStart", [&] { - return master_service_.UpsertStart(client_id, key, tenant_id, - slice_length, config); + return WithWriteTenant(tenant_id, + master_service_.IsTenantQuotaEnabled(), + [&](const TenantId& resolved_tenant_id) { + return master_service_.UpsertStart( + client_id, key, resolved_tenant_id, + slice_length, config); + }); }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, @@ -467,16 +643,22 @@ WrappedMasterService::UpsertStart(const UUID& client_id, const std::string& key, } tl::expected WrappedMasterService::UpsertEnd( - const UUID& client_id, const std::string& key, ReplicaType replica_type, - const std::string& tenant_id) { + const UUID& client_id, const ObjectMeta& object_meta, + ReplicaType replica_type, const std::string& tenant_id) { return execute_rpc( "UpsertEnd", [&] { - return master_service_.UpsertEnd(client_id, key, tenant_id, - replica_type); + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.UpsertEnd( + client_id, object_meta, + resolved_tenant_id, replica_type); + }); }, [&](auto& timer) { - timer.LogRequest("client_id=", client_id, ", key=", key, + timer.LogRequest("client_id=", client_id, ", key=", object_meta.key, ", replica_type=", replica_type); }, [] { MasterMetricManager::instance().inc_put_end_requests(); }, @@ -489,8 +671,14 @@ tl::expected WrappedMasterService::UpsertRevoke( return execute_rpc( "UpsertRevoke", [&] { - return master_service_.UpsertRevoke(client_id, key, tenant_id, - replica_type); + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.UpsertRevoke( + client_id, key, resolved_tenant_id, + replica_type); + }); }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, @@ -510,8 +698,29 @@ WrappedMasterService::BatchUpsertStart( timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); MasterMetricManager::instance().inc_batch_put_start_requests(total_keys); - auto results = master_service_.BatchUpsertStart(client_id, keys, tenant_id, - slice_lengths, config); + std::vector, ErrorCode>> + results; + if (keys.size() != slice_lengths.size()) { + LOG(ERROR) << "BatchUpsertStart: keys.size()=" << keys.size() + << " != slice_lengths.size()=" << slice_lengths.size(); + results.assign(keys.size(), + tl::make_unexpected(ErrorCode::INVALID_PARAMS)); + } else if (config.group_ids.has_value() && + config.group_ids->size() != keys.size()) { + LOG(ERROR) << "BatchUpsertStart: group_ids.size()=" + << config.group_ids->size() + << " != keys.size()=" << keys.size(); + results.assign(keys.size(), + tl::make_unexpected(ErrorCode::INVALID_PARAMS)); + } else if (auto resolved_tenant_id = ResolveTenantIdForWrite( + tenant_id, master_service_.IsTenantQuotaEnabled()); + !resolved_tenant_id) { + results.assign(keys.size(), + tl::make_unexpected(resolved_tenant_id.error())); + } else { + results = master_service_.BatchUpsertStart( + client_id, keys, resolved_tenant_id.value(), slice_lengths, config); + } size_t failure_count = 0; for (size_t i = 0; i < results.size(); ++i) { @@ -538,14 +747,20 @@ WrappedMasterService::BatchUpsertStart( } std::vector> WrappedMasterService::BatchUpsertEnd( - const UUID& client_id, const std::vector& keys, + const UUID& client_id, const std::vector& object_metas, const std::string& tenant_id) { ScopedVLogTimer timer(1, "BatchUpsertEnd"); - const size_t total_keys = keys.size(); + const size_t total_keys = object_metas.size(); timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); MasterMetricManager::instance().inc_batch_put_end_requests(total_keys); - auto results = master_service_.BatchUpsertEnd(client_id, keys, tenant_id); + auto results = WithRequestTenantBatch( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + object_metas.size(), [&](const TenantId& resolved_tenant_id) { + return master_service_.BatchUpsertEnd(client_id, object_metas, + resolved_tenant_id); + }); size_t failure_count = 0; for (size_t i = 0; i < results.size(); ++i) { @@ -553,7 +768,7 @@ std::vector> WrappedMasterService::BatchUpsertEnd( failure_count++; auto error = results[i].error(); LOG(ERROR) << "BatchUpsertEnd failed for key[" << i << "] '" - << keys[i] << "': " << toString(error); + << object_metas[i].key << "': " << toString(error); } } @@ -580,8 +795,13 @@ WrappedMasterService::BatchUpsertRevoke(const UUID& client_id, timer.LogRequest("client_id=", client_id, ", keys_count=", total_keys); MasterMetricManager::instance().inc_batch_put_revoke_requests(total_keys); - auto results = - master_service_.BatchUpsertRevoke(client_id, keys, tenant_id); + auto results = WithRequestTenantBatch( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + keys.size(), [&](const TenantId& resolved_tenant_id) { + return master_service_.BatchUpsertRevoke(client_id, keys, + resolved_tenant_id); + }); size_t failure_count = 0; for (size_t i = 0; i < results.size(); ++i) { @@ -610,7 +830,16 @@ WrappedMasterService::BatchUpsertRevoke(const UUID& client_id, tl::expected WrappedMasterService::Remove( const std::string& key, bool force, const std::string& tenant_id) { return execute_rpc( - "Remove", [&] { return master_service_.Remove(key, tenant_id, force); }, + "Remove", + [&] { + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.Remove( + key, resolved_tenant_id, force); + }); + }, [&](auto& timer) { timer.LogRequest("key=", key, ", force=", force); }, [] { MasterMetricManager::instance().inc_remove_requests(); }, [] { MasterMetricManager::instance().inc_remove_failures(); }); @@ -620,7 +849,15 @@ tl::expected WrappedMasterService::RemoveByRegex( const std::string& str, bool force, const std::string& tenant_id) { return execute_rpc( "RemoveByRegex", - [&] { return master_service_.RemoveByRegex(str, tenant_id, force); }, + [&] { + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.RemoveByRegex( + str, resolved_tenant_id, force); + }); + }, [&](auto& timer) { timer.LogRequest("regex=", str, ", force=", force); }, @@ -630,9 +867,28 @@ tl::expected WrappedMasterService::RemoveByRegex( long WrappedMasterService::RemoveAll(bool force, const std::string& tenant_id) { ScopedVLogTimer timer(1, "RemoveAll"); - timer.LogRequest("action=remove_all_objects, force=", force); + timer.LogRequest("action=remove_all_objects, force=", force, + ", tenant_id=", tenant_id); MasterMetricManager::instance().inc_remove_all_requests(); - long result = master_service_.RemoveAll(tenant_id, force); + // Empty tenant_id => clear ALL tenants (broadcast SSD signal to every + // client, overlapping with metadata deletion). A specific tenant_id => + // scoped clear (only signal clients owning that tenant's disk replicas). + if (tenant_id.empty()) { + long result = master_service_.RemoveAll(force); + timer.LogResponse("items_removed=", result); + return result; + } + auto resolved_tenant_id = ResolveRequestTenantId( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue); + if (!resolved_tenant_id) { + // Keep the legacy scalar wire result; unlike the other tenant-bearing + // RPCs, RemoveAll cannot propagate an ErrorCode without a wire change. + LOG(WARNING) << "RemoveAll rejected an invalid tenant id"; + timer.LogResponse("error=", toString(resolved_tenant_id.error())); + return 0; + } + long result = master_service_.RemoveAll(*resolved_tenant_id, force); timer.LogResponse("items_removed=", result); return result; } @@ -645,7 +901,12 @@ std::vector> WrappedMasterService::BatchRemove( timer.LogRequest("keys_count=", total_keys, ", force=", force); MasterMetricManager::instance().inc_remove_requests(total_keys); - auto results = master_service_.BatchRemove(keys, tenant_id, force); + auto results = WithRequestTenantBatch( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + keys.size(), [&](const TenantId& resolved_tenant_id) { + return master_service_.BatchRemove(keys, resolved_tenant_id, force); + }); size_t failure_count = 0; for (const auto& result : results) { @@ -804,8 +1065,13 @@ tl::expected WrappedMasterService::CopyStart( return execute_rpc( "CopyStart", [&] { - return master_service_.CopyStart(client_id, key, tenant_id, - src_segment, tgt_segments); + return WithWriteTenant(tenant_id, + master_service_.IsTenantQuotaEnabled(), + [&](const TenantId& resolved_tenant_id) { + return master_service_.CopyStart( + client_id, key, resolved_tenant_id, + src_segment, tgt_segments); + }); }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, @@ -822,7 +1088,16 @@ tl::expected WrappedMasterService::CopyEnd( const std::string& tenant_id) { return execute_rpc( "CopyEnd", - [&] { return master_service_.CopyEnd(client_id, key, tenant_id); }, + [&] { + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.CopyEnd( + client_id, key, + resolved_tenant_id); + }); + }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, ", tenant_id=", tenant_id); @@ -836,7 +1111,16 @@ tl::expected WrappedMasterService::CopyRevoke( const std::string& tenant_id) { return execute_rpc( "CopyRevoke", - [&] { return master_service_.CopyRevoke(client_id, key, tenant_id); }, + [&] { + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.CopyRevoke( + client_id, key, + resolved_tenant_id); + }); + }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, ", tenant_id=", tenant_id); @@ -851,8 +1135,13 @@ tl::expected WrappedMasterService::MoveStart( return execute_rpc( "MoveStart", [&] { - return master_service_.MoveStart(client_id, key, tenant_id, - src_segment, tgt_segment); + return WithWriteTenant(tenant_id, + master_service_.IsTenantQuotaEnabled(), + [&](const TenantId& resolved_tenant_id) { + return master_service_.MoveStart( + client_id, key, resolved_tenant_id, + src_segment, tgt_segment); + }); }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, @@ -869,7 +1158,16 @@ tl::expected WrappedMasterService::MoveEnd( const std::string& tenant_id) { return execute_rpc( "MoveEnd", - [&] { return master_service_.MoveEnd(client_id, key, tenant_id); }, + [&] { + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.MoveEnd( + client_id, key, + resolved_tenant_id); + }); + }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, ", tenant_id=", tenant_id); @@ -883,7 +1181,16 @@ tl::expected WrappedMasterService::MoveRevoke( const std::string& tenant_id) { return execute_rpc( "MoveRevoke", - [&] { return master_service_.MoveRevoke(client_id, key, tenant_id); }, + [&] { + return WithRequestTenant(master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.MoveRevoke( + client_id, key, + resolved_tenant_id); + }); + }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, ", tenant_id=", tenant_id); @@ -898,8 +1205,14 @@ tl::expected WrappedMasterService::EvictDiskReplica( return execute_rpc( "EvictDiskReplica", [&] { - return master_service_.EvictDiskReplica(client_id, key, tenant_id, - replica_type); + return WithRequestTenant( + master_service_.IsTenantQuotaEnabled() + ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.EvictDiskReplica( + client_id, key, resolved_tenant_id, replica_type); + }); }, [&](auto& timer) { timer.LogRequest("client_id=", client_id, ", key=", key, @@ -925,8 +1238,13 @@ WrappedMasterService::BatchEvictDiskReplica( ", replica_type=", replica_type); MasterMetricManager::instance().inc_evict_disk_replica_requests(); - auto results = master_service_.BatchEvictDiskReplica( - client_id, keys, tenant_id, replica_type); + auto results = WithRequestTenantBatch( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + keys.size(), [&](const TenantId& resolved_tenant_id) { + return master_service_.BatchEvictDiskReplica( + client_id, keys, resolved_tenant_id, replica_type); + }); size_t failure_count = 0; for (size_t i = 0; i < results.size(); ++i) { @@ -952,7 +1270,14 @@ tl::expected WrappedMasterService::CreateCopyTask( const std::vector& targets) { return execute_rpc( "CreateCopyTask", - [&] { return master_service_.CreateCopyTask(key, tenant_id, targets); }, + [&] { + return WithWriteTenant(tenant_id, + master_service_.IsTenantQuotaEnabled(), + [&](const TenantId& resolved_tenant_id) { + return master_service_.CreateCopyTask( + key, resolved_tenant_id, targets); + }); + }, [&](auto& timer) { timer.LogRequest("key=", key, ", tenant_id=", tenant_id, ", targets_size=", targets.size()); @@ -969,8 +1294,12 @@ tl::expected WrappedMasterService::CreateMoveTask( return execute_rpc( "CreateMoveTask", [&] { - return master_service_.CreateMoveTask(key, tenant_id, source, - target); + return WithWriteTenant( + tenant_id, master_service_.IsTenantQuotaEnabled(), + [&](const TenantId& resolved_tenant_id) { + return master_service_.CreateMoveTask( + key, resolved_tenant_id, source, target); + }); }, [&](auto& timer) { timer.LogRequest("key=", key, ", tenant_id=", tenant_id, @@ -1054,11 +1383,77 @@ tl::expected WrappedMasterService::ServiceReady() { return GetMooncakeStoreVersion(); } +tl::expected, ErrorCode> +WrappedMasterService::ListTenantQuotaSnapshots() { + if (!master_service_.IsTenantQuotaEnabled()) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + return master_service_.ListTenantQuotaSnapshots(); +} + +tl::expected +WrappedMasterService::GetTenantQuotaSnapshot(const std::string& tenant_id) { + if (!master_service_.IsTenantQuotaEnabled()) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + return WithRequestTenant( + tenant_id, + [&](const TenantId& resolved_tenant_id) + -> tl::expected { + auto snapshot = + master_service_.GetTenantQuotaSnapshot(resolved_tenant_id); + if (!snapshot.has_value()) { + return tl::make_unexpected(ErrorCode::OBJECT_NOT_FOUND); + } + return snapshot.value(); + }); +} + +tl::expected +WrappedMasterService::UpsertTenantQuotaPolicy(const std::string& tenant_id, + uint64_t requested_quota_bytes) { + if (!master_service_.IsTenantQuotaEnabled()) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + if (tenant_id.empty()) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + TenantId resolved_tenant_id(tenant_id); + if (!resolved_tenant_id.IsValid()) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + return master_service_.UpsertTenantQuotaPolicy(resolved_tenant_id, + requested_quota_bytes); +} + +tl::expected, ErrorCode> +WrappedMasterService::DeleteTenantQuotaPolicy(const std::string& tenant_id) { + if (!master_service_.IsTenantQuotaEnabled()) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + if (tenant_id.empty()) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + TenantId resolved_tenant_id(tenant_id); + if (!resolved_tenant_id.IsValid()) { + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + return master_service_.DeleteTenantQuotaPolicy(resolved_tenant_id); +} + +tl::expected +WrappedMasterService::GetTenantQuotaAllocatableCapacityBytes() { + if (!master_service_.IsTenantQuotaEnabled()) { + return tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + } + return master_service_.GetTenantQuotaAllocatableCapacityBytes(); +} + tl::expected, ErrorCode> WrappedMasterService::GetAllKeysForAdmin() { // Compatibility endpoint: /get_all_keys historically listed only the // default tenant's keys. - return master_service_.GetAllKeys("default"); + return master_service_.GetAllKeys(TenantId::Default()); } tl::expected, ErrorCode> @@ -1099,6 +1494,16 @@ WrappedMasterService::OffloadObjectHeartbeat(const UUID& client_id, return result; } +tl::expected WrappedMasterService::PollRemoveAll( + const UUID& client_id) { + ScopedVLogTimer timer(1, "PollRemoveAll"); + timer.LogRequest("client_id=", client_id); + auto result = master_service_.PollRemoveAll(client_id); + timer.LogResponse("should_remove_all=", + result.has_value() ? result.value() : false); + return result; +} + tl::expected WrappedMasterService::ReportSsdCapacity( const UUID& client_id, int64_t ssd_total_capacity_bytes) { ScopedVLogTimer timer(1, "ReportSsdCapacity"); @@ -1114,6 +1519,19 @@ tl::expected WrappedMasterService::NotifyOffloadSuccess( ScopedVLogTimer timer(1, "NotifyOffloadSuccess"); timer.LogRequest("action=notify_offload_success"); + for (const auto& task : tasks) { + auto tenant_id = + ResolveRequestTenantId(master_service_.IsTenantQuotaEnabled() + ? std::string_view(task.tenant_id) + : TenantId::kDefaultValue); + if (!tenant_id) { + auto result = tl::expected( + tl::make_unexpected(tenant_id.error())); + timer.LogResponseExpected(result); + return result; + } + } + auto result = master_service_.NotifyOffloadSuccess(client_id, tasks, metadatas); timer.LogResponseExpected(result); @@ -1133,8 +1551,12 @@ WrappedMasterService::PromotionAllocStart( uint64_t size, const std::vector& preferred_segments) { ScopedVLogTimer timer(1, "PromotionAllocStart"); timer.LogRequest("action=promotion_alloc_start"); - auto result = master_service_.PromotionAllocStart(client_id, key, tenant_id, - size, preferred_segments); + auto result = WithWriteTenant( + tenant_id, master_service_.IsTenantQuotaEnabled(), + [&](const TenantId& resolved_tenant_id) { + return master_service_.PromotionAllocStart( + client_id, key, resolved_tenant_id, size, preferred_segments); + }); timer.LogResponseExpected(result); return result; } @@ -1144,8 +1566,13 @@ tl::expected WrappedMasterService::NotifyPromotionSuccess( const std::string& tenant_id) { ScopedVLogTimer timer(1, "NotifyPromotionSuccess"); timer.LogRequest("action=notify_promotion_success"); - auto result = - master_service_.NotifyPromotionSuccess(client_id, key, tenant_id); + auto result = WithRequestTenant( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.NotifyPromotionSuccess(client_id, key, + resolved_tenant_id); + }); timer.LogResponseExpected(result); return result; } @@ -1155,8 +1582,13 @@ tl::expected WrappedMasterService::NotifyPromotionFailure( const std::string& tenant_id) { ScopedVLogTimer timer(1, "NotifyPromotionFailure"); timer.LogRequest("action=notify_promotion_failure"); - auto result = - master_service_.NotifyPromotionFailure(client_id, key, tenant_id); + auto result = WithRequestTenant( + master_service_.IsTenantQuotaEnabled() ? std::string_view(tenant_id) + : TenantId::kDefaultValue, + [&](const TenantId& resolved_tenant_id) { + return master_service_.NotifyPromotionFailure(client_id, key, + resolved_tenant_id); + }); timer.LogResponseExpected(result); return result; } @@ -1186,6 +1618,22 @@ WrappedMasterService::QuerySegmentStatusById(const UUID& segment_id) { return master_service_.QuerySegmentStatusById(segment_id); } +bool WrappedMasterService::KvEventsEnabled() const { + return master_service_.KvEventsEnabled(); +} + +KvEventPublisher::Stats WrappedMasterService::GetKvEventStats() const { + return master_service_.GetKvEventStats(); +} + +void WrappedMasterService::RestoreFromStandby( + const std::vector& objects, + uint64_t initial_oplog_sequence_id, + const std::vector& segments) { + master_service_.RestoreFromStandbySnapshot( + objects, initial_oplog_sequence_id, segments); +} + void RegisterRpcService( coro_rpc::coro_rpc_server& server, mooncake::WrappedMasterService& wrapped_master_service) { @@ -1311,6 +1759,8 @@ void RegisterRpcService( server.register_handler< &mooncake::WrappedMasterService::BatchEvictDiskReplica>( &wrapped_master_service); + server.register_handler<&mooncake::WrappedMasterService::PollRemoveAll>( + &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::CreateCopyTask>( &wrapped_master_service); server.register_handler<&mooncake::WrappedMasterService::CreateMoveTask>( diff --git a/mooncake-store/src/segment.cpp b/mooncake-store/src/segment.cpp index 57769c8c12..8ac943cc7a 100644 --- a/mooncake-store/src/segment.cpp +++ b/mooncake-store/src/segment.cpp @@ -3,6 +3,8 @@ #include "master_metric_manager.h" #include "utils/zstd_util.h" +#include + namespace mooncake { namespace { @@ -25,8 +27,88 @@ bool IsMsgpackInteger(const msgpack::object& object) { object.type == msgpack::type::NEGATIVE_INTEGER; } +void AddHostSegment(HostSegmentIndex& index, const Segment& segment) { + if (!segment.host_id.empty()) { + index[segment.host_id][segment.name].insert(segment.id); + } +} + +void RemoveHostSegment(HostSegmentIndex& index, const Segment& segment) { + if (segment.host_id.empty()) { + return; + } + auto host_it = index.find(segment.host_id); + if (host_it == index.end()) { + return; + } + auto name_it = host_it->second.find(segment.name); + if (name_it == host_it->second.end()) { + return; + } + name_it->second.erase(segment.id); + if (name_it->second.empty()) { + host_it->second.erase(name_it); + } + if (host_it->second.empty()) { + index.erase(host_it); + } +} + +std::vector BuildHostOrderedSegments( + const HostSegmentIndex& segments_by_host, const std::string& writer_host_id, + const std::string& key) { + std::vector ordered_segments; + if (writer_host_id.empty() || segments_by_host.empty()) { + return ordered_segments; + } + + auto start_it = segments_by_host.find(writer_host_id); + if (start_it == segments_by_host.end()) { + start_it = segments_by_host.lower_bound(writer_host_id); + if (start_it == segments_by_host.end()) { + start_it = segments_by_host.begin(); + } + } + + const size_t host_count = segments_by_host.size(); + auto host_it = start_it; + for (size_t host_idx = 0; host_idx < host_count; ++host_idx) { + const auto& segments_by_name = host_it->second; + if (!segments_by_name.empty()) { + std::vector host_segments; + host_segments.reserve(segments_by_name.size()); + for (const auto& [name, segment_ids] : segments_by_name) { + if (!segment_ids.empty()) { + host_segments.push_back(name); + } + } + const size_t start = + std::hash{}(key) % host_segments.size(); + for (size_t i = 0; i < host_segments.size(); ++i) { + ordered_segments.push_back( + host_segments[(start + i) % host_segments.size()]); + } + } + + ++host_it; + if (host_it == segments_by_host.end()) { + host_it = segments_by_host.begin(); + } + } + + return ordered_segments; +} + } // namespace +std::vector ScopedAllocatorAccess::GetHostOrderedSegments( + const std::string& writer_host_id, const std::string& key) const { + if (segments_by_host_ == nullptr) { + return {}; + } + return BuildHostOrderedSegments(*segments_by_host_, writer_host_id, key); +} + ErrorCode ScopedSegmentAccess::MountSegment(const Segment& segment, const UUID& client_id) { const uintptr_t buffer = segment.base; @@ -49,6 +131,7 @@ ErrorCode ScopedSegmentAccess::MountSegment(const Segment& segment, segment, SegmentStatus::OK, allocator}; segment_manager_->client_by_name_[segment.name] = client_id; segment_manager_->segment_id_by_name_[segment.name] = segment.id; + AddHostSegment(segment_manager_->segments_by_host_, segment); LOG(INFO) << "[CXL Segment Mounted Successfully] Segment name: " << segment.name @@ -130,6 +213,7 @@ ErrorCode ScopedSegmentAccess::MountSegment(const Segment& segment, segment, SegmentStatus::OK, std::move(allocator)}; segment_manager_->client_by_name_[segment.name] = client_id; segment_manager_->segment_id_by_name_[segment.name] = segment.id; + AddHostSegment(segment_manager_->segments_by_host_, segment); MasterMetricManager::instance().inc_total_mem_capacity(segment.name, size); return ErrorCode::OK; @@ -153,6 +237,10 @@ ErrorCode ScopedSegmentAccess::MountLocalDiskSegment(const UUID& client_id, ErrorCode ScopedSegmentAccess::ReMountSegment( const std::vector& segments, const UUID& client_id) { for (const auto& segment : segments) { + auto validation = ValidateRemountSegment(segment, client_id); + if (validation != ErrorCode::OK) { + return validation; + } ErrorCode err = MountSegment(segment, client_id); if (err == ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS || err == ErrorCode::INTERNAL_ERROR) { @@ -179,6 +267,74 @@ ErrorCode ScopedSegmentAccess::ReMountSegment( return ErrorCode::OK; } +ErrorCode ScopedSegmentAccess::ValidateRemountSegment( + const Segment& segment, const UUID& client_id) const { + auto mounted = segment_manager_->mounted_segments_.find(segment.id); + if (mounted == segment_manager_->mounted_segments_.end()) { + return ErrorCode::OK; + } + const auto owner = segment_manager_->client_segments_.find(client_id); + const bool owned = owner != segment_manager_->client_segments_.end() && + std::find(owner->second.begin(), owner->second.end(), + segment.id) != owner->second.end(); + const auto& authoritative = mounted->second.segment; + if (!owned || authoritative.id != segment.id || + authoritative.name != segment.name || + authoritative.base != segment.base || + authoritative.size != segment.size || + authoritative.te_endpoint != segment.te_endpoint || + authoritative.protocol != segment.protocol || + authoritative.host_id != segment.host_id) { + return ErrorCode::INVALID_PARAMS; + } + return ErrorCode::OK; +} + +bool ScopedSegmentAccess::GetSegment(const UUID& segment_id, + Segment& segment) const { + auto mounted = segment_manager_->mounted_segments_.find(segment_id); + if (mounted == segment_manager_->mounted_segments_.end()) { + return false; + } + segment = mounted->second.segment; + return true; +} + +bool ScopedSegmentAccess::ReplaceAllocators( + const std::vector& replacements) { + std::vector manager_replacements; + manager_replacements.reserve(replacements.size()); + for (const auto& replacement : replacements) { + auto mounted = + segment_manager_->mounted_segments_.find(replacement.segment_id); + if (mounted == segment_manager_->mounted_segments_.end() || + mounted->second.buf_allocator != replacement.expected || + !replacement.replacement) { + return false; + } + manager_replacements.push_back({mounted->second.segment.name, + replacement.expected, + replacement.replacement}); + } + if (!segment_manager_->allocator_manager_.replaceAllocators( + manager_replacements)) { + return false; + } + for (const auto& replacement : replacements) { + segment_manager_->mounted_segments_.at(replacement.segment_id) + .buf_allocator = replacement.replacement; + } + return true; +} + +std::shared_ptr ScopedSegmentAccess::GetAllocator( + const UUID& segment_id) const { + auto mounted = segment_manager_->mounted_segments_.find(segment_id); + return mounted == segment_manager_->mounted_segments_.end() + ? nullptr + : mounted->second.buf_allocator; +} + ErrorCode ScopedSegmentAccess::PrepareUnmountSegment( const UUID& segment_id, size_t& metrics_dec_capacity) { auto it = segment_manager_->mounted_segments_.find(segment_id); @@ -207,6 +363,7 @@ ErrorCode ScopedSegmentAccess::PrepareUnmountSegment( segment_manager_->allocator_manager_.removeAllocator(segment.name, allocator); } + RemoveHostSegment(segment_manager_->segments_by_host_, segment); // 2. Remove from mounted_segment mounted_segment.buf_allocator.reset(); @@ -252,6 +409,7 @@ ErrorCode ScopedSegmentAccess::PrepareGracefulUnmountSegment( segment_manager_->allocator_manager_.removeAllocator(segment.name, allocator); } + RemoveHostSegment(segment_manager_->segments_by_host_, segment); // Set the segment status to GRACEFULLY_UNMOUNTING mounted_segment.status = SegmentStatus::GRACEFULLY_UNMOUNTING; return ErrorCode::OK; @@ -286,6 +444,8 @@ ErrorCode ScopedSegmentAccess::CommitUnmountSegment( auto&& segment = segment_manager_->mounted_segments_.find(segment_id); if (segment != segment_manager_->mounted_segments_.end()) { segment_name = segment->second.segment.name; + RemoveHostSegment(segment_manager_->segments_by_host_, + segment->second.segment); auto segment_id_by_name_it = segment_manager_->segment_id_by_name_.find(segment_name); if (segment_id_by_name_it != @@ -303,6 +463,10 @@ ErrorCode ScopedSegmentAccess::CommitUnmountSegment( if (!is_cxl) { MasterMetricManager::instance().dec_total_mem_capacity( segment_name, metrics_dec_capacity); + // Remove per-segment metric labels entirely to avoid stale 0-value + // entries persisting in Prometheus output (e.g. after snapshot + // restore followed by client expiry / reaper cleanup). + MasterMetricManager::instance().remove_segment_metrics(segment_name); } return ErrorCode::OK; @@ -375,6 +539,12 @@ ErrorCode ScopedSegmentAccess::GetAllSegments( return ErrorCode::OK; } +std::vector ScopedSegmentAccess::GetHostOrderedSegments( + const std::string& writer_host_id, const std::string& key) const { + return BuildHostOrderedSegments(segment_manager_->segments_by_host_, + writer_host_id, key); +} + ErrorCode ScopedSegmentAccess::GetAllSegmentNames( std::vector& all_segment_names) { all_segment_names.clear(); @@ -586,7 +756,9 @@ SegmentSerializer::Serialize() { } std::sort(sorted_keys.begin(), sorted_keys.end()); - packer.pack_array(2 + sorted_keys.size() * 2); + // Trailing ssd_total_capacity_bytes so a restored master keeps the + // client-reported SSD capacity across a snapshot restore (#2783). + packer.pack_array(2 + sorted_keys.size() * 2 + 1); packer.pack(segment->enable_offloading); packer.pack(static_cast(sorted_keys.size())); @@ -598,6 +770,7 @@ SegmentSerializer::Serialize() { packer.pack(task.key); packer.pack(task.size); } + packer.pack(segment->ssd_total_capacity_bytes); } // Compress entire data @@ -657,6 +830,7 @@ tl::expected SegmentSerializer::Deserialize( // Clear existing data segment_manager_->mounted_segments_.clear(); segment_manager_->client_segments_.clear(); + segment_manager_->segments_by_host_.clear(); // Convert MessagePack map to regular map, use pointers for values to avoid // copying @@ -864,6 +1038,7 @@ tl::expected SegmentSerializer::Deserialize( // Rebuild segment indexes based on client_segments_ and mounted_segments_ segment_manager_->client_by_name_.clear(); segment_manager_->segment_id_by_name_.clear(); + segment_manager_->segments_by_host_.clear(); for (const auto& [client_id, segment_ids] : segment_manager_->client_segments_) { for (const auto& segment_id : segment_ids) { @@ -873,6 +1048,11 @@ tl::expected SegmentSerializer::Deserialize( client_id; segment_manager_->segment_id_by_name_[it->second.segment.name] = segment_id; + if (it->second.status == SegmentStatus::OK && + it->second.buf_allocator) { + AddHostSegment(segment_manager_->segments_by_host_, + it->second.segment); + } } } } @@ -979,15 +1159,25 @@ tl::expected SegmentSerializer::Deserialize( "deserialize local_disk_segments legacy " "offloading size is not integer")); } - auto [tenant_id, user_key] = - ParseTenantScopedStorageKey(key); + auto [tenant_id, user_key] = TenantId::ParseScopedKey(key); segment->offloading_objects[key] = - OffloadTaskItem{.tenant_id = std::move(tenant_id), + OffloadTaskItem{.tenant_id = tenant_id.value(), .key = std::move(user_key), .size = task_obj.as()}; } } + // ssd_total_capacity_bytes is appended after the offloading + // objects. Pre-#2783 snapshots omit it, so read it only when + // present; otherwise it keeps the default 0 until the client + // re-reports capacity. + size_t capacity_idx = 2 + count * 2; + if (client_value.via.array.size > capacity_idx && + IsMsgpackInteger(client_value.via.array.ptr[capacity_idx])) { + segment->ssd_total_capacity_bytes = + client_value.via.array.ptr[capacity_idx].as(); + } + segment_manager_->client_local_disk_segment_[client_id] = std::move(segment); } @@ -1001,6 +1191,7 @@ void SegmentSerializer::Reset() { segment_manager_->client_segments_.clear(); segment_manager_->client_by_name_.clear(); segment_manager_->segment_id_by_name_.clear(); + segment_manager_->segments_by_host_.clear(); segment_manager_->client_local_disk_segment_.clear(); segment_manager_->allocator_manager_ = AllocatorManager(); } @@ -1092,8 +1283,12 @@ ErrorCode ScopedSegmentAccess::SetSegmentStatusByName( HasAllocator(allocator_manager, name, allocator); if (should_be_allocatable && !is_allocatable && allocator) { allocator_manager.addAllocator(name, allocator); + AddHostSegment(segment_manager_->segments_by_host_, + mounted_segment.segment); } else if (!should_be_allocatable && is_allocatable) { allocator_manager.removeAllocator(name, allocator); + RemoveHostSegment(segment_manager_->segments_by_host_, + mounted_segment.segment); } mounted_segment.status = status; @@ -1283,6 +1478,7 @@ ErrorCode ScopedNoFSegmentAccess::CommitUnmountSegment( nof_segment_manager_->mounted_segments_.erase(segment_id); MasterMetricManager::instance().dec_total_nof_capacity( segment_name, metrics_dec_capacity); + MasterMetricManager::instance().remove_nof_segment_metrics(segment_name); return ErrorCode::OK; } @@ -1366,6 +1562,25 @@ void NoFSegmentManager::GetMountedSegmentsSnapshot( } } +void SegmentManager::releaseCapacityMetrics() { + // Segments that are still mounted here never went through + // CommitUnmountSegment, so their contribution to the capacity metrics + // has not been released. MasterMetricManager outlives MasterService + // instances (a new one is constructed per HA leadership term), so the + // serving instance releases it at teardown to keep the gauges + // consistent with the segments that are actually mounted. + for (const auto& [segment_id, mounted_segment] : mounted_segments_) { + const auto& segment = mounted_segment.segment; + if (segment.protocol == "cxl") { + // CXL mounts do not contribute to total_mem_capacity. + continue; + } + MasterMetricManager::instance().dec_total_mem_capacity(segment.name, + segment.size); + MasterMetricManager::instance().remove_segment_metrics(segment.name); + } +} + void SegmentManager::initializeCxlAllocator(const std::string& cxl_path, const size_t cxl_size) { LOG(INFO) << "Init CXL global allocator."; @@ -1379,4 +1594,54 @@ void SegmentManager::initializeCxlAllocator(const std::string& cxl_path, cxl_path, DEFAULT_CXL_BASE, cxl_size, cxl_path); MasterMetricManager::instance().inc_total_mem_capacity(cxl_path, cxl_size); } + +int64_t ScopedLocalDiskSegmentAccess::getSsdTotalCapacity( + const std::string& segment_name) const { + auto client_it = client_by_name_.find(segment_name); + if (client_it == client_by_name_.end()) { + return 0; + } + auto disk_it = client_local_disk_segment_.find(client_it->second); + if (disk_it == client_local_disk_segment_.end()) { + return 0; + } + return disk_it->second->ssd_total_capacity_bytes; +} + +int64_t ScopedLocalDiskSegmentAccess::getSsdUsedBytes( + const std::string& segment_name) const { + auto client_it = client_by_name_.find(segment_name); + if (client_it == client_by_name_.end()) { + return 0; + } + auto disk_it = client_local_disk_segment_.find(client_it->second); + if (disk_it == client_local_disk_segment_.end()) { + return 0; + } + return disk_it->second->ssd_used_bytes.load(std::memory_order_relaxed); +} + +bool SegmentManager::HasSegmentByEndpoint(const std::string& endpoint) const { + std::shared_lock lock(segment_mutex_); + for (const auto& [segment_id, mounted_segment] : mounted_segments_) { + if (mounted_segment.segment.te_endpoint == endpoint) { + return true; + } + } + return false; +} + +bool SegmentManager::GetSegmentBasicInfo(const UUID& segment_id, + std::string& segment_name, + std::string& te_endpoint) const { + std::shared_lock lock(segment_mutex_); + auto it = mounted_segments_.find(segment_id); + if (it == mounted_segments_.end()) { + return false; + } + const Segment& seg = it->second.segment; + segment_name = seg.name; + te_endpoint = seg.te_endpoint; + return true; +} } // namespace mooncake diff --git a/mooncake-store/src/serialize/serializer.cpp b/mooncake-store/src/serialize/serializer.cpp index 9c9b148676..ffc9ff8cc1 100644 --- a/mooncake-store/src/serialize/serializer.cpp +++ b/mooncake-store/src/serialize/serializer.cpp @@ -1,8 +1,8 @@ #include #include -#include "serialize/serializer.hpp" -#include "offset_allocator/offset_allocator.hpp" +#include "serialize/serializer.h" +#include "offset_allocator/offset_allocator.h" #include "types.h" #include "master_service.h" #include "utils/zstd_util.h" @@ -830,9 +830,10 @@ tl::expected Serializer::serialize( const MountedSegment &mounted_segment, MsgpackPacker &packer) { // Use array structure for packing, more efficient // Format: [segment_id, segment_name, segment_base, segment_size, - // te_endpoint, status, has_buffer_allocator, buffer_allocator_data...] + // te_endpoint, status, has_buffer_allocator, buffer_allocator_data, + // host_id] - packer.pack_array(8); + packer.pack_array(9); // Serialize Segment info packer.pack(UuidToString(mounted_segment.segment.id)); @@ -855,12 +856,14 @@ tl::expected Serializer::serialize( if (!result) { return tl::unexpected(result.error()); } + packer.pack(mounted_segment.segment.host_id); return {}; } } packer.pack(false); // Mark no valid buffer allocator exists packer.pack_nil(); + packer.pack(mounted_segment.segment.host_id); return {}; } @@ -917,6 +920,9 @@ Serializer::deserialize(const msgpack::object &obj) { return tl::unexpected(allocatorResult.error()); } } + if (obj.via.array.size >= 9) { + mounted_segment.segment.host_id = array[8].as(); + } } catch (const std::exception &e) { return tl::unexpected(SerializationError( ErrorCode::DESERIALIZE_FAIL, @@ -994,6 +1000,17 @@ auto Serializer::deserialize(const msgpack::object &obj) allocator->offset_allocator_ = offset_allocator_result.value(); allocator->cur_size_ = cur_size; + // The snapshot restores cur_size_ directly from persisted data + // without going through the live allocate()/adoptImportedBuffer() + // paths, so no inc_allocated_mem_size() was paired with it. The + // allocator destructor still calls dec_allocated_mem_size(cur_size_) + // to undo its contribution to the global metric; without a matching + // inc the gauge would go negative and wrap to ~16M TB when formatted + // as uint64. Pair it here so the accounting stays symmetric and the + // gauge ends at 0 after this (often throwaway) allocator is destroyed. + MasterMetricManager::instance().inc_allocated_mem_size( + segment_name, static_cast(cur_size)); + return allocator; } catch (const std::exception &e) { return tl::unexpected(SerializationError( diff --git a/mooncake-store/src/shm_helper.cpp b/mooncake-store/src/shm_helper.cpp index d1dbd424c2..fd2bbb8f64 100644 --- a/mooncake-store/src/shm_helper.cpp +++ b/mooncake-store/src/shm_helper.cpp @@ -3,7 +3,6 @@ #include #include #include -#include #include #include #include @@ -182,59 +181,4 @@ std::shared_ptr ShmHelper::get_shm(void* addr) { return nullptr; } -/* ================================================================== */ -/* ipc_send_fd / ipc_recv_fd: SCM_RIGHTS fd passing over Unix socket */ -/* ================================================================== */ - -int ipc_send_fd(int socket, int fd, void* data, size_t data_len) { - struct msghdr msg; - memset(&msg, 0, sizeof(msg)); - struct iovec iov; - char buf[CMSG_SPACE(sizeof(int))]; - memset(buf, 0, sizeof(buf)); - - iov.iov_base = data; - iov.iov_len = data_len; - - msg.msg_iov = &iov; - msg.msg_iovlen = 1; - msg.msg_control = buf; - msg.msg_controllen = sizeof(buf); - - struct cmsghdr* cmsg = CMSG_FIRSTHDR(&msg); - cmsg->cmsg_level = SOL_SOCKET; - cmsg->cmsg_type = SCM_RIGHTS; - cmsg->cmsg_len = CMSG_LEN(sizeof(int)); - memcpy(CMSG_DATA(cmsg), &fd, sizeof(int)); - - return sendmsg(socket, &msg, 0); -} - -int ipc_recv_fd(int socket, void* data, size_t data_len) { - struct msghdr msg; - memset(&msg, 0, sizeof(msg)); - struct iovec iov; - char buf[CMSG_SPACE(sizeof(int))]; - memset(buf, 0, sizeof(buf)); - - iov.iov_base = data; - iov.iov_len = data_len; - - msg.msg_iov = &iov; - msg.msg_iovlen = 1; - msg.msg_control = buf; - msg.msg_controllen = sizeof(buf); - - if (recvmsg(socket, &msg, 0) < 0) return -1; - - struct cmsghdr* cmsg = CMSG_FIRSTHDR(&msg); - if (cmsg && cmsg->cmsg_level == SOL_SOCKET && - cmsg->cmsg_type == SCM_RIGHTS) { - int fd; - memcpy(&fd, CMSG_DATA(cmsg), sizeof(int)); - return fd; - } - return -1; -} - } // namespace mooncake diff --git a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp index 7c8e6ffbb1..56483d05f6 100644 --- a/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp +++ b/mooncake-store/src/storage/distributed/distributed_storage_backend.cpp @@ -7,6 +7,7 @@ #include #include +#include "environ.h" #include "utils.h" namespace mooncake { @@ -43,16 +44,16 @@ bool DistributedStorageConfig::Validate() const { DistributedStorageConfig DistributedStorageConfig::FromEnvironment() { DistributedStorageConfig config; config.fsdir = - GetEnvStringOr("MOONCAKE_DISTRIBUTED_ROOT_DIR", config.fsdir); + Environ::GetString("MOONCAKE_DISTRIBUTED_ROOT_DIR", config.fsdir); if (!std::filesystem::path(config.fsdir).is_absolute()) { config.fsdir = std::filesystem::absolute(config.fsdir).string(); } - config.fs_adapter_type = - GetEnvStringOr("MOONCAKE_DISTRIBUTED_FS_TYPE", config.fs_adapter_type); + config.fs_adapter_type = Environ::GetString("MOONCAKE_DISTRIBUTED_FS_TYPE", + config.fs_adapter_type); config.enable_health_check = - GetEnvOr("MOONCAKE_DISTRIBUTED_HEALTH_CHECK", false); + Environ::GetBool("MOONCAKE_DISTRIBUTED_HEALTH_CHECK", false); config.hash_bucket_count = - GetEnvOr("MOONCAKE_DISTRIBUTED_HASH_BUCKET_COUNT", 256); + Environ::GetInt("MOONCAKE_DISTRIBUTED_HASH_BUCKET_COUNT", 256); return config; } @@ -139,8 +140,7 @@ tl::expected DistributedStorageBackend::BatchOffload( std::function& keys, std::vector& metadatas)> complete_handler, - std::function& evicted_keys)> - eviction_handler) { + EvictionHandler eviction_handler) { if (!initialized_) { LOG(ERROR) << "DistributedStorageBackend is not initialized"; return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); diff --git a/mooncake-store/src/storage_backend.cpp b/mooncake-store/src/storage_backend.cpp index 0820170190..d206571402 100644 --- a/mooncake-store/src/storage_backend.cpp +++ b/mooncake-store/src/storage_backend.cpp @@ -1,3 +1,4 @@ +#include "serializer.h" #include "storage_backend.h" #include @@ -5,8 +6,13 @@ #include #include #include +#include #include +#include +#include +#include +#include #include #include #include @@ -19,8 +25,31 @@ #include "mutex.h" #include "utils.h" +#include "crc32c.h" +#include "ascii_string.h" +#include "bool_parser.h" +#include "environ.h" #include + +namespace { +struct FdGuard { + int fd = -1; + explicit FdGuard(int f) : fd(f) {} + ~FdGuard() { + if (fd >= 0) close(fd); + } + FdGuard(const FdGuard&) = delete; + FdGuard& operator=(const FdGuard&) = delete; + int get() const { return fd; } + int release() { + int r = fd; + fd = -1; + return r; + } +}; +} // namespace + #include "storage/distributed/distributed_storage_backend.h" namespace mooncake { @@ -48,11 +77,11 @@ bool BucketBackendConfig::Validate() const { FilePerKeyConfig FilePerKeyConfig::FromEnvironment() { FilePerKeyConfig config; - config.fsdir = GetEnvStringOr("MOONCAKE_OFFLOAD_FSDIR", config.fsdir); + config.fsdir = Environ::GetString("MOONCAKE_OFFLOAD_FSDIR", config.fsdir); - config.enable_eviction = GetEnvOr( + config.enable_eviction = Environ::GetBool( "MOONCAKE_OFFLOAD_ENABLE_EVICTION", - GetEnvOr("ENABLE_EVICTION", config.enable_eviction)); + Environ::GetBool("ENABLE_EVICTION", config.enable_eviction)); return config; } @@ -60,20 +89,20 @@ FilePerKeyConfig FilePerKeyConfig::FromEnvironment() { BucketBackendConfig BucketBackendConfig::FromEnvironment() { BucketBackendConfig config; - config.bucket_keys_limit = GetEnvOr( + config.bucket_keys_limit = Environ::GetInt64( "MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", config.bucket_keys_limit); - config.bucket_size_limit = GetEnvOr( + config.bucket_size_limit = Environ::GetInt64( "MOONCAKE_OFFLOAD_BUCKET_SIZE_LIMIT_BYTES", config.bucket_size_limit); config.max_total_size = - GetEnvOr("MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE", - GetEnvOr("MOONCAKE_BUCKET_MAX_TOTAL_SIZE", + Environ::GetInt64("MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE", + Environ::GetInt64("MOONCAKE_BUCKET_MAX_TOTAL_SIZE", config.max_total_size)); - const auto policy_str = GetEnvStringOr( + const auto policy_str = Environ::GetString( "MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY", - GetEnvStringOr("MOONCAKE_BUCKET_EVICTION_POLICY", "fifo")); + Environ::GetString("MOONCAKE_BUCKET_EVICTION_POLICY", "fifo")); if (policy_str == "fifo") { config.eviction_policy = BucketEvictionPolicy::FIFO; } else if (policy_str == "lru") { @@ -85,6 +114,129 @@ BucketBackendConfig BucketBackendConfig::FromEnvironment() { return config; } +bool OffsetAllocatorBackendConfig::Validate() const { + if (persist_mode == OffsetPersistMode::kRelaxed) { + if (persist_interval_seconds < 5) { + LOG(ERROR) << "OffsetAllocatorBackendConfig: " + "persist_interval_seconds must be >= 5 for " + "kRelaxed mode"; + return false; + } + } + if (high_ratio <= 0.0 || high_ratio > 1.0) { + LOG(ERROR) + << "OffsetAllocatorBackendConfig: high_ratio must be in (0,1]"; + return false; + } + if (low_ratio <= 0.0 || low_ratio >= high_ratio) { + LOG(ERROR) << "OffsetAllocatorBackendConfig: low_ratio must be in (0, " + "high_ratio)"; + return false; + } + if (keys_high_ratio <= 0.0 || keys_high_ratio > 1.0) { + LOG(ERROR) + << "OffsetAllocatorBackendConfig: keys_high_ratio must be in (0,1]"; + return false; + } + if (keys_low_ratio <= 0.0 || keys_low_ratio >= keys_high_ratio) { + LOG(ERROR) << "OffsetAllocatorBackendConfig: keys_low_ratio must be in " + "(0, keys_high_ratio)"; + return false; + } + if (max_evict_per_offload == 0) { + LOG(ERROR) << "OffsetAllocatorBackendConfig: max_evict_per_offload " + "must be > 0"; + return false; + } + if (fallback_evict_batch == 0) { + LOG(ERROR) + << "OffsetAllocatorBackendConfig: fallback_evict_batch must be > 0"; + return false; + } + if (max_capacity_nodes < 0) { + LOG(ERROR) + << "OffsetAllocatorBackendConfig: max_capacity_nodes must be >= 0"; + return false; + } + return true; +} + +static std::optional GetEnvDouble(const char* name) { + const char* env = std::getenv(name); + if (!env || env[0] == '\0') return std::nullopt; + try { + return std::stod(env); + } catch (...) { + return std::nullopt; + } +} + +OffsetAllocatorBackendConfig OffsetAllocatorBackendConfig::FromEnvironment() { + OffsetAllocatorBackendConfig cfg; + + const char* pol = std::getenv("MOONCAKE_OFFSET_EVICTION_POLICY"); + if (pol) { + if (AsciiCaseInsensitiveEquals(pol, "fifo")) { + cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + } + // NONE is default; LRU reserved for phase 2 + } + + if (auto v = GetEnvDouble("MOONCAKE_OFFSET_HIGH_RATIO")) + cfg.high_ratio = *v; + if (auto v = GetEnvDouble("MOONCAKE_OFFSET_LOW_RATIO")) cfg.low_ratio = *v; + // Both byte and key watermarks derive from the same ratio pair. + cfg.keys_high_ratio = cfg.high_ratio; + cfg.keys_low_ratio = cfg.low_ratio; + + cfg.max_capacity_nodes = Environ::GetInt64( + "MOONCAKE_OFFSET_MAX_CAPACITY_NODES", cfg.max_capacity_nodes); + + // Read eviction cap as int64_t to guard against negative env values + // which would wrap to SIZE_MAX with an unsigned parser. + auto max_evict_raw = + Environ::GetInt64("MOONCAKE_OFFSET_MAX_EVICT_PER_OFFLOAD", + static_cast(cfg.max_evict_per_offload)); + if (max_evict_raw > 0) { + cfg.max_evict_per_offload = static_cast(max_evict_raw); + } else if (max_evict_raw <= 0) { + LOG(WARNING) << "MOONCAKE_OFFSET_MAX_EVICT_PER_OFFLOAD=" + << max_evict_raw << " is non-positive; using default " + << cfg.max_evict_per_offload; + } + + // Persistence mode + const char* persist = std::getenv("MOONCAKE_OFFSET_PERSIST_MODE"); + if (persist) { + std::string s(persist); + if (AsciiCaseInsensitiveEquals(s, "disabled")) { + cfg.persist_mode = OffsetPersistMode::kDisabled; + } else if (AsciiCaseInsensitiveEquals(s, "relaxed")) { + cfg.persist_mode = OffsetPersistMode::kRelaxed; + } else if (AsciiCaseInsensitiveEquals(s, "strict")) { + cfg.persist_mode = OffsetPersistMode::kStrict; + } else { + LOG(WARNING) << "Unknown MOONCAKE_OFFSET_PERSIST_MODE=" << s + << "; using default (disabled)"; + } + } + + cfg.persist_interval_seconds = + Environ::GetInt64("MOONCAKE_OFFSET_PERSIST_INTERVAL_SECONDS", + cfg.persist_interval_seconds); + + // Record CRC-32C: "0"/"false"/"off" disables per-record checksums. + const char* crc_env = std::getenv("MOONCAKE_OFFSET_RECORD_CRC"); + if (crc_env) { + const auto parsed = TryParseBool(crc_env); + if (parsed.has_value() && !*parsed) { + cfg.enable_record_crc = false; + } + } + + return cfg; +} + StorageBackendInterface::StorageBackendInterface( const FileStorageConfig& config) : file_storage_config_(config) {} @@ -107,6 +259,16 @@ void StorageBackend::RecalculateAvailableSpace() { bool StorageBackend::IsEvictionEnabled() const { return enable_eviction_; } +Mutex& StorageBackend::GetFilePathMutex(const std::string& path) { + return file_path_mutexes_[std::hash{}(path) % + kFilePathLockCount]; +} + +bool StorageBackend::IsFilePendingEviction(const std::string& path) const { + std::shared_lock lock(file_queue_mutex_); + return pending_eviction_paths_.find(path) != pending_eviction_paths_.end(); +} + tl::expected StorageBackend::Init(uint64_t quota_bytes = 0) { // Skip eviction initialization if disabled if (!IsEvictionEnabled()) { @@ -277,7 +439,8 @@ bool StorageBackend::InitQuotaEvict() { tl::expected, ErrorCode> StorageBackend::StoreObject( const std::string& path, const std::vector& slices, - const std::string& key) { + const std::string& key, + StorageBackendInterface::EvictionHandler eviction_handler) { size_t total_size = 0; for (const auto& slice : slices) { total_size += slice.size; @@ -294,7 +457,7 @@ tl::expected, ErrorCode> StorageBackend::StoreObject( return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } - auto space_result = EnsureDiskSpace(total_size); + auto space_result = EnsureDiskSpace(total_size, eviction_handler); if (!space_result) { return tl::make_unexpected(space_result.error()); } @@ -302,6 +465,12 @@ tl::expected, ErrorCode> StorageBackend::StoreObject( reserved_size = total_size; } + MutexLocker path_locker(&GetFilePathMutex(path)); + if (IsEvictionEnabled() && IsFilePendingEviction(path)) { + ReleaseSpace(reserved_size); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + // Create file and write data (common logic for both modes) auto file_result = CreateFileForWriting(path, reserved_size); if (!file_result) { @@ -323,14 +492,15 @@ tl::expected, ErrorCode> StorageBackend::StoreObject( } tl::expected, ErrorCode> StorageBackend::StoreObject( - const std::string& path, const std::string& str, const std::string& key) { - return StoreObject(path, std::span(str.data(), str.size()), - key); + const std::string& path, const std::string& str, const std::string& key, + StorageBackendInterface::EvictionHandler eviction_handler) { + return StoreObject(path, std::span(str.data(), str.size()), key, + eviction_handler); } tl::expected, ErrorCode> StorageBackend::StoreObject( - const std::string& path, std::span data, - const std::string& key) { + const std::string& path, std::span data, const std::string& key, + StorageBackendInterface::EvictionHandler eviction_handler) { size_t file_total_size = data.size(); // For eviction-enabled mode, check space and reserve @@ -344,7 +514,7 @@ tl::expected, ErrorCode> StorageBackend::StoreObject( return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } - auto space_result = EnsureDiskSpace(file_total_size); + auto space_result = EnsureDiskSpace(file_total_size, eviction_handler); if (!space_result) { return tl::make_unexpected(space_result.error()); } @@ -352,6 +522,12 @@ tl::expected, ErrorCode> StorageBackend::StoreObject( reserved_size = file_total_size; } + MutexLocker path_locker(&GetFilePathMutex(path)); + if (IsEvictionEnabled() && IsFilePendingEviction(path)) { + ReleaseSpace(reserved_size); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + // Create file and write data (common logic for both modes) auto file_result = CreateFileForWriting(path, reserved_size); if (!file_result) { @@ -485,6 +661,8 @@ void StorageBackend::RemoveFile(const std::string& path) { std::this_thread::sleep_for( std::chrono::microseconds(50)); // sleep for 50 us + MutexLocker path_locker(&GetFilePathMutex(path)); + // Eviction disabled, use simple delete (no queue tracking) if (!IsEvictionEnabled()) { if (fs::exists(path)) { @@ -564,51 +742,27 @@ void StorageBackend::RemoveByRegex(const std::string& regex_pattern) { } for (const auto& path : paths_to_remove) { - std::error_code ec; - if (fs::remove(path, ec)) { - VLOG(1) << "Removed file by regex: " << path; - } else { - LOG(ERROR) << "Failed to delete file: " << path - << ", error: " << ec.message(); - } + RemoveFile(path.string()); } return; } // Eviction-enabled logic (local mode) - std::list records_to_remove; - uint64_t total_freed_space = 0; + std::vector paths_to_remove; { std::shared_lock lock(file_queue_mutex_); for (const auto& record : file_write_queue_) { std::string filename = fs::path(record.path).filename().string(); if (std::regex_search(filename, pattern)) { - records_to_remove.push_back(record); + paths_to_remove.push_back(record.path); } } } - for (const auto& record : records_to_remove) { - std::error_code ec; - if (fs::remove(record.path, ec)) { - RemoveFileFromWriteQueue(record.path); - total_freed_space += record.size; - VLOG(1) << "Removed file by regex: " << record.path; - } else { - if (ec && ec == std::errc::no_such_file_or_directory) { - RemoveFileFromWriteQueue(record.path); - total_freed_space += record.size; - } else { - LOG(ERROR) << "Failed to delete file: " << record.path - << ", error: " << ec.message(); - } - } + for (const auto& path : paths_to_remove) { + RemoveFile(path); } - - ReleaseSpace(total_freed_space); - - return; } void StorageBackend::RemoveAll() { @@ -616,41 +770,33 @@ void StorageBackend::RemoveAll() { // Eviction disabled, use simple delete (no queue tracking) if (!IsEvictionEnabled()) { + std::vector paths_to_remove; // Iterate through the root directory and remove all files for (const auto& entry : fs::directory_iterator(root_dir_)) { if (fs::is_regular_file(entry.status())) { - std::error_code ec; - fs::remove(entry.path(), ec); - if (ec) { - LOG(ERROR) << "Failed to delete file: " << entry.path() - << ", error: " << ec.message(); - } + paths_to_remove.push_back(entry.path().string()); } } + for (const auto& path : paths_to_remove) { + RemoveFile(path); + } return; } // Eviction-enabled logic (local mode) try { - std::list records_to_remove; + std::vector paths_to_remove; { - std::unique_lock lock(file_queue_mutex_); - records_to_remove = std::move(file_write_queue_); - file_queue_map_.clear(); - } - uint64_t total_freed_space = 0; - for (const auto& record : records_to_remove) { - total_freed_space += record.size; - std::error_code ec; - fs::remove(record.path, ec); - if (ec && ec != std::errc::no_such_file_or_directory) { - LOG(ERROR) << "RemoveAll: Failed to delete file " << record.path - << ", error: " << ec.message(); + std::shared_lock lock(file_queue_mutex_); + paths_to_remove.reserve(file_write_queue_.size()); + for (const auto& record : file_write_queue_) { + paths_to_remove.push_back(record.path); } } - ReleaseSpace(total_freed_space); - + for (const auto& path : paths_to_remove) { + RemoveFile(path); + } } catch (const fs::filesystem_error& e) { LOG(ERROR) << "Filesystem error when removing all files: " << e.what(); } @@ -840,46 +986,39 @@ FileRecord StorageBackend::EvictFile() { } // Use FIFO based strategy (earliest written first out) - FileRecord record_to_evict = SelectFileToEvictByFIFO(); + FileRecord record_to_evict = PopFileToEvictByFIFO(); if (record_to_evict.path.empty()) { LOG(WARNING) << "No file selected for eviction"; return {}; } - namespace fs = std::filesystem; - std::error_code ec; - uint64_t file_size = record_to_evict.size; - - if (fs::remove(record_to_evict.path, ec)) { - RemoveFileFromWriteQueue(record_to_evict.path); - ReleaseSpace(file_size); + auto delete_result = DeleteEvictedFile(record_to_evict); + if (delete_result) { return record_to_evict; - } else { - if (!ec || ec == std::errc::no_such_file_or_directory) { - RemoveFileFromWriteQueue(record_to_evict.path); - return record_to_evict; - } else { - LOG(ERROR) << "Failed to evict file: " << record_to_evict.path - << ", error: " << ec.message(); - RemoveFileFromWriteQueue(record_to_evict.path); - return {}; - } } + RestoreFileToWriteQueueFront(record_to_evict); + return {}; } void StorageBackend::AddFileToWriteQueue(const std::string& path, uint64_t size, const std::string& key) { - std::unique_lock lock(file_queue_mutex_); + uint64_t replaced_size = 0; + { + std::unique_lock lock(file_queue_mutex_); - auto it = file_queue_map_.find(path); - if (it != file_queue_map_.end()) { - file_write_queue_.erase(it->second); - file_queue_map_.erase(it); + auto it = file_queue_map_.find(path); + if (it != file_queue_map_.end()) { + replaced_size = it->second->size; + file_write_queue_.erase(it->second); + file_queue_map_.erase(it); + } + + file_write_queue_.push_back({path, size, key}); + file_queue_map_[path] = std::prev(file_write_queue_.end()); } - file_write_queue_.push_back({path, size, key}); - file_queue_map_[path] = std::prev(file_write_queue_.end()); + ReleaseSpace(replaced_size); } void StorageBackend::RemoveFileFromWriteQueue(const std::string& path) { @@ -894,18 +1033,107 @@ void StorageBackend::RemoveFileFromWriteQueue(const std::string& path) { } } -FileRecord StorageBackend::SelectFileToEvictByFIFO() { - std::unique_lock lock(file_queue_mutex_); - if (file_write_queue_.empty()) { - LOG(WARNING) << "Queue is empty, cannot select file to evict"; +FileRecord StorageBackend::PopFileToEvictByFIFO() { + while (true) { + std::string candidate_path; + { + std::shared_lock lock(file_queue_mutex_); + if (file_write_queue_.empty()) { + LOG(WARNING) << "Queue is empty, cannot select file to evict"; + return {}; + } + candidate_path = file_write_queue_.front().path; + } + + MutexLocker path_locker(&GetFilePathMutex(candidate_path)); + std::unique_lock lock(file_queue_mutex_); + if (file_write_queue_.empty() || + file_write_queue_.front().path != candidate_path) { + continue; + } + + auto map_it = file_queue_map_.find(candidate_path); + if (map_it == file_queue_map_.end() || + map_it->second != file_write_queue_.begin()) { + continue; + } + + FileRecord record = file_write_queue_.front(); + file_queue_map_.erase(map_it); + file_write_queue_.pop_front(); + pending_eviction_paths_.insert(record.path); + return record; + } +} + +void StorageBackend::RestoreFileToWriteQueueFront(const FileRecord& record) { + if (record.path.empty()) { + return; + } + + bool was_replaced = false; + { + MutexLocker path_locker(&GetFilePathMutex(record.path)); + std::unique_lock lock(file_queue_mutex_); + pending_eviction_paths_.erase(record.path); + if (file_queue_map_.find(record.path) != file_queue_map_.end()) { + was_replaced = true; + } else { + file_write_queue_.push_front(record); + file_queue_map_[record.path] = file_write_queue_.begin(); + } + } + + if (was_replaced) { + LOG(ERROR) << "Cannot restore evicted record because the path was " + "replaced: " + << record.path; + ReleaseSpace(record.size); + } +} + +tl::expected StorageBackend::DeleteEvictedFile( + const FileRecord& record) { + namespace fs = std::filesystem; + MutexLocker path_locker(&GetFilePathMutex(record.path)); + + bool was_replaced = false; + { + std::unique_lock queue_lock(file_queue_mutex_); + if (file_queue_map_.find(record.path) != file_queue_map_.end()) { + // A newer StoreObject completed after this record was selected. + // The old file contents were replaced in place, so release the old + // accounting without deleting the new file. + was_replaced = true; + pending_eviction_paths_.erase(record.path); + } + } + if (was_replaced) { + ReleaseSpace(record.size); + return {}; + } + + std::error_code ec; + + if (fs::remove(record.path, ec) || !ec || + ec == std::errc::no_such_file_or_directory) { + { + std::unique_lock queue_lock(file_queue_mutex_); + pending_eviction_paths_.erase(record.path); + } + ReleaseSpace(record.size); return {}; } - return file_write_queue_.front(); + LOG(ERROR) << "Failed to evict file: " << record.path + << ", error: " << ec.message(); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } tl::expected, ErrorCode> -StorageBackend::EnsureDiskSpace(size_t required_size) { +StorageBackend::EnsureDiskSpace( + size_t required_size, + StorageBackendInterface::EvictionHandler eviction_handler) { std::vector evicted_keys; // If eviction is disabled, skip space checking and eviction if (!IsEvictionEnabled()) { @@ -916,21 +1144,71 @@ StorageBackend::EnsureDiskSpace(size_t required_size) { size_t attempts = 0; bool space_reserved = CheckDiskSpace(required_size); + if (space_reserved) { + return evicted_keys; + } - while (!space_reserved && attempts < kMaxEvictionAttempts) { - FileRecord evicted = EvictFile(); - if (evicted.path.empty()) { + uint64_t projected_available_space = 0; + { + std::shared_lock lock(space_mutex_); + projected_available_space = available_space_; + } + + std::vector pending_evictions; + while (projected_available_space < required_size && + attempts < kMaxEvictionAttempts) { + FileRecord pending = PopFileToEvictByFIFO(); + if (pending.path.empty()) { LOG(ERROR) << "Failed to evict file to make space."; + for (auto it = pending_evictions.rbegin(); + it != pending_evictions.rend(); ++it) { + RestoreFileToWriteQueueFront(*it); + } return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } - if (!evicted.key.empty()) { - evicted_keys.push_back(evicted.key); + if (!pending.key.empty()) { + evicted_keys.push_back(pending.key); } + projected_available_space += pending.size; + pending_evictions.push_back(std::move(pending)); attempts++; + } + + if (projected_available_space < required_size) { + for (auto it = pending_evictions.rbegin(); + it != pending_evictions.rend(); ++it) { + RestoreFileToWriteQueueFront(*it); + } + LOG(ERROR) << "Still insufficient disk space after selecting files for " + "eviction."; + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + + if (eviction_handler && !evicted_keys.empty()) { + auto notify_result = eviction_handler(evicted_keys); + if (!notify_result) { + for (auto it = pending_evictions.rbegin(); + it != pending_evictions.rend(); ++it) { + RestoreFileToWriteQueueFront(*it); + } + return tl::make_unexpected(notify_result.error()); + } + } - space_reserved = CheckDiskSpace(required_size); + bool deletion_failed = false; + for (auto& pending : pending_evictions) { + auto delete_result = DeleteEvictedFile(pending); + if (!delete_result) { + deletion_failed = true; + pending.key.clear(); + RestoreFileToWriteQueueFront(pending); + } + } + if (deletion_failed) { + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } + space_reserved = CheckDiskSpace(required_size); if (!space_reserved) { LOG(ERROR) << "Still insufficient disk space after evicting files."; return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); @@ -939,6 +1217,116 @@ StorageBackend::EnsureDiskSpace(size_t required_size) { return evicted_keys; } +tl::expected, ErrorCode> +StorageBackend::EvictAboveDiskWatermark( + double high_watermark_ratio, double low_watermark_ratio, + StorageBackendInterface::EvictionHandler eviction_handler) { + std::vector evicted_keys; + if (!IsEvictionEnabled()) { + return evicted_keys; + } + if (!initialized_.load(std::memory_order_acquire)) { + LOG(ERROR) + << "EvictAboveDiskWatermark called before StorageBackend::Init " + "was completed."; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + + uint64_t used_space = 0; + uint64_t total_space = 0; + { + std::shared_lock lock(space_mutex_); + used_space = used_space_; + total_space = total_space_; + } + if (total_space == 0) { + return evicted_keys; + } + + const uint64_t high_watermark_bytes = + static_cast(total_space * high_watermark_ratio); + if (used_space <= high_watermark_bytes) { + return evicted_keys; + } + + const uint64_t target_used_bytes = + static_cast(total_space * low_watermark_ratio); + size_t queue_size = 0; + { + std::shared_lock lock(file_queue_mutex_); + queue_size = file_write_queue_.size(); + } + + const size_t max_eviction_attempts = queue_size + 100; + size_t attempts = 0; + uint64_t projected_used_space = used_space; + std::vector pending_evictions; + while (projected_used_space > target_used_bytes && + attempts < max_eviction_attempts) { + FileRecord pending = PopFileToEvictByFIFO(); + if (pending.path.empty()) { + LOG(WARNING) << "Disk watermark eviction could not select a file " + << "for eviction; used=" << projected_used_space + << ", target=" << target_used_bytes; + break; + } + if (!pending.key.empty()) { + evicted_keys.push_back(pending.key); + } + projected_used_space = pending.size >= projected_used_space + ? 0 + : projected_used_space - pending.size; + pending_evictions.push_back(std::move(pending)); + attempts++; + } + + if (eviction_handler && !evicted_keys.empty()) { + auto notify_result = eviction_handler(evicted_keys); + if (!notify_result) { + for (auto it = pending_evictions.rbegin(); + it != pending_evictions.rend(); ++it) { + RestoreFileToWriteQueueFront(*it); + } + return tl::make_unexpected(notify_result.error()); + } + } + + bool deletion_failed = false; + for (auto& pending : pending_evictions) { + auto delete_result = DeleteEvictedFile(pending); + if (!delete_result) { + deletion_failed = true; + pending.key.clear(); + RestoreFileToWriteQueueFront(pending); + } + } + + { + std::shared_lock lock(space_mutex_); + used_space = used_space_; + } + + if (used_space > target_used_bytes) { + LOG(WARNING) << "Disk watermark eviction stopped above target: used=" + << used_space << ", target=" << target_used_bytes + << ", attempts=" << attempts; + } + if (deletion_failed) { + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + return evicted_keys; +} + +void StorageBackend::UpdateFileRecordKey(const std::string& path, + const std::string& key) { + std::unique_lock lock(file_queue_mutex_); + auto it = file_queue_map_.find(path); + if (it == file_queue_map_.end()) { + return; + } + it->second->key = key; +} + void StorageBackend::ReleaseSpace(uint64_t size_to_release) { if (size_to_release == 0) { return; @@ -969,6 +1357,22 @@ StorageBackendAdaptor::StorageBackendAdaptor( total_size(0) {} tl::expected StorageBackendAdaptor::Init() { + namespace fs = std::filesystem; + // Validate obviously invalid configuration up front. A missing root + // directory is fine: StorageBackend::Init() creates it on demand. + std::error_code ec; + const auto root_status = + fs::status(file_storage_config_.storage_filepath, ec); + if (!ec && fs::exists(root_status) && !fs::is_directory(root_status)) { + LOG(ERROR) << "Storage path exists but is not a directory: " + << file_storage_config_.storage_filepath; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + if (file_per_key_config_.fsdir.empty()) { + LOG(ERROR) << "FSDIR cannot be empty"; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + std::string storage_root = file_storage_config_.storage_filepath + file_per_key_config_.fsdir; @@ -1007,8 +1411,7 @@ tl::expected StorageBackendAdaptor::BatchOffload( std::function& keys, std::vector& metadatas)> complete_handler, - std::function& evicted_keys)> - eviction_handler) { + EvictionHandler eviction_handler) { if (batch_object.empty()) { LOG(ERROR) << "batch object is empty"; return tl::make_unexpected(ErrorCode::INVALID_KEY); @@ -1045,7 +1448,8 @@ tl::expected StorageBackendAdaptor::BatchOffload( std::string kv_buf; struct_pb::to_pb(kv, kv_buf); - auto store_result = storage_backend_->StoreObject(path, kv_buf, kv.key); + auto store_result = storage_backend_->StoreObject(path, kv_buf, kv.key, + eviction_handler); if (!store_result) { LOG(ERROR) << "Failed to store object for key: " << kv.key << ", error: " << store_result.error() @@ -1053,11 +1457,6 @@ tl::expected StorageBackendAdaptor::BatchOffload( continue; // Continue processing other keys } - // Notify eviction handler about any evicted keys (batch) - if (eviction_handler && !store_result.value().empty()) { - eviction_handler(store_result.value()); - } - { MutexLocker lock(&mutex_); total_keys++; @@ -1087,6 +1486,18 @@ tl::expected StorageBackendAdaptor::BatchOffload( return static_cast(keys.size()); } +tl::expected, ErrorCode> +StorageBackendAdaptor::EvictAboveDiskWatermark( + double high_watermark_ratio, double low_watermark_ratio, + EvictionHandler eviction_handler) { + auto eviction_result = storage_backend_->EvictAboveDiskWatermark( + high_watermark_ratio, low_watermark_ratio, eviction_handler); + if (!eviction_result) { + return tl::make_unexpected(eviction_result.error()); + } + return eviction_result; +} + tl::expected StorageBackendAdaptor::IsExist( const std::string& key) { auto path = ResolvePathFromKey(key, file_storage_config_.storage_filepath, @@ -1204,6 +1615,7 @@ tl::expected StorageBackendAdaptor::ScanMeta( KVEntry kv; struct_pb::from_pb(kv, buf); + storage_backend_->UpdateFileRecordKey(p.string(), kv.key); total_keys++; total_size += buf.size(); @@ -1229,6 +1641,20 @@ tl::expected StorageBackendAdaptor::ScanMeta( return {}; } +void StorageBackendAdaptor::RemoveAll() { + if (storage_backend_) { + storage_backend_->RemoveAll(); + } + // Reset adaptor-level counters so IsEnableOffloading() reflects the empty + // disk. Without this, total_keys/total_size keep their pre-cleanup values + // and the node stays unable to offload until a restart (which rescans + // metadata and rebuilds them). Applies to the kFilePerKey path; the + // OffsetAllocator backend resets its own atomics internally. + MutexLocker lock(&mutex_); + total_keys = 0; + total_size = 0; +} + BucketIdGenerator::BucketIdGenerator(int64_t start) { if (start <= 0) { auto cur_time_stamp = time_gen(); @@ -1280,8 +1706,7 @@ tl::expected BucketStorageBackend::BatchOffload( std::function& keys, std::vector& metadatas)> complete_handler, - std::function& evicted_keys)> - eviction_handler) { + EvictionHandler eviction_handler) { if (!initialized_.load(std::memory_order_acquire)) { LOG(ERROR) << "Storage backend is not initialized. Call Init() before use."; @@ -1313,64 +1738,104 @@ tl::expected BucketStorageBackend::BatchOffload( // Phase 1: eviction — remove oldest buckets from metadata maps to make // room. Must notify master BEFORE deleting files (Phase 2). const int64_t required_size = bucket->data_size + bucket->meta_size; - PendingEviction pending = PrepareEviction(required_size); + auto prepare_result = PrepareEviction(required_size, bucket->keys); + if (!prepare_result) { + return tl::make_unexpected(prepare_result.error()); + } + PendingEviction pending = std::move(prepare_result.value()); // Notify master about evicted keys BEFORE touching the files. if (eviction_handler && !pending.keys.empty()) { - eviction_handler(pending.keys); + auto notify_result = eviction_handler(pending.keys); + if (!notify_result) { + RestorePreparedEviction(std::move(pending)); + return tl::make_unexpected(notify_result.error()); + } } + CommitPreparedEviction(pending); // Phase 2: delete evicted files NOW, before writing the new bucket, so // that the freed disk space is available for the incoming write. - FinalizeEviction(pending); + auto finalize_result = FinalizeEviction(pending); + if (!finalize_result) { + LOG(ERROR) + << "FinalizeEviction failed after master committed eviction; " + "continuing BatchOffload: " + << finalize_result.error(); + } auto write_bucket_result = WriteBucket(bucket_id, bucket, iovs); if (!write_bucket_result) { LOG(ERROR) << "Failed to write bucket with id: " << bucket_id; + ReleasePreparedWrite(pending); return tl::make_unexpected(write_bucket_result.error()); } - if (complete_handler != nullptr) { - auto error_code = complete_handler(bucket->keys, metadatas); - if (error_code != ErrorCode::OK) { - LOG(ERROR) << "Complete handler failed: " << error_code - << ", Key count: " << bucket->keys.size() - << ", Bucket id: " << bucket_id; - return tl::make_unexpected(error_code); - } - } + // Save a copy of bucket->keys before std::move(bucket) into buckets_ + // consumes the shared_ptr. Needed for complete_handler and rollback. + const auto bucket_keys = bucket->keys; - // Commit to metadata maps under exclusive lock. - // Check for duplicate keys and rollback if any found. + // Commit to metadata maps under exclusive lock FIRST. + // This ensures any concurrent BatchLoad arriving after Master redirects + // reads to this node can find the key in object_bucket_map_. + // Even if complete_handler fails later, the read path is correct — + // RollbackCommittedBucket will undo the commit safely. { SharedMutexLocker lock(&mutex_); + ReleasePreparedWriteLocked(pending); + // Pre-check for duplicates before modifying any state - for (const auto& key : bucket->keys) { + bool duplicate_found = false; + for (const auto& key : bucket_keys) { if (object_bucket_map_.find(key) != object_bucket_map_.end()) { - LOG(WARNING) - << "Duplicate key detected in BatchOffload: " << key - << ", bucket_id=" << bucket_id - << ". Returning OBJECT_ALREADY_EXISTS."; - lock.unlock(); - CleanupOrphanedBucket(bucket_id); - return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS); + duplicate_found = true; + break; } } - // No duplicates found, safe to commit - total_size_ += bucket->data_size + bucket->meta_size; - object_bucket_map_.reserve(object_bucket_map_.size() + - bucket->keys.size()); - for (size_t i = 0; i < bucket->keys.size(); ++i) { - auto [it, inserted] = object_bucket_map_.insert( - {bucket->keys[i], std::move(metadatas[i])}); - if (!inserted) { - LOG(ERROR) << "Unexpected duplicate key after pre-check: " - << bucket->keys[i] << ", bucket_id=" << bucket_id; + if (!duplicate_found) { + total_size_ += bucket->data_size + bucket->meta_size; + object_bucket_map_.reserve(object_bucket_map_.size() + + bucket_keys.size()); + for (size_t i = 0; i < bucket_keys.size(); ++i) { + auto [it, inserted] = + object_bucket_map_.insert({bucket_keys[i], metadatas[i]}); + CHECK(inserted) + << "Reserved key became duplicated: " << bucket_keys[i]; } + buckets_.emplace(bucket_id, std::move(bucket)); + lru_index_.emplace(0LL, bucket_id); + } + if (duplicate_found) { + LOG(ERROR) << "Reserved key became duplicated before commit, " + "bucket_id=" + << bucket_id; + lock.unlock(); + CleanupOrphanedBucket(bucket_id); + return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS); + } + } + // Lock released. From this point forward, concurrent BatchLoad + // can find the keys and read from the committed bucket files. + + // Notify Master AFTER local index is committed. + // metadatas[i].transport_endpoint is empty here (it gets populated by + // complete_handler before the RPC); int64_t fields carry the metadata + // from BuildBucket unchanged. + if (complete_handler != nullptr) { + auto error_code = complete_handler(bucket_keys, metadatas); + if (error_code != ErrorCode::OK) { + LOG(ERROR) << "Complete handler failed: " << error_code + << ", Key count: " << bucket_keys.size() + << ", Bucket id: " << bucket_id; + // Master was NOT notified. The local index has entries that + // Master doesn't know about — a "client can read but Master + // doesn't know" ghost replica. Rollback the local commit + // (removes index entries + waits for inflight reads + deletes + // on-disk files). + RollbackCommittedBucket(bucket_id, bucket_keys); + return tl::make_unexpected(error_code); } - buckets_.emplace(bucket_id, std::move(bucket)); - lru_index_.emplace(0LL, bucket_id); } return bucket_id; @@ -1511,10 +1976,34 @@ tl::expected BucketStorageBackend::BatchLoad( plan.dest_slice.ptr, aligned_size, aligned_offset); if (read_res) { - // Adjust ptr to point to actual data start (no memcpy) + // Verify the aligned read returned enough bytes + // to cover the actual data region. read_aligned + // reads the full aligned range [aligned_offset, + // aligned_end); the caller-visible data starts at + // offset_in_buffer into that buffer and spans + // plan.dest_slice.size bytes. + size_t min_required = + static_cast(offset_in_buffer) + + plan.dest_slice.size; + if (read_res.value() < min_required) { + LOG(ERROR) + << "read_aligned short read for key: " << plan.key + << ", bucket_id=" << plan.bucket_id + << ", expected at least: " << min_required + << " (aligned_size=" << aligned_size + << ", data_size=" << plan.dest_slice.size << ")" + << ", got: " << read_res.value(); + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + // Adjust ptr to point to actual data start + // (zero-copy: no memcpy, buffer was oversized by + // AllocateBatch to accommodate the aligned read) batch_object.at(plan.key).ptr = static_cast(plan.dest_slice.ptr) + offset_in_buffer; + // Normalize read_res so the common validation + // below passes. The real short-read check was + // already done above (min_required). read_res = plan.dest_slice.size; } } else @@ -1571,12 +2060,22 @@ tl::expected BucketStorageBackend::Init() { lru_index_.clear(); total_size_ = 0; int64_t max_bucket_id = BucketIdGenerator::INIT_NEW_START_ID; + for (const auto& entry : fs::recursive_directory_iterator(storage_path_)) { if (entry.is_regular_file() && entry.path().extension() == BUCKET_METADATA_FILE_SUFFIX) { const auto& bucket_id_str = entry.path().stem(); - int64_t bucket_id = std::stoll(bucket_id_str); + int64_t bucket_id = 0; + try { + bucket_id = std::stoll(bucket_id_str); + } catch (const std::exception& e) { + LOG(WARNING) + << "Skipping metadata file with a non-numeric " + "bucket id: " + << entry.path().string() << " (" << e.what() << ")"; + continue; + } auto [metadata_it, success] = buckets_.try_emplace( bucket_id, std::make_shared()); if (success) lru_index_.emplace(0LL, bucket_id); @@ -1608,10 +2107,37 @@ tl::expected BucketStorageBackend::Init() { buckets_.erase(bucket_id); continue; } - auto& meta = *(metadata_it->second); - if (meta.data_size == 0 || meta.meta_size == 0 || - meta.metadatas.empty() || meta.keys.empty()) { - LOG(ERROR) << "Metadata validation failed for bucket: " + auto bucket_data_path_res = GetBucketDataPath(bucket_id); + if (!bucket_data_path_res) { + LOG(ERROR) << "Failed to get data path for bucket: " + << bucket_id_str; + return tl::make_unexpected(bucket_data_path_res.error()); + } + + std::error_code bucket_data_ec; + const auto bucket_data_status = + fs::status(bucket_data_path_res.value(), bucket_data_ec); + if (bucket_data_ec && + bucket_data_ec != std::errc::no_such_file_or_directory) { + LOG(ERROR) << "Failed to inspect bucket data file: " + << bucket_data_path_res.value() + << ", error: " << bucket_data_ec.message(); + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + if (bucket_data_ec == std::errc::no_such_file_or_directory || + !fs::is_regular_file(bucket_data_status)) { + LOG(ERROR) << "Bucket metadata has no valid data file: " + << entry.path().string() + << ", will delete the bucket's remaining files"; + CleanupOrphanedBucket(bucket_id); + lru_index_.erase({0LL, bucket_id}); + buckets_.erase(bucket_id); + continue; + } + auto& meta = *(metadata_it->second); + if (meta.data_size == 0 || meta.meta_size == 0 || + meta.metadatas.empty() || meta.keys.empty()) { + LOG(ERROR) << "Metadata validation failed for bucket: " << bucket_id_str << ", will delete the bucket's data and " "metadata. Detailed values:"; @@ -1687,7 +2213,16 @@ tl::expected BucketStorageBackend::Init() { // Extract bucket ID from filename (e.g., "12345.bucket" -> // "12345") auto bucket_id_str = entry.path().stem(); - int64_t bucket_id = std::stoll(bucket_id_str); + int64_t bucket_id = 0; + try { + bucket_id = std::stoll(bucket_id_str); + } catch (const std::exception& e) { + LOG(WARNING) + << "Skipping orphan-scan file with a non-numeric " + "bucket id: " + << entry.path().string() << " (" << e.what() << ")"; + continue; + } // Check if this bucket has valid metadata if (valid_bucket_ids.find(bucket_id) != valid_bucket_ids.end()) { @@ -1753,11 +2288,7 @@ tl::expected BucketStorageBackend::Init() { tl::expected BucketStorageBackend::IsExist( const std::string& key) { SharedMutexLocker lock(&mutex_, shared_lock); - auto bucket_id_it = object_bucket_map_.find(key); - if (bucket_id_it != object_bucket_map_.end()) { - return true; - } - return false; + return object_bucket_map_.find(key) != object_bucket_map_.end(); } tl::expected BucketStorageBackend::IsEnableOffloading() { @@ -1906,10 +2437,10 @@ tl::expected BucketStorageBackend::GroupOffloadingKeysByBucket( } if (it->second > bucket_backend_config_.bucket_size_limit) { - LOG(ERROR) << "Object size exceeds bucket size limit: " - << "key=" << it->first - << ", object_size=" << it->second << ", limit=" - << bucket_backend_config_.bucket_size_limit; + VLOG(1) << "Object size exceeds bucket size limit: " + << "key=" << it->first << ", object_size=" << it->second + << ", limit=" + << bucket_backend_config_.bucket_size_limit; ++it; continue; } @@ -1921,6 +2452,8 @@ tl::expected BucketStorageBackend::GroupOffloadingKeysByBucket( << ", error=" << is_exist_result.error(); } if (is_exist_result && is_exist_result.value()) { + VLOG(1) << "Key already exists in storage backend, skipping: " + << "key=" << it->first; ++it; continue; } @@ -2132,6 +2665,14 @@ void BucketStorageBackend::CleanupOrphanedBucket(int64_t bucket_id) { auto data_path_res = GetBucketDataPath(bucket_id); if (data_path_res) { + // Evict the cached file handle before deleting the file, matching + // the pattern in FinalizeEviction. Without this, a subsequent open + // on the same path (e.g. after bucket_id reuse) may return a stale + // handle pointing to the now-deleted file. + { + MutexLocker cache_locker(&file_cache_mutex_); + file_cache_.erase(data_path_res.value()); + } if (fs::remove(data_path_res.value(), ec)) { LOG(INFO) << "Cleaned up orphaned bucket data file: " << data_path_res.value(); @@ -2156,6 +2697,97 @@ void BucketStorageBackend::CleanupOrphanedBucket(int64_t bucket_id) { } } +void BucketStorageBackend::RollbackCommittedBucket( + int64_t bucket_id, const std::vector& keys) { + std::shared_ptr bucket_meta; + + // Phase 1: Remove from metadata maps under exclusive lock. + // This prevents new readers from finding the keys. + { + SharedMutexLocker lock(&mutex_); + + auto bucket_it = buckets_.find(bucket_id); + if (bucket_it == buckets_.end()) { + LOG(WARNING) << "RollbackCommittedBucket: bucket " << bucket_id + << " not found in buckets_ — already removed?"; + // Still clean up disk files in case they are orphaned + CleanupOrphanedBucket(bucket_id); + return; + } + + // Save a reference for inflight-read waiting + bucket_meta = bucket_it->second; + + // Remove all keys from object_bucket_map_ + for (const auto& key : keys) { + auto obj_it = object_bucket_map_.find(key); + if (obj_it != object_bucket_map_.end() && + obj_it->second.bucket_id == bucket_id) { + total_size_ -= + obj_it->second.data_size + obj_it->second.key_size; + object_bucket_map_.erase(obj_it); + } + } + + // Remove bucket metadata + total_size_ -= bucket_meta->meta_size; + lru_index_.erase({0LL, bucket_id}); + buckets_.erase(bucket_it); + } + + // Phase 2: Wait for inflight reads to drain. + // Readers that found the key before we removed it from the map + // hold a BucketReadGuard that keeps inflight_reads_ > 0. + // In practice this should never block: the bucket was committed and + // rolled back within microseconds — no reader had time to acquire a + // guard. But guard against the edge case anyway. + // + // Uses the same spin-then-sleep pattern as DeleteBucket (not the + // spin-then-yield pattern of FinalizeEviction) because rollback is a + // rare error-recovery path where CPU friendliness matters more than + // latency. + { + constexpr int kMaxSpinIterations = 1000; + constexpr auto kSleepDuration = std::chrono::microseconds(100); + constexpr auto kMaxWaitTime = std::chrono::seconds(10); + int spin_count = 0; + auto wait_start = std::chrono::steady_clock::now(); + while (bucket_meta->inflight_reads_.load(std::memory_order_acquire) > + 0) { + if (++spin_count > kMaxSpinIterations) { + std::this_thread::sleep_for(kSleepDuration); + spin_count = 0; + if (std::chrono::steady_clock::now() - wait_start > + kMaxWaitTime) { + LOG(ERROR) + << "RollbackCommittedBucket: timed out waiting " + << "for inflight reads on bucket " << bucket_id + << " (inflight=" + << bucket_meta->inflight_reads_.load( + std::memory_order_relaxed) + << "). Leaving orphaned files on disk; they will " + << "be cleaned up by Init() on next restart."; + // Return WITHOUT deleting files. A reader is still + // holding a guard, so deleting the files could cause + // I/O errors on the read path. This matches + // DeleteBucket's behavior (returns INTERNAL_ERROR + // instead of deleting). The orphan will be recovered + // by Init()'s orphan scan. + return; + } + } else { + PAUSE(); + } + } + } + + // Phase 3: Delete on-disk files now that no readers remain. + CleanupOrphanedBucket(bucket_id); + + LOG(INFO) << "RollbackCommittedBucket: rolled back bucket " << bucket_id + << " with " << keys.size() << " keys"; +} + std::map>::iterator BucketStorageBackend::SelectEvictionCandidate() { // Must be called with mutex_ held (exclusive). @@ -2203,16 +2835,31 @@ BucketStorageBackend::SelectEvictionCandidate() { } } -BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction( - int64_t required_size) { +tl::expected +BucketStorageBackend::PrepareEviction( + int64_t required_size, const std::vector& write_keys) { PendingEviction result; + SharedMutexLocker lock(&mutex_); + + if (!write_keys.empty()) { + for (const auto& key : write_keys) { + if (object_bucket_map_.find(key) != object_bucket_map_.end() || + pending_eviction_keys_.find(key) != + pending_eviction_keys_.end() || + pending_write_keys_.find(key) != pending_write_keys_.end()) { + return tl::make_unexpected(ErrorCode::OBJECT_ALREADY_EXISTS); + } + } + result.write_keys = write_keys; + result.write_size = required_size; + pending_write_size_ += required_size; + pending_write_keys_.insert(write_keys.begin(), write_keys.end()); + } if (bucket_backend_config_.eviction_policy == BucketEvictionPolicy::NONE) { return result; } - SharedMutexLocker lock(&mutex_); - // Check actual disk space once before the loop. PrepareEviction only // removes metadata -- files are deleted later in FinalizeEviction -- so // re-checking disk space inside the loop would yield the same result @@ -2231,16 +2878,19 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction( if (!ec) { uint64_t actual_available = space_info.available; constexpr uint64_t kMinFreeSpace = 256 * kMB; - uint64_t req_sz = - required_size > 0 ? static_cast(required_size) : 0; + // Watermark eviction passes a synthetic required_size to drive + // quota-based cleanup. It is not a real incoming write, so it + // should not be counted as physical disk free-space demand. + uint64_t req_sz = (!write_keys.empty() && required_size > 0) + ? static_cast(required_size) + : 0; initial_disk_full = actual_available < req_sz + kMinFreeSpace; if (initial_disk_full) { deficit = req_sz + kMinFreeSpace - actual_available; - LOG(WARNING) - << "[Evict] Actual disk space too low: available=" - << actual_available << ", required=" << required_size - << ", deficit=" << deficit - << ". Will evict buckets to free space."; + LOG(WARNING) << "[Evict] Actual disk space too low: available=" + << actual_available << ", required=" << req_sz + << ", deficit=" << deficit + << ". Will evict buckets to free space."; } } else { LOG(WARNING) << "[Evict] Failed to get disk space info for " @@ -2251,10 +2901,14 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction( size_t evict_count = 0; constexpr size_t kMaxEvictionBuckets = 1000; uint64_t accumulated_freed_space = 0; + const int64_t synthetic_required_size = + write_keys.empty() ? required_size : 0; while (!buckets_.empty() && evict_count < kMaxEvictionBuckets) { - bool quota_exceeded = - total_size_ + required_size > bucket_backend_config_.max_total_size; + bool quota_exceeded = total_size_ + pending_eviction_size_ + + pending_write_size_ + + synthetic_required_size > + bucket_backend_config_.max_total_size; bool disk_still_full = initial_disk_full && (accumulated_freed_space < deficit); @@ -2276,20 +2930,25 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction( std::move(evict_it->second); buckets_.erase(evict_it); + int64_t evicted_size = evict_meta->meta_size; // Remove all keys belonging to this bucket from the object map. for (const auto& key : evict_meta->keys) { auto obj_it = object_bucket_map_.find(key); if (obj_it != object_bucket_map_.end() && obj_it->second.bucket_id == evict_id) { - total_size_ -= + const int64_t object_size = obj_it->second.data_size + obj_it->second.key_size; + total_size_ -= object_size; + evicted_size += object_size; object_bucket_map_.erase(obj_it); } } total_size_ -= evict_meta->meta_size; + result.evicted_size += evicted_size; // Collect for notification and file deletion. for (const auto& key : evict_meta->keys) { + pending_eviction_keys_.insert(key); result.keys.push_back(key); } accumulated_freed_space += @@ -2299,6 +2958,16 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction( evict_count++; } + const bool quota_exceeded = total_size_ + pending_eviction_size_ + + pending_write_size_ + + synthetic_required_size > + bucket_backend_config_.max_total_size; + pending_eviction_size_ += result.evicted_size; + if (!write_keys.empty() && quota_exceeded) { + RestorePreparedEvictionLocked(std::move(result)); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + if (!result.buckets.empty()) { LOG(INFO) << "[Evict] prepared: buckets=" << result.buckets.size() << " keys=" << result.keys.size() @@ -2308,15 +2977,107 @@ BucketStorageBackend::PendingEviction BucketStorageBackend::PrepareEviction( return result; } -void BucketStorageBackend::FinalizeEviction(const PendingEviction& pending) { +void BucketStorageBackend::RestorePreparedEviction(PendingEviction&& pending) { + SharedMutexLocker lock(&mutex_); + RestorePreparedEvictionLocked(std::move(pending)); +} + +void BucketStorageBackend::RestorePreparedEvictionLocked( + PendingEviction&& pending) { + ReleasePreparedWriteLocked(pending); + + CHECK_GE(pending_eviction_size_, pending.evicted_size); + pending_eviction_size_ -= pending.evicted_size; + for (const auto& key : pending.keys) { + pending_eviction_keys_.erase(key); + } + for (auto& [bucket_id, bucket_meta] : pending.buckets) { + if (!bucket_meta || buckets_.find(bucket_id) != buckets_.end()) { + continue; + } + + for (size_t i = 0; i < bucket_meta->keys.size(); ++i) { + const auto& key = bucket_meta->keys[i]; + const auto& object_meta = bucket_meta->metadatas[i]; + object_bucket_map_[key] = StorageObjectMetadata{ + bucket_id, object_meta.offset, object_meta.key_size, + object_meta.data_size, ""}; + total_size_ += object_meta.data_size + object_meta.key_size; + } + total_size_ += bucket_meta->meta_size; + if (bucket_backend_config_.eviction_policy == + BucketEvictionPolicy::LRU) { + lru_index_.emplace( + bucket_meta->last_access_ns_.load(std::memory_order_relaxed), + bucket_id); + } + buckets_.emplace(bucket_id, std::move(bucket_meta)); + } +} + +void BucketStorageBackend::CommitPreparedEviction( + const PendingEviction& pending) { + if (pending.evicted_size == 0) { + return; + } + + SharedMutexLocker lock(&mutex_); + CHECK_GE(pending_eviction_size_, pending.evicted_size); + pending_eviction_size_ -= pending.evicted_size; + for (const auto& key : pending.keys) { + pending_eviction_keys_.erase(key); + } +} + +void BucketStorageBackend::ReleasePreparedWrite( + const PendingEviction& pending) { + if (pending.write_size == 0) { + return; + } + + SharedMutexLocker lock(&mutex_); + ReleasePreparedWriteLocked(pending); +} + +void BucketStorageBackend::ReleasePreparedWriteLocked( + const PendingEviction& pending) { + CHECK_GE(pending_write_size_, pending.write_size); + pending_write_size_ -= pending.write_size; + for (const auto& key : pending.write_keys) { + pending_write_keys_.erase(key); + } +} + +tl::expected BucketStorageBackend::FinalizeEviction( + const PendingEviction& pending) { namespace fs = std::filesystem; constexpr int kMaxSpinIterations = 1000; constexpr auto kMaxWaitTime = std::chrono::seconds(10); + size_t cleanup_failed_count = 0; for (const auto& [bucket_id, bucket_meta] : pending.buckets) { + bool bucket_cleanup_failed = false; + + // The master has already committed the replica removal. Attempt to + // remove persisted metadata before waiting for readers. When this + // succeeds, a later timeout or data-file deletion failure leaves an + // orphan data file instead of a bucket that Init() could recover. + std::error_code ec; + auto meta_path = GetBucketMetadataPath(bucket_id); + if (meta_path) { + fs::remove(meta_path.value(), ec); + if (ec && ec != std::errc::no_such_file_or_directory) { + LOG(ERROR) + << "FinalizeEviction: failed to remove metadata file: " + << meta_path.value() << ", error: " << ec.message(); + bucket_cleanup_failed = true; + } + } + // Wait for in-flight reads that started before the bucket was removed // from the metadata maps in PrepareEviction. + bool timed_out = false; int spin_count = 0; auto wait_start = std::chrono::steady_clock::now(); while (bucket_meta->inflight_reads_.load(std::memory_order_acquire) > @@ -2332,14 +3093,19 @@ void BucketStorageBackend::FinalizeEviction(const PendingEviction& pending) { << bucket_id << ", inflight_reads=" << bucket_meta->inflight_reads_.load( std::memory_order_relaxed); + timed_out = true; break; } } else { PAUSE(); } } + if (timed_out) { + cleanup_failed_count++; + continue; + } - std::error_code ec; + ec.clear(); auto data_path = GetBucketDataPath(bucket_id); if (data_path) { // Evict the cached file handle before deleting the file to prevent @@ -2355,29 +3121,80 @@ void BucketStorageBackend::FinalizeEviction(const PendingEviction& pending) { // up on service restart. LOG(ERROR) << "FinalizeEviction: failed to remove data file: " << data_path.value() << ", error: " << ec.message(); + bucket_cleanup_failed = true; } } - auto meta_path = GetBucketMetadataPath(bucket_id); - if (meta_path) { - ec.clear(); - fs::remove(meta_path.value(), ec); - if (ec && ec != std::errc::no_such_file_or_directory) { - LOG(ERROR) - << "FinalizeEviction: failed to remove metadata file: " - << meta_path.value() << ", error: " << ec.message(); - } + if (bucket_cleanup_failed) { + cleanup_failed_count++; } } if (!pending.buckets.empty()) { - LOG(INFO) << "[Evict] finalized: deleted " << pending.buckets.size() - << " bucket(s)"; + LOG(INFO) << "[Evict] finalized: attempted=" << pending.buckets.size() + << " cleanup_failed=" << cleanup_failed_count; + } + if (cleanup_failed_count != 0) { + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + return {}; +} + +tl::expected, ErrorCode> +BucketStorageBackend::EvictAboveDiskWatermark( + double high_watermark_ratio, double low_watermark_ratio, + EvictionHandler eviction_handler) { + std::vector evicted_keys; + if (bucket_backend_config_.eviction_policy == BucketEvictionPolicy::NONE || + bucket_backend_config_.max_total_size <= 0) { + return evicted_keys; + } + if (!initialized_.load(std::memory_order_acquire)) { + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + + int64_t total_size = 0; + int64_t max_total_size = bucket_backend_config_.max_total_size; + { + SharedMutexLocker lock(&mutex_, shared_lock); + total_size = total_size_; + } + + const auto high_watermark_bytes = + static_cast(max_total_size * high_watermark_ratio); + if (total_size <= high_watermark_bytes) { + return evicted_keys; + } + + const auto target_total_size = + static_cast(max_total_size * low_watermark_ratio); + const auto synthetic_required_size = max_total_size - target_total_size; + auto prepare_result = PrepareEviction(synthetic_required_size); + if (!prepare_result) { + return tl::make_unexpected(prepare_result.error()); + } + PendingEviction pending = std::move(prepare_result.value()); + evicted_keys = pending.keys; + + if (eviction_handler && !pending.keys.empty()) { + auto notify_result = eviction_handler(pending.keys); + if (!notify_result) { + RestorePreparedEviction(std::move(pending)); + return tl::make_unexpected(notify_result.error()); + } + } + CommitPreparedEviction(pending); + auto finalize_result = FinalizeEviction(pending); + if (!finalize_result) { + LOG(ERROR) + << "FinalizeEviction failed after master committed eviction; " + "returning evicted keys: " + << finalize_result.error(); } + return evicted_keys; } tl::expected BucketStorageBackend::DeleteBucket( int64_t bucket_id) { namespace fs = std::filesystem; - std::shared_ptr bucket_metadata; std::vector keys_to_remove; @@ -2474,6 +3291,30 @@ tl::expected BucketStorageBackend::DeleteBucket( return {}; } +void BucketStorageBackend::RemoveAll() { + namespace fs = std::filesystem; + + // Collect all bucket IDs under exclusive lock, then release. + std::vector bucket_ids; + { + SharedMutexLocker lock(&mutex_); + bucket_ids.reserve(buckets_.size()); + for (const auto& [id, _] : buckets_) { + bucket_ids.push_back(id); + } + } + + for (int64_t id : bucket_ids) { + auto result = DeleteBucket(id); + if (!result) { + LOG(WARNING) << "RemoveAll: DeleteBucket failed for bucket_id=" + << id << ", error=" << toString(result.error()); + } + } + + LOG(INFO) << "RemoveAll: removed " << bucket_ids.size() << " bucket(s)"; +} + tl::expected BucketStorageBackend::StoreBucketMetadata( int64_t id, std::shared_ptr metadata) { auto meta_path_res = GetBucketMetadataPath(id); @@ -2695,16 +3536,61 @@ BucketStorageBackend::GetFileInstance() const { // ============================================================================ OffsetAllocatorStorageBackend::OffsetAllocatorStorageBackend( - const FileStorageConfig& file_storage_config_) + const FileStorageConfig& file_storage_config_, + const OffsetAllocatorBackendConfig& offset_backend_config) : StorageBackendInterface(file_storage_config_), - storage_path_(file_storage_config_.storage_filepath) { + storage_path_(file_storage_config_.storage_filepath), + cfg_(offset_backend_config) { capacity_ = file_storage_config_.total_size_limit; } +OffsetAllocatorStorageBackend::~OffsetAllocatorStorageBackend() { + try { + if (cfg_.persist_mode == OffsetPersistMode::kDisabled) return; + if (!initialized_.load(std::memory_order_acquire)) return; + if (!data_file_) return; + if (test_skip_final_checkpoint_.load(std::memory_order_relaxed)) { + // Test hook: simulate an abrupt crash (no final checkpoint). + return; + } + + if (!metadata_dirty_.load(std::memory_order_relaxed)) { + // Nothing to persist. + return; + } + + // Best-effort final checkpoint on graceful shutdown. + auto sync_res = data_file_->datasync(); + if (!sync_res) { + LOG(ERROR) << "Final checkpoint datasync failed: " + << static_cast(sync_res.error()); + return; + } + + auto save_res = SaveMetadata(all_evicted_this_batch_); + if (!save_res) { + LOG(ERROR) << "Final checkpoint SaveMetadata failed: " + << static_cast(save_res.error()); + return; + } + + all_evicted_this_batch_.clear(); + LOG(INFO) << "Final persistence checkpoint completed on shutdown"; + } catch (const std::exception& e) { + LOG(ERROR) << "Final checkpoint threw: " << e.what(); + } catch (...) { + LOG(ERROR) << "Final checkpoint threw unknown exception"; + } +} + std::string OffsetAllocatorStorageBackend::GetDataFilePath() const { return (std::filesystem::path(storage_path_) / "kv_cache.data").string(); } +std::string OffsetAllocatorStorageBackend::GetMetaFilePath() const { + return (std::filesystem::path(storage_path_) / "kv_cache.meta").string(); +} + //----------------------------------------------------------------------------- tl::expected OffsetAllocatorStorageBackend::Init() { @@ -2722,8 +3608,130 @@ tl::expected OffsetAllocatorStorageBackend::Init() { return tl::make_unexpected(ErrorCode::INVALID_PARAMS); } + // Ensure storage path exists + { + std::error_code ec; + fs::create_directories(storage_path_, ec); + if (ec) { + LOG(ERROR) << "Failed to create storage directory: " + << storage_path_; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + } + + // ---- Resolve watermarks ---- + high_watermark_bytes_ = + cfg_.high_watermark_bytes > 0 + ? cfg_.high_watermark_bytes + : static_cast(capacity_ * cfg_.high_ratio); + low_watermark_bytes_ = + cfg_.low_watermark_bytes > 0 + ? cfg_.low_watermark_bytes + : static_cast(capacity_ * cfg_.low_ratio); + high_watermark_keys_ = + cfg_.high_watermark_keys > 0 + ? cfg_.high_watermark_keys + : static_cast(file_storage_config_.total_keys_limit * + cfg_.keys_high_ratio); + low_watermark_keys_ = + cfg_.low_watermark_keys > 0 + ? cfg_.low_watermark_keys + : static_cast(file_storage_config_.total_keys_limit * + cfg_.keys_low_ratio); + + // Auto-nudge ratio-derived low watermarks when integer + // truncation collapses them to the same value as high + if (cfg_.low_watermark_bytes == 0 && high_watermark_bytes_ > 0 && + low_watermark_bytes_ >= high_watermark_bytes_) { + low_watermark_bytes_ = + std::max(1, high_watermark_bytes_ - 1); + } + if (cfg_.low_watermark_keys == 0 && high_watermark_keys_ > 0 && + low_watermark_keys_ >= high_watermark_keys_) { + low_watermark_keys_ = + std::max(1, high_watermark_keys_ - 1); + } + + // Validate watermarks + if (low_watermark_bytes_ >= high_watermark_bytes_) { + LOG(ERROR) << "Invalid watermark: low_bytes=" + << low_watermark_bytes_ + << " >= high_bytes=" << high_watermark_bytes_; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + if (low_watermark_keys_ >= high_watermark_keys_) { + LOG(ERROR) << "Invalid watermark: low_keys=" << low_watermark_keys_ + << " >= high_keys=" << high_watermark_keys_; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + + // Clamp watermarks to not exceed capacity / total_keys_limit + if (high_watermark_bytes_ > static_cast(capacity_)) { + LOG(WARNING) << "high_watermark_bytes clamped from " + << high_watermark_bytes_ + << " to capacity=" << capacity_; + high_watermark_bytes_ = static_cast(capacity_); + low_watermark_bytes_ = + std::min(low_watermark_bytes_, high_watermark_bytes_ - 1); + } + if (high_watermark_keys_ > file_storage_config_.total_keys_limit) { + LOG(WARNING) << "high_watermark_keys clamped from " + << high_watermark_keys_ << " to total_keys_limit=" + << file_storage_config_.total_keys_limit; + high_watermark_keys_ = file_storage_config_.total_keys_limit; + low_watermark_keys_ = + std::min(low_watermark_keys_, high_watermark_keys_ - 1); + } + + // Guard against zero low-watermark on very small capacity + if (low_watermark_bytes_ <= 0 && high_watermark_bytes_ > 0) { + low_watermark_bytes_ = + std::max(1, high_watermark_bytes_ / 2); + } + if (low_watermark_keys_ <= 0 && high_watermark_keys_ > 0) { + low_watermark_keys_ = + std::max(1, high_watermark_keys_ / 2); + } + + // ---- Recovery path ---- + if (cfg_.persist_mode != OffsetPersistMode::kDisabled) { + const RecoveryResult recovery = TryRecoverFromMetadata(); + if (recovery == RecoveryResult::kTransientError) { + // A momentary resource error (fd exhaustion, OOM, I/O + // error) must not wipe a recoverable cache: fail Init + // and let the operator retry instead of falling through + // to the fresh-start path below. + LOG(ERROR) << "Transient error during recovery; failing " + "Init to protect persisted data"; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + if (recovery == RecoveryResult::kRecovered) { + last_persist_time_us_.store( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(), + std::memory_order_relaxed); + initialized_.store(true, std::memory_order_release); + LOG(INFO) << "OffsetAllocatorStorageBackend recovered: " + << total_keys_.load() << " keys, " + << total_size_.load() << " bytes"; + return {}; + } + // kNoMeta / kCorrupt: fall through to a fresh start. + } + + // ---- Fresh start path ---- + LOG(INFO) << "Fresh start path"; + + // Close any recovery fd and clean up stale metadata + data_file_.reset(); + { + std::error_code ec; + fs::remove(GetMetaFilePath(), ec); + fs::remove(GetMetaFilePath() + ".tmp", ec); + } + // Clear in-memory maps (V1: no persistence, start fresh) - // Lock all shards to ensure exclusive access during initialization { std::vector> shard_locks; shard_locks.reserve(kNumShards); @@ -2739,27 +3747,7 @@ tl::expected OffsetAllocatorStorageBackend::Init() { // Get data file path data_file_path_ = GetDataFilePath(); - // RAII wrapper to ensure fd is closed on all error paths - struct FdGuard { - int fd; - explicit FdGuard(int fd) : fd(fd) {} - ~FdGuard() { - if (fd >= 0) { - close(fd); - } - } - // Release ownership (caller takes responsibility) - int release() { - int ret = fd; - fd = -1; - return ret; - } - // Get fd without releasing (for operations) - int get() const { return fd; } - }; - // Open/truncate data file in read-write mode - // We need raw fd for fallocate, so open directly int flags = O_CLOEXEC | O_RDWR | O_CREAT | O_TRUNC; int raw_fd = open(data_file_path_.c_str(), flags, 0644); if (raw_fd < 0) { @@ -2782,25 +3770,72 @@ tl::expected OffsetAllocatorStorageBackend::Init() { // Release fd to StorageFile (takes ownership and will close it) #ifdef USE_URING if (file_storage_config_.use_uring) { - data_file_ = std::make_unique( + data_file_ = std::make_shared( data_file_path_, fd_guard.release(), 32, true); } else #endif { - data_file_ = std::make_unique(data_file_path_, + data_file_ = std::make_shared(data_file_path_, fd_guard.release()); } + if (cfg_.persist_mode != OffsetPersistMode::kDisabled) { + data_file_->SetDeleteOnWriteFail(false); + } - // Create allocator with base=0, size=capacity - allocator_ = offset_allocator::OffsetAllocator::create(0, capacity_); + // Create allocator with tuned node capacity + constexpr int64_t kMinObjectSize = 256; + constexpr int64_t kMaxNodeRamBytes = + 512LL * 1024 * 1024; // 512MB node RAM budget + constexpr uint32_t kRamBasedMaxNodes = + static_cast(kMaxNodeRamBytes / 56); + constexpr uint32_t kAbsoluteMaxNodes = + std::min(kRamBasedMaxNodes, 32U << 20); + + uint32_t max_nodes = (1U << 20); // default 1M nodes + if (cfg_.max_capacity_nodes > 0) { + if (cfg_.max_capacity_nodes > kAbsoluteMaxNodes) { + LOG(WARNING) + << "max_capacity_nodes " << cfg_.max_capacity_nodes + << " exceeds RAM budget; clamped to " << kAbsoluteMaxNodes; + max_nodes = kAbsoluteMaxNodes; + } else { + max_nodes = static_cast(cfg_.max_capacity_nodes); + } + } else { + int64_t auto_nodes = std::max( + 1LL << 20, std::min(capacity_ / kMinObjectSize, + kAbsoluteMaxNodes)); + max_nodes = static_cast(auto_nodes); + } + uint32_t init_nodes = std::min(128U * 1024, max_nodes); + allocator_ = offset_allocator::OffsetAllocator::create( + 0, capacity_, init_nodes, max_nodes); if (!allocator_) { LOG(ERROR) << "Failed to create OffsetAllocator"; return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } + // Initialize eviction index + { + MutexLocker ev(&eviction_mutex_); + fifo_index_.clear(); + insert_seq_.store(0, std::memory_order_relaxed); + } + + // Initialize persist timestamp (avoid immediate checkpoint + // on first write) + last_persist_time_us_.store( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(), + std::memory_order_relaxed); + initialized_.store(true, std::memory_order_release); LOG(INFO) << "OffsetAllocatorStorageBackend initialized, capacity: " - << capacity_ << " bytes, data file: " << data_file_path_; + << capacity_ + << " bytes, high_watermark: " << high_watermark_bytes_ + << " bytes, " << high_watermark_keys_ << " keys" + << ", data file: " << data_file_path_; } catch (const std::exception& e) { LOG(ERROR) << "OffsetAllocatorStorageBackend initialize error: " << e.what(); @@ -2810,31 +3845,977 @@ tl::expected OffsetAllocatorStorageBackend::Init() { return {}; } +// EvictToMakeRoom //----------------------------------------------------------------------------- -tl::expected OffsetAllocatorStorageBackend::BatchOffload( - const std::unordered_map>& batch_object, - std::function& keys, - std::vector& metadatas)> - complete_handler, - std::function& /*evicted_keys*/)> - /*eviction_handler*/) { - if (!initialized_.load(std::memory_order_acquire)) { - LOG(ERROR) - << "Storage backend is not initialized. Call Init() before use."; - return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); - } - if (batch_object.empty()) { - LOG(ERROR) << "BatchOffload called with empty batch"; - return tl::make_unexpected(ErrorCode::INVALID_KEY); - } +//----------------------------------------------------------------------------- +// ShouldPersistNow +//----------------------------------------------------------------------------- - auto enable_offloading_res = IsEnableOffloading(); - if (!enable_offloading_res) { - return tl::make_unexpected(enable_offloading_res.error()); - } - if (!enable_offloading_res.value()) { - return tl::make_unexpected(ErrorCode::KEYS_ULTRA_LIMIT); +bool OffsetAllocatorStorageBackend::ShouldPersistNow() const { + if (cfg_.persist_mode == OffsetPersistMode::kDisabled) return false; + if (cfg_.persist_mode == OffsetPersistMode::kStrict) return true; + using TimePointUs = std::chrono::time_point; + auto last_tp = TimePointUs(std::chrono::microseconds( + last_persist_time_us_.load(std::memory_order_relaxed))); + return (std::chrono::steady_clock::now() - last_tp) >= + std::chrono::seconds(cfg_.persist_interval_seconds); +} + +//----------------------------------------------------------------------------- +// SaveMetadata +//----------------------------------------------------------------------------- + +tl::expected OffsetAllocatorStorageBackend::SaveMetadata( + const std::unordered_set& evicted_keys_this_batch) { + namespace fs = std::filesystem; + + if (!metadata_dirty_.exchange(false, std::memory_order_relaxed) && + evicted_keys_this_batch.empty()) { + return {}; + } + + int step = test_metadata_write_failure_step_.exchange( + 0, std::memory_order_relaxed); + try { + struct DirtyGuard { + std::atomic& dirty; + std::atomic& failures; + bool active = true; + DirtyGuard(std::atomic& d, std::atomic& f) + : dirty(d), failures(f) {} + ~DirtyGuard() { + if (active) { + dirty.store(true, std::memory_order_relaxed); + failures.fetch_add(1, std::memory_order_relaxed); + } + } + void disarm() { active = false; } + } dirty_guard(metadata_dirty_, metadata_save_failures_); + + auto t_start = std::chrono::steady_clock::now(); + + // 1. Serialize allocator + std::vector alloc_buf; + ErrorCode ec = serialize_to(*allocator_, alloc_buf); + if (ec != ErrorCode::OK || alloc_buf.empty()) { + return tl::make_unexpected( + ec != ErrorCode::OK ? ec : ErrorCode::INTERNAL_ERROR); + } + if (step == 1) return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + + // 2. Build metadata struct + OffsetAllocatorPersistedMetadata meta; + meta.version = kOffsetAllocatorPersistVersion; + meta.allocator_state.assign( + reinterpret_cast(alloc_buf.data()), alloc_buf.size()); + meta.evicted_keys_this_batch.assign(evicted_keys_this_batch.begin(), + evicted_keys_this_batch.end()); + + { + MutexLocker ev(&eviction_mutex_); + meta.insert_seq = insert_seq_.load(std::memory_order_relaxed); + meta.fifo_entries.reserve(fifo_index_.size()); + for (const auto& [seq, key] : fifo_index_) { + meta.fifo_entries.push_back({seq, key}); + } + } + + // 3. Serialize metadata + std::string buf; + try { + struct_pb::to_pb(meta, buf); + } catch (const std::exception& e) { + LOG(ERROR) << "struct_pb::to_pb failed in SaveMetadata: " + << e.what(); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + if (buf.empty()) { + LOG(ERROR) << "struct_pb produced empty metadata buffer"; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + + // 4. Atomic write via tmp + rename + std::string meta_path = GetMetaFilePath(); + std::string tmp_path = meta_path + ".tmp"; + + int fd = open(tmp_path.c_str(), + O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644); + if (fd < 0) { + LOG(ERROR) << "Failed to create " << tmp_path << ": " + << strerror(errno); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + + const char* ptr = buf.data(); + size_t remaining = buf.size(); + while (remaining > 0) { + ssize_t n = write(fd, ptr, remaining); + if (n <= 0) { + if (n < 0 && errno == EINTR) continue; + LOG(ERROR) << "write failed on " << tmp_path << ": " + << (n < 0 ? strerror(errno) : "returned 0"); + close(fd); + unlink(tmp_path.c_str()); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + ptr += n; + remaining -= n; + } + if (step == 2) { + close(fd); + unlink(tmp_path.c_str()); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + + // Use fsync (not fdatasync) for the metadata file: metadata is + // small and we need inode metadata (file size) durable as well. + if (fsync(fd) != 0) { + LOG(ERROR) << "fsync failed on " << tmp_path; + close(fd); + unlink(tmp_path.c_str()); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + close(fd); + + if (rename(tmp_path.c_str(), meta_path.c_str()) != 0) { + LOG(ERROR) << "rename " << tmp_path << " -> " << meta_path + << " failed: " << strerror(errno); + unlink(tmp_path.c_str()); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + if (step == 3) return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + + // fsync parent directory + fs::path parent = fs::path(meta_path).parent_path(); + if (parent.empty()) parent = "."; + int dir_fd = open(parent.c_str(), O_RDONLY | O_DIRECTORY | O_CLOEXEC); + if (dir_fd < 0) { + LOG(ERROR) << "open parent dir failed: " << strerror(errno); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + FdGuard dir_guard(dir_fd); + if (fsync(dir_fd) != 0) { + LOG(ERROR) << "fsync parent dir failed: " << strerror(errno); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); + } + + dirty_guard.disarm(); + + last_save_metadata_cost_us_.store( + std::chrono::duration_cast( + std::chrono::steady_clock::now() - t_start) + .count(), + std::memory_order_relaxed); + + return {}; + } catch (const std::exception& e) { + // std::bad_alloc or other exception during serialization. + // DirtyGuard restores metadata_dirty_ on scope exit. + LOG(ERROR) << "SaveMetadata failed with exception: " << e.what(); + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } +} + +//----------------------------------------------------------------------------- +// LoadMetadata +//----------------------------------------------------------------------------- + +tl::expected +OffsetAllocatorStorageBackend::LoadMetadata() { + std::string meta_path = GetMetaFilePath(); + + int fd = open(meta_path.c_str(), O_RDONLY | O_CLOEXEC); + if (fd < 0) { + LOG(ERROR) << "Failed to open " << meta_path << ": " << strerror(errno); + return tl::make_unexpected(ErrorCode::FILE_OPEN_FAIL); + } + FdGuard closer(fd); + + struct stat st; + if (fstat(fd, &st) != 0) { + LOG(ERROR) << "fstat failed on " << meta_path << ": " + << strerror(errno); + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + if (st.st_size == 0) { + LOG(ERROR) << "Metadata file is empty: " << meta_path; + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + + std::string buf(static_cast(st.st_size), '\0'); + size_t read_bytes = 0; + char* ptr = buf.data(); + while (read_bytes < static_cast(st.st_size)) { + ssize_t n = pread(fd, ptr, static_cast(st.st_size) - read_bytes, + static_cast(read_bytes)); + if (n < 0) { + if (errno == EINTR) continue; + LOG(ERROR) << "Read failed from " << meta_path << ": " + << strerror(errno); + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + if (n == 0) { + LOG(ERROR) << "Unexpected EOF from " << meta_path; + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + read_bytes += static_cast(n); + ptr += n; + } + + OffsetAllocatorPersistedMetadata meta; + try { + struct_pb::from_pb(meta, buf); + } catch (const std::exception& e) { + LOG(ERROR) << "struct_pb::from_pb failed for " << meta_path << ": " + << e.what(); + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + + if (meta.version != kOffsetAllocatorPersistVersion) { + LOG(ERROR) << "Unsupported metadata version " << meta.version << " in " + << meta_path << " (expected " + << kOffsetAllocatorPersistVersion << ")"; + return tl::make_unexpected(ErrorCode::INVALID_PARAMS); + } + if (meta.allocator_state.empty()) { + LOG(ERROR) << "Empty allocator_state in " << meta_path; + return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); + } + + return meta; +} + +//----------------------------------------------------------------------------- +// RebuildShardMapsFromAllocator +//----------------------------------------------------------------------------- + +void OffsetAllocatorStorageBackend::RebuildShardMapsFromAllocator( + uint64_t checkpoint_insert_seq) { + int64_t rebuilt = 0; + int64_t skipped = 0; + + // Lock all shards during rebuild (Init is single-threaded, but + // locking documents the invariant). + std::vector> shard_locks; + shard_locks.reserve(kNumShards); + for (size_t i = 0; i < kNumShards; ++i) { + shard_locks.emplace_back( + std::make_unique(&shards_[i].mutex)); + } + total_size_.store(0, std::memory_order_relaxed); + total_keys_.store(0, std::memory_order_relaxed); + + // Reused across records for CRC streaming (avoids a fresh 1MB + // allocation per record). + std::vector chunk_buf; + + allocator_->visit_used_nodes([&](uint64_t real_offset, uint64_t alloc_size, + uint32_t node_index) { + // Helper: free a corrupt/unrecoverable node via RAII. The + // temporary handle goes out of scope immediately, returning the + // extent to the allocator. requested_size keeps the + // m_allocated_size metric consistent with the serialized state: + // when the header parsed, the exact record_size is known (it + // equals the requested size passed to allocate()); otherwise the + // node's bin size is the best available estimate. + auto free_leaked_node = [&](uint64_t requested_size) { + auto h = allocator_->createHandleAtNode(node_index, real_offset, + requested_size); + }; + + if (real_offset + RecordHeader::SIZE > + static_cast(capacity_)) { + LOG(ERROR) << "[Recover] offset " << real_offset + << " + header exceeds capacity " << capacity_; + free_leaked_node(alloc_size); + ++skipped; + return; + } + + char hdr_buf[RecordHeader::SIZE]; + iovec hdr_iov = {hdr_buf, sizeof(hdr_buf)}; + auto hdr_res = data_file_->vector_read(&hdr_iov, 1, real_offset); + if (!hdr_res || hdr_res.value() != RecordHeader::SIZE) { + LOG(ERROR) << "[Recover] Failed to read header at " << real_offset; + free_leaked_node(alloc_size); + ++skipped; + return; + } + RecordHeader header = RecordHeader::ReadFrom(hdr_buf); + + if (header.key_len > kMaxKeyLen) { + LOG(ERROR) << "[Recover] Invalid key_len " << header.key_len + << " at " << real_offset; + free_leaked_node(alloc_size); + ++skipped; + return; + } + + // Records carrying flag bits this build does not understand were + // written by a newer format; refuse them rather than guess. + if ((header.flags & ~RecordHeader::kKnownFlags) != 0) { + LOG(ERROR) << "[Recover] Unknown flags " << header.flags << " at " + << real_offset; + free_leaked_node(alloc_size); + ++skipped; + return; + } + + uint64_t record_size = + RecordHeader::RecordSize(header.key_len, header.value_len); + if (record_size > alloc_size) { + LOG(ERROR) << "[Recover] record_size " << record_size + << " > alloc_size " << alloc_size << " at " + << real_offset; + free_leaked_node(alloc_size); + ++skipped; + return; + } + if (record_size > + static_cast(std::numeric_limits::max())) { + LOG(ERROR) << "[Recover] record_size " << record_size + << " exceeds uint32_t at " << real_offset; + free_leaked_node(alloc_size); + ++skipped; + return; + } + if (real_offset + record_size > static_cast(capacity_)) { + LOG(ERROR) << "[Recover] record past capacity at " << real_offset; + free_leaked_node(alloc_size); + ++skipped; + return; + } + + // Post-checkpoint write guard: a record stamped at or after the + // checkpoint's insert_seq was written after the checkpoint fired. + // Its extent may hold a torn write (the crash hit mid-pwritev), or + // it may belong to a different key whose space was freed and reused + // after the checkpoint. Either way the checkpoint cannot vouch for + // it, so the node is dropped and freed. + if (header.seq >= checkpoint_insert_seq) { + LOG(INFO) << "[Recover] record seq " << header.seq + << " >= checkpoint insert_seq " << checkpoint_insert_seq + << " at " << real_offset + << " -- post-checkpoint write, dropping"; + free_leaked_node(record_size); + ++skipped; + return; + } + + std::string key(header.key_len, '\0'); + iovec key_iov = {key.data(), static_cast(header.key_len)}; + auto key_res = data_file_->vector_read( + &key_iov, 1, real_offset + RecordHeader::SIZE); + if (!key_res || + key_res.value() != static_cast(header.key_len)) { + LOG(ERROR) << "[Recover] Failed to read key at " << real_offset; + free_leaked_node(record_size); + ++skipped; + return; + } + + // Integrity check, only for checksummed records: stream the value + // in chunks and verify the CRC-32C over header-prefix + key + + // value. This rejects torn writes whose header survived the + // crash but whose tail is stale or zero-filled. Unchecksummed + // records (kFlagHasCrc clear, e.g. written with CRC disabled or + // by a DMA writer) skip the value scan entirely: they are trusted + // via the checkpoint's seq guard alone, and recovery stays + // O(header + key) for them. + if (header.HasCrc()) { + Crc32c crc; + crc.Extend(hdr_buf, RecordHeader::PREFIX_SIZE); + crc.Extend(key.data(), key.size()); + { + constexpr size_t kCrcChunkSize = 1 << 20; // 1MB + chunk_buf.resize( + std::min(header.value_len, kCrcChunkSize)); + uint64_t remaining = header.value_len; + uint64_t value_offset = + real_offset + + RecordHeader::ValueOffsetInRecord(header.key_len); + bool read_failed = false; + while (remaining > 0) { + const size_t chunk = + std::min(remaining, chunk_buf.size()); + iovec value_iov = {chunk_buf.data(), chunk}; + auto value_res = + data_file_->vector_read(&value_iov, 1, value_offset); + if (!value_res || value_res.value() != chunk) { + LOG(ERROR) << "[Recover] Failed to read value at " + << real_offset; + read_failed = true; + break; + } + crc.Extend(chunk_buf.data(), chunk); + value_offset += chunk; + remaining -= chunk; + } + if (read_failed) { + free_leaked_node(record_size); + ++skipped; + return; + } + } + if (crc.Final() != header.crc32) { + LOG(ERROR) << "[Recover] CRC mismatch at " << real_offset + << " (stored " << header.crc32 << ", computed " + << crc.Final() << ") -- torn or stale record"; + free_leaked_node(record_size); + ++skipped; + return; + } + } + + // Duplicate resolution BEFORE constructing any handle: a failed + // emplace would have moved allocation_ptr into a discarded + // temporary, freeing the very extent we might need to keep. + // Note: ObjectEntry.fifo_seq temporarily carries the record's seq + // for this comparison; RestoreAndRepairFifoIndex overwrites it + // with the persisted (or repaired) FIFO seq later. + size_t shard_idx = ShardForKey(key); + auto& shard_map = shards_[shard_idx].map; + auto existing = shard_map.find(key); + if (existing != shard_map.end() && + header.seq <= existing->second.fifo_seq) { + // Duplicate key (two extents were both marked used at the + // checkpoint, e.g. an overwrite whose old extent was pinned + // by a reader) and the already-inserted record is the newer + // version. Free this older extent via RAII. + LOG(INFO) << "[Recover] Duplicate key " << key << " at " + << real_offset << " (seq " << header.seq + << ") is older than existing record at " + << existing->second.offset << " (seq " + << existing->second.fifo_seq + << ") -- kept existing, freed duplicate (RAII)"; + free_leaked_node(record_size); + ++skipped; + return; + } + + auto handle = allocator_->createHandleAtNode(node_index, real_offset, + record_size); + if (!handle.has_value()) { + LOG(ERROR) << "[Recover] createHandleAtNode failed for " << key + << " at " << real_offset; + free_leaked_node(record_size); + ++skipped; + return; + } + + auto allocation_ptr = std::make_shared( + std::move(handle.value())); + + if (existing != shard_map.end()) { + // The new record is the newer version: replace the entry. + // The replaced handle frees the older extent via RAII with + // its own exact requested size. + LOG(INFO) << "[Recover] Duplicate key " << key << " at " + << real_offset << " (seq " << header.seq + << ") replaces older record at " + << existing->second.offset << " (seq " + << existing->second.fifo_seq << ")"; + const int64_t size_delta = + static_cast(record_size) - + static_cast(existing->second.total_size); + existing->second = + ObjectEntry(real_offset, static_cast(record_size), + header.value_len, std::move(allocation_ptr), + /*fifo_seq=*/header.seq); + total_size_.fetch_add(size_delta, std::memory_order_relaxed); + ++rebuilt; + return; + } + + shard_map.emplace( + key, ObjectEntry(real_offset, static_cast(record_size), + header.value_len, std::move(allocation_ptr), + /*fifo_seq=*/header.seq)); + + total_size_.fetch_add(static_cast(record_size), + std::memory_order_relaxed); + total_keys_.fetch_add(1, std::memory_order_relaxed); + ++rebuilt; + }); + + LOG(INFO) << "[Recover] RebuildShardMaps: " << rebuilt << " keys rebuilt, " + << skipped << " skipped (corrupt)"; + if (rebuilt == 0) { + LOG(WARNING) << "[Recover] No records recovered -- data file " + "may be empty or corrupt"; + } +} + +//----------------------------------------------------------------------------- +// RestoreAndRepairFifoIndex +//----------------------------------------------------------------------------- + +void OffsetAllocatorStorageBackend::RestoreAndRepairFifoIndex( + const OffsetAllocatorPersistedMetadata& meta) { + uint64_t max_seq = meta.insert_seq; + for (const auto& [seq, _] : meta.fifo_entries) { + max_seq = std::max(max_seq, seq); + } + insert_seq_.store(max_seq + 1, std::memory_order_relaxed); + + { + MutexLocker ev(&eviction_mutex_); + fifo_index_.clear(); + std::unordered_set fifo_keys; + + // Phase A: restore from persisted entries + for (const auto& [seq, key] : meta.fifo_entries) { + if (seq >= meta.insert_seq) { + LOG(WARNING) << "[Recover] seq " << seq << " >= insert_seq " + << meta.insert_seq << " -- skipping"; + continue; + } + size_t shard_idx = ShardForKey(key); + SharedMutexLocker lk(&shards_[shard_idx].mutex); + auto it = shards_[shard_idx].map.find(key); + if (it == shards_[shard_idx].map.end()) { + LOG(WARNING) << "[Recover] key " << key + << " not in shard map -- skipping"; + continue; + } + if (fifo_index_.count(seq)) { + LOG(WARNING) << "[Recover] Duplicate seq " << seq + << " -- deferred to Phase B"; + continue; + } + fifo_index_[seq] = key; + fifo_keys.insert(key); + it->second.fifo_seq = seq; + } + + // Phase B: repair -- assign new seq to any shard-map entry + // not yet in the FIFO index + for (size_t s = 0; s < kNumShards; ++s) { + SharedMutexLocker lk(&shards_[s].mutex); + for (auto& [key, entry] : shards_[s].map) { + if (!fifo_keys.count(key)) { + uint64_t new_seq = + insert_seq_.fetch_add(1, std::memory_order_relaxed); + entry.fifo_seq = new_seq; + fifo_index_[new_seq] = key; + fifo_keys.insert(key); + } + } + } + + // Phase C: delete evicted keys that were resurrected + for (const auto& key : meta.evicted_keys_this_batch) { + size_t shard_idx = ShardForKey(key); + SharedMutexLocker lk(&shards_[shard_idx].mutex); + auto it = shards_[shard_idx].map.find(key); + if (it != shards_[shard_idx].map.end()) { + total_size_.fetch_sub( + static_cast(it->second.total_size), + std::memory_order_relaxed); + total_keys_.fetch_sub(1, std::memory_order_relaxed); + fifo_index_.erase(it->second.fifo_seq); + shards_[shard_idx].map.erase(it); + } + } + } + + LOG(INFO) << "[Recover] FIFO index restored: " << fifo_index_.size() + << " entries"; +} + +//----------------------------------------------------------------------------- +// TryRecoverFromMetadata +//----------------------------------------------------------------------------- + +OffsetAllocatorStorageBackend::RecoveryResult +OffsetAllocatorStorageBackend::TryRecoverFromMetadata() { + try { + namespace fs = std::filesystem; + + struct FallbackGuard { + std::atomic& ctr; + bool disarmed = false; + FallbackGuard(std::atomic& c) : ctr(c) {} + ~FallbackGuard() { + if (!disarmed) ctr.fetch_add(1, std::memory_order_relaxed); + } + void disarm() { disarmed = true; } + } fallback_guard(metadata_load_fallbacks_); + + std::string meta_path = GetMetaFilePath(); + std::string tmp_path = meta_path + ".tmp"; + std::error_code ec_tmp, ec_meta; + + // Clean up stale tmp + if (fs::exists(tmp_path, ec_tmp)) { + if (!fs::exists(meta_path, ec_meta)) { + fs::remove(tmp_path, ec_tmp); + LOG(INFO) << "Orphaned .meta.tmp cleaned up"; + fallback_guard.disarm(); + return RecoveryResult::kNoMeta; + } + fs::remove(tmp_path, ec_tmp); + } + + if (!fs::exists(meta_path, ec_meta)) { + fallback_guard.disarm(); + return RecoveryResult::kNoMeta; + } + + LOG(INFO) << "Recovery path: " << meta_path << " found"; + + // Load metadata. An open failure here (e.g. fd exhaustion, + // permission change) is transient: the meta file exists and may + // be perfectly fine, so it must not trigger a cache-wiping fresh + // start. Parse/validation failures mean the meta is unusable. + auto meta_result = LoadMetadata(); + if (!meta_result) { + if (meta_result.error() == ErrorCode::FILE_OPEN_FAIL) { + LOG(ERROR) << "LoadMetadata could not open " << meta_path + << " -- transient error, keeping persisted data"; + fallback_guard.disarm(); + return RecoveryResult::kTransientError; + } + LOG(ERROR) << "LoadMetadata failed, falling back to fresh start"; + return RecoveryResult::kCorrupt; + } + OffsetAllocatorPersistedMetadata meta = std::move(meta_result.value()); + + // Deserialize allocator into local shared_ptr; only move to + // member on full success + std::shared_ptr alloc; + try { + const auto* data = reinterpret_cast( + meta.allocator_state.data()); + std::vector alloc_buf( + data, data + meta.allocator_state.size()); + alloc = + deserialize_from(alloc_buf); + } catch (const std::bad_alloc& e) { + LOG(ERROR) << "Allocator deserialization OOM: " << e.what() + << " -- transient error, keeping persisted data"; + fallback_guard.disarm(); + return RecoveryResult::kTransientError; + } catch (const std::exception& e) { + LOG(ERROR) << "Allocator deserialization threw: " << e.what() + << " -- fresh start"; + return RecoveryResult::kCorrupt; + } + if (!alloc) { + LOG(ERROR) << "Allocator deserialization returned null" + " -- fresh start"; + return RecoveryResult::kCorrupt; + } + + // Verify capacity match + { + auto metrics = alloc->get_metrics(); + if (static_cast(metrics.capacity) != capacity_) { + LOG(WARNING) << "Allocator capacity mismatch: serialized " + << metrics.capacity << " vs current " << capacity_ + << ". total_size_limit may have changed; " + "falling back to fresh start."; + return RecoveryResult::kCorrupt; + } + } + + // Open data file without truncation + data_file_path_ = GetDataFilePath(); + int flags = O_CLOEXEC | O_RDWR; + int raw_fd = open(data_file_path_.c_str(), flags, 0644); + if (raw_fd < 0) { + const int open_errno = errno; + LOG(ERROR) << "Failed to open data file: " << data_file_path_ + << ": " << strerror(open_errno); + // The meta references data that is genuinely gone: nothing + // recoverable remains, so a fresh start is appropriate. + if (open_errno == ENOENT) { + return RecoveryResult::kCorrupt; + } + // Anything else (fd exhaustion, permissions, I/O error) may + // be transient -- do not wipe the persisted cache for it. + fallback_guard.disarm(); + return RecoveryResult::kTransientError; + } + FdGuard fd_guard(raw_fd); + + // Verify data file size matches capacity_ + struct stat st; + if (fstat(fd_guard.get(), &st) != 0) { + LOG(ERROR) << "fstat failed on data file: " << strerror(errno); + fallback_guard.disarm(); + return RecoveryResult::kTransientError; + } + if (static_cast(st.st_size) != capacity_) { + LOG(WARNING) << "data file size mismatch: expected " << capacity_ + << ", got " << st.st_size + << " -- falling back to fresh start"; + return RecoveryResult::kCorrupt; + } + +#ifdef USE_URING + if (file_storage_config_.use_uring) { + data_file_ = std::make_shared( + data_file_path_, fd_guard.release(), 32, true); + } else +#endif + { + data_file_ = std::make_shared(data_file_path_, + fd_guard.release()); + } + data_file_->SetDeleteOnWriteFail(false); + + // Move allocator to member only after all checks pass + allocator_ = std::move(alloc); + + // Clear and rebuild shard maps + { + std::vector> locks; + locks.reserve(kNumShards); + for (size_t i = 0; i < kNumShards; ++i) { + locks.emplace_back( + std::make_unique(&shards_[i].mutex)); + shards_[i].map.clear(); + } + } + total_size_.store(0, std::memory_order_relaxed); + total_keys_.store(0, std::memory_order_relaxed); + + RebuildShardMapsFromAllocator(meta.insert_seq); + RestoreAndRepairFifoIndex(meta); + + fallback_guard.disarm(); + + LOG(INFO) << "OffsetAllocatorStorageBackend recovered: " + << total_keys_.load() << " keys, " << total_size_.load() + << " bytes"; + return RecoveryResult::kRecovered; + } catch (const std::bad_alloc& e) { + LOG(ERROR) << "Recovery OOM: " << e.what() + << " -- transient error, keeping persisted data"; + return RecoveryResult::kTransientError; + } catch (const std::exception& e) { + LOG(ERROR) << "Recovery failed with exception: " << e.what(); + return RecoveryResult::kCorrupt; + } catch (...) { + LOG(ERROR) << "Recovery failed with unknown exception"; + return RecoveryResult::kCorrupt; + } +} + +void OffsetAllocatorStorageBackend::EvictToMakeRoom( + int64_t required_bytes, size_t min_victims, + const std::unordered_set& batch_keys, + PendingEviction& out_pending) { + if (cfg_.eviction_policy == OffsetEvictionPolicy::NONE) return; + + MutexLocker ev(&eviction_mutex_); + size_t n = 0; + + while (n < cfg_.max_evict_per_offload) { + int64_t cur_size = total_size_.load(std::memory_order_relaxed); + int64_t cur_keys = total_keys_.load(std::memory_order_relaxed); + bool below_bytes = (cur_size + required_bytes <= low_watermark_bytes_); + bool below_keys = (cur_keys <= low_watermark_keys_); + // Stop when both byte and key-count watermarks are satisfied + // and we have met the minimum victim count. + if (below_bytes && below_keys && n >= min_victims) break; + + if (fifo_index_.empty()) break; + + auto oldest = fifo_index_.begin(); + uint64_t vseq = oldest->first; + std::string vkey = oldest->second; + + // Skip batch_keys prefix — keys being written in this batch. + // Do NOT erase their FIFO slots (they may fail allocate() and + // need the slot to remain in the index for future eviction). + // Worst-case comparison cost: O(|batch_keys_prefix|), bounded. + while (oldest != fifo_index_.end() && + batch_keys.count(oldest->second)) { + ++oldest; + } + if (oldest == fifo_index_.end()) break; // all are batch_keys + vkey = oldest->second; + vseq = oldest->first; + + size_t shard_idx = ShardForKey(vkey); + auto& shard = shards_[shard_idx]; + { + SharedMutexLocker lk(&shard.mutex); + auto it = shard.map.find(vkey); + if (it == shard.map.end() || it->second.fifo_seq != vseq) { + // Orphan slot in fifo_index_: the key is no longer in + // shard.map, or its fifo_seq was replaced by a newer + // overwrite. Overwrites are cleaned up in BatchOffload + // Step-4 (fifo_index_.erase(old_seq) under the lock), + // so today this branch is only reachable if a future + // per-key delete path neglects to also erase from + // fifo_index_. Keep the lazy-repair as a defense. + fifo_index_.erase(oldest); + metadata_dirty_.store(true, std::memory_order_relaxed); + ++n; // counted toward scan budget + continue; + } + // Defensive assertions against double-evict underflow. + // Precondition: single heartbeat_thread_ serialises offload; + // if concurrent offload is added, these must be re-evaluated. + DCHECK_GE(total_size_.load(std::memory_order_relaxed), + it->second.total_size); + DCHECK_GE(total_keys_.load(std::memory_order_relaxed), 1); + out_pending.objects.emplace_back(vkey, it->second); + total_size_.fetch_sub(it->second.total_size, + std::memory_order_relaxed); + total_keys_.fetch_sub(1, std::memory_order_relaxed); + shard.map.erase(it); + // NOTE: no eviction tombstone is recorded here. The + // eviction only becomes final when the master accepts the + // notification; if it fails, RestorePreparedEviction rolls + // the entry back and a tombstone would wrongly delete the + // still-valid key on recovery. Tombstones are recorded at + // commit time in NotifyAndCommitPreparedEviction. + metadata_dirty_.store(true, std::memory_order_relaxed); + } + fifo_index_.erase(oldest); + ++n; + } +} + +//----------------------------------------------------------------------------- + +void OffsetAllocatorStorageBackend::RestorePreparedEviction( + PendingEviction&& pending) { + MutexLocker ev(&eviction_mutex_); + + for (auto& [key, entry] : pending.objects) { + const uint32_t restored_size = entry.total_size; + const uint64_t restored_seq = entry.fifo_seq; + auto& shard = shards_[ShardForKey(key)]; + SharedMutexLocker lk(&shard.mutex); + + // BatchOffload's single-writer precondition guarantees that no writer + // can recreate this key or consume its FIFO sequence while the + // eviction notification is in flight. Readers do not mutate either + // index, so both entries must still be absent here. + const bool map_inserted = + shard.map.emplace(key, std::move(entry)).second; + CHECK(map_inserted) << "Failed to restore evicted key: " << key; + const bool fifo_inserted = + fifo_index_.emplace(restored_seq, key).second; + CHECK(fifo_inserted) + << "Failed to restore FIFO entry for evicted key: " << key; + total_size_.fetch_add(restored_size, std::memory_order_relaxed); + total_keys_.fetch_add(1, std::memory_order_relaxed); + } + pending.objects.clear(); +} + +//----------------------------------------------------------------------------- + +// Records restart-persistence tombstones for evictions that have become +// final (master notified, or no master to notify). Must NOT be called +// for evictions that may still be rolled back. +void OffsetAllocatorStorageBackend::RecordEvictionTombstones( + const PendingEviction& pending) { + if (cfg_.persist_mode == OffsetPersistMode::kDisabled) return; + for (const auto& [key, _] : pending.objects) { + all_evicted_this_batch_.insert(key); + } + metadata_dirty_.store(true, std::memory_order_relaxed); +} + +tl::expected +OffsetAllocatorStorageBackend::NotifyAndCommitPreparedEviction( + const EvictionHandler& eviction_handler, PendingEviction& pending) { + if (pending.objects.empty()) return {}; + if (!eviction_handler) { + // No master to notify: dropping the handles commits the eviction + // locally, so the keys still need tombstones to stay dead across + // a restart. + RecordEvictionTombstones(pending); + pending.objects.clear(); + return {}; + } + + std::vector evicted_keys; + evicted_keys.reserve(pending.objects.size()); + for (const auto& [key, _] : pending.objects) { + evicted_keys.push_back(key); + } + + auto notify_result = eviction_handler(evicted_keys); + if (!notify_result) { + const ErrorCode error = notify_result.error(); + RestorePreparedEviction(std::move(pending)); + return tl::make_unexpected(error); + } + + // The master accepted the removal: the eviction is now final, so + // record the tombstones before dropping the handles. + RecordEvictionTombstones(pending); + // Dropping the final metadata references commits the eviction and makes + // unpinned extents available to the allocator. + pending.objects.clear(); + return {}; +} + +//----------------------------------------------------------------------------- + +tl::expected OffsetAllocatorStorageBackend::BatchOffload( + const std::unordered_map>& batch_object, + std::function& keys, + std::vector& metadatas)> + complete_handler, + EvictionHandler eviction_handler) { + // ================================================================ + // SINGLE-WRITER PRECONDITION + // + // BatchOffload, EvictToMakeRoom, and the watermark accounting on + // total_size_ / total_keys_ assume only ONE thread calls + // BatchOffload at a time (currently guaranteed by FileStorage's + // single heartbeat_thread_). The atomics make individual loads + // and stores atomic, but the read-modify-write sequences are NOT + // atomic across threads. If concurrent offload is added, the + // DCHECK_GE guards in EvictToMakeRoom must also be re-evaluated. + // ================================================================ + + if (!initialized_.load(std::memory_order_acquire)) { + LOG(ERROR) + << "Storage backend is not initialized. Call Init() before use."; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + if (batch_object.empty()) { + LOG(ERROR) << "BatchOffload called with empty batch"; + return tl::make_unexpected(ErrorCode::INVALID_KEY); + } + + auto enable_offloading_res = IsEnableOffloading(); + if (!enable_offloading_res) { + return tl::make_unexpected(enable_offloading_res.error()); + } + if (!enable_offloading_res.value()) { + return tl::make_unexpected(ErrorCode::KEYS_ULTRA_LIMIT); + } + + const bool eviction_on = + (cfg_.eviction_policy != OffsetEvictionPolicy::NONE) && + (eviction_handler != nullptr); + + if (cfg_.eviction_policy != OffsetEvictionPolicy::NONE && + eviction_handler == nullptr) { + LOG_FIRST_N(WARNING, 1) + << "Eviction policy is " << static_cast(cfg_.eviction_policy) + << " but eviction_handler is null; eviction is disabled. " + "IsEnableOffloading() will still return true."; + } + + // Build the set of keys being offloaded in this batch so that + // EvictToMakeRoom does not evict them. + std::unordered_set batch_keys; + if (eviction_on) { + for (const auto& [k, _] : batch_object) batch_keys.insert(k); } std::vector keys; @@ -2842,146 +4823,361 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( keys.reserve(batch_object.size()); metadatas.reserve(batch_object.size()); - // Process each object in the batch; continue on individual failures to - // support partial success + // Pin data_file_ for the whole batch so a concurrent RemoveAll() rebuild + // (which rebinds the data_file_ member) cannot close the file while a + // write is in flight (writes do not hold shard locks during I/O). + auto data_file = data_file_; + + // Prepared victims are held here until the master accepts their removal. + // Their allocation handles prevent reuse before notification succeeds. + PendingEviction pending_eviction; + for (const auto& [key, slices] : batch_object) { - if (slices.empty()) { - // Skip empty slices (empty values are allowed but not stored) - continue; - } + if (slices.empty()) continue; - // Test-only: Check if this key should fail (deterministic failure - // injection) if (test_failure_predicate_ && test_failure_predicate_(key)) { LOG(INFO) << "[TEST] Injecting failure for key: " << key << " (test failure predicate)"; - continue; // Simulate allocation/write failure + continue; + } + + if (key.size() > kMaxKeyLen) { + // Recovery rejects key_len above kMaxKeyLen; refuse such keys + // at write time so they cannot silently vanish on restart. + LOG(ERROR) << "Key too large for SSD offload: " << key.size() + << " bytes (max " << kMaxKeyLen << ")"; + continue; } - // Calculate total value size - uint32_t value_size = 0; + uint64_t total_value_size = 0; for (const auto& slice : slices) { - value_size += static_cast(slice.size); + total_value_size += slice.size; + } + if (total_value_size > UINT32_MAX) { + LOG(ERROR) << "Object too large for SSD offload for key: " << key + << ", size: " << total_value_size; + continue; + } + uint32_t value_size = static_cast(total_value_size); + + // Stamp every write with a fresh insert_seq_ value: recovery + // rejects any record stamped at/after the checkpoint's insert_seq + // (its extent may hold a torn write), so every record must carry + // one regardless of the eviction mode. A failed write simply + // burns a seq; gaps in the sequence are harmless. + const uint64_t record_seq = + insert_seq_.fetch_add(1, std::memory_order_relaxed); + + RecordHeader header{ + .key_len = static_cast(key.size()), + .value_len = value_size, + .seq = record_seq, + .flags = cfg_.enable_record_crc ? RecordHeader::kFlagHasCrc : 0u, + .crc32 = 0}; + + // Aligned layout: header + key + zero padding + value. + const uint32_t value_padding = + RecordHeader::ValuePadding(header.key_len); + uint64_t record_size = + RecordHeader::RecordSize(header.key_len, header.value_len); + + if (record_size > UINT32_MAX) { + LOG(ERROR) << "Record too large for key: " << key + << ", record_size=" << record_size; + continue; } - // Prepare record header - RecordHeader header{.key_len = static_cast(key.size()), - .value_len = value_size}; + // Serialize the header; when enabled, the CRC covers the header + // prefix, the key and the value, and lets recovery detect + // torn/stale records. + char hdr_buf[RecordHeader::SIZE]; + header.WritePrefixTo(hdr_buf); + if (header.HasCrc()) { + Crc32c crc; + crc.Extend(hdr_buf, RecordHeader::PREFIX_SIZE); + crc.Extend(key.data(), key.size()); + for (const auto& slice : slices) { + crc.Extend(slice.ptr, slice.size); + } + header.crc32 = crc.Final(); + } + header.WriteTo(hdr_buf); + + // ---- (A) Proactive eviction (watermark-driven) ---- + if (eviction_on) { + int64_t cur_size = total_size_.load(std::memory_order_relaxed); + int64_t cur_keys = total_keys_.load(std::memory_order_relaxed); + bool over_bytes = (cur_size + static_cast(record_size) > + high_watermark_bytes_); + bool over_keys = (cur_keys > high_watermark_keys_); + if (over_bytes || over_keys) { + size_t min_v = over_keys ? cfg_.fallback_evict_batch : 0; + EvictToMakeRoom(static_cast(record_size), min_v, + batch_keys, pending_eviction); + } + } - // Use size_t for record_size to handle large objects (up to 4GB per - // RecordHeader) - size_t record_size = - RecordHeader::SIZE + header.key_len + header.value_len; + // ---- (B) Notify master of evicted keys BEFORE allocating ---- + // Layer-1 (byte safety): BatchLoad pins extents via shared_ptr, + // so the allocator cannot re-issue a still-read offset. Layer-2 + // (master metadata): the master must be told the key's local-disk + // replica is gone before we reuse its space for a new key. + // Persistence tombstones for the victims are recorded inside + // NotifyAndCommitPreparedEviction once the notification succeeds. + if (eviction_on && !pending_eviction.objects.empty()) { + auto notify_result = NotifyAndCommitPreparedEviction( + eviction_handler, pending_eviction); + if (!notify_result) { + return tl::make_unexpected(notify_result.error()); + } + } - // Step 1: Allocate space (allocator is thread-safe, ensures unique - // offsets) No locks held during allocation + // ---- (C) Allocate ---- auto allocation = allocator_->allocate(record_size); + + // ---- (D) Fallback eviction (nullopt retry loop) ---- + if (!allocation.has_value() && eviction_on) { + uint64_t prev_largest = + allocator_->get_metrics().largest_free_region_; + size_t fallback_total_evicted = 0; + const size_t kMaxFallbackEvicted = cfg_.max_evict_per_offload; + + while (!allocation.has_value() && + fallback_total_evicted < kMaxFallbackEvicted) { + EvictToMakeRoom(static_cast(record_size), + cfg_.fallback_evict_batch, batch_keys, + pending_eviction); + size_t evicted_this_turn = pending_eviction.objects.size(); + fallback_total_evicted += evicted_this_turn; + + // Notify master of fallback victims before retrying. + if (!pending_eviction.objects.empty()) { + auto notify_result = NotifyAndCommitPreparedEviction( + eviction_handler, pending_eviction); + if (!notify_result) { + return tl::make_unexpected(notify_result.error()); + } + } + + uint64_t now_largest = + allocator_->get_metrics().largest_free_region_; + if (evicted_this_turn == 0) break; // no victims at all + // Stop if the largest free region did not grow at all. + // Using `prev_largest` (rather than `prev_largest + + // record_size / 2`) allows gradual coalescence when + // many small victims must be evicted for one large + // allocation. The `fallback_total_evicted` cap still + // bounds total eviction per key. + if (now_largest <= prev_largest) break; + prev_largest = now_largest; + allocation = allocator_->allocate(record_size); + } + } + if (!allocation.has_value()) { - LOG(ERROR) << "Failed to allocate " << record_size - << " bytes for key: " << key - << " - stopping processing for this batch"; - break; // Stop processing other keys as space is likely exhausted + if (eviction_on) { + eviction_skips_.fetch_add(1, std::memory_order_relaxed); + LOG(WARNING) << "Skipping key after eviction attempts: " << key; + continue; + } else { + LOG(ERROR) << "Failed to allocate " << record_size + << " bytes for key: " << key + << " - stopping processing"; + break; + } } uint64_t offset = allocation->address(); - // Step 2: Write data to disk (no metadata locks held during I/O) + // ---- (E) Disk write ---- + // Shared zero page backing the padding iov: the value region must + // start at a kValueAlignment boundary within the record, and the + // gap between key and value is zero-filled. + static constexpr std::array + kZeroPadding{}; + std::vector iovs; - iovs.reserve(2 + slices.size()); + iovs.reserve(3 + slices.size()); - // Header - iovs.push_back( - {const_cast(reinterpret_cast(&header.key_len)), - sizeof(header.key_len)}); - iovs.push_back({const_cast( - reinterpret_cast(&header.value_len)), - sizeof(header.value_len)}); + iovs.push_back({hdr_buf, static_cast(RecordHeader::SIZE)}); - // Key iovs.push_back({const_cast(key.data()), static_cast(header.key_len)}); - // Value slices + if (value_padding > 0) { + iovs.push_back({const_cast(kZeroPadding.data()), + static_cast(value_padding)}); + } + for (const auto& slice : slices) { iovs.push_back({slice.ptr, slice.size}); } auto write_result = - data_file_->vector_write(iovs.data(), iovs.size(), offset); + data_file->vector_write(iovs.data(), iovs.size(), offset); if (!write_result) { LOG(ERROR) << "Failed to write record for key: " << key - << ", error: " << write_result.error() - << " - continuing with remaining keys"; - // Allocation handle is still local (not yet stored in the metadata - // map) and will be freed automatically when going out of scope. - continue; // Continue processing other keys + << ", error: " << write_result.error(); + continue; } - - // Handle the case where the data was written partially. - size_t written = write_result.value(); - if (written != record_size) { + if (write_result.value() != record_size) { LOG(ERROR) << "Write size mismatch for key: " << key - << ", expected: " << record_size << ", got: " << written - << " - continuing with remaining keys"; - continue; // Continue processing other keys + << ", expected: " << record_size + << ", got: " << write_result.value(); + continue; } - // Step 3: Wrap allocation in refcounted handle - auto allocation_ptr = std::make_shared( - std::move(allocation.value())); - - // Step 4: Update metadata map under exclusive shard lock - // Lock only the shard for this key (other shards can proceed in - // parallel) + // ---- (F) Metadata update with FIFO index maintenance ---- { + auto allocation_ptr = std::make_shared( + std::move(allocation.value())); size_t shard_idx = ShardForKey(key); auto& shard = shards_[shard_idx]; - SharedMutexLocker lock(&shard.mutex); - // Check if key exists to update size accounting + std::optional ev_lock; + if (eviction_on) ev_lock.emplace(&eviction_mutex_); + SharedMutexLocker shard_lock(&shard.mutex); + auto it = shard.map.find(key); int64_t size_delta = static_cast(record_size); bool is_new_key = (it == shard.map.end()); + // Stamped before the disk write; the record on disk carries + // the same seq so recovery can order writes against the + // checkpoint's insert_seq. + const uint64_t seq = record_seq; if (!is_new_key) { - // Overwrite: subtract old size size_delta -= static_cast(it->second.total_size); - // Old AllocationPtr will be dropped, refcount decremented - // Physical extent freed when last reader releases it + if (eviction_on) fifo_index_.erase(it->second.fifo_seq); } - // Update map (insert_or_assign handles both insert and overwrite) shard.map.insert_or_assign( - key, ObjectEntry(offset, record_size, value_size, - std::move(allocation_ptr))); + key, ObjectEntry(offset, static_cast(record_size), + value_size, std::move(allocation_ptr), seq)); + metadata_dirty_.store(true, std::memory_order_relaxed); + if (eviction_on) fifo_index_.emplace(seq, key); - // Update total size atomically (lock-free, separate from map - // updates) total_size_.fetch_add(size_delta, std::memory_order_relaxed); - - // Update total keys only if inserting a new key if (is_new_key) { total_keys_.fetch_add(1, std::memory_order_relaxed); } } keys.push_back(key); - metadatas.push_back(StorageObjectMetadata{ - 0, // bucket_id not used for this backend - static_cast(offset), static_cast(header.key_len), - static_cast(value_size), ""}); + metadatas.push_back( + StorageObjectMetadata{0, static_cast(offset), + static_cast(header.key_len), + static_cast(value_size), ""}); + } + + // ---- Post-loop flush: notify master of any evicted keys that + // were accumulated by the last (possibly allocate-failing) key. + // This runs BEFORE the persistence barrier so that the trailing + // victims' tombstones and freed extents are included in this + // batch's checkpoint; in kStrict mode that keeps the whole batch + // (writes AND evictions) durable, and avoids serializing victim + // extents as used only for them to be freed right after. + if (eviction_on && !pending_eviction.objects.empty()) { + auto notify_result = + NotifyAndCommitPreparedEviction(eviction_handler, pending_eviction); + if (!notify_result) { + // The eviction was rolled back. The batch is not + // checkpointed here (metadata_dirty_ stays set and a later + // batch will persist it), which is consistent with + // reporting failure to the caller. + return tl::make_unexpected(notify_result.error()); + } + } + + // ---- Persistence barrier (kRelaxed / kStrict) ---- + // Reconstruct set of keys successfully written in this batch + std::unordered_set successfully_written_keys_set(keys.begin(), + keys.end()); + + const bool persist_enabled = + (cfg_.persist_mode == OffsetPersistMode::kRelaxed || + cfg_.persist_mode == OffsetPersistMode::kStrict); + + // Filter stale tombstones on EVERY persist-enabled batch, + // regardless of whether a checkpoint is triggered. If we only + // filter inside ShouldPersistNow(), a key evicted in batch N + // and re-written in batch M (when no checkpoint fires) keeps + // its tombstone in all_evicted_this_batch_. The NEXT checkpoint + // would then write the stale tombstone to meta, causing Phase C + // to delete the valid re-written key on recovery. + if (persist_enabled && !all_evicted_this_batch_.empty() && + !successfully_written_keys_set.empty()) { + size_t removed = 0; + for (const auto& k : successfully_written_keys_set) { + removed += all_evicted_this_batch_.erase(k); + } + if (removed > 0) { + LOG(INFO) << "Cleared " << removed << " stale eviction tombstones"; + } + } + + if (persist_enabled && metadata_dirty_.load(std::memory_order_relaxed) && + ShouldPersistNow()) { + // data fsync (REQUIRED: ensures metadata points to + // durable data) + auto sync_res = data_file_->datasync(); + if (!sync_res) { + auto fails = metadata_consecutive_failures_.fetch_add( + 1, std::memory_order_relaxed) + + 1; + LOG(ERROR) << "datasync failed: " + << static_cast(sync_res.error()); + if (fails % 10 == 0) { + LOG(WARNING) << "Checkpoint datasync " << fails + << " consecutive failures (mode=" + << static_cast(cfg_.persist_mode) << ")"; + } + if (cfg_.persist_mode == OffsetPersistMode::kStrict) { + return tl::make_unexpected(sync_res.error()); + } + // kRelaxed: skip this checkpoint, retry next interval + } else { + auto save_res = SaveMetadata(all_evicted_this_batch_); + if (!save_res) { + auto fails = metadata_consecutive_failures_.fetch_add( + 1, std::memory_order_relaxed) + + 1; + // Log every 10th consecutive failure (LOG_FIRST_N + // counts by call-site, not by content — use LOG(WARNING)). + if (fails % 10 == 0) { + LOG(WARNING) << "Checkpoint " << fails + << " consecutive failures (mode=" + << static_cast(cfg_.persist_mode) + << "): " << static_cast(save_res.error()); + } + if (cfg_.persist_mode == OffsetPersistMode::kStrict) { + return tl::make_unexpected(save_res.error()); + } + // kRelaxed: keep metadata_dirty_=true (DirtyGuard + // restores it), retry next interval + } else { + metadata_consecutive_failures_.store(0, + std::memory_order_relaxed); + all_evicted_this_batch_.clear(); + last_persist_time_us_.store( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(), + std::memory_order_relaxed); + } + } } - // Invoke complete handler only if we have successful keys to report + // ---- Complete handler ---- if (complete_handler != nullptr && !keys.empty()) { auto error_code = complete_handler(keys, metadatas); if (error_code != ErrorCode::OK) { - LOG(ERROR) - << "Complete handler failed: " << error_code << " - " - << keys.size() - << " keys were successfully written to disk but master was not " - "notified. " - << "Master will learn about them via ScanMeta on next restart."; + LOG(ERROR) << "Complete handler failed: " << error_code << " - " + << keys.size() + << " keys were successfully written to disk but " + "master was not notified. " + << "Master will learn about them via ScanMeta on " + "next restart."; return tl::make_unexpected(error_code); } } @@ -2989,8 +5185,6 @@ tl::expected OffsetAllocatorStorageBackend::BatchOffload( return static_cast(keys.size()); } -//----------------------------------------------------------------------------- - tl::expected OffsetAllocatorStorageBackend::BatchLoad( std::unordered_map& batched_slices) { if (!initialized_.load(std::memory_order_acquire)) { @@ -3007,6 +5201,10 @@ tl::expected OffsetAllocatorStorageBackend::BatchLoad( uint32_t value_size; AllocationPtr allocation; // Refcounted handle keeps allocation alive Slice dest_slice; + // Pin the data file so a concurrent RemoveAll() rebuild (which rebinds + // the data_file_ member) cannot destroy this file while we still have + // pending I/O on it in phase 2 (no shard lock held there). + std::shared_ptr data_file; }; std::vector read_plans; @@ -3035,24 +5233,28 @@ tl::expected OffsetAllocatorStorageBackend::BatchLoad( } // Copy metadata and increment refcount on allocation - // This keeps the physical extent alive even if key is evicted + // This keeps the physical extent alive even if key is evicted. + // Also pin data_file_ (shared_ptr copy) so a concurrent RemoveAll + // rebuild cannot close it before our phase-2 I/O completes. read_plans.push_back( ReadPlan{key, entry.offset, entry.value_size, entry.allocation, // shared_ptr copy, increments refcount - dest_slice}); + dest_slice, data_file_}); - // Lock released here; allocation stays alive via shared_ptr + // Lock released here; allocation + data_file stay alive via shared_ptr } // Step 2: Perform disk I/O without holding any locks - // Allocations are kept alive by shared_ptr references in read_plans + // Allocations and data file are kept alive by shared_ptr refs in read_plans for (const auto& plan : read_plans) { - // Read header first - RecordHeader header; - iovec header_iovs[2] = {{&header.key_len, sizeof(header.key_len)}, - {&header.value_len, sizeof(header.value_len)}}; + // Read header first. The CRC is NOT verified here: records are + // CRC-validated once during recovery, and during normal operation + // a key only becomes visible after its write completed, so the + // hot read path stays checksum-free. + char hdr_buf[RecordHeader::SIZE]; + iovec header_iov = {hdr_buf, sizeof(hdr_buf)}; auto read_header_result = - data_file_->vector_read(header_iovs, 2, plan.offset); + plan.data_file->vector_read(&header_iov, 1, plan.offset); if (!read_header_result) { LOG(ERROR) << "Failed to read header for key: " << plan.key << ", error: " << read_header_result.error(); @@ -3063,6 +5265,7 @@ tl::expected OffsetAllocatorStorageBackend::BatchLoad( LOG(ERROR) << "Header read size mismatch for key: " << plan.key; return tl::make_unexpected(ErrorCode::FILE_READ_FAIL); } + RecordHeader header = RecordHeader::ReadFrom(hdr_buf); // Validate header matches metadata if (!header.ValidateAgainstMetadata(plan.value_size)) { @@ -3075,7 +5278,7 @@ tl::expected OffsetAllocatorStorageBackend::BatchLoad( // Read key from disk std::string stored_key(header.key_len, '\0'); iovec key_iov = {stored_key.data(), header.key_len}; - auto read_key_result = data_file_->vector_read( + auto read_key_result = plan.data_file->vector_read( &key_iov, 1, plan.offset + RecordHeader::SIZE); if (!read_key_result) { LOG(ERROR) << "Failed to read key for: " << plan.key @@ -3094,10 +5297,11 @@ tl::expected OffsetAllocatorStorageBackend::BatchLoad( return tl::make_unexpected(validate_result.error()); } - // Read value into destination slice + // Read value into destination slice (at its aligned offset) iovec value_iov = {plan.dest_slice.ptr, plan.dest_slice.size}; - auto read_value_result = data_file_->vector_read( - &value_iov, 1, plan.offset + RecordHeader::SIZE + header.key_len); + auto read_value_result = plan.data_file->vector_read( + &value_iov, 1, + plan.offset + RecordHeader::ValueOffsetInRecord(header.key_len)); if (!read_value_result) { LOG(ERROR) << "Failed to read value for key: " << plan.key << ", error: " << read_value_result.error(); @@ -3144,15 +5348,17 @@ OffsetAllocatorStorageBackend::IsEnableOffloading() { return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } - // TODO: See if free space check is needed here. - // Check quota limits only (atomic counters, completely lock-free!) + // When eviction is enabled, BatchOffload's EvictToMakeRoom is + // responsible for making room — do not block offload here. + if (cfg_.eviction_policy != OffsetEvictionPolicy::NONE) { + return true; + } + + // Eviction disabled: keep the original quota-check behavior. bool within_size_limit = total_size_.load(std::memory_order_relaxed) < file_storage_config_.total_size_limit; - - // Check keys limit (atomic counter maintained during BatchOffload) bool within_keys_limit = total_keys_.load(std::memory_order_relaxed) < file_storage_config_.total_keys_limit; - return within_size_limit && within_keys_limit; } @@ -3216,8 +5422,10 @@ tl::expected OffsetAllocatorStorageBackend::ScanMeta( metadatas.push_back(StorageObjectMetadata{ 0, // bucket_id not used static_cast(entry.offset), - static_cast(entry.total_size - RecordHeader::SIZE - - entry.value_size), // key_size + // key_size comes from the key itself: total_size also + // contains header and alignment padding, so it cannot + // be derived arithmetically. + static_cast(key.size()), static_cast(entry.value_size), ""}); // Call handler when batch limit is reached @@ -3241,6 +5449,53 @@ tl::expected OffsetAllocatorStorageBackend::ScanMeta( return {}; } +void OffsetAllocatorStorageBackend::RemoveAll() { + // Acquire exclusive locks on all shards to safely clear the maps and + // reset counters. The single-arg SharedMutexLocker constructor takes + // exclusive mode (see mutex.h), so map.clear() is safe. The shard locks + // also serialize RemoveAll vs RemoveAll (initiator sync + heartbeat + // backstop) and vs metadata updates; concurrent readers/writers keep the + // old data_file_ alive via shared_ptr pinning (see BatchLoad/BatchStore). + std::vector> shard_locks; + shard_locks.reserve(kNumShards); + for (size_t i = 0; i < kNumShards; ++i) { + shard_locks.emplace_back( + std::make_unique(&shards_[i].mutex)); + shards_[i].map.clear(); + } + total_size_.store(0, std::memory_order_relaxed); + total_keys_.store(0, std::memory_order_relaxed); + + // Truncate the data file and rebuild the allocator so the backend + // is ready for new writes (same logic as Init's fresh-start path). + // Rebinding the shared_ptr drops our ref; any in-flight + // BatchLoad/BatchStore that pinned the old data_file_ keeps it alive until + // its I/O completes — no use-after-free. + if (!data_file_path_.empty()) { + int fd = open(data_file_path_.c_str(), + O_CLOEXEC | O_RDWR | O_CREAT | O_TRUNC, 0644); + if (fd >= 0) { +#ifdef USE_URING + if (file_storage_config_.use_uring) { + data_file_ = + std::make_shared(data_file_path_, fd, 32, true); + } else +#endif + { + data_file_ = std::make_shared(data_file_path_, fd); + } + } else { + LOG(WARNING) << "RemoveAll: failed to truncate data file: " + << data_file_path_; + } + } + if (capacity_ > 0) { + allocator_ = offset_allocator::OffsetAllocator::create(0, capacity_); + } + + LOG(INFO) << "OffsetAllocatorStorageBackend::RemoveAll: cleared all data"; +} + //----------------------------------------------------------------------------- tl::expected, ErrorCode> @@ -3266,7 +5521,14 @@ CreateStorageBackend(const FileStorageConfig& config) { config, file_per_key_backend_config); } case StorageBackendType::kOffsetAllocator: { - return std::make_shared(config); + auto offset_backend_config = + OffsetAllocatorBackendConfig::FromEnvironment(); + if (!offset_backend_config.Validate()) { + throw std::invalid_argument( + "Invalid OffsetAllocatorBackendConfig"); + } + return std::make_shared( + config, offset_backend_config); } case StorageBackendType::kDistributed: { auto distributed_config = diff --git a/mooncake-store/src/store_c.map b/mooncake-store/src/store_c.map new file mode 100644 index 0000000000..3c68e1ed94 --- /dev/null +++ b/mooncake-store/src/store_c.map @@ -0,0 +1,9 @@ +# Linker version script for libmooncake_store.so (WITH_STORE_C_SHARED). +# Export only the store_c.h C ABI; localize every symbol pulled in from the +# statically-archived C++ dependency graph (asio, glog, transfer_engine, ...). +{ + global: + mooncake_store_*; + local: + *; +}; diff --git a/mooncake-store/src/store_c_shared.cpp b/mooncake-store/src/store_c_shared.cpp new file mode 100644 index 0000000000..8fdd861d0e --- /dev/null +++ b/mooncake-store/src/store_c_shared.cpp @@ -0,0 +1,20 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Intentionally empty translation unit. +// +// The `mooncake_store_*` C ABI is provided by whole-archiving the static +// `mooncake_store` library into libmooncake_store.so (see the +// WITH_STORE_C_SHARED block in CMakeLists.txt). CMake requires a shared-library +// target to have at least one source file; this file satisfies that. diff --git a/mooncake-store/src/tenant_quota.cpp b/mooncake-store/src/tenant_quota.cpp index e68d739f75..fbd7372927 100644 --- a/mooncake-store/src/tenant_quota.cpp +++ b/mooncake-store/src/tenant_quota.cpp @@ -3,15 +3,11 @@ #include #include -#include "types.h" - -#include - namespace mooncake { namespace { struct RemainderShare { - std::string tenant_id; + TenantId tenant_id; uint64_t base = 0; unsigned __int128 remainder = 0; }; @@ -23,145 +19,172 @@ uint64_t SaturatingAdd(uint64_t lhs, uint64_t rhs) { return lhs + rhs; } -bool IsLazyEmptyTenant(const TenantQuotaState& state) { - return !state.has_explicit_policy && state.used_bytes == 0 && - state.reserved_bytes == 0 && state.committed_count == 0; -} +std::map BuildEffectiveQuotaAssignmentsImpl( + const std::vector& tenants, + uint64_t allocatable_capacity_bytes) { + unsigned __int128 explicit_requested_sum = 0; + std::vector> explicit_tenants; + std::map assigned; + + for (const auto& snapshot : tenants) { + const TenantId& tenant_id = snapshot.tenant_id; + assigned.emplace(tenant_id, 0); + if (snapshot.has_explicit_policy) { + explicit_requested_sum += snapshot.requested_quota_bytes; + explicit_tenants.emplace_back(tenant_id, + snapshot.requested_quota_bytes); + } + } -TenantQuotaResult AccountingMismatch(const char* operation, - const std::string& tenant_id, - uint64_t requested, uint64_t available) { - LOG(WARNING) << operation << " accounting mismatch for tenant " << tenant_id - << ": requested=" << requested << ", available=" << available; - return tl::make_unexpected(TenantQuotaError::kAccountingMismatch); -} + if (explicit_requested_sum <= allocatable_capacity_bytes) { + for (const auto& [tenant_id, requested_quota_bytes] : + explicit_tenants) { + assigned[tenant_id] = requested_quota_bytes; + } + return assigned; + } -} // namespace + std::vector shares; + shares.reserve(explicit_tenants.size()); + uint64_t base_assigned = 0; + for (const auto& [tenant_id, requested_quota_bytes] : explicit_tenants) { + const unsigned __int128 product = + static_cast(allocatable_capacity_bytes) * + requested_quota_bytes; + const uint64_t base = + static_cast(product / explicit_requested_sum); + shares.push_back({.tenant_id = tenant_id, + .base = base, + .remainder = product % explicit_requested_sum}); + base_assigned += base; + } -void TenantQuotaTable::SetDefaultRequestedQuota(uint64_t bytes) { - default_requested_quota_bytes_ = bytes; - for (auto& [_, state] : tenants_) { - if (!state.has_explicit_policy) { - state.requested_quota_bytes = default_requested_quota_bytes_; + std::sort(shares.begin(), shares.end(), + [](const RemainderShare& lhs, const RemainderShare& rhs) { + if (lhs.remainder != rhs.remainder) { + return lhs.remainder > rhs.remainder; + } + return lhs.tenant_id < rhs.tenant_id; + }); + + uint64_t remaining = allocatable_capacity_bytes - base_assigned; + for (auto& share : shares) { + if (remaining > 0) { + ++share.base; + --remaining; } + assigned[share.tenant_id] = share.base; } + return assigned; } -uint64_t TenantQuotaTable::GetDefaultRequestedQuota() const { - return default_requested_quota_bytes_; +TenantQuotaResult AccountingMismatch() { + return tl::make_unexpected(TenantQuotaError::kAccountingMismatch); +} + +} // namespace + +std::map TenantQuotaTable::BuildEffectiveQuotaAssignments( + const std::vector& tenants, + uint64_t allocatable_capacity_bytes) { + return BuildEffectiveQuotaAssignmentsImpl(tenants, + allocatable_capacity_bytes); } TenantQuotaResult TenantQuotaTable::UpsertTenantPolicy( - std::string tenant_id, uint64_t requested_quota_bytes) { + const TenantId& tenant_id, uint64_t requested_quota_bytes) { if (requested_quota_bytes == 0) { return tl::make_unexpected(TenantQuotaError::kInvalidArgument); } - auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id)); - auto& state = GetOrCreateState(normalized_tenant_id); + auto& state = GetOrCreateState(tenant_id); state.requested_quota_bytes = requested_quota_bytes; state.has_explicit_policy = true; + RefreshOverQuota(&state); return {}; } -void TenantQuotaTable::EraseTenantPolicy(std::string tenant_id) { - auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id)); - auto it = tenants_.find(normalized_tenant_id); - if (it == tenants_.end()) { - return; +TenantQuotaPolicyResult TenantQuotaTable::DisableTenantPolicyIfEmpty( + const TenantId& tenant_id) { + auto it = tenants_.find(tenant_id); + if (it == tenants_.end() || !it->second.has_explicit_policy) { + return tl::make_unexpected(TenantQuotaError::kTenantNotFound); } + auto& state = it->second; - state.requested_quota_bytes = default_requested_quota_bytes_; + if (state.used_bytes != 0 || state.reserved_bytes != 0 || + state.committed_count != 0 || state.metadata_object_count != 0) { + return tl::make_unexpected(TenantQuotaError::kTenantNotEmpty); + } + + const uint64_t requested_quota_bytes = state.requested_quota_bytes; + state.requested_quota_bytes = 0; + state.effective_quota_bytes = 0; state.has_explicit_policy = false; + RefreshOverQuota(&state); + EraseIfLazyEmpty(it); + return requested_quota_bytes; } -void TenantQuotaTable::RecomputeEffectiveQuotas( - uint64_t allocatable_capacity_bytes) { - unsigned __int128 explicit_requested_sum = 0; - std::vector explicit_tenants; - std::vector default_tenants; +void TenantQuotaTable::ApplyTenantPolicies( + const TenantQuotaPolicyMap& policies) { + for (auto it = tenants_.begin(); it != tenants_.end();) { + auto policy_it = policies.find(it->first); + auto& state = it->second; + if (policy_it != policies.end()) { + state.requested_quota_bytes = policy_it->second; + state.has_explicit_policy = true; + RefreshOverQuota(&state); + ++it; + continue; + } - for (auto& [tenant_id, state] : tenants_) { - if (state.has_explicit_policy) { - explicit_tenants.push_back(tenant_id); - explicit_requested_sum += state.requested_quota_bytes; + state.requested_quota_bytes = 0; + state.effective_quota_bytes = 0; + state.has_explicit_policy = false; + if (IsLazyEmptyTenant(state)) { + it = tenants_.erase(it); } else { - state.requested_quota_bytes = default_requested_quota_bytes_; - if (!IsLazyEmptyTenant(state)) { - default_tenants.push_back(tenant_id); - } + RefreshOverQuota(&state); + ++it; } - state.effective_quota_bytes = 0; } - auto distribute = [&](const std::vector& tenant_ids, - uint64_t capacity, bool proportional_to_requested) { - if (tenant_ids.empty() || capacity == 0) { - return; - } - - std::vector shares; - shares.reserve(tenant_ids.size()); - uint64_t assigned = 0; - unsigned __int128 denominator = proportional_to_requested - ? explicit_requested_sum - : tenant_ids.size(); - - for (const auto& tenant_id : tenant_ids) { - auto& state = tenants_.at(tenant_id); - unsigned __int128 numerator = - proportional_to_requested ? state.requested_quota_bytes : 1; - unsigned __int128 product = - static_cast(capacity) * numerator; - uint64_t base = static_cast(product / denominator); - unsigned __int128 remainder = product % denominator; - shares.push_back({tenant_id, base, remainder}); - assigned += base; - } - - std::sort(shares.begin(), shares.end(), - [](const RemainderShare& lhs, const RemainderShare& rhs) { - if (lhs.remainder != rhs.remainder) { - return lhs.remainder > rhs.remainder; - } - return lhs.tenant_id < rhs.tenant_id; - }); - - uint64_t remaining = capacity - assigned; - for (auto& share : shares) { - if (remaining > 0) { - ++share.base; - --remaining; - } - tenants_.at(share.tenant_id).effective_quota_bytes = share.base; + for (const auto& [tenant_id, requested_quota_bytes] : policies) { + auto [it, inserted] = tenants_.try_emplace(tenant_id); + if (!inserted) { + continue; } - }; + it->second.requested_quota_bytes = requested_quota_bytes; + it->second.has_explicit_policy = true; + RefreshOverQuota(&it->second); + } +} - if (explicit_requested_sum <= allocatable_capacity_bytes) { - for (const auto& tenant_id : explicit_tenants) { - auto& state = tenants_.at(tenant_id); - state.effective_quota_bytes = state.requested_quota_bytes; +TenantQuotaPolicyMap TenantQuotaTable::GetTenantPolicies() const { + TenantQuotaPolicyMap policies; + for (const auto& [tenant_id, state] : tenants_) { + if (state.has_explicit_policy) { + policies.emplace(tenant_id, state.requested_quota_bytes); } - - const uint64_t remaining_capacity = - allocatable_capacity_bytes - - static_cast(explicit_requested_sum); - distribute(default_tenants, remaining_capacity, - /*proportional_to_requested=*/false); - } else { - distribute(explicit_tenants, allocatable_capacity_bytes, - /*proportional_to_requested=*/true); } + return policies; +} - for (auto& [_, state] : tenants_) { - RefreshOverQuota(&state); - } +void TenantQuotaTable::RecomputeEffectiveQuotas( + uint64_t allocatable_capacity_bytes) { + ApplyEffectiveQuotas(BuildEffectiveQuotaAssignments( + ListTenantSnapshots(), allocatable_capacity_bytes)); +} + +bool TenantQuotaTable::IsTenantRegistered(const TenantId& tenant_id) const { + auto it = tenants_.find(tenant_id); + return it != tenants_.end() && it->second.has_explicit_policy; } std::optional TenantQuotaTable::GetTenantSnapshot( - std::string tenant_id) const { - auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id)); - auto it = tenants_.find(normalized_tenant_id); + const TenantId& tenant_id) const { + auto it = tenants_.find(tenant_id); if (it == tenants_.end()) { return std::nullopt; } @@ -170,29 +193,47 @@ std::optional TenantQuotaTable::GetTenantSnapshot( std::vector TenantQuotaTable::ListTenantSnapshots() const { std::vector snapshots; + snapshots.reserve(tenants_.size()); for (const auto& [tenant_id, state] : tenants_) { - if (IsLazyEmptyTenant(state)) { - continue; + if (!IsLazyEmptyTenant(state)) { + snapshots.push_back(MakeSnapshot(tenant_id, state)); } - snapshots.push_back(MakeSnapshot(tenant_id, state)); } return snapshots; } -TenantQuotaResult TenantQuotaTable::Reserve(std::string tenant_id, +uint64_t TenantQuotaTable::ComputeDeficit(const TenantId& tenant_id, + uint64_t incoming_bytes) const { + auto it = tenants_.find(tenant_id); + if (it == tenants_.end()) { + return incoming_bytes; + } + + const auto& state = it->second; + const unsigned __int128 demand = + static_cast(state.used_bytes) + + state.reserved_bytes + incoming_bytes; + if (demand <= state.effective_quota_bytes) { + return 0; + } + + const unsigned __int128 deficit = demand - state.effective_quota_bytes; + return deficit > std::numeric_limits::max() + ? std::numeric_limits::max() + : static_cast(deficit); +} + +TenantQuotaResult TenantQuotaTable::Reserve(const TenantId& tenant_id, uint64_t bytes) { - auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id)); + auto it = tenants_.find(tenant_id); + if (it == tenants_.end() || !it->second.has_explicit_policy) { + return tl::make_unexpected(TenantQuotaError::kTenantNotRegistered); + } if (bytes == 0) { - GetOrCreateState(normalized_tenant_id); return {}; } - auto it = tenants_.find(normalized_tenant_id); - if (it == tenants_.end()) { - return tl::make_unexpected(TenantQuotaError::kQuotaExceeded); - } auto& state = it->second; - if (static_cast(state.used_bytes) + state.reserved_bytes + bytes > state.effective_quota_bytes) { @@ -204,98 +245,152 @@ TenantQuotaResult TenantQuotaTable::Reserve(std::string tenant_id, return {}; } -TenantQuotaResult TenantQuotaTable::Commit(std::string tenant_id, +TenantQuotaResult TenantQuotaTable::Commit(const TenantId& tenant_id, uint64_t bytes) { - auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id)); - auto& state = GetOrCreateState(normalized_tenant_id); + auto it = tenants_.find(tenant_id); + if (it == tenants_.end() || it->second.reserved_bytes < bytes) { + return AccountingMismatch(); + } if (bytes == 0) { return {}; } - if (state.reserved_bytes < bytes) { - return AccountingMismatch("Commit", normalized_tenant_id, bytes, - state.reserved_bytes); - } - + auto& state = it->second; state.reserved_bytes -= bytes; state.used_bytes = SaturatingAdd(state.used_bytes, bytes); - ++state.committed_count; + if (state.committed_count < std::numeric_limits::max()) { + ++state.committed_count; + } RefreshOverQuota(&state); return {}; } -TenantQuotaResult TenantQuotaTable::Abort(std::string tenant_id, - uint64_t bytes) { - auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id)); - auto& state = GetOrCreateState(normalized_tenant_id); +TenantQuotaResult TenantQuotaTable::CommitAdditional(const TenantId& tenant_id, + uint64_t bytes) { + auto it = tenants_.find(tenant_id); + if (it == tenants_.end() || it->second.reserved_bytes < bytes) { + return AccountingMismatch(); + } if (bytes == 0) { return {}; } - if (state.reserved_bytes < bytes) { - return AccountingMismatch("Abort", normalized_tenant_id, bytes, - state.reserved_bytes); - } - + auto& state = it->second; state.reserved_bytes -= bytes; + state.used_bytes = SaturatingAdd(state.used_bytes, bytes); RefreshOverQuota(&state); return {}; } -TenantQuotaResult TenantQuotaTable::Release(std::string tenant_id, - uint64_t bytes) { - auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id)); - auto& state = GetOrCreateState(normalized_tenant_id); +TenantQuotaResult TenantQuotaTable::Abort(const TenantId& tenant_id, + uint64_t bytes) { + auto it = tenants_.find(tenant_id); + if (it == tenants_.end() || it->second.reserved_bytes < bytes) { + return AccountingMismatch(); + } if (bytes == 0) { return {}; } - if (state.used_bytes < bytes) { - return AccountingMismatch("Release", normalized_tenant_id, bytes, - state.used_bytes); + it->second.reserved_bytes -= bytes; + RefreshOverQuota(&it->second); + EraseIfLazyEmpty(it); + return {}; +} + +TenantQuotaResult TenantQuotaTable::Release(const TenantId& tenant_id, + uint64_t bytes) { + auto it = tenants_.find(tenant_id); + if (it == tenants_.end() || it->second.used_bytes < bytes || + (bytes != 0 && it->second.committed_count == 0)) { + return AccountingMismatch(); + } + if (bytes == 0) { + return {}; } + auto& state = it->second; state.used_bytes -= bytes; - if (state.committed_count > 0) { - --state.committed_count; - } else { - LOG(WARNING) << "Release found zero committed_count for tenant " - << normalized_tenant_id; - } + --state.committed_count; RefreshOverQuota(&state); + EraseIfLazyEmpty(it); return {}; } -TenantQuotaResult TenantQuotaTable::ReleasePartial(std::string tenant_id, +TenantQuotaResult TenantQuotaTable::ReleasePartial(const TenantId& tenant_id, uint64_t bytes) { - auto normalized_tenant_id = NormalizeTenantId(std::move(tenant_id)); - auto& state = GetOrCreateState(normalized_tenant_id); + auto it = tenants_.find(tenant_id); + if (it == tenants_.end() || it->second.used_bytes < bytes) { + return AccountingMismatch(); + } if (bytes == 0) { return {}; } - DCHECK_LE(bytes, state.used_bytes); - if (state.used_bytes < bytes) { - return AccountingMismatch("ReleasePartial", normalized_tenant_id, bytes, - state.used_bytes); - } + it->second.used_bytes -= bytes; + RefreshOverQuota(&it->second); + EraseIfLazyEmpty(it); + return {}; +} - state.used_bytes -= bytes; +void TenantQuotaTable::IncrementMetadataObjectCount(const TenantId& tenant_id) { + auto& state = GetOrCreateState(tenant_id); + if (state.metadata_object_count < std::numeric_limits::max()) { + ++state.metadata_object_count; + } RefreshOverQuota(&state); +} + +TenantQuotaResult TenantQuotaTable::DecrementMetadataObjectCount( + const TenantId& tenant_id) { + auto it = tenants_.find(tenant_id); + if (it == tenants_.end() || it->second.metadata_object_count == 0) { + return AccountingMismatch(); + } + + --it->second.metadata_object_count; + RefreshOverQuota(&it->second); + EraseIfLazyEmpty(it); return {}; } -TenantQuotaState& TenantQuotaTable::GetOrCreateState( - const std::string& tenant_id) { - auto [it, inserted] = tenants_.try_emplace(tenant_id); - if (inserted) { - it->second.requested_quota_bytes = default_requested_quota_bytes_; +void TenantQuotaTable::RebuildUsage(const TenantQuotaUsageMap& usage) { + for (auto& [_, state] : tenants_) { + state.used_bytes = 0; + state.reserved_bytes = 0; + state.committed_count = 0; + state.metadata_object_count = 0; + } + + for (const auto& [tenant_id, tenant_usage] : usage) { + auto& state = GetOrCreateState(tenant_id); + if (!state.has_explicit_policy) { + state.requested_quota_bytes = 0; + state.effective_quota_bytes = 0; + } + state.used_bytes = tenant_usage.used_bytes; + state.committed_count = tenant_usage.committed_count; + state.metadata_object_count = tenant_usage.metadata_object_count; + RefreshOverQuota(&state); + } + + for (auto it = tenants_.begin(); it != tenants_.end();) { + if (IsLazyEmptyTenant(it->second)) { + it = tenants_.erase(it); + } else { + RefreshOverQuota(&it->second); + ++it; + } } - return it->second; +} + +TenantQuotaTable::TenantQuotaState& TenantQuotaTable::GetOrCreateState( + const TenantId& tenant_id) { + return tenants_.try_emplace(tenant_id).first->second; } TenantQuotaSnapshot TenantQuotaTable::MakeSnapshot( - const std::string& tenant_id, const TenantQuotaState& state) const { + const TenantId& tenant_id, const TenantQuotaState& state) const { return TenantQuotaSnapshot{ .tenant_id = tenant_id, .requested_quota_bytes = state.requested_quota_bytes, @@ -303,15 +398,40 @@ TenantQuotaSnapshot TenantQuotaTable::MakeSnapshot( .used_bytes = state.used_bytes, .reserved_bytes = state.reserved_bytes, .committed_count = state.committed_count, + .metadata_object_count = state.metadata_object_count, .has_explicit_policy = state.has_explicit_policy, .over_quota = state.over_quota, }; } -void TenantQuotaTable::RefreshOverQuota(TenantQuotaState* state) const { - state->over_quota = static_cast(state->used_bytes) + - state->reserved_bytes > - state->effective_quota_bytes; +bool TenantQuotaTable::IsLazyEmptyTenant(const TenantQuotaState& state) { + return !state.has_explicit_policy && state.used_bytes == 0 && + state.reserved_bytes == 0 && state.committed_count == 0 && + state.metadata_object_count == 0; +} + +void TenantQuotaTable::RefreshOverQuota(TenantQuotaState* state) { + state->over_quota = + (!state->has_explicit_policy && state->metadata_object_count > 0) || + static_cast(state->used_bytes) + + state->reserved_bytes > + state->effective_quota_bytes; +} + +void TenantQuotaTable::ApplyEffectiveQuotas( + const std::map& effective_quotas) { + for (auto& [tenant_id, state] : tenants_) { + auto it = effective_quotas.find(tenant_id); + state.effective_quota_bytes = + it == effective_quotas.end() ? 0 : it->second; + RefreshOverQuota(&state); + } +} + +void TenantQuotaTable::EraseIfLazyEmpty(StateMap::iterator it) { + if (IsLazyEmptyTenant(it->second)) { + tenants_.erase(it); + } } } // namespace mooncake diff --git a/mooncake-store/src/tenant_quota_policy_store.cpp b/mooncake-store/src/tenant_quota_policy_store.cpp new file mode 100644 index 0000000000..1b3db9147a --- /dev/null +++ b/mooncake-store/src/tenant_quota_policy_store.cpp @@ -0,0 +1,467 @@ +#include "tenant_quota_policy_store.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include + +#ifdef STORE_USE_ETCD +#include "etcd_helper.h" +#endif +#include "types.h" + +namespace mooncake { +namespace { + +bool IsValidTenantQuotaName(const std::string& name) { + return !name.empty() && TenantId(name).IsValid(); +} + +std::string ErrnoMessage(const std::string& action, const std::string& path) { + return action + " '" + path + "' failed: " + std::strerror(errno); +} + +#ifdef STORE_USE_ETCD +tl::expected NormalizeClusterIdForEtcdKey( + const std::string& cluster_id) { + std::string normalized = + cluster_id.empty() ? std::string(DEFAULT_CLUSTER_ID) : cluster_id; + while (!normalized.empty() && normalized.back() == '/') { + normalized.pop_back(); + } + if (normalized.empty()) { + normalized = DEFAULT_CLUSTER_ID; + } + if (!IsValidClusterIdComponent(normalized)) { + return tl::make_unexpected("invalid tenant quota etcd cluster_id '" + + cluster_id + "'"); + } + return normalized; +} + +tl::expected BuildTenantQuotaEtcdKey( + const std::string& cluster_id) { + auto normalized = NormalizeClusterIdForEtcdKey(cluster_id); + if (!normalized) { + return tl::make_unexpected(normalized.error()); + } + return "mooncake-store/" + normalized.value() + "/tenant_quota_policy"; +} +#endif + +tl::expected WriteAll(int fd, const std::string& content, + const std::string& path) { + const char* data = content.data(); + size_t remaining = content.size(); + while (remaining > 0) { + ssize_t written = ::write(fd, data, remaining); + if (written < 0) { + if (errno == EINTR) { + continue; + } + return tl::make_unexpected(ErrnoMessage("write", path)); + } + if (written == 0) { + return tl::make_unexpected("write '" + path + "' made no progress"); + } + data += written; + remaining -= static_cast(written); + } + return {}; +} + +tl::expected FsyncDirectory(const std::string& path) { + std::filesystem::path dir = std::filesystem::path(path).parent_path(); + if (dir.empty()) { + dir = "."; + } + int fd = ::open(dir.string().c_str(), O_RDONLY | O_DIRECTORY); + if (fd < 0) { + return tl::make_unexpected( + ErrnoMessage("open directory", dir.string())); + } + auto close_fd = [&] { ::close(fd); }; + if (::fsync(fd) != 0) { + std::string error = ErrnoMessage("fsync directory", dir.string()); + close_fd(); + return tl::make_unexpected(error); + } + close_fd(); + return {}; +} + +std::string MakeTempPath(const std::string& path) { + const auto now = + std::chrono::steady_clock::now().time_since_epoch().count(); + std::ostringstream oss; + oss << path << ".tmp." << ::getpid() << "." + << std::hash{}(std::this_thread::get_id()) << "." + << now; + return oss.str(); +} + +std::string QuoteYamlDoubleQuotedScalar(const std::string& value) { + std::ostringstream out; + out << '"'; + for (unsigned char c : value) { + switch (c) { + case '\\': + out << "\\\\"; + break; + case '"': + out << "\\\""; + break; + case '\0': + out << "\\0"; + break; + case '\a': + out << "\\a"; + break; + case '\b': + out << "\\b"; + break; + case '\t': + out << "\\t"; + break; + case '\n': + out << "\\n"; + break; + case '\v': + out << "\\v"; + break; + case '\f': + out << "\\f"; + break; + case '\r': + out << "\\r"; + break; + default: + if (c < 0x20 || c == 0x7f) { + out << "\\x" << std::uppercase << std::hex << std::setw(2) + << std::setfill('0') << static_cast(c) << std::dec + << std::nouppercase << std::setfill(' '); + } else { + out << static_cast(c); + } + break; + } + } + out << '"'; + return out.str(); +} + +} // namespace + +tl::expected ParseTenantQuotaBytes( + const std::string& value) { + if (value.empty()) { + return tl::make_unexpected("quota must not be empty"); + } + + size_t digits = 0; + while (digits < value.size() && value[digits] >= '0' && + value[digits] <= '9') { + ++digits; + } + if (digits == 0) { + return tl::make_unexpected("quota must start with an integer"); + } + + uint64_t number = 0; + for (size_t i = 0; i < digits; ++i) { + const uint64_t digit = static_cast(value[i] - '0'); + if (number > (std::numeric_limits::max() - digit) / 10) { + return tl::make_unexpected("quota integer overflows uint64"); + } + number = number * 10 + digit; + } + if (number == 0) { + return tl::make_unexpected("quota must be positive"); + } + + const std::string unit = value.substr(digits); + uint64_t multiplier = 1; + if (unit.empty() || unit == "B") { + multiplier = 1; + } else if (unit == "KB") { + multiplier = 1024ULL; + } else if (unit == "MB") { + multiplier = 1024ULL * 1024ULL; + } else if (unit == "GB") { + multiplier = 1024ULL * 1024ULL * 1024ULL; + } else if (unit == "TB") { + multiplier = 1024ULL * 1024ULL * 1024ULL * 1024ULL; + } else { + return tl::make_unexpected("unsupported quota unit '" + unit + "'"); + } + + if (number > std::numeric_limits::max() / multiplier) { + return tl::make_unexpected("quota byte value overflows uint64"); + } + return number * multiplier; +} + +tl::expected ParseTenantQuotaPolicyYaml( + const std::string& yaml) { + YAML::Node root; + try { + root = YAML::Load(yaml); + } catch (const YAML::Exception& e) { + return tl::make_unexpected(std::string("invalid YAML: ") + e.what()); + } + + if (!root || !root.IsMap()) { + return tl::make_unexpected("tenant quota policy must be a YAML map"); + } + const auto version_node = root["version"]; + if (!version_node || !version_node.IsScalar()) { + return tl::make_unexpected("tenant quota policy version is required"); + } + int version = 0; + try { + version = version_node.as(); + } catch (const YAML::Exception& e) { + return tl::make_unexpected(std::string("invalid version: ") + e.what()); + } + if (version != 1) { + return tl::make_unexpected("unsupported tenant quota policy version: " + + std::to_string(version)); + } + + const auto tenants_node = root["tenants"]; + if (!tenants_node || !tenants_node.IsSequence()) { + return tl::make_unexpected("tenants must be a YAML sequence"); + } + + TenantQuotaPolicySnapshot snapshot; + for (size_t i = 0; i < tenants_node.size(); ++i) { + const auto entry = tenants_node[i]; + if (!entry || !entry.IsMap()) { + return tl::make_unexpected("tenant entry must be a YAML map"); + } + const auto name_node = entry["name"]; + const auto quota_node = entry["quota"]; + if (!name_node || !name_node.IsScalar()) { + return tl::make_unexpected("tenant name is required"); + } + if (!quota_node || !quota_node.IsScalar()) { + return tl::make_unexpected("tenant quota is required"); + } + + std::string name; + std::string quota; + try { + name = name_node.as(); + quota = quota_node.as(); + } catch (const YAML::Exception& e) { + return tl::make_unexpected(std::string("invalid tenant entry: ") + + e.what()); + } + if (!IsValidTenantQuotaName(name)) { + return tl::make_unexpected("invalid tenant name '" + name + "'"); + } + const TenantId tenant_id(std::move(name)); + name = tenant_id.value(); + if (!IsValidTenantQuotaName(name)) { + return tl::make_unexpected("invalid tenant name '" + name + "'"); + } + if (snapshot.tenant_quotas.contains(name)) { + return tl::make_unexpected("duplicate tenant name '" + name + "'"); + } + + auto quota_bytes = ParseTenantQuotaBytes(quota); + if (!quota_bytes) { + return tl::make_unexpected("invalid quota for tenant '" + name + + "': " + quota_bytes.error()); + } + snapshot.tenant_quotas.emplace(std::move(name), quota_bytes.value()); + } + + return snapshot; +} + +std::string FormatTenantQuotaPolicyYaml( + const TenantQuotaPolicySnapshot& snapshot) { + std::ostringstream out; + out << "version: 1\n\n"; + if (snapshot.tenant_quotas.empty()) { + out << "tenants: []\n"; + return out.str(); + } + out << "tenants:\n"; + for (const auto& [tenant_id, quota] : snapshot.tenant_quotas) { + out << " - name: " << QuoteYamlDoubleQuotedScalar(tenant_id) << "\n"; + out << " quota: " << quota << "\n"; + } + return out.str(); +} + +YamlTenantQuotaPolicyStore::YamlTenantQuotaPolicyStore(std::string path) + : path_(std::move(path)) {} + +tl::expected +YamlTenantQuotaPolicyStore::Load() { + std::lock_guard lock(mutex_); + std::ifstream input(path_); + if (!input.is_open()) { + return tl::make_unexpected("failed to open tenant quota policy file '" + + path_ + "'"); + } + std::ostringstream buffer; + buffer << input.rdbuf(); + if (input.bad()) { + return tl::make_unexpected("failed to read tenant quota policy file '" + + path_ + "'"); + } + return ParseTenantQuotaPolicyYaml(buffer.str()); +} + +tl::expected YamlTenantQuotaPolicyStore::Save( + const TenantQuotaPolicySnapshot& snapshot) { + std::lock_guard lock(mutex_); + const std::string content = FormatTenantQuotaPolicyYaml(snapshot); + const std::string tmp_path = MakeTempPath(path_); + + int fd = ::open(tmp_path.c_str(), O_CREAT | O_TRUNC | O_WRONLY, 0644); + if (fd < 0) { + return tl::make_unexpected(ErrnoMessage("open", tmp_path)); + } + + auto cleanup = [&] { + ::close(fd); + ::unlink(tmp_path.c_str()); + }; + + auto write_result = WriteAll(fd, content, tmp_path); + if (!write_result) { + cleanup(); + return tl::make_unexpected(write_result.error()); + } + if (::fsync(fd) != 0) { + std::string error = ErrnoMessage("fsync", tmp_path); + cleanup(); + return tl::make_unexpected(error); + } + if (::close(fd) != 0) { + std::string error = ErrnoMessage("close", tmp_path); + ::unlink(tmp_path.c_str()); + return tl::make_unexpected(error); + } + fd = -1; + + if (::rename(tmp_path.c_str(), path_.c_str()) != 0) { + std::string error = ErrnoMessage("rename", path_); + ::unlink(tmp_path.c_str()); + return tl::make_unexpected(error); + } + + auto fsync_result = FsyncDirectory(path_); + if (!fsync_result) { + LOG(WARNING) << "failed to fsync tenant quota policy directory after " + "rename: " + << fsync_result.error(); + } + return {}; +} + +#ifdef STORE_USE_ETCD +EtcdTenantQuotaPolicyStore::EtcdTenantQuotaPolicyStore( + const std::string& endpoints, const std::string& cluster_id) { + auto key = BuildTenantQuotaEtcdKey(cluster_id); + if (!key) { + throw std::invalid_argument(key.error()); + } + if (endpoints.empty()) { + throw std::invalid_argument( + "tenant quota etcd connector requires a non-empty uri"); + } + ErrorCode connect_error = EtcdHelper::ConnectToEtcdStoreClient(endpoints); + if (connect_error != ErrorCode::OK) { + if (connect_error == ErrorCode::INVALID_PARAMS) { + throw std::runtime_error( + "failed to connect tenant quota etcd store: " + "tenant_quota_connector_uri must match the already connected " + "store etcd endpoints used by HA/oplog"); + } + throw std::runtime_error("failed to connect tenant quota etcd store: " + + toString(connect_error)); + } + key_ = std::move(key.value()); +} + +tl::expected +EtcdTenantQuotaPolicyStore::Load() { + std::lock_guard lock(mutex_); + std::string content; + EtcdRevisionId revision_id = 0; + ErrorCode error = + EtcdHelper::Get(key_.c_str(), key_.size(), content, revision_id); + if (error == ErrorCode::ETCD_KEY_NOT_EXIST) { + return TenantQuotaPolicySnapshot{}; + } + if (error != ErrorCode::OK) { + return tl::make_unexpected( + "failed to load tenant quota policy from " + "etcd key '" + + key_ + "': " + toString(error)); + } + return ParseTenantQuotaPolicyYaml(content); +} + +tl::expected EtcdTenantQuotaPolicyStore::Save( + const TenantQuotaPolicySnapshot& snapshot) { + std::lock_guard lock(mutex_); + const std::string content = FormatTenantQuotaPolicyYaml(snapshot); + ErrorCode error = EtcdHelper::Put(key_.c_str(), key_.size(), + content.c_str(), content.size()); + if (error != ErrorCode::OK) { + return tl::make_unexpected( + "failed to save tenant quota policy to " + "etcd key '" + + key_ + "': " + toString(error)); + } + return {}; +} +#endif + +tl::expected, std::string> +CreateTenantQuotaPolicyStore(const std::string& type, const std::string& uri, + const std::string& cluster_id) { + if (type == "file") { + if (uri.empty()) { + return tl::make_unexpected( + "tenant quota file connector requires a non-empty uri"); + } + return std::make_unique(uri); + } + if (type == "etcd") { +#ifdef STORE_USE_ETCD + try { + return std::make_unique(uri, + cluster_id); + } catch (const std::exception& e) { + return tl::make_unexpected(e.what()); + } +#else + return tl::make_unexpected( + "tenant quota etcd connector requires STORE_USE_ETCD"); +#endif + } + return tl::make_unexpected("unsupported tenant quota connector type '" + + type + "'"); +} + +} // namespace mooncake diff --git a/mooncake-store/src/transfer_task.cpp b/mooncake-store/src/transfer_task.cpp index 5995e3c4fd..ca7faef384 100644 --- a/mooncake-store/src/transfer_task.cpp +++ b/mooncake-store/src/transfer_task.cpp @@ -6,11 +6,12 @@ #include #include #include +#include #include #include #include #include -#include "gpu_staging_utils.h" +#include "device/accelerator_registry.h" #include "transfer_engine.h" #include "transport/transport.h" #ifdef USE_NOF @@ -650,26 +651,56 @@ void MemcpyWorkerPool::workerThread() { if (task.state) { try { bool ok = true; + auto runtime_accelerator = + device::GetAcceleratorRegistry().RuntimeAccelerators(); for (const auto& op : task.operations) { - int src_dev = -1, dst_dev = -1; - bool src_on_gpu = - gpu_staging::IsDevicePointer(op.src, &src_dev); - bool dst_on_gpu = - gpu_staging::IsDevicePointer(op.dest, &dst_dev); - - if (!src_on_gpu && !dst_on_gpu) { + device::PointerInfo src_info; + device::PointerInfo dst_info; + auto* src_device = runtime_accelerator.FindDeviceForPointer( + op.src, &src_info); + auto* dst_device = runtime_accelerator.FindDeviceForPointer( + op.dest, &dst_info); + + if (!src_device && !dst_device) { std::memcpy(op.dest, op.src, op.size); } else { - int dev = src_on_gpu ? src_dev : dst_dev; - gpu_staging::SetDevice(dev); - if (!gpu_staging::CopyAuto(op.dest, op.src, op.size)) { + if (src_device && dst_device && + src_device != dst_device) { LOG(ERROR) - << "GPU memcpy failed: src_dev=" << src_dev - << " dst_dev=" << dst_dev + << "GPU memcpy failed: source and destination " + "belong to different accelerator runtimes" + << " src_dev=" << src_info.device_id + << " dst_dev=" << dst_info.device_id << " size=" << op.size; ok = false; break; } + const device::AcceleratorDevice* accelerator = nullptr; + int32_t device_id = -1; + device::CopyDirection direction; + if (src_device) { + accelerator = src_device; + device_id = src_info.device_id; + direction = device::CopyDirection::kDeviceToHost; + if (dst_device) { + direction = + device::CopyDirection::kDeviceToDevice; + } + } else { + accelerator = dst_device; + device_id = dst_info.device_id; + direction = device::CopyDirection::kHostToDevice; + } + accelerator->SetContext(device_id); + if (!accelerator->Copy(op.dest, op.src, op.size, + direction)) { + LOG(ERROR) << "GPU memcpy failed: src_dev=" + << src_info.device_id + << " dst_dev=" << dst_info.device_id + << " size=" << op.size; + ok = false; + break; + } } } @@ -1046,6 +1077,11 @@ std::optional TransferSubmitter::submit_batch( return future; } +TransferEngine::ScatterTransferOperation TransferSubmitter::submitScatter( + const std::vector& transfers) { + return engine_.submitScatter(transfers); +} + std::optional TransferSubmitter::submit_batch_get_offload_object( const std::string& transfer_engine_addr, diff --git a/mooncake-store/src/types.cpp b/mooncake-store/src/types.cpp index db39474eb8..7535abccbf 100644 --- a/mooncake-store/src/types.cpp +++ b/mooncake-store/src/types.cpp @@ -36,6 +36,7 @@ const std::string& toString(ErrorCode errorCode) noexcept { {ErrorCode::REPLICA_IS_GONE, "REPLICA_IS_GONE"}, {ErrorCode::OBJECT_REPLICA_BUSY, "OBJECT_REPLICA_BUSY"}, {ErrorCode::TRANSFER_FAIL, "TRANSFER_FAIL"}, + {ErrorCode::CHECKSUM_MISMATCH, "CHECKSUM_MISMATCH"}, {ErrorCode::RPC_FAIL, "RPC_FAIL"}, {ErrorCode::RPC_TIMEOUT, "RPC_TIMEOUT"}, {ErrorCode::ETCD_OPERATION_ERROR, "ETCD_OPERATION_ERROR"}, @@ -45,6 +46,7 @@ const std::string& toString(ErrorCode errorCode) noexcept { {ErrorCode::OPLOG_ENTRY_NOT_FOUND, "OPLOG_ENTRY_NOT_FOUND"}, {ErrorCode::K8S_LEASE_OPERATION_ERROR, "K8S_LEASE_OPERATION_ERROR"}, {ErrorCode::K8S_LEASE_NOT_FOUND, "K8S_LEASE_NOT_FOUND"}, + {ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, "INCOMPLETE_OPLOG_CATCH_UP"}, {ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS, "UNAVAILABLE_IN_CURRENT_STATUS"}, {ErrorCode::UNAVAILABLE_IN_CURRENT_MODE, "UNAVAILABLE_IN_CURRENT_MODE"}, @@ -74,7 +76,9 @@ const std::string& toString(ErrorCode errorCode) noexcept { {ErrorCode::DFS_PERMISSION_DENIED, "DFS_PERMISSION_DENIED"}, {ErrorCode::DFS_STALE_HANDLE, "DFS_STALE_HANDLE"}, {ErrorCode::DFS_PARTIAL_WRITE, "DFS_PARTIAL_WRITE"}, - {ErrorCode::TENANT_QUOTA_EXCEEDED, "TENANT_QUOTA_EXCEEDED"}}; + {ErrorCode::TENANT_QUOTA_EXCEEDED, "TENANT_QUOTA_EXCEEDED"}, + {ErrorCode::TENANT_NOT_REGISTERED, "TENANT_NOT_REGISTERED"}, + {ErrorCode::TENANT_NOT_EMPTY, "TENANT_NOT_EMPTY"}}; auto it = errorCodeMap.find(errorCode); static const std::string unknownError = "UNKNOWN_ERROR"; diff --git a/mooncake-store/src/uds_transport.cpp b/mooncake-store/src/uds_transport.cpp new file mode 100644 index 0000000000..2b8c2a27ac --- /dev/null +++ b/mooncake-store/src/uds_transport.cpp @@ -0,0 +1,344 @@ +#include "uds_transport.h" + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace mooncake { +namespace { + +std::string errnoMessage(const std::string &operation) { + return operation + ": " + strerror(errno); +} + +tl::expected makeAbstractAddress( + const std::string &socket_name, sockaddr_un &addr, socklen_t &addr_len) { + if (socket_name.size() > sizeof(addr.sun_path) - 2) { + return tl::make_unexpected("UDS socket name too long: " + socket_name); + } + + memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + addr.sun_path[0] = '\0'; + strncpy(addr.sun_path + 1, socket_name.c_str(), sizeof(addr.sun_path) - 2); + addr_len = sizeof(sa_family_t) + 1 + socket_name.length(); + return {}; +} + +tl::expected setSendTimeout( + int fd, std::chrono::milliseconds timeout) { + timeval tv = { + .tv_sec = + std::chrono::duration_cast(timeout).count(), + .tv_usec = std::chrono::duration_cast( + timeout % std::chrono::seconds(1)) + .count(), + }; + if (setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv)) < 0) { + return tl::make_unexpected( + errnoMessage("Failed to set UDS send timeout")); + } + return {}; +} + +tl::expected createConnectedSocket( + const std::string &socket_name, std::chrono::milliseconds connect_timeout) { + if (connect_timeout.count() <= 0) { + return tl::make_unexpected("UDS connect timeout must be positive"); + } + + int sock_fd = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + if (sock_fd < 0) { + return tl::make_unexpected(errnoMessage("Failed to create UDS socket")); + } + + sockaddr_un addr; + socklen_t addr_len = 0; + auto addr_result = makeAbstractAddress(socket_name, addr, addr_len); + if (!addr_result) { + close(sock_fd); + return tl::make_unexpected(addr_result.error()); + } + + auto timeout_result = setSendTimeout(sock_fd, connect_timeout); + if (!timeout_result) { + close(sock_fd); + return tl::make_unexpected(timeout_result.error()); + } + + if (::connect(sock_fd, reinterpret_cast(&addr), addr_len) < 0) { + auto error = + errno == EAGAIN || errno == EINPROGRESS + ? "Timed out connecting UDS socket '" + socket_name + "'" + : errnoMessage("Failed to connect UDS socket '" + socket_name + + "'"); + close(sock_fd); + return tl::make_unexpected(error); + } + + auto clear_timeout_result = + setSendTimeout(sock_fd, std::chrono::milliseconds(0)); + if (!clear_timeout_result) { + close(sock_fd); + return tl::make_unexpected(clear_timeout_result.error()); + } + + return sock_fd; +} + +} // namespace + +UdsConnection::UdsConnection(int fd) : fd_(fd) {} + +UdsConnection::~UdsConnection() { close(); } + +UdsConnection::UdsConnection(UdsConnection &&other) noexcept + : fd_(other.release()) {} + +UdsConnection &UdsConnection::operator=(UdsConnection &&other) noexcept { + if (this != &other) { + close(); + fd_ = other.release(); + } + return *this; +} + +bool UdsConnection::valid() const { return fd_ >= 0; } + +int UdsConnection::fd() const { return fd_; } + +int UdsConnection::release() { + int fd = fd_; + fd_ = -1; + return fd; +} + +void UdsConnection::close() { + if (fd_ >= 0) { + ::close(fd_); + fd_ = -1; + } +} + +tl::expected UdsConnection::setRecvTimeout( + std::chrono::seconds timeout) { + timeval tv = {.tv_sec = timeout.count(), .tv_usec = 0}; + if (setsockopt(fd_, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv)) < 0) { + return tl::make_unexpected( + errnoMessage("Failed to set UDS recv timeout")); + } + return {}; +} + +int UdsConnection::sendRaw(const void *data, size_t len) { + const char *pos = static_cast(data); + size_t remaining = len; + while (remaining > 0) { + ssize_t sent = ::send(fd_, pos, remaining, 0); + if (sent < 0 && errno == EINTR) continue; + if (sent <= 0) return -1; + pos += sent; + remaining -= static_cast(sent); + } + return 0; +} + +int UdsConnection::recvRaw(void *data, size_t len) { + char *pos = static_cast(data); + size_t remaining = len; + while (remaining > 0) { + ssize_t received = ::recv(fd_, pos, remaining, 0); + if (received < 0 && errno == EINTR) continue; + if (received <= 0) return -1; + pos += received; + remaining -= static_cast(received); + } + return 0; +} + +int UdsConnection::sendFd(int fd, void *data, size_t data_len) { + msghdr msg; + memset(&msg, 0, sizeof(msg)); + iovec iov; + char buf[CMSG_SPACE(sizeof(int))]; + memset(buf, 0, sizeof(buf)); + + iov.iov_base = data; + iov.iov_len = data_len; + + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + msg.msg_control = buf; + msg.msg_controllen = sizeof(buf); + + cmsghdr *cmsg = CMSG_FIRSTHDR(&msg); + cmsg->cmsg_level = SOL_SOCKET; + cmsg->cmsg_type = SCM_RIGHTS; + cmsg->cmsg_len = CMSG_LEN(sizeof(int)); + memcpy(CMSG_DATA(cmsg), &fd, sizeof(int)); + + while (true) { + ssize_t sent = sendmsg(fd_, &msg, 0); + if (sent < 0 && errno == EINTR) continue; + if (sent < 0) return -1; + return sent == static_cast(data_len) ? 0 : -1; + } +} + +int UdsConnection::recvFd(void *data, size_t data_len) { + msghdr msg; + memset(&msg, 0, sizeof(msg)); + iovec iov; + char buf[CMSG_SPACE(sizeof(int))]; + memset(buf, 0, sizeof(buf)); + + iov.iov_base = data; + iov.iov_len = data_len; + + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + msg.msg_control = buf; + msg.msg_controllen = sizeof(buf); + + ssize_t received = 0; + while (true) { + received = recvmsg(fd_, &msg, 0); + if (received < 0 && errno == EINTR) continue; + if (received <= 0) return -1; + break; + } + + cmsghdr *cmsg = CMSG_FIRSTHDR(&msg); + if (cmsg && cmsg->cmsg_level == SOL_SOCKET && + cmsg->cmsg_type == SCM_RIGHTS) { + int received_fd = -1; + memcpy(&received_fd, CMSG_DATA(cmsg), sizeof(int)); + // Even for malformed messages, close any delivered fd instead of + // leaking it into this process. + if (received != static_cast(data_len) || + (msg.msg_flags & (MSG_TRUNC | MSG_CTRUNC))) { + ::close(received_fd); + return -1; + } + return received_fd; + } + return -1; +} + +UdsConnector::UdsConnector(std::string socket_name, + std::chrono::milliseconds connect_timeout) + : socket_name_(std::move(socket_name)), connect_timeout_(connect_timeout) {} + +tl::expected, std::string> +UdsConnector::connect() { + auto fd = createConnectedSocket(socket_name_, connect_timeout_); + if (!fd) return tl::make_unexpected(fd.error()); + return std::make_unique(*fd); +} + +UdsAcceptor::UdsAcceptor(std::string socket_name) + : socket_name_(std::move(socket_name)) {} + +UdsAcceptor::~UdsAcceptor() { stop(); } + +void UdsAcceptor::registerHandler(Handler handler) { + handler_ = std::move(handler); +} + +tl::expected UdsAcceptor::start() { + if (running_.load()) return {}; + if (!handler_) { + return tl::make_unexpected("UDS acceptor handler is not registered"); + } + + listen_fd_ = socket(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0); + if (listen_fd_ < 0) { + return tl::make_unexpected(errnoMessage("Failed to create UDS socket")); + } + + sockaddr_un addr; + socklen_t addr_len = 0; + auto addr_result = makeAbstractAddress(socket_name_, addr, addr_len); + if (!addr_result) { + ::close(listen_fd_); + listen_fd_ = -1; + return tl::make_unexpected(addr_result.error()); + } + + if (bind(listen_fd_, reinterpret_cast(&addr), addr_len) < 0) { + auto error = + errnoMessage("Failed to bind UDS socket '" + socket_name_ + "'"); + ::close(listen_fd_); + listen_fd_ = -1; + return tl::make_unexpected(error); + } + + if (listen(listen_fd_, 5) < 0) { + auto error = errnoMessage("Failed to listen on UDS socket '" + + socket_name_ + "'"); + ::close(listen_fd_); + listen_fd_ = -1; + return tl::make_unexpected(error); + } + + running_ = true; + thread_ = std::jthread([this]() { acceptLoop(); }); + return {}; +} + +void UdsAcceptor::stop() { + if (!running_.exchange(false)) return; + wakeAccept(); + int active_client_fd = active_client_fd_.load(); + if (active_client_fd >= 0) { + // Stop should unblock handlers waiting on either read or write. + ::shutdown(active_client_fd, SHUT_RDWR); + } + if (thread_.joinable()) thread_.join(); + if (listen_fd_ >= 0) { + ::close(listen_fd_); + listen_fd_ = -1; + } +} + +void UdsAcceptor::acceptLoop() { + while (running_.load()) { + int client_sock = accept(listen_fd_, nullptr, nullptr); + if (client_sock < 0) { + if (running_.load()) { + LOG(ERROR) << "Accept failed: " << strerror(errno); + } + continue; + } + + UdsConnection connection(client_sock); + if (!running_.load()) break; + + active_client_fd_ = connection.fd(); + if (handler_) handler_(connection); + active_client_fd_ = -1; + } + + if (listen_fd_ >= 0) { + ::close(listen_fd_); + listen_fd_ = -1; + } +} + +void UdsAcceptor::wakeAccept() { + // accept() is a blocking syscall on the listener thread. Shutting down the + // listener wakes accept() so the loop can observe running_ == false. + if (listen_fd_ >= 0) { + ::shutdown(listen_fd_, SHUT_RDWR); + } +} + +} // namespace mooncake diff --git a/mooncake-store/src/uring_file.cpp b/mooncake-store/src/uring_file.cpp index 870eeba2f0..02b9137124 100644 --- a/mooncake-store/src/uring_file.cpp +++ b/mooncake-store/src/uring_file.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -143,6 +144,7 @@ class SharedUringRing { tl::expected batch_read(int fd, const ReadDesc* descs, int cnt) { ensure_buf_registered(); + uint64_t op = ++op_id_; size_t total = 0; int remaining = cnt; int idx = 0; @@ -161,9 +163,10 @@ class SharedUringRing { io_uring_prep_read_fixed(sqe, fd, d.buf, d.len, d.off, 0); else io_uring_prep_read(sqe, fd, d.buf, d.len, d.off); + sqe->user_data = op; } - auto res = collect(batch); + auto res = collect(batch, op); if (!res) return res; total += res.value(); idx += batch; @@ -182,9 +185,11 @@ class SharedUringRing { LOG(ERROR) << "[SharedUringRing] SQ full (fsync)"; return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); } + uint64_t op = ++op_id_; io_uring_prep_fsync(sqe, fd, IORING_FSYNC_DATASYNC); + sqe->user_data = op; - auto res = collect(1); + auto res = collect(1, op); if (!res) return tl::make_unexpected(res.error()); return {}; } @@ -242,8 +247,26 @@ class SharedUringRing { return std::max(s, MIN_CHUNK); } - // Drain exactly @expected CQEs and accumulate bytes. - tl::expected collect(int expected) { + static size_t max_rw_count() { + static const size_t value = [] { + long page_size = sysconf(_SC_PAGESIZE); + if (page_size <= 0) page_size = 4096; + return static_cast(INT_MAX) & + ~(static_cast(page_size) - 1); + }(); + return value; + } + + // Drain exactly @expected CQEs matching @op_id and + // accumulate bytes. Stale CQEs (user_data != op_id, e.g. from + // a previous batch_read / submit_rw that hit an error and left + // in-flight SQEs) are silently consumed and discarded. + // + // Uses peek-then-wait: after io_uring_submit_and_wait most CQEs + // are already available, so io_uring_peek_cqe() succeeds without + // a syscall in the common case. io_uring_wait_cqe() only kicks + // in when the ring is unexpectedly drained (stale CQE storms). + tl::expected collect(int expected, uint64_t op_id) { int ret = io_uring_submit_and_wait(&ring_, expected); if (ret < 0) { LOG(ERROR) << "[SharedUringRing] io_uring_submit_and_wait: " @@ -252,9 +275,24 @@ class SharedUringRing { } size_t total = 0; bool err = false; - unsigned head, cnt = 0; - struct io_uring_cqe* cqe; - io_uring_for_each_cqe(&ring_, head, cqe) { + int processed = 0; + + while (processed < expected) { + struct io_uring_cqe* cqe; + int wait_ret = io_uring_peek_cqe(&ring_, &cqe); + if (wait_ret == -EAGAIN) { + wait_ret = io_uring_wait_cqe(&ring_, &cqe); + } + if (wait_ret < 0) { + LOG(ERROR) << "[SharedUringRing] CQE wait error: " + << strerror(-wait_ret); + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + if (cqe->user_data != op_id) { + // Stale CQE from a previous failed operation. + io_uring_cq_advance(&ring_, 1); + continue; + } if (cqe->res < 0) { LOG(ERROR) << "[SharedUringRing] CQE error: " << strerror(-cqe->res); @@ -262,9 +300,9 @@ class SharedUringRing { } else { total += static_cast(cqe->res); } - ++cnt; + io_uring_cq_advance(&ring_, 1); + ++processed; } - io_uring_cq_advance(&ring_, cnt); if (err) return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); return total; } @@ -277,6 +315,7 @@ class SharedUringRing { is_write ? ErrorCode::FILE_WRITE_FAIL : ErrorCode::FILE_READ_FAIL; const bool fix_buf = (use_fixed_buf && buf_registered_); + uint64_t op = ++op_id_; char* ptr = static_cast(buf); size_t total = 0; size_t remaining = len; @@ -288,7 +327,7 @@ class SharedUringRing { static_cast((remaining + cs - 1) / cs), QUEUE_DEPTH); for (unsigned i = 0; i < n; ++i) { - size_t chunk = std::min(cs, remaining); + size_t chunk = std::min({cs, remaining, max_rw_count()}); struct io_uring_sqe* sqe = io_uring_get_sqe(&ring_); if (!sqe) { @@ -307,6 +346,7 @@ class SharedUringRing { else io_uring_prep_read(sqe, fd, ptr, chunk, cur); } + sqe->user_data = op; ptr += chunk; cur += static_cast(chunk); @@ -314,7 +354,7 @@ class SharedUringRing { if (remaining == 0) break; } - auto res = collect(static_cast(n)); + auto res = collect(static_cast(n), op); if (!res) return res; total += res.value(); if (!is_write && res.value() == 0) break; // read EOF @@ -332,14 +372,32 @@ class SharedUringRing { off_t off) { const ErrorCode err_code = is_write ? ErrorCode::FILE_WRITE_FAIL : ErrorCode::FILE_READ_FAIL; + const size_t max_io = max_rw_count(); + uint64_t op = ++op_id_; size_t total = 0; off_t cur = off; int remaining = cnt; int idx = 0; while (remaining > 0) { - int batch = std::min(remaining, static_cast(QUEUE_DEPTH)); + if (iovs[idx].iov_len > max_io) { + auto res = submit_rw(is_write, fd, iovs[idx].iov_base, + iovs[idx].iov_len, cur, + /*use_fixed_buf=*/false); + if (!res) return res; + total += res.value(); + cur += static_cast(iovs[idx].iov_len); + ++idx; + --remaining; + continue; + } + + int batch = 0; + while (batch < remaining && batch < static_cast(QUEUE_DEPTH) && + iovs[idx + batch].iov_len <= max_io) { + ++batch; + } for (int i = 0; i < batch; ++i) { struct io_uring_sqe* sqe = io_uring_get_sqe(&ring_); @@ -354,12 +412,13 @@ class SharedUringRing { else io_uring_prep_read(sqe, fd, iovs[idx].iov_base, iovs[idx].iov_len, cur); + sqe->user_data = op; cur += static_cast(iovs[idx].iov_len); ++idx; } - auto res = collect(batch); + auto res = collect(batch, op); if (!res) return res; total += res.value(); remaining -= batch; @@ -377,6 +436,7 @@ class SharedUringRing { bool buf_register_failed_ = false; // set on first failure; skip retries void* buf_base_ = nullptr; size_t buf_size_ = 0; + uint64_t op_id_ = 0; // monotonic ID for stale CQE filtering }; // ============================================================================ @@ -406,7 +466,8 @@ UringFile::~UringFile() { if (close(fd_) != 0) { LOG(WARNING) << "[UringFile] close failed: " << filename_; } - if (error_code_ == ErrorCode::FILE_WRITE_FAIL) { + if (delete_on_write_fail_ && + error_code_ == ErrorCode::FILE_WRITE_FAIL) { if (::unlink(filename_.c_str()) == -1) LOG(ERROR) << "[UringFile] failed to delete corrupted file: " << filename_; @@ -651,7 +712,7 @@ tl::expected UringFile::datasync() { auto res = SharedUringRing::instance().fsync(fd_); if (!res) { LOG(ERROR) << "[UringFile::datasync] fsync failed for: " << filename_; - return make_error(ErrorCode::FILE_WRITE_FAIL); + return tl::make_unexpected(ErrorCode::FILE_WRITE_FAIL); } return {}; } diff --git a/mooncake-store/src/utils.cpp b/mooncake-store/src/utils.cpp index d1f0354917..6b52484d74 100644 --- a/mooncake-store/src/utils.cpp +++ b/mooncake-store/src/utils.cpp @@ -1,9 +1,14 @@ #include "utils.h" +#include "random.h" #include "mmap_arena.h" #include "config.h" #include "common.h" #include "ub_allocator.h" +#include "ascii_string.h" +#include "bool_parser.h" +#include "environ.h" + #include #include #include @@ -13,19 +18,25 @@ #include #include #include -#include - #include #include -#include -#include #include #include +#include #include -#include +#include +#include #include #include -#include +#include +#include +#include +#ifdef USE_CUDA +#include +#endif +#ifdef USE_INTRA_NVLINK +#include "gpu_vendor/intra_nvlink.h" +#endif // Feature flag to enable/disable arena allocator. Disabled by default so the // library does not pre-map a large pool unless the operator opts in via gflag @@ -43,6 +54,9 @@ DEFINE_uint64(mmap_arena_pool_size, 8ULL * 1024 * 1024 * 1024, #if defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) #include "ascend_allocator.h" #endif +#if defined(USE_SUNRISE) +#include "sunrise_allocator.h" +#endif #ifdef USE_NOF #include "spdk/spdk_wrapper.h" @@ -73,12 +87,8 @@ bool isPortAvailable(int port) { // AutoPortBinder implementation AutoPortBinder::AutoPortBinder(int min_port, int max_port) : socket_fd_(-1), port_(-1) { - static std::random_device rand_gen; - std::mt19937 gen(rand_gen()); - std::uniform_int_distribution<> rand_dist(min_port, max_port); - for (int attempt = 0; attempt < 20; ++attempt) { - int port = rand_dist(gen); + int port = randomUniform(min_port, max_port); socket_fd_ = socket(AF_INET, SOCK_STREAM, 0); if (socket_fd_ < 0) continue; @@ -104,6 +114,36 @@ AutoPortBinder::~AutoPortBinder() { } } +#ifdef USE_VRAM_SEGMENT +tl::expected allocate_vram_memory( + size_t total_size, const std::string &protocol) { + cudaError_t res; + int device; + void *ptr = nullptr; + res = cudaGetDevice(&device); + if (res != cudaSuccess) { + LOG(ERROR) << "VRAM Segment cudaGetDevice failed."; + return tl::make_unexpected("VRAM Segment cudaGetDevice failed."); + } + if (protocol == "nvlink_intra") { +#ifdef USE_INTRA_NVLINK + ptr = allocateFabricMemory_intra(total_size); + return ptr; +#else + LOG(ERROR) << "Protocol nvlink_intra need USE_INTRA_NVLINK=ON. Please " + "rebuild mooncake from source."; + return tl::make_unexpected("Protocol not supported"); +#endif + } + res = cudaMalloc((void **)&ptr, total_size); + if (res != cudaSuccess) { + LOG(ERROR) << "VRAM Segment cudaMalloc failed."; + return tl::make_unexpected("VRAM Segment cudaMalloc failed."); + } + return ptr; +} +#endif + void *allocate_buffer_allocator_memory(size_t total_size, const std::string &protocol, size_t alignment, bool use_spdk_dma) { @@ -118,6 +158,13 @@ void *allocate_buffer_allocator_memory(size_t total_size, return ascend_allocate_memory(total_size, protocol); } #endif +#if defined(USE_SUNRISE) + if (protocol == "sunrise_link") { + return sunrise_allocate_memory( + total_size, alignment, + mooncake::globalConfig().sunrise_use_device_mem); + } +#endif #if defined(USE_UB) if (protocol == "ub") { return mooncake::ub_allocate_memory(alignment, total_size); @@ -128,6 +175,14 @@ void *allocate_buffer_allocator_memory(size_t total_size, return mooncake::SpdkWrapper::GetInstance().Alloc(total_size, alignment, -1); } +#endif +#ifdef USE_VRAM_SEGMENT + auto ret = allocate_vram_memory(total_size, protocol); + if (!ret) { + LOG(ERROR) << ret.error(); + return nullptr; + } + return *ret; #endif // Allocate aligned memory return aligned_alloc(alignment, total_size); @@ -141,12 +196,13 @@ static std::atomic g_arena_noop_free_count{0}; static void initializeGlobalArena() { const std::string env_pool_size = - GetEnvStringOr("MC_MMAP_ARENA_POOL_SIZE", ""); + Environ::GetString("MC_MMAP_ARENA_POOL_SIZE", ""); // Allow env var to override the gflag (useful when loaded as .so from // Python). An explicit pool-size env var is also treated as an opt-in, // because pybind11 users cannot easily pass gflags. - const std::string env_disable = GetEnvStringOr("MC_DISABLE_MMAP_ARENA", ""); - const std::optional disable_override = string_to_bool(env_disable); + const std::string env_disable = + Environ::GetString("MC_DISABLE_MMAP_ARENA", ""); + const std::optional disable_override = TryParseBool(env_disable); if (!env_disable.empty() && !disable_override.has_value()) { LOG(WARNING) << "Ignoring invalid MC_DISABLE_MMAP_ARENA='" << env_disable @@ -227,7 +283,133 @@ static inline size_t mmap_map_size(size_t total_size, size_t hugepage_size) { return align_up(total_size, page_size); } +namespace { + +size_t touch_thread_count(size_t page_count) { + const unsigned int hardware_threads = std::thread::hardware_concurrency(); + const size_t available_threads = + hardware_threads == 0 ? 1 : std::min(hardware_threads, 16); + return std::min(available_threads, page_count); +} + +void touch_page_range(volatile char *data, size_t page_size, size_t begin_page, + size_t end_page, int numa_node) { + if (numa_node >= 0 && numa_run_on_node(numa_node) != 0) { + LOG(WARNING) << "Failed to bind HugeTLB population worker to NUMA node " + << numa_node << ": " << std::strerror(errno); + } + for (size_t page = begin_page; page < end_page; ++page) { + data[page * page_size] = 0; + } +} + +void touch_mmap_pages(void *ptr, size_t map_size, size_t page_size) { + if (ptr == nullptr || map_size == 0 || page_size == 0) { + return; + } + + const size_t page_count = (map_size + page_size - 1) / page_size; + const size_t num_threads = touch_thread_count(page_count); + + auto *data = static_cast(ptr); + if (num_threads <= 1) { + touch_page_range(data, page_size, 0, page_count, -1); + return; + } + + std::vector threads; + threads.reserve(num_threads); + const size_t pages_per_thread = + (page_count + num_threads - 1) / num_threads; + for (size_t thread_index = 0; thread_index < num_threads; ++thread_index) { + const size_t begin_page = thread_index * pages_per_thread; + const size_t end_page = + std::min(begin_page + pages_per_thread, page_count); + if (begin_page >= end_page) { + break; + } + threads.emplace_back(touch_page_range, data, page_size, begin_page, + end_page, -1); + } +} + +void touch_numa_mmap_pages(void *ptr, size_t map_size, size_t page_size, + const std::vector &numa_nodes) { + if (ptr == nullptr || map_size == 0 || page_size == 0 || + numa_nodes.empty()) { + return; + } + + const size_t node_count = numa_nodes.size(); + if (map_size % node_count != 0 || + (map_size / node_count) % page_size != 0) { + LOG(ERROR) << "Invalid NUMA HugeTLB mapping layout: size=" << map_size + << ", page_size=" << page_size << ", nodes=" << node_count; + return; + } + + const size_t region_size = map_size / node_count; + const size_t pages_per_region = region_size / page_size; + const size_t page_count = pages_per_region * node_count; + const size_t num_threads = + std::max(node_count, touch_thread_count(page_count)); + const size_t base_threads_per_node = num_threads / node_count; + const size_t extra_threads = num_threads % node_count; + + auto *data = static_cast(ptr); + std::vector threads; + threads.reserve(num_threads); + for (size_t node_index = 0; node_index < node_count; ++node_index) { + const size_t node_threads = + base_threads_per_node + (node_index < extra_threads ? 1 : 0); + const size_t pages_per_thread = + (pages_per_region + node_threads - 1) / node_threads; + const size_t region_begin_page = node_index * pages_per_region; + for (size_t thread_index = 0; thread_index < node_threads; + ++thread_index) { + const size_t begin_page = + region_begin_page + thread_index * pages_per_thread; + const size_t end_page = + std::min(begin_page + pages_per_thread, + region_begin_page + pages_per_region); + if (begin_page >= end_page) { + continue; + } + threads.emplace_back(touch_page_range, data, page_size, begin_page, + end_page, numa_nodes[node_index]); + } + } +} + +} // namespace + +void populate_hugetlb_mapping(void *ptr, size_t total_size) { + const size_t hugepage_size = get_hugepage_size_from_env(); + if (ptr == nullptr || total_size == 0 || hugepage_size == 0) { + return; + } + + touch_mmap_pages(ptr, mmap_map_size(total_size, hugepage_size), + hugepage_size); +} + +void populate_hugetlb_numa_mapping(void *ptr, size_t total_size, + const std::vector &numa_nodes) { + const size_t hugepage_size = get_hugepage_size_from_env(); + if (ptr == nullptr || total_size == 0 || hugepage_size == 0 || + numa_nodes.empty()) { + return; + } + + touch_numa_mmap_pages(ptr, total_size, hugepage_size, numa_nodes); +} + void *allocate_buffer_mmap_memory(size_t total_size, size_t alignment) { + return allocate_buffer_mmap_memory(total_size, alignment, false); +} + +void *allocate_buffer_mmap_memory(size_t total_size, size_t alignment, + bool defer_hugetlb_population) { if (total_size == 0) { LOG(ERROR) << "Total size must be greater than 0 for mmap"; return nullptr; @@ -256,7 +438,12 @@ void *allocate_buffer_mmap_memory(size_t total_size, size_t alignment) { } // Traditional mmap allocation (fallback or arena disabled). - unsigned int flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_POPULATE; + const bool defer_direct_population = + defer_hugetlb_population && get_hugepage_size_from_env() > 0; + unsigned int flags = MAP_PRIVATE | MAP_ANONYMOUS; + if (!defer_direct_population) { + flags |= MAP_POPULATE; + } const size_t hugepage_size = get_hugepage_size_from_env(&flags); const size_t map_size = mmap_map_size(total_size, hugepage_size); const size_t guaranteed_alignment = @@ -280,6 +467,11 @@ void *allocate_buffer_mmap_memory(size_t total_size, size_t alignment) { return ptr; } +bool is_mmap_arena_allocation(const void *ptr) { + return ptr != nullptr && g_mmap_arena && g_mmap_arena->isInitialized() && + g_mmap_arena->owns(ptr); +} + void free_buffer_mmap_memory(void *ptr, size_t total_size) { if (!ptr || total_size == 0) { return; @@ -327,12 +519,21 @@ void *allocate_buffer_numa_segments(size_t total_size, size_t region_size = align_up(total_size / n, page_size); size_t map_size = region_size * n; - // reserve contiguous VMA, no physical pages yet - void *ptr = mmap(nullptr, map_size, PROT_READ | PROT_WRITE, - MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + // reserve contiguous VMA; use hugepages if page_size indicates so + unsigned int flags = MAP_PRIVATE | MAP_ANONYMOUS; + if (page_size == SZ_2MB) { + flags |= MAP_HUGETLB | MAP_HUGE_2MB; + } else if (page_size == SZ_1GB) { + flags |= MAP_HUGETLB | MAP_HUGE_1GB; + } else if (page_size != static_cast(getpagesize())) { + flags |= MAP_HUGETLB; + } + void *ptr = mmap(nullptr, map_size, PROT_READ | PROT_WRITE, flags, -1, 0); if (ptr == MAP_FAILED) { - LOG(ERROR) << "mmap failed, size=" << map_size << ", errno=" << errno - << " (" << strerror(errno) << ")"; + LOG(ERROR) << "mmap failed (hugepage=" + << ((flags & MAP_HUGETLB) ? "yes" : "no") + << "), size=" << map_size << ", errno=" << errno << " (" + << strerror(errno) << ")"; return nullptr; } @@ -353,10 +554,9 @@ void *allocate_buffer_numa_segments(size_t total_size, } } - // No explicit prefault needed — ibv_reg_mr() will call get_user_pages() - // which triggers page faults that respect the mbind NUMA policy. - // Pages are allocated directly on the target NUMA during MR registration, - // avoiding a redundant full-buffer traversal. + // Leave the mapping lazy. The caller may explicitly populate it with + // NUMA-local workers before registration; otherwise ibv_reg_mr() calls + // get_user_pages(), whose faults respect the mbind policy. LOG(INFO) << "Allocated NUMA-segmented buffer: " << map_size << " bytes, " << n << " regions, page_size=" << page_size << ", nodes=[" << @@ -371,54 +571,41 @@ void *allocate_buffer_numa_segments(size_t total_size, return ptr; } -void free_memory(const std::string &protocol, void *ptr) { +void free_memory(const std::string &protocol, void *ptr, bool use_spdk_dma) { #if defined(USE_ASCEND_DIRECT) || defined(USE_UBSHMEM) if (protocol == "ascend" || protocol == "ubshmem") { return ascend_free_memory(protocol, ptr); } #endif +#if defined(USE_SUNRISE) + if (protocol == "sunrise_link") { + return sunrise_free_memory(ptr); + } +#endif #if defined(USE_UB) if (protocol == "ub") { mooncake::ub_free_memory(ptr); return; } #endif - free(ptr); -} - -std::string formatDeviceNames(const std::string &device_names) { - std::stringstream ss(device_names); - std::string item; - std::vector tokens; - while (getline(ss, item, ',')) { - tokens.push_back(item); - } - - std::string formatted; - for (size_t i = 0; i < tokens.size(); ++i) { - formatted += "\"" + tokens[i] + "\""; - if (i < tokens.size() - 1) { - formatted += ","; - } - } - return formatted; -} - -std::vector splitString(const std::string &str, char delimiter, - bool trim_spaces, bool keep_empty) { - std::vector result; - - boost::split( - result, str, boost::is_any_of(std::string(1, delimiter)), - keep_empty ? boost::token_compress_off : boost::token_compress_on); - - if (trim_spaces) { - for (auto &token : result) { - boost::trim(token); - } +#ifdef USE_NOF + // Mirror allocate_buffer_allocator_memory(): a buffer taken from the SPDK + // hugepage pool (spdk_zmalloc) must be released with spdk_free, not glibc + // free(), which would abort with "free(): invalid pointer". + if (use_spdk_dma) { + mooncake::SpdkWrapper::GetInstance().Free(ptr); + return; } - - return result; +#endif +#ifdef USE_VRAM_SEGMENT +#ifdef USE_INTRA_NVLINK + freeFabricMemory_intra(ptr); +#else + cudaFree(ptr); +#endif + return; +#endif + free(ptr); } tl::expected httpGet(const std::string &url) { @@ -546,9 +733,22 @@ int64_t time_gen() { .count(); } -std::string GetEnvStringOr(const char *name, const std::string &default_value) { - const char *env_val = std::getenv(name); - return env_val ? std::string(env_val) : default_value; +std::string ResolveMooncakeHostId(const std::string &local_hostname) { + const std::string hostname(TrimAsciiWhitespace(local_hostname)); + const std::string host_id = (hostname == "::1" || hostname == "::") + ? hostname + : std::string(TrimAsciiWhitespace( + getHostNameWithoutPort(hostname))); + if (host_id.empty()) { + return ""; + } + + if (AsciiCaseInsensitiveEquals(host_id, "localhost") || + host_id == "127.0.0.1" || host_id == "0.0.0.0" || host_id == "::1" || + host_id == "[::1]" || host_id == "::" || host_id == "[::]") { + return ""; + } + return host_id; } static std::string SanitizeKey(const std::string &key) { diff --git a/mooncake-store/src/utils/s3_helper.cpp b/mooncake-store/src/utils/s3_helper.cpp index f951d4bbee..fd2f9ead24 100644 --- a/mooncake-store/src/utils/s3_helper.cpp +++ b/mooncake-store/src/utils/s3_helper.cpp @@ -6,15 +6,15 @@ #include #include #include -#include +#include #include #include #include +#include #include #include #include #include -#include #include #include #include @@ -27,69 +27,30 @@ #include #include #include -#include "utils/type_util.h" +#include +#include "environ.h" #include "fmt/format.h" namespace mooncake { namespace { -constexpr int64_t kDefaultS3ConnectTimeoutMs = 10000; -constexpr int64_t kDefaultS3RequestTimeoutMs = 30000; - -struct S3Env { - std::string region; - - std::string endpoint; - - std::string bucket; - - std::string access_key; - - std::string secret_key; - - bool use_virtual_addressing = true; - - int64_t connect_timeout_ms = kDefaultS3ConnectTimeoutMs; - - int64_t request_timeout_ms = kDefaultS3RequestTimeoutMs; -}; - -S3Env s3_env; - -void AssignStringFromEnv(const char *env_name, std::string &target) { - const char *env_value = std::getenv(env_name); - if (env_value && *env_value) { - target = env_value; - } else { - target.clear(); - } -} - -void AssignBoolFromEnv(const char *env_name, bool &target) { - const char *env_value = std::getenv(env_name); - if (env_value && *env_value) { - bool parsed = true; - if (TypeUtil::ParseBool(env_value, parsed)) { - target = parsed; - return; - } - LOG(WARNING) << "Invalid " << env_name << " value: " << env_value; - } -} - -void AssignTimeoutFromEnv(const char *env_name, int64_t default_value, - int64_t &target) { - const char *env_value = std::getenv(env_name); - if (env_value && *env_value) { - int64_t parsed; - if (TypeUtil::ParseInt64(env_value, parsed)) { - target = parsed; - return; - } - LOG(WARNING) << "Invalid " << env_name << " value: " << env_value; - } - target = default_value; +// Parse a checksum-mode string ("when_supported" | "when_required") into the +// AWS SDK enum. Returns std::nullopt when the value is unset or invalid — the +// caller should then leave the AWS SDK's default in place. This is purely +// string-to-enum logic (not getenv), so it lives next to the AWS types it +// produces. +template +std::optional ParseChecksumMode(const std::string &value) { + if (value.empty()) return std::nullopt; + std::string lower = value; + std::transform(lower.begin(), lower.end(), lower.begin(), + [](unsigned char ch) { return std::tolower(ch); }); + if (lower == "when_required") return T::WHEN_REQUIRED; + if (lower == "when_supported") return T::WHEN_SUPPORTED; + LOG(WARNING) << "Invalid value: " << value + << ", ignoring (keeping AWS SDK default)"; + return std::nullopt; } } // namespace @@ -102,21 +63,10 @@ void S3Helper::InitAPI() { Aws::InitAPI(options_); aws_initialized = true; - // Read environment variables once during initialization (fallback as needed - // if not set) - AssignStringFromEnv("MOONCAKE_AWS_REGION", s3_env.region); - AssignStringFromEnv("MOONCAKE_AWS_S3_ENDPOINT", s3_env.endpoint); - AssignStringFromEnv("MOONCAKE_AWS_BUCKET_NAME", s3_env.bucket); - AssignStringFromEnv("MOONCAKE_AWS_ACCESS_KEY_ID", s3_env.access_key); - AssignStringFromEnv("MOONCAKE_AWS_SECRET_ACCESS_KEY", s3_env.secret_key); - - AssignBoolFromEnv("MOONCAKE_AWS_USE_VIRTUAL_ADDRESSING", - s3_env.use_virtual_addressing); - - AssignTimeoutFromEnv("MOONCAKE_AWS_CONNECT_TIMEOUT_MS", - kDefaultS3ConnectTimeoutMs, s3_env.connect_timeout_ms); - AssignTimeoutFromEnv("MOONCAKE_AWS_REQUEST_TIMEOUT_MS", - kDefaultS3RequestTimeoutMs, s3_env.request_timeout_ms); + // Force Environ initialization so all MOONCAKE_AWS_* env vars are read + // once during startup, matching the previous "read once at InitAPI" + // behavior. + (void)Environ::Get(); } void S3Helper::ShutdownAPI() { @@ -128,30 +78,43 @@ void S3Helper::ShutdownAPI() { S3Helper::S3Helper(const std::string &endpoint, const std::string &bucket, const std::string ®ion) { + const auto &env = Environ::Get(); Aws::Client::ClientConfiguration config(true); - config.connectTimeoutMs = s3_env.connect_timeout_ms; - config.requestTimeoutMs = s3_env.request_timeout_ms; - config.scheme = Aws::Http::Scheme::HTTPS; + config.connectTimeoutMs = env.GetAwsConnectTimeoutMs(); + config.requestTimeoutMs = env.GetAwsRequestTimeoutMs(); + config.scheme = env.GetAwsUseHttps() ? Aws::Http::Scheme::HTTPS + : Aws::Http::Scheme::HTTP; + if (auto request_checksum = + ParseChecksumMode( + env.GetAwsRequestChecksumCalculation())) { + config.checksumConfig.requestChecksumCalculation = *request_checksum; + } + if (auto response_checksum = + ParseChecksumMode( + env.GetAwsResponseChecksumValidation())) { + config.checksumConfig.responseChecksumValidation = *response_checksum; + } if (!region.empty()) { config.region = region; } else { - config.region = s3_env.region; + config.region = env.GetAwsRegion(); } if (!endpoint.empty()) { config.endpointOverride = endpoint; } else { - config.endpointOverride = s3_env.endpoint; + config.endpointOverride = env.GetAwsS3Endpoint(); } - bucket_ = s3_env.bucket; + bucket_ = env.GetAwsBucketName(); if (!bucket.empty()) { bucket_ = bucket; } - Aws::Auth::AWSCredentials credentials(s3_env.access_key, s3_env.secret_key); + Aws::Auth::AWSCredentials credentials(env.GetAwsAccessKeyId(), + env.GetAwsSecretAccessKey()); // Concatenate log information into member variable connection_info_ connection_info_ = fmt::format( @@ -169,14 +132,14 @@ S3Helper::S3Helper(const std::string &endpoint, const std::string &bucket, bucket_.empty() ? "unset" : bucket_, config.connectTimeoutMs, config.requestTimeoutMs, config.scheme == Aws::Http::Scheme::HTTPS ? "HTTPS" : "HTTP", - !s3_env.access_key.empty() ? "set" : "unset", - !s3_env.secret_key.empty() ? "set" : "unset", - s3_env.use_virtual_addressing ? "true" : "false"); + !env.GetAwsAccessKeyId().empty() ? "set" : "unset", + !env.GetAwsSecretAccessKey().empty() ? "set" : "unset", + env.GetAwsUseVirtualAddressing() ? "true" : "false"); s3_client_ = Aws::S3::S3Client( credentials, config, Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::Never, - s3_env.use_virtual_addressing); + env.GetAwsUseVirtualAddressing()); } S3Helper::~S3Helper() = default; @@ -818,7 +781,7 @@ tl::expected S3Helper::ListObjectsWithPrefix( const std::string &prefix, std::vector &object_keys) { object_keys.clear(); - Aws::S3::Model::ListObjectsRequest request; + Aws::S3::Model::ListObjectsV2Request request; request.WithBucket(bucket_); request.WithPrefix(prefix); @@ -828,10 +791,10 @@ tl::expected S3Helper::ListObjectsWithPrefix( bool done = false; while (!done) { - auto outcome = s3_client_.ListObjects(request); + auto outcome = s3_client_.ListObjectsV2(request); if (!outcome.IsSuccess()) { return tl::make_unexpected(fmt::format( - "ListObjects error: {}", outcome.GetError().GetMessage())); + "ListObjectsV2 error: {}", outcome.GetError().GetMessage())); } const auto &result = outcome.GetResult(); @@ -843,8 +806,13 @@ tl::expected S3Helper::ListObjectsWithPrefix( // Check if there are more objects to fetch if (result.GetIsTruncated()) { - // Set marker to get next page - request.WithMarker(result.GetNextMarker()); + const auto &next_token = result.GetNextContinuationToken(); + if (next_token.empty()) { + return tl::make_unexpected( + "ListObjectsV2 error: truncated response missing next " + "continuation token"); + } + request.SetContinuationToken(next_token); } else { done = true; } @@ -872,4 +840,4 @@ tl::expected S3Helper::DeleteObjectsWithPrefix( return DeleteObjects(object_keys); } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/src/utils/type_util.cpp b/mooncake-store/src/utils/type_util.cpp deleted file mode 100644 index 7cfcd933b7..0000000000 --- a/mooncake-store/src/utils/type_util.cpp +++ /dev/null @@ -1,48 +0,0 @@ -#include "utils/type_util.h" - -#include -#include -#include -#include - -namespace mooncake { - -bool TypeUtil::ParseBool(std::string_view value, bool& out) { - std::string lower; - lower.reserve(value.size()); - for (char ch : value) { - lower.push_back( - static_cast(std::tolower(static_cast(ch)))); - } - if (lower == "1" || lower == "true" || lower == "yes") { - out = true; - return true; - } - if (lower == "0" || lower == "false" || lower == "no") { - out = false; - return true; - } - return false; -} - -bool TypeUtil::ParseUint64(std::string_view value, uint64_t& out) { - const char* begin = value.data(); - const char* end = begin + value.size(); - if (begin == end) { - return false; - } - auto result = std::from_chars(begin, end, out); - return result.ec == std::errc{} && result.ptr == end; -} - -bool TypeUtil::ParseInt64(std::string_view value, int64_t& out) { - const char* begin = value.data(); - const char* end = begin + value.size(); - if (begin == end) { - return false; - } - auto result = std::from_chars(begin, end, out); - return result.ec == std::errc{} && result.ptr == end; -} - -} // namespace mooncake diff --git a/mooncake-store/tests/CMakeLists.txt b/mooncake-store/tests/CMakeLists.txt index 1b4d8f0076..201ab04c87 100644 --- a/mooncake-store/tests/CMakeLists.txt +++ b/mooncake-store/tests/CMakeLists.txt @@ -34,10 +34,46 @@ function(add_ha_test name) endfunction() add_store_test(buffer_allocator_test buffer_allocator_test.cpp) +add_store_test(runtime_accelerator_test runtime_accelerator_test.cpp) +add_store_test(registered_pinned_memory_test registered_pinned_memory_test.cpp) add_store_test(allocation_strategy_test allocation_strategy_test.cpp) +add_store_test(replica_selection_test replica_selection_test.cpp) +add_test( + NAME replica_selection_env_opt_in_test + COMMAND replica_selection_test + --gtest_filter=ReplicaSelectionTest.EnvironmentOptInUsesBuiltinScorer) +set_tests_properties(replica_selection_env_opt_in_test + PROPERTIES ENVIRONMENT "MC_STORE_REPLICA_SCORING=1") add_store_test(eviction_strategy_test eviction_strategy_test.cpp) add_store_test(deadline_scheduler_test deadline_scheduler_test.cpp) +add_store_test(kv_event_publisher_test kv_event_publisher_test.cpp) +if(ENABLE_KV_EVENTS) + find_library( + KV_EVENT_TEST_ZMQ_LIBRARY + NAMES zmq libzmq + PATHS /usr/lib /usr/local/lib /usr/lib64) + find_path( + KV_EVENT_TEST_ZMQ_INCLUDE_DIR + NAMES zmq.h + PATHS /usr/include /usr/local/include) + target_compile_definitions(kv_event_publisher_test + PRIVATE MOONCAKE_ENABLE_KV_EVENTS=1) + if(KV_EVENT_TEST_ZMQ_LIBRARY AND KV_EVENT_TEST_ZMQ_INCLUDE_DIR) + target_include_directories(kv_event_publisher_test + PRIVATE ${KV_EVENT_TEST_ZMQ_INCLUDE_DIR}) + target_link_libraries(kv_event_publisher_test + PRIVATE ${KV_EVENT_TEST_ZMQ_LIBRARY}) + endif() +endif() add_store_test(master_service_test master_service_test.cpp) +add_store_test(batch_evict_test batch_evict_test.cpp) +add_store_test(master_scenario_test master_scenario.cpp + master_scenario_test.cpp) +add_store_test(master_service_scenario_test master_scenario.cpp + master_service_scenario_test.cpp) +add_store_test(master_service_config_test master_service_config_test.cpp) +add_store_test(master_service_processing_key_double_erase_test + master_service_processing_key_double_erase_test.cpp) add_store_test(master_service_tenant_quota_test master_service_tenant_quota_test.cpp) add_store_test(batch_remove_test batch_remove_test.cpp) @@ -53,6 +89,14 @@ if(USE_NOF) add_store_test(nof_heartbeat_test nof_heartbeat_test.cpp) endif() add_store_test(client_integration_test client_integration_test.cpp) +add_test( + NAME object_checksum_client_test + COMMAND + client_integration_test + "--gtest_filter=ObjectChecksumTest.*:ClientIntegrationTest.ObjectChecksumRejectsCorruptedObject:ClientIntegrationTest.BatchPutPreservesObjectChecksumPairing" +) +set_tests_properties(object_checksum_client_test + PROPERTIES ENVIRONMENT "MOONCAKE_STORE_CHECKSUM=1") add_store_test(rpc_timeout_test rpc_timeout_test.cpp) if(USE_CXL) add_store_test(cxl_client_integration_test cxl_client_integration_test.cpp) @@ -62,7 +106,21 @@ add_store_test(master_admin_server_test master_admin_server_test.cpp) add_store_test(posix_file_test posix_file_test.cpp) add_store_test(thread_pool_test thread_pool_test.cpp) add_store_test(transfer_task_test transfer_task_test.cpp) +if(USE_CUDA) + find_package(CUDAToolkit REQUIRED) + target_compile_definitions(transfer_task_test PRIVATE USE_CUDA) + target_link_libraries(transfer_task_test PRIVATE CUDA::cudart) + if(USE_TENT) + add_test( + NAME transfer_scatter_tent_test + COMMAND transfer_task_test + --gtest_filter=TransferTaskTest.TransferScatterWritesGpuDestinationDirectly) + set_tests_properties(transfer_scatter_tent_test + PROPERTIES ENVIRONMENT "MC_USE_TENT=1") + endif() +endif() add_store_test(tenant_quota_test tenant_quota_test.cpp) +add_store_test(tenant_id_test tenant_id_test.cpp) add_store_test(segment_test segment_test.cpp) add_store_test(offset_allocator_test offset_allocator_test.cpp) add_store_test(utils_test utils_test.cpp) @@ -71,6 +129,7 @@ add_store_test(client_local_hot_cache_test client_local_hot_cache_test.cpp) add_store_test(client_tcp_local_memcpy_test client_tcp_local_memcpy_test.cpp) add_store_test(pybind_client_test pybind_client_test.cpp) add_store_test(ipv6_client_test ipv6_client_test.cpp) +add_store_test(host_port_fix_test host_port_fix_test.cpp) add_store_test(client_metrics_test client_metrics_test.cpp) add_store_test(ssd_metrics_test ssd_metrics_test.cpp) add_store_test(serializer_test serializer_test.cpp) @@ -85,16 +144,28 @@ add_store_test( add_store_test(file_util_test file_util_test.cpp) add_store_test(snapshot_child_process_test ha/snapshot/snapshot_child_process_test.cpp) +add_store_test(master_snapshot_codec_test + ha/snapshot/master_snapshot_codec_test.cpp) add_store_test(master_service_test_for_snapshot ha/snapshot/master_service_test_for_snapshot.cpp) add_store_test(non_ha_reconnect_test non_ha_reconnect_test.cpp) add_store_test(storage_backend_test storage_backend_test.cpp) +add_store_test(client_storage_backend_test client_storage_backend_test.cpp) add_store_test(mutex_test mutex_test.cpp) add_store_test(file_storage_test file_storage_test.cpp) add_store_test(task_manager_test task_manager_test.cpp) add_store_test(task_executor_test task_executor_test.cpp) add_store_test(task_integration_test task_integration_test.cpp) add_store_test(dummy_client_get_buffer_test dummy_client_get_buffer_test.cpp) +add_test( + NAME dummy_client_object_checksum_test + COMMAND + dummy_client_get_buffer_test + "--gtest_filter=DummyClientGetBufferTest.GetIntoRejectsCorruptedObjectWithChecksum:DummyClientGetBufferTest.BatchQueryPreservesObjectChecksum" +) +set_tests_properties(dummy_client_object_checksum_test + PROPERTIES ENVIRONMENT "MOONCAKE_STORE_CHECKSUM=1") +add_store_test(uds_transport_test uds_transport_test.cpp) add_store_test(health_check_test health_check_test.cpp) add_store_test(http_metadata_server_test http_metadata_server_test.cpp) add_store_test(mmap_arena_test mmap_arena_test.cpp) @@ -160,6 +231,14 @@ if(STORE_USE_ETCD endif() add_ha_test(ha_backend_availability_test ha/leadership/ha_backend_availability_test.cpp) +add_ha_test(leader_label_reconciler_test + ha/leadership/leader_label_reconciler_test.cpp) +add_ha_test(oplog_batch_storage_test ha/oplog/oplog_batch_storage_test.cpp) +add_ha_test(oplog_batch_auditor_test ha/oplog/oplog_batch_auditor_test.cpp) +target_sources(oplog_batch_auditor_test + PRIVATE ../tools/oplog_batch_auditor.cpp) +target_include_directories(oplog_batch_auditor_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/..) add_executable(stress_workload_test stress_workload_test.cpp) target_link_libraries( @@ -181,14 +260,26 @@ add_ha_test(hot_standby_service_test ha/standby/hot_standby_service_test.cpp) add_ha_test(hot_standby_snapshot_bootstrap_test ha/standby/hot_standby_snapshot_bootstrap_test.cpp) add_ha_test(oplog_applier_test ha/oplog/oplog_applier_test.cpp) -add_ha_test(oplog_manager_test ha/oplog/oplog_manager_test.cpp) -add_ha_test(etcd_oplog_store_test ha/oplog/etcd_oplog_store_test.cpp) -add_ha_test(oplog_replicator_test ha/oplog/oplog_replicator_test.cpp) -add_ha_test(oplog_serializer_test ha/oplog/oplog_serializer_test.cpp) -add_ha_test(localfs_oplog_store_test ha/oplog/localfs_oplog_store_test.cpp) -add_ha_test(ha_recovery_test ha/oplog/ha_recovery_test.cpp) -add_ha_test(localfs_hot_standby_integration_test - ha/oplog/localfs_hot_standby_integration_test.cpp) +add_ha_test(oplog_types_test ha/oplog/oplog_types_test.cpp) +add_ha_test(oplog_batch_codec_test ha/oplog/oplog_batch_codec_test.cpp) +add_ha_test(oplog_batch_standby_reader_test + ha/oplog/oplog_batch_standby_reader_test.cpp) +add_ha_test(ordered_oplog_writer_test ha/oplog/ordered_oplog_writer_test.cpp) +if(MOONCAKE_ENABLE_OPLOG_PERF_METRICS) + target_compile_definitions(ha_metric_manager_test + PRIVATE MOONCAKE_ENABLE_OPLOG_PERF_METRICS) + target_compile_definitions(ordered_oplog_writer_test + PRIVATE MOONCAKE_ENABLE_OPLOG_PERF_METRICS) +endif() +add_executable( + oplog_test_failpoint_test ha/oplog/oplog_test_failpoint_test.cpp + ../src/ha/oplog/oplog_test_failpoint.cpp) +target_compile_definitions(oplog_test_failpoint_test + PRIVATE MOONCAKE_ENABLE_TEST_FAILPOINTS) +target_link_libraries(oplog_test_failpoint_test PRIVATE glog gtest gtest_main + pthread) +add_test(NAME oplog_test_failpoint_test COMMAND oplog_test_failpoint_test) +add_ha_test(master_service_ha_test ha/master_service_ha_test.cpp) add_ha_test(catalog_backed_snapshot_provider_test ha/snapshot/catalog_backed_snapshot_provider_test.cpp) if(STORE_USE_REDIS) diff --git a/mooncake-store/tests/allocation_strategy_test.cpp b/mooncake-store/tests/allocation_strategy_test.cpp index 293e57f081..7129244498 100644 --- a/mooncake-store/tests/allocation_strategy_test.cpp +++ b/mooncake-store/tests/allocation_strategy_test.cpp @@ -89,6 +89,9 @@ INSTANTIATE_TEST_SUITE_P( case AllocationStrategyType::FREE_RATIO_FIRST: strategy_str = "FreeRatioFirst"; break; + case AllocationStrategyType::SSD_FREE_RATIO_FIRST: + strategy_str = "SsdFreeRatioFirst"; + break; default: strategy_str = "Unknown"; } @@ -752,4 +755,255 @@ TEST_F(AllocationStrategyTest, PerformanceTest) { // public API. The functionality is now encapsulated within the Allocate() // method. +// Mock SsdMetricsProvider for testing SSD-aware allocation +class MockSsdMetricsProvider : public SsdMetricsProvider { + public: + std::unordered_map total_capacity; + std::unordered_map used_bytes; + + int64_t getSsdTotalCapacity(const std::string& name) const override { + auto it = total_capacity.find(name); + return it != total_capacity.end() ? it->second : 0; + } + + int64_t getSsdUsedBytes(const std::string& name) const override { + auto it = used_bytes.find(name); + return it != used_bytes.end() ? it->second : 0; + } +}; + +TEST_F(AllocationStrategyTest, SsdFreeRatioFirstChoosesHighestFreeRatio) { + auto ssd_strategy = std::make_unique(); + + const int kNumSegments = 3; + const size_t kSegmentSize = 64 * MiB; + + AllocatorManager allocator_manager; + for (int i = 0; i < kNumSegments; i++) { + const auto name = std::to_string(i) + "-segment"; + allocator_manager.addAllocator( + name, + std::make_shared( + name, 0x100000000ULL + i * kSegmentSize, kSegmentSize, name)); + } + + // SSD free ratios: segment 0 = 20%, segment 1 = 60%, segment 2 = 90% + MockSsdMetricsProvider ssd_provider; + for (int i = 0; i < kNumSegments; i++) { + const auto name = std::to_string(i) + "-segment"; + ssd_provider.total_capacity[name] = 1000 * MiB; + } + ssd_provider.used_bytes["0-segment"] = 800 * MiB; // 20% free + ssd_provider.used_bytes["1-segment"] = 400 * MiB; // 60% free + ssd_provider.used_bytes["2-segment"] = 100 * MiB; // 90% free + + auto result = + ssd_strategy->Allocate(allocator_manager, 64 * 1024, 1, {}, {}, + ReplicaType::MEMORY, &ssd_provider); + ASSERT_TRUE(result.has_value()); + ASSERT_EQ(result.value().size(), 1u); + + const auto& replica = result.value()[0]; + auto descriptor = replica.get_descriptor(); + ASSERT_TRUE(descriptor.is_memory_replica()); + const auto& mem_desc = descriptor.get_memory_descriptor(); + EXPECT_EQ(mem_desc.buffer_descriptor.transport_endpoint_, "2-segment"); +} + +// Test that SsdFreeRatioFirstAllocationStrategy works without an +// SsdMetricsProvider (delegates to base class random allocation). +TEST_F(AllocationStrategyTest, + SsdFreeRatioFirstWithoutMetricsProviderAllocates) { + auto ssd_strategy = std::make_unique(); + + auto allocator = std::make_shared( + "segment1", 0x100000000ULL, 64 * MiB, "segment1"); + AllocatorManager allocator_manager; + allocator_manager.addAllocator("segment1", allocator); + + auto result = ssd_strategy->Allocate(allocator_manager, 1024, 1, {}, {}); + ASSERT_TRUE(result.has_value()); + EXPECT_EQ(result.value().size(), 1u); +} + +// Test that SsdFreeRatioFirstAllocationStrategy skips excluded segments even +// when they have the highest SSD free ratio. +TEST_F(AllocationStrategyTest, SsdFreeRatioFirstExcludedSegmentsSkipped) { + auto ssd_strategy = std::make_unique(); + + const int kNumSegments = 3; + const size_t kSegmentSize = 64 * MiB; + + AllocatorManager allocator_manager; + for (int i = 0; i < kNumSegments; i++) { + const auto name = std::to_string(i) + "-segment"; + allocator_manager.addAllocator( + name, + std::make_shared( + name, 0x100000000ULL + i * kSegmentSize, kSegmentSize, name)); + } + + // SSD free ratios: 0-segment=95%, 1-segment=50%, 2-segment=30% + MockSsdMetricsProvider ssd_provider; + for (int i = 0; i < kNumSegments; i++) { + const auto name = std::to_string(i) + "-segment"; + ssd_provider.total_capacity[name] = 1000 * MiB; + } + ssd_provider.used_bytes["0-segment"] = 50 * MiB; // 95% free (highest) + ssd_provider.used_bytes["1-segment"] = 500 * MiB; // 50% free + ssd_provider.used_bytes["2-segment"] = 700 * MiB; // 30% free + + // Exclude 0-segment which has the highest SSD free ratio + std::set excluded = {"0-segment"}; + + auto result = + ssd_strategy->Allocate(allocator_manager, 64 * 1024, 1, {}, excluded, + ReplicaType::MEMORY, &ssd_provider); + ASSERT_TRUE(result.has_value()); + ASSERT_EQ(result.value().size(), 1u); + + const auto& replica = result.value()[0]; + auto descriptor = replica.get_descriptor(); + ASSERT_TRUE(descriptor.is_memory_replica()); + const auto& mem_desc = descriptor.get_memory_descriptor(); + // Must NOT be allocated to the excluded segment despite its high free ratio + EXPECT_NE(mem_desc.buffer_descriptor.transport_endpoint_, "0-segment"); +} + +// Test that ssd_used_bytes > ssd_total_capacity is clamped to 0% free, +// so the other segment (with normal usage) is preferred. +TEST_F(AllocationStrategyTest, SsdFreeRatioFirstUsedExceedsTotalIsClamped) { + auto ssd_strategy = std::make_unique(); + + const size_t kSegmentSize = 64 * MiB; + + AllocatorManager allocator_manager; + allocator_manager.addAllocator( + "0-segment", + std::make_shared("0-segment", 0x100000000ULL, + kSegmentSize, "0-segment")); + allocator_manager.addAllocator( + "1-segment", std::make_shared( + "1-segment", 0x100000000ULL + kSegmentSize, + kSegmentSize, "1-segment")); + + // 0-segment: used exceeds total (concurrent drift) → clamped to 0% free + // 1-segment: used=100, total=1000 → 90% free + MockSsdMetricsProvider ssd_provider; + ssd_provider.total_capacity["0-segment"] = 1000 * MiB; + ssd_provider.total_capacity["1-segment"] = 1000 * MiB; + ssd_provider.used_bytes["0-segment"] = 1500 * MiB; // exceeds total + ssd_provider.used_bytes["1-segment"] = 100 * MiB; // 90% free + + auto result = + ssd_strategy->Allocate(allocator_manager, 64 * 1024, 1, {}, {}, + ReplicaType::MEMORY, &ssd_provider); + ASSERT_TRUE(result.has_value()); + ASSERT_EQ(result.value().size(), 1u); + + const auto& replica = result.value()[0]; + auto descriptor = replica.get_descriptor(); + ASSERT_TRUE(descriptor.is_memory_replica()); + const auto& mem_desc = descriptor.get_memory_descriptor(); + // 0-segment is clamped to 0% free; 1-segment is 90% free → pick 1-segment + EXPECT_EQ(mem_desc.buffer_descriptor.transport_endpoint_, "1-segment"); +} + +// Strategy-level performance comparison: Random vs SsdFreeRatioFirst. +// Uses OffsetBufferAllocator (real physical memory via aligned_alloc) and +// MockSsdMetricsProvider. Measures pure strategy overhead without MasterService +// locking, reflecting only the allocation algorithm cost. +TEST_F(AllocationStrategyTest, SsdFreeRatioFirstVsRandomStrategyPerformance) { + constexpr size_t kNumSegments = 64; + constexpr size_t kSegmentSize = 8 * MiB; // 64 * 8MB = 512MB physical + constexpr size_t kAllocSize = 128 * 1024; // 128KB per allocation + constexpr int kWarmupRounds = 200; + constexpr int kBenchmarkRounds = 2000; + + // Build AllocatorManager: 64 segments, each 8 MiB + AllocatorManager allocator_manager; + for (size_t i = 0; i < kNumSegments; i++) { + const auto name = "perf_seg_" + std::to_string(i); + allocator_manager.addAllocator( + name, + std::make_shared( + name, 0x200000000ULL + i * kSegmentSize, kSegmentSize, name)); + } + + // SSD free ratios vary uniformly from ~10% to ~90% across segments + MockSsdMetricsProvider ssd_provider; + for (size_t i = 0; i < kNumSegments; i++) { + const auto name = "perf_seg_" + std::to_string(i); + ssd_provider.total_capacity[name] = 1000 * MiB; + // used: 100 MiB (10% used) to 900 MiB (90% used) + ssd_provider.used_bytes[name] = + static_cast((100 + i * 800 / kNumSegments)) * MiB; + } + + // -------- Random strategy -------- + auto random_strategy = std::make_unique(); + { + // Warmup (results discarded) + for (int i = 0; i < kWarmupRounds; i++) { + (void)random_strategy->Allocate(allocator_manager, kAllocSize); + } + } + std::vector> random_replicas; + random_replicas.reserve(kBenchmarkRounds); + + auto t_rand_start = std::chrono::steady_clock::now(); + for (int i = 0; i < kBenchmarkRounds; i++) { + auto r = random_strategy->Allocate(allocator_manager, kAllocSize); + ASSERT_TRUE(r.has_value()); + random_replicas.emplace_back(std::move(r.value())); + } + auto rand_us = std::chrono::duration_cast( + std::chrono::steady_clock::now() - t_rand_start); + random_replicas.clear(); // deallocate so SSD strategy starts fresh + + // -------- SsdFreeRatioFirst strategy -------- + auto ssd_strategy = std::make_unique(); + { + // Warmup + for (int i = 0; i < kWarmupRounds; i++) { + (void)ssd_strategy->Allocate(allocator_manager, kAllocSize, 1, {}, + {}, ReplicaType::MEMORY, + &ssd_provider); + } + } + std::vector> ssd_replicas; + ssd_replicas.reserve(kBenchmarkRounds); + + auto t_ssd_start = std::chrono::steady_clock::now(); + for (int i = 0; i < kBenchmarkRounds; i++) { + auto r = ssd_strategy->Allocate(allocator_manager, kAllocSize, 1, {}, + {}, ReplicaType::MEMORY, &ssd_provider); + ASSERT_TRUE(r.has_value()); + ssd_replicas.emplace_back(std::move(r.value())); + } + auto ssd_us = std::chrono::duration_cast( + std::chrono::steady_clock::now() - t_ssd_start); + ssd_replicas.clear(); + + double rand_us_per_op = + static_cast(rand_us.count()) / kBenchmarkRounds; + double ssd_us_per_op = + static_cast(ssd_us.count()) / kBenchmarkRounds; + double overhead_ratio = + static_cast(ssd_us.count()) / rand_us.count(); + + std::cout + << "\n=== Strategy-Level Performance: Random vs SsdFreeRatioFirst ===\n" + << "Segments: " << kNumSegments + << " | Alloc size: " << (kAllocSize / 1024) << " KB" + << " | Rounds: " << kBenchmarkRounds << "\n" + << "Random: " << rand_us.count() << " us total | " + << std::fixed << std::setprecision(3) << rand_us_per_op << " us/op\n" + << "SsdFreeRatioFirst: " << ssd_us.count() << " us total | " + << ssd_us_per_op << " us/op\n" + << "Overhead ratio: " << std::setprecision(2) << overhead_ratio + << "x (" << std::setprecision(1) << (overhead_ratio - 1.0) * 100.0 + << "% slower)\n\n"; +} + } // namespace mooncake diff --git a/mooncake-store/tests/batch_evict_test.cpp b/mooncake-store/tests/batch_evict_test.cpp new file mode 100644 index 0000000000..9086b63fc1 --- /dev/null +++ b/mooncake-store/tests/batch_evict_test.cpp @@ -0,0 +1,467 @@ +#include +#include + +#include +#include +#include +#include +#include + +#include "master_service.h" +#include "mutex.h" +#include "types.h" +#include "utils.h" + +namespace mooncake::test { + +// Deterministic correctness tests for MasterService::BatchEvict. +// +// These tests drive BatchEvict directly instead of filling a segment and +// waiting for the background eviction thread. Lease timestamps are written +// explicitly so that the eviction order, the exact eviction count, the +// soft-pin fallback and the whole-group semantics are all observable without +// sleeps, background threads or timing assumptions. +class BatchEvictTest : public ::testing::Test { + protected: + static constexpr const char* kSegmentName = "batch_evict_test_segment"; + static constexpr size_t kSegmentBase = 0x500000000; + static constexpr size_t kSegmentSize = 256ULL * 1024 * 1024; + static constexpr uint64_t kObjectSize = 1024; + + void SetUp() override { + google::InitGoogleLogging("BatchEvictTest"); + FLAGS_logtostderr = true; + } + + void TearDown() override { google::ShutdownGoogleLogging(); } + + // Eviction ratio 0 and a 100% high watermark keep the background eviction + // thread from interfering: every eviction in these tests is the one the + // test itself requests. + static MasterServiceConfig MakeConfig(bool allow_soft_pin_eviction) { + return MasterServiceConfig::builder() + .set_memory_allocator(BufferAllocatorType::OFFSET) + .set_default_kv_lease_ttl(0) + .set_default_kv_soft_pin_ttl(60 * 60 * 1000) + .set_allow_evict_soft_pinned_objects(allow_soft_pin_eviction) + .set_eviction_ratio(0.0) + .set_eviction_high_watermark_ratio(1.0) + .set_client_live_ttl_sec(3600) + .build(); + } + + static Segment MakeSegment() { + Segment segment; + segment.id = generate_uuid(); + segment.name = kSegmentName; + segment.base = kSegmentBase; + segment.size = kSegmentSize; + segment.te_endpoint = segment.name; + return segment; + } + + static UUID MountSegment(MasterService& service) { + const UUID client_id = generate_uuid(); + auto result = service.MountSegment(MakeSegment(), client_id); + EXPECT_TRUE(result.has_value()); + return client_id; + } + + static std::string Key(size_t index) { + return "batch_evict_key_" + std::to_string(index); + } + + static void PutObject(MasterService& service, const UUID& client_id, + const std::string& key, bool with_soft_pin = false, + const std::string& group_id = std::string()) { + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segment = kSegmentName; + config.with_soft_pin = with_soft_pin; + if (!group_id.empty()) { + config.group_ids = std::vector{group_id}; + } + + auto put_start = service.PutStart(client_id, key, TenantId::Default(), + kObjectSize, config); + ASSERT_TRUE(put_start.has_value()) + << "PutStart failed for key=" << key + << ", error=" << toString(put_start.error()); + ASSERT_TRUE(service + .PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); + } + + // Ungrouped objects are sharded by key hash, but grouped objects are + // routed by their group id, so the direct shard index is not always the + // one holding the key. Fall back to a shard scan in that case. + template + static bool WithMetadata(MasterService& service, const std::string& key, + Fn&& fn) { + const TenantId tenant = TenantId::Default(); + + auto try_shard = [&](size_t shard_idx) -> bool { + MasterService::MetadataShardAccessorRW shard(&service, shard_idx); + auto tenant_it = shard->tenants.find(tenant); + if (tenant_it == shard->tenants.end()) { + return false; + } + auto metadata_it = tenant_it->second.metadata.find(key); + if (metadata_it == tenant_it->second.metadata.end()) { + return false; + } + fn(metadata_it->second); + return true; + }; + + if (try_shard(service.getMetadataShardIndex(tenant, key))) { + return true; + } + for (size_t shard_idx = 0; shard_idx < MasterService::kNumShards; + ++shard_idx) { + if (try_shard(shard_idx)) { + return true; + } + } + ADD_FAILURE() << "metadata not found for key=" << key; + return false; + } + + static void SetLease(MasterService& service, const std::string& key, + std::chrono::system_clock::time_point lease_timeout, + std::optional + soft_pin_timeout = std::nullopt) { + WithMetadata(service, key, + [&](MasterService::ObjectMetadata& metadata) { + SpinLocker locker(&metadata.lock); + metadata.lease_timeout = lease_timeout; + metadata.soft_pin_timeout = soft_pin_timeout; + }); + } + + static bool Exists(MasterService& service, const std::string& key) { + auto result = service.ExistKey(key, TenantId::Default()); + if (!result.has_value()) { + ADD_FAILURE() << "ExistKey failed for key=" << key + << ", error=" << toString(result.error()); + return false; + } + return result.value(); + } + + static void RunBatchEvict(MasterService& service, double target, + double lowerbound) { + service.BatchEvict(target, lowerbound); + } + + static std::chrono::system_clock::time_point ExpiredBase() { + return std::chrono::system_clock::now() - std::chrono::hours(1); + } + + // Populates `count` expired objects whose lease timestamps are strictly + // increasing, so Key(i) is always older than Key(i + 1) and no timestamp + // ties exist at the eviction boundary. + static void PopulateOldestFirst(MasterService& service, + const UUID& client_id, size_t count) { + const auto base = ExpiredBase(); + for (size_t i = 0; i < count; ++i) { + PutObject(service, client_id, Key(i)); + SetLease(service, Key(i), base + std::chrono::nanoseconds(i)); + } + } + // Builds a population whose oldest `blocked_count` candidates belong to a + // group that also holds one member under an active lease. Those candidates + // pass the census — their own lease has expired — but cannot be evicted + // during execution, because a group is only evictable once every member's + // lease has expired. That is an execution-stage failure without any hook: + // it exercises the same recovery path as metadata churn between the census + // and the eviction pass. + struct BlockedGroupPopulation { + size_t blocked_count; + size_t keeper_index; + size_t plain_begin; + size_t plain_count; + size_t total_objects; + }; + + static BlockedGroupPopulation PopulateBlockedGroup( + MasterService& service, const UUID& client_id, + const std::string& group_id, size_t blocked_count, size_t plain_count) { + const auto base = ExpiredBase(); + const auto active_lease = + std::chrono::system_clock::now() + std::chrono::hours(1); + + // Oldest leases: the blocked group members are always selected first. + for (size_t i = 0; i < blocked_count; ++i) { + PutObject(service, client_id, Key(i), /*with_soft_pin=*/false, + group_id); + SetLease(service, Key(i), base + std::chrono::nanoseconds(i)); + } + + // The keeper shares the group but holds an active lease, so no member + // of the group can be evicted. + const size_t keeper_index = blocked_count; + PutObject(service, client_id, Key(keeper_index), + /*with_soft_pin=*/false, group_id); + SetLease(service, Key(keeper_index), active_lease); + + // Plain objects are strictly newer than every blocked member. + const size_t plain_begin = keeper_index + 1; + for (size_t i = 0; i < plain_count; ++i) { + const size_t index = plain_begin + i; + PutObject(service, client_id, Key(index)); + SetLease(service, Key(index), + base + std::chrono::nanoseconds(index)); + } + + return {blocked_count, keeper_index, plain_begin, plain_count, + plain_begin + plain_count}; + } + + static size_t CountAlive(MasterService& service, size_t begin, size_t end) { + size_t alive = 0; + for (size_t i = begin; i < end; ++i) { + if (Exists(service, Key(i))) { + ++alive; + } + } + return alive; + } +}; + +// Oldest-first: with distinct lease timestamps the evicted set must be exactly +// the oldest ceil(N * ratio) objects, and nothing newer. +TEST_F(BatchEvictTest, EvictsExactOldestObjectsAtLowRatio) { + constexpr size_t kObjectCount = 400; + constexpr size_t kExpectedEvicted = 20; // ceil(400 * 0.05) + + MasterService service(MakeConfig(/*allow_soft_pin_eviction=*/false)); + const UUID client_id = MountSegment(service); + PopulateOldestFirst(service, client_id, kObjectCount); + ASSERT_EQ(service.GetKeyCount(), kObjectCount); + + RunBatchEvict(service, /*target=*/0.05, /*lowerbound=*/0.05); + + EXPECT_EQ(service.GetKeyCount(), kObjectCount - kExpectedEvicted); + for (size_t i = 0; i < kObjectCount; ++i) { + EXPECT_EQ(Exists(service, Key(i)), i >= kExpectedEvicted) + << "unexpected eviction outcome at index=" << i; + } +} + +// target == lowerbound: the first pass already satisfies the lower bound, so +// the second pass must not evict anything extra. +TEST_F(BatchEvictTest, TargetEqualsLowerBoundEvictsExactCount) { + constexpr size_t kObjectCount = 250; + constexpr size_t kExpectedEvicted = 25; // ceil(250 * 0.10) + + MasterService service(MakeConfig(/*allow_soft_pin_eviction=*/false)); + const UUID client_id = MountSegment(service); + PopulateOldestFirst(service, client_id, kObjectCount); + ASSERT_EQ(service.GetKeyCount(), kObjectCount); + + RunBatchEvict(service, /*target=*/0.10, /*lowerbound=*/0.10); + + EXPECT_EQ(service.GetKeyCount(), kObjectCount - kExpectedEvicted); + EXPECT_FALSE(Exists(service, Key(kExpectedEvicted - 1))); + EXPECT_TRUE(Exists(service, Key(kExpectedEvicted))); +} + +// Soft-pin fallback: unpinned objects go first; soft-pinned objects are only +// evicted by the second pass, oldest first, and only up to the lower bound. +TEST_F(BatchEvictTest, SoftPinnedEvictedOnlyAfterUnpinned) { + constexpr size_t kNoPinCount = 10; + constexpr size_t kSoftPinCount = 10; + // ceil(20 * 0.80) == 16; the first pass can only evict the 10 unpinned + // objects, leaving 6 for the soft-pinned second pass. + constexpr size_t kExpectedSoftPinEvicted = 6; + + MasterService service(MakeConfig(/*allow_soft_pin_eviction=*/true)); + const UUID client_id = MountSegment(service); + + const auto base = ExpiredBase(); + const auto active_soft_pin = + std::chrono::system_clock::now() + std::chrono::hours(1); + + for (size_t i = 0; i < kNoPinCount; ++i) { + PutObject(service, client_id, Key(i)); + SetLease(service, Key(i), base + std::chrono::nanoseconds(i)); + } + for (size_t i = 0; i < kSoftPinCount; ++i) { + const size_t index = kNoPinCount + i; + PutObject(service, client_id, Key(index), /*with_soft_pin=*/true); + SetLease(service, Key(index), base + std::chrono::nanoseconds(index), + active_soft_pin); + } + ASSERT_EQ(service.GetKeyCount(), kNoPinCount + kSoftPinCount); + + RunBatchEvict(service, /*target=*/0.80, /*lowerbound=*/0.80); + + for (size_t i = 0; i < kNoPinCount; ++i) { + EXPECT_FALSE(Exists(service, Key(i))) + << "unpinned object survived at index=" << i; + } + for (size_t i = 0; i < kSoftPinCount; ++i) { + const size_t index = kNoPinCount + i; + EXPECT_EQ(Exists(service, Key(index)), i >= kExpectedSoftPinEvicted) + << "unexpected soft-pin outcome at index=" << index; + } + EXPECT_EQ(service.GetKeyCount(), kSoftPinCount - kExpectedSoftPinEvicted); +} + +// Whole-group: selecting one group member evicts the entire group as a unit, +// even when the target only calls for a single object. +TEST_F(BatchEvictTest, WholeGroupEvictedTogether) { + constexpr size_t kObjectCount = 10; + constexpr size_t kGroupSize = 3; + const std::string group_id = "batch_evict_test_group"; + + MasterService service(MakeConfig(/*allow_soft_pin_eviction=*/false)); + const UUID client_id = MountSegment(service); + + const auto base = ExpiredBase(); + for (size_t i = 0; i < kObjectCount; ++i) { + PutObject(service, client_id, Key(i), /*with_soft_pin=*/false, + i < kGroupSize ? group_id : std::string()); + SetLease(service, Key(i), base + std::chrono::nanoseconds(i)); + } + + // ceil(10 * 0.10) == 1 and the oldest candidate is a group member, so the + // whole group is evicted even though only one object was requested. + RunBatchEvict(service, /*target=*/0.10, /*lowerbound=*/0.10); + + EXPECT_EQ(service.GetKeyCount(), kObjectCount - kGroupSize); + for (size_t i = 0; i < kGroupSize; ++i) { + EXPECT_FALSE(Exists(service, Key(i))) + << "group member survived at index=" << i; + } + for (size_t i = kGroupSize; i < kObjectCount; ++i) { + EXPECT_TRUE(Exists(service, Key(i))) + << "non-group object evicted at index=" << i; + } +} + +// Whole-group safety: a single group member under an active lease keeps every +// member of that group resident. +TEST_F(BatchEvictTest, UnexpiredGroupMemberBlocksWholeGroup) { + constexpr size_t kObjectCount = 10; + constexpr size_t kGroupSize = 3; + const std::string group_id = "batch_evict_test_blocked_group"; + + MasterService service(MakeConfig(/*allow_soft_pin_eviction=*/false)); + const UUID client_id = MountSegment(service); + + const auto base = ExpiredBase(); + const auto unexpired = + std::chrono::system_clock::now() + std::chrono::hours(1); + + for (size_t i = 0; i < kObjectCount; ++i) { + PutObject(service, client_id, Key(i), /*with_soft_pin=*/false, + i < kGroupSize ? group_id : std::string()); + SetLease(service, Key(i), base + std::chrono::nanoseconds(i)); + } + // Hold one member of the group under an active lease. + SetLease(service, Key(kGroupSize - 1), unexpired); + + RunBatchEvict(service, /*target=*/0.10, /*lowerbound=*/0.10); + + for (size_t i = 0; i < kGroupSize; ++i) { + EXPECT_TRUE(Exists(service, Key(i))) + << "blocked group member was evicted at index=" << i; + } + // The blocked group yields nothing, so exactly one ungrouped object is + // evicted instead. + EXPECT_EQ(service.GetKeyCount(), kObjectCount - 1); +} + +// High ratio: the same oldest-first and exact-count guarantees must hold when +// the requested ratio covers most of the population. +TEST_F(BatchEvictTest, HighRatioEvictsExactOldestCount) { + constexpr size_t kObjectCount = 200; + constexpr size_t kExpectedEvicted = 160; // ceil(200 * 0.80) + + MasterService service(MakeConfig(/*allow_soft_pin_eviction=*/false)); + const UUID client_id = MountSegment(service); + PopulateOldestFirst(service, client_id, kObjectCount); + ASSERT_EQ(service.GetKeyCount(), kObjectCount); + + RunBatchEvict(service, /*target=*/0.80, /*lowerbound=*/0.80); + + EXPECT_EQ(service.GetKeyCount(), kObjectCount - kExpectedEvicted); + for (size_t i = 0; i < kObjectCount; ++i) { + EXPECT_EQ(Exists(service, Key(i)), i >= kExpectedEvicted) + << "unexpected eviction outcome at index=" << i; + } +} + +// Reserve: a small number of candidates that pass the census but fail during +// execution is absorbed by the reserve slack, so the requested target is still +// met exactly and no refill scan is needed. +TEST_F(BatchEvictTest, ReserveAbsorbsExecutionFailuresAndMeetsTarget) { + constexpr size_t kBlockedCount = 12; + constexpr size_t kPlainCount = 1200; + // 1213 objects in total, ceil(1213 * 0.05) == 61. + constexpr size_t kExpectedEvicted = 61; + // The 61 oldest candidates are the 12 blocked members plus the 49 oldest + // plain objects, so those 49 are always among the evicted set. The + // remaining 12 evictions come from the reserve, whose internal order is + // unspecified, so only the total is asserted for them. + constexpr size_t kAlwaysEvictedPlain = 49; + + MasterService service(MakeConfig(/*allow_soft_pin_eviction=*/false)); + const UUID client_id = MountSegment(service); + const auto population = + PopulateBlockedGroup(service, client_id, "batch_evict_reserve_group", + kBlockedCount, kPlainCount); + ASSERT_EQ(service.GetKeyCount(), population.total_objects); + + RunBatchEvict(service, /*target=*/0.05, /*lowerbound=*/0.05); + + EXPECT_EQ(service.GetKeyCount(), + population.total_objects - kExpectedEvicted); + EXPECT_EQ(CountAlive(service, 0, kBlockedCount), kBlockedCount) + << "blocked group members must survive"; + EXPECT_TRUE(Exists(service, Key(population.keeper_index))); + for (size_t i = 0; i < kAlwaysEvictedPlain; ++i) { + EXPECT_FALSE(Exists(service, Key(population.plain_begin + i))) + << "oldest plain object survived at offset=" << i; + } + EXPECT_EQ( + CountAlive(service, population.plain_begin, population.total_objects), + kPlainCount - kExpectedEvicted); +} + +// Refill: when every candidate inside the reserve frontier fails during +// execution the reserve is exhausted, and the refill scan must recover the +// remaining objects so the requested target is still met exactly. +TEST_F(BatchEvictTest, RefillAfterReserveExhaustionStillMeetsTarget) { + // The reserve frontier spans target + max(1024, 10% of target), so more + // than 1024 blocked candidates are required to exhaust it. + constexpr size_t kBlockedCount = 1160; + constexpr size_t kPlainCount = 200; + // 1361 objects in total, ceil(1361 * 0.05) == 69. + constexpr size_t kExpectedEvicted = 69; + + MasterService service(MakeConfig(/*allow_soft_pin_eviction=*/false)); + const UUID client_id = MountSegment(service); + const auto population = + PopulateBlockedGroup(service, client_id, "batch_evict_refill_group", + kBlockedCount, kPlainCount); + ASSERT_EQ(service.GetKeyCount(), population.total_objects); + + RunBatchEvict(service, /*target=*/0.05, /*lowerbound=*/0.05); + + // Exact target attainment is the property refill exists to protect: the + // whole frontier yielded nothing, so every eviction came from the refill. + EXPECT_EQ(service.GetKeyCount(), + population.total_objects - kExpectedEvicted); + EXPECT_EQ(CountAlive(service, 0, kBlockedCount), kBlockedCount) + << "blocked group members must survive"; + EXPECT_TRUE(Exists(service, Key(population.keeper_index))); + EXPECT_EQ( + CountAlive(service, population.plain_begin, population.total_objects), + kPlainCount - kExpectedEvicted); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/buffer_allocator_test.cpp b/mooncake-store/tests/buffer_allocator_test.cpp index 5471cf5d0c..cefaf21080 100644 --- a/mooncake-store/tests/buffer_allocator_test.cpp +++ b/mooncake-store/tests/buffer_allocator_test.cpp @@ -111,6 +111,294 @@ TEST_F(BufferAllocatorTest, AllocateMultiple) { } } +TEST_F(BufferAllocatorTest, OffsetLargestFreeRegionRemainsExact) { + constexpr size_t CAPACITY = 16 * 1024 * 1024; + auto allocator = std::make_shared( + "exact-largest-free-region", 0x140000000ULL, CAPACITY, + "exact-largest-free-region"); + + EXPECT_EQ(allocator->getLargestFreeRegion(), CAPACITY); + + auto buffer = allocator->allocate(CAPACITY / 2); + ASSERT_NE(buffer, nullptr); + const auto internal_allocator = allocator->getOffsetAllocator(); + + // Successful allocations intentionally leave the fast-fail hint high, but + // the segment-selection query must still return the authoritative value. + EXPECT_GT(internal_allocator->getLargestFreeRegion(), + allocator->getLargestFreeRegion()); + EXPECT_EQ(allocator->getLargestFreeRegion(), + internal_allocator->storageReport().largestFreeRegion); + + buffer.reset(); + EXPECT_EQ(allocator->getLargestFreeRegion(), CAPACITY); +} + +TEST_F(BufferAllocatorTest, RestoreOffsetAllocationsAtOriginalAddresses) { + constexpr uintptr_t kBase = 0x180000000ULL; + constexpr size_t kCapacity = 16 * 1024 * 1024; + const std::string segment = "restore-segment"; + const std::string endpoint = "restore-endpoint"; + + auto original = std::make_shared( + segment, kBase, kCapacity, endpoint); + auto first = original->allocate(123); + auto removed = original->allocate(5003); + auto last = original->allocate(777); + ASSERT_NE(first, nullptr); + ASSERT_NE(removed, nullptr); + ASSERT_NE(last, nullptr); + + std::vector descriptors = { + first->get_descriptor(), last->get_descriptor()}; + removed.reset(); + + auto restored = RestoreOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, descriptors); + ASSERT_TRUE(restored.has_value()); + ASSERT_EQ(restored->buffers.size(), descriptors.size()); + EXPECT_EQ(restored->buffers[0]->get_descriptor().buffer_address_, + descriptors[0].buffer_address_); + EXPECT_EQ(restored->buffers[1]->get_descriptor().buffer_address_, + descriptors[1].buffer_address_); + + auto new_buffer = restored->allocator->allocate(1024); + ASSERT_NE(new_buffer, nullptr); + const auto new_address = reinterpret_cast(new_buffer->data()); + for (const auto& descriptor : descriptors) { + EXPECT_TRUE( + new_address + new_buffer->size() <= descriptor.buffer_address_ || + descriptor.buffer_address_ + descriptor.size_ <= new_address); + } + + auto wrong_endpoint = descriptors; + wrong_endpoint[0].transport_endpoint_ = "other-endpoint"; + EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, wrong_endpoint) + .has_value()); + + auto duplicate = descriptors; + duplicate.push_back(descriptors.front()); + EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, duplicate) + .has_value()); + + auto out_of_range = descriptors; + out_of_range[0].buffer_address_ = kBase + kCapacity; + EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, out_of_range) + .has_value()); +} + +TEST_F(BufferAllocatorTest, RestoreOffsetAllocationsValidatesRangesAndOrder) { + constexpr uintptr_t kBase = 0x190000000ULL; + constexpr size_t kCapacity = 4096; + const std::string segment = "restore-validation"; + const std::string endpoint = "restore-validation-endpoint"; + auto descriptor = [&](uintptr_t address, uint64_t size) { + return AllocatedBuffer::Descriptor{size, address, "tcp", endpoint}; + }; + + std::vector unsorted = { + descriptor(kBase + 512, 64), descriptor(kBase + 128, 64)}; + auto restored = RestoreOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, unsorted); + ASSERT_TRUE(restored.has_value()); + ASSERT_EQ(restored->buffers.size(), unsorted.size()); + EXPECT_EQ(reinterpret_cast(restored->buffers[0]->data()), + unsorted[0].buffer_address_); + EXPECT_EQ(reinterpret_cast(restored->buffers[1]->data()), + unsorted[1].buffer_address_); + + std::vector overlapping = { + descriptor(kBase + 128, 100), descriptor(kBase + 200, 32)}; + EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, overlapping) + .has_value()); + + std::vector normalized_past_end = { + descriptor(kBase + kCapacity - 100, 100)}; + EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, normalized_past_end) + .has_value()); + + EXPECT_FALSE(RestoreOffsetBufferAllocator( + segment, std::numeric_limits::max() - 100, 200, + endpoint, {}) + .has_value()); + std::vector descriptor_overflow = { + descriptor(std::numeric_limits::max() - 10, 20)}; + EXPECT_FALSE(RestoreOffsetBufferAllocator(segment, kBase, kCapacity, + endpoint, descriptor_overflow) + .has_value()); +} + +TEST_F(BufferAllocatorTest, RestoredOffsetHandleReleasesItsExactAddress) { + constexpr uintptr_t kBase = 0x1A0000000ULL; + constexpr size_t kCapacity = 4096; + const std::string endpoint = "restore-release"; + std::vector descriptors = { + {64, kBase + 128, "tcp", endpoint}, {64, kBase + 512, "tcp", endpoint}}; + auto restored = RestoreOffsetBufferAllocator( + "restore-release", kBase, kCapacity, endpoint, descriptors); + ASSERT_TRUE(restored.has_value()); + + restored->buffers[0].reset(); + auto replacement = restored->allocator->allocate(64); + ASSERT_NE(replacement, nullptr); + EXPECT_EQ(reinterpret_cast(replacement->data()), + descriptors[0].buffer_address_); +} + +TEST_F(BufferAllocatorTest, RestoreOffsetAllocationsHasNoArbitraryGapLimit) { + constexpr uintptr_t kBase = 0x1B0000000ULL; + constexpr size_t kGapCount = 65537; + const std::string endpoint = "restore-many-gaps"; + std::vector descriptors; + descriptors.reserve(kGapCount); + for (size_t i = 0; i < kGapCount; ++i) { + descriptors.push_back({1, kBase + 1 + i * 2, "tcp", endpoint}); + } + + auto restored = RestoreOffsetBufferAllocator( + "restore-many-gaps", kBase, kGapCount * 2 + 1, endpoint, descriptors); + ASSERT_TRUE(restored.has_value()); + EXPECT_EQ(restored->buffers.size(), descriptors.size()); + EXPECT_EQ(reinterpret_cast(restored->buffers.back()->data()), + descriptors.back().buffer_address_); +} + +TEST_F(BufferAllocatorTest, RestoreCachelibAllocationsAtOriginalAddresses) { + constexpr uintptr_t kBase = 0x1C0000000ULL; + constexpr size_t kCapacity = 4 * facebook::cachelib::Slab::kSize; + const std::string segment = "cachelib-restore"; + const std::string endpoint = "cachelib-restore-endpoint"; + auto original = std::make_shared( + segment, kBase, kCapacity, endpoint); + + auto small_first = original->allocate(64); + auto small_hole = original->allocate(64); + auto small_last = original->allocate(64); + auto large_first = original->allocate(4096); + auto large_hole = original->allocate(4096); + auto large_last = original->allocate(4096); + ASSERT_NE(small_first, nullptr); + ASSERT_NE(small_hole, nullptr); + ASSERT_NE(small_last, nullptr); + ASSERT_NE(large_first, nullptr); + ASSERT_NE(large_hole, nullptr); + ASSERT_NE(large_last, nullptr); + + std::vector descriptors = { + large_last->get_descriptor(), small_first->get_descriptor(), + large_first->get_descriptor(), small_last->get_descriptor()}; + small_hole.reset(); + large_hole.reset(); + + auto restored = RestoreCachelibBufferAllocator(segment, kBase, kCapacity, + endpoint, descriptors); + ASSERT_TRUE(restored.has_value()); + ASSERT_EQ(restored->buffers.size(), descriptors.size()); + for (size_t i = 0; i < descriptors.size(); ++i) { + EXPECT_EQ(reinterpret_cast(restored->buffers[i]->data()), + descriptors[i].buffer_address_); + } + + auto new_buffer = restored->allocator->allocate(64); + ASSERT_NE(new_buffer, nullptr); + const auto new_address = reinterpret_cast(new_buffer->data()); + for (const auto& descriptor : descriptors) { + EXPECT_NE(new_address, descriptor.buffer_address_); + } + + const uintptr_t released = descriptors[1].buffer_address_; + restored->buffers[1].reset(); + auto replacement = restored->allocator->allocate(descriptors[1].size_); + ASSERT_NE(replacement, nullptr); + EXPECT_EQ(reinterpret_cast(replacement->data()), released); +} + +TEST_F(BufferAllocatorTest, RestoreCachelibAllocationsRejectsInvalidLayouts) { + constexpr uintptr_t kBase = 0x1D0000000ULL; + constexpr size_t kCapacity = 4 * facebook::cachelib::Slab::kSize; + constexpr size_t kSlabSize = facebook::cachelib::Slab::kSize; + const std::string endpoint = "cachelib-invalid-endpoint"; + auto descriptor = [&](uintptr_t address, uint64_t size) { + return AllocatedBuffer::Descriptor{size, address, "tcp", endpoint}; + }; + auto restore = [&](const std::vector& descs) { + return RestoreCachelibBufferAllocator("cachelib-invalid", kBase, + kCapacity, endpoint, descs); + }; + + EXPECT_FALSE( + restore({descriptor(kBase, 64), descriptor(kBase, 4096)}).has_value()); + EXPECT_FALSE(restore({descriptor(kBase + 1, 64)}).has_value()); + EXPECT_FALSE( + restore({descriptor(kBase, 64), descriptor(kBase, 64)}).has_value()); + + auto wrong_endpoint = descriptor(kBase, 64); + wrong_endpoint.transport_endpoint_ = "wrong"; + EXPECT_FALSE(restore({wrong_endpoint}).has_value()); + EXPECT_FALSE(restore({descriptor(kBase + kCapacity, 64)}).has_value()); + EXPECT_FALSE(RestoreCachelibBufferAllocator("cachelib-invalid", kBase + 1, + kCapacity, endpoint, {}) + .has_value()); + EXPECT_FALSE(RestoreCachelibBufferAllocator( + "cachelib-invalid", + std::numeric_limits::max() - kSlabSize, + 2 * kSlabSize, endpoint, {}) + .has_value()); + + auto valid_after_fail = restore({descriptor(kBase + kSlabSize, 4096)}); + ASSERT_TRUE(valid_after_fail.has_value()); + EXPECT_EQ(reinterpret_cast(valid_after_fail->buffers[0]->data()), + kBase + kSlabSize); +} + +TEST_F(BufferAllocatorTest, CachelibImportRejectsChunkInSlabTail) { + constexpr uintptr_t kBase = 0x1E0000000ULL; + constexpr size_t kCapacity = 2 * facebook::cachelib::Slab::kSize; + constexpr uint32_t kAllocSize = facebook::cachelib::Slab::kSize - 16; + const size_t header_size = sizeof(facebook::cachelib::SlabHeader) * 2 + 1; + auto headers = std::make_unique(header_size); + facebook::cachelib::MemoryAllocator allocator( + facebook::cachelib::MemoryAllocator::Config({kAllocSize}), + headers.get(), header_size, reinterpret_cast(kBase), kCapacity); + const auto pool = allocator.addPool("main", kCapacity); + + EXPECT_FALSE(allocator.importAllocations( + pool, {{reinterpret_cast(kBase + kAllocSize), kAllocSize}})); +} + +TEST_F(BufferAllocatorTest, RestoreCachelibRejectsNonMemoryDescriptors) { + constexpr uintptr_t kBase = 0x1F0000000ULL; + constexpr size_t kCapacity = 2 * facebook::cachelib::Slab::kSize; + const std::string endpoint = "cachelib-memory-only"; + std::vector descriptors = { + {64, kBase, "tcp", endpoint}}; + + EXPECT_FALSE(RestoreCachelibBufferAllocator( + "cachelib-memory-only", kBase, kCapacity, endpoint, + descriptors, ReplicaType::NOF_SSD) + .has_value()); + + descriptors[0].protocol_ = "cxl"; + EXPECT_FALSE(RestoreCachelibBufferAllocator("cachelib-memory-only", kBase, + kCapacity, endpoint, + descriptors) + .has_value()); + + descriptors[0].protocol_ = "rdma"; + auto rdma = RestoreCachelibBufferAllocator( + "cachelib-memory-only", kBase, kCapacity, endpoint, descriptors); + ASSERT_TRUE(rdma.has_value()); + const auto restored = rdma->buffers[0]->get_descriptor(); + EXPECT_EQ(restored.protocol_, descriptors[0].protocol_); + EXPECT_EQ(restored.buffer_address_, descriptors[0].buffer_address_); + EXPECT_EQ(restored.transport_endpoint_, descriptors[0].transport_endpoint_); +} + // Test allocation request larger than available space TEST_F(BufferAllocatorTest, AllocateTooLarge) { for (const auto& allocator_type : allocator_types_) { diff --git a/mooncake-store/tests/client_buffer_test.cpp b/mooncake-store/tests/client_buffer_test.cpp index e48069eca3..d144f3c561 100644 --- a/mooncake-store/tests/client_buffer_test.cpp +++ b/mooncake-store/tests/client_buffer_test.cpp @@ -1,5 +1,5 @@ // client_buffer_test.cpp -#include "client_buffer.hpp" +#include "client_buffer.h" #include #include @@ -7,6 +7,7 @@ #include #include #include +#include #include #include @@ -59,6 +60,35 @@ TEST_F(ClientBufferTest, ZeroSizeAllocator) { EXPECT_FALSE(handle_opt.has_value()); } +#if defined(USE_NOF) +TEST_F(ClientBufferTest, SpdkDmaAllocatorDestroysWithSpdkFree) { + constexpr size_t buffer_size = 4096; + constexpr size_t alloc_size = 64; + + std::shared_ptr allocator; + try { + allocator = ClientBufferAllocator::create( + buffer_size, "tcp", /*use_hugepage=*/false, + /*use_spdk_dma=*/true); + } catch (const std::bad_alloc&) { + GTEST_SKIP() + << "SPDK DMA allocation is unavailable in this environment"; + } + + ASSERT_NE(allocator, nullptr); + VerifyAlignment(allocator->getBase()); + + { + auto handle_opt = allocator->allocate(alloc_size); + ASSERT_TRUE(handle_opt.has_value()); + BufferHandle handle = std::move(handle_opt.value()); + VerifyBufferHandle(handle, alloc_size); + } + + allocator.reset(); +} +#endif + // Test multiple allocations TEST_F(ClientBufferTest, MultipleAllocations) { const size_t buffer_size = 1024 * 1024; // 1MB diff --git a/mooncake-store/tests/client_integration_test.cpp b/mooncake-store/tests/client_integration_test.cpp index e30b6df405..6710096d87 100644 --- a/mooncake-store/tests/client_integration_test.cpp +++ b/mooncake-store/tests/client_integration_test.cpp @@ -2,11 +2,15 @@ #include #include +#include +#include #include #include #include #include +#include #include +#include #include #include #include @@ -20,6 +24,8 @@ #include "utils.h" #include "test_server_helpers.h" #include "default_config.h" +#include "crc_checksum.h" +#include "environ.h" DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp"); DEFINE_string(device_name, "", "Device name to use, valid if protocol=rdma"); @@ -52,6 +58,54 @@ UUID ParseClientId(const std::string& client_id_str) { return client_id; } +class ObjectChecksumClient : public Client { + public: + ObjectChecksumClient() : Client("localhost", "", "tcp") {} +}; + +TEST(ObjectChecksumTest, ClientVerifiesOnlyLogicalObjectBytes) { + if (!Environ::Get().GetStoreChecksumEnabled()) { + GTEST_SKIP() << "MOONCAKE_STORE_CHECKSUM is not enabled"; + } + + ObjectChecksumClient client; + std::array value = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, + 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, + 0x0E, 0x0F, 0xAA, 0xBB, 0xCC, 0xDD}; + constexpr size_t object_size = 16; + std::vector slices{{value.data(), value.size()}}; + const uint64_t checksum = ComputeCrcChecksum(value.data(), object_size); + + EXPECT_TRUE( + client.VerifyObjectChecksum("key", slices, object_size, checksum) + .has_value()); + value[16] ^= 0xFF; + EXPECT_TRUE( + client.VerifyObjectChecksum("key", slices, object_size, checksum) + .has_value()); + value[0] ^= 0xFF; + auto mismatch = + client.VerifyObjectChecksum("key", slices, object_size, checksum); + ASSERT_FALSE(mismatch.has_value()); + EXPECT_EQ(mismatch.error(), ErrorCode::CHECKSUM_MISMATCH); +} + +TEST(ObjectChecksumTest, BatchPutStopsAfterChecksumPrecomputeFailure) { + if (!Environ::Get().GetStoreChecksumEnabled()) { + GTEST_SKIP() << "MOONCAKE_STORE_CHECKSUM is not enabled"; + } + + ObjectChecksumClient client; + const std::vector keys{"invalid-slice"}; + std::vector> slices{{Slice{nullptr, 1}}}; + + auto results = client.BatchPut(keys, slices, ReplicateConfig{}); + + ASSERT_EQ(results.size(), 1); + ASSERT_FALSE(results[0].has_value()); + EXPECT_EQ(results[0].error(), ErrorCode::INVALID_PARAMS); +} + class ClientIdCaptureSink : public google::LogSink { public: std::string captured_client_id; @@ -342,6 +396,82 @@ TEST_F(ClientIntegrationTest, BasicPutGetOperations) { client_buffer_allocator_->deallocate(buffer, test_data.size()); } +TEST_F(ClientIntegrationTest, ObjectChecksumRejectsCorruptedObject) { + if (!Environ::Get().GetStoreChecksumEnabled()) { + GTEST_SKIP() << "MOONCAKE_STORE_CHECKSUM is not enabled"; + } + + const std::string key = "checksum_corruption_key"; + const std::string test_data = "0123456789abcdef"; + void* source = client_buffer_allocator_->allocate(test_data.size()); + memcpy(source, test_data.data(), test_data.size()); + std::vector slices{{source, test_data.size()}}; + + ReplicateConfig config; + config.replica_num = 1; + auto put_result = test_client_->Put(key, slices, config); + ASSERT_TRUE(put_result.has_value()) << toString(put_result.error()); + client_buffer_allocator_->deallocate(source, test_data.size()); + + auto query_result = test_client_->Query(key); + ASSERT_TRUE(query_result.has_value()) << toString(query_result.error()); + ASSERT_TRUE(query_result->object_checksum.has_value()); + ASSERT_EQ(query_result->replicas.size(), 1); + ASSERT_TRUE(query_result->replicas[0].is_memory_replica()); + + auto& descriptor = + query_result->replicas[0].get_memory_descriptor().buffer_descriptor; + auto* stored_data = reinterpret_cast(descriptor.buffer_address_); + stored_data[0] ^= 0x01; + + void* target = client_buffer_allocator_->allocate(test_data.size()); + slices = {{target, test_data.size()}}; + auto get_result = test_client_->Get(key, query_result.value(), slices); + ASSERT_FALSE(get_result.has_value()); + EXPECT_EQ(get_result.error(), ErrorCode::CHECKSUM_MISMATCH); + + stored_data[0] ^= 0x01; + client_buffer_allocator_->deallocate(target, test_data.size()); +} + +TEST_F(ClientIntegrationTest, BatchPutPreservesObjectChecksumPairing) { + if (!Environ::Get().GetStoreChecksumEnabled()) { + GTEST_SKIP() << "MOONCAKE_STORE_CHECKSUM is not enabled"; + } + + const std::vector keys = {"checksum_batch_a", + "checksum_batch_b"}; + const std::vector values = {"first-object", + "different-second-object"}; + std::vector> batched_slices; + std::vector buffers; + batched_slices.reserve(values.size()); + buffers.reserve(values.size()); + for (const auto& value : values) { + void* buffer = client_buffer_allocator_->allocate(value.size()); + memcpy(buffer, value.data(), value.size()); + buffers.emplace_back(buffer); + batched_slices.push_back({Slice{buffer, value.size()}}); + } + + ReplicateConfig config; + config.replica_num = 1; + auto results = test_client_->BatchPut(keys, batched_slices, config); + ASSERT_EQ(results.size(), keys.size()); + for (const auto& result : results) { + ASSERT_TRUE(result.has_value()) << toString(result.error()); + } + + for (size_t i = 0; i < keys.size(); ++i) { + auto query_result = test_client_->Query(keys[i]); + ASSERT_TRUE(query_result.has_value()) << toString(query_result.error()); + ASSERT_TRUE(query_result->object_checksum.has_value()); + EXPECT_EQ(*query_result->object_checksum, + ComputeCrcChecksum(values[i].data(), values[i].size())); + client_buffer_allocator_->deallocate(buffers[i], values[i].size()); + } +} + // Test Remove operation TEST_F(ClientIntegrationTest, RemoveOperation) { const std::string test_data = "Test data for removal"; diff --git a/mooncake-store/tests/client_local_hot_cache_test.cpp b/mooncake-store/tests/client_local_hot_cache_test.cpp index 86927c2814..70fbf7c162 100644 --- a/mooncake-store/tests/client_local_hot_cache_test.cpp +++ b/mooncake-store/tests/client_local_hot_cache_test.cpp @@ -1,6 +1,6 @@ // client_local_hot_cache_test.cpp #include "client_service.h" -#include "client_buffer.hpp" +#include "client_buffer.h" #include "count_min_sketch.h" #include "local_hot_cache.h" #include "replica.h" diff --git a/mooncake-store/tests/client_metrics_test.cpp b/mooncake-store/tests/client_metrics_test.cpp index 5c3fc5bbfc..1f23bc0059 100644 --- a/mooncake-store/tests/client_metrics_test.cpp +++ b/mooncake-store/tests/client_metrics_test.cpp @@ -1,12 +1,77 @@ #include #include +// csignal must precede coro_http_client.hpp: the bundled ylt's coro_io.hpp +// calls std::signal without including itself. +#include #include +#include #include +#include +#include #include "client_metric.h" +#include "real_client.h" +#include "test_server_helpers.h" +#include "utils.h" namespace mooncake::test { +namespace { + +struct HttpResponse { + int status; + std::string body; +}; + +HttpResponse FetchUrl(const std::string& url) { + coro_http::coro_http_client client; + auto res = client.get(url); + return HttpResponse{res.status, std::string(res.resp_body)}; +} + +int GetTestPort(std::unordered_set& used_ports) { + for (int i = 0; i < 100; ++i) { + int port = getFreeTcpPort(); + if (port > 0 && port < 65535 && !used_ports.contains(port)) { + used_ports.insert(port); + return port; + } + } + return -1; +} + +class ScopedEnv { + public: + explicit ScopedEnv(const char* name) : name_(name) { + const char* value = std::getenv(name); + if (value) old_value_ = value; + } + + ~ScopedEnv() { + if (old_value_) { + setenv(name_, old_value_->c_str(), 1); + } else { + unsetenv(name_); + } + } + + private: + const char* name_; + std::optional old_value_; +}; + +tl::expected SetupClientWithHttp( + const std::shared_ptr& client, const std::string& client_addr, + const std::string& master_addr, bool enable_http, int http_port) { + return client->setup_internal( + client_addr, "P2PHANDSHAKE", /*global_segment_size=*/0, + /*local_buffer_size=*/0, "tcp", "", master_addr, nullptr, "", + /*local_rpc_port=*/50052, /*enable_ssd_offload=*/false, + /*start_offload_rpc_server=*/false, /*ssd_offload_path=*/"", + /*tenant_id=*/"default", enable_http, http_port); +} + +} // namespace class ClientMetricsTest : public ::testing::Test { protected: @@ -304,4 +369,178 @@ TEST_F(ClientMetricsTest, SerializeWithoutDynamicLabels) { } } +TEST_F(ClientMetricsTest, HttpMetricsEndpointsReturnData) { + std::unordered_set used_ports; + int master_rpc_port = GetTestPort(used_ports); + int master_http_port = GetTestPort(used_ports); + int http_port = GetTestPort(used_ports); + int client_port = GetTestPort(used_ports); + ASSERT_GT(master_rpc_port, 0); + ASSERT_GT(master_http_port, 0); + ASSERT_GT(http_port, 0); + ASSERT_GT(client_port, 0); + + mooncake::testing::InProcMaster master; + ASSERT_TRUE(master.Start(mooncake::InProcMasterConfigBuilder() + .set_rpc_port(master_rpc_port) + .set_http_metrics_port(master_http_port) + .set_http_metadata_port(0) + .build())); + + auto client = RealClient::create(); + auto setup_result = SetupClientWithHttp( + client, "127.0.0.1:" + std::to_string(client_port), + master.master_address(), /*enable_http=*/true, http_port); + ASSERT_TRUE(setup_result.has_value()) << toString(setup_result.error()); + + auto metrics = + FetchUrl("http://127.0.0.1:" + std::to_string(http_port) + "/metrics"); + EXPECT_EQ(metrics.status, 200); + EXPECT_EQ(metrics.body.find("metrics not available"), std::string::npos); + + auto summary = FetchUrl("http://127.0.0.1:" + std::to_string(http_port) + + "/metrics/summary"); + EXPECT_EQ(summary.status, 200); + EXPECT_NE(summary.body.find("Client Metrics Summary"), std::string::npos); + + EXPECT_EQ(client->tearDownAll(), 0); +} + +TEST_F(ClientMetricsTest, HttpMetricsConfigParserTrimsWhitespace) { + std::unordered_set used_ports; + int master_rpc_port = GetTestPort(used_ports); + int master_http_port = GetTestPort(used_ports); + int http_port = GetTestPort(used_ports); + int client_port = GetTestPort(used_ports); + ASSERT_GT(master_rpc_port, 0); + ASSERT_GT(master_http_port, 0); + ASSERT_GT(http_port, 0); + ASSERT_GT(client_port, 0); + + mooncake::testing::InProcMaster master; + ASSERT_TRUE(master.Start(mooncake::InProcMasterConfigBuilder() + .set_rpc_port(master_rpc_port) + .set_http_metrics_port(master_http_port) + .set_http_metadata_port(0) + .build())); + + ConfigDict config = { + {CONFIG_KEY_LOCAL_HOSTNAME, "127.0.0.1:" + std::to_string(client_port)}, + {CONFIG_KEY_METADATA_SERVER, "P2PHANDSHAKE"}, + {CONFIG_KEY_GLOBAL_SEGMENT_SIZE, "0"}, + {CONFIG_KEY_LOCAL_BUFFER_SIZE, "0"}, + {CONFIG_KEY_PROTOCOL, "tcp"}, + {CONFIG_KEY_MASTER_SERVER_ADDR, master.master_address()}, + {CONFIG_KEY_ENABLE_CLIENT_HTTP_SERVER, " true "}, + {CONFIG_KEY_CLIENT_HTTP_PORT, " " + std::to_string(http_port) + " "}, + }; + + auto client = RealClient::create(); + auto setup_result = client->setup_internal(config); + ASSERT_TRUE(setup_result.has_value()) << toString(setup_result.error()); + + auto health = + FetchUrl("http://127.0.0.1:" + std::to_string(http_port) + "/health"); + EXPECT_EQ(health.status, 200); + EXPECT_NE(health.body.find("\"status\":\"healthy\""), std::string::npos); + + EXPECT_EQ(client->tearDownAll(), 0); +} + +TEST_F(ClientMetricsTest, HttpMetricsConfigParserRejectsInvalidIntegers) { + const char* invalid_ports[] = { + "9300x", + "999999999999999999999999", + }; + + for (const char* invalid_port : invalid_ports) { + ConfigDict config = { + {CONFIG_KEY_LOCAL_HOSTNAME, "127.0.0.1:1"}, + {CONFIG_KEY_METADATA_SERVER, "P2PHANDSHAKE"}, + {CONFIG_KEY_GLOBAL_SEGMENT_SIZE, "0"}, + {CONFIG_KEY_LOCAL_BUFFER_SIZE, "0"}, + {CONFIG_KEY_PROTOCOL, "tcp"}, + {CONFIG_KEY_ENABLE_CLIENT_HTTP_SERVER, "true"}, + {CONFIG_KEY_CLIENT_HTTP_PORT, invalid_port}, + }; + + auto client = RealClient::create(); + auto setup_result = client->setup_internal(config); + ASSERT_FALSE(setup_result.has_value()) << invalid_port; + EXPECT_EQ(setup_result.error(), ErrorCode::INVALID_PARAMS) + << invalid_port; + } +} + +TEST_F(ClientMetricsTest, HttpMetricsEndpointReturns503WhenMetricsDisabled) { + ScopedEnv metrics_env("MC_STORE_CLIENT_METRIC"); + setenv("MC_STORE_CLIENT_METRIC", "0", 1); + + std::unordered_set used_ports; + int master_rpc_port = GetTestPort(used_ports); + int master_http_port = GetTestPort(used_ports); + int http_port = GetTestPort(used_ports); + int client_port = GetTestPort(used_ports); + ASSERT_GT(master_rpc_port, 0); + ASSERT_GT(master_http_port, 0); + ASSERT_GT(http_port, 0); + ASSERT_GT(client_port, 0); + + mooncake::testing::InProcMaster master; + ASSERT_TRUE(master.Start(mooncake::InProcMasterConfigBuilder() + .set_rpc_port(master_rpc_port) + .set_http_metrics_port(master_http_port) + .set_http_metadata_port(0) + .build())); + + auto client = RealClient::create(); + auto setup_result = SetupClientWithHttp( + client, "127.0.0.1:" + std::to_string(client_port), + master.master_address(), /*enable_http=*/true, http_port); + ASSERT_TRUE(setup_result.has_value()) << toString(setup_result.error()); + + auto metrics = + FetchUrl("http://127.0.0.1:" + std::to_string(http_port) + "/metrics"); + EXPECT_EQ(metrics.status, 503); + EXPECT_NE(metrics.body.find("metrics not available"), std::string::npos); + + EXPECT_EQ(client->tearDownAll(), 0); +} + +TEST_F(ClientMetricsTest, HttpMetricsPortConflictDoesNotFailSetup) { + std::unordered_set used_ports; + int master_rpc_port = GetTestPort(used_ports); + int master_http_port = GetTestPort(used_ports); + int http_port = GetTestPort(used_ports); + int first_client_port = GetTestPort(used_ports); + int second_client_port = GetTestPort(used_ports); + ASSERT_GT(master_rpc_port, 0); + ASSERT_GT(master_http_port, 0); + ASSERT_GT(http_port, 0); + ASSERT_GT(first_client_port, 0); + ASSERT_GT(second_client_port, 0); + + mooncake::testing::InProcMaster master; + ASSERT_TRUE(master.Start(mooncake::InProcMasterConfigBuilder() + .set_rpc_port(master_rpc_port) + .set_http_metrics_port(master_http_port) + .set_http_metadata_port(0) + .build())); + + auto first_client = RealClient::create(); + auto first_setup = SetupClientWithHttp( + first_client, "127.0.0.1:" + std::to_string(first_client_port), + master.master_address(), /*enable_http=*/true, http_port); + ASSERT_TRUE(first_setup.has_value()) << toString(first_setup.error()); + + auto second_client = RealClient::create(); + auto second_setup = SetupClientWithHttp( + second_client, "127.0.0.1:" + std::to_string(second_client_port), + master.master_address(), /*enable_http=*/true, http_port); + EXPECT_TRUE(second_setup.has_value()) << toString(second_setup.error()); + + EXPECT_EQ(second_client->tearDownAll(), 0); + EXPECT_EQ(first_client->tearDownAll(), 0); +} + } // namespace mooncake::test diff --git a/mooncake-store/tests/client_storage_backend_test.cpp b/mooncake-store/tests/client_storage_backend_test.cpp new file mode 100644 index 0000000000..07f15662a0 --- /dev/null +++ b/mooncake-store/tests/client_storage_backend_test.cpp @@ -0,0 +1,54 @@ +// Regression tests for Client::PrepareStorageBackend (issue #3134): an +// invalid storage configuration must surface as an error instead of +// dereferencing a null backend or leaving a half-initialized one behind. + +#include +#include + +#include + +#include "client_service.h" + +namespace mooncake { +namespace { + +class TestableClient : public Client { + public: + TestableClient() + : Client(/*local_hostname=*/"localhost:9003", + /*metadata_connstring=*/"", + /*protocol=*/"tcp", + /*labels=*/{}) {} + + using Client::PrepareStorageBackend; +}; + +TEST(ClientPrepareStorageBackendTest, InvalidRootDirReturnsErrorWithoutCrash) { + TestableClient client; + // Before the fix this dereferenced a null StorageBackend and crashed. + ErrorCode err = client.PrepareStorageBackend( + "/nonexistent_mooncake_store_test_path/12345", "fsdir", true, 0); + EXPECT_NE(err, ErrorCode::OK); +} + +TEST(ClientPrepareStorageBackendTest, EmptyFsdirReturnsErrorWithoutCrash) { + TestableClient client; + ErrorCode err = client.PrepareStorageBackend( + std::filesystem::current_path().string(), "", true, 0); + EXPECT_NE(err, ErrorCode::OK); +} + +TEST(ClientPrepareStorageBackendTest, ValidRootDirSucceeds) { + std::string root = std::filesystem::current_path().string() + + "/data/client_prepare_storage_backend_test"; + std::filesystem::create_directories(root); + + TestableClient client; + EXPECT_EQ(client.PrepareStorageBackend(root, "fsdir", true, 0), + ErrorCode::OK); + + std::filesystem::remove_all(root); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/dummy_client_get_buffer_test.cpp b/mooncake-store/tests/dummy_client_get_buffer_test.cpp index b1b1dc5328..4ee6063b17 100644 --- a/mooncake-store/tests/dummy_client_get_buffer_test.cpp +++ b/mooncake-store/tests/dummy_client_get_buffer_test.cpp @@ -16,6 +16,7 @@ #include #include "dummy_client.h" +#include "environ.h" #include "real_client.h" #include "default_config.h" #include "test_server_helpers.h" @@ -41,6 +42,7 @@ static void RegisterRpcHandlers(coro_rpc::coro_rpc_server &server, server.register_handler<&RealClient::removeAll_internal>(&rc); server.register_handler<&RealClient::isExist_internal>(&rc); server.register_handler<&RealClient::getSize_internal>(&rc); + server.register_handler<&RealClient::get_into_range_shm_helper>(&rc); server.register_handler<&RealClient::batch_get_into_dummy_helper>(&rc); server.register_handler<&RealClient::batch_put_from_dummy_helper>(&rc); server.register_handler<&RealClient::acquire_hot_cache>(&rc); @@ -50,6 +52,7 @@ static void RegisterRpcHandlers(coro_rpc::coro_rpc_server &server, server.register_handler<&RealClient::acquire_buffer_dummy>(&rc); server.register_handler<&RealClient::release_buffer_dummy>(&rc); server.register_handler<&RealClient::batch_acquire_buffer_dummy>(&rc); + server.register_handler<&RealClient::batch_get_query_results>(&rc); } static constexpr size_t kMB = 1024ULL * 1024; @@ -375,6 +378,58 @@ TEST_F(DummyClientGetBufferTest, GetBuffer_AllocatorFallback) { EXPECT_EQ(got, data) << "Data mismatch on allocator fallback path"; } +TEST_F(DummyClientGetBufferTest, GetIntoRejectsCorruptedObjectWithChecksum) { + if (!Environ::Get().GetStoreChecksumEnabled()) { + GTEST_SKIP() << "MOONCAKE_STORE_CHECKSUM is not enabled"; + } + ASSERT_TRUE(SetupStack()) << "Failed to bring up real+dummy stack"; + + const std::string key = "dummy_checksum_corruption"; + const std::string data = "0123456789abcdef"; + PutData(key, data); + + auto replicas = real_client_->get_replica_desc(key); + ASSERT_EQ(replicas.size(), 1); + ASSERT_TRUE(replicas[0].is_memory_replica()); + auto &descriptor = replicas[0].get_memory_descriptor().buffer_descriptor; + auto *stored_data = reinterpret_cast(descriptor.buffer_address_); + stored_data[0] ^= 0x01; + + const uint64_t target_addr = + dummy_client_->alloc_from_mem_pool(data.size()); + ASSERT_NE(target_addr, 0); + void *target = reinterpret_cast(target_addr); + ASSERT_EQ(dummy_client_->register_buffer(target, data.size()), 0); + EXPECT_EQ(dummy_client_->get_into(key, target, data.size()), + toInt(ErrorCode::CHECKSUM_MISMATCH)); + + stored_data[0] ^= 0x01; + EXPECT_EQ(dummy_client_->unregister_buffer(target), 0); + EXPECT_EQ(ShmHelper::getInstance()->free(target), 0); +} + +TEST_F(DummyClientGetBufferTest, BatchQueryPreservesObjectChecksum) { + if (!Environ::Get().GetStoreChecksumEnabled()) { + GTEST_SKIP() << "MOONCAKE_STORE_CHECKSUM is not enabled"; + } + ASSERT_TRUE(SetupStack()) << "Failed to bring up real+dummy stack"; + + const std::string key = "dummy_batch_query_checksum"; + const std::string data = "0123456789abcdef"; + PutData(key, data); + + auto real_results = real_client_->batch_query({key}); + ASSERT_EQ(real_results.size(), 1); + ASSERT_TRUE(real_results[0].has_value()); + ASSERT_TRUE(real_results[0]->object_checksum.has_value()); + + auto dummy_results = dummy_client_->batch_query({key}); + ASSERT_EQ(dummy_results.size(), 1); + ASSERT_TRUE(dummy_results[0].has_value()); + EXPECT_EQ(dummy_results[0]->object_checksum, + real_results[0]->object_checksum); +} + // ---- Test: get_buffer via hot cache shm zero-copy path ---- TEST_F(DummyClientGetBufferTest, GetBuffer_HotCachePath) { ASSERT_TRUE(SetupStack()) << "Failed to bring up real+dummy stack"; diff --git a/mooncake-store/tests/e2e/CMakeLists.txt b/mooncake-store/tests/e2e/CMakeLists.txt index ce23d3062b..7f477bd227 100644 --- a/mooncake-store/tests/e2e/CMakeLists.txt +++ b/mooncake-store/tests/e2e/CMakeLists.txt @@ -44,6 +44,16 @@ target_link_libraries(client_runner PUBLIC ${ETCD_WRAPPER_LIB} ) +add_executable(oplog_ha_client oplog_ha_client.cpp client_wrapper.cpp) +target_link_libraries(oplog_ha_client PUBLIC + mooncake_store + transfer_engine + cachelib_memory_allocator + glog + pthread + ${ETCD_WRAPPER_LIB} +) + add_executable(chaos_test chaos_test.cpp ${MOONCAKE_E2E_TEST_SOURCES}) target_link_libraries(chaos_test PUBLIC mooncake_store @@ -81,3 +91,8 @@ target_link_libraries(storage_backend_e2e_test PUBLIC ) add_test(NAME storage_backend_e2e_test COMMAND storage_backend_e2e_test) + +add_executable(oplog_batch_e2e_test oplog_batch_e2e_test.cpp process_handler.cpp) +target_link_libraries(oplog_batch_e2e_test PUBLIC glog gtest pthread) +add_test(NAME oplog_batch_e2e_test COMMAND oplog_batch_e2e_test) +set_tests_properties(oplog_batch_e2e_test PROPERTIES LABELS oplog_batch_smoke) diff --git a/mooncake-store/tests/e2e/client_wrapper.cpp b/mooncake-store/tests/e2e/client_wrapper.cpp index b8509bdac7..0f76a5a43e 100644 --- a/mooncake-store/tests/e2e/client_wrapper.cpp +++ b/mooncake-store/tests/e2e/client_wrapper.cpp @@ -155,6 +155,94 @@ ErrorCode ClientTestWrapper::Delete(const std::string& key) { return remove_result.has_value() ? ErrorCode::OK : remove_result.error(); } +ErrorCode ClientTestWrapper::BatchSmoke(const std::string& key_prefix) { + constexpr size_t kCount = 3; + std::vector keys; + std::vector values; + std::vector> put_guards; + std::vector> put_slices; + for (size_t i = 0; i < kCount; ++i) { + keys.push_back(key_prefix + "-" + std::to_string(i)); + values.push_back("batch-value-" + std::to_string(i)); + auto guard = + std::make_unique(values.back().size(), allocator_); + size_t offset = 0; + for (const auto& slice : guard->slices_) { + memcpy(slice.ptr, values.back().data() + offset, slice.size); + offset += slice.size; + } + put_slices.push_back(guard->slices_); + put_guards.push_back(std::move(guard)); + } + + ReplicateConfig config; + config.replica_num = 1; + auto put_results = client_->BatchPut(keys, put_slices, config); + if (put_results.size() != kCount) return ErrorCode::INTERNAL_ERROR; + for (const auto& result : put_results) { + if (!result) return result.error(); + } + + std::vector mixed_keys = keys; + mixed_keys.push_back(key_prefix + "-missing"); + auto exist_results = client_->BatchIsExist(mixed_keys); + if (exist_results.size() != mixed_keys.size()) { + return ErrorCode::INTERNAL_ERROR; + } + for (size_t i = 0; i < kCount; ++i) { + if (!exist_results[i] || !*exist_results[i]) { + return ErrorCode::INTERNAL_ERROR; + } + } + if (!exist_results.back() || *exist_results.back()) { + return ErrorCode::INTERNAL_ERROR; + } + + std::vector> get_guards; + std::unordered_map> get_slices; + for (size_t i = 0; i < mixed_keys.size(); ++i) { + const size_t size = i < kCount ? values[i].size() : 1; + auto guard = std::make_unique(size, allocator_); + get_slices.emplace(mixed_keys[i], guard->slices_); + get_guards.push_back(std::move(guard)); + } + auto get_results = client_->BatchGet(mixed_keys, get_slices); + if (get_results.size() != mixed_keys.size()) { + return ErrorCode::INTERNAL_ERROR; + } + for (size_t i = 0; i < kCount; ++i) { + if (!get_results[i]) return get_results[i].error(); + std::string actual; + for (const auto& slice : get_slices.at(keys[i])) { + actual.append(static_cast(slice.ptr), slice.size); + } + if (actual != values[i]) return ErrorCode::INTERNAL_ERROR; + } + if (get_results.back() || + get_results.back().error() != ErrorCode::OBJECT_NOT_FOUND) { + return ErrorCode::INTERNAL_ERROR; + } + + auto remove_results = client_->BatchRemove(mixed_keys, /*force=*/true); + if (remove_results.size() != mixed_keys.size()) { + return ErrorCode::INTERNAL_ERROR; + } + for (size_t i = 0; i < kCount; ++i) { + if (!remove_results[i]) return remove_results[i].error(); + } + if (remove_results.back() || + remove_results.back().error() != ErrorCode::OBJECT_NOT_FOUND) { + return ErrorCode::INTERNAL_ERROR; + } + + exist_results = client_->BatchIsExist(keys); + if (exist_results.size() != kCount) return ErrorCode::INTERNAL_ERROR; + for (const auto& result : exist_results) { + if (!result || *result) return ErrorCode::INTERNAL_ERROR; + } + return ErrorCode::OK; +} + bool ClientTestWrapper::HasDiskReplica(const std::string& key) { auto query_result = client_->Query(key); if (!query_result.has_value()) return false; diff --git a/mooncake-store/tests/e2e/client_wrapper.h b/mooncake-store/tests/e2e/client_wrapper.h index ed37f2a83f..68f3d12b71 100644 --- a/mooncake-store/tests/e2e/client_wrapper.h +++ b/mooncake-store/tests/e2e/client_wrapper.h @@ -69,6 +69,7 @@ class ClientTestWrapper { ErrorCode Get(const std::string& key, std::string& value); ErrorCode Put(const std::string& key, const std::string& value); ErrorCode Delete(const std::string& key); + ErrorCode BatchSmoke(const std::string& key_prefix); // Returns true if the key has a DISK replica // (master-assigned, written by PutToLocalFile). @@ -117,4 +118,4 @@ class ClientTestWrapper { }; } // namespace testing -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/tests/e2e/clientctl.cpp b/mooncake-store/tests/e2e/clientctl.cpp index ec1b7d69bf..29555a3f81 100644 --- a/mooncake-store/tests/e2e/clientctl.cpp +++ b/mooncake-store/tests/e2e/clientctl.cpp @@ -39,8 +39,14 @@ class ClientCtl { HandlePut(iss); } else if (cmd == "get") { HandleGet(iss); + } else if (cmd == "delete") { + HandleDelete(iss); + } else if (cmd == "batch-smoke") { + HandleBatchSmoke(iss); } else if (cmd == "mount") { HandleMount(iss); + } else if (cmd == "unmount") { + HandleUnmount(iss); } else if (cmd == "remove") { HandleRemove(iss); } else if (cmd == "sleep") { @@ -135,6 +141,53 @@ class ClientCtl { std::cout << "Get value: " << value << std::endl; } + void HandleDelete(std::istringstream& iss) { + std::string name, key; + iss >> name >> key; + + auto it = clients_.find(name); + if (it == clients_.end()) { + std::cout << "Client not found: " << name << std::endl; + return; + } + if (key.empty()) { + std::cout << "Empty key" << std::endl; + return; + } + + ErrorCode error_code = it->second.client->Delete(key); + if (error_code != ErrorCode::OK) { + std::cout << "Failed to delete value: " << toString(error_code) + << std::endl; + return; + } + std::cout << "Successfully deleted value for key: " << key << std::endl; + } + + void HandleBatchSmoke(std::istringstream& iss) { + std::string name, key_prefix; + iss >> name >> key_prefix; + + auto it = clients_.find(name); + if (it == clients_.end()) { + std::cout << "Client not found: " << name << std::endl; + return; + } + if (key_prefix.empty()) { + std::cout << "Empty key prefix" << std::endl; + return; + } + + ErrorCode error_code = it->second.client->BatchSmoke(key_prefix); + if (error_code != ErrorCode::OK) { + std::cout << "Failed batch smoke: " << toString(error_code) + << std::endl; + return; + } + std::cout << "Successfully completed batch smoke: " << key_prefix + << std::endl; + } + void HandleMount(std::istringstream& iss) { std::string client_name; std::string segment_name; @@ -175,6 +228,34 @@ class ClientCtl { << std::endl; } + void HandleUnmount(std::istringstream& iss) { + std::string client_name; + std::string segment_name; + iss >> client_name >> segment_name; + + auto client_it = clients_.find(client_name); + if (client_it == clients_.end()) { + std::cout << "Client not found: " << client_name << std::endl; + return; + } + auto segment_it = client_it->second.segments.find(segment_name); + if (segment_it == client_it->second.segments.end()) { + std::cout << "Segment not found: " << segment_name << std::endl; + return; + } + + ErrorCode error_code = + client_it->second.client->Unmount(segment_it->second); + if (error_code != ErrorCode::OK) { + std::cout << "Failed to unmount segment: " << toString(error_code) + << std::endl; + return; + } + client_it->second.segments.erase(segment_it); + std::cout << "Successfully unmounted segment from client " + << client_name << std::endl; + } + void HandleRemove(std::istringstream& iss) { std::string name; iss >> name; diff --git a/mooncake-store/tests/e2e/oplog_batch_e2e_test.cpp b/mooncake-store/tests/e2e/oplog_batch_e2e_test.cpp new file mode 100644 index 0000000000..592a87651a --- /dev/null +++ b/mooncake-store/tests/e2e/oplog_batch_e2e_test.cpp @@ -0,0 +1,87 @@ +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "process_handler.h" + +namespace mooncake::testing { +namespace { + +std::string SelfExecutable() { + return std::filesystem::read_symlink("/proc/self/exe").string(); +} + +TEST(ProcessHandlerLifecycleTest, WaitSignalAndRepeatedStop) { + const auto out_dir = + "/tmp/mooncake-process-handler-" + std::to_string(::getpid()); + MasterRunnerConfig config; + config.enable_ha = false; + config.extra_args = {"--process-handler-helper"}; + MasterProcessHandler child(SelfExecutable(), config, 0, 0, out_dir); + + ASSERT_TRUE(child.start()); + EXPECT_GT(child.pid(), 0); + EXPECT_TRUE(child.is_running()); + EXPECT_EQ(child.stdout_path(), out_dir + "/master_0.out"); + EXPECT_EQ(child.stderr_path(), out_dir + "/master_0.err"); + + int status = 0; + EXPECT_FALSE(child.wait_for_exit(std::chrono::milliseconds(20), &status)); + ASSERT_TRUE(child.signal(SIGTERM)); + ASSERT_TRUE(child.wait_for_exit(std::chrono::seconds(2), &status)); + EXPECT_TRUE(WIFSIGNALED(status)); + EXPECT_EQ(WTERMSIG(status), SIGTERM); + EXPECT_FALSE(child.is_running()); + EXPECT_TRUE(child.stop(std::chrono::milliseconds(20))); +} + +TEST(ProcessHandlerLifecycleTest, BuildsUniqueMetricsAndExtraArguments) { + const auto out_dir = + "/tmp/mooncake-process-handler-args-" + std::to_string(::getpid()); + MasterRunnerConfig config; + config.enable_ha = false; + config.extra_args = {"--process-handler-print-args"}; + MasterProcessHandler child(SelfExecutable(), config, 0, 7, out_dir); + + ASSERT_TRUE(child.start()); + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (child.is_running() && std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + EXPECT_FALSE(child.is_running()); + int status = 0; + ASSERT_TRUE(child.wait_for_exit(std::chrono::seconds(2), &status)); + + std::ifstream input(child.stdout_path()); + std::stringstream output; + output << input.rdbuf(); + EXPECT_NE(output.str().find("--metrics-port=9010"), std::string::npos); + EXPECT_TRUE(output.str().ends_with("--process-handler-print-args\n")); +} + +} // namespace +} // namespace mooncake::testing + +int main(int argc, char** argv) { + for (int i = 1; i < argc; ++i) { + if (std::string(argv[i]) == "--process-handler-helper") { + std::this_thread::sleep_for(std::chrono::minutes(5)); + return 0; + } + if (std::string(argv[i]) == "--process-handler-print-args") { + for (int j = 1; j < argc; ++j) { + std::cout << argv[j] << '\n'; + } + return 0; + } + } + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-store/tests/e2e/oplog_fault_ctl.sh b/mooncake-store/tests/e2e/oplog_fault_ctl.sh new file mode 100755 index 0000000000..778d324655 --- /dev/null +++ b/mooncake-store/tests/e2e/oplog_fault_ctl.sh @@ -0,0 +1,123 @@ +#!/usr/bin/env bash +set -euo pipefail + +die() { + echo "error: $*" >&2 + exit 1 +} + +usage() { + cat >&2 < --run-dir DIR --name NAME [--timeout-sec N] + $0 process --run-dir DIR --name NAME +EOF +} + +parse_options() { + RUN_DIR="" + NAME="" + TIMEOUT_SEC=30 + while (($#)); do + case "$1" in + --run-dir) + (($# >= 2)) || die "--run-dir requires a value" + RUN_DIR=$2 + shift 2 + ;; + --name) + (($# >= 2)) || die "--name requires a value" + NAME=$2 + shift 2 + ;; + --timeout-sec) + (($# >= 2)) || die "--timeout-sec requires a value" + TIMEOUT_SEC=$2 + shift 2 + ;; + *) die "unknown option: $1" ;; + esac + done + [[ -d "$RUN_DIR" ]] || die "run directory does not exist: $RUN_DIR" + [[ "$NAME" =~ ^[a-z0-9][a-z0-9_-]*$ ]] || die "invalid name: $NAME" + [[ "$TIMEOUT_SEC" =~ ^[1-9][0-9]*$ ]] || die "timeout-sec must be positive" +} + +record_event() { + mkdir -p "$RUN_DIR/faults" + printf '%s\t%s\t%s\n' "$(date +%s)" "$1" "$NAME" >>"$RUN_DIR/faults/events.tsv" +} + +failpoint_ctl() { + local action=$1 + [[ "$NAME" =~ ^[a-z0-9][a-z0-9_]*$ ]] || die "invalid failpoint name: $NAME" + local directory=${FAILPOINT_DIR:-"$RUN_DIR/failpoints"} + mkdir -p "$directory" + local base="$directory/$NAME" + case "$action" in + arm) + [[ ! -e "$base.hit" && ! -e "$base.release" ]] || + die "stale failpoint files exist for $NAME" + : >"$base.arm" + record_event "failpoint_arm" + ;; + wait) + local deadline=$((SECONDS + TIMEOUT_SEC)) + while [[ ! -e "$base.hit" ]] && ((SECONDS < deadline)); do + sleep 0.01 + done + [[ -e "$base.hit" ]] || die "timed out waiting for failpoint: $NAME" + record_event "failpoint_hit" + ;; + release) + [[ -e "$base.hit" ]] || die "failpoint has not been hit: $NAME" + : >"$base.release" + record_event "failpoint_release" + ;; + *) die "unknown failpoint action: $action" ;; + esac + echo "ok: failpoint $action $NAME" +} + +process_ctl() { + local action=$1 + local pid_file="$RUN_DIR/pids/$NAME.pid" + local cmd_file="$RUN_DIR/pids/$NAME.cmd" + [[ -r "$pid_file" && -r "$cmd_file" ]] || die "missing process files for $NAME" + local pid + pid=$(<"$pid_file") + [[ "$pid" =~ ^[0-9]+$ && -r "/proc/$pid/cmdline" ]] || + die "process is not running: $NAME" + local actual + actual=$(tr '\0' ' ' <"/proc/$pid/cmdline") + grep -Fqx "$actual" "$cmd_file" || die "PID command does not match: $NAME" + + local signal + case "$action" in + stop) signal=STOP ;; + continue) signal=CONT ;; + kill) signal=KILL ;; + *) die "unknown process action: $action" ;; + esac + kill -"$signal" "$pid" + record_event "process_$action" + echo "ok: process $action $NAME pid=$pid" +} + +main() { + (($# >= 2)) || { + usage + exit 1 + } + local kind=$1 + local action=$2 + shift 2 + parse_options "$@" + case "$kind" in + failpoint) failpoint_ctl "$action" ;; + process) process_ctl "$action" ;; + *) die "unknown control type: $kind" ;; + esac +} + +main "$@" diff --git a/mooncake-store/tests/e2e/oplog_fault_ctl_test.sh b/mooncake-store/tests/e2e/oplog_fault_ctl_test.sh new file mode 100755 index 0000000000..8fd6015827 --- /dev/null +++ b/mooncake-store/tests/e2e/oplog_fault_ctl_test.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +SCRIPT="$SCRIPT_DIR/oplog_fault_ctl.sh" +TEST_ROOT=${TEST_ROOT:-"/tmp/mooncake-oplog-fault-ctl-test-$$"} +mkdir -p "$TEST_ROOT"/{failpoints,pids} + +fail() { + echo "FAIL: $*" >&2 + exit 1 +} + +expect_failure() { + local expected=$1 + shift + local output="$TEST_ROOT/output-$RANDOM.log" + if "$@" >"$output" 2>&1; then + fail "command unexpectedly succeeded" + fi + grep -Fq "$expected" "$output" || fail "missing error '$expected'" +} + +NAME=batch_txn_succeeded_before_callback +"$SCRIPT" failpoint arm --run-dir "$TEST_ROOT" --name "$NAME" +[[ -e "$TEST_ROOT/failpoints/$NAME.arm" ]] || fail "arm file missing" +: >"$TEST_ROOT/failpoints/$NAME.hit" +"$SCRIPT" failpoint wait --run-dir "$TEST_ROOT" --name "$NAME" --timeout-sec 1 +"$SCRIPT" failpoint release --run-dir "$TEST_ROOT" --name "$NAME" +[[ -e "$TEST_ROOT/failpoints/$NAME.release" ]] || fail "release file missing" +expect_failure "invalid name" "$SCRIPT" failpoint arm \ + --run-dir "$TEST_ROOT" --name ../escape +expect_failure "invalid failpoint name" "$SCRIPT" failpoint arm \ + --run-dir "$TEST_ROOT" --name bad-name + +sleep 30 & +PID=$! +trap 'kill -CONT "$PID" 2>/dev/null || true; kill "$PID" 2>/dev/null || true' EXIT +printf '%s\n' "$PID" >"$TEST_ROOT/pids/master-0.pid" +tr '\0' ' ' <"/proc/$PID/cmdline" >"$TEST_ROOT/pids/master-0.cmd" +"$SCRIPT" process stop --run-dir "$TEST_ROOT" --name master-0 +grep -Eq '^State:[[:space:]]+T' "/proc/$PID/status" || fail "process was not stopped" +"$SCRIPT" process continue --run-dir "$TEST_ROOT" --name master-0 +for _ in {1..100}; do + grep -Eq '^State:[[:space:]]+T' "/proc/$PID/status" || break + sleep 0.01 +done +grep -Eq '^State:[[:space:]]+T' "/proc/$PID/status" && fail "process was not continued" +[[ $(wc -l <"$TEST_ROOT/faults/events.tsv") -eq 5 ]] || fail "unexpected event count" + +echo "PASS" diff --git a/mooncake-store/tests/e2e/oplog_ha_client.cpp b/mooncake-store/tests/e2e/oplog_ha_client.cpp new file mode 100644 index 0000000000..d4cdce2318 --- /dev/null +++ b/mooncake-store/tests/e2e/oplog_ha_client.cpp @@ -0,0 +1,297 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "client_wrapper.h" +#include "e2e_utils.h" +#include "types.h" +#include "utils.h" + +USE_engine_flags; +DEFINE_string(master_server_entry, "etcd://0.0.0.0:2379", + "Master server entry"); +DEFINE_string(mode, "", "provider, seed, verify, or pressure"); +DEFINE_int32(port, 19001, "Client transfer-engine port"); +DEFINE_string(key_prefix, "ha-e2e", "Key prefix"); +DEFINE_uint64(start_index, 0, "First object index"); +DEFINE_uint64(count, 1000, "Number of objects for seed"); +DEFINE_uint64(payload_size, 4096, "Payload bytes per object"); +DEFINE_string(payload_sizes, "", "Comma-separated payload sizes"); +DEFINE_string(manifest, "", "Acknowledged-write manifest path"); +DEFINE_uint64(duration_sec, 45, "Pressure duration"); +DEFINE_uint64(sleep_ms, 25, "Delay between pressure operations"); +DEFINE_uint64(connect_timeout_sec, 30, "Client creation timeout"); +DEFINE_uint64(segment_size, 134217728, "Provider memory-segment bytes"); + +namespace mooncake::testing { +namespace { + +std::vector payload_sizes; + +std::string Key(uint64_t index) { + return FLAGS_key_prefix + "-" + std::to_string(index); +} + +std::string Payload(uint64_t index) { + const uint64_t size = payload_sizes.empty() + ? FLAGS_payload_size + : payload_sizes[index % payload_sizes.size()]; + std::string value(size, '\0'); + for (uint64_t offset = 0; offset < size; ++offset) { + value[offset] = static_cast( + (index * 1315423911ULL + offset * 2654435761ULL) & 0xff); + } + return value; +} + +bool ParsePayloadSizes() { + if (FLAGS_payload_sizes.empty()) return true; + size_t begin = 0; + while (begin <= FLAGS_payload_sizes.size()) { + const size_t end = FLAGS_payload_sizes.find(',', begin); + const std::string item = FLAGS_payload_sizes.substr(begin, end - begin); + try { + size_t parsed = 0; + if (item.empty() || + item.find_first_not_of("0123456789") != std::string::npos) + return false; + const uint64_t size = std::stoull(item, &parsed); + if (parsed != item.size() || size == 0) return false; + payload_sizes.push_back(size); + } catch (const std::exception&) { + return false; + } + if (end == std::string::npos) break; + begin = end + 1; + } + return !payload_sizes.empty(); +} + +std::shared_ptr CreateClient() { + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::seconds(FLAGS_connect_timeout_sec); + do { + auto client = ClientTestWrapper::CreateClientWrapper( + "localhost:" + std::to_string(FLAGS_port), FLAGS_engine_meta_url, + FLAGS_protocol, FLAGS_device_name, FLAGS_master_server_entry); + if (client) return *client; + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } while (std::chrono::steady_clock::now() < deadline); + return nullptr; +} + +class AckManifest { + public: + explicit AckManifest(const std::string& path) + : fd_(open(path.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644)) {} + + ~AckManifest() { + if (fd_ >= 0) close(fd_); + } + + bool valid() const { return fd_ >= 0; } + + bool Append(uint64_t index) { + const std::string line = std::to_string(index) + "\n"; + size_t written = 0; + while (written < line.size()) { + const ssize_t n = + write(fd_, line.data() + written, line.size() - written); + if (n <= 0) return false; + written += static_cast(n); + } + return fsync(fd_) == 0; + } + + private: + int fd_; +}; + +bool ReadManifest(std::vector& indexes) { + std::ifstream input(FLAGS_manifest); + if (!input) return false; + std::set seen; + std::string line; + while (std::getline(input, line)) { + if (line.empty() || + line.find_first_not_of("0123456789") != std::string::npos) { + return false; + } + try { + size_t parsed = 0; + uint64_t index = std::stoull(line, &parsed); + if (parsed != line.size() || !seen.insert(index).second) { + return false; + } + indexes.push_back(index); + } catch (const std::exception&) { + return false; + } + } + return !indexes.empty(); +} + +bool VerifyOne(ClientTestWrapper& client, uint64_t index) { + std::string actual; + const ErrorCode error = client.Get(Key(index), actual); + if (error != ErrorCode::OK) { + std::cerr << "get_failed index=" << index + << " error=" << toString(error) << '\n'; + return false; + } + if (actual != Payload(index)) { + std::cerr << "value_mismatch index=" << index + << " actual_size=" << actual.size() << '\n'; + return false; + } + return true; +} + +int Seed(ClientTestWrapper& client) { + AckManifest manifest(FLAGS_manifest); + if (!manifest.valid()) return 2; + uint64_t put_ok = 0; + for (uint64_t offset = 0; offset < FLAGS_count; ++offset) { + const uint64_t index = FLAGS_start_index + offset; + const ErrorCode error = client.Put(Key(index), Payload(index)); + if (error != ErrorCode::OK || !manifest.Append(index)) { + std::cerr << "seed_failed index=" << index + << " error=" << toString(error) << '\n'; + return 20; + } + ++put_ok; + } + uint64_t get_ok = 0; + for (uint64_t offset = 0; offset < FLAGS_count; ++offset) { + if (!VerifyOne(client, FLAGS_start_index + offset)) return 21; + ++get_ok; + } + std::cout << "summary mode=seed put_ok=" << put_ok + << " put_fail=0 get_ok=" << get_ok << " get_fail=0 mismatch=0\n"; + return 0; +} + +int Verify(ClientTestWrapper& client) { + std::vector indexes; + if (!ReadManifest(indexes)) return 2; + uint64_t get_ok = 0; + for (uint64_t index : indexes) { + if (!VerifyOne(client, index)) return 21; + ++get_ok; + } + std::cout << "summary mode=verify put_ok=0 put_fail=0 get_ok=" << get_ok + << " get_fail=0 mismatch=0\n"; + return 0; +} + +int Delete(ClientTestWrapper& client) { + std::vector indexes; + if (!ReadManifest(indexes)) return 2; + for (uint64_t index : indexes) { + const ErrorCode error = client.Delete(Key(index)); + if (error != ErrorCode::OK) { + std::cerr << "delete_failed index=" << index + << " error=" << toString(error) << '\n'; + return 20; + } + } + std::cout << "summary mode=delete delete_ok=" << indexes.size() + << " delete_fail=0\n"; + return 0; +} + +int VerifyAbsent(ClientTestWrapper& client) { + std::vector indexes; + if (!ReadManifest(indexes)) return 2; + for (uint64_t index : indexes) { + std::string value; + const ErrorCode error = client.Get(Key(index), value); + if (error != ErrorCode::OBJECT_NOT_FOUND) { + std::cerr << "unexpected_present index=" << index + << " error=" << toString(error) << '\n'; + return 21; + } + } + std::cout << "summary mode=verify-absent absent_ok=" << indexes.size() + << " absent_fail=0\n"; + return 0; +} + +int Pressure(ClientTestWrapper& client) { + AckManifest manifest(FLAGS_manifest); + if (!manifest.valid()) return 2; + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::seconds(FLAGS_duration_sec); + uint64_t index = FLAGS_start_index; + uint64_t put_ok = 0; + uint64_t get_ok = 0; + while (std::chrono::steady_clock::now() < deadline) { + const ErrorCode error = client.Put(Key(index), Payload(index)); + if (error != ErrorCode::OK || !manifest.Append(index)) { + std::cerr << "pressure_put_failed index=" << index + << " error=" << toString(error) << '\n'; + return 20; + } + ++put_ok; + if (!VerifyOne(client, index)) return 21; + ++get_ok; + ++index; + std::this_thread::sleep_for(std::chrono::milliseconds(FLAGS_sleep_ms)); + } + if (put_ok == 0 || get_ok == 0) return 22; + std::cout << "summary mode=pressure put_ok=" << put_ok + << " put_fail=0 get_ok=" << get_ok << " get_fail=0 mismatch=0\n"; + return 0; +} + +int Provider(ClientTestWrapper& client) { + void* buffer = nullptr; + const ErrorCode error = client.Mount(FLAGS_segment_size, buffer); + if (error != ErrorCode::OK) { + std::cerr << "provider_mount_failed error=" << toString(error) << '\n'; + return 20; + } + std::cout << "provider_ready size=" << FLAGS_segment_size << std::endl; + while (true) pause(); +} + +} // namespace +} // namespace mooncake::testing + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = 1; + if (!mooncake::testing::ParsePayloadSizes() || + (FLAGS_mode != "provider" && FLAGS_mode != "seed" && + FLAGS_mode != "verify" && FLAGS_mode != "delete" && + FLAGS_mode != "verify-absent" && FLAGS_mode != "pressure") || + (FLAGS_mode != "provider" && FLAGS_manifest.empty()) || + FLAGS_key_prefix.empty() || FLAGS_payload_size == 0 || + FLAGS_connect_timeout_sec == 0 || + (FLAGS_mode == "provider" && FLAGS_segment_size == 0) || + (FLAGS_mode == "seed" && FLAGS_count == 0) || + (FLAGS_mode == "pressure" && FLAGS_duration_sec == 0)) { + std::cerr << "invalid arguments\n"; + return 2; + } + auto client = mooncake::testing::CreateClient(); + if (!client) return 10; + if (FLAGS_mode == "provider") return mooncake::testing::Provider(*client); + if (FLAGS_mode == "seed") return mooncake::testing::Seed(*client); + if (FLAGS_mode == "verify") return mooncake::testing::Verify(*client); + if (FLAGS_mode == "delete") return mooncake::testing::Delete(*client); + if (FLAGS_mode == "verify-absent") + return mooncake::testing::VerifyAbsent(*client); + return mooncake::testing::Pressure(*client); +} diff --git a/mooncake-store/tests/e2e/process_handler.cpp b/mooncake-store/tests/e2e/process_handler.cpp index 85f51e62c3..25d5ea2dfc 100644 --- a/mooncake-store/tests/e2e/process_handler.cpp +++ b/mooncake-store/tests/e2e/process_handler.cpp @@ -9,8 +9,10 @@ #include #include +#include #include #include +#include #include namespace mooncake { @@ -25,6 +27,56 @@ std::string ResolveMasterBackendConnstring(const MasterRunnerConfig& config) { return config.etcd_endpoints; } +bool SignalProcess(pid_t pid, int signo) { + if (pid == 0) { + return false; + } + if (::kill(-pid, signo) == 0) { + return true; + } + return errno == ESRCH && ::kill(pid, signo) == 0; +} + +bool WaitForProcess(pid_t& pid, std::chrono::milliseconds timeout, + int* status) { + if (pid == 0) { + return true; + } + const auto deadline = std::chrono::steady_clock::now() + timeout; + int child_status = 0; + do { + const pid_t result = waitpid(pid, &child_status, WNOHANG); + if (result == pid) { + pid = 0; + if (status != nullptr) { + *status = child_status; + } + return true; + } + if (result == -1 && errno != EINTR) { + if (errno == ECHILD) { + pid = 0; + return true; + } + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } while (std::chrono::steady_clock::now() < deadline); + return false; +} + +bool StopProcess(pid_t& pid, std::chrono::milliseconds timeout) { + if (pid == 0) { + return true; + } + SignalProcess(pid, SIGTERM); + if (WaitForProcess(pid, timeout, nullptr)) { + return true; + } + SignalProcess(pid, SIGKILL); + return WaitForProcess(pid, std::chrono::seconds(2), nullptr); +} + } // namespace MasterProcessHandler::MasterProcessHandler(const std::string& path, @@ -37,6 +89,12 @@ MasterProcessHandler::MasterProcessHandler(const std::string& path, .ha_backend_type = "etcd", .ha_backend_connstring = "", .etcd_endpoints = etcd_endpoints, + .cluster_id = "", + .oplog_store_type = "", + .metrics_port = 0, + .oplog_batch_max_entries = 1024, + .batch_oplog_retry_timeout_sec = 180, + .extra_args = {}, }), port_(port), index_(index), @@ -54,7 +112,7 @@ MasterProcessHandler::MasterProcessHandler(const std::string& path, MasterProcessHandler::~MasterProcessHandler() { if (master_pid_ != 0) { - kill(); + stop(std::chrono::seconds(2)); } } @@ -74,9 +132,8 @@ bool MasterProcessHandler::start() { return false; } - std::stringstream stdout_file, stderr_file; - stdout_file << out_dir_ << "/master_" + std::to_string(index_) + ".out"; - stderr_file << out_dir_ << "/master_" + std::to_string(index_) + ".err"; + stdout_path_ = out_dir_ + "/master_" + std::to_string(index_) + ".out"; + stderr_path_ = out_dir_ + "/master_" + std::to_string(index_) + ".err"; pid_t pid = fork(); @@ -88,6 +145,7 @@ bool MasterProcessHandler::start() { if (pid == 0) { // Child process + setpgid(0, 0); // Determine file opening mode based on whether it's the first start int open_flags = O_WRONLY | O_CREAT; @@ -98,19 +156,19 @@ bool MasterProcessHandler::start() { } // Open stdout file - int stdout_fd = open(stdout_file.str().c_str(), open_flags, 0644); + int stdout_fd = open(stdout_path_.c_str(), open_flags, 0644); if (stdout_fd == -1) { LOG(ERROR) << "[m" << index_ - << "] Failed to open stdout file: " << stdout_file.str() + << "] Failed to open stdout file: " << stdout_path_ << ", error: " << strerror(errno); exit(1); } // Open stderr file - int stderr_fd = open(stderr_file.str().c_str(), open_flags, 0644); + int stderr_fd = open(stderr_path_.c_str(), open_flags, 0644); if (stderr_fd == -1) { LOG(ERROR) << "[m" << index_ - << "] Failed to open stderr file: " << stderr_file.str() + << "] Failed to open stderr file: " << stderr_path_ << ", error: " << strerror(errno); close(stdout_fd); exit(1); @@ -153,8 +211,23 @@ bool MasterProcessHandler::start() { if (!config_.etcd_endpoints.empty()) { args.push_back("--etcd-endpoints=" + config_.etcd_endpoints); } + if (!config_.cluster_id.empty()) { + args.push_back("--cluster-id=" + config_.cluster_id); + } + if (!config_.oplog_store_type.empty()) { + args.push_back("--oplog-store-type=" + config_.oplog_store_type); + } + const uint32_t metrics_port = + config_.metrics_port == 0 ? 9003 + index_ : config_.metrics_port; + args.push_back("--metrics-port=" + std::to_string(metrics_port)); + args.push_back("--oplog-batch-max-entries=" + + std::to_string(config_.oplog_batch_max_entries)); + args.push_back("--batch-oplog-retry-timeout-sec=" + + std::to_string(config_.batch_oplog_retry_timeout_sec)); args.push_back(rpc_address_arg); args.push_back(rpc_port_arg); + args.insert(args.end(), config_.extra_args.begin(), + config_.extra_args.end()); std::vector argv; argv.reserve(args.size() + 1); @@ -180,6 +253,7 @@ bool MasterProcessHandler::start() { exit(1); } else { // Parent process - store the PID + setpgid(pid, pid); master_pid_ = pid; LOG(INFO) << "[m" << index_ << "] Started master process with PID: " << pid; @@ -199,15 +273,13 @@ bool MasterProcessHandler::kill() { bool success = false; // Kill the process forcefully to simulate a crash - if (::kill(master_pid_, SIGKILL) == 0) { + if (SignalProcess(master_pid_, SIGKILL)) { LOG(INFO) << "[m" << index_ << "] Force killed master process with PID: " << master_pid_ << " (simulating crash)"; // Wait for the process to be reaped - int status; - waitpid(master_pid_, &status, 0); - success = true; + success = WaitForProcess(master_pid_, std::chrono::seconds(2), nullptr); } else { LOG(ERROR) << TEST_ERROR_STR << " [m" << index_ << "] Failed to kill master process with PID: " @@ -215,11 +287,37 @@ bool MasterProcessHandler::kill() { success = false; } - master_pid_ = 0; + if (!success) { + master_pid_ = 0; + } return success; } -bool MasterProcessHandler::is_running() const { return master_pid_ != 0; } +bool MasterProcessHandler::stop(std::chrono::milliseconds timeout) { + return StopProcess(master_pid_, timeout); +} + +bool MasterProcessHandler::signal(int signo) { + return SignalProcess(master_pid_, signo); +} + +bool MasterProcessHandler::wait_for_exit(std::chrono::milliseconds timeout, + int* status) { + return WaitForProcess(master_pid_, timeout, status); +} + +bool MasterProcessHandler::is_running() const { + if (master_pid_ == 0) { + return false; + } + int status = 0; + const pid_t result = waitpid(master_pid_, &status, WNOHANG); + if (result == master_pid_ || (result == -1 && errno == ECHILD)) { + master_pid_ = 0; + return false; + } + return result == 0 || (result == -1 && errno == EINTR); +} ClientProcessHandler::ClientProcessHandler(const std::string& path, const int index, @@ -229,7 +327,7 @@ ClientProcessHandler::ClientProcessHandler(const std::string& path, ClientProcessHandler::~ClientProcessHandler() { if (client_pid_ != 0) { - kill(); + stop(std::chrono::seconds(2)); } } @@ -249,9 +347,8 @@ bool ClientProcessHandler::start() { return false; } - std::stringstream stdout_file, stderr_file; - stdout_file << out_dir_ << "/client_" + std::to_string(index_) + ".out"; - stderr_file << out_dir_ << "/client_" + std::to_string(index_) + ".err"; + stdout_path_ = out_dir_ + "/client_" + std::to_string(index_) + ".out"; + stderr_path_ = out_dir_ + "/client_" + std::to_string(index_) + ".err"; pid_t pid = fork(); @@ -263,6 +360,7 @@ bool ClientProcessHandler::start() { if (pid == 0) { // Child process + setpgid(0, 0); // Determine file opening mode based on whether it's the first start int open_flags = O_WRONLY | O_CREAT; @@ -273,19 +371,19 @@ bool ClientProcessHandler::start() { } // Open stdout file - int stdout_fd = open(stdout_file.str().c_str(), open_flags, 0644); + int stdout_fd = open(stdout_path_.c_str(), open_flags, 0644); if (stdout_fd == -1) { LOG(ERROR) << "[c" << index_ - << "] Failed to open stdout file: " << stdout_file.str() + << "] Failed to open stdout file: " << stdout_path_ << ", error: " << strerror(errno); exit(1); } // Open stderr file - int stderr_fd = open(stderr_file.str().c_str(), open_flags, 0644); + int stderr_fd = open(stderr_path_.c_str(), open_flags, 0644); if (stderr_fd == -1) { LOG(ERROR) << "[c" << index_ - << "] Failed to open stderr file: " << stderr_file.str() + << "] Failed to open stderr file: " << stderr_path_ << ", error: " << strerror(errno); close(stdout_fd); exit(1); @@ -375,6 +473,7 @@ bool ClientProcessHandler::start() { exit(1); } else { // Parent process - store the PID + setpgid(pid, pid); client_pid_ = pid; LOG(INFO) << "[c" << index_ << "] Started client process with PID: " << pid; @@ -393,15 +492,13 @@ bool ClientProcessHandler::kill() { bool success = false; // Kill the process forcefully to simulate a crash - if (::kill(client_pid_, SIGKILL) == 0) { + if (SignalProcess(client_pid_, SIGKILL)) { LOG(INFO) << "[c" << index_ << "] Force killed client process with PID: " << client_pid_ << " (simulating crash)"; // Wait for the process to be reaped - int status; - waitpid(client_pid_, &status, 0); - success = true; + success = WaitForProcess(client_pid_, std::chrono::seconds(2), nullptr); } else { LOG(ERROR) << TEST_ERROR_STR << " [c" << index_ << "] Failed to kill client process with PID: " @@ -409,11 +506,37 @@ bool ClientProcessHandler::kill() { success = false; } - client_pid_ = 0; + if (!success) { + client_pid_ = 0; + } return success; } -bool ClientProcessHandler::is_running() const { return client_pid_ != 0; } +bool ClientProcessHandler::stop(std::chrono::milliseconds timeout) { + return StopProcess(client_pid_, timeout); +} + +bool ClientProcessHandler::signal(int signo) { + return SignalProcess(client_pid_, signo); +} + +bool ClientProcessHandler::wait_for_exit(std::chrono::milliseconds timeout, + int* status) { + return WaitForProcess(client_pid_, timeout, status); +} + +bool ClientProcessHandler::is_running() const { + if (client_pid_ == 0) { + return false; + } + int status = 0; + const pid_t result = waitpid(client_pid_, &status, WNOHANG); + if (result == client_pid_ || (result == -1 && errno == ECHILD)) { + client_pid_ = 0; + return false; + } + return result == 0 || (result == -1 && errno == EINTR); +} } // namespace testing } // namespace mooncake diff --git a/mooncake-store/tests/e2e/process_handler.h b/mooncake-store/tests/e2e/process_handler.h index a92e183413..bff064089d 100644 --- a/mooncake-store/tests/e2e/process_handler.h +++ b/mooncake-store/tests/e2e/process_handler.h @@ -2,8 +2,11 @@ #include +#include +#include #include #include +#include namespace mooncake { namespace testing { @@ -26,6 +29,12 @@ struct MasterRunnerConfig { std::string ha_backend_type = "etcd"; std::string ha_backend_connstring; std::string etcd_endpoints; + std::string cluster_id{}; + std::string oplog_store_type{}; + uint32_t metrics_port{0}; + uint32_t oplog_batch_max_entries{1024}; + uint32_t batch_oplog_retry_timeout_sec{180}; + std::vector extra_args{}; }; /** @@ -66,12 +75,19 @@ class MasterProcessHandler { // Kill the master process. bool kill(); + bool stop(std::chrono::milliseconds timeout); + bool signal(int signo); + bool wait_for_exit(std::chrono::milliseconds timeout, int* status); + pid_t pid() const { return master_pid_; } + const std::string& stdout_path() const { return stdout_path_; } + const std::string& stderr_path() const { return stderr_path_; } + // Check if the process is started and is not killed yet. bool is_running() const; private: // The PID of the master process. - pid_t master_pid_{0}; + mutable pid_t master_pid_{0}; // The path to the master executable. std::string master_path_; // The startup configuration for the master process. @@ -84,6 +100,8 @@ class MasterProcessHandler { bool first_start_{true}; // The directory to store the log files. std::string out_dir_; + std::string stdout_path_; + std::string stderr_path_; }; /** @@ -136,18 +154,27 @@ class ClientProcessHandler { // Kill the client process. bool kill(); + bool stop(std::chrono::milliseconds timeout); + bool signal(int signo); + bool wait_for_exit(std::chrono::milliseconds timeout, int* status); + pid_t pid() const { return client_pid_; } + const std::string& stdout_path() const { return stdout_path_; } + const std::string& stderr_path() const { return stderr_path_; } + // Check if the process is started and is not killed yet. bool is_running() const; private: // The PID of the client process. - pid_t client_pid_{0}; + mutable pid_t client_pid_{0}; // The path to the client executable. std::string client_path_; // The index of the client, used for log. int index_; // The directory to store the log files. std::string out_dir_; + std::string stdout_path_; + std::string stderr_path_; // The start parameters for the client runner. ClientRunnerConfig config_; // Whether the client is started for the first time. diff --git a/mooncake-store/tests/e2e/run_oplog_batch_cluster.sh b/mooncake-store/tests/e2e/run_oplog_batch_cluster.sh new file mode 100755 index 0000000000..c6ea928bf8 --- /dev/null +++ b/mooncake-store/tests/e2e/run_oplog_batch_cluster.sh @@ -0,0 +1,1918 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +REPO_ROOT=$(cd -- "$SCRIPT_DIR/../../.." && pwd) + +die() { + echo "error: $*" >&2 + exit 1 +} + +usage() { + echo "Usage: $0 [options]" >&2 +} + +parse_up_options() { + BUILD_DIR="$REPO_ROOT/build" + RUN_DIR="" + MASTER_COUNT=3 + ENABLE_HA=true + CLIENT_COUNT=0 + CLUSTER_ID="oplog-test-$$" + PROTOCOL=tcp + OPLOG_STORE_TYPE=etcd_batch_record + BATCH_ENTRIES=1024 + RETRY_TIMEOUT_SEC=180 + MASTER_CONFIG="" + USE_ETCD_OBSERVER=true + NON_HA_WORKERS=4 + FAILPOINT_DIR="" + FAILPOINT_TIMEOUT_SEC=30 + START_TIMEOUT_SEC=30 + HA_OBJECTS=1000 + HA_PAYLOAD_BYTES=4096 + HA_PRESSURE_SEC=45 + MEMORY_ALLOCATOR="" + RECOVERY_PAYLOAD_SIZES=64,128,4096,65536,1048576 + RECOVERY_SEED_OBJECTS=240 + RECOVERY_REFILL_OBJECTS=120 + RECOVERY_PRESSURE_SEC=10 + RECOVERY_SEGMENT_BYTES=268435456 + ETCD_ENDPOINTS="" + ETCD_BIN=${ETCD_BIN:-etcd} + while (($#)); do + case "$1" in + --build-dir) + (($# >= 2)) || die "--build-dir requires a value" + BUILD_DIR=$2 + shift 2 + ;; + --run-dir) + (($# >= 2)) || die "--run-dir requires a value" + RUN_DIR=$2 + shift 2 + ;; + --masters) + (($# >= 2)) || die "--masters requires a value" + MASTER_COUNT=$2 + shift 2 + ;; + --clients) + (($# >= 2)) || die "--clients requires a value" + CLIENT_COUNT=$2 + shift 2 + ;; + --cluster-id) + (($# >= 2)) || die "--cluster-id requires a value" + CLUSTER_ID=$2 + shift 2 + ;; + --protocol) + (($# >= 2)) || die "--protocol requires a value" + PROTOCOL=$2 + shift 2 + ;; + --oplog-store) + (($# >= 2)) || die "--oplog-store requires a value" + OPLOG_STORE_TYPE=$2 + shift 2 + ;; + --batch-entries) + (($# >= 2)) || die "--batch-entries requires a value" + BATCH_ENTRIES=$2 + shift 2 + ;; + --retry-timeout-sec) + (($# >= 2)) || die "--retry-timeout-sec requires a value" + RETRY_TIMEOUT_SEC=$2 + shift 2 + ;; + --master-config) + (($# >= 2)) || die "--master-config requires a value" + MASTER_CONFIG=$2 + shift 2 + ;; + --no-etcd-observer) + USE_ETCD_OBSERVER=false + shift + ;; + --non-ha-workers) + (($# >= 2)) || die "--non-ha-workers requires a value" + NON_HA_WORKERS=$2 + shift 2 + ;; + --timeout-sec) + (($# >= 2)) || die "--timeout-sec requires a value" + START_TIMEOUT_SEC=$2 + shift 2 + ;; + --ha-objects) + (($# >= 2)) || die "--ha-objects requires a value" + HA_OBJECTS=$2 + shift 2 + ;; + --ha-payload-bytes) + (($# >= 2)) || die "--ha-payload-bytes requires a value" + HA_PAYLOAD_BYTES=$2 + shift 2 + ;; + --ha-pressure-sec) + (($# >= 2)) || die "--ha-pressure-sec requires a value" + HA_PRESSURE_SEC=$2 + shift 2 + ;; + --memory-allocator) + (($# >= 2)) || die "--memory-allocator requires a value" + MEMORY_ALLOCATOR=$2 + shift 2 + ;; + --recovery-payload-sizes) + (($# >= 2)) || die "--recovery-payload-sizes requires a value" + RECOVERY_PAYLOAD_SIZES=$2 + shift 2 + ;; + --recovery-seed-objects) + (($# >= 2)) || die "--recovery-seed-objects requires a value" + RECOVERY_SEED_OBJECTS=$2 + shift 2 + ;; + --recovery-refill-objects) + (($# >= 2)) || die "--recovery-refill-objects requires a value" + RECOVERY_REFILL_OBJECTS=$2 + shift 2 + ;; + --recovery-pressure-sec) + (($# >= 2)) || die "--recovery-pressure-sec requires a value" + RECOVERY_PRESSURE_SEC=$2 + shift 2 + ;; + --recovery-segment-bytes) + (($# >= 2)) || die "--recovery-segment-bytes requires a value" + RECOVERY_SEGMENT_BYTES=$2 + shift 2 + ;; + --failpoint-dir) + (($# >= 2)) || die "--failpoint-dir requires a value" + FAILPOINT_DIR=$2 + shift 2 + ;; + --failpoint-timeout-sec) + (($# >= 2)) || die "--failpoint-timeout-sec requires a value" + FAILPOINT_TIMEOUT_SEC=$2 + shift 2 + ;; + --etcd-endpoints) + (($# >= 2)) || die "--etcd-endpoints requires a value" + ETCD_ENDPOINTS=$2 + shift 2 + ;; + --etcd-bin) + (($# >= 2)) || die "--etcd-bin requires a value" + ETCD_BIN=$2 + shift 2 + ;; + *) die "unknown option: $1" ;; + esac + done + [[ -d "$BUILD_DIR" ]] || die "build directory does not exist: $BUILD_DIR" + [[ "$MASTER_COUNT" =~ ^[1-9][0-9]*$ ]] || die "masters must be positive" + [[ "$CLIENT_COUNT" =~ ^[0-9]+$ ]] || die "clients must be non-negative" + [[ "$BATCH_ENTRIES" =~ ^[1-9][0-9]*$ ]] || + die "batch-entries must be positive" + [[ "$START_TIMEOUT_SEC" =~ ^[1-9][0-9]*$ ]] || + die "timeout-sec must be positive" + [[ "$FAILPOINT_TIMEOUT_SEC" =~ ^[1-9][0-9]*$ ]] || + die "failpoint-timeout-sec must be positive" + [[ "$NON_HA_WORKERS" =~ ^[1-9][0-9]*$ ]] || + die "non-ha-workers must be positive" + [[ "$HA_OBJECTS" =~ ^[1-9][0-9]*$ ]] || die "ha-objects must be positive" + [[ "$HA_PAYLOAD_BYTES" =~ ^[1-9][0-9]*$ ]] || + die "ha-payload-bytes must be positive" + [[ "$HA_PRESSURE_SEC" =~ ^[1-9][0-9]*$ ]] || + die "ha-pressure-sec must be positive" + [[ -z "$MEMORY_ALLOCATOR" || "$MEMORY_ALLOCATOR" == offset || + "$MEMORY_ALLOCATOR" == cachelib ]] || + die "memory-allocator must be offset or cachelib" + [[ "$RECOVERY_SEED_OBJECTS" =~ ^[1-9][0-9]*$ ]] || + die "recovery-seed-objects must be positive" + [[ "$RECOVERY_REFILL_OBJECTS" =~ ^[1-9][0-9]*$ ]] || + die "recovery-refill-objects must be positive" + [[ "$RECOVERY_PRESSURE_SEC" =~ ^[1-9][0-9]*$ ]] || + die "recovery-pressure-sec must be positive" + [[ "$RECOVERY_SEGMENT_BYTES" =~ ^[1-9][0-9]*$ ]] || + die "recovery-segment-bytes must be positive" + [[ "$RECOVERY_PAYLOAD_SIZES" =~ ^[1-9][0-9]*(,[1-9][0-9]*)*$ ]] || + die "recovery-payload-sizes must be positive CSV integers" + [[ "$PROTOCOL" == tcp || "$PROTOCOL" == rdma ]] || + die "protocol must be tcp or rdma" + [[ -z "$MASTER_CONFIG" || -f "$MASTER_CONFIG" ]] || + die "master config does not exist: $MASTER_CONFIG" + if [[ -z "$RUN_DIR" ]]; then + RUN_DIR="/tmp/mooncake-oplog-test/$(date +%Y%m%d-%H%M%S)-$$" + fi +} + +require_executable() { + [[ -x "$1" ]] || die "missing executable: $1" +} + +up_cluster() { + MASTER_BIN="$BUILD_DIR/mooncake-store/src/mooncake_master" + INSPECTOR_BIN="$BUILD_DIR/mooncake-store/tools/oplog_batch_inspector" + METADATA_SCRIPT="$REPO_ROOT/mooncake-wheel/mooncake/http_metadata_server.py" + require_executable "$MASTER_BIN" + [[ "$USE_ETCD_OBSERVER" != true ]] || require_executable "$INSPECTOR_BIN" + [[ -f "$METADATA_SCRIPT" ]] || die "missing executable: $METADATA_SCRIPT" + command -v python3 >/dev/null || die "missing executable: python3" + command -v curl >/dev/null || die "missing executable: curl" + command -v setsid >/dev/null || die "missing executable: setsid" + if [[ "$USE_ETCD_OBSERVER" == true && -z "$ETCD_ENDPOINTS" ]]; then + command -v "$ETCD_BIN" >/dev/null || die "missing executable: $ETCD_BIN" + fi + python3 -c 'import aiohttp' >/dev/null 2>&1 || + die "missing Python dependency aiohttp; activate the project Python environment" + + mkdir -p "$RUN_DIR"/{audit,configs,etcd,logs,metrics,pids,workload} + METADATA_PORT=$(find_free_port) + local etcd_client_port + local etcd_peer_port + if [[ "$USE_ETCD_OBSERVER" == true && -z "$ETCD_ENDPOINTS" ]]; then + etcd_client_port=$(find_free_port) + etcd_peer_port=$(find_free_port) + ETCD_ENDPOINTS="127.0.0.1:$etcd_client_port" + start_process etcd \ + "$ETCD_BIN" --name "$CLUSTER_ID-etcd" \ + --data-dir "$RUN_DIR/etcd/data" \ + --listen-client-urls "http://127.0.0.1:$etcd_client_port" \ + --advertise-client-urls "http://127.0.0.1:$etcd_client_port" \ + --listen-peer-urls "http://127.0.0.1:$etcd_peer_port" \ + --initial-advertise-peer-urls "http://127.0.0.1:$etcd_peer_port" \ + --initial-cluster "$CLUSTER_ID-etcd=http://127.0.0.1:$etcd_peer_port" + wait_http "http://127.0.0.1:$etcd_client_port/health" "$START_TIMEOUT_SEC" || + die "local etcd did not become healthy; see $RUN_DIR/logs/etcd.err" + fi + + start_process metadata python3 "$METADATA_SCRIPT" --host 127.0.0.1 \ + --port "$METADATA_PORT" + wait_port 127.0.0.1 "$METADATA_PORT" "$START_TIMEOUT_SEC" || + die "metadata server did not start; see $RUN_DIR/logs/metadata.err" + + RPC_PORTS=() + ADMIN_PORTS=() + local index + for ((index = 0; index < MASTER_COUNT; ++index)); do + local rpc_port + local admin_port + rpc_port=$(find_free_port) + admin_port=$(find_free_port) + RPC_PORTS+=("$rpc_port") + ADMIN_PORTS+=("$admin_port") + start_master "$index" + done + + write_cluster_env + write_manifest + wait_for_single_leader "$START_TIMEOUT_SEC" || { + collect_cluster || true + die "cluster did not elect exactly one ready leader; see $RUN_DIR/logs" + } + + CLIENT_BIN="$BUILD_DIR/mooncake-store/tests/e2e/client_runner" + if ((CLIENT_COUNT > 0)); then + require_executable "$CLIENT_BIN" + for ((index = 0; index < CLIENT_COUNT; ++index)); do + start_process "client-$index" "$CLIENT_BIN" \ + --port="$((17812 + index))" \ + --master_server_entry="etcd://$ETCD_ENDPOINTS" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" + done + fi + echo "cluster started: $RUN_DIR" + status_cluster +} + +find_free_port() { + python3 -c 'import socket; s=socket.socket(); s.bind(("127.0.0.1", 0)); print(s.getsockname()[1]); s.close()' +} + +wait_http() { + local url=$1 + local timeout=$2 + local deadline=$((SECONDS + timeout)) + while ((SECONDS < deadline)); do + curl -fsS "$url" >/dev/null 2>&1 && return 0 + sleep 0.1 + done + return 1 +} + +wait_port() { + local host=$1 + local port=$2 + local timeout=$3 + python3 -c 'import socket,sys,time +host,port,timeout=sys.argv[1],int(sys.argv[2]),float(sys.argv[3]) +deadline=time.time()+timeout +while time.time()>"$RUN_DIR/logs/$name.out" \ + 2>>"$RUN_DIR/logs/$name.err" & + local pid=$! + printf '%s\n' "$pid" >"$RUN_DIR/pids/$name.pid" + local deadline=$((SECONDS + 2)) + while [[ ! -r "/proc/$pid/cmdline" ]] && ((SECONDS < deadline)); do + sleep 0.05 + done + sleep 0.05 + if kill -0 "$pid" 2>/dev/null && [[ -r "/proc/$pid/cmdline" ]]; then + tr '\0' ' ' <"/proc/$pid/cmdline" >"$RUN_DIR/pids/$name.cmd" + else + wait "$pid" || true + die "$name exited during startup; see $RUN_DIR/logs/$name.err" + fi +} + +start_master() { + local index=$1 + local -a environment=() + if [[ -n "$FAILPOINT_DIR" ]]; then + mkdir -p "$FAILPOINT_DIR" + environment=(env "MOONCAKE_TEST_FAILPOINT_DIR=$FAILPOINT_DIR" + "MOONCAKE_TEST_FAILPOINT_TIMEOUT_SEC=$FAILPOINT_TIMEOUT_SEC") + fi + local -a ha_args=(--enable_ha=false) + local -a non_ha_args=() + local -a config_args=() + local -a allocator_args=() + [[ -z "$MASTER_CONFIG" ]] || config_args=(--config_path="$MASTER_CONFIG") + [[ -z "$MEMORY_ALLOCATOR" ]] || + allocator_args=(--memory_allocator="$MEMORY_ALLOCATOR") + if [[ "$ENABLE_HA" == true ]]; then + ha_args=(--enable_ha=true --ha_backend_type=etcd + --ha_backend_connstring="$ETCD_ENDPOINTS" + --etcd_endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" + --enable_oplog=true + --oplog_batch_max_entries="$BATCH_ENTRIES" + --batch_oplog_retry_timeout_sec="$RETRY_TIMEOUT_SEC") + else + non_ha_args=(--default_kv_lease_ttl=5s) + fi + start_process "master-$index" "${environment[@]}" "$MASTER_BIN" \ + "${config_args[@]}" "${ha_args[@]}" "${non_ha_args[@]}" \ + "${allocator_args[@]}" \ + --rpc_address=127.0.0.1 --rpc_port="${RPC_PORTS[$index]}" \ + --metrics_port="${ADMIN_PORTS[$index]}" \ + --enable_http_metadata_server=false --logtostderr=true +} + +write_cluster_env() { + { + printf 'RUN_DIR=%q\n' "$RUN_DIR" + printf 'BUILD_DIR=%q\n' "$BUILD_DIR" + printf 'CLUSTER_ID=%q\n' "$CLUSTER_ID" + printf 'ETCD_ENDPOINTS=%q\n' "$ETCD_ENDPOINTS" + printf 'METADATA_PORT=%q\n' "$METADATA_PORT" + printf 'MASTER_COUNT=%q\n' "$MASTER_COUNT" + printf 'ENABLE_HA=%q\n' "$ENABLE_HA" + printf 'PROTOCOL=%q\n' "$PROTOCOL" + printf 'OPLOG_STORE_TYPE=%q\n' "$OPLOG_STORE_TYPE" + printf 'BATCH_ENTRIES=%q\n' "$BATCH_ENTRIES" + printf 'RETRY_TIMEOUT_SEC=%q\n' "$RETRY_TIMEOUT_SEC" + printf 'MASTER_CONFIG=%q\n' "$MASTER_CONFIG" + printf 'USE_ETCD_OBSERVER=%q\n' "$USE_ETCD_OBSERVER" + printf 'NON_HA_WORKERS=%q\n' "$NON_HA_WORKERS" + printf 'START_TIMEOUT_SEC=%q\n' "$START_TIMEOUT_SEC" + printf 'FAILPOINT_DIR=%q\n' "$FAILPOINT_DIR" + printf 'FAILPOINT_TIMEOUT_SEC=%q\n' "$FAILPOINT_TIMEOUT_SEC" + printf 'MEMORY_ALLOCATOR=%q\n' "$MEMORY_ALLOCATOR" + printf 'RPC_PORTS=%q\n' "${RPC_PORTS[*]}" + printf 'ADMIN_PORTS=%q\n' "${ADMIN_PORTS[*]}" + printf 'INSPECTOR_BIN=%q\n' "$INSPECTOR_BIN" + } >"$RUN_DIR/cluster.env" +} + +write_manifest() { + python3 -c 'import json,subprocess,sys,time +path,run_dir,build,cluster,endpoints,masters=sys.argv[1:] +try: commit=subprocess.check_output(["git","rev-parse","HEAD"],text=True).strip() +except Exception: commit="unknown" +with open(path,"w",encoding="utf-8") as f: + json.dump({"schema_version":1,"run_id":run_dir.rsplit("/",1)[-1],"repo_commit":commit,"build_dir":build,"cluster_id":cluster,"etcd_endpoints":endpoints,"masters":int(masters),"created_at_epoch":time.time()},f,sort_keys=True) + f.write("\n")' "$RUN_DIR/manifest.json" "$RUN_DIR" "$BUILD_DIR" \ + "$CLUSTER_ID" "$ETCD_ENDPOINTS" "$MASTER_COUNT" +} + +load_cluster_env() { + [[ -f "$RUN_DIR/cluster.env" ]] || die "missing cluster.env: $RUN_DIR" + # This file is generated by this script with shell-escaped values. + source "$RUN_DIR/cluster.env" + read -r -a RPC_PORTS <<<"$RPC_PORTS" + read -r -a ADMIN_PORTS <<<"$ADMIN_PORTS" + MASTER_COUNT=${MASTER_COUNT:-${#RPC_PORTS[@]}} + ENABLE_HA=${ENABLE_HA:-true} + PROTOCOL=${PROTOCOL:-tcp} + OPLOG_STORE_TYPE=${OPLOG_STORE_TYPE:-etcd_batch_record} + BATCH_ENTRIES=${BATCH_ENTRIES:-1024} + RETRY_TIMEOUT_SEC=${RETRY_TIMEOUT_SEC:-180} + MASTER_CONFIG=${MASTER_CONFIG:-} + USE_ETCD_OBSERVER=${USE_ETCD_OBSERVER:-true} + NON_HA_WORKERS=${NON_HA_WORKERS:-4} + START_TIMEOUT_SEC=${START_TIMEOUT_SEC:-30} + FAILPOINT_DIR=${FAILPOINT_DIR:-} + FAILPOINT_TIMEOUT_SEC=${FAILPOINT_TIMEOUT_SEC:-30} + MEMORY_ALLOCATOR=${MEMORY_ALLOCATOR:-} + MASTER_BIN="$BUILD_DIR/mooncake-store/src/mooncake_master" +} + +restart_masters() { + require_executable "$MASTER_BIN" + local pid_file + shopt -s nullglob + for pid_file in "$RUN_DIR"/pids/master-*.pid; do + stop_pid_file "$pid_file" + done + local index + for ((index = 0; index < MASTER_COUNT; ++index)); do + start_master "$index" + done + wait_for_single_leader "$START_TIMEOUT_SEC" || { + collect_cluster || true + die "restarted cluster did not elect exactly one ready leader; see $RUN_DIR/logs" + } + echo "masters restarted: $RUN_DIR" +} + +wait_for_single_leader() { + local timeout=$1 + local deadline=$((SECONDS + timeout)) + while ((SECONDS < deadline)); do + local ready=0 + local port + for port in "${ADMIN_PORTS[@]}"; do + local health + health=$(curl -fsS "http://127.0.0.1:$port/health" 2>/dev/null || true) + grep -Eq '"service_ready"[[:space:]]*:[[:space:]]*true' <<<"$health" && + ready=$((ready + 1)) + done + ((ready == 1)) && return 0 + sleep 0.2 + done + return 1 +} + +status_cluster() { + [[ ${#ADMIN_PORTS[@]} -gt 0 ]] || load_cluster_env + echo "run_dir=$RUN_DIR cluster_id=$CLUSTER_ID etcd=$ETCD_ENDPOINTS" + local index + for index in "${!ADMIN_PORTS[@]}"; do + local health + health=$(curl -fsS "http://127.0.0.1:${ADMIN_PORTS[$index]}/health" 2>/dev/null || echo unavailable) + echo "master-$index rpc=127.0.0.1:${RPC_PORTS[$index]} admin=127.0.0.1:${ADMIN_PORTS[$index]} $health" + done + if [[ "$USE_ETCD_OBSERVER" == true ]]; then + LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" summary \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json || true + fi +} + +collect_cluster() { + [[ ${#ADMIN_PORTS[@]} -gt 0 ]] || load_cluster_env + local index + for index in "${!ADMIN_PORTS[@]}"; do + curl -fsS "http://127.0.0.1:${ADMIN_PORTS[$index]}/health" \ + >"$RUN_DIR/metrics/master-$index-health.json" 2>/dev/null || true + curl -fsS "http://127.0.0.1:${ADMIN_PORTS[$index]}/metrics" \ + >"$RUN_DIR/metrics/master-$index.prom" 2>/dev/null || true + done + if [[ "$USE_ETCD_OBSERVER" == true ]]; then + LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" summary \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/summary.json" || true + LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/verify.json" || true + fi + echo "artifacts collected: $RUN_DIR" +} + +run_client_smoke() { + local label=$1 + local hold_seconds=${2:-0} + local clientctl="$BUILD_DIR/mooncake-store/tests/e2e/clientctl" + require_executable "$clientctl" + local client_port master_entry + client_port=$(find_free_port) + master_entry="etcd://$ETCD_ENDPOINTS" + [[ "$ENABLE_HA" == true ]] || master_entry="127.0.0.1:${RPC_PORTS[0]}" + if ! { + printf 'create %s %s\n' "$label" "$client_port" + printf 'mount %s segment-%s %s\n' "$label" "$label" \ + "$((64 * 1024 * 1024))" + printf 'put %s oplog-%s-%s value-%s\n' \ + "$label" "$label" "$CLUSTER_ID" "$label" + printf 'get %s oplog-%s-%s\n' "$label" "$label" "$CLUSTER_ID" + if ((hold_seconds > 0)); then + printf 'sleep %s\n' "$hold_seconds" + fi + printf 'terminate\n' + } | MC_STORE_CLUSTER_ID="$CLUSTER_ID" LSAN_OPTIONS=detect_leaks=0 \ + "$clientctl" \ + --master_server_entry="$master_entry" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" >"$RUN_DIR/workload/$label.log" 2>&1; then + return 1 + fi + grep -Fq "Successfully put value" "$RUN_DIR/workload/$label.log" && + grep -Fq "Get value: value-$label" "$RUN_DIR/workload/$label.log" +} + +read_durable_sequence() { + LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" summary \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json | + python3 -c 'import json,sys; print(json.load(sys.stdin)["durable_prefix"]["last_seq"])' +} + +ready_leader_index() { + local index health found="" + for index in "${!ADMIN_PORTS[@]}"; do + health=$(curl --max-time 1 -fsS \ + "http://127.0.0.1:${ADMIN_PORTS[$index]}/health" 2>/dev/null || true) + if grep -Eq '"service_ready"[[:space:]]*:[[:space:]]*true' \ + <<<"$health"; then + [[ -z "$found" ]] || return 1 + found=$index + fi + done + [[ -n "$found" ]] || return 1 + printf '%s\n' "$found" +} + +capture_ha_stage() { + local stage=$1 index + mkdir -p "$RUN_DIR/audit/$stage" "$RUN_DIR/metrics/$stage" + for index in "${!ADMIN_PORTS[@]}"; do + curl --max-time 1 -fsS \ + "http://127.0.0.1:${ADMIN_PORTS[$index]}/health" \ + >"$RUN_DIR/metrics/$stage/master-$index-health.json" 2>/dev/null || true + curl --max-time 1 -fsS \ + "http://127.0.0.1:${ADMIN_PORTS[$index]}/metrics" \ + >"$RUN_DIR/metrics/$stage/master-$index.prom" 2>/dev/null || true + done + LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" summary \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/$stage/summary.json" + LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/$stage/verify.json" +} + +smoke_cluster() { + up_cluster + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/before-workload.json"; then + smoke_failed "initial OpLog verification failed" + return 1 + fi + + if ! run_client_smoke smoke; then + smoke_failed "minimal put/get workload failed" + return 1 + fi + + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" wait \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" \ + --last_seq=1 --timeout_sec="$START_TIMEOUT_SEC"; then + smoke_failed "OpLog did not become durable after workload" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/after-workload.json"; then + smoke_failed "final OpLog verification failed" + return 1 + fi + + collect_cluster + down_cluster + echo "PASS: smoke completed; artifacts: $RUN_DIR" +} + +restart_smoke_cluster() { + up_cluster + if ! run_client_smoke before-restart; then + smoke_failed "pre-restart put/get workload failed" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" wait \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" \ + --last_seq=1 --timeout_sec="$START_TIMEOUT_SEC"; then + smoke_failed "pre-restart OpLog did not become durable" + return 1 + fi + local before_sequence + before_sequence=$(read_durable_sequence) || { + smoke_failed "failed to read pre-restart durable sequence" + return 1 + } + + restart_masters + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/after-restart.json"; then + smoke_failed "existing OpLog history failed verification after restart" + return 1 + fi + if ! run_client_smoke after-restart; then + smoke_failed "post-restart put/get workload failed" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" wait \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" \ + --last_seq="$((before_sequence + 1))" \ + --timeout_sec="$START_TIMEOUT_SEC"; then + smoke_failed "durable sequence did not advance after restart" + return 1 + fi + local after_sequence + after_sequence=$(read_durable_sequence) || { + smoke_failed "failed to read post-restart durable sequence" + return 1 + } + if ((after_sequence <= before_sequence)); then + smoke_failed "durable sequence did not increase after restart" + return 1 + fi + + collect_cluster + down_cluster + echo "PASS: restart smoke completed; before_seq=$before_sequence after_seq=$after_sequence artifacts: $RUN_DIR" +} + +failpoint_smoke_cluster() { + [[ -n "$FAILPOINT_DIR" ]] || FAILPOINT_DIR="$RUN_DIR/failpoints" + local fault_ctl="$SCRIPT_DIR/oplog_fault_ctl.sh" + require_executable "$fault_ctl" + up_cluster + + local before_sequence + before_sequence=$(read_durable_sequence) || { + smoke_failed "failed to read initial durable sequence" + return 1 + } + local name=batch_txn_succeeded_before_callback + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint arm \ + --run-dir "$RUN_DIR" --name "$name" + + run_client_smoke failpoint-smoke & + local client_pid=$! + if ! FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint wait \ + --run-dir "$RUN_DIR" --name "$name" \ + --timeout-sec "$START_TIMEOUT_SEC"; then + wait "$client_pid" || true + smoke_failed "failpoint was not hit; ensure the build enables MOONCAKE_ENABLE_TEST_FAILPOINTS" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/during-failpoint.json"; then + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$name" || true + wait "$client_pid" || true + smoke_failed "OpLog verification failed while writer was paused" + return 1 + fi + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$name" + if ! wait "$client_pid"; then + smoke_failed "client workload failed after failpoint release" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" wait \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" \ + --last_seq="$((before_sequence + 1))" \ + --timeout_sec="$START_TIMEOUT_SEC"; then + smoke_failed "durable sequence did not advance after failpoint release" + return 1 + fi + local after_sequence + after_sequence=$(read_durable_sequence) || { + smoke_failed "failed to read final durable sequence" + return 1 + } + + collect_cluster + down_cluster + echo "PASS: failpoint smoke completed; before_seq=$before_sequence after_seq=$after_sequence artifacts: $RUN_DIR" +} + +failpoint_crash_smoke_cluster() { + [[ -n "$FAILPOINT_DIR" ]] || FAILPOINT_DIR="$RUN_DIR/failpoints" + local fault_ctl="$SCRIPT_DIR/oplog_fault_ctl.sh" + require_executable "$fault_ctl" + up_cluster + + local before_sequence + before_sequence=$(read_durable_sequence) || { + smoke_failed "failed to read initial durable sequence" + return 1 + } + local name=batch_txn_succeeded_before_callback + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint arm \ + --run-dir "$RUN_DIR" --name "$name" + run_client_smoke before-leader-crash & + local client_pid=$! + if ! FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint wait \ + --run-dir "$RUN_DIR" --name "$name" \ + --timeout-sec "$START_TIMEOUT_SEC"; then + wait "$client_pid" || true + smoke_failed "failpoint was not hit; ensure the build enables MOONCAKE_ENABLE_TEST_FAILPOINTS" + return 1 + fi + + local hit_pid + hit_pid=$(<"$FAILPOINT_DIR/$name.hit") + local master_name="" + local pid_file + for pid_file in "$RUN_DIR"/pids/master-*.pid; do + if [[ $(<"$pid_file") == "$hit_pid" ]]; then + master_name=$(basename "$pid_file" .pid) + break + fi + done + if [[ -z "$master_name" ]]; then + wait "$client_pid" || true + smoke_failed "failpoint PID does not match a managed master: $hit_pid" + return 1 + fi + "$fault_ctl" process kill --run-dir "$RUN_DIR" --name "$master_name" + wait "$client_pid" || true + + if ! wait_for_single_leader "$START_TIMEOUT_SEC"; then + smoke_failed "cluster did not elect a replacement leader" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/after-leader-crash.json"; then + smoke_failed "durable OpLog failed verification after leader crash" + return 1 + fi + if ! run_client_smoke after-leader-crash; then + smoke_failed "replacement leader did not complete a new client workload" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" wait \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" \ + --last_seq="$((before_sequence + 2))" \ + --timeout_sec="$START_TIMEOUT_SEC"; then + smoke_failed "durable sequence did not advance after leader recovery" + return 1 + fi + local after_sequence + after_sequence=$(read_durable_sequence) || { + smoke_failed "failed to read recovered durable sequence" + return 1 + } + + collect_cluster + down_cluster + echo "PASS: failpoint crash smoke completed; killed=$master_name before_seq=$before_sequence after_seq=$after_sequence artifacts: $RUN_DIR" +} + +remove_boundary_smoke_cluster() { + [[ -n "$FAILPOINT_DIR" ]] || FAILPOINT_DIR="$RUN_DIR/failpoints" + local fault_ctl="$SCRIPT_DIR/oplog_fault_ctl.sh" + local clientctl="$BUILD_DIR/mooncake-store/tests/e2e/clientctl" + require_executable "$fault_ctl" + require_executable "$clientctl" + up_cluster + + local fifo="$RUN_DIR/workload/remove-boundary.fifo" + local log="$RUN_DIR/workload/remove-boundary.log" + local client_port key name + client_port=$(find_free_port) + key="remove-boundary-$CLUSTER_ID" + name=remove-boundary + mkfifo "$fifo" + exec 3<>"$fifo" + MC_STORE_CLUSTER_ID="$CLUSTER_ID" LSAN_OPTIONS=detect_leaks=0 \ + "$clientctl" --master_server_entry="etcd://$ETCD_ENDPOINTS" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" <"$fifo" >"$log" 2>&1 & + local client_pid=$! + + printf 'create %s %s\nmount %s segment-%s %s\nput %s %s value-before-delete\n' \ + "$name" "$client_port" "$name" "$name" "$((64 * 1024 * 1024))" \ + "$name" "$key" >&3 + if ! wait_file_text "$log" "Successfully put value for key: $key" \ + "$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "failed to prepare object for remove boundary test" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" wait \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" \ + --last_seq=1 --timeout_sec="$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "initial put did not become durable" + return 1 + fi + local before_sequence + before_sequence=$(read_durable_sequence) || { + stop_boundary_client "$client_pid" + smoke_failed "failed to read durable sequence after initial put" + return 1 + } + + local point=batch_before_txn + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint arm \ + --run-dir "$RUN_DIR" --name "$point" + printf 'delete %s %s\n' "$name" "$key" >&3 + if ! FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint wait \ + --run-dir "$RUN_DIR" --name "$point" \ + --timeout-sec "$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "pre-transaction failpoint was not hit; ensure the build enables MOONCAKE_ENABLE_TEST_FAILPOINTS" + return 1 + fi + printf 'get %s %s\n' "$name" "$key" >&3 + if ! wait_file_text "$log" "Failed to get value: OBJECT_NOT_FOUND" \ + "$START_TIMEOUT_SEC"; then + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + stop_boundary_client "$client_pid" + smoke_failed "removed object remained visible before durable commit" + return 1 + fi + local during_sequence + during_sequence=$(read_durable_sequence) || { + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + stop_boundary_client "$client_pid" + smoke_failed "failed to read durable sequence while failpoint was held" + return 1 + } + if [[ "$during_sequence" != "$before_sequence" ]]; then + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + stop_boundary_client "$client_pid" + smoke_failed "durable sequence advanced while pre-transaction failpoint was held" + return 1 + fi + + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" wait \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" \ + --last_seq="$((before_sequence + 1))" \ + --timeout_sec="$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "remove did not become durable after failpoint release" + return 1 + fi + stop_boundary_client "$client_pid" + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/after-remove-boundary.json"; then + smoke_failed "OpLog verification failed after durable remove" + return 1 + fi + local after_sequence + after_sequence=$(read_durable_sequence) || { + smoke_failed "failed to read durable sequence after remove" + return 1 + } + + collect_cluster + down_cluster + echo "PASS: remove boundary smoke completed; before_seq=$before_sequence during_seq=$during_sequence after_seq=$after_sequence artifacts: $RUN_DIR" +} + +standby_read_smoke_cluster() { + [[ -n "$FAILPOINT_DIR" ]] || FAILPOINT_DIR="$RUN_DIR/failpoints" + local fault_ctl="$SCRIPT_DIR/oplog_fault_ctl.sh" + require_executable "$fault_ctl" + up_cluster + + local point=standby_prefix_read_before_batch + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint arm \ + --run-dir "$RUN_DIR" --name "$point" + run_client_smoke standby-read & + local client_pid=$! + if ! FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint wait \ + --run-dir "$RUN_DIR" --name "$point" \ + --timeout-sec "$START_TIMEOUT_SEC"; then + wait "$client_pid" || true + smoke_failed "standby read failpoint was not hit; ensure the build enables MOONCAKE_ENABLE_TEST_FAILPOINTS" + return 1 + fi + if ! wait "$client_pid"; then + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + smoke_failed "client workload failed while standby reader was paused" + return 1 + fi + + local hit_pid master_index="" + hit_pid=$(<"$FAILPOINT_DIR/$point.hit") + local index + for index in "${!RPC_PORTS[@]}"; do + if [[ $(<"$RUN_DIR/pids/master-$index.pid") == "$hit_pid" ]]; then + master_index=$index + break + fi + done + if [[ -z "$master_index" ]]; then + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + smoke_failed "failpoint PID does not match a managed master: $hit_pid" + return 1 + fi + local health + health=$(curl -fsS "http://127.0.0.1:${ADMIN_PORTS[$master_index]}/health") + if ! grep -Eq '"role"[[:space:]]*:[[:space:]]*"standby"' <<<"$health"; then + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + smoke_failed "standby reader failpoint was hit by non-standby master-$master_index" + return 1 + fi + + local durable_sequence applied_during + durable_sequence=$(read_durable_sequence) || { + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + smoke_failed "failed to read durable sequence while standby was paused" + return 1 + } + applied_during=$(read_master_metric "$master_index" \ + ha_oplog_applied_sequence_id) || { + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + smoke_failed "failed to read paused standby applied sequence" + return 1 + } + if ((applied_during >= durable_sequence)); then + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + smoke_failed "paused standby was not behind durable sequence" + return 1 + fi + + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" + if ! wait_master_metric_at_least "$master_index" \ + ha_oplog_applied_sequence_id "$durable_sequence" "$START_TIMEOUT_SEC"; then + smoke_failed "standby did not catch up after failpoint release" + return 1 + fi + local applied_after + applied_after=$(read_master_metric "$master_index" \ + ha_oplog_applied_sequence_id) || return 1 + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/after-standby-read.json"; then + smoke_failed "OpLog verification failed after standby catch-up" + return 1 + fi + + collect_cluster + down_cluster + echo "PASS: standby read smoke completed; master=master-$master_index durable_seq=$durable_sequence applied_during=$applied_during applied_after=$applied_after artifacts: $RUN_DIR" +} + +promotion_catchup_smoke_cluster() { + [[ -n "$FAILPOINT_DIR" ]] || FAILPOINT_DIR="$RUN_DIR/failpoints" + local fault_ctl="$SCRIPT_DIR/oplog_fault_ctl.sh" + local clientctl="$BUILD_DIR/mooncake-store/tests/e2e/clientctl" + require_executable "$fault_ctl" + require_executable "$clientctl" + up_cluster + + local fifo="$RUN_DIR/workload/promotion-business.fifo" + local log="$RUN_DIR/workload/promotion-business.log" + local client_port client_name=promotion-business + local keep_key="promotion-keep-$CLUSTER_ID" + local removed_key="promotion-removed-$CLUSTER_ID" + client_port=$(find_free_port) + mkfifo "$fifo" + exec 3<>"$fifo" + MC_STORE_CLUSTER_ID="$CLUSTER_ID" LSAN_OPTIONS=detect_leaks=0 \ + "$clientctl" --master_server_entry="etcd://$ETCD_ENDPOINTS" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" <"$fifo" >"$log" 2>&1 & + local client_pid=$! + printf 'create %s %s\nmount %s segment-%s %s\n' \ + "$client_name" "$client_port" "$client_name" "$client_name" \ + "$((64 * 1024 * 1024))" >&3 + printf 'put %s %s keep-value\nput %s %s removed-value\n' \ + "$client_name" "$keep_key" "$client_name" "$removed_key" >&3 + printf 'sleep 6\ndelete %s %s\nget %s %s\n' \ + "$client_name" "$removed_key" "$client_name" "$removed_key" >&3 + if ! wait_file_text "$log" \ + "Successfully deleted value for key: $removed_key" \ + "$START_TIMEOUT_SEC" || + ! wait_file_count "$log" "Failed to get value: OBJECT_NOT_FOUND" 1 \ + "$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "failed to prepare promotion business state" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" wait \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" \ + --last_seq=1 --timeout_sec="$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "pre-promotion workload did not become durable" + return 1 + fi + local durable_before leader_index="" + durable_before=$(read_durable_sequence) || { + stop_boundary_client "$client_pid" + smoke_failed "failed to read pre-promotion durable sequence" + return 1 + } + local index health + for index in "${!ADMIN_PORTS[@]}"; do + health=$(curl --max-time 1 -fsS \ + "http://127.0.0.1:${ADMIN_PORTS[$index]}/health" 2>/dev/null || true) + if grep -Eq '"service_ready"[[:space:]]*:[[:space:]]*true' <<<"$health"; then + leader_index=$index + else + wait_master_metric_at_least "$index" ha_oplog_applied_sequence_id \ + "$durable_before" "$START_TIMEOUT_SEC" || { + stop_boundary_client "$client_pid" + smoke_failed "master-$index did not catch up before promotion test" + return 1 + } + fi + done + [[ -n "$leader_index" ]] || { + stop_boundary_client "$client_pid" + smoke_failed "failed to identify ready leader" + return 1 + } + + local point=promotion_final_catch_up_before_complete + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint arm \ + --run-dir "$RUN_DIR" --name "$point" + "$fault_ctl" process kill --run-dir "$RUN_DIR" \ + --name "master-$leader_index" + if ! FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint wait \ + --run-dir "$RUN_DIR" --name "$point" \ + --timeout-sec "$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "promotion final catch-up failpoint was not hit" + return 1 + fi + + local ready=0 + for index in "${!ADMIN_PORTS[@]}"; do + health=$(curl --max-time 1 -fsS \ + "http://127.0.0.1:${ADMIN_PORTS[$index]}/health" 2>/dev/null || true) + grep -Eq '"service_ready"[[:space:]]*:[[:space:]]*true' <<<"$health" && + ready=$((ready + 1)) + done + if ((ready != 0)); then + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" || true + stop_boundary_client "$client_pid" + smoke_failed "a master became ready before promotion final catch-up completed" + return 1 + fi + + local promoted_pid + promoted_pid=$(<"$FAILPOINT_DIR/$point.hit") + FAILPOINT_DIR="$FAILPOINT_DIR" "$fault_ctl" failpoint release \ + --run-dir "$RUN_DIR" --name "$point" + if ! wait_for_single_leader "$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "cluster did not finish promotion after failpoint release" + return 1 + fi + local promoted_index="" + for index in "${!ADMIN_PORTS[@]}"; do + health=$(curl --max-time 1 -fsS \ + "http://127.0.0.1:${ADMIN_PORTS[$index]}/health" 2>/dev/null || true) + if grep -Eq '"service_ready"[[:space:]]*:[[:space:]]*true' \ + <<<"$health"; then + promoted_index=$index + break + fi + done + if [[ -z "$promoted_index" ]] || + ! wait_file_text "$log" "Reconnected to master" "$START_TIMEOUT_SEC" || + ! wait_master_metric_at_least "$promoted_index" master_active_clients 1 \ + "$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "client did not reconnect to the promoted leader" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/after-promotion-catchup.json"; then + stop_boundary_client "$client_pid" + smoke_failed "durable OpLog failed verification after promotion" + return 1 + fi + + printf 'get %s %s\nget %s %s\n' "$client_name" "$keep_key" \ + "$client_name" "$removed_key" >&3 + printf 'batch-smoke %s promotion-batch-%s\n' \ + "$client_name" "$CLUSTER_ID" >&3 + if ! wait_file_text "$log" "Get value: keep-value" \ + "$START_TIMEOUT_SEC" || + ! wait_file_count "$log" "Failed to get value: OBJECT_NOT_FOUND" 2 \ + "$START_TIMEOUT_SEC" || + ! wait_file_text "$log" \ + "Successfully completed batch smoke: promotion-batch-$CLUSTER_ID" \ + "$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "promoted leader business-state verification failed" + return 1 + fi + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" wait \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" \ + --last_seq="$((durable_before + 6))" \ + --timeout_sec="$START_TIMEOUT_SEC"; then + stop_boundary_client "$client_pid" + smoke_failed "post-promotion batch mutations did not become durable" + return 1 + fi + local durable_after + durable_after=$(read_durable_sequence) || { + stop_boundary_client "$client_pid" + smoke_failed "failed to read post-promotion durable sequence" + return 1 + } + stop_boundary_client "$client_pid" + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/after-promotion-business.json"; then + smoke_failed "post-promotion business OpLog verification failed" + return 1 + fi + + collect_cluster + down_cluster + echo "PASS: promotion catch-up smoke completed; put value survived; remove did not resurrect; batch smoke passed; killed=master-$leader_index promoted_pid=$promoted_pid before_seq=$durable_before after_seq=$durable_after artifacts: $RUN_DIR" +} + +ha_failover_smoke_cluster() { + local ha_client="$BUILD_DIR/mooncake-store/tests/e2e/oplog_ha_client" + local fault_ctl="$SCRIPT_DIR/oplog_fault_ctl.sh" + require_executable "$ha_client" + require_executable "$fault_ctl" + CLIENT_COUNT=0 + up_cluster + export MC_STORE_CLUSTER_ID="$CLUSTER_ID" + + local seed_manifest="$RUN_DIR/workload/seed.ack" + local pre_kill_manifest="$RUN_DIR/workload/seed.pre-kill.ack" + local seed_prefix="ha-seed-$CLUSTER_ID" + local pressure_prefix="ha-pressure-$CLUSTER_ID" + local pressure_pids=() + local leader_index durable_before index pid + + for index in 0 1; do + start_process "provider-$index" "$ha_client" --mode=provider \ + --port="$(find_free_port)" \ + --master_server_entry="etcd://$ETCD_ENDPOINTS" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" + wait_file_text "$RUN_DIR/logs/provider-$index.out" "provider_ready" \ + "$START_TIMEOUT_SEC" || { + smoke_failed "provider-$index did not mount a storage segment" + return 1 + } + done + if ! "$ha_client" --mode=seed --port="$(find_free_port)" \ + --master_server_entry="etcd://$ETCD_ENDPOINTS" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" --key_prefix="$seed_prefix" \ + --count="$HA_OBJECTS" --payload_size="$HA_PAYLOAD_BYTES" \ + --manifest="$seed_manifest" \ + >"$RUN_DIR/workload/seed.log" 2>&1; then + smoke_failed "deterministic seed failed" + return 1 + fi + if ! "$INSPECTOR_BIN" wait --endpoints="$ETCD_ENDPOINTS" \ + --cluster_id="$CLUSTER_ID" --last_seq="$HA_OBJECTS" \ + --timeout_sec="$START_TIMEOUT_SEC"; then + smoke_failed "seed OpLog did not become durable" + return 1 + fi + durable_before=$(read_durable_sequence) || { + smoke_failed "failed to read pre-kill durable sequence" + return 1 + } + cp "$seed_manifest" "$pre_kill_manifest" + if [[ $(wc -l <"$pre_kill_manifest") -ne "$HA_OBJECTS" ]]; then + smoke_failed "pre-kill manifest does not contain every acknowledged seed" + return 1 + fi + leader_index=$(ready_leader_index) || { + smoke_failed "failed to identify pre-kill leader" + return 1 + } + for index in "${!ADMIN_PORTS[@]}"; do + [[ "$index" == "$leader_index" ]] && continue + wait_master_metric_at_least "$index" ha_oplog_applied_sequence_id \ + "$durable_before" "$START_TIMEOUT_SEC" || { + smoke_failed "master-$index did not apply pre-kill durable sequence" + return 1 + } + done + capture_ha_stage pre-kill || { + smoke_failed "pre-kill evidence collection failed" + return 1 + } + printf 'leader_index=%s\ndurable_before=%s\n' "$leader_index" \ + "$durable_before" >"$RUN_DIR/audit/pre-kill-leader.txt" + + "$fault_ctl" process kill --run-dir "$RUN_DIR" \ + --name "master-$leader_index" + pid=$(<"$RUN_DIR/pids/master-$leader_index.pid") + wait "$pid" 2>/dev/null || true + if kill -0 "$pid" 2>/dev/null; then + smoke_failed "killed leader process is still alive" + return 1 + fi + if ! wait_for_single_leader "$START_TIMEOUT_SEC"; then + smoke_failed "cluster did not promote a new leader" + return 1 + fi + local promoted_index + promoted_index=$(ready_leader_index) || { + smoke_failed "promotion did not produce exactly one ready leader" + return 1 + } + [[ "$promoted_index" != "$leader_index" ]] || { + smoke_failed "old leader remained ready after kill" + return 1 + } + wait_master_metric_at_least "$promoted_index" master_active_clients 2 \ + "$START_TIMEOUT_SEC" || { + smoke_failed "storage providers did not reconnect to promoted leader" + return 1 + } + printf 'promoted_index=%s\n' "$promoted_index" \ + >"$RUN_DIR/audit/post-promotion-leader.txt" + capture_ha_stage post-promotion || { + smoke_failed "post-promotion evidence collection failed" + return 1 + } + if ! "$ha_client" --mode=verify --port="$(find_free_port)" \ + --master_server_entry="etcd://$ETCD_ENDPOINTS" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" --key_prefix="$seed_prefix" \ + --payload_size="$HA_PAYLOAD_BYTES" --manifest="$pre_kill_manifest" \ + >"$RUN_DIR/workload/post-promotion-verify.log" 2>&1; then + smoke_failed "acknowledged data was not readable after promotion" + return 1 + fi + + for ((index = 0; index < 2; ++index)); do + local pressure_manifest="$RUN_DIR/workload/pressure-$index.ack" + "$ha_client" --mode=pressure --port="$(find_free_port)" \ + --master_server_entry="etcd://$ETCD_ENDPOINTS" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" --key_prefix="$pressure_prefix-$index" \ + --start_index=1000000 --payload_size="$HA_PAYLOAD_BYTES" \ + --duration_sec="$HA_PRESSURE_SEC" --sleep_ms=25 \ + --manifest="$pressure_manifest" \ + >"$RUN_DIR/workload/pressure-$index.log" 2>&1 & + pressure_pids+=("$!") + done + local pressure_status=0 pressure_pid + for pressure_pid in "${pressure_pids[@]}"; do + wait "$pressure_pid" || pressure_status=1 + done + if ((pressure_status != 0)); then + smoke_failed "post-promotion pressure workload failed" + return 1 + fi + for index in 0 1; do + "$ha_client" --mode=verify --port="$(find_free_port)" \ + --master_server_entry="etcd://$ETCD_ENDPOINTS" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" --key_prefix="$pressure_prefix-$index" \ + --payload_size="$HA_PAYLOAD_BYTES" \ + --manifest="$RUN_DIR/workload/pressure-$index.ack" \ + >"$RUN_DIR/workload/pressure-$index-verify.log" 2>&1 || { + smoke_failed "post-promotion pressure data verification failed" + return 1 + } + done + if ! "$ha_client" --mode=verify --port="$(find_free_port)" \ + --master_server_entry="etcd://$ETCD_ENDPOINTS" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" --key_prefix="$seed_prefix" \ + --payload_size="$HA_PAYLOAD_BYTES" --manifest="$pre_kill_manifest" \ + >"$RUN_DIR/workload/post-pressure-seed-verify.log" 2>&1; then + smoke_failed "pressure corrupted pre-kill acknowledged data" + return 1 + fi + capture_ha_stage post-pressure || { + smoke_failed "post-pressure evidence collection failed" + return 1 + } + collect_cluster + down_cluster + echo "PASS: real HA failover completed; leader=$leader_index promoted=$promoted_index durable_before=$durable_before artifacts=$RUN_DIR" +} + +allocator_recovery_smoke_cluster() { + local ha_client="$BUILD_DIR/mooncake-store/tests/e2e/oplog_ha_client" + local fault_ctl="$SCRIPT_DIR/oplog_fault_ctl.sh" + require_executable "$ha_client" + require_executable "$fault_ctl" + CLIENT_COUNT=0 + [[ -n "$MEMORY_ALLOCATOR" ]] || MEMORY_ALLOCATOR=offset + up_cluster + export MC_STORE_CLUSTER_ID="$CLUSTER_ID" + + local seed_prefix="allocator-seed-$CLUSTER_ID" + local refill_prefix="allocator-refill-$CLUSTER_ID" + local pressure_prefix="allocator-pressure-$CLUSTER_ID" + local seed_manifest="$RUN_DIR/workload/seed.ack" + local survivor_manifest="$RUN_DIR/workload/survivor.ack" + local deleted_manifest="$RUN_DIR/workload/deleted.ack" + local refill_manifest="$RUN_DIR/workload/refill.ack" + local pressure_manifest="$RUN_DIR/workload/pressure.ack" + local -a client_args=( + --master_server_entry="etcd://$ETCD_ENDPOINTS" + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" + --protocol="$PROTOCOL") + local index leader_index promoted_index pid durable_start durable_before expected + + for index in 0 1; do + start_process "provider-$index" "$ha_client" --mode=provider \ + --port="$(find_free_port)" --segment_size="$RECOVERY_SEGMENT_BYTES" \ + "${client_args[@]}" + wait_file_text "$RUN_DIR/logs/provider-$index.out" provider_ready \ + "$START_TIMEOUT_SEC" || { + smoke_failed "provider-$index did not mount a storage segment" + return 1 + } + done + durable_start=$(read_durable_sequence) || { + smoke_failed "failed to read initial durable sequence" + return 1 + } + + "$ha_client" --mode=seed --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$seed_prefix" \ + --count="$RECOVERY_SEED_OBJECTS" \ + --payload_sizes="$RECOVERY_PAYLOAD_SIZES" --manifest="$seed_manifest" \ + >"$RUN_DIR/workload/seed.log" 2>&1 || { + smoke_failed "mixed-size seed failed" + return 1 + } + awk 'NR % 3 == 0' "$seed_manifest" >"$deleted_manifest" + awk 'NR % 3 != 0' "$seed_manifest" >"$survivor_manifest" + [[ -s "$deleted_manifest" && -s "$survivor_manifest" ]] || { + smoke_failed "failed to split seed manifest into holes and survivors" + return 1 + } + "$ha_client" --mode=delete --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$seed_prefix" \ + --manifest="$deleted_manifest" >"$RUN_DIR/workload/delete.log" 2>&1 || { + smoke_failed "failed to create allocator holes" + return 1 + } + "$ha_client" --mode=verify --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$seed_prefix" \ + --payload_sizes="$RECOVERY_PAYLOAD_SIZES" \ + --manifest="$survivor_manifest" \ + >"$RUN_DIR/workload/pre-kill-survivor.log" 2>&1 || { + smoke_failed "survivors were corrupted while creating holes" + return 1 + } + "$ha_client" --mode=verify-absent --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$seed_prefix" \ + --manifest="$deleted_manifest" \ + >"$RUN_DIR/workload/pre-kill-absent.log" 2>&1 || { + smoke_failed "deleted objects remained visible" + return 1 + } + + expected=$((durable_start + RECOVERY_SEED_OBJECTS + $(wc -l <"$deleted_manifest"))) + "$INSPECTOR_BIN" wait --endpoints="$ETCD_ENDPOINTS" \ + --cluster_id="$CLUSTER_ID" --last_seq="$expected" \ + --timeout_sec="$START_TIMEOUT_SEC" || { + smoke_failed "seed and delete OpLog did not become durable" + return 1 + } + durable_before=$(read_durable_sequence) || { + smoke_failed "failed to read pre-kill durable sequence" + return 1 + } + leader_index=$(ready_leader_index) || { + smoke_failed "failed to identify pre-kill leader" + return 1 + } + for index in "${!ADMIN_PORTS[@]}"; do + [[ "$index" == "$leader_index" ]] && continue + wait_master_metric_at_least "$index" ha_oplog_applied_sequence_id \ + "$durable_before" "$START_TIMEOUT_SEC" || { + smoke_failed "master-$index did not apply allocator state" + return 1 + } + done + capture_ha_stage allocator-pre-kill || { + smoke_failed "pre-kill evidence collection failed" + return 1 + } + + "$fault_ctl" process kill --run-dir "$RUN_DIR" \ + --name "master-$leader_index" + pid=$(<"$RUN_DIR/pids/master-$leader_index.pid") + wait "$pid" 2>/dev/null || true + wait_for_single_leader "$START_TIMEOUT_SEC" || { + smoke_failed "cluster did not promote a new leader" + return 1 + } + promoted_index=$(ready_leader_index) || { + smoke_failed "promotion did not produce exactly one ready leader" + return 1 + } + [[ "$promoted_index" != "$leader_index" ]] || { + smoke_failed "old leader remained ready after kill" + return 1 + } + wait_master_metric_at_least "$promoted_index" master_active_clients 2 \ + "$START_TIMEOUT_SEC" || { + smoke_failed "storage providers did not reconnect to promoted leader" + return 1 + } + + "$ha_client" --mode=verify --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$seed_prefix" \ + --payload_sizes="$RECOVERY_PAYLOAD_SIZES" \ + --manifest="$survivor_manifest" \ + >"$RUN_DIR/workload/post-promotion-survivor.log" 2>&1 || { + smoke_failed "survivors were not readable after allocator recovery" + return 1 + } + "$ha_client" --mode=verify-absent --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$seed_prefix" \ + --manifest="$deleted_manifest" \ + >"$RUN_DIR/workload/post-promotion-absent.log" 2>&1 || { + smoke_failed "deleted objects reappeared after allocator recovery" + return 1 + } + "$ha_client" --mode=seed --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$refill_prefix" \ + --start_index=100000 --count="$RECOVERY_REFILL_OBJECTS" \ + --payload_sizes="$RECOVERY_PAYLOAD_SIZES" --manifest="$refill_manifest" \ + >"$RUN_DIR/workload/refill.log" 2>&1 || { + smoke_failed "post-promotion mixed-size refill failed" + return 1 + } + "$ha_client" --mode=verify --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$seed_prefix" \ + --payload_sizes="$RECOVERY_PAYLOAD_SIZES" --manifest="$survivor_manifest" \ + >"$RUN_DIR/workload/pre-pressure-survivor.log" 2>&1 || { + smoke_failed "survivors failed after refill" + return 1 + } + "$ha_client" --mode=verify --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$refill_prefix" \ + --payload_sizes="$RECOVERY_PAYLOAD_SIZES" --manifest="$refill_manifest" \ + >"$RUN_DIR/workload/pre-pressure-refill.log" 2>&1 || { + smoke_failed "refill data failed before pressure" + return 1 + } + "$ha_client" --mode=verify-absent --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$seed_prefix" \ + --manifest="$deleted_manifest" \ + >"$RUN_DIR/workload/pre-pressure-absent.log" 2>&1 || { + smoke_failed "deleted objects reappeared after refill" + return 1 + } + "$ha_client" --mode=pressure --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$pressure_prefix" \ + --start_index=1000000 --payload_sizes="$RECOVERY_PAYLOAD_SIZES" \ + --duration_sec="$RECOVERY_PRESSURE_SEC" --sleep_ms=25 \ + --manifest="$pressure_manifest" \ + >"$RUN_DIR/workload/pressure.log" 2>&1 || { + smoke_failed "post-promotion mixed-size pressure failed" + return 1 + } + + for index in survivor refill pressure; do + local prefix=$seed_prefix + local manifest=$survivor_manifest + [[ "$index" != refill ]] || { + prefix=$refill_prefix + manifest=$refill_manifest + } + [[ "$index" != pressure ]] || { + prefix=$pressure_prefix + manifest=$pressure_manifest + } + "$ha_client" --mode=verify --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$prefix" \ + --payload_sizes="$RECOVERY_PAYLOAD_SIZES" --manifest="$manifest" \ + >"$RUN_DIR/workload/final-$index.log" 2>&1 || { + smoke_failed "final $index verification failed" + return 1 + } + done + "$ha_client" --mode=verify-absent --port="$(find_free_port)" \ + "${client_args[@]}" --key_prefix="$seed_prefix" \ + --manifest="$deleted_manifest" \ + >"$RUN_DIR/workload/final-absent.log" 2>&1 || { + smoke_failed "deleted objects reappeared during post-promotion writes" + return 1 + } + expected=$((durable_before + RECOVERY_REFILL_OBJECTS + $(wc -l <"$pressure_manifest"))) + "$INSPECTOR_BIN" wait --endpoints="$ETCD_ENDPOINTS" \ + --cluster_id="$CLUSTER_ID" --last_seq="$expected" \ + --timeout_sec="$START_TIMEOUT_SEC" || { + smoke_failed "post-promotion writes did not become durable" + return 1 + } + capture_ha_stage allocator-final || { + smoke_failed "final OpLog verification failed" + return 1 + } + collect_cluster + down_cluster + echo "PASS: allocator recovery completed; allocator=$MEMORY_ALLOCATOR leader=$leader_index promoted=$promoted_index artifacts=$RUN_DIR" +} + +allocator_recovery_matrix() { + local base_run_dir=$RUN_DIR + local base_cluster_id=$CLUSTER_ID + local base_etcd_endpoints=$ETCD_ENDPOINTS + local allocator + for allocator in offset cachelib; do + RUN_DIR="$base_run_dir-$allocator" + CLUSTER_ID="$base_cluster_id-$allocator" + ETCD_ENDPOINTS=$base_etcd_endpoints + MEMORY_ALLOCATOR=$allocator + RPC_PORTS=() + ADMIN_PORTS=() + allocator_recovery_smoke_cluster || return 1 + done + echo "PASS: allocator recovery matrix completed; artifacts=$base_run_dir-{offset,cachelib}" +} + +non_ha_smoke_cluster() { + local fault_ctl="$SCRIPT_DIR/oplog_fault_ctl.sh" + require_executable "$fault_ctl" + up_cluster + + if [[ "$USE_ETCD_OBSERVER" == true ]]; then + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/non-ha-before.json"; then + smoke_failed \ + "observer etcd namespace was not empty before non-HA workload" + return 1 + fi + "$fault_ctl" process stop --run-dir "$RUN_DIR" --name etcd + fi + + local label=non-ha key clientctl client_port + key="non-ha-$CLUSTER_ID" + clientctl="$BUILD_DIR/mooncake-store/tests/e2e/clientctl" + client_port=$(find_free_port) + require_executable "$clientctl" + if ! { + printf 'create %s %s\n' "$label" "$client_port" + printf 'mount %s segment-%s %s\n' "$label" "$label" \ + "$((64 * 1024 * 1024))" + printf 'put %s %s non-ha-value\n' "$label" "$key" + printf 'get %s %s\n' "$label" "$key" + printf 'get %s missing-%s\n' "$label" "$key" + printf 'put %s %s duplicate-value\n' "$label" "$key" + printf 'get %s %s\n' "$label" "$key" + printf 'delete %s %s\n' "$label" "$key" + printf 'sleep 6\n' + printf 'delete %s %s\n' "$label" "$key" + printf 'get %s %s\n' "$label" "$key" + printf 'batch-smoke %s batch-%s\n' "$label" "$CLUSTER_ID" + printf 'unmount %s segment-%s\n' "$label" "$label" + printf 'mount %s segment-%s %s\n' "$label" "$label" \ + "$((64 * 1024 * 1024))" + printf 'put %s remount-%s remount-value\n' "$label" "$CLUSTER_ID" + printf 'sleep 6\n' + printf 'delete %s remount-%s\n' "$label" "$CLUSTER_ID" + printf 'unmount %s segment-%s\n' "$label" "$label" + printf 'terminate\n' + } | MC_STORE_CLUSTER_ID="$CLUSTER_ID" LSAN_OPTIONS=detect_leaks=0 \ + "$clientctl" --master_server_entry="127.0.0.1:${RPC_PORTS[0]}" \ + --engine_meta_url="http://127.0.0.1:$METADATA_PORT/metadata" \ + --protocol="$PROTOCOL" >"$RUN_DIR/workload/non-ha.log" 2>&1; then + resume_observer_etcd "$fault_ctl" + smoke_failed "non-HA client workflow failed while etcd was stopped" + return 1 + fi + if ! grep -Fq "Successfully put value for key: $key" \ + "$RUN_DIR/workload/non-ha.log" || + ! grep -Fq "Get value: non-ha-value" "$RUN_DIR/workload/non-ha.log" || + ! grep -Fq "Failed to get value: OBJECT_NOT_FOUND" \ + "$RUN_DIR/workload/non-ha.log" || + ! grep -Fq "Failed to delete value: OBJECT_HAS_LEASE" \ + "$RUN_DIR/workload/non-ha.log" || + ! grep -Fq "Successfully deleted value for key: $key" \ + "$RUN_DIR/workload/non-ha.log" || + ! grep -Fq "Failed to get value: OBJECT_NOT_FOUND" \ + "$RUN_DIR/workload/non-ha.log" || + ! grep -Fq "Successfully completed batch smoke: batch-$CLUSTER_ID" \ + "$RUN_DIR/workload/non-ha.log" || + ! grep -Fq "Successfully put value for key: remount-$CLUSTER_ID" \ + "$RUN_DIR/workload/non-ha.log" || + [[ $(grep -Fc "Successfully mounted segment on client $label" \ + "$RUN_DIR/workload/non-ha.log") -ne 2 ]] || + [[ $(grep -Fc "Successfully unmounted segment from client $label" \ + "$RUN_DIR/workload/non-ha.log") -ne 2 ]]; then + resume_observer_etcd "$fault_ctl" + smoke_failed "non-HA client workflow returned unexpected results" + return 1 + fi + if [[ $(grep -Fc "Get value: non-ha-value" \ + "$RUN_DIR/workload/non-ha.log") -ne 2 ]]; then + resume_observer_etcd "$fault_ctl" + smoke_failed "repeated non-HA put did not preserve the original value" + return 1 + fi + if ! wait_master_metric_equal 0 master_active_clients 0 \ + "$START_TIMEOUT_SEC"; then + resume_observer_etcd "$fault_ctl" + smoke_failed "non-HA lifecycle client was not cleaned up" + return 1 + fi + + local -a worker_pids=() + local worker + for ((worker = 0; worker < NON_HA_WORKERS; ++worker)); do + run_client_smoke "non-ha-worker-$worker" 2 & + worker_pids+=("$!") + done + local worker_failed=0 pid + for pid in "${worker_pids[@]}"; do + wait "$pid" || worker_failed=1 + done + if ((worker_failed)); then + resume_observer_etcd "$fault_ctl" + smoke_failed "concurrent non-HA client workload failed" + return 1 + fi + + if [[ "$USE_ETCD_OBSERVER" == true ]]; then + "$fault_ctl" process continue --run-dir "$RUN_DIR" --name etcd + wait_http "http://$ETCD_ENDPOINTS/health" "$START_TIMEOUT_SEC" || { + smoke_failed "observer etcd did not resume" + return 1 + } + if ! LSAN_OPTIONS=detect_leaks=0 "$INSPECTOR_BIN" verify \ + --endpoints="$ETCD_ENDPOINTS" --cluster_id="$CLUSTER_ID" --json \ + >"$RUN_DIR/audit/non-ha-after.json"; then + smoke_failed "non-HA workload created invalid OpLog state" + return 1 + fi + if ! python3 -c 'import json,sys +d=json.load(open(sys.argv[1], encoding="utf-8")) +assert d["batch_count"] == 0 +assert d["entry_count"] == 0 +assert d["legacy_max_seq"] == 0 +assert d["durable_prefix"] is None' "$RUN_DIR/audit/non-ha-after.json"; then + smoke_failed "non-HA workload unexpectedly created an OpLog namespace" + return 1 + fi + fi + + collect_cluster + down_cluster + echo "PASS: non-HA smoke completed without an HA backend; $NON_HA_WORKERS concurrent clients passed; observer_etcd=$USE_ETCD_OBSERVER; artifacts: $RUN_DIR" +} + +resume_observer_etcd() { + local fault_ctl=$1 + [[ "$USE_ETCD_OBSERVER" != true ]] || + "$fault_ctl" process continue --run-dir "$RUN_DIR" --name etcd || true +} + +stop_boundary_client() { + local pid=$1 + printf 'terminate\n' >&3 2>/dev/null || true + exec 3>&- + wait "$pid" 2>/dev/null || true +} + +wait_file_text() { + local file=$1 + local expected=$2 + local timeout=$3 + local deadline=$((SECONDS + timeout)) + while ((SECONDS < deadline)); do + [[ -f "$file" ]] && grep -Fq "$expected" "$file" && return 0 + sleep 0.1 + done + return 1 +} + +wait_file_count() { + local file=$1 + local expected=$2 + local count=$3 + local timeout=$4 + local deadline=$((SECONDS + timeout)) + while ((SECONDS < deadline)); do + [[ -f "$file" ]] && + [[ $(grep -Fc "$expected" "$file") -ge "$count" ]] && return 0 + sleep 0.1 + done + return 1 +} + +read_master_metric() { + local index=$1 + local metric=$2 + curl -fsS "http://127.0.0.1:${ADMIN_PORTS[$index]}/metrics" | + awk -v name="$metric" '$1 == name { print $2; found = 1; exit } END { if (!found) exit 1 }' +} + +wait_master_metric_at_least() { + local index=$1 + local metric=$2 + local target=$3 + local timeout=$4 + local deadline=$((SECONDS + timeout)) + while ((SECONDS < deadline)); do + local value + value=$(read_master_metric "$index" "$metric" 2>/dev/null || true) + [[ "$value" =~ ^[0-9]+$ ]] && ((value >= target)) && return 0 + sleep 0.1 + done + return 1 +} + +wait_master_metric_equal() { + local index=$1 + local metric=$2 + local target=$3 + local timeout=$4 + local deadline=$((SECONDS + timeout)) + while ((SECONDS < deadline)); do + local value + value=$(read_master_metric "$index" "$metric" 2>/dev/null || true) + [[ "$value" == "$target" ]] && return 0 + sleep 0.1 + done + return 1 +} + +smoke_failed() { + echo "error: $1" >&2 + collect_cluster || true + printf 'environment preserved; stop it with: %q down --run-dir %q\n' \ + "$0" "$RUN_DIR" >&2 +} + +parse_run_dir() { + RUN_DIR="" + while (($#)); do + case "$1" in + --run-dir) + (($# >= 2)) || die "--run-dir requires a value" + RUN_DIR=$2 + shift 2 + ;; + *) die "unknown option: $1" ;; + esac + done + [[ -n "$RUN_DIR" ]] || die "--run-dir is required" + [[ -d "$RUN_DIR" ]] || die "run directory does not exist: $RUN_DIR" +} + +pid_matches() { + local pid=$1 + local expected_file=$2 + [[ -r "/proc/$pid/cmdline" && -r "$expected_file" ]] || return 1 + local actual + actual=$(tr '\0' ' ' <"/proc/$pid/cmdline") + grep -Fqx "$actual" "$expected_file" +} + +stop_pid_file() { + local pid_file=$1 + local expected_file=${pid_file%.pid}.cmd + local pid + pid=$(<"$pid_file") + [[ "$pid" =~ ^[0-9]+$ ]] || return 0 + kill -0 "$pid" 2>/dev/null || return 0 + pid_matches "$pid" "$expected_file" || { + echo "warning: refusing to stop PID $pid because command does not match" >&2 + return 0 + } + kill -TERM "$pid" 2>/dev/null || return 0 + local deadline=$((SECONDS + 5)) + while kill -0 "$pid" 2>/dev/null && ((SECONDS < deadline)); do + sleep 0.1 + done + if kill -0 "$pid" 2>/dev/null; then + kill -KILL "$pid" + fi + return 0 +} + +down_cluster() { + local pid_file + shopt -s nullglob + for pid_file in "$RUN_DIR"/pids/client-*.pid; do + stop_pid_file "$pid_file" + done + for pid_file in "$RUN_DIR"/pids/provider-*.pid; do + stop_pid_file "$pid_file" + done + for pid_file in "$RUN_DIR"/pids/master-*.pid; do + stop_pid_file "$pid_file" + done + for pid_file in "$RUN_DIR"/pids/metadata.pid "$RUN_DIR"/pids/etcd.pid; do + [[ -f "$pid_file" ]] && stop_pid_file "$pid_file" + done + echo "cluster is stopped: $RUN_DIR" +} + +main() { + (($# >= 1)) || { + usage + exit 1 + } + local command=$1 + shift + case "$command" in + up | smoke | restart-smoke | failpoint-smoke | failpoint-crash-smoke | remove-boundary-smoke | standby-read-smoke | promotion-catchup-smoke | ha-failover-smoke | allocator-recovery-smoke | allocator-recovery-matrix | non-ha-smoke) + parse_up_options "$@" + if [[ "$command" == allocator-recovery-smoke || + "$command" == allocator-recovery-matrix ]] && + [[ "$MASTER_COUNT" != 3 ]]; then + die "allocator recovery requires exactly 3 masters" + fi + if [[ "$command" == non-ha-smoke ]]; then + MASTER_COUNT=1 + ENABLE_HA=false + fi + if [[ "$USE_ETCD_OBSERVER" != true && "$ENABLE_HA" == true ]]; then + die "--no-etcd-observer is only supported by non-ha-smoke" + fi + if [[ "$command" == up ]]; then + up_cluster + elif [[ "$command" == smoke ]]; then + smoke_cluster + elif [[ "$command" == restart-smoke ]]; then + restart_smoke_cluster + elif [[ "$command" == failpoint-smoke ]]; then + failpoint_smoke_cluster + elif [[ "$command" == remove-boundary-smoke ]]; then + remove_boundary_smoke_cluster + elif [[ "$command" == standby-read-smoke ]]; then + standby_read_smoke_cluster + elif [[ "$command" == promotion-catchup-smoke ]]; then + promotion_catchup_smoke_cluster + elif [[ "$command" == ha-failover-smoke ]]; then + ha_failover_smoke_cluster + elif [[ "$command" == allocator-recovery-smoke ]]; then + allocator_recovery_smoke_cluster + elif [[ "$command" == allocator-recovery-matrix ]]; then + allocator_recovery_matrix + elif [[ "$command" == non-ha-smoke ]]; then + non_ha_smoke_cluster + else + failpoint_crash_smoke_cluster + fi + ;; + status | collect | restart) + parse_run_dir "$@" + load_cluster_env + if [[ "$command" == status ]]; then + status_cluster + elif [[ "$command" == collect ]]; then + collect_cluster + else + restart_masters + fi + ;; + down) + parse_run_dir "$@" + down_cluster + ;; + *) die "unknown command: $command" ;; + esac +} + +main "$@" diff --git a/mooncake-store/tests/e2e/run_oplog_batch_cluster_test.sh b/mooncake-store/tests/e2e/run_oplog_batch_cluster_test.sh new file mode 100755 index 0000000000..4631729ca6 --- /dev/null +++ b/mooncake-store/tests/e2e/run_oplog_batch_cluster_test.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +SCRIPT="$SCRIPT_DIR/run_oplog_batch_cluster.sh" +TEST_ROOT=${TEST_ROOT:-"/tmp/mooncake-oplog-cluster-script-test-$$"} +mkdir -p "$TEST_ROOT" + +fail() { + echo "FAIL: $*" >&2 + exit 1 +} + +expect_failure() { + local expected_status=$1 + local expected_text=$2 + shift 2 + + local output_file="$TEST_ROOT/output-$RANDOM.log" + local status=0 + "$@" >"$output_file" 2>&1 || status=$? + [[ "$status" -eq "$expected_status" ]] || + fail "expected status $expected_status, got $status: $(<"$output_file")" + grep -Fq "$expected_text" "$output_file" || + fail "expected '$expected_text': $(<"$output_file")" +} + +expect_success() { + local expected_text=$1 + shift + + local output_file="$TEST_ROOT/output-$RANDOM.log" + "$@" >"$output_file" 2>&1 || + fail "expected success: $(<"$output_file")" + grep -Fq "$expected_text" "$output_file" || + fail "expected '$expected_text': $(<"$output_file")" +} + +expect_failure 1 "unknown command" "$SCRIPT" unknown +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" up --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" failpoint-smoke --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" failpoint-crash-smoke --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" remove-boundary-smoke --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" standby-read-smoke --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" promotion-catchup-smoke --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" ha-failover-smoke --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" allocator-recovery-smoke --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" allocator-recovery-matrix --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "build directory does not exist" \ + "$SCRIPT" non-ha-smoke --build-dir "$TEST_ROOT/missing-build" +expect_failure 1 "master config does not exist" \ + "$SCRIPT" non-ha-smoke --build-dir "$TEST_ROOT" \ + --master-config "$TEST_ROOT/missing-master.yaml" +expect_failure 1 "non-ha-workers must be positive" \ + "$SCRIPT" non-ha-smoke --build-dir "$TEST_ROOT" --non-ha-workers 0 +expect_failure 1 "ha-objects must be positive" \ + "$SCRIPT" ha-failover-smoke --build-dir "$TEST_ROOT" --ha-objects 0 +expect_failure 1 "ha-payload-bytes must be positive" \ + "$SCRIPT" ha-failover-smoke --build-dir "$TEST_ROOT" --ha-payload-bytes 0 +expect_failure 1 "ha-pressure-sec must be positive" \ + "$SCRIPT" ha-failover-smoke --build-dir "$TEST_ROOT" --ha-pressure-sec 0 +expect_failure 1 "memory-allocator must be offset or cachelib" \ + "$SCRIPT" allocator-recovery-smoke --build-dir "$TEST_ROOT" \ + --memory-allocator invalid +expect_failure 1 "recovery-seed-objects must be positive" \ + "$SCRIPT" allocator-recovery-smoke --build-dir "$TEST_ROOT" \ + --recovery-seed-objects 0 +expect_failure 1 "recovery-refill-objects must be positive" \ + "$SCRIPT" allocator-recovery-smoke --build-dir "$TEST_ROOT" \ + --recovery-refill-objects 0 +expect_failure 1 "recovery-pressure-sec must be positive" \ + "$SCRIPT" allocator-recovery-smoke --build-dir "$TEST_ROOT" \ + --recovery-pressure-sec 0 +expect_failure 1 "recovery-segment-bytes must be positive" \ + "$SCRIPT" allocator-recovery-smoke --build-dir "$TEST_ROOT" \ + --recovery-segment-bytes 0 +expect_failure 1 "recovery-payload-sizes must be positive CSV integers" \ + "$SCRIPT" allocator-recovery-smoke --build-dir "$TEST_ROOT" \ + --recovery-payload-sizes 64,0,4096 +expect_failure 1 "allocator recovery requires exactly 3 masters" \ + "$SCRIPT" allocator-recovery-smoke --build-dir "$TEST_ROOT" --masters 2 +EMPTY_BUILD="$TEST_ROOT/empty-build" +mkdir -p "$EMPTY_BUILD" +expect_failure 1 "missing executable" "$SCRIPT" up --build-dir "$EMPTY_BUILD" +expect_failure 1 "missing executable" "$SCRIPT" allocator-recovery-smoke \ + --build-dir "$EMPTY_BUILD" +expect_failure 1 "missing executable" "$SCRIPT" non-ha-smoke \ + --build-dir "$EMPTY_BUILD" --no-etcd-observer +expect_failure 1 "run directory does not exist" \ + "$SCRIPT" status --run-dir "$TEST_ROOT/missing-run" +expect_failure 1 "run directory does not exist" \ + "$SCRIPT" restart --run-dir "$TEST_ROOT/missing-run" + +EMPTY_RUN="$TEST_ROOT/empty-run" +mkdir -p "$EMPTY_RUN/pids" +expect_success "cluster is stopped" "$SCRIPT" down --run-dir "$EMPTY_RUN" +expect_success "cluster is stopped" "$SCRIPT" down --run-dir "$EMPTY_RUN" + +MISMATCH_RUN="$TEST_ROOT/mismatch-run" +mkdir -p "$MISMATCH_RUN/pids" +sleep 30 & +SLEEP_PID=$! +trap 'kill "$SLEEP_PID" 2>/dev/null || true' EXIT +printf '%s\n' "$SLEEP_PID" >"$MISMATCH_RUN/pids/master-0.pid" +printf '%s\n' "definitely-not-the-sleep-command" \ + >"$MISMATCH_RUN/pids/master-0.cmd" +expect_success "refusing to stop PID" \ + "$SCRIPT" down --run-dir "$MISMATCH_RUN" +kill -0 "$SLEEP_PID" 2>/dev/null || fail "mismatched PID was stopped" + +MATCH_RUN="$TEST_ROOT/match-run" +mkdir -p "$MATCH_RUN/pids" +sleep 30 & +MATCH_PID=$! +printf '%s\n' "$MATCH_PID" >"$MATCH_RUN/pids/master-0.pid" +tr '\0' ' ' <"/proc/$MATCH_PID/cmdline" >"$MATCH_RUN/pids/master-0.cmd" +expect_success "cluster is stopped" "$SCRIPT" down --run-dir "$MATCH_RUN" +kill -0 "$MATCH_PID" 2>/dev/null && fail "matching PID was not stopped" +expect_success "cluster is stopped" "$SCRIPT" down --run-dir "$MATCH_RUN" + +echo "PASS" diff --git a/mooncake-store/tests/e2e/store_client_e2e.py b/mooncake-store/tests/e2e/store_client_e2e.py index 9e06167774..c922f154fd 100644 --- a/mooncake-store/tests/e2e/store_client_e2e.py +++ b/mooncake-store/tests/e2e/store_client_e2e.py @@ -1,7 +1,5 @@ #!/usr/bin/env python3 import argparse -import os -import sys import time diff --git a/mooncake-store/tests/file_storage_test.cpp b/mooncake-store/tests/file_storage_test.cpp index 24b2f67243..ef29ccae8b 100644 --- a/mooncake-store/tests/file_storage_test.cpp +++ b/mooncake-store/tests/file_storage_test.cpp @@ -1,14 +1,18 @@ #include #include +#include #include +#include #include #include "allocator.h" -#include "storage_backend.h" +#include "client_metric.h" #include "file_storage.h" +#include "storage_backend.h" +#include "tenant_id.h" +#include "test_server_helpers.h" #include "utils/common.h" -#include "client_metric.h" namespace mooncake { @@ -33,11 +37,19 @@ class FileStorageTest : public ::testing::Test { UnsetEnv("MOONCAKE_OFFLOAD_TOTAL_KEYS_LIMIT"); UnsetEnv("MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES"); UnsetEnv("MOONCAKE_OFFLOAD_HEARTBEAT_INTERVAL_SECONDS"); + UnsetEnv("MOONCAKE_OFFLOAD_ENABLE_DISK_WATERMARK_EVICTION"); + UnsetEnv("MOONCAKE_OFFLOAD_DISK_EVICTION_HIGH_WATERMARK_RATIO"); + UnsetEnv("MOONCAKE_OFFLOAD_DISK_EVICTION_LOW_WATERMARK_RATIO"); + UnsetEnv("MOONCAKE_DISK_EVICTION_HIGH_WATERMARK_RATIO"); + UnsetEnv("MOONCAKE_DISK_EVICTION_LOW_WATERMARK_RATIO"); data_path = std::filesystem::current_path().string() + "/data"; fs::create_directories(data_path); for (const auto& entry : fs::directory_iterator(data_path)) { + std::error_code ec; if (entry.is_regular_file()) { - fs::remove(entry.path()); + fs::remove(entry.path(), ec); + } else if (entry.is_directory()) { + fs::remove_all(entry.path(), ec); } } } @@ -69,6 +81,12 @@ class FileStorageTest : public ::testing::Test { return fileStorage.IsEnableOffloading(); } + tl::expected FileStorageNotifyEvictedDiskReplicas( + FileStorage& fileStorage, + const std::vector& evicted_keys) { + return fileStorage.NotifyEvictedDiskReplicas(evicted_keys); + } + tl::expected FileStorageGroupOffloadingKeysByBucket( FileStorage& fileStorage, const std::unordered_map& offloading_objects, @@ -91,6 +109,64 @@ class FileStorageTest : public ::testing::Test { return bucket_backend->UngroupedOffloadingObjectsSize(); } + // Static funnel to the private FileStorage::IsPerBucketSoftOffloadError. + // FileStorageTest is friended; TEST_F-generated subclasses are not. + static bool CallIsPerBucketSoftOffloadError(ErrorCode error) { + return FileStorage::IsPerBucketSoftOffloadError(error); + } + + void AssertHeartbeatEvictsAllKeys( + FileStorage& fileStorage, const std::vector& expected_keys, + const std::unordered_map>& + batch_object) { + ASSERT_TRUE(fileStorage.storage_backend_->Init()); + { + MutexLocker locker(&fileStorage.offloading_mutex_); + fileStorage.enable_offloading_ = true; + } + + auto offload_result = fileStorage.storage_backend_->BatchOffload( + batch_object, + [&fileStorage](const std::vector& keys, + std::vector& metadatas) { + for (auto& metadata : metadatas) { + metadata.transport_endpoint = fileStorage.local_rpc_addr_; + } + auto result = + fileStorage.client_->NotifyOffloadSuccess(keys, metadatas); + if (!result) { + return result.error(); + } + return ErrorCode::OK; + }); + ASSERT_TRUE(offload_result.has_value()); + ASSERT_EQ(offload_result.value(), + static_cast(expected_keys.size())); + + for (const auto& key : expected_keys) { + auto query_result = fileStorage.client_->Query(key); + ASSERT_TRUE(query_result.has_value()); + bool has_local_disk_replica = false; + for (const auto& replica : query_result->replicas) { + has_local_disk_replica |= replica.is_local_disk_replica(); + } + EXPECT_TRUE(has_local_disk_replica); + } + + auto heartbeat_result = fileStorage.Heartbeat(); + ASSERT_TRUE(heartbeat_result.has_value()); + + for (const auto& key : expected_keys) { + auto exists = fileStorage.storage_backend_->IsExist(key); + ASSERT_TRUE(exists.has_value()); + EXPECT_FALSE(exists.value()); + + auto query_result = fileStorage.client_->Query(key); + ASSERT_FALSE(query_result.has_value()); + EXPECT_EQ(query_result.error(), ErrorCode::OBJECT_NOT_FOUND); + } + } + void TearDown() override { google::ShutdownGoogleLogging(); LOG(INFO) << "Clear test data..."; @@ -280,6 +356,9 @@ TEST_F(FileStorageTest, DefaultValuesWhenNoEnvSet) { EXPECT_EQ(config.total_keys_limit, 10'000'000); EXPECT_EQ(config.total_size_limit, 2ULL * 1024 * 1024 * 1024 * 1024); EXPECT_EQ(config.heartbeat_interval_seconds, 10u); + EXPECT_TRUE(config.enable_disk_watermark_eviction); + EXPECT_DOUBLE_EQ(config.disk_eviction_high_watermark_ratio, 0.90); + EXPECT_DOUBLE_EQ(config.disk_eviction_low_watermark_ratio, 0.80); } TEST_F(FileStorageTest, ReadStringFromEnv) { @@ -309,6 +388,165 @@ TEST_F(FileStorageTest, ReadUint32FromEnv) { EXPECT_EQ(config.heartbeat_interval_seconds, 5u); } +TEST_F(FileStorageTest, ReadDiskWatermarkConfigFromEnv) { + SetEnv("MOONCAKE_DISK_EVICTION_HIGH_WATERMARK_RATIO", "0.77"); + SetEnv("MOONCAKE_DISK_EVICTION_LOW_WATERMARK_RATIO", "0.55"); + + auto alias_config = FileStorageConfig::FromEnvironment(); + EXPECT_DOUBLE_EQ(alias_config.disk_eviction_high_watermark_ratio, 0.77); + EXPECT_DOUBLE_EQ(alias_config.disk_eviction_low_watermark_ratio, 0.55); + + SetEnv("MOONCAKE_OFFLOAD_ENABLE_DISK_WATERMARK_EVICTION", "0"); + SetEnv("MOONCAKE_OFFLOAD_DISK_EVICTION_HIGH_WATERMARK_RATIO", "0.75"); + SetEnv("MOONCAKE_OFFLOAD_DISK_EVICTION_LOW_WATERMARK_RATIO", "0.50"); + + auto config = FileStorageConfig::FromEnvironment(); + EXPECT_FALSE(config.enable_disk_watermark_eviction); + EXPECT_DOUBLE_EQ(config.disk_eviction_high_watermark_ratio, 0.75); + EXPECT_DOUBLE_EQ(config.disk_eviction_low_watermark_ratio, 0.50); + + SetEnv("MOONCAKE_OFFLOAD_DISK_EVICTION_HIGH_WATERMARK_RATIO", "0,75"); + SetEnv("MOONCAKE_OFFLOAD_DISK_EVICTION_LOW_WATERMARK_RATIO", "nan"); + + auto invalid_config = FileStorageConfig::FromEnvironment(); + EXPECT_DOUBLE_EQ(invalid_config.disk_eviction_high_watermark_ratio, 0.90); + EXPECT_DOUBLE_EQ(invalid_config.disk_eviction_low_watermark_ratio, 0.80); +} + +TEST_F(FileStorageTest, HeartbeatRunsDiskWatermarkEvictionWithoutOffloadWork) { + std::filesystem::path master_root = + std::filesystem::path(data_path) / "heartbeat_master"; + std::filesystem::create_directories(master_root); + + testing::InProcMaster master; + auto master_config = InProcMasterConfigBuilder() + .set_enable_offload(true) + .set_root_fs_dir(master_root.string()) + .build(); + ASSERT_TRUE(master.Start(master_config)); + + std::string local_rpc_addr = + "127.0.0.1:" + std::to_string(getFreeTcpPort()); + auto client = Client::Create(local_rpc_addr, master.metadata_url(), "tcp", + std::nullopt, master.master_address()); + ASSERT_TRUE(client.has_value()); + auto mount_result = client.value()->MountLocalDiskSegment(true); + ASSERT_TRUE(mount_result.has_value()) + << "MountLocalDiskSegment failed: " << toString(mount_result.error()); + + FileStorageConfig config = FileStorageConfig::FromEnvironment(); + config.storage_backend_type = StorageBackendType::kFilePerKey; + config.storage_filepath = data_path + "/heartbeat_watermark"; + config.local_buffer_size = 4 * 1024 * 1024; + config.disk_eviction_high_watermark_ratio = 1e-12; + config.disk_eviction_low_watermark_ratio = 0.5e-12; + fs::create_directories(config.storage_filepath); + + FileStorage file_storage(config, client.value(), local_rpc_addr); + + std::unordered_map> batch_object; + std::vector> buffers; + std::vector expected_keys = { + "heartbeat_key_1", "heartbeat_key_2", "heartbeat_key_3"}; + for (size_t i = 0; i < expected_keys.size(); ++i) { + auto buffer = std::make_unique(512); + std::memset(buffer.get(), static_cast('a' + i), 512); + batch_object.emplace(expected_keys[i], + std::vector{Slice{buffer.get(), 512}}); + buffers.push_back(std::move(buffer)); + } + + AssertHeartbeatEvictsAllKeys(file_storage, expected_keys, batch_object); +} + +TEST_F(FileStorageTest, NotifyEvictedDiskReplicasUsesTenantScopedKeys) { + std::filesystem::path master_root = + std::filesystem::path(data_path) / "tenant_notify_master"; + std::filesystem::create_directories(master_root); + + testing::InProcMaster master; + auto master_config = InProcMasterConfigBuilder() + .set_enable_offload(true) + .set_root_fs_dir(master_root.string()) + .build(); + ASSERT_TRUE(master.Start(master_config)); + + std::string local_rpc_addr = + "127.0.0.1:" + std::to_string(getFreeTcpPort()); + auto client = Client::Create(local_rpc_addr, master.metadata_url(), "tcp", + std::nullopt, master.master_address()); + ASSERT_TRUE(client.has_value()); + auto mount_result = client.value()->MountLocalDiskSegment(true); + ASSERT_TRUE(mount_result.has_value()) + << "MountLocalDiskSegment failed: " << toString(mount_result.error()); + + FileStorageConfig config = FileStorageConfig::FromEnvironment(); + config.storage_backend_type = StorageBackendType::kFilePerKey; + config.storage_filepath = data_path + "/tenant_notify"; + fs::create_directories(config.storage_filepath); + FileStorage file_storage(config, client.value(), local_rpc_addr); + + const std::string key = "shared_key"; + std::vector tasks = { + {.tenant_id = "tenant_a", .key = key, .size = 128}, + {.tenant_id = "tenant_b", .key = key, .size = 128}, + }; + std::vector metadatas; + metadatas.reserve(tasks.size()); + for (const auto& task : tasks) { + metadatas.push_back(StorageObjectMetadata{ + .bucket_id = 0, + .offset = 0, + .key_size = static_cast(task.key.size()), + .data_size = task.size, + .transport_endpoint = local_rpc_addr, + }); + } + ASSERT_TRUE(client.value()->NotifyOffloadSuccess(tasks, metadatas)); + + for (const auto& task : tasks) { + auto before = client.value()->BatchQuery({key}, task.tenant_id); + ASSERT_EQ(before.size(), 1); + ASSERT_TRUE(before[0].has_value()); + bool has_local_disk_replica = false; + for (const auto& replica : before[0]->replicas) { + has_local_disk_replica |= replica.is_local_disk_replica(); + } + ASSERT_TRUE(has_local_disk_replica); + } + + auto notify_result = FileStorageNotifyEvictedDiskReplicas( + file_storage, {TenantId("tenant_a").MakeScopedKey(key), + TenantId("tenant_b").MakeScopedKey(key)}); + ASSERT_TRUE(notify_result.has_value()); + + for (const auto& task : tasks) { + auto after = client.value()->BatchQuery({key}, task.tenant_id); + ASSERT_EQ(after.size(), 1); + ASSERT_FALSE(after[0].has_value()); + EXPECT_EQ(after[0].error(), ErrorCode::OBJECT_NOT_FOUND); + } +} + +// Regression test for issue #2827: under concurrent/repeat offload of the same +// keys, the bucket backend rejects a whole bucket atomically with +// OBJECT_ALREADY_EXISTS (see BucketStorageBackend duplicate-key tests). That +// error must be treated as a per-bucket soft failure so OffloadObjects reports +// the keys back to the master and continues, rather than aborting the whole +// offload cycle and leaving master/SSD metadata inconsistent (which surfaced as +// spurious INVALID_KEY on the read path). It stays alongside INVALID_READ, +// while genuinely fatal errors (e.g. KEYS_ULTRA_LIMIT, INTERNAL_ERROR) do not. +TEST_F(FileStorageTest, DuplicateOffloadErrorIsPerBucketSoftError) { + EXPECT_TRUE( + CallIsPerBucketSoftOffloadError(ErrorCode::OBJECT_ALREADY_EXISTS)); + EXPECT_TRUE(CallIsPerBucketSoftOffloadError(ErrorCode::INVALID_READ)); + + EXPECT_FALSE(CallIsPerBucketSoftOffloadError(ErrorCode::KEYS_ULTRA_LIMIT)); + EXPECT_FALSE(CallIsPerBucketSoftOffloadError(ErrorCode::INTERNAL_ERROR)); + EXPECT_FALSE(CallIsPerBucketSoftOffloadError(ErrorCode::INVALID_KEY)); + EXPECT_FALSE(CallIsPerBucketSoftOffloadError(ErrorCode::OK)); +} + TEST_F(FileStorageTest, InvalidIntValueUsesDefault) { SetEnv("MOONCAKE_OFFLOAD_BUCKET_KEYS_LIMIT", "abc"); SetEnv("MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES", "sdfsdf"); @@ -320,6 +558,8 @@ TEST_F(FileStorageTest, InvalidIntValueUsesDefault) { EXPECT_EQ(bucket_backend_config.bucket_keys_limit, 500); EXPECT_EQ(config.total_size_limit, 2ULL * 1024 * 1024 * 1024 * 1024); EXPECT_EQ(config.heartbeat_interval_seconds, 10u); + EXPECT_DOUBLE_EQ(config.disk_eviction_high_watermark_ratio, 0.90); + EXPECT_DOUBLE_EQ(config.disk_eviction_low_watermark_ratio, 0.80); } TEST_F(FileStorageTest, OutOfRangeValueUsesDefault) { @@ -385,6 +625,19 @@ TEST_F(FileStorageTest, ValidateFailsOnInvalidLimits) { config.total_size_limit = 1; config.heartbeat_interval_seconds = 0; EXPECT_FALSE(config.Validate()); + + config.heartbeat_interval_seconds = 1; + config.disk_eviction_low_watermark_ratio = 0.9; + config.disk_eviction_high_watermark_ratio = 0.8; + EXPECT_FALSE(config.Validate()); + + config.disk_eviction_low_watermark_ratio = 0.0; + config.disk_eviction_high_watermark_ratio = 0.8; + EXPECT_FALSE(config.Validate()); + + config.disk_eviction_low_watermark_ratio = 0.8; + config.disk_eviction_high_watermark_ratio = 1.1; + EXPECT_FALSE(config.Validate()); } TEST_F(FileStorageTest, BatchLoad_WithStorageBackendAdaptor) { @@ -542,4 +795,4 @@ TEST_F(FileStorageTest, NullSsdMetricDoesNotCrash) { // No crash = success. No metrics pointer, so nothing to verify. } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-store/tests/ha/leadership/ha_backend_availability_test.cpp b/mooncake-store/tests/ha/leadership/ha_backend_availability_test.cpp index 047ae862bb..e301a05f56 100644 --- a/mooncake-store/tests/ha/leadership/ha_backend_availability_test.cpp +++ b/mooncake-store/tests/ha/leadership/ha_backend_availability_test.cpp @@ -36,16 +36,26 @@ TEST(HABackendAvailabilityTest, RedisAvailabilityMatchesBuildFlag) { #endif } -TEST(HABackendAvailabilityTest, K8sLeaseIsRejectedUntilCoordinatorExists) { +TEST(HABackendAvailabilityTest, K8sAvailabilityMatchesBuildFlag) { +#ifdef STORE_USE_K8S_LEASE + EXPECT_EQ(ErrorCode::OK, ValidateHABackendAvailability(HABackendType::K8S)); +#else EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE, ValidateHABackendAvailability(HABackendType::K8S)); +#endif } -TEST(HABackendAvailabilityTest, - ClientSpecParsingRejectsUnavailableBackendBeforeCoordinatorCreation) { +TEST(HABackendAvailabilityTest, ClientSpecParsingMatchesK8sBuildFlag) { auto spec = ParseHABackendSpec("k8s://default/master"); +#ifdef STORE_USE_K8S_LEASE + ASSERT_TRUE(spec.has_value()); + ASSERT_TRUE(spec->has_value()); + EXPECT_EQ(HABackendType::K8S, (*spec)->type); + EXPECT_EQ("default/master", (*spec)->connstring); +#else ASSERT_FALSE(spec.has_value()); EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_MODE, spec.error()); +#endif } } // namespace diff --git a/mooncake-store/tests/ha/leadership/leader_label_reconciler_test.cpp b/mooncake-store/tests/ha/leadership/leader_label_reconciler_test.cpp new file mode 100644 index 0000000000..799fd74bce --- /dev/null +++ b/mooncake-store/tests/ha/leadership/leader_label_reconciler_test.cpp @@ -0,0 +1,186 @@ +#include "ha/leadership/leader_label_reconciler.h" + +#include + +#include +#include +#include +#include +#include + +namespace mooncake { +namespace ha { +namespace { + +using namespace std::chrono_literals; + +// Records every apply call and can be told to fail the first N attempts to +// simulate a flaky / unreachable K8s API server. +class FakeLabelBackend { + public: + LeaderLabelReconciler::ApplyFn Fn() { + return [this](bool desired) { return Apply(desired); }; + } + + void FailNext(int count) { + std::lock_guard lock(mutex_); + fail_remaining_ = count; + } + + void CommitThenFailNext(int count) { + std::lock_guard lock(mutex_); + commit_then_fail_remaining_ = count; + } + + // Waits until the last successful apply matches `want`, or times out. + bool WaitForApplied(bool want, std::chrono::milliseconds timeout) { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, timeout, [&] { return applied_ == want; }); + } + + bool WaitForCommitted(bool want, std::chrono::milliseconds timeout) { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, timeout, [&] { return committed_ == want; }); + } + + int CallCount() { + std::lock_guard lock(mutex_); + return call_count_; + } + + std::optional Applied() { + std::lock_guard lock(mutex_); + return applied_; + } + + std::optional Committed() { + std::lock_guard lock(mutex_); + return committed_; + } + + private: + ErrorCode Apply(bool desired) { + std::lock_guard lock(mutex_); + ++call_count_; + if (commit_then_fail_remaining_ > 0) { + --commit_then_fail_remaining_; + committed_ = desired; + cv_.notify_all(); + return ErrorCode::K8S_LEASE_OPERATION_ERROR; + } + if (fail_remaining_ > 0) { + --fail_remaining_; + cv_.notify_all(); + return ErrorCode::K8S_LEASE_OPERATION_ERROR; + } + committed_ = desired; + applied_ = desired; + cv_.notify_all(); + return ErrorCode::OK; + } + + std::mutex mutex_; + std::condition_variable cv_; + int call_count_ = 0; + int fail_remaining_ = 0; + int commit_then_fail_remaining_ = 0; + std::optional applied_; + std::optional committed_; +}; + +constexpr auto kRetry = 5ms; +constexpr auto kTimeout = 2s; +constexpr auto kSlowRetry = 10s; + +TEST(LeaderLabelReconcilerTest, ConvergesToLeaderAfterTransientFailures) { + FakeLabelBackend backend; + backend.FailNext(3); + LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), kRetry); + + reconciler.SetLeader(true); + + ASSERT_TRUE(backend.WaitForApplied(true, kTimeout)); + EXPECT_GT(backend.CallCount(), 3); +} + +TEST(LeaderLabelReconcilerTest, ClearsLeaderAfterTransientFailures) { + FakeLabelBackend backend; + LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), kRetry); + + reconciler.SetLeader(true); + ASSERT_TRUE(backend.WaitForApplied(true, kTimeout)); + + backend.FailNext(3); + reconciler.SetLeader(false); + + ASSERT_TRUE(backend.WaitForApplied(false, kTimeout)); +} + +TEST(LeaderLabelReconcilerTest, LatestDesiredStateWins) { + FakeLabelBackend backend; + LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), kRetry); + + reconciler.SetLeader(true); + reconciler.SetLeader(false); + + ASSERT_TRUE(backend.WaitForApplied(false, kTimeout)); + EXPECT_EQ(backend.Applied(), std::make_optional(false)); +} + +TEST(LeaderLabelReconcilerTest, StopsCallingOnceConverged) { + FakeLabelBackend backend; + LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), kRetry); + + reconciler.SetLeader(true); + ASSERT_TRUE(backend.WaitForApplied(true, kTimeout)); + + int settled = backend.CallCount(); + std::this_thread::sleep_for(kRetry * 20); + EXPECT_EQ(backend.CallCount(), settled); +} + +TEST(LeaderLabelReconcilerTest, DisabledNeverApplies) { + FakeLabelBackend backend; + LeaderLabelReconciler reconciler(/*enabled=*/false, backend.Fn(), kRetry); + + reconciler.SetLeader(true); + std::this_thread::sleep_for(kRetry * 20); + + EXPECT_EQ(backend.CallCount(), 0); + EXPECT_EQ(backend.Applied(), std::nullopt); +} + +TEST(LeaderLabelReconcilerTest, AmbiguousSetFailureDoesNotLeaveStaleLeader) { + FakeLabelBackend backend; + backend.CommitThenFailNext(1); + LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), + kSlowRetry); + + reconciler.SetLeader(true); + ASSERT_TRUE(backend.WaitForCommitted(true, kTimeout)); + + reconciler.SetLeader(false); + + ASSERT_TRUE(backend.WaitForCommitted(false, kTimeout)); +} + +TEST(LeaderLabelReconcilerTest, AmbiguousClearFailureDoesNotLeaveMissingLabel) { + FakeLabelBackend backend; + LeaderLabelReconciler reconciler(/*enabled=*/true, backend.Fn(), + kSlowRetry); + + reconciler.SetLeader(true); + ASSERT_TRUE(backend.WaitForCommitted(true, kTimeout)); + + backend.CommitThenFailNext(1); + reconciler.SetLeader(false); + ASSERT_TRUE(backend.WaitForCommitted(false, kTimeout)); + + reconciler.SetLeader(true); + + ASSERT_TRUE(backend.WaitForCommitted(true, kTimeout)); +} + +} // namespace +} // namespace ha +} // namespace mooncake diff --git a/mooncake-store/tests/ha/master_service_ha_test.cpp b/mooncake-store/tests/ha/master_service_ha_test.cpp new file mode 100644 index 0000000000..a898480109 --- /dev/null +++ b/mooncake-store/tests/ha/master_service_ha_test.cpp @@ -0,0 +1,3660 @@ +#include "master_service.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "hot_standby_service.h" +#include "ha/kv/ha_kv_backend.h" +#include "ha/oplog/mock_metadata_store.h" +#include "ha/oplog/oplog_batch_codec.h" +#include "ha/oplog/oplog_batch_storage.h" +#include "ha/oplog/oplog_batch_standby_reader.h" +#include "ha/oplog/oplog_batch_types.h" +#include "ha/oplog/oplog_applier.h" +#include "ha/oplog/ordered_oplog_writer.h" +#include "types.h" + +namespace mooncake::test { + +class FakeBatchHaKvBackend : public HaKvBackend { + public: + ErrorCode Get(std::string_view key, std::string& value) override { + auto it = kvs_.find(std::string(key)); + if (it == kvs_.end()) { + return ErrorCode::ETCD_KEY_NOT_EXIST; + } + value = it->second; + return ErrorCode::OK; + } + + ErrorCode Put(std::string_view key, std::string_view value) override { + kvs_[std::string(key)] = std::string(value); + return ErrorCode::OK; + } + + ErrorCode Range(std::string_view begin_key, std::string_view end_key, + size_t limit, std::vector& kvs) override { + kvs.clear(); + for (auto it = kvs_.lower_bound(std::string(begin_key)); + it != kvs_.end() && it->first < end_key; ++it) { + kvs.push_back({.key = it->first, .value = it->second}); + if (limit != 0 && kvs.size() >= limit) { + break; + } + } + return ErrorCode::OK; + } + + bool SupportsTxn() const override { return true; } + + ErrorCode Txn(const KvTxn& txn) override { + for (const auto& compare : txn.compares) { + auto it = kvs_.find(compare.key); + if (compare.kind == KvCompareKind::kKeyNotExists) { + if (it != kvs_.end()) { + return ErrorCode::ETCD_TRANSACTION_FAIL; + } + } else if (it == kvs_.end() || + it->second != compare.expected_value) { + return ErrorCode::ETCD_TRANSACTION_FAIL; + } + } + for (const auto& put : txn.puts) { + kvs_[put.key] = put.value; + } + return ErrorCode::OK; + } + + private: + std::map kvs_; +}; + +class BlockingBatchHaKvBackend : public FakeBatchHaKvBackend { + public: + void BlockTxn() { + std::lock_guard lock(mutex_); + blocked_ = true; + } + + void AllowTxn() { + { + std::lock_guard lock(mutex_); + blocked_ = false; + } + cv_.notify_all(); + } + + ErrorCode Txn(const KvTxn& txn) override { + { + std::unique_lock lock(mutex_); + cv_.wait(lock, [this] { return !blocked_; }); + } + return FakeBatchHaKvBackend::Txn(txn); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + bool blocked_{false}; +}; + +class FailingBatchHaKvBackend : public FakeBatchHaKvBackend { + public: + void SetTxnError(ErrorCode error) { + std::lock_guard lock(mutex_); + txn_error_ = error; + txn_calls_ = 0; + } + + bool WaitForTxnCalls(size_t count, std::chrono::milliseconds timeout = + std::chrono::milliseconds(1000)) { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, timeout, [&] { return txn_calls_ >= count; }); + } + + ErrorCode Txn(const KvTxn& txn) override { + ErrorCode txn_error; + { + std::lock_guard lock(mutex_); + ++txn_calls_; + txn_error = txn_error_; + } + cv_.notify_all(); + if (txn_error != ErrorCode::OK) { + return txn_error; + } + return FakeBatchHaKvBackend::Txn(txn); + } + + private: + std::mutex mutex_; + std::condition_variable cv_; + ErrorCode txn_error_{ErrorCode::OK}; + size_t txn_calls_{0}; +}; + +class MasterServiceHATest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + google::InitGoogleLogging("MasterServiceHATest"); + FLAGS_logtostderr = 1; + } + + static void TearDownTestSuite() { google::ShutdownGoogleLogging(); } + + static constexpr size_t kDefaultSegmentBase = 0x300000000; + static constexpr size_t kDefaultSegmentSize = 1024 * 1024 * 16; + static constexpr uint64_t kStrictTenantQuotaBytes = 4 * 1024 * 1024; + inline static const TenantId kDefaultTenant = TenantId::Default(); + + void SetUp() override { + ::setenv("MOONCAKE_SNAPSHOT_LOCAL_PATH", LegacyOpLogRootDir().c_str(), + 1); + } + + void TearDown() override { + for (const auto& path : policy_files_) { + std::error_code ec; + std::filesystem::remove(path, ec); + } + policy_files_.clear(); + std::error_code ec; + std::filesystem::remove_all(LegacyOpLogRootDir(), ec); + ::unsetenv("MOONCAKE_SNAPSHOT_LOCAL_PATH"); + } + + std::string LegacyOpLogRootDir() const { + return (std::filesystem::temp_directory_path() / + ("mooncake_master_service_ha_oplog_" + + std::to_string(::getpid()))) + .string(); + } + + std::string WriteTenantPolicyFile( + const std::map& tenant_quotas) { + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = tenant_quotas; + auto path = + std::filesystem::temp_directory_path() / + ("mooncake_master_service_ha_test_" + std::to_string(::getpid()) + + "_" + std::to_string(next_policy_file_++) + ".yaml"); + std::ofstream out(path); + out << FormatTenantQuotaPolicyYaml(snapshot); + out.close(); + policy_files_.push_back(path.string()); + return path.string(); + } + + MasterServiceConfig MakeStrictHAConfig( + const std::vector& tenants = { + std::string(kDefaultTenant.value()), "tenant_a"}) { + std::map tenant_quotas; + for (const auto& tenant : tenants) { + tenant_quotas.emplace(tenant, kStrictTenantQuotaBytes); + } + return MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id("test_cluster") + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile(tenant_quotas)) + .build(); + } + + static bool HasOpLogWriter(const MasterService& service) { + return service.ordered_oplog_writer_ != nullptr; + } + + static bool IsOpLogEnabled(const MasterService& service) { + return service.enable_oplog_; + } + + static bool HasBatchOpLogStorage(const MasterService& service) { + return service.batch_oplog_storage_ != nullptr; + } + + Segment MakeSegment(std::string name = "test_segment", + size_t base = kDefaultSegmentBase, + size_t size = kDefaultSegmentSize) const { + Segment segment; + segment.id = generate_uuid(); + segment.name = std::move(name); + segment.base = base; + segment.size = size; + segment.te_endpoint = segment.name; + return segment; + } + +#ifdef USE_NOF + NoFSegment MakeNoFSegment( + std::string name = "test_nof_segment", + std::string endpoint = "test_nof_segment_endpoint", + size_t base = kDefaultSegmentBase + kDefaultSegmentSize, + size_t size = kDefaultSegmentSize) const { + NoFSegment segment; + segment.id = generate_uuid(); + segment.name = std::move(name); + segment.base = base; + segment.size = size; + segment.te_endpoint = std::move(endpoint); + return segment; + } +#endif + + struct MountedSegmentContext { + UUID segment_id; + UUID client_id; + }; + + MountedSegmentContext PrepareSimpleSegment( + MasterService& service, std::string name = "test_segment", + size_t base = kDefaultSegmentBase, + size_t size = kDefaultSegmentSize) const { + Segment segment = MakeSegment(std::move(name), base, size); + UUID client_id = generate_uuid(); + auto mount_result = service.MountSegment(segment, client_id); + EXPECT_TRUE(mount_result.has_value()); + return {.segment_id = segment.id, .client_id = client_id}; + } + + std::string PutObject(MasterService& service, const UUID& client_id, + const std::string& key, + size_t slice_length = 1024) const { + ReplicateConfig config; + config.replica_num = 1; + auto put_start = service.PutStart(client_id, key, kDefaultTenant, + slice_length, config); + EXPECT_TRUE(put_start.has_value()); + EXPECT_TRUE( + service.PutEnd(client_id, key, kDefaultTenant, ReplicaType::MEMORY) + .has_value()); + return key; + } + + std::string PutObjectOnSegment(MasterService& service, + const UUID& client_id, + const std::string& key, + const std::string& segment_name, + size_t slice_length = 1024) const { + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {segment_name}; + auto put_start = service.PutStart(client_id, key, kDefaultTenant, + slice_length, config); + EXPECT_TRUE(put_start.has_value()); + EXPECT_TRUE( + service.PutEnd(client_id, key, kDefaultTenant, ReplicaType::MEMORY) + .has_value()); + return key; + } + + void ReadBatchEventually(OpLogBatchStorage& storage, uint64_t batch_id, + OpLogBatchRecord& batch) const { + ErrorCode read_err = ErrorCode::ETCD_KEY_NOT_EXIST; + for (int i = 0; i < 50; ++i) { + read_err = storage.ReadBatch(batch_id, batch); + if (read_err == ErrorCode::OK) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_EQ(ErrorCode::OK, read_err); + } + + void ReadRemoveBatchEventually(OpLogBatchStorage& storage, + uint64_t first_batch_id, + const std::string& key, + OpLogBatchRecord& batch) const { + bool saw_remove = false; + for (int attempt = 0; attempt < 50 && !saw_remove; ++attempt) { + for (uint64_t batch_id = first_batch_id; + batch_id < first_batch_id + 8 && !saw_remove; ++batch_id) { + if (storage.ReadBatch(batch_id, batch) != ErrorCode::OK) { + continue; + } + for (const auto& entry : batch.entries) { + if (entry.op_type == OpType::REMOVE && + entry.object_key == key) { + saw_remove = true; + break; + } + } + } + if (!saw_remove) { + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + } + ASSERT_TRUE(saw_remove); + } + + std::string PutObjectWithTenant(MasterService& service, + const UUID& client_id, + const std::string& key, + const std::string& tenant_id, + size_t slice_length = 1024) const { + ReplicateConfig config; + config.replica_num = 1; + const TenantId tenant(tenant_id); + auto put_start = + service.PutStart(client_id, key, tenant, slice_length, config); + EXPECT_TRUE(put_start.has_value()); + EXPECT_TRUE(service.PutEnd(client_id, key, tenant, ReplicaType::MEMORY) + .has_value()); + return key; + } + + std::string PutObjectOnSegmentWithTenant(MasterService& service, + const UUID& client_id, + const std::string& key, + const std::string& segment_name, + const std::string& tenant_id, + size_t slice_length = 1024) const { + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {segment_name}; + const TenantId tenant(tenant_id); + auto put_start = + service.PutStart(client_id, key, tenant, slice_length, config); + EXPECT_TRUE(put_start.has_value()); + EXPECT_TRUE(service.PutEnd(client_id, key, tenant, ReplicaType::MEMORY) + .has_value()); + return key; + } + + Replica::Descriptor MakeStandbyMemoryReplica(const std::string& endpoint, + size_t size = 1024) const { + Replica::Descriptor replica; + replica.id = 1; + replica.status = ReplicaStatus::COMPLETE; + + MemoryDescriptor mem_desc; + mem_desc.buffer_descriptor.transport_endpoint_ = endpoint; + mem_desc.buffer_descriptor.buffer_address_ = 0; + mem_desc.buffer_descriptor.size_ = size; + replica.descriptor_variant = std::move(mem_desc); + return replica; + } + + StandbyObjectEntry MakeStandbyObject(const std::string& key, + const std::string& endpoint, + size_t size = 1024) const { + StandbyObjectMetadata metadata; + metadata.client_id = generate_uuid(); + metadata.size = size; + metadata.last_sequence_id = 1; + metadata.replicas.push_back(MakeStandbyMemoryReplica(endpoint, size)); + return StandbyObjectEntry{"default", key, std::move(metadata)}; + } + + StandbySegmentInfo MakeStandbyMemorySegment( + const std::string& endpoint, + size_t capacity = kDefaultSegmentSize) const { + StandbySegmentInfo segment; + segment.segment_name = endpoint; + segment.transport_endpoint = endpoint; + segment.capacity = capacity; + segment.is_memory_segment = true; + return segment; + } + + // Friend access to MasterService::metadata_shards_ and + // getMetadataShardIndex, which are otherwise private. + // MasterServiceHATest is friended; TEST_F-generated subclasses are not, + // hence this static funnel. Seeds an in-flight PromotionTask for a + // given (tenant, key) so NotifyPromotionSuccess can proceed without + // going through the on-hit admission gate (which is currently + // restricted to the "default" tenant). Used only by the non-default + // tenant promotion tests. + static void SeedPromotionTaskForTesting( + MasterService* service, const TenantId& tenant, const std::string& key, + const UUID& holder_id, ReplicaID alloc_id, uint64_t object_size) { + const size_t shard_idx = service->getMetadataShardIndex(tenant, key); + auto shard_access = + MasterService::MetadataShardAccessorRW(service, shard_idx); + auto& tenant_state = shard_access->tenants[tenant]; + tenant_state.promotion_tasks.emplace( + key, MasterService::PromotionTask{ + .source_id = 0, + .alloc_id = alloc_id, + .object_size = object_size, + .start_time = std::chrono::system_clock::now(), + .holder_id = holder_id}); + } + + static bool SnapshotManagerCreatedForTesting(const MasterService& service) { + return service.snapshot_manager_ != nullptr; + } + + static tl::expected AppendVisibleForTesting( + MasterService& service, OpType type, const std::string& tenant_id, + const std::string& key, const std::string& payload) { + return service.AppendOpLogVisibleBeforeDurable(type, tenant_id, key, + payload); + } + + static tl::expected AppendVisibleForTesting( + MasterService& service, OpType type, const TenantId& tenant_id, + const std::string& key, const std::string& payload) { + return service.AppendOpLogVisibleBeforeDurable(type, tenant_id.value(), + key, payload); + } + + static tl::expected AppendFinalizeForTesting( + MasterService& service, OpType type, const std::string& tenant_id, + const std::string& key, const std::string& payload, + MasterService::DurableFinalizeCallback callback) { + return service.AppendOpLogWithDurableFinalize( + type, tenant_id, key, payload, std::move(callback)); + } + + static tl::expected AppendFinalizeForTesting( + MasterService& service, OpType type, const TenantId& tenant_id, + const std::string& key, const std::string& payload, + MasterService::DurableFinalizeCallback callback) { + return service.AppendOpLogWithDurableFinalize( + type, tenant_id.value(), key, payload, std::move(callback)); + } + + static tl::expected + ReserveBatchSlotForTesting(MasterService& service) { + if (!service.ordered_oplog_writer_) { + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + } + return service.ordered_oplog_writer_->Reserve(); + } + + static void ClearInvalidHandlesForTesting( + MasterService& service, + const std::unordered_set>& alive_clients) { + service.ClearInvalidHandles(alive_clients); + } + + static size_t ReplicaCountForTesting(MasterService& service, + const TenantId& tenant_id, + const std::string& key) { + MasterService::MetadataAccessorRO accessor( + &service, MasterService::ObjectIdentity{tenant_id, key}); + return accessor.Exists() ? accessor.Get().CountReplicas() : 0; + } + + static std::vector ReplicaDescriptorsForTesting( + MasterService& service, const TenantId& tenant_id, + const std::string& key) { + MasterService::MetadataAccessorRO accessor( + &service, MasterService::ObjectIdentity{tenant_id, key}); + if (!accessor.Exists()) { + return {}; + } + std::vector descriptors; + for (const auto& replica : accessor.Get().GetAllReplicas()) { + descriptors.push_back(replica.get_descriptor()); + } + return descriptors; + } + + static bool HasInvalidMemoryHandleForTesting(MasterService& service, + const TenantId& tenant_id, + const std::string& key) { + MasterService::MetadataAccessorRO accessor( + &service, MasterService::ObjectIdentity{tenant_id, key}); + if (!accessor.Exists()) { + return true; + } + for (const auto& replica : accessor.Get().GetAllReplicas()) { + if (replica.has_invalid_mem_handle()) { + return true; + } + } + return false; + } + + static size_t SegmentAllocatedSizeForTesting(MasterService& service, + const std::string& name) { + auto access = service.segment_manager_.getAllocatorAccess(); + const auto* allocators = + access.getAllocatorManager().getAllocators(name); + EXPECT_NE(allocators, nullptr); + EXPECT_EQ(allocators == nullptr ? 0 : allocators->size(), 1); + return allocators == nullptr || allocators->empty() + ? 0 + : allocators->front()->size(); + } + + static void EraseObjectForTesting(MasterService& service, + const TenantId& tenant_id, + const std::string& key) { + MasterService::MetadataAccessorRW accessor( + &service, MasterService::ObjectIdentity{tenant_id, key}); + ASSERT_TRUE(accessor.Exists()); + accessor.Erase(); + } + + static void PrepareUnmountSegmentForTesting(MasterService& service, + const UUID& segment_id) { + auto segment_access = service.segment_manager_.getSegmentAccess(); + size_t metrics_dec_capacity = 0; + ASSERT_EQ(ErrorCode::OK, segment_access.PrepareUnmountSegment( + segment_id, metrics_dec_capacity)); + } + + static std::vector MarkCompletedReplicasRemovedForTesting( + MasterService& service, const TenantId& tenant_id, + const std::string& key) { + MasterService::MetadataAccessorRW accessor( + &service, MasterService::ObjectIdentity{tenant_id, key}); + if (!accessor.Exists()) { + return {}; + } + std::vector ids; + accessor.Get().VisitReplicas(&Replica::fn_is_completed, + [&ids](Replica& replica) { + ids.push_back(replica.id()); + replica.mark_removed(); + }); + return ids; + } + + static void FinalizeRemovedReplicasForTesting( + MasterService& service, const OpLogEntry& durable_entry, + const std::vector& replica_ids) { + service.FinalizeRemovedReplicasAfterDurable( + durable_entry, replica_ids, MasterService::QuotaEraseMode::kFull); + } + + static void SetLocalDiskUsedBytesForTesting(MasterService& service, + const UUID& client_id, + int64_t used_bytes) { + auto access = service.segment_manager_.getLocalDiskSegmentAccess(); + access.getClientLocalDiskSegment().at(client_id)->ssd_used_bytes.store( + used_bytes, std::memory_order_relaxed); + } + + static int64_t GetLocalDiskUsedBytesForTesting( + MasterService& service, const std::string& segment_name) { + auto access = service.segment_manager_.getLocalDiskSegmentAccess(); + return access.getSsdUsedBytes(segment_name); + } + + std::vector policy_files_; + int next_policy_file_{0}; +}; + +class MasterServiceBatchRecordE2ETest : public MasterServiceHATest {}; + +TEST_F(MasterServiceHATest, RestoreFromStandbyPreservesMemoryBufferDescriptor) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + + const std::string endpoint = "standby_restore_segment"; + const size_t size = 4096; + const uintptr_t address = 0x12345000; + auto object = MakeStandbyObject("standby_restore_key", endpoint, size); + auto& descriptor = object.metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor; + descriptor.buffer_address_ = address; + descriptor.protocol_ = "tcp"; + + service.RestoreFromStandbySnapshot({object}, 7, + {MakeStandbyMemorySegment(endpoint)}); + + auto replicas = ReplicaDescriptorsForTesting(service, kDefaultTenant, + "standby_restore_key"); + ASSERT_EQ(replicas.size(), 1); + ASSERT_TRUE(replicas.front().is_memory_replica()); + const auto& restored = + replicas.front().get_memory_descriptor().buffer_descriptor; + EXPECT_EQ(restored.size_, size); + EXPECT_EQ(restored.buffer_address_, address); + EXPECT_EQ(restored.protocol_, "tcp"); + EXPECT_EQ(restored.transport_endpoint_, endpoint); +} + +TEST_F(MasterServiceHATest, RemountMakesRestoredMemoryReplicaReady) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + + const std::string endpoint = "standby_remount_segment"; + const auto metric_before = + MasterMetricManager::instance().get_allocated_mem_size(); + const std::string first_key = "standby_remount_first_key"; + const std::string second_key = "standby_remount_second_key"; + auto first_object = MakeStandbyObject(first_key, endpoint); + first_object.metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_ = kDefaultSegmentBase; + auto second_object = MakeStandbyObject(second_key, endpoint); + second_object.metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_ = kDefaultSegmentBase + 4096; + service.RestoreFromStandbySnapshot({first_object, second_object}, 7, + {MakeStandbyMemorySegment(endpoint)}); + EXPECT_EQ(MasterMetricManager::instance().get_allocated_mem_size() - + metric_before, + 2048); + + auto before = service.GetReplicaList(first_key, kDefaultTenant); + ASSERT_FALSE(before.has_value()); + EXPECT_EQ(before.error(), ErrorCode::REPLICA_IS_NOT_READY); + auto batch_before = + service.BatchGetReplicaList({first_key, second_key}, kDefaultTenant); + ASSERT_EQ(batch_before.size(), 2); + for (const auto& result : batch_before) { + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::REPLICA_IS_NOT_READY); + } + + Segment segment = MakeSegment(endpoint); + ASSERT_TRUE(service.ReMountSegment({segment}, generate_uuid()).has_value()); + + auto after = service.GetReplicaList(first_key, kDefaultTenant); + ASSERT_TRUE(after.has_value()) << toString(after.error()); + ASSERT_EQ(after->replicas.size(), 1); + EXPECT_TRUE(after->replicas.front().is_memory_replica()); + EXPECT_FALSE( + HasInvalidMemoryHandleForTesting(service, kDefaultTenant, first_key)); + EXPECT_FALSE( + HasInvalidMemoryHandleForTesting(service, kDefaultTenant, second_key)); + auto batch_after = + service.BatchGetReplicaList({first_key, second_key}, kDefaultTenant); + ASSERT_EQ(batch_after.size(), 2); + ASSERT_TRUE(batch_after[0].has_value()); + ASSERT_TRUE(batch_after[1].has_value()); + EXPECT_EQ(batch_after[0] + ->replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_, + kDefaultSegmentBase); + EXPECT_EQ(batch_after[1] + ->replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_, + kDefaultSegmentBase + 4096); + EXPECT_EQ(SegmentAllocatedSizeForTesting(service, endpoint), 5120); + EXPECT_EQ(MasterMetricManager::instance().get_allocated_mem_size() - + metric_before, + 5120); + + const std::string new_key = "post_remount_allocation"; + PutObjectOnSegment(service, generate_uuid(), new_key, endpoint); + auto new_replicas = + ReplicaDescriptorsForTesting(service, kDefaultTenant, new_key); + ASSERT_EQ(new_replicas.size(), 1); + EXPECT_GE(new_replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_, + kDefaultSegmentBase + 1024); + EXPECT_EQ(MasterMetricManager::instance().get_allocated_mem_size() - + metric_before, + 6144); + + EraseObjectForTesting(service, kDefaultTenant, first_key); + EXPECT_EQ(MasterMetricManager::instance().get_allocated_mem_size() - + metric_before, + 5120); + const std::string replacement_key = "post_remount_replacement"; + PutObjectOnSegment(service, generate_uuid(), replacement_key, endpoint); + auto replacement = + ReplicaDescriptorsForTesting(service, kDefaultTenant, replacement_key); + ASSERT_EQ(replacement.size(), 1); + EXPECT_EQ(replacement.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_, + kDefaultSegmentBase); + EXPECT_EQ(MasterMetricManager::instance().get_allocated_mem_size() - + metric_before, + 6144); +} + +TEST_F(MasterServiceHATest, RemountRestoresCachelibMemoryReplica) { + MasterService service( + MasterServiceConfig::builder() + .set_enable_ha(false) + .set_memory_allocator(BufferAllocatorType::CACHELIB) + .build()); + + const std::string endpoint = "standby_cachelib_remount_segment"; + const auto metric_before = + MasterMetricManager::instance().get_allocated_mem_size(); + const std::string key = "standby_cachelib_remount_key"; + auto object = MakeStandbyObject(key, endpoint, 64); + object.metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_ = kDefaultSegmentBase; + service.RestoreFromStandbySnapshot({object}, 7, + {MakeStandbyMemorySegment(endpoint)}); + EXPECT_EQ(MasterMetricManager::instance().get_allocated_mem_size() - + metric_before, + 64); + + auto single_before = service.GetReplicaList(key, kDefaultTenant); + ASSERT_FALSE(single_before.has_value()); + EXPECT_EQ(single_before.error(), ErrorCode::REPLICA_IS_NOT_READY); + auto batch_before = service.BatchGetReplicaList({key}, kDefaultTenant); + ASSERT_EQ(batch_before.size(), 1); + ASSERT_FALSE(batch_before[0].has_value()); + EXPECT_EQ(batch_before[0].error(), ErrorCode::REPLICA_IS_NOT_READY); + Segment segment = MakeSegment(endpoint); + ASSERT_TRUE(service.ReMountSegment({segment}, generate_uuid()).has_value()); + ASSERT_TRUE(service.GetReplicaList(key, kDefaultTenant).has_value()); + auto batch_after = service.BatchGetReplicaList({key}, kDefaultTenant); + ASSERT_EQ(batch_after.size(), 1); + ASSERT_TRUE(batch_after[0].has_value()); + EXPECT_FALSE( + HasInvalidMemoryHandleForTesting(service, kDefaultTenant, key)); + EXPECT_EQ(SegmentAllocatedSizeForTesting(service, endpoint), 64); + EXPECT_EQ(MasterMetricManager::instance().get_allocated_mem_size() - + metric_before, + 64); + + PutObjectOnSegment(service, generate_uuid(), "cachelib_after_remount", + endpoint, 64); + const auto old_descriptor = + ReplicaDescriptorsForTesting(service, kDefaultTenant, key)[0] + .get_memory_descriptor() + .buffer_descriptor; + const auto new_descriptor = + ReplicaDescriptorsForTesting(service, kDefaultTenant, + "cachelib_after_remount")[0] + .get_memory_descriptor() + .buffer_descriptor; + EXPECT_NE(new_descriptor.buffer_address_, old_descriptor.buffer_address_); +} + +TEST_F(MasterServiceHATest, FailedRemountKeepsReplicaInvalidAndCanBeRetried) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + + const std::string endpoint = "standby_retry_remount_segment"; + auto first = MakeStandbyObject("standby_retry_first", endpoint); + auto conflicting = MakeStandbyObject("standby_retry_conflict", endpoint); + first.metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_ = kDefaultSegmentBase; + conflicting.metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_ = kDefaultSegmentBase; + service.RestoreFromStandbySnapshot({first, conflicting}, 7, + {MakeStandbyMemorySegment(endpoint)}); + + Segment segment = MakeSegment(endpoint); + auto failed = service.ReMountSegment({segment}, generate_uuid()); + ASSERT_FALSE(failed.has_value()); + EXPECT_EQ(failed.error(), ErrorCode::INVALID_PARAMS); + auto get_after_failure = + service.GetReplicaList("standby_retry_first", kDefaultTenant); + ASSERT_FALSE(get_after_failure.has_value()); + EXPECT_EQ(get_after_failure.error(), ErrorCode::REPLICA_IS_NOT_READY); + auto batch_after_failure = + service.BatchGetReplicaList({"standby_retry_first"}, kDefaultTenant); + ASSERT_EQ(batch_after_failure.size(), 1); + ASSERT_FALSE(batch_after_failure[0].has_value()); + EXPECT_EQ(batch_after_failure[0].error(), ErrorCode::REPLICA_IS_NOT_READY); + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {endpoint}; + auto allocation = service.PutStart(generate_uuid(), "must_not_allocate", + kDefaultTenant, 1024, config); + EXPECT_FALSE(allocation.has_value()); + + EraseObjectForTesting(service, kDefaultTenant, "standby_retry_conflict"); + ASSERT_TRUE(service.ReMountSegment({segment}, generate_uuid()).has_value()); + EXPECT_FALSE(HasInvalidMemoryHandleForTesting(service, kDefaultTenant, + "standby_retry_first")); + EXPECT_TRUE(service.GetReplicaList("standby_retry_first", kDefaultTenant) + .has_value()); +} + +TEST_F(MasterServiceHATest, MultiSegmentRemountFailurePublishesNeitherSegment) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + + const std::string good_endpoint = "standby_atomic_good_segment"; + const std::string bad_endpoint = "standby_atomic_bad_segment"; + auto good = MakeStandbyObject("standby_atomic_good", good_endpoint); + good.metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_ = kDefaultSegmentBase; + auto bad_first = + MakeStandbyObject("standby_atomic_bad_first", bad_endpoint); + auto bad_second = + MakeStandbyObject("standby_atomic_bad_second", bad_endpoint); + for (auto* object : {&bad_first, &bad_second}) { + object->metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor.buffer_address_ = + kDefaultSegmentBase + kDefaultSegmentSize; + } + service.RestoreFromStandbySnapshot( + {good, bad_first, bad_second}, 7, + {MakeStandbyMemorySegment(good_endpoint), + MakeStandbyMemorySegment(bad_endpoint)}); + + Segment good_segment = MakeSegment(good_endpoint); + Segment bad_segment = + MakeSegment(bad_endpoint, kDefaultSegmentBase + kDefaultSegmentSize); + auto remount = + service.ReMountSegment({good_segment, bad_segment}, generate_uuid()); + ASSERT_FALSE(remount.has_value()); + auto good_get = + service.GetReplicaList("standby_atomic_good", kDefaultTenant); + ASSERT_FALSE(good_get.has_value()); + EXPECT_EQ(good_get.error(), ErrorCode::REPLICA_IS_NOT_READY); + auto bad_batch = service.BatchGetReplicaList({"standby_atomic_bad_first"}, + kDefaultTenant); + ASSERT_EQ(bad_batch.size(), 1); + ASSERT_FALSE(bad_batch[0].has_value()); + EXPECT_EQ(bad_batch[0].error(), ErrorCode::REPLICA_IS_NOT_READY); + + for (const auto& endpoint : {good_endpoint, bad_endpoint}) { + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {endpoint}; + auto allocation = + service.PutStart(generate_uuid(), "must_not_allocate_" + endpoint, + kDefaultTenant, 1024, config); + EXPECT_FALSE(allocation.has_value()); + } +} + +TEST_F(MasterServiceHATest, EmptyStandbySegmentCanRemount) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + const std::string endpoint = "standby_empty_segment"; + service.RestoreFromStandbySnapshot({}, 7, + {MakeStandbyMemorySegment(endpoint)}); + + Segment segment = MakeSegment(endpoint); + ASSERT_TRUE(service.ReMountSegment({segment}, generate_uuid()).has_value()); + PutObjectOnSegment(service, generate_uuid(), "empty_segment_new_object", + endpoint); + EXPECT_EQ(SegmentAllocatedSizeForTesting(service, endpoint), 1024); +} + +TEST_F(MasterServiceHATest, RemountRejectsStandbySegmentNameMismatch) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + const std::string name = "standby_identity_name"; + const std::string endpoint = "standby_identity_endpoint"; + StandbySegmentInfo standby = MakeStandbyMemorySegment(endpoint); + standby.segment_name = name; + service.RestoreFromStandbySnapshot({}, 7, {standby}); + + Segment mismatched = MakeSegment("wrong_name"); + mismatched.te_endpoint = endpoint; + auto failed = service.ReMountSegment({mismatched}, generate_uuid()); + ASSERT_FALSE(failed.has_value()); + EXPECT_EQ(failed.error(), ErrorCode::INVALID_PARAMS); + + Segment correct = MakeSegment(name); + correct.te_endpoint = endpoint; + ASSERT_TRUE(service.ReMountSegment({correct}, generate_uuid()).has_value()); +} + +TEST_F(MasterServiceHATest, RemountRejectsStandbySegmentEndpointMismatch) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + const std::string name = "standby_endpoint_name"; + const std::string endpoint = "standby_endpoint_value"; + StandbySegmentInfo standby = MakeStandbyMemorySegment(endpoint); + standby.segment_name = name; + service.RestoreFromStandbySnapshot({}, 7, {standby}); + + Segment mismatched = MakeSegment(name); + mismatched.te_endpoint = "wrong_endpoint"; + auto failed = service.ReMountSegment({mismatched}, generate_uuid()); + ASSERT_FALSE(failed.has_value()); + EXPECT_EQ(failed.error(), ErrorCode::INVALID_PARAMS); + + Segment correct = MakeSegment(name); + correct.te_endpoint = endpoint; + ASSERT_TRUE(service.ReMountSegment({correct}, generate_uuid()).has_value()); +} + +TEST_F(MasterServiceHATest, RemountRejectsCxlForStandbyMemorySegment) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + const std::string endpoint = "standby_protocol_segment"; + service.RestoreFromStandbySnapshot({}, 7, + {MakeStandbyMemorySegment(endpoint)}); + + Segment mismatched = MakeSegment(endpoint); + mismatched.protocol = "cxl"; + auto failed = service.ReMountSegment({mismatched}, generate_uuid()); + ASSERT_FALSE(failed.has_value()); + EXPECT_EQ(failed.error(), ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + + Segment correct = MakeSegment(endpoint); + correct.protocol = "tcp"; + ASSERT_TRUE(service.ReMountSegment({correct}, generate_uuid()).has_value()); +} + +TEST_F(MasterServiceHATest, RestoreFromStandbyPreservesCxlBufferDescriptor) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + + const std::string segment_name = "standby_restore_cxl_segment"; + const std::string transport_endpoint = "standby_restore_tcp_endpoint"; + const size_t size = 4096; + const uintptr_t address = 0x12345000; + auto object = + MakeStandbyObject("standby_restore_cxl_key", segment_name, size); + auto& descriptor = object.metadata.replicas.front() + .get_memory_descriptor() + .buffer_descriptor; + descriptor.buffer_address_ = address; + descriptor.protocol_ = "cxl"; + + StandbySegmentInfo segment = MakeStandbyMemorySegment(transport_endpoint); + segment.segment_name = segment_name; + service.RestoreFromStandbySnapshot({object}, 7, {segment}); + + auto replicas = ReplicaDescriptorsForTesting(service, kDefaultTenant, + "standby_restore_cxl_key"); + ASSERT_EQ(replicas.size(), 1); + ASSERT_TRUE(replicas.front().is_memory_replica()); + const auto& restored = + replicas.front().get_memory_descriptor().buffer_descriptor; + EXPECT_EQ(restored.size_, size); + EXPECT_EQ(restored.buffer_address_, address); + EXPECT_EQ(restored.protocol_, "cxl"); + EXPECT_EQ(restored.transport_endpoint_, segment_name); + + auto public_replicas = + service.GetReplicaList("standby_restore_cxl_key", kDefaultTenant); + ASSERT_FALSE(public_replicas.has_value()); + EXPECT_EQ(public_replicas.error(), ErrorCode::REPLICA_IS_NOT_READY); + + Segment remount = MakeSegment(segment_name); + remount.te_endpoint = transport_endpoint; + auto remount_result = service.ReMountSegment({remount}, generate_uuid()); + ASSERT_FALSE(remount_result.has_value()); + EXPECT_EQ(remount_result.error(), ErrorCode::UNAVAILABLE_IN_CURRENT_MODE); + auto single_after = + service.GetReplicaList("standby_restore_cxl_key", kDefaultTenant); + ASSERT_FALSE(single_after.has_value()); + EXPECT_EQ(single_after.error(), ErrorCode::REPLICA_IS_NOT_READY); + auto batch_after = service.BatchGetReplicaList({"standby_restore_cxl_key"}, + kDefaultTenant); + ASSERT_EQ(batch_after.size(), 1); + ASSERT_FALSE(batch_after[0].has_value()); + EXPECT_EQ(batch_after[0].error(), ErrorCode::REPLICA_IS_NOT_READY); +} + +TEST_F(MasterServiceHATest, RemountRejectsExistingSegmentFromDifferentClient) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + Segment segment = MakeSegment("remount_owner_segment"); + const UUID owner = generate_uuid(); + ASSERT_TRUE(service.MountSegment(segment, owner).has_value()); + + auto result = service.ReMountSegment({segment}, generate_uuid()); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); + EXPECT_EQ(SegmentAllocatedSizeForTesting(service, segment.name), 0); +} + +TEST_F(MasterServiceHATest, RemountRejectsMismatchedExistingSegmentIdentity) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + Segment segment = MakeSegment("remount_identity_segment"); + const UUID owner = generate_uuid(); + ASSERT_TRUE(service.MountSegment(segment, owner).has_value()); + + Segment mismatched = segment; + mismatched.base += 4096; + auto result = service.ReMountSegment({mismatched}, owner); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); + EXPECT_EQ(SegmentAllocatedSizeForTesting(service, segment.name), 0); +} + +TEST_F(MasterServiceHATest, RestoreFromStandbyPreservesNoFBufferDescriptor) { + MasterService service( + MasterServiceConfig::builder().set_enable_ha(false).build()); + + const std::string endpoint = "standby_restore_nof_endpoint"; + const size_t size = 4096; + const uintptr_t address = 0x12345000; + Replica::Descriptor replica; + replica.id = 1; + replica.status = ReplicaStatus::COMPLETE; + NoFDescriptor nof_descriptor; + nof_descriptor.buffer_descriptor = {static_cast(size), address, + "nvmeof", endpoint}; + replica.descriptor_variant = std::move(nof_descriptor); + StandbyObjectMetadata metadata; + metadata.client_id = generate_uuid(); + metadata.size = size; + metadata.last_sequence_id = 1; + metadata.replicas.push_back(std::move(replica)); + StandbyObjectEntry object{kDefaultTenant.value(), "standby_restore_nof_key", + std::move(metadata)}; + + service.RestoreFromStandbySnapshot({object}, 7, {}); + + auto replicas = ReplicaDescriptorsForTesting(service, kDefaultTenant, + "standby_restore_nof_key"); + ASSERT_EQ(replicas.size(), 1); + ASSERT_TRUE(replicas.front().is_nof_replica()); + const auto& restored = + replicas.front().get_nof_descriptor().buffer_descriptor; + EXPECT_EQ(restored.size_, size); + EXPECT_EQ(restored.buffer_address_, address); + EXPECT_EQ(restored.protocol_, "nvmeof"); + EXPECT_EQ(restored.transport_endpoint_, endpoint); + + auto public_replicas = + service.GetReplicaList("standby_restore_nof_key", kDefaultTenant); + ASSERT_TRUE(public_replicas.has_value()); + ASSERT_EQ(public_replicas->replicas.size(), 1); + EXPECT_TRUE(public_replicas->replicas.front().is_nof_replica()); +} + +TEST_F(MasterServiceHATest, OplogDisabledByDefaultDoesNotCreateWriter) { + auto config = MasterServiceConfig::builder() + .set_enable_ha(true) + .set_cluster_id("oplog_disabled_by_default") + .build(); + + MasterService service(config); + EXPECT_FALSE(HasOpLogWriter(service)); +} + +TEST_F(MasterServiceHATest, OplogExplicitEnableCreatesWriter) { + auto backend = std::make_shared(); + auto config = MasterServiceConfig::builder() + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id("oplog_explicit_enable") + .build(); + + MasterService service(config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + EXPECT_TRUE(IsOpLogEnabled(service)); + EXPECT_TRUE(HasOpLogWriter(service)); + EXPECT_TRUE(HasBatchOpLogStorage(service)); +} + +TEST_F(MasterServiceHATest, OplogDoesNotStartWithUnsupportedHABackend) { + auto config = MasterServiceConfig::builder() + .set_enable_ha(true) + .set_enable_oplog(true) + .set_ha_backend_type("redis") + .set_cluster_id("oplog_unsupported_ha_backend") + .build(); + + MasterService service(config); + EXPECT_FALSE(HasOpLogWriter(service)); +} + +TEST_F(MasterServiceHATest, BatchPrimaryDoesNotStartSnapshotWorker) { + auto config = + MasterServiceConfig::builder() + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id("batch_snapshot_gate") + .set_enable_snapshot(true) + .set_snapshot_backup_dir(LegacyOpLogRootDir() + "/batch_snapshot") + .set_snapshot_object_store_type("local") + .build(); + + MasterService service(config); + EXPECT_FALSE(SnapshotManagerCreatedForTesting(service)); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + BatchRecordConstructorThrowsWhenProductionWriterInitFails) { + auto service_config = + MasterServiceConfig::builder() + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id("test_batch_record_writer_init_fail") + .set_ha_backend_type("etcd") + .set_ha_backend_connstring("127.0.0.1:1") + .build(); + + EXPECT_THROW( + { MasterService service(service_config); }, std::runtime_error); +} + +TEST_F(MasterServiceHATest, GetReplicaListClassifiesRemovedReplicaStates) { + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(false) + .build(); + MasterService service(service_config); + auto mounted = PrepareSimpleSegment(service, "get_readiness_segment"); + + const std::string complete_key = "get_readiness_complete_key"; + PutObjectOnSegment(service, mounted.client_id, complete_key, + "get_readiness_segment"); + EXPECT_TRUE( + service.GetReplicaList(complete_key, kDefaultTenant).has_value()); + + const std::string removed_key = "get_readiness_removed_key"; + PutObjectOnSegment(service, mounted.client_id, removed_key, + "get_readiness_segment"); + MarkCompletedReplicasRemovedForTesting(service, kDefaultTenant, + removed_key); + auto removed = service.GetReplicaList(removed_key, kDefaultTenant); + ASSERT_FALSE(removed.has_value()); + EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, removed.error()); + + const std::string processing_key = "get_readiness_processing_key"; + ReplicateConfig config; + config.replica_num = 1; + ASSERT_TRUE(service + .PutStart(mounted.client_id, processing_key, kDefaultTenant, + 1024, config) + .has_value()); + auto processing = service.GetReplicaList(processing_key, kDefaultTenant); + ASSERT_FALSE(processing.has_value()); + EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, processing.error()); + + auto missing = + service.GetReplicaList("get_readiness_missing_key", kDefaultTenant); + ASSERT_FALSE(missing.has_value()); + EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, missing.error()); +} + +TEST_F(MasterServiceHATest, BatchGetReplicaListClassifiesRemovedReplicaStates) { + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(false) + .build(); + MasterService service(service_config); + auto mounted = PrepareSimpleSegment(service, "batch_get_readiness_segment"); + + const std::string complete_key = "batch_get_readiness_complete_key"; + PutObjectOnSegment(service, mounted.client_id, complete_key, + "batch_get_readiness_segment"); + + const std::string removed_key = "batch_get_readiness_removed_key"; + PutObjectOnSegment(service, mounted.client_id, removed_key, + "batch_get_readiness_segment"); + MarkCompletedReplicasRemovedForTesting(service, kDefaultTenant, + removed_key); + + const std::string processing_key = "batch_get_readiness_processing_key"; + ReplicateConfig config; + config.replica_num = 1; + ASSERT_TRUE(service + .PutStart(mounted.client_id, processing_key, kDefaultTenant, + 1024, config) + .has_value()); + + auto results = + service.BatchGetReplicaList({complete_key, removed_key, processing_key, + "batch_get_readiness_missing_key"}, + kDefaultTenant); + ASSERT_EQ(4u, results.size()); + EXPECT_TRUE(results[0].has_value()); + ASSERT_FALSE(results[1].has_value()); + EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, results[1].error()); + ASSERT_FALSE(results[2].has_value()); + EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, results[2].error()); + ASSERT_FALSE(results[3].has_value()); + EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, results[3].error()); +} + +TEST_F(MasterServiceHATest, ExistKeyRequiresCompletedReplica) { + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(false) + .build(); + MasterService service(service_config); + auto mounted = PrepareSimpleSegment(service, "exist_readiness_segment"); + + const std::string complete_key = "exist_readiness_complete_key"; + PutObjectOnSegment(service, mounted.client_id, complete_key, + "exist_readiness_segment"); + auto complete = service.ExistKey(complete_key, kDefaultTenant); + ASSERT_TRUE(complete.has_value()); + EXPECT_TRUE(complete.value()); + + const std::string removed_key = "exist_readiness_removed_key"; + PutObjectOnSegment(service, mounted.client_id, removed_key, + "exist_readiness_segment"); + MarkCompletedReplicasRemovedForTesting(service, kDefaultTenant, + removed_key); + auto removed = service.ExistKey(removed_key, kDefaultTenant); + ASSERT_TRUE(removed.has_value()); + EXPECT_FALSE(removed.value()); + + const std::string processing_key = "exist_readiness_processing_key"; + ReplicateConfig config; + config.replica_num = 1; + ASSERT_TRUE(service + .PutStart(mounted.client_id, processing_key, kDefaultTenant, + 1024, config) + .has_value()); + auto processing = service.ExistKey(processing_key, kDefaultTenant); + ASSERT_TRUE(processing.has_value()); + EXPECT_FALSE(processing.value()); + + auto missing = + service.ExistKey("exist_readiness_missing_key", kDefaultTenant); + ASSERT_TRUE(missing.has_value()); + EXPECT_FALSE(missing.value()); +} + +TEST_F(MasterServiceHATest, BatchExistKeyRequiresCompletedReplica) { + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(false) + .build(); + MasterService service(service_config); + auto mounted = + PrepareSimpleSegment(service, "batch_exist_readiness_segment"); + + const std::string complete_key = "batch_exist_readiness_complete_key"; + PutObjectOnSegment(service, mounted.client_id, complete_key, + "batch_exist_readiness_segment"); + + const std::string removed_key = "batch_exist_readiness_removed_key"; + PutObjectOnSegment(service, mounted.client_id, removed_key, + "batch_exist_readiness_segment"); + MarkCompletedReplicasRemovedForTesting(service, kDefaultTenant, + removed_key); + + const std::string processing_key = "batch_exist_readiness_processing_key"; + ReplicateConfig config; + config.replica_num = 1; + ASSERT_TRUE(service + .PutStart(mounted.client_id, processing_key, kDefaultTenant, + 1024, config) + .has_value()); + + auto results = + service.BatchExistKey({complete_key, removed_key, processing_key, + "batch_exist_readiness_missing_key"}, + kDefaultTenant); + ASSERT_EQ(4u, results.size()); + ASSERT_TRUE(results[0].has_value()); + EXPECT_TRUE(results[0].value()); + ASSERT_TRUE(results[1].has_value()); + EXPECT_FALSE(results[1].value()); + ASSERT_TRUE(results[2].has_value()); + EXPECT_FALSE(results[2].value()); + ASSERT_TRUE(results[3].has_value()); + EXPECT_FALSE(results[3].value()); +} + +TEST_F(MasterServiceHATest, BatchRecordSubmissionHelpersUseOrderedWriter) { + const std::string cluster_id = "test_batch_record_helpers_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto visible = AppendVisibleForTesting(service, OpType::PUT_END, "tenant", + "visible_key", "visible_payload"); + ASSERT_TRUE(visible.has_value()); + EXPECT_EQ(1u, visible.value()); + + std::promise finalized_promise; + auto finalized_future = finalized_promise.get_future(); + auto finalized = AppendFinalizeForTesting( + service, OpType::REMOVE, "tenant", "remove_key", {}, + [&finalized_promise](const OpLogEntry& durable_entry) { + finalized_promise.set_value(durable_entry); + }); + ASSERT_TRUE(finalized.has_value()); + EXPECT_EQ(2u, finalized->sequence_id); + + ASSERT_EQ(std::future_status::ready, + finalized_future.wait_for(std::chrono::seconds(5))); + OpLogEntry finalized_entry = finalized_future.get(); + EXPECT_EQ(2u, finalized_entry.sequence_id); + EXPECT_EQ(OpType::REMOVE, finalized_entry.op_type); + EXPECT_EQ("remove_key", finalized_entry.object_key); + + OpLogBatchStorage storage(cluster_id, *backend); + DurablePrefix prefix; + ASSERT_EQ(ErrorCode::OK, storage.ReadDurablePrefix(prefix)); + EXPECT_EQ(2u, prefix.batch_id); + EXPECT_EQ(2u, prefix.last_seq); +} + +TEST_F(MasterServiceHATest, + BatchRecordWriterResolvesEmptyTenantWhenMultiTenantDisabled) { + const std::string cluster_id = "test_batch_record_default_tenant"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(false) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto appended = AppendFinalizeForTesting(service, OpType::REMOVE, "", + "default_tenant_key", {}, nullptr); + ASSERT_TRUE(appended.has_value()); + + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ("default", batch.entries[0].tenant_id); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + PrimaryWritesBatchRecordAndDurablePrefix) { + const std::string cluster_id = "test_batch_record_e2e_primary"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_e2e_primary_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::SEGMENT_MOUNT, batch.entries[0].op_type); + EXPECT_EQ(1u, batch.entries[0].sequence_id); + + const std::string key = "batch_e2e_primary_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_e2e_primary_segment"); + ReadBatchEventually(storage, 2, batch); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(2u, batch.entries[0].sequence_id); + EXPECT_FALSE(batch.entries[0].payload.empty()); + + DurablePrefix prefix; + ASSERT_EQ(ErrorCode::OK, storage.ReadDurablePrefix(prefix)); + EXPECT_EQ(2u, prefix.batch_id); + EXPECT_EQ(2u, prefix.last_seq); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + ASSERT_TRUE(replicas.has_value()) << toString(replicas.error()); + ASSERT_EQ(1u, replicas->replicas.size()); + EXPECT_TRUE(replicas->replicas.front().is_memory_replica()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, StandbyAppliesPrimaryBatchRecords) { + const std::string cluster_id = "test_batch_record_e2e_standby_apply"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_e2e_standby_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_e2e_standby_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_e2e_standby_segment"); + ReadBatchEventually(storage, 2, batch); + + MockMetadataStore standby_metadata; + OpLogApplier applier(&standby_metadata, cluster_id); + OpLogBatchStandbyReader reader(cluster_id, *backend, applier); + + auto result = reader.PollOnce(); + ASSERT_EQ(ErrorCode::OK, result.error); + EXPECT_EQ(2u, result.applied_entries); + EXPECT_EQ(3u, applier.GetExpectedSequenceId()); + EXPECT_TRUE(standby_metadata.Exists(kDefaultTenant.value(), key)); +} + +TEST_F(MasterServiceBatchRecordE2ETest, PromotionCatchesUpToDurablePrefix) { + const std::string cluster_id = "test_batch_record_e2e_promotion_catchup"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_e2e_promotion_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_e2e_promotion_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_e2e_promotion_segment"); + ReadBatchEventually(storage, 2, batch); + + HotStandbyConfig standby_config; + standby_config.enable_verification = false; + standby_config.max_replication_lag_entries = 1000; + standby_config.enable_oplog_following = true; + standby_config.enable_snapshot_bootstrap = false; + standby_config.oplog_poll_interval_ms = 50; + + HotStandbyService standby(standby_config); + standby.SetCatchUpBatchKvBackendForTesting(backend); + auto start_err = standby.Start("", "", cluster_id); + ASSERT_EQ(ErrorCode::OK, start_err); + ASSERT_EQ(StandbyState::WATCHING, standby.GetState()); + + StandbySnapshot snapshot; + ASSERT_EQ(ErrorCode::OK, standby.PromoteAndExportSnapshot(snapshot)); + EXPECT_EQ(2u, snapshot.oplog_sequence_id); + + bool found_key = false; + for (const auto& object : snapshot.objects) { + found_key = found_key || object.key == key; + } + EXPECT_TRUE(found_key); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + BackendFailureStopsNewBatchReservations) { + const std::string cluster_id = "test_batch_record_e2e_backend_failure"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + backend->SetTxnError(ErrorCode::PERSISTENT_FAIL); + auto first = AppendVisibleForTesting(service, OpType::PUT_END, + kDefaultTenant, "failing_key", {}); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(backend->WaitForTxnCalls(1)); + + tl::expected second; + for (int i = 0; i < 100; ++i) { + second = AppendVisibleForTesting(service, OpType::PUT_END, + kDefaultTenant, "rejected_key", {}); + if (!second.has_value()) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + ASSERT_FALSE(second.has_value()); + EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS, second.error()); + backend->SetTxnError(ErrorCode::OK); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + RetryRecoveryRestoresBatchWriterAccepting) { + const std::string cluster_id = "test_batch_record_e2e_retry_recovery"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + + backend->SetTxnError(ErrorCode::PERSISTENT_FAIL); + auto first = AppendVisibleForTesting(service, OpType::PUT_END, + kDefaultTenant, "retry_key", {}); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(backend->WaitForTxnCalls(1)); + + backend->SetTxnError(ErrorCode::OK); + ReadBatchEventually(storage, 1, batch); + + auto recovered = AppendVisibleForTesting( + service, OpType::PUT_END, kDefaultTenant, "recovered_key", {}); + ASSERT_TRUE(recovered.has_value()) << toString(recovered.error()); + ReadBatchEventually(storage, 2, batch); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + RemoveBeforeDurableHidesReplicasFromPrimaryReads) { + const std::string cluster_id = "test_batch_record_e2e_remove_visibility"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_e2e_remove_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_e2e_remove_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_e2e_remove_segment"); + ReadBatchEventually(storage, 2, batch); + + backend->BlockTxn(); + auto removed = service.Remove(key, kDefaultTenant, /*force=*/true); + EXPECT_TRUE(removed.has_value()); + EXPECT_FALSE(service.GetReplicaList(key, kDefaultTenant).has_value()); + DurablePrefix prefix; + auto prefix_read = storage.ReadDurablePrefix(prefix); + EXPECT_EQ(ErrorCode::OK, prefix_read); + if (prefix_read == ErrorCode::OK) { + EXPECT_EQ(2u, prefix.batch_id); + EXPECT_EQ(2u, prefix.last_seq); + } + + backend->AllowTxn(); + ReadBatchEventually(storage, 3, batch); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::REMOVE, batch.entries[0].op_type); + EXPECT_EQ(key, batch.entries[0].object_key); +} + +TEST_F(MasterServiceBatchRecordE2ETest, RemoveByRegexWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_remove_regex"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_regex_remove_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string removed_key = "batch_regex_remove_key"; + const std::string kept_key = "batch_regex_keep_key"; + PutObjectOnSegment(service, mounted.client_id, removed_key, + "batch_regex_remove_segment"); + ReadBatchEventually(storage, 2, batch); + PutObjectOnSegment(service, mounted.client_id, kept_key, + "batch_regex_remove_segment"); + ReadBatchEventually(storage, 3, batch); + + auto removed = service.RemoveByRegex("^batch_regex_remove_", kDefaultTenant, + /*force=*/true); + ASSERT_TRUE(removed.has_value()) << toString(removed.error()); + EXPECT_EQ(1, removed.value()); + + ReadBatchEventually(storage, 4, batch); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::REMOVE, batch.entries[0].op_type); + EXPECT_EQ(removed_key, batch.entries[0].object_key); + auto removed_exists = service.ExistKey(removed_key, kDefaultTenant); + ASSERT_TRUE(removed_exists.has_value()) << toString(removed_exists.error()); + EXPECT_FALSE(removed_exists.value()); + auto kept_exists = service.ExistKey(kept_key, kDefaultTenant); + ASSERT_TRUE(kept_exists.has_value()) << toString(kept_exists.error()); + EXPECT_TRUE(kept_exists.value()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + RemoveFailsBeforeMutationWhenBatchReservationUnavailable) { + const std::string cluster_id = "test_batch_record_remove_reserve_full"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_remove_reserve_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_remove_reserve_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_remove_reserve_seg"); + ReadBatchEventually(storage, 2, batch); + + auto held_reservation = ReserveBatchSlotForTesting(service); + ASSERT_TRUE(held_reservation.has_value()) + << toString(held_reservation.error()); + + auto removed = service.Remove(key, kDefaultTenant, /*force=*/true); + ASSERT_FALSE(removed.has_value()); + EXPECT_EQ(ErrorCode::TASK_PENDING_LIMIT_EXCEEDED, removed.error()); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + ASSERT_TRUE(replicas.has_value()) << toString(replicas.error()); + EXPECT_EQ(1u, replicas->replicas.size()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + RemoveDurableCallbackReleasesResources) { + const std::string cluster_id = "test_batch_record_e2e_remove_finalize"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 1024}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = + PrepareSimpleSegment(service, "batch_e2e_remove_finalize_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_e2e_remove_finalize_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_e2e_remove_finalize_segment"); + ReadBatchEventually(storage, 2, batch); + + backend->BlockTxn(); + ASSERT_TRUE( + service.Remove(key, kDefaultTenant, /*force=*/true).has_value()); + + ReplicateConfig config; + config.replica_num = 1; + auto before_finalize = service.PutStart( + mounted.client_id, "batch_e2e_before_remove_finalize_key", + kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 3, batch); + + auto removed = service.ExistKey(key, kDefaultTenant); + ASSERT_TRUE(removed.has_value()); + EXPECT_FALSE(removed.value()); + + auto after_finalize = service.PutStart( + mounted.client_id, "batch_e2e_after_remove_finalize_key", + kDefaultTenant, 1024, config); + EXPECT_TRUE(after_finalize.has_value()) << toString(after_finalize.error()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + ClearInvalidHandlesReleasesResourcesAfterDurable) { + const std::string cluster_id = "test_batch_record_stale_finalize"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 1024}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_stale_finalize_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + auto spare = + PrepareSimpleSegment(service, "batch_stale_finalize_spare", + kDefaultSegmentBase + kDefaultSegmentSize); + ReadBatchEventually(storage, 2, batch); + + const std::string key = "batch_stale_finalize_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_stale_finalize_seg"); + ReadBatchEventually(storage, 3, batch); + + backend->BlockTxn(); + auto unmounted = + service.UnmountSegment(mounted.segment_id, mounted.client_id); + ASSERT_TRUE(unmounted.has_value()) << toString(unmounted.error()); + + EXPECT_FALSE(service.GetReplicaList(key, kDefaultTenant).has_value()); + + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_stale_finalize_spare"}; + auto before_finalize = + service.PutStart(spare.client_id, "batch_stale_before_finalize", + kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadRemoveBatchEventually(storage, 4, key, batch); + + tl::expected, ErrorCode> after_finalize = + tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); + for (int i = 0; i < 50; ++i) { + after_finalize = + service.PutStart(spare.client_id, "batch_stale_after_finalize", + kDefaultTenant, 1024, config); + if (after_finalize.has_value()) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + EXPECT_TRUE(after_finalize.has_value()) << toString(after_finalize.error()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + UpsertStartStaleCleanupReleasesResourcesAfterDurable) { + const std::string cluster_id = "test_batch_record_upsert_stale_finalize"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 1024}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_upsert_stale_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + auto spare = + PrepareSimpleSegment(service, "batch_upsert_stale_spare", + kDefaultSegmentBase + kDefaultSegmentSize); + ReadBatchEventually(storage, 2, batch); + + const std::string key = "batch_upsert_stale_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_upsert_stale_seg"); + ReadBatchEventually(storage, 3, batch); + + PrepareUnmountSegmentForTesting(service, mounted.segment_id); + backend->BlockTxn(); + + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_upsert_stale_spare"}; + auto upsert_result = + service.UpsertStart(spare.client_id, key, kDefaultTenant, 1024, config); + EXPECT_FALSE(upsert_result.has_value()); + + auto before_finalize = + service.PutStart(spare.client_id, "batch_upsert_stale_before_finalize", + kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadRemoveBatchEventually(storage, 4, key, batch); + + tl::expected, ErrorCode> after_finalize = + tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); + for (int i = 0; i < 50; ++i) { + after_finalize = + service.PutStart(spare.client_id, "batch_upsert_stale_after", + kDefaultTenant, 1024, config); + if (after_finalize.has_value()) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + EXPECT_TRUE(after_finalize.has_value()) << toString(after_finalize.error()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + BatchRemoveStaleCleanupReleasesResourcesAfterDurable) { + const std::string cluster_id = + "test_batch_record_batch_remove_stale_finalize"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 1024}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = + PrepareSimpleSegment(service, "batch_remove_stale_finalize_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + auto spare = + PrepareSimpleSegment(service, "batch_remove_stale_finalize_spare", + kDefaultSegmentBase + kDefaultSegmentSize); + ReadBatchEventually(storage, 2, batch); + + const std::string key = "batch_remove_stale_finalize_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_remove_stale_finalize_seg"); + ReadBatchEventually(storage, 3, batch); + + PrepareUnmountSegmentForTesting(service, mounted.segment_id); + backend->BlockTxn(); + + auto remove_result = service.BatchRemove({key}, kDefaultTenant, + /*force=*/true); + ASSERT_EQ(1u, remove_result.size()); + EXPECT_FALSE(remove_result[0].has_value()); + EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, remove_result[0].error()); + + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_remove_stale_finalize_spare"}; + auto before_finalize = + service.PutStart(spare.client_id, "batch_remove_stale_before_finalize", + kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadRemoveBatchEventually(storage, 4, key, batch); + + tl::expected, ErrorCode> after_finalize = + tl::make_unexpected(ErrorCode::TENANT_QUOTA_EXCEEDED); + for (int i = 0; i < 50; ++i) { + after_finalize = + service.PutStart(spare.client_id, "batch_remove_stale_after", + kDefaultTenant, 1024, config); + if (after_finalize.has_value()) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + EXPECT_TRUE(after_finalize.has_value()) << toString(after_finalize.error()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + PartialEvictLeavesRemainingCompleteReplicasReadable) { + const std::string cluster_id = "test_batch_record_e2e_partial_evict"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_e2e_partial_evict_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_e2e_partial_evict_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_e2e_partial_evict_seg"); + ReadBatchEventually(storage, 2, batch); + + OffloadTaskItem task{ + .tenant_id = kDefaultTenant.value(), .key = key, .size = 1024}; + StorageObjectMetadata metadata; + metadata.data_size = 1024; + metadata.transport_endpoint = "batch_e2e_partial_evict_disk"; + ASSERT_TRUE( + service.NotifyOffloadSuccess(mounted.client_id, {task}, {metadata}) + .has_value()); + ReadBatchEventually(storage, 3, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + service.RunBatchEvictForTesting(/*evict_ratio_target=*/1.0, + /*evict_ratio_lowerbound=*/1.0); + ReadBatchEventually(storage, 4, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(key, batch.entries[0].object_key); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + ASSERT_TRUE(replicas.has_value()) << toString(replicas.error()); + ASSERT_EQ(1u, replicas->replicas.size()); + EXPECT_TRUE(replicas->replicas.front().is_local_disk_replica()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + EvictAllReadableReplicasReturnsObjectNotFound) { + const std::string cluster_id = "test_batch_record_e2e_evict_all"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_e2e_evict_all_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_e2e_evict_all_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_e2e_evict_all_seg"); + ReadBatchEventually(storage, 2, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + service.RunBatchEvictForTesting(/*evict_ratio_target=*/1.0, + /*evict_ratio_lowerbound=*/1.0); + ReadBatchEventually(storage, 3, batch); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + ASSERT_FALSE(replicas.has_value()); + EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, replicas.error()); + + auto exists = service.ExistKey(key, kDefaultTenant); + ASSERT_TRUE(exists.has_value()); + EXPECT_FALSE(exists.value()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + ProcessingOnlyObjectExistKeyReturnsFalse) { + const std::string cluster_id = "test_batch_record_e2e_processing_only"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_e2e_processing_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_e2e_processing_seg"}; + const std::string key = "batch_e2e_processing_key"; + ASSERT_TRUE( + service.PutStart(mounted.client_id, key, kDefaultTenant, 1024, config) + .has_value()); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + ASSERT_FALSE(replicas.has_value()); + EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, replicas.error()); + + auto exists = service.ExistKey(key, kDefaultTenant); + ASSERT_TRUE(exists.has_value()); + EXPECT_FALSE(exists.value()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + OffloadAndPromotionSuccessRemainFunctional) { + const std::string cluster_id = "test_batch_record_e2e_offload_promotion"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = + PrepareSimpleSegment(service, "batch_e2e_offload_promotion_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string offload_key = "batch_e2e_offload_key"; + PutObjectOnSegment(service, mounted.client_id, offload_key, + "batch_e2e_offload_promotion_seg"); + ReadBatchEventually(storage, 2, batch); + + OffloadTaskItem task{ + .tenant_id = kDefaultTenant.value(), .key = offload_key, .size = 1024}; + StorageObjectMetadata metadata; + metadata.data_size = 1024; + metadata.transport_endpoint = "batch_e2e_offload_endpoint"; + ASSERT_TRUE( + service.NotifyOffloadSuccess(mounted.client_id, {task}, {metadata}) + .has_value()); + ReadBatchEventually(storage, 3, batch); + + auto offload_replicas = service.GetReplicaList(offload_key, kDefaultTenant); + ASSERT_TRUE(offload_replicas.has_value()) + << toString(offload_replicas.error()); + bool has_local_disk = false; + for (const auto& replica : offload_replicas->replicas) { + has_local_disk = has_local_disk || replica.is_local_disk_replica(); + } + EXPECT_TRUE(has_local_disk); + + const std::string promotion_key = "batch_e2e_promotion_success_key"; + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_e2e_offload_promotion_seg"}; + auto put_start = service.PutStart(mounted.client_id, promotion_key, + kDefaultTenant, 1024, config); + ASSERT_TRUE(put_start.has_value()); + ASSERT_EQ(1u, put_start->size()); + SeedPromotionTaskForTesting(&service, kDefaultTenant, promotion_key, + mounted.client_id, put_start->front().id, 1024); + + ASSERT_TRUE(service + .NotifyPromotionSuccess(mounted.client_id, promotion_key, + kDefaultTenant) + .has_value()); + ReadBatchEventually(storage, 4, batch); + + auto promotion_replicas = + service.GetReplicaList(promotion_key, kDefaultTenant); + ASSERT_TRUE(promotion_replicas.has_value()) + << toString(promotion_replicas.error()); + ASSERT_EQ(1u, promotion_replicas->replicas.size()); + EXPECT_TRUE(promotion_replicas->replicas.front().is_memory_replica()); +} + +TEST_F(MasterServiceBatchRecordE2ETest, + SegmentLifecycleEntriesRemainFunctional) { + const std::string cluster_id = "test_batch_record_e2e_segment_lifecycle"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + const UUID mount_client = generate_uuid(); + Segment mounted = MakeSegment("batch_e2e_lifecycle_mount"); + ASSERT_TRUE(service.MountSegment(mounted, mount_client).has_value()); + ReadBatchEventually(storage, 1, batch); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::SEGMENT_MOUNT, batch.entries[0].op_type); + + const UUID remount_client = generate_uuid(); + Segment remounted = MakeSegment("batch_e2e_lifecycle_remount", + kDefaultSegmentBase + kDefaultSegmentSize); + ASSERT_TRUE( + service.ReMountSegment({remounted}, remount_client).has_value()); + ReadBatchEventually(storage, 2, batch); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::SEGMENT_MOUNT, batch.entries[0].op_type); + + const std::string key = "batch_e2e_lifecycle_key"; + PutObjectOnSegment(service, remount_client, key, + "batch_e2e_lifecycle_remount"); + ReadBatchEventually(storage, 3, batch); + ASSERT_TRUE(service.GetReplicaList(key, kDefaultTenant).has_value()); + + ASSERT_TRUE(service.UnmountSegment(mounted.id, mount_client).has_value()); + bool saw_unmount = false; + for (uint64_t batch_id = 4; batch_id <= 6 && !saw_unmount; ++batch_id) { + ReadBatchEventually(storage, batch_id, batch); + ASSERT_EQ(1u, batch.entries.size()); + saw_unmount = batch.entries[0].op_type == OpType::SEGMENT_UNMOUNT; + } + EXPECT_TRUE(saw_unmount); +} + +TEST_F(MasterServiceHATest, PutEndWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_put_end_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_put_end_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ErrorCode read_err = ErrorCode::ETCD_KEY_NOT_EXIST; + for (int i = 0; i < 50; ++i) { + read_err = storage.ReadBatch(1, batch); + if (read_err == ErrorCode::OK) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_EQ(ErrorCode::OK, read_err); + + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_put_end_segment"}; + const std::string key = "batch_put_end_key"; + auto put_start = + service.PutStart(mounted.client_id, key, kDefaultTenant, 1024, config); + ASSERT_TRUE(put_start.has_value()); + ASSERT_TRUE( + service + .PutEnd(mounted.client_id, key, kDefaultTenant, ReplicaType::MEMORY) + .has_value()); + + for (int i = 0; i < 50; ++i) { + read_err = storage.ReadBatch(2, batch); + if (read_err == ErrorCode::OK) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_EQ(ErrorCode::OK, read_err); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(2u, batch.entries[0].sequence_id); + + DurablePrefix prefix; + ASSERT_EQ(ErrorCode::OK, storage.ReadDurablePrefix(prefix)); + EXPECT_EQ(2u, prefix.batch_id); + EXPECT_EQ(2u, prefix.last_seq); +} + +TEST_F(MasterServiceHATest, PutEndVisibleBeforeBatchRecordDurable) { + const std::string cluster_id = "test_batch_record_put_end_visible"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_put_visible_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_put_visible_segment"}; + const std::string key = "batch_put_visible_key"; + auto put_start = + service.PutStart(mounted.client_id, key, kDefaultTenant, 1024, config); + ASSERT_TRUE(put_start.has_value()); + + backend->BlockTxn(); + auto put_end = service.PutEnd(mounted.client_id, key, kDefaultTenant, + ReplicaType::MEMORY); + EXPECT_TRUE(put_end.has_value()); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + EXPECT_TRUE(replicas.has_value()); + if (replicas.has_value()) { + EXPECT_EQ(1u, replicas->replicas.size()); + if (!replicas->replicas.empty()) { + EXPECT_TRUE(replicas->replicas.front().is_memory_replica()); + } + } + DurablePrefix prefix; + auto prefix_read = storage.ReadDurablePrefix(prefix); + EXPECT_EQ(ErrorCode::OK, prefix_read); + if (prefix_read == ErrorCode::OK) { + EXPECT_EQ(1u, prefix.batch_id); + EXPECT_EQ(1u, prefix.last_seq); + } + + backend->AllowTxn(); + ReadBatchEventually(storage, 2, batch); +} + +TEST_F(MasterServiceHATest, CopyEndVisibleBeforeBatchRecordDurable) { + const std::string cluster_id = "test_batch_record_copy_end_visible"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto source = PrepareSimpleSegment(service, "batch_copy_visible_src"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_copy_visible_key"; + PutObjectOnSegment(service, source.client_id, key, + "batch_copy_visible_src"); + ReadBatchEventually(storage, 2, batch); + + PrepareSimpleSegment(service, "batch_copy_visible_dst", + kDefaultSegmentBase + kDefaultSegmentSize); + ReadBatchEventually(storage, 3, batch); + + ASSERT_TRUE(service + .CopyStart(source.client_id, key, kDefaultTenant, + "batch_copy_visible_src", + {"batch_copy_visible_dst"}) + .has_value()); + + backend->BlockTxn(); + auto copy_future = std::async(std::launch::async, [&] { + return service.CopyEnd(source.client_id, key, kDefaultTenant); + }); + const auto status = copy_future.wait_for(std::chrono::milliseconds(200)); + EXPECT_EQ(std::future_status::ready, status) + << "CopyEnd must queue batch OpLog and return before durable"; + if (status != std::future_status::ready) { + backend->AllowTxn(); + } + auto copy_end = copy_future.get(); + ASSERT_TRUE(copy_end.has_value()) << toString(copy_end.error()); + + DurablePrefix prefix; + ASSERT_EQ(ErrorCode::OK, storage.ReadDurablePrefix(prefix)); + EXPECT_EQ(3u, prefix.batch_id); + EXPECT_EQ(3u, prefix.last_seq); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + ASSERT_TRUE(replicas.has_value()); + EXPECT_EQ(2u, replicas->replicas.size()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 4, batch); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(key, batch.entries[0].object_key); +} + +TEST_F(MasterServiceHATest, + MoveEndHidesSourceBeforeDurableAndReleasesAfterDurable) { + const std::string cluster_id = "test_batch_record_move_end_visible"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_eviction_high_watermark_ratio(1.0) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto source = PrepareSimpleSegment(service, "batch_move_visible_src", + kDefaultSegmentBase, 1024); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_move_visible_key"; + PutObjectOnSegment(service, source.client_id, key, "batch_move_visible_src", + 1024); + ReadBatchEventually(storage, 2, batch); + + PrepareSimpleSegment(service, "batch_move_visible_dst", + kDefaultSegmentBase + kDefaultSegmentSize, 1024); + ReadBatchEventually(storage, 3, batch); + + ASSERT_TRUE(service + .MoveStart(source.client_id, key, kDefaultTenant, + "batch_move_visible_src", + "batch_move_visible_dst") + .has_value()); + + backend->BlockTxn(); + auto move_future = std::async(std::launch::async, [&] { + return service.MoveEnd(source.client_id, key, kDefaultTenant); + }); + const auto status = move_future.wait_for(std::chrono::milliseconds(200)); + EXPECT_EQ(std::future_status::ready, status) + << "MoveEnd must queue batch OpLog and return before durable"; + if (status != std::future_status::ready) { + backend->AllowTxn(); + } + auto move_end = move_future.get(); + ASSERT_TRUE(move_end.has_value()) << toString(move_end.error()); + + DurablePrefix prefix; + ASSERT_EQ(ErrorCode::OK, storage.ReadDurablePrefix(prefix)); + EXPECT_EQ(3u, prefix.batch_id); + EXPECT_EQ(3u, prefix.last_seq); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + ASSERT_TRUE(replicas.has_value()); + EXPECT_EQ(1u, replicas->replicas.size()); + + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_move_visible_src"}; + auto before_finalize = + service.PutStart(source.client_id, "batch_move_before_finalize", + kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 4, batch); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(key, batch.entries[0].object_key); + + tl::expected, ErrorCode> after_finalize = + tl::make_unexpected(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS); + for (int i = 0; i < 50; ++i) { + after_finalize = + service.PutStart(source.client_id, "batch_move_after_finalize", + kDefaultTenant, 1024, config); + if (after_finalize.has_value()) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + EXPECT_TRUE(after_finalize.has_value()) << toString(after_finalize.error()); +} + +TEST_F(MasterServiceHATest, + NotifyOffloadSuccessFallbackWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_offload_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_offload_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ErrorCode read_err = ErrorCode::ETCD_KEY_NOT_EXIST; + for (int i = 0; i < 50; ++i) { + read_err = storage.ReadBatch(1, batch); + if (read_err == ErrorCode::OK) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_EQ(ErrorCode::OK, read_err); + + const std::string key = "batch_offload_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_offload_segment"); + for (int i = 0; i < 50; ++i) { + read_err = storage.ReadBatch(2, batch); + if (read_err == ErrorCode::OK) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_EQ(ErrorCode::OK, read_err); + + OffloadTaskItem task{ + .tenant_id = kDefaultTenant.value(), .key = key, .size = 1024}; + StorageObjectMetadata metadata; + metadata.data_size = 1024; + metadata.transport_endpoint = "local_disk_endpoint"; + ASSERT_TRUE( + service.NotifyOffloadSuccess(mounted.client_id, {task}, {metadata}) + .has_value()); + auto replicas = service.GetReplicaList(key, kDefaultTenant); + ASSERT_TRUE(replicas.has_value()); + bool has_local_disk = false; + for (const auto& replica : replicas->replicas) { + has_local_disk = has_local_disk || replica.is_local_disk_replica(); + } + EXPECT_TRUE(has_local_disk); + + for (int i = 0; i < 50; ++i) { + read_err = storage.ReadBatch(3, batch); + if (read_err == ErrorCode::OK) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_EQ(ErrorCode::OK, read_err); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(3u, batch.entries[0].sequence_id); + EXPECT_FALSE(batch.entries[0].payload.empty()); +} + +TEST_F(MasterServiceHATest, + NotifyOffloadSuccessFallbackVisibleBeforeBatchRecordDurable) { + const std::string cluster_id = "test_batch_record_offload_visible"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_offload_visible_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_offload_visible_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_offload_visible_seg"); + ReadBatchEventually(storage, 2, batch); + + OffloadTaskItem task{ + .tenant_id = kDefaultTenant.value(), .key = key, .size = 1024}; + StorageObjectMetadata metadata; + metadata.data_size = 1024; + metadata.transport_endpoint = "local_disk_visible_endpoint"; + + backend->BlockTxn(); + auto offload = + service.NotifyOffloadSuccess(mounted.client_id, {task}, {metadata}); + EXPECT_TRUE(offload.has_value()); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + EXPECT_TRUE(replicas.has_value()); + if (replicas.has_value()) { + bool has_local_disk = false; + for (const auto& replica : replicas->replicas) { + has_local_disk = has_local_disk || replica.is_local_disk_replica(); + } + EXPECT_TRUE(has_local_disk); + } + DurablePrefix prefix; + auto prefix_read = storage.ReadDurablePrefix(prefix); + EXPECT_EQ(ErrorCode::OK, prefix_read); + if (prefix_read == ErrorCode::OK) { + EXPECT_EQ(2u, prefix.batch_id); + EXPECT_EQ(2u, prefix.last_seq); + } + + backend->AllowTxn(); + ReadBatchEventually(storage, 3, batch); +} + +TEST_F(MasterServiceHATest, SegmentLifecycleWritesBatchRecordOpLogs) { + const std::string cluster_id = "test_batch_record_segment_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + const UUID client_id = generate_uuid(); + Segment mounted = MakeSegment("batch_segment_mount"); + ASSERT_TRUE(service.MountSegment(mounted, client_id).has_value()); + + const UUID remount_client_id = generate_uuid(); + Segment remounted = MakeSegment("batch_segment_remount", + kDefaultSegmentBase + kDefaultSegmentSize); + ASSERT_TRUE( + service.ReMountSegment({remounted}, remount_client_id).has_value()); + + ASSERT_TRUE(service.UnmountSegment(mounted.id, client_id).has_value()); + + OpLogBatchStorage storage(cluster_id, *backend); + auto read_batch = [&](uint64_t batch_id) { + OpLogBatchRecord batch; + ErrorCode read_err = ErrorCode::ETCD_KEY_NOT_EXIST; + for (int i = 0; i < 50; ++i) { + read_err = storage.ReadBatch(batch_id, batch); + if (read_err == ErrorCode::OK) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + EXPECT_EQ(ErrorCode::OK, read_err); + return batch; + }; + + auto mount_batch = read_batch(1); + ASSERT_EQ(1u, mount_batch.entries.size()); + EXPECT_EQ(OpType::SEGMENT_MOUNT, mount_batch.entries[0].op_type); + EXPECT_EQ(1u, mount_batch.entries[0].sequence_id); + EXPECT_FALSE(mount_batch.entries[0].payload.empty()); + + auto remount_batch = read_batch(2); + ASSERT_EQ(1u, remount_batch.entries.size()); + EXPECT_EQ(OpType::SEGMENT_MOUNT, remount_batch.entries[0].op_type); + EXPECT_EQ(2u, remount_batch.entries[0].sequence_id); + EXPECT_FALSE(remount_batch.entries[0].payload.empty()); + + auto unmount_batch = read_batch(3); + ASSERT_EQ(1u, unmount_batch.entries.size()); + EXPECT_EQ(OpType::SEGMENT_UNMOUNT, unmount_batch.entries[0].op_type); + EXPECT_EQ(3u, unmount_batch.entries[0].sequence_id); + EXPECT_FALSE(unmount_batch.entries[0].payload.empty()); +} + +TEST_F(MasterServiceHATest, NotifyPromotionSuccessWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_promotion_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + const auto mounted = PrepareSimpleSegment(service, "batch_promotion_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ErrorCode read_err = ErrorCode::ETCD_KEY_NOT_EXIST; + for (int i = 0; i < 50; ++i) { + read_err = storage.ReadBatch(1, batch); + if (read_err == ErrorCode::OK) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_EQ(ErrorCode::OK, read_err); + + const std::string key = "batch_promotion_key"; + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_promotion_seg"}; + auto put_start = + service.PutStart(mounted.client_id, key, kDefaultTenant, 1024, config); + ASSERT_TRUE(put_start.has_value()); + ASSERT_EQ(1u, put_start->size()); + SeedPromotionTaskForTesting(&service, kDefaultTenant, key, + mounted.client_id, put_start->front().id, 1024); + + auto res = + service.NotifyPromotionSuccess(mounted.client_id, key, kDefaultTenant); + ASSERT_TRUE(res.has_value()); + + for (int i = 0; i < 50; ++i) { + read_err = storage.ReadBatch(2, batch); + if (read_err == ErrorCode::OK) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_EQ(ErrorCode::OK, read_err); + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(2u, batch.entries[0].sequence_id); + EXPECT_FALSE(batch.entries[0].payload.empty()); +} + +TEST_F(MasterServiceHATest, + NotifyPromotionSuccessVisibleBeforeBatchRecordDurable) { + const std::string cluster_id = "test_batch_record_promotion_visible"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + const auto mounted = + PrepareSimpleSegment(service, "batch_promotion_visible_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_promotion_visible_key"; + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segments = {"batch_promotion_visible_seg"}; + auto put_start = + service.PutStart(mounted.client_id, key, kDefaultTenant, 1024, config); + ASSERT_TRUE(put_start.has_value()); + ASSERT_EQ(1u, put_start->size()); + SeedPromotionTaskForTesting(&service, kDefaultTenant, key, + mounted.client_id, put_start->front().id, 1024); + + backend->BlockTxn(); + auto promotion = + service.NotifyPromotionSuccess(mounted.client_id, key, kDefaultTenant); + EXPECT_TRUE(promotion.has_value()); + + auto replicas = service.GetReplicaList(key, kDefaultTenant); + EXPECT_TRUE(replicas.has_value()); + if (replicas.has_value()) { + EXPECT_EQ(1u, replicas->replicas.size()); + if (!replicas->replicas.empty()) { + EXPECT_TRUE(replicas->replicas.front().is_memory_replica()); + } + } + DurablePrefix prefix; + auto prefix_read = storage.ReadDurablePrefix(prefix); + EXPECT_EQ(ErrorCode::OK, prefix_read); + if (prefix_read == ErrorCode::OK) { + EXPECT_EQ(1u, prefix.batch_id); + EXPECT_EQ(1u, prefix.last_seq); + } + + backend->AllowTxn(); + ReadBatchEventually(storage, 2, batch); +} + +TEST_F(MasterServiceHATest, RemoveWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_remove_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_remove_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_remove_key"; + PutObjectOnSegment(service, mounted.client_id, key, "batch_remove_segment"); + ReadBatchEventually(storage, 2, batch); + + ASSERT_TRUE( + service.Remove(key, kDefaultTenant, /*force=*/true).has_value()); + ReadBatchEventually(storage, 3, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::REMOVE, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(3u, batch.entries[0].sequence_id); +} + +TEST_F(MasterServiceHATest, RemoveHidesBeforeDurableAndReleasesAfterFinalize) { + const std::string cluster_id = "test_batch_record_remove_finalize_cluster"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 1024}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_remove_finalize_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_remove_finalize_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_remove_finalize_seg"); + ReadBatchEventually(storage, 2, batch); + + backend->BlockTxn(); + ASSERT_TRUE( + service.Remove(key, kDefaultTenant, /*force=*/true).has_value()); + EXPECT_FALSE(service.GetReplicaList(key, kDefaultTenant).has_value()); + + ReplicateConfig config; + config.replica_num = 1; + const std::string before_finalize_key = "before_remove_finalize_key"; + auto before_finalize = service.PutStart( + mounted.client_id, before_finalize_key, kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 3, batch); + + if (!before_finalize.has_value()) { + const std::string after_finalize_key = "after_remove_finalize_key"; + auto after_finalize = + service.PutStart(mounted.client_id, after_finalize_key, + kDefaultTenant, 1024, config); + EXPECT_TRUE(after_finalize.has_value()) + << toString(after_finalize.error()); + } +} + +TEST_F(MasterServiceHATest, BatchRemoveWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_batch_remove_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_remove_many_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_remove_many_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_remove_many_segment"); + ReadBatchEventually(storage, 2, batch); + + auto results = service.BatchRemove({key}, kDefaultTenant, /*force=*/true); + ASSERT_EQ(1u, results.size()); + ASSERT_TRUE(results[0].has_value()); + ReadBatchEventually(storage, 3, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::REMOVE, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(3u, batch.entries[0].sequence_id); +} + +TEST_F(MasterServiceHATest, BatchRemoveFinalizesEachObjectAfterDurable) { + const std::string cluster_id = "test_batch_record_batch_remove_finalize"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 1024}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = + PrepareSimpleSegment(service, "batch_remove_finalize_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_remove_finalize_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_remove_finalize_segment"); + ReadBatchEventually(storage, 2, batch); + + backend->BlockTxn(); + auto results = service.BatchRemove({key}, kDefaultTenant, /*force=*/true); + ASSERT_EQ(1u, results.size()); + ASSERT_TRUE(results[0].has_value()); + EXPECT_FALSE(service.GetReplicaList(key, kDefaultTenant).has_value()); + + ReplicateConfig config; + config.replica_num = 1; + const std::string before_finalize_key = "before_batch_remove_finalize_key"; + auto before_finalize = service.PutStart( + mounted.client_id, before_finalize_key, kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 3, batch); + + if (!before_finalize.has_value()) { + const std::string after_finalize_key = + "after_batch_remove_finalize_key"; + auto after_finalize = + service.PutStart(mounted.client_id, after_finalize_key, + kDefaultTenant, 1024, config); + EXPECT_TRUE(after_finalize.has_value()) + << toString(after_finalize.error()); + } +} + +TEST_F(MasterServiceHATest, RemoveAllWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_remove_all_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_remove_all_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_remove_all_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_remove_all_segment"); + ReadBatchEventually(storage, 2, batch); + + EXPECT_EQ(1, service.RemoveAll(kDefaultTenant, /*force=*/true)); + ReadBatchEventually(storage, 3, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::REMOVE, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(3u, batch.entries[0].sequence_id); +} + +TEST_F(MasterServiceHATest, RemoveAllFinalizesAfterDurable) { + const std::string cluster_id = "test_batch_record_remove_all_finalize"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 1024}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "remove_all_finalize_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "remove_all_finalize_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "remove_all_finalize_segment"); + ReadBatchEventually(storage, 2, batch); + + backend->BlockTxn(); + EXPECT_EQ(1, service.RemoveAll(kDefaultTenant, /*force=*/true)); + EXPECT_FALSE(service.GetReplicaList(key, kDefaultTenant).has_value()); + + ReplicateConfig config; + config.replica_num = 1; + const std::string before_finalize_key = "before_remove_all_finalize_key"; + auto before_finalize = service.PutStart( + mounted.client_id, before_finalize_key, kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 3, batch); + + if (!before_finalize.has_value()) { + const std::string after_finalize_key = "after_remove_all_finalize_key"; + auto after_finalize = + service.PutStart(mounted.client_id, after_finalize_key, + kDefaultTenant, 1024, config); + EXPECT_TRUE(after_finalize.has_value()) + << toString(after_finalize.error()); + } +} + +TEST_F(MasterServiceHATest, BatchReplicaClearAllWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_clear_all_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_clear_all_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_clear_all_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_clear_all_segment"); + ReadBatchEventually(storage, 2, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + auto res = service.BatchReplicaClear({key}, mounted.client_id, ""); + ASSERT_TRUE(res.has_value()); + ASSERT_EQ(1u, res->size()); + EXPECT_EQ(key, (*res)[0]); + ReadBatchEventually(storage, 3, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::REMOVE, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(3u, batch.entries[0].sequence_id); +} + +TEST_F(MasterServiceHATest, BatchReplicaClearSegmentWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_clear_segment_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_clear_seg1"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_clear_segment_key"; + PutObjectOnSegment(service, mounted.client_id, key, "batch_clear_seg1"); + ReadBatchEventually(storage, 2, batch); + + PrepareSimpleSegment(service, "batch_clear_seg2", + kDefaultSegmentBase + kDefaultSegmentSize); + ReadBatchEventually(storage, 3, batch); + + auto copy_start = + service.CopyStart(mounted.client_id, key, kDefaultTenant, + "batch_clear_seg1", {"batch_clear_seg2"}); + ASSERT_TRUE(copy_start.has_value()); + auto copy_end = service.CopyEnd(mounted.client_id, key, kDefaultTenant); + ASSERT_TRUE(copy_end.has_value()); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + auto res = + service.BatchReplicaClear({key}, mounted.client_id, "batch_clear_seg1"); + ASSERT_TRUE(res.has_value()); + ASSERT_EQ(1u, res->size()); + EXPECT_EQ(key, (*res)[0]); + ReadBatchEventually(storage, 4, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(4u, batch.entries[0].sequence_id); + EXPECT_FALSE(batch.entries[0].payload.empty()); +} + +TEST_F(MasterServiceHATest, BatchReplicaClearAllReleasesAfterDurable) { + const std::string cluster_id = "test_batch_record_clear_all_finalize"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 1024}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "clear_all_finalize_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "clear_all_finalize_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "clear_all_finalize_segment"); + ReadBatchEventually(storage, 2, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + backend->BlockTxn(); + auto res = service.BatchReplicaClear({key}, mounted.client_id, ""); + ASSERT_TRUE(res.has_value()); + ASSERT_EQ(1u, res->size()); + EXPECT_EQ(key, (*res)[0]); + EXPECT_FALSE(service.GetReplicaList(key, kDefaultTenant).has_value()); + + ReplicateConfig config; + config.replica_num = 1; + const std::string before_finalize_key = "before_clear_all_finalize_key"; + auto before_finalize = service.PutStart( + mounted.client_id, before_finalize_key, kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 3, batch); + + if (!before_finalize.has_value()) { + const std::string after_finalize_key = "after_clear_all_finalize_key"; + auto after_finalize = + service.PutStart(mounted.client_id, after_finalize_key, + kDefaultTenant, 1024, config); + EXPECT_TRUE(after_finalize.has_value()) + << toString(after_finalize.error()); + } +} + +TEST_F(MasterServiceHATest, BatchReplicaClearSegmentReleasesAfterDurable) { + const std::string cluster_id = "test_batch_record_clear_segment_finalize"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 2048}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "clear_finalize_seg1"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "clear_segment_finalize_key"; + ReplicateConfig pinned_config; + pinned_config.replica_num = 1; + pinned_config.preferred_segments = {"clear_finalize_seg1"}; + pinned_config.with_hard_pin = true; + ASSERT_TRUE(service + .PutStart(mounted.client_id, key, kDefaultTenant, 1024, + pinned_config) + .has_value()); + ASSERT_TRUE( + service + .PutEnd(mounted.client_id, key, kDefaultTenant, ReplicaType::MEMORY) + .has_value()); + ReadBatchEventually(storage, 2, batch); + + PrepareSimpleSegment(service, "clear_finalize_seg2", + kDefaultSegmentBase + kDefaultSegmentSize); + ReadBatchEventually(storage, 3, batch); + + auto copy_start = + service.CopyStart(mounted.client_id, key, kDefaultTenant, + "clear_finalize_seg1", {"clear_finalize_seg2"}); + ASSERT_TRUE(copy_start.has_value()); + ASSERT_TRUE( + service.CopyEnd(mounted.client_id, key, kDefaultTenant).has_value()); + ReadBatchEventually(storage, 4, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + backend->BlockTxn(); + auto res = service.BatchReplicaClear({key}, mounted.client_id, + "clear_finalize_seg1"); + ASSERT_TRUE(res.has_value()); + ASSERT_EQ(1u, res->size()); + EXPECT_EQ(key, (*res)[0]); + + ReplicateConfig config; + config.replica_num = 1; + const std::string before_finalize_key = "before_clear_segment_finalize_key"; + auto before_finalize = service.PutStart( + mounted.client_id, before_finalize_key, kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 5, batch); + + if (!before_finalize.has_value()) { + const std::string after_finalize_key = + "after_clear_segment_finalize_key"; + auto after_finalize = + service.PutStart(mounted.client_id, after_finalize_key, + kDefaultTenant, 1024, config); + EXPECT_TRUE(after_finalize.has_value()) + << toString(after_finalize.error()); + } +} + +TEST_F(MasterServiceHATest, BatchEvictWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_batch_evict_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_evict_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_evict_key"; + PutObjectOnSegment(service, mounted.client_id, key, "batch_evict_segment"); + ReadBatchEventually(storage, 2, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + service.RunBatchEvictForTesting(/*evict_ratio_target=*/1.0, + /*evict_ratio_lowerbound=*/1.0); + ReadBatchEventually(storage, 3, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::REMOVE, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(3u, batch.entries[0].sequence_id); +} + +TEST_F(MasterServiceHATest, BatchEvictReleasesMemoryAfterDurable) { + const std::string cluster_id = "test_batch_record_batch_evict_finalize"; + auto backend = std::make_shared(); + auto service_config = + MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile({{kDefaultTenant.value(), 1024}})) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_evict_finalize_seg"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_evict_finalize_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_evict_finalize_seg"); + ReadBatchEventually(storage, 2, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + backend->BlockTxn(); + service.RunBatchEvictForTesting(/*evict_ratio_target=*/1.0, + /*evict_ratio_lowerbound=*/1.0); + EXPECT_FALSE(service.GetReplicaList(key, kDefaultTenant).has_value()); + + ReplicateConfig config; + config.replica_num = 1; + const std::string before_finalize_key = "before_batch_evict_finalize_key"; + auto before_finalize = service.PutStart( + mounted.client_id, before_finalize_key, kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 3, batch); + + if (!before_finalize.has_value()) { + const std::string after_finalize_key = "after_batch_evict_finalize_key"; + auto after_finalize = + service.PutStart(mounted.client_id, after_finalize_key, + kDefaultTenant, 1024, config); + EXPECT_TRUE(after_finalize.has_value()) + << toString(after_finalize.error()); + } +} + +TEST_F(MasterServiceHATest, EvictDiskReplicaWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_disk_evict_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto mounted = PrepareSimpleSegment(service, "batch_disk_evict_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_disk_evict_key"; + PutObjectOnSegment(service, mounted.client_id, key, + "batch_disk_evict_segment"); + ReadBatchEventually(storage, 2, batch); + + Replica local_disk_replica(mounted.client_id, 1024, "local_disk_endpoint", + ReplicaStatus::COMPLETE); + ASSERT_TRUE(service + .AddReplica(mounted.client_id, key, kDefaultTenant, + local_disk_replica) + .has_value()); + ReadBatchEventually(storage, 3, batch); + + ASSERT_TRUE(service + .EvictDiskReplica(mounted.client_id, key, kDefaultTenant, + ReplicaType::LOCAL_DISK) + .has_value()); + ReadBatchEventually(storage, 4, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(4u, batch.entries[0].sequence_id); + EXPECT_FALSE(batch.entries[0].payload.empty()); +} + +TEST_F(MasterServiceHATest, EvictDiskReplicaReleasesLocalDiskAfterDurable) { + const std::string cluster_id = + "test_batch_record_disk_evict_finalize_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_enable_offload(true) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + const std::string segment_name = "batch_disk_evict_finalize_segment"; + auto mounted = PrepareSimpleSegment(service, segment_name); + ASSERT_TRUE(service + .MountLocalDiskSegment(mounted.client_id, + /*enable_offloading=*/false) + .has_value()); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_disk_evict_finalize_key"; + PutObjectOnSegment(service, mounted.client_id, key, segment_name); + ReadBatchEventually(storage, 2, batch); + + Replica local_disk_replica(mounted.client_id, 1024, "local_disk_endpoint", + ReplicaStatus::COMPLETE); + ASSERT_TRUE(service + .AddReplica(mounted.client_id, key, kDefaultTenant, + local_disk_replica) + .has_value()); + ReadBatchEventually(storage, 3, batch); + SetLocalDiskUsedBytesForTesting(service, mounted.client_id, 1024); + + backend->BlockTxn(); + ASSERT_TRUE(service + .EvictDiskReplica(mounted.client_id, key, kDefaultTenant, + ReplicaType::LOCAL_DISK) + .has_value()); + EXPECT_EQ(1024, GetLocalDiskUsedBytesForTesting(service, segment_name)); + + auto before_finalize = service.GetReplicaList(key, kDefaultTenant); + ASSERT_TRUE(before_finalize.has_value()); + EXPECT_FALSE(std::any_of(before_finalize->replicas.begin(), + before_finalize->replicas.end(), + [](const Replica::Descriptor& desc) { + return desc.is_local_disk_replica(); + })); + + backend->AllowTxn(); + ReadBatchEventually(storage, 4, batch); + EXPECT_EQ(0, GetLocalDiskUsedBytesForTesting(service, segment_name)); +} + +#ifdef USE_NOF +TEST_F(MasterServiceHATest, NoFBatchEvictWritesBatchRecordOpLog) { + const std::string cluster_id = "test_batch_record_nof_evict_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + NoFSegment nof_segment = + MakeNoFSegment("batch_nof_evict_segment", "batch_nof_evict_endpoint"); + const UUID client_id = generate_uuid(); + ASSERT_TRUE(service.MountNoFSegment(nof_segment, client_id).has_value()); + + const std::string key = "batch_nof_evict_key"; + ReplicateConfig config; + config.nof_replica_num = 1; + auto put_start = + service.PutStart(client_id, key, kDefaultTenant, 1024, config); + ASSERT_TRUE(put_start.has_value()); + ASSERT_TRUE( + service.PutEnd(client_id, key, kDefaultTenant, ReplicaType::NOF_SSD) + .has_value()); + + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + service.RunNoFBatchEvictForTesting(/*evict_ratio_target=*/1.0, + /*evict_ratio_lowerbound=*/1.0); + ReadBatchEventually(storage, 2, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::REMOVE, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(2u, batch.entries[0].sequence_id); +} + +TEST_F(MasterServiceHATest, NoFBatchEvictReleasesNoFSpaceAfterDurable) { + const std::string cluster_id = + "test_batch_record_nof_evict_finalize_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + NoFSegment nof_segment = MakeNoFSegment("batch_nof_evict_finalize_segment", + "batch_nof_evict_finalize_endpoint", + kDefaultSegmentBase, 1024); + const UUID client_id = generate_uuid(); + ASSERT_TRUE(service.MountNoFSegment(nof_segment, client_id).has_value()); + + ReplicateConfig config; + config.nof_replica_num = 1; + const std::string key = "batch_nof_evict_finalize_key"; + ASSERT_TRUE(service.PutStart(client_id, key, kDefaultTenant, 1024, config) + .has_value()); + ASSERT_TRUE( + service.PutEnd(client_id, key, kDefaultTenant, ReplicaType::NOF_SSD) + .has_value()); + + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(60)); + backend->BlockTxn(); + service.RunNoFBatchEvictForTesting(/*evict_ratio_target=*/1.0, + /*evict_ratio_lowerbound=*/1.0); + EXPECT_FALSE(service.GetReplicaList(key, kDefaultTenant).has_value()); + + const std::string before_finalize_key = + "before_batch_nof_evict_finalize_key"; + auto before_finalize = service.PutStart(client_id, before_finalize_key, + kDefaultTenant, 1024, config); + EXPECT_FALSE(before_finalize.has_value()); + + backend->AllowTxn(); + ReadBatchEventually(storage, 2, batch); + + if (!before_finalize.has_value()) { + const std::string after_finalize_key = + "after_batch_nof_evict_finalize_key"; + auto after_finalize = service.PutStart(client_id, after_finalize_key, + kDefaultTenant, 1024, config); + EXPECT_TRUE(after_finalize.has_value()) + << toString(after_finalize.error()); + } +} +#endif + +TEST_F(MasterServiceHATest, PutStartExpiredOverwriteWritesBatchRecordOpLog) { + const std::string cluster_id = + "test_batch_record_put_start_cleanup_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_put_start_discard_timeout_sec(1) + .set_put_start_release_timeout_sec(2) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + [[maybe_unused]] const auto mounted = + PrepareSimpleSegment(service, "batch_put_start_cleanup_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + ReplicateConfig config; + config.replica_num = 1; + const std::string key = "batch_put_start_cleanup_key"; + const UUID client_id = generate_uuid(); + ASSERT_TRUE(service.PutStart(client_id, key, kDefaultTenant, 1024, config) + .has_value()); + + std::this_thread::sleep_for(std::chrono::milliseconds(1100)); + const UUID next_client_id = generate_uuid(); + ASSERT_TRUE( + service.PutStart(next_client_id, key, kDefaultTenant, 1024, config) + .has_value()); + ReadBatchEventually(storage, 2, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::REMOVE, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(2u, batch.entries[0].sequence_id); +} + +TEST_F(MasterServiceHATest, + DiscardExpiredProcessingReplicasWritesBatchRecordOpLog) { + const std::string cluster_id = + "test_batch_record_discard_processing_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_put_start_discard_timeout_sec(1) + .set_put_start_release_timeout_sec(2) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + [[maybe_unused]] const auto mounted = + PrepareSimpleSegment(service, "batch_discard_processing_segment"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + ReplicateConfig config; + config.replica_num = 1; + const UUID client_id = generate_uuid(); + const std::string key = "batch_discard_processing_key"; + ASSERT_TRUE(service.PutStart(client_id, key, kDefaultTenant, 1024, config) + .has_value()); + + Replica local_disk_replica(client_id, 1024, "local_disk_endpoint", + ReplicaStatus::COMPLETE); + ASSERT_TRUE( + service.AddReplica(client_id, key, kDefaultTenant, local_disk_replica) + .has_value()); + ReadBatchEventually(storage, 2, batch); + + std::this_thread::sleep_for(std::chrono::milliseconds(2100)); + service.RunBatchEvictForTesting(/*evict_ratio_target=*/1.0, + /*evict_ratio_lowerbound=*/1.0); + ReadBatchEventually(storage, 3, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(3u, batch.entries[0].sequence_id); + EXPECT_FALSE(batch.entries[0].payload.empty()); +} + +TEST_F(MasterServiceHATest, + DiscardExpiredReplicationTaskWritesBatchRecordOpLog) { + const std::string cluster_id = + "test_batch_record_discard_replication_cluster"; + auto backend = std::make_shared(); + auto service_config = MasterServiceConfig::builder() + .set_default_kv_lease_ttl(50) + .set_enable_ha(true) + .set_enable_oplog(true) + .set_cluster_id(cluster_id) + .set_oplog_batch_max_entries(1) + .set_put_start_discard_timeout_sec(1) + .set_put_start_release_timeout_sec(2) + .build(); + MasterService service(service_config); + ASSERT_EQ(ErrorCode::OK, service.SetBatchOpLogBackendForTesting(backend)); + + auto src = PrepareSimpleSegment(service, "batch_replication_src"); + OpLogBatchStorage storage(cluster_id, *backend); + OpLogBatchRecord batch; + ReadBatchEventually(storage, 1, batch); + + const std::string key = "batch_discard_replication_key"; + PutObjectOnSegment(service, src.client_id, key, "batch_replication_src"); + ReadBatchEventually(storage, 2, batch); + + [[maybe_unused]] const auto target = + PrepareSimpleSegment(service, "batch_replication_target", + kDefaultSegmentBase + kDefaultSegmentSize); + ReadBatchEventually(storage, 3, batch); + + ASSERT_TRUE(service + .MoveStart(src.client_id, key, kDefaultTenant, + "batch_replication_src", + "batch_replication_target") + .has_value()); + + std::this_thread::sleep_for(std::chrono::milliseconds(2100)); + service.RunBatchEvictForTesting(/*evict_ratio_target=*/1.0, + /*evict_ratio_lowerbound=*/1.0); + ReadBatchEventually(storage, 4, batch); + + ASSERT_EQ(1u, batch.entries.size()); + EXPECT_EQ(OpType::PUT_END, batch.entries[0].op_type); + EXPECT_EQ(kDefaultTenant.value(), batch.entries[0].tenant_id); + EXPECT_EQ(key, batch.entries[0].object_key); + EXPECT_EQ(4u, batch.entries[0].sequence_id); + EXPECT_FALSE(batch.entries[0].payload.empty()); +} + +} // namespace mooncake::test + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-store/tests/ha/oplog/etcd_oplog_store_test.cpp b/mooncake-store/tests/ha/oplog/etcd_oplog_store_test.cpp deleted file mode 100644 index 25931bfafa..0000000000 --- a/mooncake-store/tests/ha/oplog/etcd_oplog_store_test.cpp +++ /dev/null @@ -1,407 +0,0 @@ -#include "ha/oplog/etcd_oplog_store.h" - -#include -#include -#include - -#include -#include -#include -#include - -#include "etcd_helper.h" - -DEFINE_string(etcd_endpoints, "0.0.0.0:2379", - "Etcd endpoints for EtcdOpLogStoreTest"); - -namespace mooncake::test { - -class EtcdOpLogStoreTest : public ::testing::Test { - protected: - static void SetUpTestSuite() { -#ifdef STORE_USE_ETCD - google::InitGoogleLogging("EtcdOpLogStoreTest"); - FLAGS_logtostderr = 1; - - ASSERT_EQ(ErrorCode::OK, - EtcdHelper::ConnectToEtcdStoreClient(FLAGS_etcd_endpoints)) - << "Failed to connect to etcd at " << FLAGS_etcd_endpoints; -#endif - } - - static void TearDownTestSuite() { -#ifdef STORE_USE_ETCD - google::ShutdownGoogleLogging(); -#endif - } - - void SetUp() override { -#ifndef STORE_USE_ETCD - GTEST_SKIP() - << "STORE_USE_ETCD is disabled, skipping EtcdOpLogStore tests."; -#else - cluster_id_ = "test_cluster_etcd_oplog_store"; - store_ = std::make_unique( - cluster_id_, - /*enable_latest_seq_batch_update=*/false, - /*enable_batch_write=*/true); - ASSERT_EQ(ErrorCode::OK, store_->Init()); - CleanupTestData(); -#endif - } - - void TearDown() override { -#ifdef STORE_USE_ETCD - CleanupTestData(); - store_.reset(); -#endif - } - - std::string cluster_id_; - std::unique_ptr store_; - - void CleanupTestData() { -#ifdef STORE_USE_ETCD - // Delete all keys under /oplog/{cluster_id_}/ prefix - std::string prefix = std::string("/oplog/") + cluster_id_ + "/"; - - auto prefix_end = [](std::string p) -> std::string { - for (int i = static_cast(p.size()) - 1; i >= 0; --i) { - unsigned char c = static_cast(p[i]); - if (c < 0xFF) { - p[i] = static_cast(c + 1); - p.resize(i + 1); - return p; - } - } - return std::string(1, '\0'); - }; - std::string end_key = prefix_end(prefix); - - (void)EtcdHelper::DeleteRange(prefix.c_str(), prefix.size(), - end_key.c_str(), end_key.size()); -#endif - } - - static OpLogEntry MakeEntry(uint64_t seq, OpType type, - const std::string& key, - const std::string& payload) { - OpLogEntry e; - e.sequence_id = seq; - e.timestamp_ms = 123456; - e.op_type = type; - e.object_key = key; - e.payload = payload; - e.checksum = 0; - e.prefix_hash = 0; - return e; - } -}; - -// ========== 3.1.1 Basic CRUD tests ========== - -TEST_F(EtcdOpLogStoreTest, TestWriteOpLog) { - OpLogEntry e = MakeEntry(1, OpType::PUT_END, "key1", "value1"); - - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); - - // Latest sequence ID should be updated to 1 - uint64_t latest = 0; - ASSERT_EQ(ErrorCode::OK, store_->GetLatestSequenceId(latest)); - EXPECT_EQ(1u, latest); -} - -TEST_F(EtcdOpLogStoreTest, TestReadOpLog) { - OpLogEntry e = MakeEntry(2, OpType::PUT_END, "key2", "value2"); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); - - OpLogEntry out; - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLog(2, out)); - - EXPECT_EQ(2u, out.sequence_id); - EXPECT_EQ(OpType::PUT_END, out.op_type); - EXPECT_EQ("key2", out.object_key); - EXPECT_EQ("value2", out.payload); -} - -TEST_F(EtcdOpLogStoreTest, TestReadOpLogSince) { - // Write multiple entries - for (uint64_t i = 10; i < 15; ++i) { - OpLogEntry e = MakeEntry(i, OpType::PUT_END, "key_" + std::to_string(i), - "value_" + std::to_string(i)); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); - } - - std::vector entries; - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLogSince(11, 10, entries)); - - // Expect entries with seq > 11: 12,13,14 - ASSERT_EQ(3u, entries.size()); - EXPECT_EQ(12u, entries[0].sequence_id); - EXPECT_EQ(13u, entries[1].sequence_id); - EXPECT_EQ(14u, entries[2].sequence_id); -} - -TEST_F(EtcdOpLogStoreTest, TestReadOpLogSince_Empty) { - std::vector entries; - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLogSince(1000, 10, entries)); - EXPECT_TRUE(entries.empty()); -} - -TEST_F(EtcdOpLogStoreTest, TestReadOpLogSince_Limit) { - for (uint64_t i = 1; i <= 5; ++i) { - OpLogEntry e = MakeEntry(i, OpType::PUT_END, "key_" + std::to_string(i), - "value_" + std::to_string(i)); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); - } - - std::vector entries; - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLogSince(0, 3, entries)); - ASSERT_EQ(3u, entries.size()); - EXPECT_EQ(1u, entries[0].sequence_id); - EXPECT_EQ(2u, entries[1].sequence_id); - EXPECT_EQ(3u, entries[2].sequence_id); -} - -// ========== 3.1.2 Serialization tests ========== - -TEST_F(EtcdOpLogStoreTest, TestSerializeDeserializeRoundTrip) { - OpLogEntry in = - MakeEntry(42, OpType::PUT_END, "roundtrip-key", "roundtrip-value"); - - // Indirectly verify serialization / deserialization via WriteOpLog + - // ReadOpLog - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(in)); - - OpLogEntry out; - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLog(42, out)); - - EXPECT_EQ(in.sequence_id, out.sequence_id); - EXPECT_EQ(in.op_type, out.op_type); - EXPECT_EQ(in.object_key, out.object_key); - EXPECT_EQ(in.payload, out.payload); -} - -TEST_F(EtcdOpLogStoreTest, TestDeserializeInvalidJson) { - // Write invalid JSON directly into etcd; subsequent ReadOpLog should return - // INTERNAL_ERROR - std::string key = "/oplog/" + cluster_id_ + "/00000000000000000077"; - std::string bad_json = "{ this is not valid json }"; - ASSERT_EQ(ErrorCode::OK, - EtcdHelper::Put(key.c_str(), key.size(), bad_json.c_str(), - bad_json.size())); - - OpLogEntry out; - ASSERT_EQ(ErrorCode::INTERNAL_ERROR, store_->ReadOpLog(77, out)); -} - -// ========== 3.1.3 Fencing tests ========== - -TEST_F(EtcdOpLogStoreTest, TestWriteOpLog_Fencing) { - OpLogEntry e1 = MakeEntry(100, OpType::PUT_END, "key_fence", "value1"); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e1)); - - // Same seq, same content => idempotent (OK) - OpLogEntry e2 = e1; - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e2)); - - // NOTE: Duplicate sequence_id with different content is not a supported - // production scenario (sequence_id is monotonic). The batching write path - // does not guarantee conflict detection for that case, so we don't assert - // on it here. -} - -TEST_F(EtcdOpLogStoreTest, TestWriteOpLog_Idempotent) { - OpLogEntry e = MakeEntry(200, OpType::PUT_END, "key_idem", "v"); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); - - // Repeatedly writing the exact same entry should return OK (idempotent) - // even if the underlying Create operation reports a transaction failure. - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); -} - -// ========== 3.1.4 Sequence ID management tests ========== - -TEST_F(EtcdOpLogStoreTest, TestGetLatestSequenceId) { - OpLogEntry e1 = MakeEntry(1, OpType::PUT_END, "k1", "v1"); - OpLogEntry e2 = MakeEntry(2, OpType::PUT_END, "k2", "v2"); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e1)); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e2)); - - uint64_t latest = 0; - ASSERT_EQ(ErrorCode::OK, store_->GetLatestSequenceId(latest)); - EXPECT_EQ(2u, latest); -} - -TEST_F(EtcdOpLogStoreTest, TestGetMaxSequenceIdAndEmpty) { - uint64_t max_seq = 0; - - // Empty cluster: after cleanup, GetMaxSequenceId should return - // OPLOG_ENTRY_NOT_FOUND - CleanupTestData(); - EXPECT_EQ(ErrorCode::OPLOG_ENTRY_NOT_FOUND, - store_->GetMaxSequenceId(max_seq)); - - // After writing several entries, MaxSequenceId should equal the last - // entry's seq - for (uint64_t i = 10; i <= 15; ++i) { - OpLogEntry e = MakeEntry(i, OpType::PUT_END, "key_" + std::to_string(i), - "value_" + std::to_string(i)); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); - } - - ASSERT_EQ(ErrorCode::OK, store_->GetMaxSequenceId(max_seq)); - EXPECT_EQ(15u, max_seq); -} - -TEST_F(EtcdOpLogStoreTest, TestUpdateLatestSequenceId) { - // Directly call UpdateLatestSequenceId, then GetLatestSequenceId should - // match - ASSERT_EQ(ErrorCode::OK, store_->UpdateLatestSequenceId(12345)); - - uint64_t latest = 0; - ASSERT_EQ(ErrorCode::OK, store_->GetLatestSequenceId(latest)); - EXPECT_EQ(12345u, latest); -} - -// ========== 3.1.5 Batch update tests ========== - -TEST_F(EtcdOpLogStoreTest, TestBatchUpdate_EnabledAndThreshold) { - // Use a store with batch enabled, then verify /latest is updated to the max - // seq - EtcdOpLogStore writer(cluster_id_, - /*enable_latest_seq_batch_update=*/true, - /*enable_batch_write=*/true); - ASSERT_EQ(ErrorCode::OK, writer.Init()); - - const uint64_t base_seq = 1000; - const int kEntries = 5; - for (int i = 0; i < kEntries; ++i) { - OpLogEntry e = MakeEntry(base_seq + i, OpType::PUT_END, - "batch_key_" + std::to_string(i), - "batch_val_" + std::to_string(i)); - ASSERT_EQ(ErrorCode::OK, writer.WriteOpLog(e)); - } - - // Wait a short period to give the batch thread a chance to flush `/latest` - std::this_thread::sleep_for(std::chrono::milliseconds(2 * 1000)); - - uint64_t latest = 0; - ASSERT_EQ(ErrorCode::OK, store_->GetLatestSequenceId(latest)); - EXPECT_EQ(base_seq + kEntries - 1, latest); -} - -TEST_F(EtcdOpLogStoreTest, TestBatchUpdate_FailurePlaceholder) { - GTEST_SKIP() << "Batch failure scenarios are better tested with a " - "fault-injection etcd wrapper."; -} - -// ========== 3.1.6 Cleanup tests ========== - -TEST_F(EtcdOpLogStoreTest, TestCleanupOpLogBeforeAndBoundary) { - // Write seq 1..5 - for (uint64_t i = 1; i <= 5; ++i) { - OpLogEntry e = - MakeEntry(i, OpType::PUT_END, "cleanup_key_" + std::to_string(i), - "cleanup_val_" + std::to_string(i)); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); - } - - // Cleanup seq < 3 => 1,2 should be deleted; 3,4,5 should remain - ASSERT_EQ(ErrorCode::OK, store_->CleanupOpLogBefore(3)); - - OpLogEntry out; - EXPECT_EQ(ErrorCode::OPLOG_ENTRY_NOT_FOUND, store_->ReadOpLog(1, out)); - EXPECT_EQ(ErrorCode::OPLOG_ENTRY_NOT_FOUND, store_->ReadOpLog(2, out)); - - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLog(3, out)); - EXPECT_EQ(3u, out.sequence_id); - - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLog(5, out)); - EXPECT_EQ(5u, out.sequence_id); -} - -TEST_F(EtcdOpLogStoreTest, TestCleanupOpLogBefore_Empty) { - // Cleanup on an empty cluster should return OK - CleanupTestData(); - EXPECT_EQ(ErrorCode::OK, store_->CleanupOpLogBefore(100)); -} - -// ========== 3.1.7 Cluster ID validation tests ========== - -TEST_F(EtcdOpLogStoreTest, TestInvalidClusterId_Rejected) { - // Invalid cluster_id (containing slashes) should trigger LOG(FATAL) and - // terminate - EXPECT_DEATH( - { - EtcdOpLogStore bad_store("invalid/cluster", false); - (void)bad_store; - }, - "Invalid cluster_id"); -} - -TEST_F(EtcdOpLogStoreTest, TestClusterIdNormalization) { - // Trailing slashes should be normalized away from the cluster_id - std::string raw_cluster = cluster_id_ + "///"; - EtcdOpLogStore normalized_store(raw_cluster, - /*enable_latest_seq_batch_update=*/false, - /*enable_batch_write=*/true); - ASSERT_EQ(ErrorCode::OK, normalized_store.Init()); - - OpLogEntry e = MakeEntry(999, OpType::PUT_END, "norm-key", "norm-val"); - ASSERT_EQ(ErrorCode::OK, normalized_store.WriteOpLog(e)); - - // Read the same seq via the current store_ to confirm the normalized - // cluster_id is used - OpLogEntry out; - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLog(999, out)); - EXPECT_EQ("norm-key", out.object_key); - EXPECT_EQ("norm-val", out.payload); -} - -// ========== 3.1.8 Pagination tests ========== - -TEST_F(EtcdOpLogStoreTest, TestReadOpLogSince_Pagination) { - // Write 20 entries and verify pagination via limit - for (uint64_t i = 1; i <= 20; ++i) { - OpLogEntry e = - MakeEntry(i, OpType::PUT_END, "page_key_" + std::to_string(i), - "page_val_" + std::to_string(i)); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); - } - - std::vector entries; - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLogSince(0, 20, entries)); - ASSERT_EQ(20u, entries.size()); - for (uint64_t i = 0; i < entries.size(); ++i) { - EXPECT_EQ(i + 1, entries[i].sequence_id); - } -} - -TEST_F(EtcdOpLogStoreTest, TestReadOpLogSince_LargeDataset) { - // Write a larger number of entries to verify ReadOpLogSince returns the - // first N correctly - const uint64_t total = 200; - const uint64_t limit = 150; - CleanupTestData(); - for (uint64_t i = 1; i <= total; ++i) { - OpLogEntry e = - MakeEntry(i, OpType::PUT_END, "large_key_" + std::to_string(i), - "large_val_" + std::to_string(i)); - ASSERT_EQ(ErrorCode::OK, store_->WriteOpLog(e)); - } - - std::vector entries; - ASSERT_EQ(ErrorCode::OK, store_->ReadOpLogSince(0, limit, entries)); - ASSERT_EQ(limit, entries.size()); - for (uint64_t i = 0; i < limit; ++i) { - EXPECT_EQ(i + 1, entries[i].sequence_id); - } -} - -} // namespace mooncake::test - -int main(int argc, char** argv) { - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/mooncake-store/tests/ha/oplog/ha_recovery_test.cpp b/mooncake-store/tests/ha/oplog/ha_recovery_test.cpp deleted file mode 100644 index d9fb1a8f8f..0000000000 --- a/mooncake-store/tests/ha/oplog/ha_recovery_test.cpp +++ /dev/null @@ -1,359 +0,0 @@ -// mooncake-store/tests/hot_standby_ut/ha_recovery_test.cpp -// -// Mock-layer tests for HA recovery scenarios. -// Tests at component level: OpLogApplier + MockOpLogStore + -// MockMetadataStore + MockSnapshotProvider. -// Does NOT go through HotStandbyService (OpLogStoreFactory does not -// support creating MockOpLogStore). - -#include -#include -#include -#include - -#include -#include - -#include "mock_metadata_store.h" -#include "mock_oplog_store.h" -#include "mock_snapshot_provider.h" -#include "ha/oplog/oplog_applier.h" -#include "ha/oplog/oplog_manager.h" - -using namespace mooncake; -using namespace mooncake::test; - -// Helper: create a struct_pack-serialized payload (same format as OpLogManager) -static std::string MakeValidPayload(uint64_t client_first = 1, - uint64_t client_second = 2, - uint64_t size = 1024) { - MetadataPayload payload; - payload.client_id = {client_first, client_second}; - payload.size = size; - auto result = struct_pack::serialize(payload); - return std::string(result.begin(), result.end()); -} - -// Helper: create an OpLogEntry with PUT_END and valid checksum -static OpLogEntry MakeEntry(uint64_t seq, const std::string& key) { - OpLogEntry entry; - entry.sequence_id = seq; - entry.timestamp_ms = - std::chrono::duration_cast( - std::chrono::steady_clock::now().time_since_epoch()) - .count(); - entry.op_type = OpType::PUT_END; - entry.object_key = key; - entry.payload = MakeValidPayload(); - entry.checksum = static_cast( - XXH32(entry.payload.data(), entry.payload.size(), 0)); - entry.prefix_hash = - key.empty() ? 0 - : static_cast(XXH32(key.data(), key.size(), 0)); - return entry; -} - -class HaRecoveryTest : public ::testing::Test { - protected: - void SetUp() override { - mock_store_ = std::make_shared(); - mock_store_->Init(); - mock_snapshot_ = std::make_shared(); - metadata_store_ = std::make_shared(); - applier_ = std::make_unique( - metadata_store_.get(), "test_cluster", mock_store_.get()); - } - - // Write N entries to MockOpLogStore with seq [start, start+count) - void WriteEntriesToStore(uint64_t start, uint64_t count) { - for (uint64_t i = 0; i < count; ++i) { - uint64_t seq = start + i; - mock_store_->WriteOpLog( - MakeEntry(seq, "key_" + std::to_string(seq))); - } - } - - // Build snapshot data for keys key_1..key_count - std::vector> - BuildSnapshotData(uint64_t count) { - std::vector> data; - for (uint64_t i = 1; i <= count; ++i) { - StandbyObjectMetadata meta; - meta.client_id = {1, 2}; - meta.size = 1024; - meta.last_sequence_id = i; - data.emplace_back("key_" + std::to_string(i), meta); - } - return data; - } - - // Load snapshot into metadata_store_ and recover applier - bool LoadSnapshotAndRecover() { - auto result = mock_snapshot_->LoadLatestSnapshot("test_cluster"); - - if (!result.has_value() || !result->has_value()) { - return false; - } - const auto& snap = result->value(); - for (const auto& [key, meta] : snap.metadata) { - metadata_store_->PutMetadata(key, meta); - } - applier_->Recover(snap.snapshot_sequence_id); - return true; - } - - std::shared_ptr mock_store_; - std::shared_ptr mock_snapshot_; - std::shared_ptr metadata_store_; - std::unique_ptr applier_; -}; - -// ============================================================ -// Scenario 1: OpLog + Snapshot Joint Recovery -// ============================================================ - -TEST_F(HaRecoveryTest, SnapshotThenOpLogReplay) { - // Snapshot @ seq=10 with 10 entries - mock_snapshot_->SetSnapshot("snap1", 10, BuildSnapshotData(10)); - - // OpLog store has seq 11~20 - WriteEntriesToStore(11, 10); - - // Load snapshot, recover applier to seq=10 - ASSERT_TRUE(LoadSnapshotAndRecover()); - EXPECT_EQ(metadata_store_->Size(), 10u); - EXPECT_EQ(applier_->GetExpectedSequenceId(), 11u); - - // Read and apply entries 11~20 - std::vector entries; - ASSERT_EQ(ErrorCode::OK, mock_store_->ReadOpLogSince(10, 1000, entries)); - EXPECT_EQ(entries.size(), 10u); - EXPECT_EQ(applier_->ApplyOpLogEntries(entries), 10u); - - // 10 snapshot + 10 OpLog = 20 total - EXPECT_EQ(metadata_store_->Size(), 20u); - EXPECT_EQ(applier_->GetExpectedSequenceId(), 21u); -} - -TEST_F(HaRecoveryTest, SnapshotLoadFail_FallbackToFullReplay) { - mock_snapshot_->SetLoadFail(true); - - // OpLog store has seq 1~20 (full history) - WriteEntriesToStore(1, 20); - - // Snapshot fails -> recover from 0 - ASSERT_FALSE(LoadSnapshotAndRecover()); - applier_->Recover(0); - - // Read all and apply - std::vector entries; - ASSERT_EQ(ErrorCode::OK, mock_store_->ReadOpLogSince(0, 1000, entries)); - EXPECT_EQ(entries.size(), 20u); - EXPECT_EQ(applier_->ApplyOpLogEntries(entries), 20u); - - EXPECT_EQ(metadata_store_->Size(), 20u); -} - -TEST_F(HaRecoveryTest, SnapshotWithNoSubsequentOpLog) { - // Snapshot @ seq=10, no subsequent OpLog - mock_snapshot_->SetSnapshot("snap1", 10, BuildSnapshotData(10)); - - ASSERT_TRUE(LoadSnapshotAndRecover()); - EXPECT_EQ(metadata_store_->Size(), 10u); - - std::vector entries; - ASSERT_EQ(ErrorCode::OK, mock_store_->ReadOpLogSince(10, 1000, entries)); - EXPECT_EQ(entries.size(), 0u); - EXPECT_EQ(metadata_store_->Size(), 10u); -} - -TEST_F(HaRecoveryTest, SnapshotSeqMismatch_GapInOpLog) { - // Snapshot @ seq=10, but OpLog starts at seq=15 (11~14 missing) - mock_snapshot_->SetSnapshot("snap1", 10, BuildSnapshotData(10)); - WriteEntriesToStore(15, 6); // seq 15~20 - - ASSERT_TRUE(LoadSnapshotAndRecover()); - EXPECT_EQ(applier_->GetExpectedSequenceId(), 11u); - - // Feed seq 15~20 (expected=11, so gap for 11~14) - std::vector entries; - ASSERT_EQ(ErrorCode::OK, mock_store_->ReadOpLogSince(10, 1000, entries)); - EXPECT_EQ(entries.size(), 6u); - - for (const auto& e : entries) { - applier_->ApplyOpLogEntry(e); - } - - // Wait for gap timeout (kMissingEntrySkipSeconds = 3s) - std::this_thread::sleep_for(std::chrono::seconds(4)); - applier_->ProcessPendingEntries(); - - // Snapshot data preserved, gaps eventually skipped, pending applied - EXPECT_GE(metadata_store_->Size(), 10u); -} - -// ============================================================ -// Scenario 2: OpLog GC + Snapshot Safety -// ============================================================ - -TEST_F(HaRecoveryTest, GC_BasicCleanupSemantics) { - WriteEntriesToStore(1, 20); - EXPECT_EQ(mock_store_->EntryCount(), 20u); - - // Cleanup before seq=15 -> deletes 1~14 - ASSERT_EQ(ErrorCode::OK, mock_store_->CleanupOpLogBefore(15)); - EXPECT_EQ(mock_store_->EntryCount(), 6u); - - // Verify deleted vs preserved - OpLogEntry entry; - EXPECT_EQ(ErrorCode::OPLOG_ENTRY_NOT_FOUND, - mock_store_->ReadOpLog(14, entry)); - EXPECT_EQ(ErrorCode::OK, mock_store_->ReadOpLog(15, entry)); - EXPECT_EQ(entry.sequence_id, 15u); - - // Snapshot metadata is independent - ASSERT_EQ(ErrorCode::OK, - mock_store_->RecordSnapshotSequenceId("snap1", 10)); - uint64_t snap_seq = 0; - ASSERT_EQ(ErrorCode::OK, - mock_store_->GetSnapshotSequenceId("snap1", snap_seq)); - EXPECT_EQ(snap_seq, 10u); -} - -TEST_F(HaRecoveryTest, GC_NewStandbyAfterCleanup) { - WriteEntriesToStore(1, 20); - ASSERT_EQ(ErrorCode::OK, mock_store_->CleanupOpLogBefore(10)); - - // Snapshot @ seq=10 covers keys 1~10 - mock_snapshot_->SetSnapshot("snap1", 10, BuildSnapshotData(10)); - - ASSERT_TRUE(LoadSnapshotAndRecover()); - EXPECT_EQ(metadata_store_->Size(), 10u); - EXPECT_EQ(applier_->GetExpectedSequenceId(), 11u); - - // ReadOpLogSince(10) returns seq 11~20 - std::vector entries; - ASSERT_EQ(ErrorCode::OK, mock_store_->ReadOpLogSince(10, 1000, entries)); - EXPECT_EQ(entries.size(), 10u); - EXPECT_EQ(applier_->ApplyOpLogEntries(entries), 10u); - - EXPECT_EQ(metadata_store_->Size(), 20u); -} - -TEST_F(HaRecoveryTest, GC_NoSnapshot_IncompleteRecovery) { - WriteEntriesToStore(1, 20); - ASSERT_EQ(ErrorCode::OK, mock_store_->CleanupOpLogBefore(10)); - - // No snapshot - mock_snapshot_->SetLoadFail(true); - ASSERT_FALSE(LoadSnapshotAndRecover()); - applier_->Recover(0); - - // ReadOpLogSince(0) returns seq 10~20 (11 entries) - std::vector entries; - ASSERT_EQ(ErrorCode::OK, mock_store_->ReadOpLogSince(0, 1000, entries)); - EXPECT_EQ(entries.size(), 11u); - - // Applier expects seq=1 but first entry is seq=10. This creates a gap. - // Entries go to pending buffer — applier cannot apply them immediately. - // This demonstrates data incompleteness without snapshot protection: - // the applier sees a gap it cannot fill (seq 1~9 are GC'd). - size_t applied = applier_->ApplyOpLogEntries(entries); - EXPECT_EQ(applied, 0); - - // Wait for gap timeout so pending entries eventually drain - std::this_thread::sleep_for(std::chrono::seconds(4)); - applier_->ProcessPendingEntries(); - - // After timeout, gaps 1~9 are skipped, entries 10~20 apply. - // Key insight: without snapshot, recovery is delayed and data for - // seq 1~9 is permanently lost. - EXPECT_LE(metadata_store_->Size(), 11u); -} - -// ============================================================ -// Scenario 3: Promotion Write Consistency -// ============================================================ - -TEST_F(HaRecoveryTest, PromotionCatchUp_AllEntriesApplied) { - applier_->Recover(0); - for (uint64_t i = 1; i <= 10; ++i) { - EXPECT_TRUE(applier_->ApplyOpLogEntry( - MakeEntry(i, "key_" + std::to_string(i)))); - } - EXPECT_EQ(applier_->GetExpectedSequenceId(), 11u); - EXPECT_EQ(metadata_store_->Size(), 10u); - - // seq 11~15 in store but not applied (simulates lag) - WriteEntriesToStore(11, 5); - - // Final catch-up - std::vector entries; - ASSERT_EQ(ErrorCode::OK, mock_store_->ReadOpLogSince(10, 1000, entries)); - EXPECT_EQ(entries.size(), 5u); - EXPECT_EQ(applier_->ApplyOpLogEntries(entries), 5u); - - EXPECT_EQ(metadata_store_->Size(), 15u); - EXPECT_EQ(applier_->GetExpectedSequenceId(), 16u); -} - -TEST_F(HaRecoveryTest, PromotionWithPendingGaps) { - applier_->Recover(0); - for (uint64_t i = 1; i <= 7; ++i) { - EXPECT_TRUE(applier_->ApplyOpLogEntry( - MakeEntry(i, "key_" + std::to_string(i)))); - } - EXPECT_EQ(applier_->GetExpectedSequenceId(), 8u); - - // Skip seq=8, apply seq 9~10 -> go to pending - applier_->ApplyOpLogEntry(MakeEntry(9, "key_9")); - applier_->ApplyOpLogEntry(MakeEntry(10, "key_10")); - EXPECT_EQ(applier_->GetExpectedSequenceId(), 8u); - EXPECT_EQ(metadata_store_->Size(), 7u); - - // Write seq=8 to store (late arrival) - mock_store_->WriteOpLog(MakeEntry(8, "key_8")); - - // First ProcessPendingEntries: detect gap, register in - // missing_sequence_ids_ - applier_->ProcessPendingEntries(); - - // Wait for kMissingEntryRequestSeconds (1s) so next call fetches seq=8 - std::this_thread::sleep_for(std::chrono::seconds(2)); - - // Second ProcessPendingEntries: request seq=8 from store, apply it, - // then drain pending 9~10 - applier_->ProcessPendingEntries(); - - EXPECT_EQ(applier_->GetExpectedSequenceId(), 11u); - EXPECT_EQ(metadata_store_->Size(), 10u); -} - -TEST_F(HaRecoveryTest, PromotionGapUnresolvable_SkipAndPromote) { - applier_->Recover(0); - for (uint64_t i = 1; i <= 7; ++i) { - EXPECT_TRUE(applier_->ApplyOpLogEntry( - MakeEntry(i, "key_" + std::to_string(i)))); - } - - // Skip seq=8, apply seq 9~10 -> pending - applier_->ApplyOpLogEntry(MakeEntry(9, "key_9")); - applier_->ApplyOpLogEntry(MakeEntry(10, "key_10")); - EXPECT_EQ(applier_->GetExpectedSequenceId(), 8u); - - // seq=8 NOT in store (GC'd). Let gap detection register. - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - applier_->ProcessPendingEntries(); - - auto result = applier_->TryResolveGapsOnceForPromotion(1024); - // seq=8 not in store, so fetched=0 - EXPECT_EQ(result.fetched, 0u); - - // Wait for skip timeout (3s) so gap is eventually skipped - std::this_thread::sleep_for(std::chrono::seconds(4)); - applier_->ProcessPendingEntries(); - - // After skip timeout, gap is skipped, pending 9~10 should drain - // No crash, no hang. Original data preserved. - EXPECT_GE(metadata_store_->Size(), 7u); -} diff --git a/mooncake-store/tests/ha/oplog/localfs_hot_standby_integration_test.cpp b/mooncake-store/tests/ha/oplog/localfs_hot_standby_integration_test.cpp deleted file mode 100644 index 7dd8d40336..0000000000 --- a/mooncake-store/tests/ha/oplog/localfs_hot_standby_integration_test.cpp +++ /dev/null @@ -1,463 +0,0 @@ -// mooncake-store/tests/localfs_hot_standby_integration_test.cpp -// -// End-to-end integration tests for the HA replication flow using LocalFS -// backend. Mirrors the structure of hot_standby_integration_test.cpp but -// replaces etcd with a shared temp directory so the tests run without any -// external service. -// -// Test flow: -// Primary OpLogManager -> LocalFsOpLogStore (WRITER) -// -> shared filesystem directory <- -// HotStandbyService (PollingOpLogChangeNotifier -> OpLogReplicator -// -> OpLogApplier -> StandbyMetadataStore) - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "hot_standby_service.h" -#include "ha/oplog/oplog_manager.h" -#include "ha/oplog/oplog_store_factory.h" -#include "standby_state_machine.h" - -namespace mooncake { -namespace testing { - -// ============================================================ -// RAII helper to ensure HotStandbyService is always stopped -// ============================================================ - -class StandbyServiceGuard { - public: - explicit StandbyServiceGuard(HotStandbyService* service) - : service_(service) {} - ~StandbyServiceGuard() { - if (service_) { - service_->Stop(); - // LocalFS has no goroutines to drain; a short sleep suffices to - // let background polling threads exit cleanly. - std::this_thread::sleep_for(std::chrono::milliseconds(500)); - } - } - StandbyServiceGuard(const StandbyServiceGuard&) = delete; - StandbyServiceGuard& operator=(const StandbyServiceGuard&) = delete; - - private: - HotStandbyService* service_; -}; - -// ============================================================ -// Fixture -// ============================================================ - -class LocalFsHotStandbyIntegrationTest : public ::testing::Test { - protected: - void SetUp() override { - // Generate a unique temp directory per test - static std::atomic counter{0}; - test_dir_ = "/tmp/localfs_ha_test_" + std::to_string(getpid()) + "_" + - std::to_string(counter.fetch_add(1)); - cluster_id_ = "localfs_ha_cluster"; - poll_interval_ms_ = 100; // fast polling for tests - - std::filesystem::create_directories(test_dir_); - } - - void TearDown() override { - std::error_code ec; - std::filesystem::remove_all(test_dir_, ec); - if (ec) { - LOG(WARNING) << "Failed to remove test dir " << test_dir_ << ": " - << ec.message(); - } - } - - // -- Helpers -- - - std::unique_ptr CreatePrimaryOpLogManager() { - auto store = OpLogStoreFactory::Create( - OpLogStoreType::LOCAL_FS, cluster_id_, OpLogStoreRole::WRITER, - test_dir_, poll_interval_ms_); - if (!store) return nullptr; - auto mgr = std::make_unique(); - mgr->SetOpLogStore(std::shared_ptr(std::move(store))); - return mgr; - } - - HotStandbyConfig MakeHotStandbyConfig() { - HotStandbyConfig cfg; - cfg.enable_verification = false; - cfg.max_replication_lag_entries = 1000; - cfg.oplog_store_type = OpLogStoreType::LOCAL_FS; - cfg.oplog_store_root_dir = test_dir_; - cfg.oplog_poll_interval_ms = poll_interval_ms_; - return cfg; - } - - bool WaitForSync(HotStandbyService& standby, uint64_t target_seq, - int timeout_sec = 30) { - auto deadline = std::chrono::steady_clock::now() + - std::chrono::seconds(timeout_sec); - while (std::chrono::steady_clock::now() < deadline) { - auto status = standby.GetSyncStatus(); - LOG(INFO) << "Standby: state=" << StandbyStateToString(status.state) - << ", applied_seq_id=" << status.applied_seq_id - << ", primary_seq_id=" << status.primary_seq_id - << ", lag_entries=" << status.lag_entries; - if (status.state == StandbyState::WATCHING && - status.lag_entries == 0 && - status.applied_seq_id >= target_seq) { - return true; - } - std::this_thread::sleep_for(std::chrono::milliseconds(200)); - } - return false; - } - - std::string test_dir_; - std::string cluster_id_; - int poll_interval_ms_; -}; - -// ============================================================ -// Test cases -// ============================================================ - -TEST_F(LocalFsHotStandbyIntegrationTest, TestPrimaryStandbySync) { - // 1. Primary writes 10 entries (last one via AppendAndPersist to flush) - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - std::vector test_keys; - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - - for (int i = 0; i < 10; ++i) { - std::string key = "test_key_" + std::to_string(i); - if (i < 9) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()) - << "AppendAndPersist failed for last entry"; - } - test_keys.push_back(key); - } - - uint64_t last_seq_id = primary->GetLastSequenceId(); - LOG(INFO) << "Primary wrote " << last_seq_id << " OpLog entries"; - - // 2. Start the standby - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, - standby.Start("", /*oplog_endpoints=*/"", cluster_id_)); - - // 3. Wait for sync - ASSERT_TRUE(WaitForSync(standby, last_seq_id)) - << "Standby failed to sync within timeout"; - - // 4. Verify metadata snapshot - std::vector> snapshot; - ASSERT_TRUE(standby.ExportMetadataSnapshot(snapshot)); - LOG(INFO) << "Standby metadata snapshot size: " << snapshot.size(); - - std::set snapshot_keys; - for (const auto& kv : snapshot) { - snapshot_keys.insert(kv.first); - } - for (const auto& key : test_keys) { - EXPECT_NE(snapshot_keys.end(), snapshot_keys.find(key)) - << "Key " << key << " not found in Standby snapshot"; - } - EXPECT_GE(snapshot.size(), test_keys.size()); - - // 5. Verify sequence IDs - EXPECT_GE(standby.GetLatestAppliedSequenceId(), last_seq_id); -} - -TEST_F(LocalFsHotStandbyIntegrationTest, TestStandbyPromotion) { - // 1. Write entries - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - std::vector test_keys; - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - - for (int i = 0; i < 5; ++i) { - std::string key = "promote_test_key_" + std::to_string(i); - if (i < 4) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()); - } - test_keys.push_back(key); - } - - uint64_t last_seq_id = primary->GetLastSequenceId(); - - // 2. Start standby and wait for sync - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, standby.Start("", "", cluster_id_)); - - ASSERT_TRUE(WaitForSync(standby, last_seq_id)); - - // 3. Verify ready for promotion - ASSERT_TRUE(standby.IsReadyForPromotion()); - - // 4. Promote - EXPECT_EQ(ErrorCode::OK, standby.Promote()); - - // 5. Verify sequence ID - EXPECT_GE(standby.GetLatestAppliedSequenceId(), last_seq_id); - - // 6. Verify metadata preserved - std::vector> snapshot; - ASSERT_TRUE(standby.ExportMetadataSnapshot(snapshot)); - - std::set snapshot_keys; - for (const auto& kv : snapshot) { - snapshot_keys.insert(kv.first); - } - for (const auto& key : test_keys) { - EXPECT_NE(snapshot_keys.end(), snapshot_keys.find(key)) - << "Key " << key << " should be in snapshot after promotion"; - } -} - -TEST_F(LocalFsHotStandbyIntegrationTest, TestFailoverScenario) { - // 1. Primary writes data - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - std::vector test_keys; - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - - for (int i = 0; i < 10; ++i) { - std::string key = "failover_key_" + std::to_string(i); - if (i < 9) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()); - } - test_keys.push_back(key); - } - - uint64_t last_seq_id = primary->GetLastSequenceId(); - - // 2. Start standby - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, standby.Start("", "", cluster_id_)); - - ASSERT_TRUE(WaitForSync(standby, last_seq_id)); - - // 3. Simulate primary failure (destroy the OpLogManager) - primary.reset(); - - // 4. Verify standby data integrity - std::vector> snapshot; - ASSERT_TRUE(standby.ExportMetadataSnapshot(snapshot)); - - std::set snapshot_keys; - for (const auto& kv : snapshot) { - snapshot_keys.insert(kv.first); - } - for (const auto& key : test_keys) { - EXPECT_NE(snapshot_keys.end(), snapshot_keys.find(key)) - << "Key " << key << " should be in Standby after Primary failure"; - } - - // 5. Promote - ASSERT_TRUE(standby.IsReadyForPromotion()); - EXPECT_EQ(ErrorCode::OK, standby.Promote()); - - // 6. Verify state - EXPECT_GE(standby.GetLatestAppliedSequenceId(), last_seq_id); -} - -TEST_F(LocalFsHotStandbyIntegrationTest, TestDataConsistency) { - // 1. Mixed PUT and REMOVE operations - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - std::map expected_keys; // key -> should_exist - - // PUT 5 keys - for (int i = 0; i < 5; ++i) { - std::string key = "put_key_" + std::to_string(i); - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - primary->Append(OpType::PUT_END, key, payload); - expected_keys[key] = true; - } - - // REMOVE first 2 - for (int i = 0; i < 2; ++i) { - std::string key = "put_key_" + std::to_string(i); - primary->Append(OpType::REMOVE, key, ""); - expected_keys[key] = false; - } - - // PUT 3 more (last one sync) - for (int i = 5; i < 8; ++i) { - std::string key = "put_key_" + std::to_string(i); - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":2048,"replicas":[]})"; - if (i < 7) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()); - } - expected_keys[key] = true; - } - - uint64_t last_seq_id = primary->GetLastSequenceId(); - LOG(INFO) << "Primary wrote " << last_seq_id << " OpLog entries"; - - // 2. Start standby - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, standby.Start("", "", cluster_id_)); - - // 3. Wait for sync - ASSERT_TRUE(WaitForSync(standby, last_seq_id)); - - // 4. Verify consistency - std::vector> snapshot; - ASSERT_TRUE(standby.ExportMetadataSnapshot(snapshot)); - - std::set actual_keys; - for (const auto& kv : snapshot) { - actual_keys.insert(kv.first); - } - - for (const auto& kv : expected_keys) { - if (kv.second) { - EXPECT_NE(actual_keys.end(), actual_keys.find(kv.first)) - << "Key " << kv.first << " should exist but not found"; - } else { - EXPECT_EQ(actual_keys.end(), actual_keys.find(kv.first)) - << "Key " << kv.first << " should be removed but still exists"; - } - } - - size_t expected_count = 0; - for (const auto& kv : expected_keys) { - if (kv.second) expected_count++; - } - EXPECT_EQ(expected_count, actual_keys.size()); -} - -TEST_F(LocalFsHotStandbyIntegrationTest, TestHighThroughputSync) { - // 1. Create primary first (initializes directory structure) - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - // 2. Start standby (will poll for new entries) - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, standby.Start("", "", cluster_id_)); - - // Wait for WATCHING state - { - auto deadline = - std::chrono::steady_clock::now() + std::chrono::seconds(10); - while (std::chrono::steady_clock::now() < deadline) { - if (standby.GetSyncStatus().state == StandbyState::WATCHING) { - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - } - } - - // 3. High-throughput writes (last one sync) - - const int num_writes = 100; - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - - auto write_start = std::chrono::steady_clock::now(); - for (int i = 0; i < num_writes; ++i) { - std::string key = "throughput_key_" + std::to_string(i); - if (i < num_writes - 1) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()); - } - } - auto write_end = std::chrono::steady_clock::now(); - auto write_duration = std::chrono::duration_cast( - write_end - write_start); - - uint64_t last_seq_id = primary->GetLastSequenceId(); - LOG(INFO) << "Wrote " << num_writes << " entries in " - << write_duration.count() << "ms, last_seq_id=" << last_seq_id; - - // 3. Monitor lag - auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); - uint64_t max_lag = 0; - - while (std::chrono::steady_clock::now() < deadline) { - auto status = standby.GetSyncStatus(); - if (status.lag_entries > max_lag) { - max_lag = status.lag_entries; - } - LOG(INFO) << "Standby lag: " << status.lag_entries - << " entries, applied_seq_id=" << status.applied_seq_id - << ", primary_seq_id=" << status.primary_seq_id; - if (status.applied_seq_id >= last_seq_id && status.lag_entries == 0) { - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(200)); - } - - // 4. Verify - auto final_status = standby.GetSyncStatus(); - EXPECT_GE(final_status.applied_seq_id, last_seq_id) - << "Standby should have applied all entries"; - EXPECT_EQ(0u, final_status.lag_entries) - << "Standby lag should be zero after sync"; - - LOG(INFO) << "Max lag observed: " << max_lag << " entries"; -} - -} // namespace testing -} // namespace mooncake - -int main(int argc, char** argv) { - gflags::ParseCommandLineFlags(&argc, &argv, true); - google::InitGoogleLogging(argv[0]); - google::SetVLOGLevel("*", 1); - FLAGS_logtostderr = 1; - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/mooncake-store/tests/ha/oplog/localfs_oplog_store_test.cpp b/mooncake-store/tests/ha/oplog/localfs_oplog_store_test.cpp deleted file mode 100644 index 5b202a2c77..0000000000 --- a/mooncake-store/tests/ha/oplog/localfs_oplog_store_test.cpp +++ /dev/null @@ -1,567 +0,0 @@ -#include -#include -#include -#include - -#include "ha/oplog/localfs_oplog_store.h" -#include "ha/oplog/oplog_store_factory.h" - -namespace fs = std::filesystem; - -namespace mooncake { -namespace { - -class LocalFsOpLogStoreTest : public ::testing::Test { - protected: - void SetUp() override { - test_dir_ = "/tmp/localfs_oplog_test_" + std::to_string(::getpid()) + - "_" + std::to_string(test_counter_++); - fs::create_directories(test_dir_); - cluster_id_ = "test_cluster"; - } - - void TearDown() override { fs::remove_all(test_dir_); } - - std::unique_ptr CreateWriter() { - auto store = std::make_unique( - cluster_id_, test_dir_, /*enable_batch_write=*/true); - EXPECT_EQ(ErrorCode::OK, store->Init()); - return store; - } - - std::unique_ptr CreateReader() { - auto store = std::make_unique( - cluster_id_, test_dir_, /*enable_batch_write=*/false); - EXPECT_EQ(ErrorCode::OK, store->Init()); - return store; - } - - OpLogEntry MakeEntry(uint64_t seq_id, const std::string& key = "test_key") { - OpLogEntry entry; - entry.sequence_id = seq_id; - entry.op_type = OpType::PUT_END; - entry.object_key = key; - entry.payload = "payload_" + std::to_string(seq_id); - return entry; - } - - std::string test_dir_; - std::string cluster_id_; - static inline int test_counter_ = 0; -}; - -// --- Init tests --- - -TEST_F(LocalFsOpLogStoreTest, InitCreatesDirectoryStructure) { - auto store = CreateWriter(); - EXPECT_TRUE(fs::exists(test_dir_ + "/" + cluster_id_ + "/segments")); - // snapshots dir is created lazily on first RecordSnapshotSequenceId - EXPECT_FALSE(fs::exists(test_dir_ + "/" + cluster_id_ + "/snapshots")); - EXPECT_TRUE(fs::exists(test_dir_ + "/" + cluster_id_ + "/latest")); -} - -TEST_F(LocalFsOpLogStoreTest, InitWriterCreatesLatestFile) { - auto store = CreateWriter(); - std::string latest_path = test_dir_ + "/" + cluster_id_ + "/latest"; - std::ifstream f(latest_path); - std::string content; - f >> content; - EXPECT_EQ("0", content); -} - -TEST_F(LocalFsOpLogStoreTest, InitReaderDoesNotCreateLatestFile) { - // Remove the test dir so Init starts fresh - fs::remove_all(test_dir_); - fs::create_directories(test_dir_); - auto store = std::make_unique( - cluster_id_, test_dir_, /*enable_batch_write=*/false); - EXPECT_EQ(ErrorCode::OK, store->Init()); - // Reader should create dirs but NOT create latest file - EXPECT_TRUE(fs::exists(test_dir_ + "/" + cluster_id_ + "/segments")); - EXPECT_FALSE(fs::exists(test_dir_ + "/" + cluster_id_ + "/latest")); -} - -TEST_F(LocalFsOpLogStoreTest, InitCleansTmpFiles) { - // Create directory structure manually - std::string seg_dir = test_dir_ + "/" + cluster_id_ + "/segments"; - fs::create_directories(seg_dir); - // Create a stale .tmp file - std::ofstream(seg_dir + - "/seg_00000000000000000001_00000000000000000010.tmp") - << "stale data"; - EXPECT_TRUE(fs::exists( - seg_dir + "/seg_00000000000000000001_00000000000000000010.tmp")); - - auto store = CreateWriter(); - EXPECT_FALSE(fs::exists( - seg_dir + "/seg_00000000000000000001_00000000000000000010.tmp")); -} - -// --- Write tests --- - -TEST_F(LocalFsOpLogStoreTest, WriteSyncEntry) { - auto store = CreateWriter(); - auto entry = MakeEntry(1); - EXPECT_EQ(ErrorCode::OK, store->WriteOpLog(entry, /*sync=*/true)); - - // Verify it was flushed to a segment file - auto segments_path = test_dir_ + "/" + cluster_id_ + "/segments"; - int file_count = 0; - for (auto& p : fs::directory_iterator(segments_path)) { - if (p.path().extension() != ".tmp") file_count++; - } - EXPECT_GE(file_count, 1); -} - -TEST_F(LocalFsOpLogStoreTest, WriteAsyncBatchFlush) { - auto store = CreateWriter(); - // Write several async entries - for (uint64_t i = 1; i <= 5; i++) { - EXPECT_EQ(ErrorCode::OK, - store->WriteOpLog(MakeEntry(i), /*sync=*/false)); - } - // Force flush by writing a sync entry - EXPECT_EQ(ErrorCode::OK, store->WriteOpLog(MakeEntry(6), /*sync=*/true)); - - // All entries should be readable - OpLogEntry read_entry; - for (uint64_t i = 1; i <= 6; i++) { - EXPECT_EQ(ErrorCode::OK, store->ReadOpLog(i, read_entry)); - EXPECT_EQ(i, read_entry.sequence_id); - } -} - -TEST_F(LocalFsOpLogStoreTest, WriteReaderReturnsError) { - auto store = CreateReader(); - auto entry = MakeEntry(1); - EXPECT_EQ(ErrorCode::INVALID_PARAMS, store->WriteOpLog(entry, true)); -} - -TEST_F(LocalFsOpLogStoreTest, WriteUpdatesLatestFile) { - auto store = CreateWriter(); - for (uint64_t i = 1; i <= 3; i++) { - store->WriteOpLog(MakeEntry(i), /*sync=*/false); - } - store->WriteOpLog(MakeEntry(4), /*sync=*/true); - - uint64_t latest = 0; - EXPECT_EQ(ErrorCode::OK, store->GetLatestSequenceId(latest)); - EXPECT_EQ(4, latest); -} - -TEST_F(LocalFsOpLogStoreTest, WriteOutOfOrderBatchStillReadableBySequenceId) { - auto store = CreateWriter(); - - EXPECT_EQ(ErrorCode::OK, store->WriteOpLog(MakeEntry(2), /*sync=*/false)); - EXPECT_EQ(ErrorCode::OK, store->WriteOpLog(MakeEntry(1), /*sync=*/true)); - - OpLogEntry entry1; - OpLogEntry entry2; - EXPECT_EQ(ErrorCode::OK, store->ReadOpLog(1, entry1)); - EXPECT_EQ(ErrorCode::OK, store->ReadOpLog(2, entry2)); - EXPECT_EQ(1u, entry1.sequence_id); - EXPECT_EQ(2u, entry2.sequence_id); - - std::vector entries; - EXPECT_EQ(ErrorCode::OK, store->ReadOpLogSince(0, 10, entries)); - ASSERT_EQ(2u, entries.size()); - EXPECT_EQ(1u, entries[0].sequence_id); - EXPECT_EQ(2u, entries[1].sequence_id); -} - -TEST_F(LocalFsOpLogStoreTest, - RewritePersistedSequenceWithDifferentPayloadFails) { - auto store = CreateWriter(); - auto entry = MakeEntry(1); - ASSERT_EQ(ErrorCode::OK, store->WriteOpLog(entry, /*sync=*/true)); - - auto conflicting = entry; - conflicting.payload = "different_payload"; - EXPECT_NE(ErrorCode::OK, store->WriteOpLog(conflicting, /*sync=*/true)); - - OpLogEntry persisted; - ASSERT_EQ(ErrorCode::OK, store->ReadOpLog(1, persisted)); - EXPECT_EQ(entry.payload, persisted.payload); -} - -TEST_F(LocalFsOpLogStoreTest, - RewritePersistedSequenceAfterRestartWithDifferentPayloadFails) { - { - auto store = CreateWriter(); - auto entry = MakeEntry(1); - ASSERT_EQ(ErrorCode::OK, store->WriteOpLog(entry, /*sync=*/true)); - } - - auto reopened = CreateWriter(); - auto conflicting = MakeEntry(1); - conflicting.payload = "payload_after_restart_conflict"; - EXPECT_NE(ErrorCode::OK, reopened->WriteOpLog(conflicting, /*sync=*/true)); - - OpLogEntry persisted; - ASSERT_EQ(ErrorCode::OK, reopened->ReadOpLog(1, persisted)); - EXPECT_EQ("payload_1", persisted.payload); -} - -// --- Read tests --- - -TEST_F(LocalFsOpLogStoreTest, ReadOpLogSingleEntry) { - auto store = CreateWriter(); - store->WriteOpLog(MakeEntry(1), true); - - OpLogEntry entry; - EXPECT_EQ(ErrorCode::OK, store->ReadOpLog(1, entry)); - EXPECT_EQ(1, entry.sequence_id); - EXPECT_EQ("test_key", entry.object_key); -} - -TEST_F(LocalFsOpLogStoreTest, ReadOpLogNotFound) { - auto store = CreateWriter(); - OpLogEntry entry; - EXPECT_EQ(ErrorCode::OPLOG_ENTRY_NOT_FOUND, store->ReadOpLog(999, entry)); -} - -TEST_F(LocalFsOpLogStoreTest, ReadOpLogSinceBasic) { - auto store = CreateWriter(); - for (uint64_t i = 1; i <= 10; i++) { - store->WriteOpLog(MakeEntry(i), /*sync=*/(i == 10)); - } - - std::vector entries; - EXPECT_EQ(ErrorCode::OK, store->ReadOpLogSince(0, 100, entries)); - EXPECT_EQ(10, entries.size()); - EXPECT_EQ(1, entries.front().sequence_id); - EXPECT_EQ(10, entries.back().sequence_id); -} - -TEST_F(LocalFsOpLogStoreTest, ReadOpLogSinceWithOffset) { - auto store = CreateWriter(); - for (uint64_t i = 1; i <= 10; i++) { - store->WriteOpLog(MakeEntry(i), /*sync=*/(i == 10)); - } - - std::vector entries; - EXPECT_EQ(ErrorCode::OK, store->ReadOpLogSince(5, 100, entries)); - EXPECT_EQ(5, entries.size()); - EXPECT_EQ(6, entries.front().sequence_id); -} - -TEST_F(LocalFsOpLogStoreTest, ReadOpLogSinceWithLimit) { - auto store = CreateWriter(); - for (uint64_t i = 1; i <= 10; i++) { - store->WriteOpLog(MakeEntry(i), /*sync=*/(i == 10)); - } - - std::vector entries; - EXPECT_EQ(ErrorCode::OK, store->ReadOpLogSince(0, 3, entries)); - EXPECT_EQ(3, entries.size()); -} - -TEST_F(LocalFsOpLogStoreTest, ReadOpLogSinceAcrossSegments) { - auto store = CreateWriter(); - // Write entries with sync=true to force separate segments - for (uint64_t i = 1; i <= 3; i++) { - store->WriteOpLog(MakeEntry(i), /*sync=*/true); - } - - std::vector entries; - EXPECT_EQ(ErrorCode::OK, store->ReadOpLogSince(0, 100, entries)); - EXPECT_EQ(3, entries.size()); -} - -TEST_F(LocalFsOpLogStoreTest, GetLatestSequenceIdFirstBoot) { - auto store = CreateWriter(); - uint64_t seq = 999; - EXPECT_EQ(ErrorCode::OK, store->GetLatestSequenceId(seq)); - EXPECT_EQ(0, seq); -} - -TEST_F(LocalFsOpLogStoreTest, GetMaxSequenceIdEmpty) { - auto store = CreateWriter(); - uint64_t seq = 0; - EXPECT_EQ(ErrorCode::OPLOG_ENTRY_NOT_FOUND, store->GetMaxSequenceId(seq)); -} - -TEST_F(LocalFsOpLogStoreTest, GetMaxSequenceIdAfterWrite) { - auto store = CreateWriter(); - for (uint64_t i = 1; i <= 5; i++) { - store->WriteOpLog(MakeEntry(i), /*sync=*/(i == 5)); - } - uint64_t seq = 0; - EXPECT_EQ(ErrorCode::OK, store->GetMaxSequenceId(seq)); - EXPECT_EQ(5, seq); -} - -// --- Corrupted segment file tests --- - -TEST_F(LocalFsOpLogStoreTest, ReadCorruptedSegmentBadMagic) { - auto store = CreateWriter(); - store->WriteOpLog(MakeEntry(1), /*sync=*/true); - - // Corrupt the segment file by overwriting magic bytes - auto seg_dir = test_dir_ + "/" + cluster_id_ + "/segments"; - for (auto& p : fs::directory_iterator(seg_dir)) { - if (p.path().extension() != ".tmp") { - std::fstream f(p.path(), - std::ios::in | std::ios::out | std::ios::binary); - f.write("XXXX", 4); // corrupt magic - f.close(); - break; - } - } - - OpLogEntry entry; - EXPECT_EQ(ErrorCode::INTERNAL_ERROR, store->ReadOpLog(1, entry)); -} - -TEST_F(LocalFsOpLogStoreTest, ReadCorruptedSegmentWrongVersion) { - auto store = CreateWriter(); - store->WriteOpLog(MakeEntry(1), /*sync=*/true); - - auto seg_dir = test_dir_ + "/" + cluster_id_ + "/segments"; - for (auto& p : fs::directory_iterator(seg_dir)) { - if (p.path().extension() != ".tmp") { - std::fstream f(p.path(), - std::ios::in | std::ios::out | std::ios::binary); - f.seekp(4); // skip magic, write bad version - uint32_t bad_version = 99; - f.write(reinterpret_cast(&bad_version), - sizeof(bad_version)); - f.close(); - break; - } - } - - OpLogEntry entry; - EXPECT_EQ(ErrorCode::INTERNAL_ERROR, store->ReadOpLog(1, entry)); -} - -TEST_F(LocalFsOpLogStoreTest, ReadCorruptedSegmentTruncated) { - auto store = CreateWriter(); - store->WriteOpLog(MakeEntry(1), /*sync=*/true); - - auto seg_dir = test_dir_ + "/" + cluster_id_ + "/segments"; - for (auto& p : fs::directory_iterator(seg_dir)) { - if (p.path().extension() != ".tmp") { - // Truncate to just 16 bytes (less than header size) - fs::resize_file(p.path(), 16); - break; - } - } - - OpLogEntry entry; - EXPECT_EQ(ErrorCode::INTERNAL_ERROR, store->ReadOpLog(1, entry)); -} - -// --- Snapshot tests --- - -TEST_F(LocalFsOpLogStoreTest, SnapshotRecordAndGet) { - auto store = CreateWriter(); - EXPECT_EQ(ErrorCode::OK, store->RecordSnapshotSequenceId("snap1", 42)); - - uint64_t seq = 0; - EXPECT_EQ(ErrorCode::OK, store->GetSnapshotSequenceId("snap1", seq)); - EXPECT_EQ(42, seq); -} - -TEST_F(LocalFsOpLogStoreTest, SnapshotNotFound) { - auto store = CreateWriter(); - uint64_t seq = 0; - EXPECT_NE(ErrorCode::OK, store->GetSnapshotSequenceId("nonexistent", seq)); -} - -TEST_F(LocalFsOpLogStoreTest, SnapshotIdValidation) { - auto store = CreateWriter(); - EXPECT_EQ(ErrorCode::INVALID_PARAMS, - store->RecordSnapshotSequenceId("../escape", 1)); - EXPECT_EQ(ErrorCode::INVALID_PARAMS, - store->RecordSnapshotSequenceId("path/slash", 1)); - EXPECT_EQ(ErrorCode::INVALID_PARAMS, - store->RecordSnapshotSequenceId(std::string("null\0byte", 9), 1)); -} - -// --- Cleanup tests --- - -TEST_F(LocalFsOpLogStoreTest, CleanupRemovesOldSegments) { - auto store = CreateWriter(); - // Write entries in separate segments (sync=true each time) - for (uint64_t i = 1; i <= 5; i++) { - store->WriteOpLog(MakeEntry(i), /*sync=*/true); - } - - // Count segments before cleanup - auto seg_dir = test_dir_ + "/" + cluster_id_ + "/segments"; - int before_count = 0; - for (auto& p : fs::directory_iterator(seg_dir)) { - (void)p; - before_count++; - } - EXPECT_EQ(5, before_count); - - // Cleanup entries before seq 4 (should remove segments with max_seq < 4) - EXPECT_EQ(ErrorCode::OK, store->CleanupOpLogBefore(4)); - - // Entries 4 and 5 should still be readable - OpLogEntry entry; - EXPECT_EQ(ErrorCode::OK, store->ReadOpLog(4, entry)); - EXPECT_EQ(ErrorCode::OK, store->ReadOpLog(5, entry)); -} - -TEST_F(LocalFsOpLogStoreTest, CleanupEmptyDir) { - auto store = CreateWriter(); - EXPECT_EQ(ErrorCode::OK, store->CleanupOpLogBefore(100)); -} - -// --- ChangeNotifier tests --- - -TEST_F(LocalFsOpLogStoreTest, ChangeNotifierBasic) { - auto writer = CreateWriter(); - auto reader = CreateReader(); - - auto notifier = reader->CreateChangeNotifier(cluster_id_); - ASSERT_NE(nullptr, notifier); - - std::vector received; - std::mutex received_mutex; - std::condition_variable received_cv; - - auto on_entry = [&](const OpLogEntry& entry) { - std::lock_guard lock(received_mutex); - received.push_back(entry); - received_cv.notify_one(); - }; - auto on_error = [](ErrorCode) {}; - - EXPECT_EQ(ErrorCode::OK, notifier->Start(0, on_entry, on_error)); - EXPECT_TRUE(notifier->IsHealthy()); - - // Write an entry from writer - writer->WriteOpLog(MakeEntry(1), /*sync=*/true); - - // Wait for notifier to deliver it (with timeout) - { - std::unique_lock lock(received_mutex); - EXPECT_TRUE(received_cv.wait_for(lock, std::chrono::seconds(5), - [&] { return received.size() >= 1; })); - } - EXPECT_EQ(1, received.size()); - EXPECT_EQ(1, received[0].sequence_id); - - notifier->Stop(); -} - -TEST_F(LocalFsOpLogStoreTest, ChangeNotifierCatchUp) { - auto writer = CreateWriter(); - // Write entries BEFORE starting notifier - for (uint64_t i = 1; i <= 5; i++) { - writer->WriteOpLog(MakeEntry(i), /*sync=*/(i == 5)); - } - - auto reader = CreateReader(); - auto notifier = reader->CreateChangeNotifier(cluster_id_); - ASSERT_NE(nullptr, notifier); - - std::vector received; - std::mutex received_mutex; - std::condition_variable received_cv; - - auto on_entry = [&](const OpLogEntry& entry) { - std::lock_guard lock(received_mutex); - received.push_back(entry); - received_cv.notify_one(); - }; - auto on_error = [](ErrorCode) {}; - - // Start from 0 should catch up with existing entries - EXPECT_EQ(ErrorCode::OK, notifier->Start(0, on_entry, on_error)); - - { - std::unique_lock lock(received_mutex); - EXPECT_TRUE(received_cv.wait_for(lock, std::chrono::seconds(5), - [&] { return received.size() >= 5; })); - } - EXPECT_EQ(5, received.size()); - EXPECT_EQ(1, received[0].sequence_id); - EXPECT_EQ(5, received[4].sequence_id); - - notifier->Stop(); -} - -TEST_F(LocalFsOpLogStoreTest, ChangeNotifierStopNoMoreCallbacks) { - auto writer = CreateWriter(); - auto reader = CreateReader(); - - auto notifier = reader->CreateChangeNotifier(cluster_id_); - ASSERT_NE(nullptr, notifier); - - std::atomic callback_count{0}; - auto on_entry = [&](const OpLogEntry&) { callback_count++; }; - auto on_error = [](ErrorCode) {}; - - EXPECT_EQ(ErrorCode::OK, notifier->Start(0, on_entry, on_error)); - notifier->Stop(); - - // Write after stop - should not trigger callback - writer->WriteOpLog(MakeEntry(1), /*sync=*/true); - std::this_thread::sleep_for(std::chrono::milliseconds(200)); - EXPECT_EQ(0, callback_count.load()); -} - -// --- Factory tests --- - -TEST_F(LocalFsOpLogStoreTest, FactoryCreateWriter) { - auto store = - OpLogStoreFactory::Create(OpLogStoreType::LOCAL_FS, cluster_id_, - OpLogStoreRole::WRITER, test_dir_, 100); - ASSERT_NE(nullptr, store); - - // Should be able to write - OpLogEntry entry; - entry.sequence_id = 1; - entry.op_type = OpType::PUT_END; - entry.object_key = "factory_test"; - entry.payload = "data"; - EXPECT_EQ(ErrorCode::OK, store->WriteOpLog(entry, /*sync=*/true)); - - // Should be able to read back - OpLogEntry read_entry; - EXPECT_EQ(ErrorCode::OK, store->ReadOpLog(1, read_entry)); - EXPECT_EQ(1, read_entry.sequence_id); -} - -TEST_F(LocalFsOpLogStoreTest, FactoryCreateReader) { - // First create some data with a writer - { - auto writer = - OpLogStoreFactory::Create(OpLogStoreType::LOCAL_FS, cluster_id_, - OpLogStoreRole::WRITER, test_dir_, 100); - ASSERT_NE(nullptr, writer); - OpLogEntry entry; - entry.sequence_id = 1; - entry.op_type = OpType::PUT_END; - entry.object_key = "key"; - entry.payload = "data"; - writer->WriteOpLog(entry, /*sync=*/true); - } - - // Reader should be able to read but not write - auto reader = - OpLogStoreFactory::Create(OpLogStoreType::LOCAL_FS, cluster_id_, - OpLogStoreRole::READER, test_dir_, 100); - ASSERT_NE(nullptr, reader); - - OpLogEntry read_entry; - EXPECT_EQ(ErrorCode::OK, reader->ReadOpLog(1, read_entry)); - EXPECT_EQ(1, read_entry.sequence_id); - - // Writer should fail on reader - OpLogEntry new_entry; - new_entry.sequence_id = 2; - new_entry.op_type = OpType::PUT_END; - new_entry.object_key = "key2"; - new_entry.payload = "data2"; - EXPECT_EQ(ErrorCode::INVALID_PARAMS, - reader->WriteOpLog(new_entry, /*sync=*/true)); -} - -} // namespace -} // namespace mooncake diff --git a/mooncake-store/tests/ha/oplog/mock_metadata_store.h b/mooncake-store/tests/ha/oplog/mock_metadata_store.h index 3c1b9a5697..1b8ef4d190 100644 --- a/mooncake-store/tests/ha/oplog/mock_metadata_store.h +++ b/mooncake-store/tests/ha/oplog/mock_metadata_store.h @@ -6,6 +6,7 @@ #include #include "metadata_store.h" +#include "types.h" namespace mooncake::test { @@ -16,54 +17,98 @@ class MockMetadataStore : public MetadataStore { MockMetadataStore() = default; ~MockMetadataStore() override = default; - bool PutMetadata(const std::string& key, - const StandbyObjectMetadata& metadata) override { - metadata_map_[key] = metadata; - return true; - } + // Bring base-class key-only overloads into scope (name hiding) + using MetadataStore::Exists; + using MetadataStore::GetMetadata; + using MetadataStore::PutMetadata; + using MetadataStore::Remove; - bool Put(const std::string& key, const std::string& payload) override { - // For testing, we can use PutMetadata with empty metadata - StandbyObjectMetadata meta; - metadata_map_[key] = meta; + // Tenant-aware methods (primary API) + bool PutMetadata(const std::string& tenant_id, const std::string& key, + const StandbyObjectMetadata& metadata) override { + const auto normalized = NormalizeTenantId(tenant_id); + metadata_map_[normalized][key] = metadata; return true; } std::optional GetMetadata( - const std::string& key) const override { - auto it = metadata_map_.find(key); - if (it != metadata_map_.end()) { + const std::string& tenant_id, const std::string& key) const override { + const auto normalized = NormalizeTenantId(tenant_id); + auto tenant_it = metadata_map_.find(normalized); + if (tenant_it == metadata_map_.end()) { + return std::nullopt; + } + auto it = tenant_it->second.find(key); + if (it != tenant_it->second.end()) { return it->second; } return std::nullopt; } - bool Remove(const std::string& key) override { - auto it = metadata_map_.find(key); - if (it != metadata_map_.end()) { - metadata_map_.erase(it); + bool Remove(const std::string& tenant_id, const std::string& key) override { + const auto normalized = NormalizeTenantId(tenant_id); + auto tenant_it = metadata_map_.find(normalized); + if (tenant_it == metadata_map_.end()) { + return false; + } + auto it = tenant_it->second.find(key); + if (it != tenant_it->second.end()) { + tenant_it->second.erase(it); + if (tenant_it->second.empty()) { + metadata_map_.erase(tenant_it); + } return true; } return false; } - bool Exists(const std::string& key) const override { - return metadata_map_.find(key) != metadata_map_.end(); + bool Exists(const std::string& tenant_id, + const std::string& key) const override { + const auto normalized = NormalizeTenantId(tenant_id); + auto tenant_it = metadata_map_.find(normalized); + if (tenant_it == metadata_map_.end()) { + return false; + } + return tenant_it->second.find(key) != tenant_it->second.end(); } - size_t GetKeyCount() const override { return metadata_map_.size(); } + size_t GetKeyCountForTenant(const std::string& tenant_id) const override { + const auto normalized = NormalizeTenantId(tenant_id); + auto tenant_it = metadata_map_.find(normalized); + if (tenant_it == metadata_map_.end()) { + return 0; + } + return tenant_it->second.size(); + } + + // Legacy key-only Put delegates to "default" tenant + bool Put(const std::string& key, const std::string& payload) override { + StandbyObjectMetadata meta; + metadata_map_["default"][key] = meta; + return true; + } + + // Total count across ALL tenants + size_t GetKeyCount() const override { + size_t total = 0; + for (const auto& [tenant_id, tenant_map] : metadata_map_) { + total += tenant_map.size(); + } + return total; + } // Test helper methods void Clear() { metadata_map_.clear(); } - size_t Size() const { return metadata_map_.size(); } + size_t Size() const { return GetKeyCount(); } bool Contains(const std::string& key) const { - return metadata_map_.count(key) > 0; + return Exists("default", key); } private: - std::map metadata_map_; + std::map> + metadata_map_; }; } // namespace mooncake::test diff --git a/mooncake-store/tests/ha/oplog/mock_oplog_store.h b/mooncake-store/tests/ha/oplog/mock_oplog_store.h deleted file mode 100644 index 2e5da68ee2..0000000000 --- a/mooncake-store/tests/ha/oplog/mock_oplog_store.h +++ /dev/null @@ -1,139 +0,0 @@ -// mooncake-store/tests/hot_standby_ut/mock_oplog_store.h -#pragma once - -#include -#include -#include -#include - -#include "ha/oplog/oplog_manager.h" -#include "ha/oplog/oplog_store.h" -#include "types.h" - -namespace mooncake::test { - -// In-memory OpLog store for unit tests. -class MockOpLogStore : public OpLogStore { - public: - MockOpLogStore() = default; - - ErrorCode Init() override { return ErrorCode::OK; } - - ErrorCode WriteOpLog(const OpLogEntry& entry, - bool /*sync*/ = true) override { - std::lock_guard lock(mutex_); - if (write_error_ != ErrorCode::OK) { - return write_error_; - } - entries_[entry.sequence_id] = entry; - if (entry.sequence_id > latest_seq_id_) { - latest_seq_id_ = entry.sequence_id; - } - return ErrorCode::OK; - } - - ErrorCode ReadOpLog(uint64_t sequence_id, OpLogEntry& entry) override { - std::lock_guard lock(mutex_); - if (read_error_ != ErrorCode::OK) { - return read_error_; - } - auto it = entries_.find(sequence_id); - if (it == entries_.end()) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - entry = it->second; - return ErrorCode::OK; - } - - ErrorCode ReadOpLogSince(uint64_t start_sequence_id, size_t limit, - std::vector& entries) override { - std::lock_guard lock(mutex_); - if (read_error_ != ErrorCode::OK) { - return read_error_; - } - entries.clear(); - for (const auto& [seq, entry] : entries_) { - if (seq > start_sequence_id) { - entries.push_back(entry); - if (entries.size() >= limit) break; - } - } - return ErrorCode::OK; - } - - ErrorCode GetLatestSequenceId(uint64_t& sequence_id) override { - std::lock_guard lock(mutex_); - sequence_id = latest_seq_id_; - return ErrorCode::OK; - } - - ErrorCode GetMaxSequenceId(uint64_t& sequence_id) override { - std::lock_guard lock(mutex_); - if (entries_.empty()) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - sequence_id = entries_.rbegin()->first; - return ErrorCode::OK; - } - - ErrorCode UpdateLatestSequenceId(uint64_t sequence_id) override { - std::lock_guard lock(mutex_); - latest_seq_id_ = sequence_id; - return ErrorCode::OK; - } - - ErrorCode RecordSnapshotSequenceId(const std::string& snapshot_id, - uint64_t sequence_id) override { - std::lock_guard lock(mutex_); - snapshots_[snapshot_id] = sequence_id; - return ErrorCode::OK; - } - - ErrorCode GetSnapshotSequenceId(const std::string& snapshot_id, - uint64_t& sequence_id) override { - std::lock_guard lock(mutex_); - auto it = snapshots_.find(snapshot_id); - if (it == snapshots_.end()) { - return ErrorCode::OPLOG_ENTRY_NOT_FOUND; - } - sequence_id = it->second; - return ErrorCode::OK; - } - - ErrorCode CleanupOpLogBefore(uint64_t before_sequence_id) override { - std::lock_guard lock(mutex_); - for (auto it = entries_.begin(); it != entries_.end();) { - if (it->first < before_sequence_id) { - it = entries_.erase(it); - } else { - break; - } - } - return ErrorCode::OK; - } - - // === Test control methods === - - void SetWriteError(ErrorCode err) { write_error_ = err; } - void SetReadError(ErrorCode err) { read_error_ = err; } - void Clear() { - std::lock_guard lock(mutex_); - entries_.clear(); - snapshots_.clear(); - latest_seq_id_ = 0; - } - size_t EntryCount() const { - std::lock_guard lock(mutex_); - return entries_.size(); - } - - private: - mutable std::mutex mutex_; - std::map entries_; - std::map snapshots_; - uint64_t latest_seq_id_{0}; - ErrorCode write_error_{ErrorCode::OK}; - ErrorCode read_error_{ErrorCode::OK}; -}; - -} // namespace mooncake::test diff --git a/mooncake-store/tests/ha/oplog/mock_snapshot_provider.h b/mooncake-store/tests/ha/oplog/mock_snapshot_provider.h index ff0256002d..a827bb21e6 100644 --- a/mooncake-store/tests/ha/oplog/mock_snapshot_provider.h +++ b/mooncake-store/tests/ha/oplog/mock_snapshot_provider.h @@ -4,10 +4,10 @@ #include #include #include -#include #include #include "ha/snapshot/snapshot_provider.h" +#include "metadata_store.h" namespace mooncake::test { @@ -15,9 +15,8 @@ namespace mooncake::test { // Allows pre-filling snapshot data and simulating load failures. class MockSnapshotProvider : public SnapshotProvider { public: - void SetSnapshot( - const std::string& snap_id, uint64_t seq_id, - std::vector> data) { + void SetSnapshot(const std::string& snap_id, uint64_t seq_id, + std::vector data) { snapshot_id_ = snap_id; snapshot_seq_id_ = seq_id; snapshot_data_ = std::move(data); @@ -43,7 +42,7 @@ class MockSnapshotProvider : public SnapshotProvider { private: std::string snapshot_id_; uint64_t snapshot_seq_id_ = 0; - std::vector> snapshot_data_; + std::vector snapshot_data_; bool should_fail_ = false; }; diff --git a/mooncake-store/tests/ha/oplog/oplog_applier_test.cpp b/mooncake-store/tests/ha/oplog/oplog_applier_test.cpp index b9fd0cbd5a..9bbc975230 100644 --- a/mooncake-store/tests/ha/oplog/oplog_applier_test.cpp +++ b/mooncake-store/tests/ha/oplog/oplog_applier_test.cpp @@ -4,27 +4,23 @@ #include #include -#include #include #include -#include #include #include #include "metadata_store.h" #include "mock_metadata_store.h" -#include "mock_oplog_store.h" -#include "ha/oplog/oplog_manager.h" +#include "ha/oplog/oplog_types.h" #include "types.h" using mooncake::test::MockMetadataStore; -using mooncake::test::MockOpLogStore; namespace mooncake::test { // Helper function to create a valid OpLogEntry with checksum -// Uses the same checksum algorithm as OpLogManager (XXH32) +// Uses the production checksum algorithm (XXH32). OpLogEntry MakeEntry(uint64_t seq, OpType type, const std::string& key, const std::string& payload) { OpLogEntry e; @@ -35,7 +31,7 @@ OpLogEntry MakeEntry(uint64_t seq, OpType type, const std::string& key, e.op_type = type; e.object_key = key; e.payload = payload; - // Compute checksum and prefix_hash using the same algorithm as OpLogManager + // Compute checksum and prefix_hash using the production algorithm. e.checksum = static_cast(XXH32(payload.data(), payload.size(), 0)); e.prefix_hash = @@ -144,56 +140,28 @@ TEST_F(OpLogApplierTest, TestApplyInOrder) { EXPECT_TRUE(mock_metadata_store_->Exists("key3")); } -TEST_F(OpLogApplierTest, TestApplyOutOfOrder) { +TEST_F(OpLogApplierTest, FutureSequenceFailsWithoutBuffering) { std::string payload = MakeValidPayload(); OpLogEntry entry1 = MakeEntry(1, OpType::PUT_END, "key1", payload); - OpLogEntry entry3 = MakeEntry(3, OpType::PUT_END, "key3", payload); OpLogEntry entry2 = MakeEntry(2, OpType::PUT_END, "key2", payload); + OpLogEntry entry3 = MakeEntry(3, OpType::PUT_END, "key3", payload); // Apply entry1 (seq=1) - should succeed EXPECT_TRUE(applier_->ApplyOpLogEntry(entry1)); EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); - // Apply entry3 (seq=3) - should be cached (out of order) + // A future entry fails and is not retained for later application. EXPECT_FALSE(applier_->ApplyOpLogEntry(entry3)); - EXPECT_EQ(2u, - applier_->GetExpectedSequenceId()); // Still waiting for seq=2 + EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); EXPECT_FALSE(mock_metadata_store_->Exists("key3")); - // Apply entry2 (seq=2) - should succeed and trigger processing of entry3 - // ApplyOpLogEntry internally calls ProcessPendingEntries(), so entry3 - // should be processed + // Applying the expected entry does not resurrect the rejected future one. EXPECT_TRUE(applier_->ApplyOpLogEntry(entry2)); - EXPECT_EQ(4u, applier_->GetExpectedSequenceId()); // Now at seq=4 + EXPECT_EQ(3u, applier_->GetExpectedSequenceId()); - // entry3 should already be processed by ApplyOpLogEntry EXPECT_TRUE(mock_metadata_store_->Exists("key1")); EXPECT_TRUE(mock_metadata_store_->Exists("key2")); - EXPECT_TRUE(mock_metadata_store_->Exists("key3")); - - // ProcessPendingEntries may return 0 if entry3 was already processed - (void)applier_->ProcessPendingEntries(); - EXPECT_EQ(4u, applier_->GetExpectedSequenceId()); -} - -TEST_F(OpLogApplierTest, TestApplyWithGap) { - std::string payload = MakeValidPayload(); - OpLogEntry entry1 = MakeEntry(1, OpType::PUT_END, "key1", payload); - OpLogEntry entry4 = MakeEntry(4, OpType::PUT_END, "key4", payload); - - // Apply entry1 (seq=1) - EXPECT_TRUE(applier_->ApplyOpLogEntry(entry1)); - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); - - // Apply entry4 (seq=4) - gap at seq=2,3 - EXPECT_FALSE(applier_->ApplyOpLogEntry(entry4)); - EXPECT_EQ(2u, - applier_->GetExpectedSequenceId()); // Still waiting for seq=2 - - // Process pending entries - should detect gap and schedule wait - (void)applier_->ProcessPendingEntries(); - // May process 0 entries if gap resolution is still waiting - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); + EXPECT_FALSE(mock_metadata_store_->Exists("key3")); } TEST_F(OpLogApplierTest, TestApplyDuplicateSequenceId) { @@ -213,93 +181,6 @@ TEST_F(OpLogApplierTest, TestApplyDuplicateSequenceId) { EXPECT_FALSE(mock_metadata_store_->Exists("key1_dup")); } -// ========== 4.1.3 Gap resolution tests ========== - -class OpLogApplierGapTest : public ::testing::Test { - protected: - void SetUp() override { - google::InitGoogleLogging("OpLogApplierGapTest"); - FLAGS_logtostderr = 1; - mock_metadata_store_ = std::make_unique(); - mock_oplog_store_ = std::make_unique(); - applier_ = std::make_unique(mock_metadata_store_.get(), - "test_cluster", - mock_oplog_store_.get()); - } - void TearDown() override { google::ShutdownGoogleLogging(); } - - std::unique_ptr mock_metadata_store_; - std::unique_ptr mock_oplog_store_; - std::unique_ptr applier_; -}; - -TEST_F(OpLogApplierGapTest, RequestMissingOpLog_Success) { - std::string payload = MakeValidPayload(); - // Pre-populate seq=2 in mock store - OpLogEntry missing = MakeEntry(2, OpType::PUT_END, "key2", payload); - mock_oplog_store_->WriteOpLog(missing); - - // Apply seq=1 - EXPECT_TRUE(applier_->ApplyOpLogEntry( - MakeEntry(1, OpType::PUT_END, "key1", payload))); - // Apply seq=3 (gap at seq=2) - EXPECT_FALSE(applier_->ApplyOpLogEntry( - MakeEntry(3, OpType::PUT_END, "key3", payload))); - - // First call registers the gap in missing_sequence_ids_ - applier_->ProcessPendingEntries(); - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); - - // Wait for gap resolution trigger (kMissingEntryRequestSeconds = 1s) - std::this_thread::sleep_for(std::chrono::milliseconds(1500)); - applier_->ProcessPendingEntries(); - - // After gap resolution, all 3 should be applied - EXPECT_EQ(4u, applier_->GetExpectedSequenceId()); - EXPECT_TRUE(mock_metadata_store_->Exists("key1")); - EXPECT_TRUE(mock_metadata_store_->Exists("key2")); - EXPECT_TRUE(mock_metadata_store_->Exists("key3")); -} - -TEST_F(OpLogApplierGapTest, RequestMissingOpLog_StoreError) { - std::string payload = MakeValidPayload(); - mock_oplog_store_->SetReadError(ErrorCode::ETCD_OPERATION_ERROR); - - EXPECT_TRUE(applier_->ApplyOpLogEntry( - MakeEntry(1, OpType::PUT_END, "key1", payload))); - EXPECT_FALSE(applier_->ApplyOpLogEntry( - MakeEntry(3, OpType::PUT_END, "key3", payload))); - - // First call registers the gap - applier_->ProcessPendingEntries(); - - // Gap resolution should fail gracefully - std::this_thread::sleep_for(std::chrono::milliseconds(1500)); - applier_->ProcessPendingEntries(); - - // seq=2 not resolved, expected still at 2 - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); -} - -TEST_F(OpLogApplierGapTest, RequestMissingOpLog_NotFound) { - std::string payload = MakeValidPayload(); - // Don't preload seq=2 in mock store - - EXPECT_TRUE(applier_->ApplyOpLogEntry( - MakeEntry(1, OpType::PUT_END, "key1", payload))); - EXPECT_FALSE(applier_->ApplyOpLogEntry( - MakeEntry(3, OpType::PUT_END, "key3", payload))); - - // First call registers the gap - applier_->ProcessPendingEntries(); - - std::this_thread::sleep_for(std::chrono::milliseconds(1500)); - applier_->ProcessPendingEntries(); - - // Not found, should not advance - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); -} - // ========== 4.1.4 Checksum tests ========== TEST_F(OpLogApplierTest, TestApplyOpLogEntry_ValidChecksum) { @@ -340,7 +221,7 @@ TEST_F(OpLogApplierTest, TestApplyOpLogEntry_ValidSize) { std::string payload = MakeValidPayload(); OpLogEntry entry = MakeEntry(1, OpType::PUT_END, "key1", payload); - EXPECT_TRUE(OpLogManager::ValidateEntrySize(entry)); + EXPECT_TRUE(ValidateOpLogEntrySize(entry)); EXPECT_TRUE(applier_->ApplyOpLogEntry(entry)); } @@ -348,9 +229,9 @@ TEST_F(OpLogApplierTest, TestApplyOpLogEntry_InvalidSize) { OpLogEntry entry = MakeEntry(1, OpType::PUT_END, "key1", ""); // Make key too large - entry.object_key.assign(OpLogManager::kMaxObjectKeySize + 1, 'k'); + entry.object_key.assign(kMaxOpLogObjectKeySize + 1, 'k'); - EXPECT_FALSE(OpLogManager::ValidateEntrySize(entry)); + EXPECT_FALSE(ValidateOpLogEntrySize(entry)); EXPECT_FALSE(applier_->ApplyOpLogEntry(entry)); EXPECT_EQ(1u, applier_->GetExpectedSequenceId()); // Should not advance EXPECT_FALSE(mock_metadata_store_->Exists("key1")); @@ -360,9 +241,9 @@ TEST_F(OpLogApplierTest, TestApplyOpLogEntry_PayloadTooLarge) { OpLogEntry entry = MakeEntry(1, OpType::PUT_END, "key1", ""); // Make payload too large - entry.payload.assign(OpLogManager::kMaxPayloadSize + 1, 'p'); + entry.payload.assign(kMaxOpLogPayloadSize + 1, 'p'); - EXPECT_FALSE(OpLogManager::ValidateEntrySize(entry)); + EXPECT_FALSE(ValidateOpLogEntrySize(entry)); EXPECT_FALSE(applier_->ApplyOpLogEntry(entry)); EXPECT_EQ(1u, applier_->GetExpectedSequenceId()); // Should not advance } @@ -393,123 +274,6 @@ TEST_F(OpLogApplierTest, TestRecover_ZeroSequenceId) { EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); } -TEST_F(OpLogApplierTest, TestRecover_AfterGap) { - std::string payload = MakeValidPayload(); - OpLogEntry entry1 = MakeEntry(1, OpType::PUT_END, "key1", payload); - OpLogEntry entry3 = MakeEntry(3, OpType::PUT_END, "key3", payload); - - // Apply entry1 - EXPECT_TRUE(applier_->ApplyOpLogEntry(entry1)); - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); - - // Apply entry3 (creates gap) - EXPECT_FALSE(applier_->ApplyOpLogEntry(entry3)); - - // Recover from seq=3 (skip the gap) - applier_->Recover(3); - EXPECT_EQ(4u, applier_->GetExpectedSequenceId()); - - // Now entry3 should be processable - (void)applier_->ProcessPendingEntries(); - // entry3 should be in pending, but expected_seq is now 4, so it won't be - // processed This tests that recovery resets the expected sequence -} - -// ========== 4.1.7 Pending entries tests ========== - -TEST_F(OpLogApplierTest, TestProcessPendingEntries) { - std::string payload = MakeValidPayload(); - OpLogEntry entry1 = MakeEntry(1, OpType::PUT_END, "key1", payload); - OpLogEntry entry3 = MakeEntry(3, OpType::PUT_END, "key3", payload); - OpLogEntry entry2 = MakeEntry(2, OpType::PUT_END, "key2", payload); - - // Apply entry1 - EXPECT_TRUE(applier_->ApplyOpLogEntry(entry1)); - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); - - // Apply entry3 (out of order) - EXPECT_FALSE(applier_->ApplyOpLogEntry(entry3)); - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); - - // Process pending - should not process entry3 yet (waiting for seq=2) - size_t processed1 = applier_->ProcessPendingEntries(); - EXPECT_EQ(0u, processed1); - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); - - // Apply entry2 - this will internally call ProcessPendingEntries() and - // process entry3 - EXPECT_TRUE(applier_->ApplyOpLogEntry(entry2)); - EXPECT_EQ(4u, applier_->GetExpectedSequenceId()); - - // entry3 should already be processed by ApplyOpLogEntry - EXPECT_TRUE(mock_metadata_store_->Exists("key1")); - EXPECT_TRUE(mock_metadata_store_->Exists("key2")); - EXPECT_TRUE(mock_metadata_store_->Exists("key3")); - - // ProcessPendingEntries may return 0 if entry3 was already processed - (void)applier_->ProcessPendingEntries(); - EXPECT_EQ(4u, applier_->GetExpectedSequenceId()); -} - -TEST_F(OpLogApplierTest, TestPendingEntriesTimeout) { - std::string payload = MakeValidPayload(); - OpLogEntry entry1 = MakeEntry(1, OpType::PUT_END, "key1", payload); - OpLogEntry entry3 = MakeEntry(3, OpType::PUT_END, "key3", payload); - - // Apply entry1 - EXPECT_TRUE(applier_->ApplyOpLogEntry(entry1)); - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); - - // Apply entry3 (creates gap at seq=2) - EXPECT_FALSE(applier_->ApplyOpLogEntry(entry3)); - - // Process pending entries multiple times to trigger timeout - // After kMissingEntrySkipSeconds (3s), the gap should be skipped - size_t processed = 0; - for (int i = 0; i < 10; ++i) { - processed = applier_->ProcessPendingEntries(); - if (processed > 0 || applier_->GetExpectedSequenceId() > 2) { - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(500)); - } - - // After timeout, gap should be skipped and entry3 should be processed - // Note: This test may be flaky due to timing, but it tests the timeout - // logic - EXPECT_GE(applier_->GetExpectedSequenceId(), 2u); -} - -TEST_F(OpLogApplierTest, TestPendingEntriesSkip) { - std::string payload = MakeValidPayload(); - OpLogEntry entry1 = MakeEntry(1, OpType::PUT_END, "key1", payload); - OpLogEntry entry4 = MakeEntry(4, OpType::PUT_END, "key4", payload); - - // Apply entry1 - EXPECT_TRUE(applier_->ApplyOpLogEntry(entry1)); - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); - - // Apply entry4 (creates gap at seq=2,3) - EXPECT_FALSE(applier_->ApplyOpLogEntry(entry4)); - - // Process pending entries to trigger skip logic - // After timeout (3 seconds), gaps should be skipped - for (int i = 0; i < 10; ++i) { - applier_->ProcessPendingEntries(); - uint64_t expected = applier_->GetExpectedSequenceId(); - if (expected >= 3) { // Gap at seq=2 is skipped, expected becomes 3 - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(500)); - } - - // After skip, expected_seq should advance to 3 (gap at seq=2 is skipped) - // entry4 is still pending, waiting for seq=3 - EXPECT_GE(applier_->GetExpectedSequenceId(), 3u); - // entry4 should still be pending (not applied yet) - EXPECT_FALSE(mock_metadata_store_->Exists("key4")); -} - // ========== 4.1.8 Payload deserialization tests ========== TEST_F(OpLogApplierTest, TestApplyPutEnd_ValidPayload) { @@ -583,21 +347,11 @@ TEST_F(OpLogApplierTest, TestApplyOpLogEntries_WithGaps) { entries.push_back(MakeEntry(2, OpType::PUT_END, "key2", payload)); size_t applied = applier_->ApplyOpLogEntries(entries); - // entry1 should be applied, entry3 should be pending, entry2 should be - // applied and trigger processing of entry3 - EXPECT_GE(applied, 2u); // entry1 and entry2 are applied - EXPECT_LE(applied, 3u); - - // entry2's ApplyOpLogEntry internally calls ProcessPendingEntries(), so - // entry3 should be processed - EXPECT_GE(applier_->GetExpectedSequenceId(), 4u); + EXPECT_EQ(2u, applied); + EXPECT_EQ(3u, applier_->GetExpectedSequenceId()); EXPECT_TRUE(mock_metadata_store_->Exists("key1")); EXPECT_TRUE(mock_metadata_store_->Exists("key2")); - EXPECT_TRUE(mock_metadata_store_->Exists("key3")); - - // ProcessPendingEntries may return 0 if entry3 was already processed - (void)applier_->ProcessPendingEntries(); - EXPECT_GE(applier_->GetExpectedSequenceId(), 4u); + EXPECT_FALSE(mock_metadata_store_->Exists("key3")); } TEST_F(OpLogApplierTest, TestGetExpectedSequenceId) { @@ -609,6 +363,168 @@ TEST_F(OpLogApplierTest, TestGetExpectedSequenceId) { EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); } +// ========== Segment OpLog apply tests ========== + +std::string MakeSegmentMountPayload(const std::string& segment_name, + const std::string& transport_endpoint, + uint64_t capacity = 1024, + bool is_memory = true, + const std::string& file_path = "") { + SegmentMountOp op; + op.segment_name = segment_name; + op.transport_endpoint = transport_endpoint; + op.capacity = capacity; + op.is_memory_segment = is_memory; + op.file_path = file_path; + auto result = struct_pack::serialize(op); + return std::string(result.begin(), result.end()); +} + +std::string MakeSegmentUnmountPayload(const std::string& transport_endpoint) { + SegmentUnmountOp op; + op.transport_endpoint = transport_endpoint; + auto result = struct_pack::serialize(op); + return std::string(result.begin(), result.end()); +} + +std::string MakeSegmentUpdatePayload(const std::string& segment_name, + const std::string& transport_endpoint, + uint64_t capacity = 2048, + bool is_memory = true, + const std::string& file_path = "") { + SegmentUpdateOp op; + op.segment_name = segment_name; + op.transport_endpoint = transport_endpoint; + op.capacity = capacity; + op.is_memory_segment = is_memory; + op.file_path = file_path; + auto result = struct_pack::serialize(op); + return std::string(result.begin(), result.end()); +} + +TEST_F(OpLogApplierTest, TestApplySegmentMount) { + std::string payload = + MakeSegmentMountPayload("seg1", "192.168.1.1:12345", 1024, true); + OpLogEntry entry = MakeEntry(1, OpType::SEGMENT_MOUNT, "seg1", payload); + + EXPECT_TRUE(applier_->ApplyOpLogEntry(entry)); + EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); + + const auto& registry = applier_->GetSegmentRegistry(); + EXPECT_TRUE(registry.HasSegment("192.168.1.1:12345")); + auto info = registry.GetSegment("192.168.1.1:12345"); + ASSERT_TRUE(info.has_value()); + EXPECT_EQ("seg1", info->segment_name); + EXPECT_EQ(1024u, info->capacity); + EXPECT_TRUE(info->is_memory_segment); +} + +TEST_F(OpLogApplierTest, TestApplySegmentUnmount) { + // Mount first + std::string mount_payload = + MakeSegmentMountPayload("seg1", "192.168.1.1:12345", 1024, true); + OpLogEntry mount_entry = + MakeEntry(1, OpType::SEGMENT_MOUNT, "seg1", mount_payload); + EXPECT_TRUE(applier_->ApplyOpLogEntry(mount_entry)); + EXPECT_TRUE(applier_->GetSegmentRegistry().HasSegment("192.168.1.1:12345")); + + // Then unmount + std::string unmount_payload = + MakeSegmentUnmountPayload("192.168.1.1:12345"); + OpLogEntry unmount_entry = + MakeEntry(2, OpType::SEGMENT_UNMOUNT, "seg1", unmount_payload); + EXPECT_TRUE(applier_->ApplyOpLogEntry(unmount_entry)); + EXPECT_EQ(3u, applier_->GetExpectedSequenceId()); + + EXPECT_FALSE( + applier_->GetSegmentRegistry().HasSegment("192.168.1.1:12345")); +} + +TEST_F(OpLogApplierTest, TestApplySegmentUpdate) { + // Mount first + std::string mount_payload = + MakeSegmentMountPayload("seg1", "192.168.1.1:12345", 1024, true); + OpLogEntry mount_entry = + MakeEntry(1, OpType::SEGMENT_MOUNT, "seg1", mount_payload); + EXPECT_TRUE(applier_->ApplyOpLogEntry(mount_entry)); + + auto info_before = + applier_->GetSegmentRegistry().GetSegment("192.168.1.1:12345"); + ASSERT_TRUE(info_before.has_value()); + EXPECT_EQ(1024u, info_before->capacity); + + // Update + std::string update_payload = + MakeSegmentUpdatePayload("seg1", "192.168.1.1:12345", 2048, true); + OpLogEntry update_entry = + MakeEntry(2, OpType::SEGMENT_UPDATE, "seg1", update_payload); + EXPECT_TRUE(applier_->ApplyOpLogEntry(update_entry)); + EXPECT_EQ(3u, applier_->GetExpectedSequenceId()); + + auto info_after = + applier_->GetSegmentRegistry().GetSegment("192.168.1.1:12345"); + ASSERT_TRUE(info_after.has_value()); + EXPECT_EQ(2048u, info_after->capacity); +} + +TEST_F(OpLogApplierTest, TestApplySegmentMount_InvalidPayload) { + OpLogEntry entry = MakeEntry(1, OpType::SEGMENT_MOUNT, "seg1", "garbage"); + + // Should not crash, but should not add segment either + EXPECT_TRUE(applier_->ApplyOpLogEntry(entry)); + EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); + EXPECT_FALSE(applier_->GetSegmentRegistry().HasSegment("seg1")); +} + +TEST_F(OpLogApplierTest, TestApplySegmentUnmount_NonExistent) { + // Unmount a segment that was never mounted - should not crash + std::string unmount_payload = + MakeSegmentUnmountPayload("192.168.1.1:12345"); + OpLogEntry unmount_entry = + MakeEntry(1, OpType::SEGMENT_UNMOUNT, "seg1", unmount_payload); + EXPECT_TRUE(applier_->ApplyOpLogEntry(unmount_entry)); + EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); + EXPECT_FALSE( + applier_->GetSegmentRegistry().HasSegment("192.168.1.1:12345")); +} + +TEST_F(OpLogApplierTest, TestApplySegmentOperations_Mixed) { + // Mount multiple segments + OpLogEntry mount1 = + MakeEntry(1, OpType::SEGMENT_MOUNT, "seg1", + MakeSegmentMountPayload("seg1", "192.168.1.1:12345", 1024)); + OpLogEntry mount2 = + MakeEntry(2, OpType::SEGMENT_MOUNT, "seg2", + MakeSegmentMountPayload("seg2", "192.168.1.2:12345", 2048)); + OpLogEntry mount3 = + MakeEntry(3, OpType::SEGMENT_MOUNT, "seg3", + MakeSegmentMountPayload("seg3", "192.168.1.3:12345", 4096)); + + EXPECT_TRUE(applier_->ApplyOpLogEntry(mount1)); + EXPECT_TRUE(applier_->ApplyOpLogEntry(mount2)); + EXPECT_TRUE(applier_->ApplyOpLogEntry(mount3)); + + const auto& registry = applier_->GetSegmentRegistry(); + EXPECT_EQ(3u, registry.GetAllSegments().size()); + + // Unmount one + OpLogEntry unmount2 = + MakeEntry(4, OpType::SEGMENT_UNMOUNT, "seg2", + MakeSegmentUnmountPayload("192.168.1.2:12345")); + EXPECT_TRUE(applier_->ApplyOpLogEntry(unmount2)); + EXPECT_EQ(2u, registry.GetAllSegments().size()); + EXPECT_FALSE(registry.HasSegment("192.168.1.2:12345")); + + // Update another + OpLogEntry update3 = + MakeEntry(5, OpType::SEGMENT_UPDATE, "seg3", + MakeSegmentUpdatePayload("seg3", "192.168.1.3:12345", 8192)); + EXPECT_TRUE(applier_->ApplyOpLogEntry(update3)); + auto info3 = registry.GetSegment("192.168.1.3:12345"); + ASSERT_TRUE(info3.has_value()); + EXPECT_EQ(8192u, info3->capacity); +} + } // namespace mooncake::test int main(int argc, char** argv) { diff --git a/mooncake-store/tests/ha/oplog/oplog_batch_auditor_test.cpp b/mooncake-store/tests/ha/oplog/oplog_batch_auditor_test.cpp new file mode 100644 index 0000000000..f376d04bb3 --- /dev/null +++ b/mooncake-store/tests/ha/oplog/oplog_batch_auditor_test.cpp @@ -0,0 +1,193 @@ +#include + +#if __has_include() +#include // Ubuntu +#else +#include // CentOS +#endif + +#include +#include +#include +#include + +#include "ha/oplog/oplog_batch_codec.h" +#include "tools/oplog_batch_auditor.h" + +namespace mooncake { +namespace { + +class RangeBackend : public HaKvBackend { + public: + explicit RangeBackend(std::vector values) + : values_(std::move(values)) {} + + ErrorCode Get(std::string_view, std::string&) override { + return ErrorCode::ETCD_KEY_NOT_EXIST; + } + ErrorCode Put(std::string_view, std::string_view) override { + return ErrorCode::OK; + } + ErrorCode Range(std::string_view, std::string_view, size_t limit, + std::vector& kvs) override { + observed_limit = limit; + kvs.assign(values_.begin(), + values_.begin() + std::min(limit, values_.size())); + return ErrorCode::OK; + } + bool SupportsTxn() const override { return false; } + ErrorCode Txn(const KvTxn&) override { return ErrorCode::INVALID_PARAMS; } + + size_t observed_limit{0}; + + private: + std::vector values_; +}; + +OpLogBatchRecord MakeBatch(uint64_t batch_id, uint64_t first_seq, + size_t entry_count) { + OpLogBatchRecord batch; + batch.batch_id = batch_id; + batch.first_seq = first_seq; + batch.last_seq = first_seq + entry_count - 1; + for (size_t i = 0; i < entry_count; ++i) { + batch.entries.push_back({.sequence_id = first_seq + i, + .op_type = OpType::REMOVE, + .tenant_id = "default", + .object_key = "key-" + std::to_string(i), + .payload = ""}); + } + return batch; +} + +TEST(OpLogBatchAuditorTest, AcceptsContiguousHistoryAfterLegacyCutover) { + const std::string cluster = "audit-test"; + std::vector kvs = { + {.key = "/oplog/audit-test/00000000000000000007", .value = "legacy"}, + {.key = BuildDurablePrefixKey(cluster), + .value = EncodeDurablePrefix({.batch_id = 2, .last_seq = 10})}, + {.key = BuildBatchRecordKey(cluster, 1), + .value = EncodeOpLogBatchRecord(MakeBatch(1, 8, 2))}, + {.key = BuildBatchRecordKey(cluster, 2), + .value = EncodeOpLogBatchRecord(MakeBatch(2, 10, 1))}, + }; + + OpLogAuditReport report = AuditOpLogNamespace(cluster, kvs, 100); + + EXPECT_TRUE(report.ok); + EXPECT_EQ(report.legacy_max_seq, 7); + ASSERT_TRUE(report.durable_prefix.has_value()); + EXPECT_EQ(report.durable_prefix->batch_id, 2); + EXPECT_EQ(report.durable_prefix->last_seq, 10); + EXPECT_EQ(report.batch_count, 2); + EXPECT_EQ(report.entry_count, 3); + EXPECT_TRUE(report.errors.empty()); +} + +TEST(OpLogBatchAuditorTest, SerializesMachineReadableReport) { + OpLogAuditReport report; + report.ok = false; + report.cluster_id = "audit-test"; + report.legacy_max_seq = 7; + report.durable_prefix = {.batch_id = 2, .last_seq = 10}; + report.batch_count = 2; + report.entry_count = 3; + report.orphan_batches = {3}; + report.warnings = {"warning"}; + report.errors = {"error"}; + report.truncated_errors = true; + + const std::string encoded = OpLogAuditReportToJson(report); + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::istringstream input(encoded); + + ASSERT_TRUE(Json::parseFromStream(builder, input, &root, &errors)) + << errors; + EXPECT_EQ(root["schema_version"].asUInt(), 1); + EXPECT_FALSE(root["ok"].asBool()); + EXPECT_EQ(root["durable_prefix"]["last_seq"].asUInt64(), 10); + EXPECT_EQ(root["orphan_batches"][0].asUInt64(), 3); + EXPECT_EQ(root["warnings"][0].asString(), "warning"); + EXPECT_EQ(root["errors"][0].asString(), "error"); + EXPECT_TRUE(root["truncated_errors"].asBool()); +} + +TEST(OpLogBatchAuditorTest, ReportsUnknownSidecarWithoutFailingAudit) { + std::vector kvs = { + {.key = "/oplog/audit-test/custom-sidecar", .value = "value"}, + }; + + OpLogAuditReport report = AuditOpLogNamespace("audit-test", kvs, 100); + + EXPECT_TRUE(report.ok); + ASSERT_EQ(report.warnings.size(), 1); + EXPECT_NE(report.warnings[0].find("custom-sidecar"), std::string::npos); +} + +TEST(OpLogBatchAuditorTest, RejectsMalformedLegacyLikeKey) { + std::vector kvs = { + {.key = "/oplog/audit-test/0000000000000000000x", .value = "value"}, + }; + + OpLogAuditReport report = AuditOpLogNamespace("audit-test", kvs, 100); + + EXPECT_FALSE(report.ok); + ASSERT_EQ(report.errors.size(), 1); + EXPECT_NE(report.errors[0].find("malformed legacy key"), std::string::npos); +} + +TEST(OpLogBatchAuditorTest, CapsErrorsAndCountsMalformedBatchKeys) { + std::vector kvs; + for (size_t i = 0; i < 105; ++i) { + kvs.push_back( + {.key = "/oplog/audit-test/batches/bad-" + std::to_string(i), + .value = "value"}); + } + + OpLogAuditReport report = AuditOpLogNamespace("audit-test", kvs, 100); + + EXPECT_FALSE(report.ok); + EXPECT_EQ(report.batch_count, 105); + EXPECT_EQ(report.errors.size(), 100); + EXPECT_TRUE(report.truncated_errors); +} + +TEST(OpLogBatchAuditorTest, ClassifiesBatchesWithoutPrefixAsOrphans) { + const std::string cluster = "audit-test"; + std::vector kvs = { + {.key = BuildBatchRecordKey(cluster, 1), + .value = EncodeOpLogBatchRecord(MakeBatch(1, 1, 1))}, + }; + + OpLogAuditReport report = AuditOpLogNamespace(cluster, kvs, 100); + + EXPECT_FALSE(report.ok); + EXPECT_EQ(report.batch_count, 1); + EXPECT_EQ(report.entry_count, 1); + ASSERT_EQ(report.orphan_batches.size(), 1); + EXPECT_EQ(report.orphan_batches[0], 1); +} + +TEST(OpLogBatchAuditorTest, CapsNamespaceReadBeforeAudit) { + std::vector values(3); + RangeBackend backend(std::move(values)); + OpLogNamespaceRead result; + + ErrorCode err = ReadOpLogNamespace("audit-test", backend, 2, &result); + + EXPECT_EQ(err, ErrorCode::OK); + EXPECT_EQ(backend.observed_limit, 3); + EXPECT_TRUE(result.truncated); + EXPECT_EQ(result.kvs.size(), 2); +} + +TEST(OpLogBatchAuditorTest, RejectsOverflowingDumpRange) { + EXPECT_FALSE(ComputeOpLogDumpLimit(1, UINT64_MAX, 100).has_value()); + EXPECT_EQ(ComputeOpLogDumpLimit(1, 0, 100), 100); + EXPECT_EQ(ComputeOpLogDumpLimit(1, 2, 100), 2); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/ha/oplog/oplog_batch_codec_test.cpp b/mooncake-store/tests/ha/oplog/oplog_batch_codec_test.cpp new file mode 100644 index 0000000000..22a1bf8b64 --- /dev/null +++ b/mooncake-store/tests/ha/oplog/oplog_batch_codec_test.cpp @@ -0,0 +1,349 @@ +#include "ha/oplog/oplog_batch_codec.h" +#include "ha/oplog/oplog_batch_types.h" + +#include +#include + +#include +#include +#include +#include + +#if __has_include() +#include +#else +#include +#endif + +namespace mooncake::test { + +namespace { + +OpLogEntry MakeEntry(uint64_t seq, OpType type = OpType::PUT_END, + std::string key = "key", std::string payload = "value") { + OpLogEntry entry; + entry.sequence_id = seq; + entry.timestamp_ms = 1234567890; + entry.op_type = type; + entry.tenant_id = "tenant"; + entry.object_key = std::move(key); + entry.payload = std::move(payload); + entry.checksum = static_cast( + XXH32(entry.payload.data(), entry.payload.size(), 0)); + entry.prefix_hash = static_cast( + XXH32(entry.object_key.data(), entry.object_key.size(), 0)); + return entry; +} + +OpLogBatchRecord MakeBatch(uint64_t batch_id, std::vector entries) { + OpLogBatchRecord batch; + batch.batch_id = batch_id; + batch.entries = std::move(entries); + batch.first_seq = batch.entries.front().sequence_id; + batch.last_seq = batch.entries.back().sequence_id; + return batch; +} + +Json::Value ParseJson(const std::string& value) { + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::istringstream stream(value); + EXPECT_TRUE(Json::parseFromStream(builder, stream, &root, &errors)) + << errors; + return root; +} + +std::string WriteJson(const Json::Value& root) { + Json::StreamWriterBuilder builder; + builder["indentation"] = ""; + return Json::writeString(builder, root); +} + +} // namespace + +TEST(OpLogBatchTypesTest, RejectsEmptyEntries) { + OpLogBatchRecord batch; + batch.batch_id = 1; + batch.first_seq = 1; + batch.last_seq = 1; + + std::string reason; + EXPECT_FALSE(ValidateOpLogBatchRecordShape(batch, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchTypesTest, RejectsFirstSeqMismatch) { + auto batch = MakeBatch(1, {MakeEntry(2)}); + batch.first_seq = 1; + + std::string reason; + EXPECT_FALSE(ValidateOpLogBatchRecordShape(batch, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchTypesTest, RejectsLastSeqMismatch) { + auto batch = MakeBatch(1, {MakeEntry(1)}); + batch.last_seq = 2; + + std::string reason; + EXPECT_FALSE(ValidateOpLogBatchRecordShape(batch, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchTypesTest, RejectsNonContiguousEntrySequence) { + auto batch = MakeBatch(1, {MakeEntry(1), MakeEntry(3)}); + batch.last_seq = 3; + + std::string reason; + EXPECT_FALSE(ValidateOpLogBatchRecordShape(batch, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchTypesTest, RejectsZeroSchemaVersion) { + auto batch = MakeBatch(1, {MakeEntry(1)}); + batch.schema_version = 0; + + std::string reason; + EXPECT_FALSE(ValidateOpLogBatchRecordShape(batch, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchTypesTest, AcceptsValidContiguousBatch) { + auto batch = MakeBatch(1, {MakeEntry(1), MakeEntry(2), MakeEntry(3)}); + + std::string reason; + EXPECT_TRUE(ValidateOpLogBatchRecordShape(batch, &reason)); + EXPECT_TRUE(reason.empty()); +} + +TEST(OpLogBatchKeyLayoutTest, BuildsPaddedBatchRecordKey) { + EXPECT_EQ("/oplog/clusterA/batches/00000000000000000001", + BuildBatchRecordKey("clusterA", 1)); +} + +TEST(OpLogBatchKeyLayoutTest, BuildsDurablePrefixKey) { + EXPECT_EQ("/oplog/clusterA/durable_prefix", + BuildDurablePrefixKey("clusterA")); +} + +TEST(OpLogBatchKeyLayoutTest, BuildsBatchRangeBounds) { + auto bounds = BuildBatchRecordRange("clusterA", 7); + EXPECT_EQ("/oplog/clusterA/batches/00000000000000000008", bounds.begin_key); + EXPECT_EQ("/oplog/clusterA/batches0", bounds.end_key); +} + +TEST(OpLogBatchKeyLayoutTest, MaxBatchRangeIsEmpty) { + auto bounds = BuildBatchRecordRange("clusterA", UINT64_MAX); + EXPECT_EQ(bounds.begin_key, bounds.end_key); +} + +TEST(OpLogBatchKeyLayoutTest, RejectsInvalidClusterId) { + std::string reason; + EXPECT_FALSE(ValidateOpLogBatchClusterId("bad/cluster", &reason)); + EXPECT_FALSE(reason.empty()); + EXPECT_TRUE(BuildBatchRecordKey("bad/cluster", 1).empty()); + EXPECT_TRUE(BuildDurablePrefixKey("bad/cluster").empty()); + auto bounds = BuildBatchRecordRange("bad/cluster", 1); + EXPECT_TRUE(bounds.begin_key.empty()); + EXPECT_TRUE(bounds.end_key.empty()); +} + +TEST(OpLogBatchKeyLayoutTest, NormalizesTrailingSlashClusterId) { + EXPECT_EQ("/oplog/clusterA/batches/00000000000000000001", + BuildBatchRecordKey("clusterA/", 1)); +} + +TEST(OpLogDurablePrefixCodecTest, RoundTripsZeroPrefix) { + DurablePrefix in; + auto encoded = EncodeDurablePrefix(in); + + DurablePrefix out; + std::string reason; + ASSERT_TRUE(DecodeDurablePrefix(encoded, &out, &reason)); + EXPECT_EQ(in.batch_id, out.batch_id); + EXPECT_EQ(in.last_seq, out.last_seq); + EXPECT_TRUE(reason.empty()); +} + +TEST(OpLogDurablePrefixCodecTest, RoundTripsNonZeroPrefix) { + DurablePrefix in{.batch_id = 9, .last_seq = 1024}; + auto encoded = EncodeDurablePrefix(in); + + DurablePrefix out; + std::string reason; + ASSERT_TRUE(DecodeDurablePrefix(encoded, &out, &reason)); + EXPECT_EQ(in.batch_id, out.batch_id); + EXPECT_EQ(in.last_seq, out.last_seq); + EXPECT_TRUE(reason.empty()); +} + +TEST(OpLogDurablePrefixCodecTest, RejectsMalformedPayload) { + DurablePrefix out; + std::string reason; + EXPECT_FALSE(DecodeDurablePrefix("{", &out, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogDurablePrefixCodecTest, RejectsMalformedTypedFieldsWithoutThrowing) { + DurablePrefix out; + std::string reason; + EXPECT_NO_THROW(EXPECT_FALSE(DecodeDurablePrefix( + R"({"schema_version":"x","batch_id":"x","last_seq":"x"})", &out, + &reason))); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogDurablePrefixCodecTest, RejectsUnsupportedSchemaVersion) { + DurablePrefix out; + std::string reason; + EXPECT_FALSE(DecodeDurablePrefix( + R"({"schema_version":2,"batch_id":1,"last_seq":1})", &out, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchRecordCodecTest, RoundTripsSingleEntryBatch) { + auto in = MakeBatch(3, {MakeEntry(10)}); + auto encoded = EncodeOpLogBatchRecord(in); + + OpLogBatchRecord out; + std::string reason; + ASSERT_TRUE(DecodeOpLogBatchRecord(encoded, &out, &reason)); + EXPECT_EQ(in.batch_id, out.batch_id); + EXPECT_EQ(in.first_seq, out.first_seq); + EXPECT_EQ(in.last_seq, out.last_seq); + ASSERT_EQ(1u, out.entries.size()); + EXPECT_EQ(in.entries[0].sequence_id, out.entries[0].sequence_id); + EXPECT_EQ(in.entries[0].payload, out.entries[0].payload); + EXPECT_TRUE(reason.empty()); +} + +TEST(OpLogBatchRecordCodecTest, EncodesCompactArraySchema) { + const auto encoded = EncodeOpLogBatchRecord(MakeBatch( + 3, + {MakeEntry(10, OpType::PUT_END, "key", std::string("\0binary", 7))})); + const auto root = ParseJson(encoded); + + ASSERT_TRUE(root.isArray()); + ASSERT_EQ(6u, root.size()); + EXPECT_EQ(kOpLogBatchRecordSchemaVersion, root[0].asUInt()); + EXPECT_EQ(3u, root[1].asUInt64()); + EXPECT_EQ(10u, root[2].asUInt64()); + EXPECT_EQ(10u, root[3].asUInt64()); + ASSERT_TRUE(root[4].isArray()); + ASSERT_EQ(1u, root[4].size()); + ASSERT_TRUE(root[4][0].isArray()); + EXPECT_EQ(4u, root[4][0].size()); + EXPECT_EQ(static_cast(OpType::PUT_END), root[4][0][0].asUInt()); + EXPECT_EQ("tenant", root[4][0][1].asString()); + EXPECT_EQ("key", root[4][0][2].asString()); + EXPECT_EQ("AGJpbmFyeQ==", root[4][0][3].asString()); + EXPECT_EQ(std::string::npos, encoded.find("timestamp_ms")); + EXPECT_EQ(std::string::npos, encoded.find("prefix_hash")); + EXPECT_EQ(std::string::npos, encoded.find("sequence_id")); +} + +TEST(OpLogBatchRecordCodecTest, RebuildsNonPersistedEntryFields) { + auto in = MakeBatch(3, {MakeEntry(10)}); + ASSERT_NE(0u, in.entries[0].timestamp_ms); + ASSERT_NE(0u, in.entries[0].prefix_hash); + + OpLogBatchRecord out; + std::string reason; + ASSERT_TRUE( + DecodeOpLogBatchRecord(EncodeOpLogBatchRecord(in), &out, &reason)); + ASSERT_EQ(1u, out.entries.size()); + EXPECT_EQ(0u, out.entries[0].timestamp_ms); + EXPECT_EQ(0u, out.entries[0].prefix_hash); + EXPECT_TRUE(VerifyOpLogChecksum(out.entries[0])); +} + +TEST(OpLogBatchRecordCodecTest, RoundTripsMultiEntryBatch) { + auto in = MakeBatch( + 3, {MakeEntry(10), MakeEntry(11, OpType::REMOVE, "dead-key", "")}); + auto encoded = EncodeOpLogBatchRecord(in); + + OpLogBatchRecord out; + std::string reason; + ASSERT_TRUE(DecodeOpLogBatchRecord(encoded, &out, &reason)); + EXPECT_EQ(in.batch_id, out.batch_id); + EXPECT_EQ(in.first_seq, out.first_seq); + EXPECT_EQ(in.last_seq, out.last_seq); + ASSERT_EQ(2u, out.entries.size()); + EXPECT_EQ(in.entries[1].op_type, out.entries[1].op_type); + EXPECT_EQ(in.entries[1].object_key, out.entries[1].object_key); + EXPECT_EQ(in.entries[1].payload, out.entries[1].payload); +} + +TEST(OpLogBatchRecordCodecTest, RejectsCorruptedChecksum) { + auto root = + ParseJson(EncodeOpLogBatchRecord(MakeBatch(3, {MakeEntry(10)}))); + ASSERT_TRUE(root.isArray()); + ASSERT_EQ(6u, root.size()); + root[5] = root[5].asUInt() + 1; + + OpLogBatchRecord out; + std::string reason; + EXPECT_FALSE(DecodeOpLogBatchRecord(WriteJson(root), &out, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchRecordCodecTest, RejectsMalformedTypedFieldsWithoutThrowing) { + OpLogBatchRecord out; + std::string reason; + EXPECT_NO_THROW(EXPECT_FALSE( + DecodeOpLogBatchRecord(R"(["x","x","x","x",[],"x"])", &out, &reason))); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchRecordCodecTest, RejectsWrongElementCount) { + OpLogBatchRecord out; + std::string reason; + EXPECT_FALSE(DecodeOpLogBatchRecord(R"([1,3,10,10,[]])", &out, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchRecordCodecTest, RejectsCorruptedSequenceContinuity) { + auto in = MakeBatch(3, {MakeEntry(10), MakeEntry(11)}); + auto root = ParseJson(EncodeOpLogBatchRecord(in)); + ASSERT_TRUE(root.isArray()); + root[3] = Json::UInt64(12); + + OpLogBatchRecord out; + std::string reason; + EXPECT_FALSE(DecodeOpLogBatchRecord(WriteJson(root), &out, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchRecordCodecTest, RejectsUnsupportedSchemaVersion) { + auto root = + ParseJson(EncodeOpLogBatchRecord(MakeBatch(3, {MakeEntry(10)}))); + ASSERT_TRUE(root.isArray()); + root[0] = kOpLogBatchRecordSchemaVersion + 1; + + OpLogBatchRecord out; + std::string reason; + EXPECT_FALSE(DecodeOpLogBatchRecord(WriteJson(root), &out, &reason)); + EXPECT_FALSE(reason.empty()); +} + +TEST(OpLogBatchRecordCodecTest, RejectsInvalidBase64Payload) { + auto root = + ParseJson(EncodeOpLogBatchRecord(MakeBatch(3, {MakeEntry(10)}))); + ASSERT_TRUE(root.isArray()); + root[4][0][3] = "%%%"; + Json::Value checksum_payload(Json::arrayValue); + for (Json::ArrayIndex i = 0; i < 5; ++i) { + checksum_payload.append(root[i]); + } + const auto serialized = WriteJson(checksum_payload); + root[5] = + static_cast(XXH32(serialized.data(), serialized.size(), 0)); + + OpLogBatchRecord out; + std::string reason; + EXPECT_FALSE(DecodeOpLogBatchRecord(WriteJson(root), &out, &reason)); + EXPECT_NE(std::string::npos, reason.find("base64")); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/ha/oplog/oplog_batch_standby_reader_test.cpp b/mooncake-store/tests/ha/oplog/oplog_batch_standby_reader_test.cpp new file mode 100644 index 0000000000..b2de07ba5a --- /dev/null +++ b/mooncake-store/tests/ha/oplog/oplog_batch_standby_reader_test.cpp @@ -0,0 +1,481 @@ +#include "ha/oplog/oplog_batch_standby_reader.h" + +#include +#include + +#include +#include +#include +#include + +#include "ha/kv/ha_kv_backend.h" +#include "ha/oplog/oplog_applier.h" +#include "ha/oplog/oplog_batch_codec.h" +#include "mock_metadata_store.h" + +namespace mooncake::test { +namespace { + +class FakeHaKvBackend : public HaKvBackend { + public: + ErrorCode Get(std::string_view key, std::string& value) override { + if (next_get_error_ != ErrorCode::OK) { + auto error = next_get_error_; + next_get_error_ = ErrorCode::OK; + return error; + } + auto it = values_.find(std::string(key)); + if (it == values_.end()) { + return ErrorCode::ETCD_KEY_NOT_EXIST; + } + value = it->second; + return ErrorCode::OK; + } + ErrorCode Put(std::string_view key, std::string_view value) override { + values_[std::string(key)] = std::string(value); + return ErrorCode::OK; + } + ErrorCode Range(std::string_view begin_key, std::string_view end_key, + size_t limit, std::vector& kvs) override { + if (next_range_error_ != ErrorCode::OK) { + auto error = next_range_error_; + next_range_error_ = ErrorCode::OK; + return error; + } + kvs.clear(); + ++range_calls_; + for (const auto& [key, value] : values_) { + if (key >= begin_key && key < end_key) { + kvs.push_back({.key = key, .value = value}); + if (limit != 0 && kvs.size() >= limit) break; + } + } + return ErrorCode::OK; + } + bool SupportsTxn() const override { return true; } + ErrorCode Txn(const KvTxn&) override { return ErrorCode::OK; } + + int range_calls() const { return range_calls_; } + void Erase(std::string_view key) { values_.erase(std::string(key)); } + void FailNextGet(ErrorCode error) { next_get_error_ = error; } + void FailNextRange(ErrorCode error) { next_range_error_ = error; } + + private: + std::map values_; + int range_calls_{0}; + ErrorCode next_get_error_{ErrorCode::OK}; + ErrorCode next_range_error_{ErrorCode::OK}; +}; + +OpLogEntry MakeEntry(uint64_t sequence_id) { + OpLogEntry entry; + entry.sequence_id = sequence_id; + entry.op_type = OpType::REMOVE; + entry.tenant_id = "tenant"; + entry.object_key = "key" + std::to_string(sequence_id); + entry.checksum = static_cast( + XXH32(entry.payload.data(), entry.payload.size(), 0)); + entry.prefix_hash = static_cast( + XXH32(entry.object_key.data(), entry.object_key.size(), 0)); + return entry; +} + +OpLogBatchRecord MakeBatch(uint64_t batch_id, uint64_t first_seq, + uint64_t last_seq) { + OpLogBatchRecord batch; + batch.batch_id = batch_id; + batch.first_seq = first_seq; + batch.last_seq = last_seq; + for (uint64_t seq = first_seq; seq <= last_seq; ++seq) { + batch.entries.push_back(MakeEntry(seq)); + } + return batch; +} + +} // namespace + +TEST(OpLogBatchStandbyReaderTest, MissingPrefixWaitsWithoutLegacyFallback) { + FakeHaKvBackend backend; + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + ASSERT_EQ(ErrorCode::OK, result.error); + EXPECT_FALSE(result.durable_prefix_present); + EXPECT_EQ(0u, result.applied_entries); + EXPECT_EQ(1u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, MissingPrefixAfterObservationFailsClosed) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + ASSERT_EQ(ErrorCode::OK, reader.PollOnce().error); + + backend.Erase(BuildDurablePrefixKey("clusterA")); + auto result = reader.PollOnce(); + + EXPECT_NE(ErrorCode::OK, result.error); + EXPECT_EQ(OpLogBatchStandbyPollDisposition::FATAL, result.disposition); +} + +TEST(OpLogBatchStandbyReaderTest, PrefixTransportErrorIsRetryable) { + FakeHaKvBackend backend; + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + backend.FailNextGet(ErrorCode::ETCD_OPERATION_ERROR); + auto result = reader.PollOnce(); + + EXPECT_EQ(ErrorCode::ETCD_OPERATION_ERROR, result.error); + EXPECT_EQ(OpLogBatchStandbyPollDisposition::RETRYABLE, result.disposition); +} + +TEST(OpLogBatchStandbyReaderTest, RangeTransportErrorIsRetryable) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + backend.FailNextRange(ErrorCode::ETCD_OPERATION_ERROR); + + auto result = reader.PollOnce(); + + EXPECT_EQ(ErrorCode::ETCD_OPERATION_ERROR, result.error); + EXPECT_EQ(OpLogBatchStandbyPollDisposition::RETRYABLE, result.disposition); +} + +TEST(OpLogBatchStandbyReaderTest, MissingTargetBatchIsFatal) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + EXPECT_EQ(OpLogBatchStandbyPollDisposition::FATAL, result.disposition); + EXPECT_EQ(ErrorCode::ETCD_KEY_NOT_EXIST, result.error); +} + +TEST(OpLogBatchStandbyReaderTest, DurablePrefixRegressionFailsClosed) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + ASSERT_EQ(ErrorCode::OK, reader.PollOnce().error); + + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 0, .last_seq = 0}))); + EXPECT_NE(ErrorCode::OK, reader.PollOnce().error); +} + +TEST(OpLogBatchStandbyReaderTest, BatchZeroMustHaveZeroSequence) { + FakeHaKvBackend backend; + ASSERT_EQ( + ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 0, .last_seq = 42}))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + EXPECT_EQ(ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, result.error); + EXPECT_EQ(OpLogBatchStandbyPollDisposition::FATAL, result.disposition); +} + +TEST(OpLogBatchStandbyReaderTest, RangeReadsBatchesWhenDurablePrefixAdvances) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 2}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 2)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + ASSERT_EQ(ErrorCode::OK, result.error); + EXPECT_EQ(1, backend.range_calls()); + EXPECT_EQ(2u, result.applied_entries); + EXPECT_EQ(3u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, DoesNothingWhenDurablePrefixDoesNotAdvance) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + ASSERT_EQ(ErrorCode::OK, reader.PollOnce().error); + int range_calls_after_first_poll = backend.range_calls(); + auto second = reader.PollOnce(); + + ASSERT_EQ(ErrorCode::OK, second.error); + EXPECT_EQ(0u, second.applied_entries); + EXPECT_EQ(range_calls_after_first_poll, backend.range_calls()); + EXPECT_EQ(2u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, FullPageContinuesOnNextPoll) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 3, .last_seq = 3}))); + for (uint64_t id = 1; id <= 3; ++id) { + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", id), + EncodeOpLogBatchRecord(MakeBatch(id, id, id)))); + } + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto first = reader.PollOnce(/*max_batches=*/2); + ASSERT_EQ(ErrorCode::OK, first.error); + EXPECT_EQ(2u, first.applied_entries); + EXPECT_EQ(3u, applier.GetExpectedSequenceId()); + + auto second = reader.PollOnce(/*max_batches=*/2); + ASSERT_EQ(ErrorCode::OK, second.error); + EXPECT_EQ(1u, second.applied_entries); + EXPECT_EQ(4u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, RejectsPrefixPointingIntoBatch) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 2}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 3)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + EXPECT_NE(ErrorCode::OK, result.error); + EXPECT_EQ(0u, result.applied_entries); + EXPECT_EQ(1u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, AppliesExpandedEntriesInSequenceOrder) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 2, .last_seq = 3}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 2), + EncodeOpLogBatchRecord(MakeBatch(2, 2, 3)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + ASSERT_EQ(ErrorCode::OK, result.error); + EXPECT_EQ(3u, result.applied_entries); + EXPECT_EQ(4u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, AcceptsFirstBatchAfterSnapshotBaseline) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 4}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 3, 4)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + applier.Recover(2); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + ASSERT_EQ(ErrorCode::OK, result.error); + EXPECT_EQ(2u, result.applied_entries); + EXPECT_EQ(5u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, FailsWhenLaterBatchHasSequenceGap) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 2, .last_seq = 4}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 2)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 2), + EncodeOpLogBatchRecord(MakeBatch(2, 4, 4)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + EXPECT_NE(ErrorCode::OK, result.error); + EXPECT_EQ(2u, result.applied_entries); + EXPECT_EQ(3u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, MissingBatchMarksReaderUnhealthy) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 2, .last_seq = 2}))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + EXPECT_NE(ErrorCode::OK, result.error); + EXPECT_EQ(0u, result.applied_entries); + EXPECT_EQ(1u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, StartsFromEarliestExistingBatch) { + FakeHaKvBackend backend; + ASSERT_EQ( + ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 42, .last_seq = 44}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 42), + EncodeOpLogBatchRecord(MakeBatch(42, 43, 44)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + applier.Recover(42); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + EXPECT_EQ(ErrorCode::OK, result.error); + EXPECT_EQ(2u, result.applied_entries); + EXPECT_EQ(45u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, RejectsSequenceOverlapAcrossBatches) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 2, .last_seq = 3}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 2)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 2), + EncodeOpLogBatchRecord(MakeBatch(2, 2, 3)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + EXPECT_NE(ErrorCode::OK, result.error); + EXPECT_EQ(2u, result.applied_entries); + EXPECT_EQ(3u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, ChecksumFailureMarksReaderUnhealthy) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + std::string encoded = EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)); + encoded.back() = static_cast(encoded.back() ^ 0x1); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), encoded)); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + EXPECT_NE(ErrorCode::OK, result.error); + EXPECT_EQ(0u, result.applied_entries); + EXPECT_EQ(1u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, RejectsFutureFirstSequence) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 2}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 2, 2)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + EXPECT_EQ(ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, result.error); + EXPECT_EQ(OpLogBatchStandbyPollDisposition::FATAL, result.disposition); + EXPECT_EQ(0u, result.applied_entries); + EXPECT_EQ(1u, applier.GetExpectedSequenceId()); +} + +TEST(OpLogBatchStandbyReaderTest, AlreadyAppliedEntriesAreSkipped) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildDurablePrefixKey("clusterA"), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 3}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put(BuildBatchRecordKey("clusterA", 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 3)))); + MockMetadataStore metadata_store; + OpLogApplier applier(&metadata_store, "clusterA"); + applier.Recover(2); + OpLogBatchStandbyReader reader("clusterA", backend, applier); + + auto result = reader.PollOnce(); + + ASSERT_EQ(ErrorCode::OK, result.error); + EXPECT_EQ(1u, result.applied_entries); + EXPECT_EQ(4u, applier.GetExpectedSequenceId()); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/ha/oplog/oplog_batch_storage_test.cpp b/mooncake-store/tests/ha/oplog/oplog_batch_storage_test.cpp new file mode 100644 index 0000000000..c1df73d21a --- /dev/null +++ b/mooncake-store/tests/ha/oplog/oplog_batch_storage_test.cpp @@ -0,0 +1,668 @@ +#include "ha/oplog/oplog_batch_storage.h" + +#include +#include + +#include +#include +#include + +#include "ha/kv/ha_kv_backend.h" +#include "ha/oplog/oplog_batch_codec.h" +#include "ha/oplog/oplog_batch_types.h" + +namespace mooncake::test { +namespace { + +class FakeHaKvBackend : public HaKvBackend { + public: + ErrorCode Get(std::string_view key, std::string& value) override { + if (next_get_key_error_ != ErrorCode::OK && + key == next_get_error_key_) { + ErrorCode err = next_get_key_error_; + next_get_key_error_ = ErrorCode::OK; + next_get_error_key_.clear(); + return err; + } + if (next_get_error_ != ErrorCode::OK) { + ErrorCode err = next_get_error_; + next_get_error_ = ErrorCode::OK; + return err; + } + auto it = kvs_.find(std::string(key)); + if (it == kvs_.end()) { + return ErrorCode::ETCD_KEY_NOT_EXIST; + } + value = it->second; + return ErrorCode::OK; + } + + ErrorCode Put(std::string_view key, std::string_view value) override { + kvs_[std::string(key)] = std::string(value); + return ErrorCode::OK; + } + + ErrorCode Range(std::string_view begin_key, std::string_view end_key, + size_t limit, std::vector& kvs) override { + range_limits_.push_back(limit); + if (next_range_error_ != ErrorCode::OK) { + ErrorCode err = next_range_error_; + next_range_error_ = ErrorCode::OK; + return err; + } + kvs.clear(); + for (auto it = kvs_.lower_bound(std::string(begin_key)); + it != kvs_.end() && it->first < end_key; ++it) { + kvs.push_back({.key = it->first, .value = it->second}); + if (limit != 0 && kvs.size() >= limit) { + break; + } + } + return ErrorCode::OK; + } + + bool SupportsTxn() const override { return supports_txn_; } + + ErrorCode Txn(const KvTxn& txn) override { + if (!supports_txn_) { + return ErrorCode::INVALID_PARAMS; + } + if (next_txn_error_ != ErrorCode::OK) { + ErrorCode err = next_txn_error_; + next_txn_error_ = ErrorCode::OK; + return err; + } + if (race_before_next_txn_) { + kvs_[race_key_] = race_value_; + race_before_next_txn_ = false; + } + for (const auto& compare : txn.compares) { + auto it = kvs_.find(compare.key); + if (compare.kind == KvCompareKind::kKeyNotExists) { + if (it != kvs_.end()) { + return ErrorCode::ETCD_TRANSACTION_FAIL; + } + } else if (it == kvs_.end() || + it->second != compare.expected_value) { + return ErrorCode::ETCD_TRANSACTION_FAIL; + } + } + for (const auto& put : txn.puts) { + kvs_[put.key] = put.value; + } + return ErrorCode::OK; + } + + void SetSupportsTxn(bool supports_txn) { supports_txn_ = supports_txn; } + void CreateBeforeNextTxn(std::string key, std::string value) { + race_before_next_txn_ = true; + race_key_ = std::move(key); + race_value_ = std::move(value); + } + void FailNextGet(ErrorCode err) { next_get_error_ = err; } + void FailNextGetForKey(std::string key, ErrorCode err) { + next_get_error_key_ = std::move(key); + next_get_key_error_ = err; + } + void FailNextRange(ErrorCode err) { next_range_error_ = err; } + void FailNextTxn(ErrorCode err) { next_txn_error_ = err; } + const std::vector& range_limits() const { return range_limits_; } + + private: + std::map kvs_; + bool supports_txn_{true}; + bool race_before_next_txn_{false}; + std::string race_key_; + std::string race_value_; + ErrorCode next_get_error_{ErrorCode::OK}; + std::string next_get_error_key_; + ErrorCode next_get_key_error_{ErrorCode::OK}; + ErrorCode next_range_error_{ErrorCode::OK}; + ErrorCode next_txn_error_{ErrorCode::OK}; + std::vector range_limits_; +}; + +OpLogEntry MakeEntry(uint64_t seq) { + OpLogEntry entry; + entry.sequence_id = seq; + entry.timestamp_ms = 1234567890; + entry.op_type = OpType::PUT_END; + entry.tenant_id = "tenant"; + entry.object_key = "key" + std::to_string(seq); + entry.payload = "value" + std::to_string(seq); + entry.checksum = static_cast( + XXH32(entry.payload.data(), entry.payload.size(), 0)); + entry.prefix_hash = static_cast( + XXH32(entry.object_key.data(), entry.object_key.size(), 0)); + return entry; +} + +OpLogBatchRecord MakeBatch(uint64_t batch_id, uint64_t first_seq, + size_t count) { + OpLogBatchRecord batch; + batch.batch_id = batch_id; + batch.first_seq = first_seq; + batch.last_seq = first_seq + count - 1; + for (size_t i = 0; i < count; ++i) { + batch.entries.push_back(MakeEntry(first_seq + i)); + } + return batch; +} + +} // namespace + +TEST(OpLogBatchStorageTest, InitializesEmptyNamespaceAtZero) { + FakeHaKvBackend backend; + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::OK, storage.InitDurablePrefix(prefix)); + + EXPECT_EQ(0u, prefix.batch_id); + EXPECT_EQ(0u, prefix.last_seq); + std::string encoded; + ASSERT_EQ(ErrorCode::OK, + backend.Get("/oplog/clusterA/durable_prefix", encoded)); + DurablePrefix stored; + ASSERT_TRUE(DecodeDurablePrefix(encoded, &stored)); + EXPECT_EQ(prefix.batch_id, stored.batch_id); + EXPECT_EQ(prefix.last_seq, stored.last_seq); +} + +TEST(OpLogBatchStorageTest, RejectsLegacyLatest) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, backend.Put("/oplog/clusterA/latest", "42")); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_NE(ErrorCode::OK, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, RejectsLegacyEntry) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/00000000000000000042", "entry")); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_NE(ErrorCode::OK, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, RejectsLegacySnapshotSidecar) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/snapshot/old/sequence_id", "42")); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_NE(ErrorCode::OK, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, InitDurablePrefixFailsClosedWhenBatchesExist) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, RejectsInvalidClusterId) { + FakeHaKvBackend backend; + OpLogBatchStorage storage("bad/cluster", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::INVALID_PARAMS, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, RereadsDurablePrefixWhenCreateIfAbsentLosesRace) { + FakeHaKvBackend backend; + backend.CreateBeforeNextTxn( + "/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 0, .last_seq = 0})); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::OK, storage.InitDurablePrefix(prefix)); + + EXPECT_EQ(0u, prefix.batch_id); + EXPECT_EQ(0u, prefix.last_seq); +} + +TEST(OpLogBatchStorageTest, ValidatesExistingDurablePrefix) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000003", + EncodeOpLogBatchRecord(MakeBatch(3, 7, 3)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 3, .last_seq = 9}))); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::OK, storage.InitDurablePrefix(prefix)); + + EXPECT_EQ(3u, prefix.batch_id); + EXPECT_EQ(9u, prefix.last_seq); +} + +TEST(OpLogBatchStorageTest, RejectsZeroPrefixWithExistingBatch) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 0, .last_seq = 0}))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, RejectsInconsistentDurablePrefixZeroFields) { + for (const DurablePrefix stored : + {DurablePrefix{.batch_id = 0, .last_seq = 9}, + DurablePrefix{.batch_id = 3, .last_seq = 0}}) { + SCOPED_TRACE("batch_id=" + std::to_string(stored.batch_id) + + " last_seq=" + std::to_string(stored.last_seq)); + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix(stored))); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.InitDurablePrefix(prefix)); + } +} + +TEST(OpLogBatchStorageTest, RejectsPrefixWithoutTerminalBatch) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 3, .last_seq = 9}))); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, RejectsTerminalBatchIdMismatch) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000003", + EncodeOpLogBatchRecord(MakeBatch(2, 7, 3)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 3, .last_seq = 9}))); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, RejectsTerminalBatchLastSequenceMismatch) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000003", + EncodeOpLogBatchRecord(MakeBatch(3, 7, 3)))); + ASSERT_EQ( + ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 3, .last_seq = 10}))); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, RejectsCorruptedTerminalBatch) { + FakeHaKvBackend backend; + std::string encoded = EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)); + const auto pos = + encoded.find_first_of("0123456789", encoded.rfind(',') + 1); + ASSERT_NE(std::string::npos, pos); + encoded[pos] = encoded[pos] == '0' ? '1' : '0'; + ASSERT_EQ( + ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", encoded)); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, RejectsInvalidTerminalBatchSequenceRange) { + FakeHaKvBackend backend; + auto batch = MakeBatch(1, 1, 1); + batch.last_seq = 2; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", + EncodeOpLogBatchRecord(batch))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 1, .last_seq = 2}))); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageTest, WriteBatchAndAdvancePrefixCommitsAtomically) { + FakeHaKvBackend backend; + const auto old_prefix = EncodeDurablePrefix({.batch_id = 1, .last_seq = 3}); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", old_prefix)); + OpLogBatchStorage storage("clusterA", backend); + + auto batch = MakeBatch(/*batch_id=*/2, /*first_seq=*/4, /*count=*/2); + EXPECT_EQ(ErrorCode::OK, storage.WriteBatchAndAdvancePrefix( + batch, {.batch_id = 1, .last_seq = 3})); + + OpLogBatchRecord stored_batch; + ASSERT_EQ(ErrorCode::OK, storage.ReadBatch(2, stored_batch)); + EXPECT_EQ(2u, stored_batch.batch_id); + EXPECT_EQ(4u, stored_batch.first_seq); + EXPECT_EQ(5u, stored_batch.last_seq); + + std::string encoded_prefix; + ASSERT_EQ(ErrorCode::OK, + backend.Get("/oplog/clusterA/durable_prefix", encoded_prefix)); + DurablePrefix prefix; + ASSERT_TRUE(DecodeDurablePrefix(encoded_prefix, &prefix)); + EXPECT_EQ(2u, prefix.batch_id); + EXPECT_EQ(5u, prefix.last_seq); +} + +TEST(OpLogBatchStorageTest, CompareFailureDoesNotWriteBatchOrAdvancePrefix) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 2, .last_seq = 6}))); + OpLogBatchStorage storage("clusterA", backend); + + auto batch = MakeBatch(/*batch_id=*/2, /*first_seq=*/4, /*count=*/1); + EXPECT_EQ(ErrorCode::ETCD_TRANSACTION_FAIL, + storage.WriteBatchAndAdvancePrefix( + batch, {.batch_id = 1, .last_seq = 3})); + + OpLogBatchRecord missing; + EXPECT_EQ(ErrorCode::ETCD_KEY_NOT_EXIST, storage.ReadBatch(2, missing)); + std::string encoded_prefix; + ASSERT_EQ(ErrorCode::OK, + backend.Get("/oplog/clusterA/durable_prefix", encoded_prefix)); + DurablePrefix prefix; + ASSERT_TRUE(DecodeDurablePrefix(encoded_prefix, &prefix)); + EXPECT_EQ(2u, prefix.batch_id); + EXPECT_EQ(6u, prefix.last_seq); +} + +TEST(OpLogBatchStorageTest, CompareFailureIsOkWhenTargetBatchAlreadyDurable) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 2, .last_seq = 5}))); + auto batch = MakeBatch(/*batch_id=*/2, /*first_seq=*/4, /*count=*/2); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000002", + EncodeOpLogBatchRecord(batch))); + backend.FailNextTxn(ErrorCode::ETCD_TRANSACTION_FAIL); + OpLogBatchStorage storage("clusterA", backend); + + EXPECT_EQ(ErrorCode::OK, storage.WriteBatchAndAdvancePrefix( + batch, {.batch_id = 1, .last_seq = 3})); +} + +TEST(OpLogBatchStorageTest, RejectsSkippedBatchId) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 1, .last_seq = 3}))); + OpLogBatchStorage storage("clusterA", backend); + + EXPECT_EQ(ErrorCode::INVALID_PARAMS, + storage.WriteBatchAndAdvancePrefix( + MakeBatch(/*batch_id=*/3, /*first_seq=*/4, /*count=*/1), + {.batch_id = 1, .last_seq = 3})); +} + +TEST(OpLogBatchStorageTest, RejectsSkippedSequenceRange) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 1, .last_seq = 3}))); + OpLogBatchStorage storage("clusterA", backend); + + EXPECT_EQ(ErrorCode::INVALID_PARAMS, + storage.WriteBatchAndAdvancePrefix( + MakeBatch(/*batch_id=*/2, /*first_seq=*/5, /*count=*/1), + {.batch_id = 1, .last_seq = 3})); +} + +TEST(OpLogBatchStorageTest, RejectsRegressedSequenceRange) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 1, .last_seq = 3}))); + OpLogBatchStorage storage("clusterA", backend); + + EXPECT_EQ(ErrorCode::INVALID_PARAMS, + storage.WriteBatchAndAdvancePrefix( + MakeBatch(/*batch_id=*/2, /*first_seq=*/3, /*count=*/1), + {.batch_id = 1, .last_seq = 3})); +} + +TEST(OpLogBatchStorageTest, RejectsMaxDurablePrefixAdvance) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = UINT64_MAX, + .last_seq = UINT64_MAX}))); + OpLogBatchStorage storage("clusterA", backend); + + EXPECT_EQ(ErrorCode::INVALID_PARAMS, + storage.WriteBatchAndAdvancePrefix( + MakeBatch(/*batch_id=*/0, /*first_seq=*/1, /*count=*/1), + {.batch_id = UINT64_MAX, .last_seq = UINT64_MAX})); +} + +TEST(OpLogBatchStorageTest, NonTxnBackendRejectsWriteBatchAndAdvancePrefix) { + FakeHaKvBackend backend; + backend.SetSupportsTxn(false); + OpLogBatchStorage storage("clusterA", backend); + + EXPECT_EQ(ErrorCode::INVALID_PARAMS, + storage.WriteBatchAndAdvancePrefix( + MakeBatch(/*batch_id=*/1, /*first_seq=*/1, /*count=*/1), + {.batch_id = 0, .last_seq = 0})); +} + +TEST(OpLogBatchStorageTest, ReadBatchRejectsCorruptedRecord) { + FakeHaKvBackend backend; + auto batch = MakeBatch(/*batch_id=*/1, /*first_seq=*/1, /*count=*/1); + std::string encoded = EncodeOpLogBatchRecord(batch); + auto pos = encoded.find_first_of("0123456789", encoded.rfind(',') + 1); + ASSERT_NE(std::string::npos, pos); + encoded[pos] = encoded[pos] == '0' ? '1' : '0'; + ASSERT_EQ( + ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", encoded)); + OpLogBatchStorage storage("clusterA", backend); + + OpLogBatchRecord out; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.ReadBatch(1, out)); +} + +TEST(OpLogBatchStorageTest, ReadBatchRejectsMismatchedPayloadBatchId) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", + EncodeOpLogBatchRecord(MakeBatch(2, 1, 1)))); + OpLogBatchStorage storage("clusterA", backend); + + OpLogBatchRecord out; + EXPECT_EQ(ErrorCode::INTERNAL_ERROR, storage.ReadBatch(1, out)); +} + +TEST(OpLogBatchStorageTest, ReadBatchesAfterReturnsOrderedBatches) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000002", + EncodeOpLogBatchRecord(MakeBatch(2, 4, 2)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", + EncodeOpLogBatchRecord(MakeBatch(1, 1, 3)))); + OpLogBatchStorage storage("clusterA", backend); + + std::vector batches; + EXPECT_EQ(ErrorCode::OK, storage.ReadBatchesAfter(/*after_batch_id=*/0, + /*limit=*/10, batches)); + + ASSERT_EQ(2u, batches.size()); + EXPECT_EQ(1u, batches[0].batch_id); + EXPECT_EQ(2u, batches[1].batch_id); +} + +TEST(OpLogBatchStorageTest, ReadBatchesAfterSkipsNonBatchKeysInRange) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/sidecar", "not a batch")); + OpLogBatchStorage storage("clusterA", backend); + + std::vector batches; + EXPECT_EQ(ErrorCode::OK, storage.ReadBatchesAfter(/*after_batch_id=*/0, + /*limit=*/10, batches)); + + ASSERT_EQ(1u, batches.size()); + EXPECT_EQ(1u, batches[0].batch_id); +} + +TEST(OpLogBatchStorageTest, ReadBatchesAfterLimitCountsOnlyBatchKeys) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000002", + EncodeOpLogBatchRecord(MakeBatch(2, 2, 1)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001.meta", + "sidecar")); + OpLogBatchStorage storage("clusterA", backend); + + std::vector batches; + EXPECT_EQ(ErrorCode::OK, storage.ReadBatchesAfter(/*after_batch_id=*/0, + /*limit=*/2, batches)); + + ASSERT_EQ(2u, batches.size()); + EXPECT_EQ(1u, batches[0].batch_id); + EXPECT_EQ(2u, batches[1].batch_id); + ASSERT_EQ(2u, backend.range_limits().size()); + EXPECT_EQ(2u, backend.range_limits()[0]); + EXPECT_EQ(1u, backend.range_limits()[1]); +} + +TEST(OpLogBatchStorageTest, ReadBatchesAfterHonorsLimit) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000001", + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/batches/00000000000000000002", + EncodeOpLogBatchRecord(MakeBatch(2, 2, 1)))); + OpLogBatchStorage storage("clusterA", backend); + + std::vector batches; + EXPECT_EQ(ErrorCode::OK, storage.ReadBatchesAfter(/*after_batch_id=*/0, + /*limit=*/1, batches)); + + ASSERT_EQ(1u, batches.size()); + EXPECT_EQ(1u, batches[0].batch_id); +} + +TEST(OpLogBatchStorageBackendErrorTest, PropagatesReadDurablePrefixError) { + FakeHaKvBackend backend; + backend.FailNextGet(ErrorCode::ETCD_OPERATION_ERROR); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::ETCD_OPERATION_ERROR, + storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageBackendErrorTest, + PropagatesZeroPrefixValidationRangeError) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 0, .last_seq = 0}))); + backend.FailNextRange(ErrorCode::ETCD_OPERATION_ERROR); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::ETCD_OPERATION_ERROR, + storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageBackendErrorTest, + PropagatesTerminalBatchReadErrorAtStartup) { + FakeHaKvBackend backend; + const std::string batch_key = + "/oplog/clusterA/batches/00000000000000000001"; + ASSERT_EQ( + ErrorCode::OK, + backend.Put(batch_key, EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + backend.FailNextGetForKey(batch_key, ErrorCode::ETCD_OPERATION_ERROR); + OpLogBatchStorage storage("clusterA", backend); + + DurablePrefix prefix; + EXPECT_EQ(ErrorCode::ETCD_OPERATION_ERROR, + storage.InitDurablePrefix(prefix)); +} + +TEST(OpLogBatchStorageBackendErrorTest, PropagatesRangeError) { + FakeHaKvBackend backend; + backend.FailNextRange(ErrorCode::ETCD_OPERATION_ERROR); + OpLogBatchStorage storage("clusterA", backend); + + std::vector batches; + EXPECT_EQ( + ErrorCode::ETCD_OPERATION_ERROR, + storage.ReadBatchesAfter(/*after_batch_id=*/0, /*limit=*/1, batches)); + EXPECT_TRUE(batches.empty()); +} + +TEST(OpLogBatchStorageBackendErrorTest, TxnErrorDoesNotAdvanceDurablePrefix) { + FakeHaKvBackend backend; + ASSERT_EQ(ErrorCode::OK, + backend.Put("/oplog/clusterA/durable_prefix", + EncodeDurablePrefix({.batch_id = 1, .last_seq = 3}))); + backend.FailNextTxn(ErrorCode::ETCD_OPERATION_ERROR); + OpLogBatchStorage storage("clusterA", backend); + + EXPECT_EQ(ErrorCode::ETCD_OPERATION_ERROR, + storage.WriteBatchAndAdvancePrefix( + MakeBatch(/*batch_id=*/2, /*first_seq=*/4, /*count=*/1), + {.batch_id = 1, .last_seq = 3})); + + std::string encoded_prefix; + ASSERT_EQ(ErrorCode::OK, + backend.Get("/oplog/clusterA/durable_prefix", encoded_prefix)); + DurablePrefix prefix; + ASSERT_TRUE(DecodeDurablePrefix(encoded_prefix, &prefix)); + EXPECT_EQ(1u, prefix.batch_id); + EXPECT_EQ(3u, prefix.last_seq); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/ha/oplog/oplog_manager_test.cpp b/mooncake-store/tests/ha/oplog/oplog_manager_test.cpp deleted file mode 100644 index ed4a69d1c9..0000000000 --- a/mooncake-store/tests/ha/oplog/oplog_manager_test.cpp +++ /dev/null @@ -1,323 +0,0 @@ -#include "ha/oplog/oplog_manager.h" - -#include -#include - -#include -#include -#include -#include -#include -#include - -namespace mooncake::test { - -class OpLogManagerTest : public ::testing::Test { - protected: - void SetUp() override { - google::InitGoogleLogging("OpLogManagerTest"); - FLAGS_logtostderr = 1; - manager_ = std::make_unique(); - } - - void TearDown() override { google::ShutdownGoogleLogging(); } - - OpLogManager& M() { return *manager_; } - - std::unique_ptr manager_; -}; - -// ========== 2.1.1 Basic functionality tests ========== - -TEST_F(OpLogManagerTest, TestAppendEntry) { - uint64_t id = M().Append(OpType::PUT_END, "key1", "value1"); - - EXPECT_LT(0u, id); - EXPECT_EQ(id, M().GetLastSequenceId()); - EXPECT_EQ(1u, M().GetEntryCount()); -} - -TEST_F(OpLogManagerTest, TestSequenceIdIncrement) { - uint64_t id1 = M().Append(OpType::PUT_END, "key1", "value1"); - uint64_t id2 = M().Append(OpType::PUT_END, "key2", "value2"); - uint64_t id3 = M().Append(OpType::REMOVE, "key3", ""); - - EXPECT_LT(0u, id1); - EXPECT_EQ(id1 + 1, id2); - EXPECT_EQ(id2 + 1, id3); - EXPECT_EQ(id3, M().GetLastSequenceId()); - EXPECT_EQ(3u, M().GetEntryCount()); -} - -TEST_F(OpLogManagerTest, TestAllocateEntry) { - OpLogEntry e1 = M().AllocateEntry(OpType::PUT_END, "key1", "value1"); - OpLogEntry e2 = M().AllocateEntry(OpType::PUT_END, "key2", "value2"); - - EXPECT_LT(0u, e1.sequence_id); - EXPECT_EQ(e1.sequence_id + 1, e2.sequence_id); - EXPECT_EQ(e2.sequence_id, M().GetLastSequenceId()); - EXPECT_EQ(2u, M().GetEntryCount()); - - // Basic field validation - EXPECT_EQ(OpType::PUT_END, e1.op_type); - EXPECT_EQ("key1", e1.object_key); - EXPECT_EQ("value1", e1.payload); - EXPECT_NE(0u, e1.timestamp_ms); - EXPECT_NE(0u, e1.checksum); - EXPECT_NE(0u, e1.prefix_hash); -} - -TEST_F(OpLogManagerTest, TestPersistEntry) { - // Without an OpLogStore configured, PersistEntry should return an - // error - OpLogEntry entry = - M().AllocateEntry(OpType::PUT_END, "key", "payload-data"); - - ErrorCode err = M().PersistEntry(entry); - EXPECT_EQ(ErrorCode::INTERNAL_ERROR, err); -} - -TEST_F(OpLogManagerTest, TestAppendAndPersist) { - // Without an EtcdOpLogStore configured, AppendAndPersist should return an - // error - auto res = M().AppendAndPersist(OpType::REMOVE, "key", ""); - ASSERT_FALSE(res.has_value()); - EXPECT_EQ(ErrorCode::INTERNAL_ERROR, res.error()); -} - -// ========== Initial Sequence Id Tests ========== - -TEST_F(OpLogManagerTest, SetInitialSequenceIdOnEmptyManager) { - EXPECT_EQ(0u, M().GetLastSequenceId()); - EXPECT_EQ(0u, M().GetEntryCount()); - - M().SetInitialSequenceId(100); - EXPECT_EQ(100u, M().GetLastSequenceId()); - - // The first appended entry should have sequence_id 101 - uint64_t id = M().Append(OpType::PUT_END, "key", "value"); - EXPECT_EQ(101u, id); -} - -TEST_F(OpLogManagerTest, SetInitialSequenceIdIgnoredWhenNotEmpty) { - uint64_t id1 = M().Append(OpType::PUT_END, "key1", "value1"); - EXPECT_EQ(id1, M().GetLastSequenceId()); - - // Setting initial sequence id on a non-empty manager should be ignored - M().SetInitialSequenceId(500); - EXPECT_EQ(id1, M().GetLastSequenceId()); -} - -// ========== 2.1.2 Checksum tests ========== - -TEST_F(OpLogManagerTest, TestChecksumComputation) { - // Same payload => same checksum; different payload => different checksum - OpLogEntry e1 = M().AllocateEntry(OpType::PUT_END, "k1", "payload-X"); - OpLogEntry e2 = M().AllocateEntry(OpType::PUT_END, "k2", "payload-X"); - OpLogEntry e3 = M().AllocateEntry(OpType::PUT_END, "k3", "payload-Y"); - - EXPECT_EQ(e1.checksum, e2.checksum); - EXPECT_NE(e1.checksum, e3.checksum); -} - -TEST_F(OpLogManagerTest, TestPrefixHashComputation) { - // Same key => same prefix_hash; different key => (with high probability) - // different prefix_hash - OpLogEntry e1 = M().AllocateEntry(OpType::PUT_END, "same-key", "v1"); - OpLogEntry e2 = M().AllocateEntry(OpType::PUT_END, "same-key", "v2"); - OpLogEntry e3 = M().AllocateEntry(OpType::PUT_END, "other-key", "v3"); - - EXPECT_EQ(e1.prefix_hash, e2.prefix_hash); - EXPECT_NE(e1.prefix_hash, e3.prefix_hash); -} - -TEST_F(OpLogManagerTest, TestVerifyChecksum) { - OpLogEntry entry = - M().AllocateEntry(OpType::PUT_END, "key", "payload-data"); - EXPECT_TRUE(OpLogManager::VerifyChecksum(entry)); - - // Verification should fail after tampering with the payload - entry.payload = "tampered"; - EXPECT_FALSE(OpLogManager::VerifyChecksum(entry)); -} - -// ========== 2.1.3 Size validation tests ========== - -TEST_F(OpLogManagerTest, TestValidateEntrySize_Valid) { - OpLogEntry entry; - entry.object_key = "normal-key"; - entry.payload = "small-payload"; - - std::string reason; - EXPECT_TRUE(OpLogManager::ValidateEntrySize(entry, &reason)); - EXPECT_TRUE(reason.empty()); -} - -TEST_F(OpLogManagerTest, TestValidateEntrySize_KeyTooLarge) { - OpLogEntry entry; - entry.object_key.assign(OpLogManager::kMaxObjectKeySize + 1, 'k'); - entry.payload = "payload"; - - std::string reason; - EXPECT_FALSE(OpLogManager::ValidateEntrySize(entry, &reason)); - EXPECT_FALSE(reason.empty()); -} - -TEST_F(OpLogManagerTest, TestValidateEntrySize_PayloadTooLarge) { - OpLogEntry entry; - entry.object_key = "key"; - entry.payload.assign(OpLogManager::kMaxPayloadSize + 1, 'p'); - - std::string reason; - EXPECT_FALSE(OpLogManager::ValidateEntrySize(entry, &reason)); - EXPECT_FALSE(reason.empty()); -} - -TEST_F(OpLogManagerTest, TestValidateEntrySize_EmptyKey) { - // Current implementation only enforces upper bounds; empty keys are - // accepted - OpLogEntry entry; - entry.object_key = ""; - entry.payload = "payload"; - - std::string reason; - EXPECT_TRUE(OpLogManager::ValidateEntrySize(entry, &reason)); -} - -// ========== 2.1.4 Etcd integration tests (placeholder) ========== - -TEST_F(OpLogManagerTest, TestWriteToEtcd_Success) { -#if defined(STORE_USE_ETCD) - GTEST_SKIP() - << "TODO: requires real EtcdOpLogStore and running etcd cluster."; -#else - GTEST_SKIP() << "STORE_USE_ETCD is disabled."; -#endif -} - -TEST_F(OpLogManagerTest, TestWriteToEtcd_Failure) { - OpLogEntry entry = - M().AllocateEntry(OpType::PUT_END, "key", "payload-data"); - ErrorCode err = M().PersistEntry(entry); - EXPECT_EQ(ErrorCode::INTERNAL_ERROR, err); -} - -TEST_F(OpLogManagerTest, TestWriteToEtcd_Retry) { - GTEST_SKIP() << "TODO: retry / idempotent semantics are tested at " - "EtcdOpLogStore level."; -} - -TEST_F(OpLogManagerTest, TestIdempotentWrite) { - GTEST_SKIP() << "TODO: idempotent write belongs to " - "EtcdOpLogStore::WriteOpLog tests."; -} - -// ========== 2.1.5 Boundary condition tests ========== - -TEST_F(OpLogManagerTest, TestSequenceIdWrapAround) { - // Theoretical wrap-around test: set initial value near UINT64_MAX and - // verify wrap-around semantics - uint64_t near_max = std::numeric_limits::max() - 2; - M().SetInitialSequenceId(near_max); - - std::vector ids; - ids.push_back(M().Append(OpType::PUT_END, "k1", "v1")); // max-1 - ids.push_back(M().Append(OpType::PUT_END, "k2", "v2")); // max - ids.push_back(M().Append(OpType::PUT_END, "k3", "v3")); // 0 (wrap) - - ASSERT_EQ(3u, ids.size()); - // Use wrap-around-safe comparison helpers to verify monotonic increase - EXPECT_TRUE(IsSequenceNewer(ids[1], ids[0])); - EXPECT_TRUE(IsSequenceNewer( - ids[2], ids[1])); // 0 is considered newer than UINT64_MAX -} - -TEST_F(OpLogManagerTest, TestConcurrentAppend) { - constexpr int kThreads = 8; - constexpr int kPerThread = 1000; - - std::vector ids; - ids.reserve(kThreads * kPerThread); - std::mutex m; - - auto worker = [&]() { - for (int i = 0; i < kPerThread; ++i) { - uint64_t id = M().Append(OpType::PUT_END, "key", "value"); - std::lock_guard lock(m); - ids.push_back(id); - } - }; - - std::vector threads; - for (int i = 0; i < kThreads; ++i) { - threads.emplace_back(worker); - } - for (auto& t : threads) { - t.join(); - } - - EXPECT_EQ(static_cast(kThreads * kPerThread), ids.size()); - - std::sort(ids.begin(), ids.end()); - // Ensure there are no duplicates and sequence IDs are strictly increasing - for (size_t i = 1; i < ids.size(); ++i) { - EXPECT_GT(ids[i], ids[i - 1]); - } -} - -TEST_F(OpLogManagerTest, TestLargePayload) { - // Construct a payload close to the upper limit and verify it passes - // validation and appends successfully - std::string key = "large-payload-key"; - std::string payload(OpLogManager::kMaxPayloadSize - 1, 'x'); - - OpLogEntry entry; - entry.object_key = key; - entry.payload = payload; - - std::string reason; - EXPECT_TRUE(OpLogManager::ValidateEntrySize(entry, &reason)); - - uint64_t id = M().Append(OpType::PUT_END, key, payload); - EXPECT_LT(0u, id); - EXPECT_EQ(1u, M().GetEntryCount()); -} - -TEST_F(OpLogManagerTest, TestAppendMultipleTypes) { - uint64_t id1 = M().Append(OpType::PUT_END, "k1", "payload"); - uint64_t id2 = M().Append(OpType::PUT_REVOKE, "k2", ""); - uint64_t id3 = M().Append(OpType::REMOVE, "k3", ""); - EXPECT_EQ(id1 + 1, id2); - EXPECT_EQ(id2 + 1, id3); - EXPECT_EQ(3u, M().GetEntryCount()); -} - -TEST_F(OpLogManagerTest, TestBufferEviction) { - for (int i = 0; i < 10; ++i) { - M().Append(OpType::PUT_END, "key_" + std::to_string(i), "val"); - } - EXPECT_EQ(10u, M().GetEntryCount()); -} - -TEST_F(OpLogManagerTest, TestAllocateEntry_ChecksumValid) { - OpLogEntry e = M().AllocateEntry(OpType::REMOVE, "del-key", "some-data"); - EXPECT_TRUE(OpLogManager::VerifyChecksum(e)); -} - -TEST_F(OpLogManagerTest, TestAllocateEntry_FieldsComplete) { - OpLogEntry e = M().AllocateEntry(OpType::REMOVE, "del-key", ""); - EXPECT_EQ(OpType::REMOVE, e.op_type); - EXPECT_EQ("del-key", e.object_key); - EXPECT_EQ("", e.payload); - EXPECT_NE(0u, e.timestamp_ms); - EXPECT_TRUE(OpLogManager::VerifyChecksum(e)); - // "del-key" prefix hash - EXPECT_NE(0u, e.prefix_hash); -} - -} // namespace mooncake::test - -int main(int argc, char** argv) { - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/mooncake-store/tests/ha/oplog/oplog_replicator_test.cpp b/mooncake-store/tests/ha/oplog/oplog_replicator_test.cpp deleted file mode 100644 index 7cb148622b..0000000000 --- a/mooncake-store/tests/ha/oplog/oplog_replicator_test.cpp +++ /dev/null @@ -1,180 +0,0 @@ -#include "ha/oplog/oplog_replicator.h" - -#include -#include - -#include -#include -#include -#include -#include - -#include - -#include "metadata_store.h" -#include "ha/oplog/oplog_applier.h" -#include "ha/oplog/oplog_change_notifier.h" -#include "ha/oplog/oplog_manager.h" -#include "ha/oplog/oplog_serializer.h" -#include "types.h" - -namespace mooncake::test { - -// Minimal MetadataStore implementation for OpLogApplier -class MinimalMockMetadataStore : public MetadataStore { - public: - bool PutMetadata(const std::string&, - const StandbyObjectMetadata&) override { - return true; - } - bool Put(const std::string&, const std::string&) override { return true; } - std::optional GetMetadata( - const std::string&) const override { - return std::nullopt; - } - bool Remove(const std::string&) override { return true; } - bool Exists(const std::string&) const override { return false; } - size_t GetKeyCount() const override { return 0; } -}; - -// In-memory MockOpLogChangeNotifier for tests -class MockOpLogChangeNotifier : public OpLogChangeNotifier { - public: - ErrorCode Start(uint64_t /*start_seq*/, EntryCallback on_entry, - ErrorCallback on_error) override { - on_entry_ = std::move(on_entry); - on_error_ = std::move(on_error); - healthy_ = true; - return ErrorCode::OK; - } - void Stop() override { healthy_ = false; } - bool IsHealthy() const override { return healthy_; } - - // Test helpers - void InjectEntry(const OpLogEntry& entry) { - if (on_entry_) on_entry_(entry); - } - void InjectError(ErrorCode err) { - if (on_error_) on_error_(err); - } - - private: - EntryCallback on_entry_; - ErrorCallback on_error_; - bool healthy_{false}; -}; - -// Helper function to create a valid OpLogEntry with checksum -OpLogEntry MakeEntry(uint64_t seq, OpType type, const std::string& key, - const std::string& payload) { - OpLogEntry e; - e.sequence_id = seq; - e.timestamp_ms = std::chrono::duration_cast( - std::chrono::steady_clock::now().time_since_epoch()) - .count(); - e.op_type = type; - e.object_key = key; - e.payload = payload; - e.checksum = - static_cast(XXH32(payload.data(), payload.size(), 0)); - e.prefix_hash = - key.empty() ? 0 - : static_cast(XXH32(key.data(), key.size(), 0)); - return e; -} - -class OpLogReplicatorTest : public ::testing::Test { - protected: - void SetUp() override { - google::InitGoogleLogging("OpLogReplicatorTest"); - FLAGS_logtostderr = 1; - metadata_store_ = std::make_unique(); - applier_ = - std::make_unique(metadata_store_.get(), "test"); - notifier_ = std::make_unique(); - replicator_ = - std::make_unique(notifier_.get(), applier_.get()); - } - - void TearDown() override { - if (replicator_) { - replicator_->Stop(); - } - google::ShutdownGoogleLogging(); - } - - MockOpLogChangeNotifier& Notifier() { return *notifier_; } - - std::unique_ptr metadata_store_; - std::unique_ptr applier_; - std::unique_ptr notifier_; - std::unique_ptr replicator_; -}; - -// ========== Start/Stop tests ========== - -TEST_F(OpLogReplicatorTest, TestStartStop) { - EXPECT_TRUE(replicator_->StartFromSequenceId(0)); - EXPECT_TRUE(replicator_->IsHealthy()); - replicator_->Stop(); - EXPECT_FALSE(replicator_->IsHealthy()); -} - -TEST_F(OpLogReplicatorTest, TestStartFromSequenceId) { - EXPECT_TRUE(replicator_->StartFromSequenceId(100)); - EXPECT_TRUE(replicator_->IsHealthy()); -} - -// ========== Entry delivery tests ========== - -TEST_F(OpLogReplicatorTest, InjectEntry_Applied) { - replicator_->StartFromSequenceId(0); - - OpLogEntry entry = MakeEntry(1, OpType::PUT_END, "key1", "payload1"); - Notifier().InjectEntry(entry); - - EXPECT_EQ(1u, replicator_->GetLastProcessedSequenceId()); - EXPECT_EQ(2u, applier_->GetExpectedSequenceId()); -} - -TEST_F(OpLogReplicatorTest, InjectMultipleEntries_SequenceTracking) { - replicator_->StartFromSequenceId(0); - - Notifier().InjectEntry(MakeEntry(1, OpType::PUT_END, "k1", "p1")); - Notifier().InjectEntry(MakeEntry(2, OpType::PUT_END, "k2", "p2")); - Notifier().InjectEntry(MakeEntry(3, OpType::REMOVE, "k3", "")); - - EXPECT_EQ(3u, replicator_->GetLastProcessedSequenceId()); - EXPECT_EQ(4u, applier_->GetExpectedSequenceId()); -} - -TEST_F(OpLogReplicatorTest, InjectError_NotifiesCallback) { - bool error_received = false; - replicator_->SetStateCallback([&](StandbyEvent event) { - if (event == StandbyEvent::WATCH_BROKEN) { - error_received = true; - } - }); - - replicator_->StartFromSequenceId(0); - Notifier().InjectError(ErrorCode::ETCD_OPERATION_ERROR); - - EXPECT_TRUE(error_received); -} - -// ========== Utility tests ========== - -TEST_F(OpLogReplicatorTest, GetLastProcessedSequenceId_InitiallyZero) { - EXPECT_EQ(0u, replicator_->GetLastProcessedSequenceId()); -} - -TEST_F(OpLogReplicatorTest, IsHealthy_FalseBeforeStart) { - EXPECT_FALSE(replicator_->IsHealthy()); -} - -} // namespace mooncake::test - -int main(int argc, char** argv) { - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/mooncake-store/tests/ha/oplog/oplog_serializer_test.cpp b/mooncake-store/tests/ha/oplog/oplog_serializer_test.cpp deleted file mode 100644 index c27dc7496b..0000000000 --- a/mooncake-store/tests/ha/oplog/oplog_serializer_test.cpp +++ /dev/null @@ -1,142 +0,0 @@ -#include "ha/oplog/oplog_serializer.h" - -#include -#include -#include - -#include - -namespace mooncake::test { - -class OpLogSerializerTest : public ::testing::Test { - protected: - void SetUp() override { - google::InitGoogleLogging("OpLogSerializerTest"); - FLAGS_logtostderr = 1; - } - void TearDown() override { google::ShutdownGoogleLogging(); } - - static OpLogEntry MakeEntry(uint64_t seq, OpType type, - const std::string& key, - const std::string& payload) { - OpLogEntry e; - e.sequence_id = seq; - e.timestamp_ms = 1234567890; - e.op_type = type; - e.object_key = key; - e.payload = payload; - e.checksum = - static_cast(XXH32(payload.data(), payload.size(), 0)); - e.prefix_hash = - key.empty() - ? 0 - : static_cast(XXH32(key.data(), key.size(), 0)); - return e; - } -}; - -TEST_F(OpLogSerializerTest, RoundTrip_PutEnd) { - OpLogEntry in = MakeEntry(1, OpType::PUT_END, "key1", "value1"); - std::string json = SerializeOpLogEntry(in); - OpLogEntry out; - ASSERT_TRUE(DeserializeOpLogEntry(json, out)); - EXPECT_EQ(in.sequence_id, out.sequence_id); - EXPECT_EQ(in.timestamp_ms, out.timestamp_ms); - EXPECT_EQ(in.op_type, out.op_type); - EXPECT_EQ(in.object_key, out.object_key); - EXPECT_EQ(in.payload, out.payload); - EXPECT_EQ(in.checksum, out.checksum); - EXPECT_EQ(in.prefix_hash, out.prefix_hash); -} - -TEST_F(OpLogSerializerTest, RoundTrip_Remove) { - OpLogEntry in = MakeEntry(42, OpType::REMOVE, "obj/to/remove", ""); - std::string json = SerializeOpLogEntry(in); - OpLogEntry out; - ASSERT_TRUE(DeserializeOpLogEntry(json, out)); - EXPECT_EQ(in.sequence_id, out.sequence_id); - EXPECT_EQ(in.op_type, out.op_type); - EXPECT_EQ(in.object_key, out.object_key); - EXPECT_EQ(in.payload, out.payload); -} - -TEST_F(OpLogSerializerTest, RoundTrip_PutRevoke) { - OpLogEntry in = MakeEntry(99, OpType::PUT_REVOKE, "revoked_key", "meta"); - std::string json = SerializeOpLogEntry(in); - OpLogEntry out; - ASSERT_TRUE(DeserializeOpLogEntry(json, out)); - EXPECT_EQ(in.op_type, out.op_type); - EXPECT_EQ(in.payload, out.payload); -} - -TEST_F(OpLogSerializerTest, RoundTrip_BinaryPayload) { - // Payload with null bytes, high bytes — must survive base64 round-trip - std::string binary_payload; - for (int i = 0; i < 256; ++i) { - binary_payload.push_back(static_cast(i)); - } - OpLogEntry in = MakeEntry(7, OpType::PUT_END, "bin_key", binary_payload); - std::string json = SerializeOpLogEntry(in); - OpLogEntry out; - ASSERT_TRUE(DeserializeOpLogEntry(json, out)); - EXPECT_EQ(in.payload, out.payload); -} - -TEST_F(OpLogSerializerTest, RoundTrip_EmptyPayload) { - OpLogEntry in = MakeEntry(10, OpType::REMOVE, "key", ""); - std::string json = SerializeOpLogEntry(in); - OpLogEntry out; - ASSERT_TRUE(DeserializeOpLogEntry(json, out)); - EXPECT_EQ("", out.payload); -} - -TEST_F(OpLogSerializerTest, RoundTrip_EmptyKey) { - OpLogEntry in = MakeEntry(11, OpType::PUT_END, "", "payload"); - std::string json = SerializeOpLogEntry(in); - OpLogEntry out; - ASSERT_TRUE(DeserializeOpLogEntry(json, out)); - EXPECT_EQ("", out.object_key); - EXPECT_EQ(0u, out.prefix_hash); -} - -TEST_F(OpLogSerializerTest, Deserialize_InvalidJson) { - OpLogEntry out; - EXPECT_FALSE(DeserializeOpLogEntry("{ not valid json }", out)); -} - -TEST_F(OpLogSerializerTest, Deserialize_EmptyString) { - OpLogEntry out; - EXPECT_FALSE(DeserializeOpLogEntry("", out)); -} - -TEST_F(OpLogSerializerTest, Deserialize_MissingFields) { - // JSON with only partial fields — should fail or produce defaults - std::string partial = R"({"sequence_id": 1})"; - OpLogEntry out; - // Behavior depends on JsonCpp defaults for missing fields. - // At minimum, should not crash. - (void)DeserializeOpLogEntry(partial, out); -} - -TEST_F(OpLogSerializerTest, Deserialize_KeyTooLarge) { - OpLogEntry in = MakeEntry(1, OpType::PUT_END, "k", "v"); - in.object_key.assign(OpLogManager::kMaxObjectKeySize + 1, 'k'); - std::string json = SerializeOpLogEntry(in); - OpLogEntry out; - EXPECT_FALSE(DeserializeOpLogEntry(json, out)); -} - -TEST_F(OpLogSerializerTest, Deserialize_PayloadTooLarge) { - OpLogEntry in = MakeEntry(1, OpType::PUT_END, "k", ""); - in.payload.assign(OpLogManager::kMaxPayloadSize + 1, 'p'); - std::string json = SerializeOpLogEntry(in); - OpLogEntry out; - EXPECT_FALSE(DeserializeOpLogEntry(json, out)); -} - -} // namespace mooncake::test - -int main(int argc, char** argv) { - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/mooncake-store/tests/ha/oplog/oplog_test_failpoint_test.cpp b/mooncake-store/tests/ha/oplog/oplog_test_failpoint_test.cpp new file mode 100644 index 0000000000..ebc630bfa7 --- /dev/null +++ b/mooncake-store/tests/ha/oplog/oplog_test_failpoint_test.cpp @@ -0,0 +1,69 @@ +#include "ha/oplog/oplog_test_failpoint.h" + +#include + +#include +#include +#include +#include +#include +#include + +#include + +namespace mooncake::test { +namespace { + +class OpLogTestFailPointTest : public ::testing::Test { + protected: + void SetUp() override { + dir_ = "/tmp/mooncake-oplog-failpoint-" + std::to_string(::getpid()); + std::filesystem::create_directories(dir_); + setenv("MOONCAKE_TEST_FAILPOINT_DIR", dir_.c_str(), 1); + setenv("MOONCAKE_TEST_FAILPOINT_TIMEOUT_SEC", "1", 1); + } + + void TearDown() override { + unsetenv("MOONCAKE_TEST_FAILPOINT_DIR"); + unsetenv("MOONCAKE_TEST_FAILPOINT_TIMEOUT_SEC"); + std::filesystem::remove_all(dir_); + } + + std::string dir_; +}; + +TEST_F(OpLogTestFailPointTest, UnarmedPointReturnsImmediately) { + const auto started = std::chrono::steady_clock::now(); + EXPECT_FALSE(TestFailPoint::Wait("batch_txn_succeeded_before_callback")); + EXPECT_LT(std::chrono::steady_clock::now() - started, + std::chrono::milliseconds(100)); +} + +TEST_F(OpLogTestFailPointTest, CreatesHitAndWaitsForRelease) { + const std::string name = "batch_txn_succeeded_before_callback"; + std::ofstream(dir_ + "/" + name + ".arm").put('\n'); + std::atomic completed{false}; + std::thread waiter([&] { + EXPECT_TRUE(TestFailPoint::Wait(name)); + completed = true; + }); + + const auto hit = std::filesystem::path(dir_) / (name + ".hit"); + for (int i = 0; i < 100 && !std::filesystem::exists(hit); ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_TRUE(std::filesystem::exists(hit)); + const auto claim = std::filesystem::path(dir_) / + (name + ".claimed." + std::to_string(::getpid())); + EXPECT_FALSE(std::filesystem::exists(dir_ + "/" + name + ".arm")); + EXPECT_TRUE(std::filesystem::exists(claim)); + EXPECT_FALSE(completed.load()); + std::ofstream(dir_ + "/" + name + ".release").put('\n'); + waiter.join(); + EXPECT_TRUE(completed.load()); + EXPECT_FALSE(std::filesystem::exists(hit)); + EXPECT_FALSE(std::filesystem::exists(claim)); +} + +} // namespace +} // namespace mooncake::test diff --git a/mooncake-store/tests/ha/oplog/oplog_types_test.cpp b/mooncake-store/tests/ha/oplog/oplog_types_test.cpp new file mode 100644 index 0000000000..641dec4cd5 --- /dev/null +++ b/mooncake-store/tests/ha/oplog/oplog_types_test.cpp @@ -0,0 +1,26 @@ +#include "ha/oplog/oplog_types.h" + +#include + +namespace mooncake::test { + +TEST(OpLogTypesTest, ChecksumRoundTrips) { + OpLogEntry entry; + entry.payload = "payload"; + entry.checksum = ComputeOpLogChecksum(entry.payload); + EXPECT_TRUE(VerifyOpLogChecksum(entry)); +} + +TEST(OpLogTypesTest, RejectsOversizedEntry) { + OpLogEntry entry; + entry.object_key.assign(kMaxOpLogObjectKeySize + 1, 'x'); + EXPECT_FALSE(ValidateOpLogEntrySize(entry)); +} + +TEST(OpLogTypesTest, NormalizesTrailingClusterSlashes) { + std::string cluster_id = "cluster///"; + EXPECT_TRUE(NormalizeAndValidateClusterId(cluster_id)); + EXPECT_EQ("cluster", cluster_id); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/ha/oplog/ordered_oplog_writer_test.cpp b/mooncake-store/tests/ha/oplog/ordered_oplog_writer_test.cpp new file mode 100644 index 0000000000..29c2c011a7 --- /dev/null +++ b/mooncake-store/tests/ha/oplog/ordered_oplog_writer_test.cpp @@ -0,0 +1,1090 @@ +#include "ha/oplog/ordered_oplog_writer.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS +#include "ha_metric_manager.h" +#endif + +namespace mooncake::test { +namespace { + +class FakeBatchWriter { + public: + using Clock = std::chrono::steady_clock; + + ErrorCode Write(const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + std::unique_lock lock(mutex_); + attempt_times_.push_back(Clock::now()); + attempt_cv_.notify_all(); + while (blocked_) { + blocked_write_active_ = true; + blocked_cv_.notify_all(); + blocked_cv_.wait(lock); + } + blocked_write_active_ = false; + if (next_error_ != ErrorCode::OK) { + ErrorCode err = next_error_; + next_error_ = ErrorCode::OK; + return err; + } + if (failures_remaining_ > 0) { + --failures_remaining_; + return repeated_error_; + } + writes_.push_back({.batch = batch, .expected_prefix = expected_prefix}); + cv_.notify_all(); + return ErrorCode::OK; + } + + std::vector Batches() const { + std::lock_guard lock(mutex_); + std::vector batches; + batches.reserve(writes_.size()); + for (const auto& write : writes_) { + batches.push_back(write.batch); + } + return batches; + } + + bool WaitForWrites(size_t count, std::chrono::milliseconds timeout = + std::chrono::milliseconds(1000)) { + std::unique_lock lock(mutex_); + return cv_.wait_for(lock, timeout, + [&] { return writes_.size() >= count; }); + } + + bool WaitForAttempts(size_t count, std::chrono::milliseconds timeout = + std::chrono::milliseconds(5000)) { + std::unique_lock lock(mutex_); + return attempt_cv_.wait_for( + lock, timeout, [&] { return attempt_times_.size() >= count; }); + } + + std::vector AttemptTimes() const { + std::lock_guard lock(mutex_); + return attempt_times_; + } + + void FailNextWrite(ErrorCode err) { + std::lock_guard lock(mutex_); + next_error_ = err; + } + void FailNextWrites(size_t count, ErrorCode err) { + std::lock_guard lock(mutex_); + failures_remaining_ = count; + repeated_error_ = err; + } + void AllowWrites() { + std::lock_guard lock(mutex_); + next_error_ = ErrorCode::OK; + failures_remaining_ = 0; + } + void BlockWrites() { + std::lock_guard lock(mutex_); + blocked_ = true; + } + void UnblockWrites() { + { + std::lock_guard lock(mutex_); + blocked_ = false; + } + blocked_cv_.notify_all(); + } + bool WaitForBlockedWrite( + std::chrono::milliseconds timeout = std::chrono::milliseconds(1000)) { + std::unique_lock lock(mutex_); + return blocked_cv_.wait_for(lock, timeout, + [&] { return blocked_write_active_; }); + } + + private: + struct WriteRecord { + OpLogBatchRecord batch; + DurablePrefix expected_prefix; + }; + + mutable std::mutex mutex_; + std::condition_variable cv_; + std::condition_variable attempt_cv_; + std::condition_variable blocked_cv_; + std::vector writes_; + std::vector attempt_times_; + ErrorCode next_error_{ErrorCode::OK}; + ErrorCode repeated_error_{ErrorCode::OK}; + size_t failures_remaining_{0}; + bool blocked_{false}; + bool blocked_write_active_{false}; +}; + +OpLogEntry MakeEntry(std::string key = "key", std::string payload = "value") { + OpLogEntry entry; + entry.timestamp_ms = 1234567890; + entry.op_type = OpType::PUT_END; + entry.tenant_id = "tenant"; + entry.object_key = std::move(key); + entry.payload = std::move(payload); + entry.checksum = static_cast( + XXH32(entry.payload.data(), entry.payload.size(), 0)); + entry.prefix_hash = static_cast( + XXH32(entry.object_key.data(), entry.object_key.size(), 0)); + return entry; +} + +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS +bool WaitForMetric(const std::function& predicate) { + for (int i = 0; i < 100; ++i) { + if (predicate()) { + return true; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + return false; +} +#endif + +} // namespace + +TEST(OrderedOpLogWriterAdmissionTest, AbortLeavesNoSequenceGap) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + auto first = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + writer.Abort(std::move(*first)); + + auto second = writer.Reserve(); + ASSERT_TRUE(second.has_value()); + auto pending = writer.Commit(std::move(*second), MakeEntry(), + [](const OpLogEntry&) {}); + ASSERT_TRUE(pending.has_value()); + + EXPECT_EQ(1u, pending->sequence_id()); +} + +TEST(OrderedOpLogWriterAdmissionTest, CommitAssignsContiguousSequences) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 3}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + auto first = writer.Reserve(); + auto second = writer.Reserve(); + auto third = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(second.has_value()); + ASSERT_TRUE(third.has_value()); + + auto p1 = + writer.Commit(std::move(*first), MakeEntry("k1"), [](const auto&) {}); + auto p2 = + writer.Commit(std::move(*second), MakeEntry("k2"), [](const auto&) {}); + auto p3 = + writer.Commit(std::move(*third), MakeEntry("k3"), [](const auto&) {}); + + ASSERT_TRUE(p1.has_value()); + ASSERT_TRUE(p2.has_value()); + ASSERT_TRUE(p3.has_value()); + EXPECT_EQ(1u, p1->sequence_id()); + EXPECT_EQ(2u, p2->sequence_id()); + EXPECT_EQ(3u, p3->sequence_id()); +} + +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS +TEST(OrderedOpLogWriterMetricTest, RecordsDurabilityQueuesAndRetry) { + FakeBatchWriter storage; + storage.FailNextWrites(1, ErrorCode::ETCD_TRANSACTION_FAIL); + auto& metrics = HAMetricManager::instance(); + const auto batches_before = + metrics.get_batch_record_durable_batches_total(); + const auto entries_before = + metrics.get_batch_record_durable_entries_total(); + const auto retries_before = metrics.get_batch_record_retries_total(); + + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + std::atomic callback_done{false}; + ASSERT_TRUE(writer + .Commit(std::move(*reservation), MakeEntry(), + [&](const OpLogEntry&) { callback_done = true; }) + .has_value()); + + ASSERT_TRUE(WaitForMetric([&] { return callback_done.load(); })); + writer.Stop(); + + EXPECT_EQ(batches_before + 1, + metrics.get_batch_record_durable_batches_total()); + EXPECT_EQ(entries_before + 1, + metrics.get_batch_record_durable_entries_total()); + EXPECT_EQ(retries_before + 1, metrics.get_batch_record_retries_total()); + EXPECT_EQ(1, metrics.get_batch_record_last_batch_id()); + EXPECT_EQ(1, metrics.get_batch_record_durable_sequence()); + EXPECT_EQ(0, metrics.get_batch_record_committed_queue_depth()); + EXPECT_EQ(0, metrics.get_batch_record_callback_queue_depth()); +} +#endif + +TEST(OrderedOpLogWriterAdmissionTest, + InvalidEntryDoesNotConsumeSequenceOrInvokeCallback) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + auto invalid_reservation = writer.Reserve(); + ASSERT_TRUE(invalid_reservation.has_value()); + auto invalid = MakeEntry("bad"); + invalid.tenant_id = "_reserved"; + bool callback_called = false; + auto rejected = + writer.Commit(std::move(*invalid_reservation), std::move(invalid), + [&](const OpLogEntry&) { callback_called = true; }); + ASSERT_FALSE(rejected.has_value()); + EXPECT_EQ(ErrorCode::INVALID_PARAMS, rejected.error()); + EXPECT_FALSE(callback_called); + + auto valid_reservation = writer.Reserve(); + ASSERT_TRUE(valid_reservation.has_value()); + auto accepted = writer.Commit(std::move(*valid_reservation), MakeEntry(), + [](const OpLogEntry&) {}); + ASSERT_TRUE(accepted.has_value()); + EXPECT_EQ(1u, accepted->sequence_id()); +} + +TEST(OrderedOpLogWriterAdmissionTest, RejectsOpTypeOutsideEnumRange) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 1}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + auto invalid = MakeEntry(); + invalid.op_type = static_cast(255); + auto rejected = writer.Commit(std::move(*reservation), std::move(invalid), + [](const OpLogEntry&) {}); + + ASSERT_FALSE(rejected.has_value()); + EXPECT_EQ(ErrorCode::INVALID_PARAMS, rejected.error()); +} + +TEST(OrderedOpLogWriterAdmissionTest, + ReserveFailsWhenOpenWaitingSlotsReachMax) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + auto first = writer.Reserve(); + auto second = writer.Reserve(); + EXPECT_TRUE(first.has_value()); + EXPECT_TRUE(second.has_value()); + auto third = writer.Reserve(); + + ASSERT_FALSE(third.has_value()); + EXPECT_EQ(ErrorCode::TASK_PENDING_LIMIT_EXCEEDED, third.error()); +} + +TEST(OrderedOpLogWriterAdmissionTest, + ReserveFailsWhenInitialDurablePrefixCannotAdvance) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{ + .max_entries_per_batch = 2, + .initial_durable_prefix = {.batch_id = 1, .last_seq = UINT64_MAX}}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + auto reservation = writer.Reserve(); + + ASSERT_FALSE(reservation.has_value()); + EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS, reservation.error()); + EXPECT_EQ(ErrorCode::INVALID_PARAMS, writer.LastError()); +} + +TEST(OrderedOpLogWriterAdmissionTest, MoveAssigningReservationReleasesOldSlot) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + auto first = writer.Reserve(); + auto second = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(second.has_value()); + + *first = std::move(*second); + EXPECT_TRUE(writer.Reserve().has_value()); +} + +TEST(OrderedOpLogWriterAdmissionTest, DestroyingReservationReleasesSlot) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 1}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + { + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + } + + EXPECT_TRUE(writer.Reserve().has_value()); +} + +TEST(OrderedOpLogWriterAdmissionTest, + SealingCommittedEntriesFreesWaitingSlots) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 1}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + auto first = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*first), MakeEntry("k1"), [](const auto&) {}) + .has_value()); + ASSERT_TRUE(storage.WaitForWrites(1)); + + auto second = writer.Reserve(); + EXPECT_TRUE(second.has_value()); + writer.Stop(); +} + +TEST(OrderedOpLogWriterAdmissionTest, + FirstCommitFreesOpenSlotBeforeWriterThreadRuns) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 1}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + auto first = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*first), MakeEntry("k1"), [](const auto&) {}) + .has_value()); + + auto second = writer.Reserve(); + ASSERT_TRUE(second.has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*second), MakeEntry("k2"), [](const auto&) {}) + .has_value()); + + writer.Start(); + ASSERT_TRUE(storage.WaitForWrites(2)); + writer.Stop(); +} + +TEST(OrderedOpLogWriterAdmissionTest, StopClosesAdmission) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 1}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + writer.Start(); + writer.Stop(); + + EXPECT_FALSE(writer.IsAccepting()); + auto reservation = writer.Reserve(); + ASSERT_FALSE(reservation.has_value()); + EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS, reservation.error()); +} + +TEST(OrderedOpLogWriterAdmissionTest, StopRejectsOutstandingReservation) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 1}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + writer.Start(); + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + writer.Stop(); + + auto pending = + writer.Commit(std::move(*reservation), MakeEntry(), [](const auto&) {}); + ASSERT_FALSE(pending.has_value()); + EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS, pending.error()); + EXPECT_TRUE(storage.Batches().empty()); +} + +TEST(OrderedOpLogWriterAdmissionTest, + OpenBatchCapacityRecoversAfterInflightWrite) { + FakeBatchWriter storage; + storage.BlockWrites(); + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + auto first = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*first), MakeEntry("k1"), [](const auto&) {}) + .has_value()); + ASSERT_TRUE(storage.WaitForBlockedWrite()); + + auto second = writer.Reserve(); + auto third = writer.Reserve(); + ASSERT_TRUE(second.has_value()); + ASSERT_TRUE(third.has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*second), MakeEntry("k2"), [](const auto&) {}) + .has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*third), MakeEntry("k3"), [](const auto&) {}) + .has_value()); + + auto full = writer.Reserve(); + ASSERT_FALSE(full.has_value()); + EXPECT_EQ(ErrorCode::TASK_PENDING_LIMIT_EXCEEDED, full.error()); + + storage.UnblockWrites(); + ASSERT_TRUE(storage.WaitForWrites(2)); + EXPECT_TRUE(writer.Reserve().has_value()); + writer.Stop(); +} + +TEST(OrderedOpLogWriterAdmissionTest, + ExistingReservationCanCommitAfterAcceptingFalse) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + auto existing = writer.Reserve(); + ASSERT_TRUE(existing.has_value()); + auto failing = writer.Reserve(); + ASSERT_TRUE(failing.has_value()); + storage.FailNextWrites(100000, ErrorCode::PERSISTENT_FAIL); + ASSERT_TRUE( + writer + .Commit(std::move(*failing), MakeEntry("fail"), [](const auto&) {}) + .has_value()); + + for (int i = 0; i < 100 && writer.IsAccepting(); ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_FALSE(writer.IsAccepting()); + + auto pending = writer.Commit(std::move(*existing), MakeEntry("existing"), + [](const auto&) {}); + EXPECT_TRUE(pending.has_value()); + writer.Stop(); +} + +TEST(OrderedOpLogWriterLoopTest, + WritesSingleEntryWithoutWaitingForMaxBatchSize) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 1024}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + ASSERT_TRUE(writer + .Commit(std::move(*reservation), MakeEntry("k1"), + [](const auto&) {}) + .has_value()); + + ASSERT_TRUE(storage.WaitForWrites(1)); + auto batches = storage.Batches(); + ASSERT_EQ(1u, batches.size()); + ASSERT_EQ(1u, batches[0].entries.size()); + EXPECT_EQ(1u, batches[0].first_seq); + EXPECT_EQ(1u, batches[0].last_seq); + writer.Stop(); +} + +TEST(OrderedOpLogWriterLoopTest, ContinuesFromInitialDurablePrefix) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{ + .max_entries_per_batch = 4, + .initial_durable_prefix = {.batch_id = 7, .last_seq = 42}}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + auto pending = writer.Commit(std::move(*reservation), MakeEntry("k1"), + [](const auto&) {}); + ASSERT_TRUE(pending.has_value()); + EXPECT_EQ(43u, pending->sequence_id()); + + ASSERT_TRUE(storage.WaitForWrites(1)); + auto batches = storage.Batches(); + ASSERT_EQ(1u, batches.size()); + EXPECT_EQ(8u, batches[0].batch_id); + EXPECT_EQ(43u, batches[0].first_seq); + EXPECT_EQ(43u, batches[0].last_seq); + writer.Stop(); +} + +TEST(OrderedOpLogWriterLoopTest, CommitWhileReadyBatchExistsFormsNextBatch) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 4}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + auto first = writer.Reserve(); + auto second = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(second.has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*first), MakeEntry("k1"), [](const auto&) {}) + .has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*second), MakeEntry("k2"), [](const auto&) {}) + .has_value()); + writer.Start(); + + ASSERT_TRUE(storage.WaitForWrites(2)); + auto batches = storage.Batches(); + ASSERT_EQ(2u, batches.size()); + ASSERT_EQ(1u, batches[0].entries.size()); + EXPECT_EQ(1u, batches[0].first_seq); + EXPECT_EQ(1u, batches[0].last_seq); + ASSERT_EQ(1u, batches[1].entries.size()); + EXPECT_EQ(2u, batches[1].first_seq); + EXPECT_EQ(2u, batches[1].last_seq); + writer.Stop(); +} + +TEST(OrderedOpLogWriterLoopTest, CommitsDuringInflightWriteFormNextBatch) { + FakeBatchWriter storage; + storage.BlockWrites(); + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 4}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + auto first = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*first), MakeEntry("k1"), [](const auto&) {}) + .has_value()); + ASSERT_TRUE(storage.WaitForBlockedWrite()); + + auto second = writer.Reserve(); + ASSERT_TRUE(second.has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*second), MakeEntry("k2"), [](const auto&) {}) + .has_value()); + + storage.UnblockWrites(); + ASSERT_TRUE(storage.WaitForWrites(2)); + auto batches = storage.Batches(); + ASSERT_EQ(2u, batches.size()); + EXPECT_EQ(1u, batches[0].entries.size()); + EXPECT_EQ(1u, batches[1].entries.size()); + EXPECT_EQ(1u, batches[0].first_seq); + EXPECT_EQ(2u, batches[1].first_seq); + writer.Stop(); +} + +TEST(OrderedOpLogWriterLoopTest, + DoesNotWaitForUncommittedReservationsBeforeDraining) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 4}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + auto uncommitted = writer.Reserve(); + auto committed = writer.Reserve(); + ASSERT_TRUE(uncommitted.has_value()); + ASSERT_TRUE(committed.has_value()); + ASSERT_TRUE( + writer + .Commit(std::move(*committed), MakeEntry("k1"), [](const auto&) {}) + .has_value()); + + ASSERT_TRUE(storage.WaitForWrites(1)); + auto batches = storage.Batches(); + ASSERT_EQ(1u, batches.size()); + EXPECT_EQ(1u, batches[0].entries.size()); + writer.Abort(std::move(*uncommitted)); + writer.Stop(); +} + +TEST(OrderedOpLogWriterFailureTest, FirstStorageFailureStopsNewReservations) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + storage.FailNextWrites(100000, ErrorCode::PERSISTENT_FAIL); + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + ASSERT_TRUE(writer + .Commit(std::move(*reservation), MakeEntry("k1"), + [](const auto&) {}) + .has_value()); + + for (int i = 0; i < 100 && writer.IsAccepting(); ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + + EXPECT_FALSE(writer.IsAccepting()); + EXPECT_EQ(ErrorCode::PERSISTENT_FAIL, writer.LastError()); + EXPECT_FALSE(writer.Reserve().has_value()); + writer.Stop(); +} + +TEST(OrderedOpLogWriterFailureTest, + DoesNotInvokeCallbacksWhileBatchIsUndurable) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + storage.FailNextWrites(100000, ErrorCode::PERSISTENT_FAIL); + std::atomic callbacks{0}; + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + ASSERT_TRUE(writer + .Commit(std::move(*reservation), MakeEntry("k1"), + [&](const auto&) { ++callbacks; }) + .has_value()); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + + EXPECT_EQ(0, callbacks.load()); + writer.Stop(); +} + +TEST(OrderedOpLogWriterFailureTest, RetriesSameBatchUntilSuccess) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + storage.FailNextWrites(2, ErrorCode::PERSISTENT_FAIL); + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + ASSERT_TRUE(writer + .Commit(std::move(*reservation), MakeEntry("k1"), + [](const auto&) {}) + .has_value()); + + ASSERT_TRUE(storage.WaitForWrites(1)); + auto batches = storage.Batches(); + ASSERT_EQ(1u, batches.size()); + EXPECT_EQ(1u, batches[0].batch_id); + EXPECT_EQ(1u, batches[0].first_seq); + EXPECT_EQ(1u, batches[0].last_seq); + writer.Stop(); +} + +TEST(OrderedOpLogWriterFailureTest, RetryBackoffGrowsAndCaps) { + using namespace std::chrono_literals; + + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + storage.FailNextWrites(12, ErrorCode::PERSISTENT_FAIL); + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + ASSERT_TRUE(writer + .Commit(std::move(*reservation), MakeEntry("k1"), + [](const auto&) {}) + .has_value()); + + ASSERT_TRUE(storage.WaitForAttempts(13)); + const auto attempts = storage.AttemptTimes(); + ASSERT_GE(attempts.size(), 13u); + EXPECT_GE(attempts[9] - attempts[0], 400ms); + EXPECT_GE(attempts[11] - attempts[10], 800ms); + EXPECT_LT(attempts[12] - attempts[11], 1500ms); + writer.Stop(); +} + +TEST(OrderedOpLogWriterFailureTest, RetryBackoffResetsAfterSuccess) { + using namespace std::chrono_literals; + + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + auto first = writer.Reserve(); + auto second = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(second.has_value()); + + storage.FailNextWrites(10, ErrorCode::PERSISTENT_FAIL); + ASSERT_TRUE( + writer.Commit(std::move(*first), MakeEntry("k1"), [](const auto&) {}) + .has_value()); + ASSERT_TRUE(storage.WaitForWrites(1, 5s)); + const auto first_batch_attempts = storage.AttemptTimes(); + ASSERT_GE(first_batch_attempts.size(), 11u); + EXPECT_GE(first_batch_attempts[9] - first_batch_attempts[0], 400ms); + + storage.FailNextWrites(1, ErrorCode::PERSISTENT_FAIL); + ASSERT_TRUE( + writer.Commit(std::move(*second), MakeEntry("k2"), [](const auto&) {}) + .has_value()); + ASSERT_TRUE(storage.WaitForAttempts(13)); + + const auto attempts = storage.AttemptTimes(); + ASSERT_GE(attempts.size(), 13u); + EXPECT_LT(attempts[12] - attempts[11], 250ms); + writer.Stop(); +} + +TEST(OrderedOpLogWriterFailureTest, StopInterruptsRetryBackoff) { + using namespace std::chrono_literals; + + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 1}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + storage.FailNextWrites(100000, ErrorCode::PERSISTENT_FAIL); + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + ASSERT_TRUE(writer + .Commit(std::move(*reservation), MakeEntry("k1"), + [](const auto&) {}) + .has_value()); + ASSERT_TRUE(storage.WaitForAttempts(11)); + const auto attempts = storage.AttemptTimes(); + ASSERT_GE(attempts.size(), 11u); + EXPECT_GE(attempts[10] - attempts[0], 800ms); + + const auto started_at = FakeBatchWriter::Clock::now(); + writer.Stop(); + EXPECT_LT(FakeBatchWriter::Clock::now() - started_at, 250ms); +} + +TEST(OrderedOpLogWriterFailureTest, SuccessAfterRetryRestoresAccepting) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + storage.FailNextWrites(1, ErrorCode::PERSISTENT_FAIL); + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + ASSERT_TRUE(writer + .Commit(std::move(*reservation), MakeEntry("k1"), + [](const auto&) {}) + .has_value()); + + ASSERT_TRUE(storage.WaitForWrites(1)); + EXPECT_TRUE(writer.IsAccepting()); + EXPECT_TRUE(writer.Reserve().has_value()); + writer.Stop(); +} + +TEST(OrderedOpLogWriterFailureTest, LaterBatchDoesNotOvertakeStuckBatch) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 2}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + auto first = writer.Reserve(); + auto second = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(second.has_value()); + + storage.FailNextWrites(100000, ErrorCode::PERSISTENT_FAIL); + ASSERT_TRUE( + writer.Commit(std::move(*first), MakeEntry("k1"), [](const auto&) {}) + .has_value()); + for (int i = 0; i < 100 && writer.IsAccepting(); ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + ASSERT_FALSE(writer.IsAccepting()); + + ASSERT_TRUE( + writer.Commit(std::move(*second), MakeEntry("k2"), [](const auto&) {}) + .has_value()); + storage.AllowWrites(); + + ASSERT_TRUE(storage.WaitForWrites(2)); + auto batches = storage.Batches(); + ASSERT_EQ(2u, batches.size()); + EXPECT_EQ(1u, batches[0].first_seq); + EXPECT_EQ(1u, batches[0].last_seq); + EXPECT_EQ(2u, batches[1].first_seq); + EXPECT_EQ(2u, batches[1].last_seq); + writer.Stop(); +} + +TEST(OrderedOpLogWriterCallbackTest, DispatchesCallbacksInBatchSequenceOrder) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 4}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + std::mutex mutex; + std::condition_variable cv; + std::vector callback_sequences; + auto callback = [&](const OpLogEntry& entry) { + { + std::lock_guard lock(mutex); + callback_sequences.push_back(entry.sequence_id); + } + cv.notify_all(); + }; + + auto first = writer.Reserve(); + auto second = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(second.has_value()); + ASSERT_TRUE(writer.Commit(std::move(*first), MakeEntry("k1"), callback) + .has_value()); + ASSERT_TRUE(writer.Commit(std::move(*second), MakeEntry("k2"), callback) + .has_value()); + writer.Start(); + + { + std::unique_lock lock(mutex); + ASSERT_TRUE(cv.wait_for(lock, std::chrono::seconds(1), [&] { + return callback_sequences.size() == 2; + })); + EXPECT_EQ((std::vector{1, 2}), callback_sequences); + } + writer.Stop(); +} + +TEST(OrderedOpLogWriterCallbackTest, + DispatchesCallbacksAcrossBatchesInGlobalOrder) { + FakeBatchWriter storage; + storage.BlockWrites(); + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 4}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + std::mutex mutex; + std::condition_variable cv; + std::vector callback_sequences; + auto callback = [&](const OpLogEntry& entry) { + { + std::lock_guard lock(mutex); + callback_sequences.push_back(entry.sequence_id); + } + cv.notify_all(); + }; + writer.Start(); + + auto first = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(writer.Commit(std::move(*first), MakeEntry("k1"), callback) + .has_value()); + ASSERT_TRUE(storage.WaitForBlockedWrite()); + + auto second = writer.Reserve(); + ASSERT_TRUE(second.has_value()); + ASSERT_TRUE(writer.Commit(std::move(*second), MakeEntry("k2"), callback) + .has_value()); + storage.UnblockWrites(); + + { + std::unique_lock lock(mutex); + ASSERT_TRUE(cv.wait_for(lock, std::chrono::seconds(1), [&] { + return callback_sequences.size() == 2; + })); + EXPECT_EQ((std::vector{1, 2}), callback_sequences); + } + writer.Stop(); +} + +TEST(OrderedOpLogWriterCallbackTest, + StopDrainsCallbacksFromInflightSuccessfulWrite) { + FakeBatchWriter storage; + storage.BlockWrites(); + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 4}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + writer.Start(); + + std::atomic callbacks{0}; + auto reservation = writer.Reserve(); + ASSERT_TRUE(reservation.has_value()); + ASSERT_TRUE(writer + .Commit(std::move(*reservation), MakeEntry("k1"), + [&](const auto&) { ++callbacks; }) + .has_value()); + ASSERT_TRUE(storage.WaitForBlockedWrite()); + + std::thread stopper([&] { writer.Stop(); }); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + storage.UnblockWrites(); + stopper.join(); + + EXPECT_EQ(1, callbacks.load()); +} + +TEST(OrderedOpLogWriterCallbackTest, SlowCallbackDoesNotPreventNextBatchWrite) { + FakeBatchWriter storage; + OrderedOpLogWriter writer( + OrderedOpLogWriterConfig{.max_entries_per_batch = 4}, + [&](const OpLogBatchRecord& batch, + const DurablePrefix& expected_prefix) { + return storage.Write(batch, expected_prefix); + }); + + std::mutex mutex; + std::condition_variable callback_started_cv; + std::condition_variable release_callback_cv; + bool callback_started = false; + bool release_callback = false; + writer.Start(); + + auto first = writer.Reserve(); + ASSERT_TRUE(first.has_value()); + ASSERT_TRUE(writer + .Commit(std::move(*first), MakeEntry("k1"), + [&](const auto&) { + std::unique_lock lock(mutex); + callback_started = true; + callback_started_cv.notify_all(); + release_callback_cv.wait( + lock, [&] { return release_callback; }); + }) + .has_value()); + { + std::unique_lock lock(mutex); + ASSERT_TRUE(callback_started_cv.wait_for( + lock, std::chrono::seconds(1), [&] { return callback_started; })); + } + + auto second = writer.Reserve(); + ASSERT_TRUE(second.has_value()); + ASSERT_TRUE( + writer.Commit(std::move(*second), MakeEntry("k2"), [](const auto&) {}) + .has_value()); + + EXPECT_TRUE(storage.WaitForWrites(2, std::chrono::milliseconds(100))); + { + std::lock_guard lock(mutex); + release_callback = true; + } + release_callback_cv.notify_all(); + writer.Stop(); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp b/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp index 071940c826..c81d5eccf8 100644 --- a/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp +++ b/mooncake-store/tests/ha/snapshot/catalog_backed_snapshot_provider_test.cpp @@ -80,7 +80,9 @@ class CatalogBackedSnapshotProviderTest ASSERT_TRUE(snapshot->has_value()); ASSERT_EQ(snapshot->value().metadata.size(), 1u); - const auto& [key, metadata] = snapshot->value().metadata.front(); + const auto& [tenant_id, key, metadata] = + snapshot->value().metadata.front(); + EXPECT_EQ(tenant_id, "default"); EXPECT_EQ(key, kDefaultTestObjectKey); EXPECT_EQ(metadata.client_id, (UUID{1, 2})); EXPECT_EQ(metadata.size, kDefaultTestObjectSize); @@ -134,14 +136,25 @@ TEST_P(CatalogBackedSnapshotProviderTest, LoadLatestSnapshotRoundTrip) { descriptor_.last_included_seq); ASSERT_EQ(snapshot->value().metadata.size(), 1u); - const auto& [key, metadata] = snapshot->value().metadata.front(); - EXPECT_EQ(key, kDefaultTestObjectKey); - EXPECT_EQ(metadata.client_id, (UUID{1, 2})); - EXPECT_EQ(metadata.size, kDefaultTestObjectSize); - EXPECT_EQ(metadata.last_sequence_id, descriptor_.last_included_seq); - ASSERT_EQ(metadata.replicas.size(), 1u); - - const auto& replica = metadata.replicas.front(); + // The test snapshot's segment payload is built from an empty + // SegmentManager (BuildSegmentsPayload), so the loaded snapshot + // must report no StandbySegmentInfo entries. This pins the contract + // documented at the extraction site in + // catalog_backed_snapshot_provider.cpp: only memory segments + // (mounted_segments_) populate snapshot.segments; local-disk and + // NoF segments arrive via SEGMENT_MOUNT OpLog replay instead. + EXPECT_TRUE(snapshot->value().segments.empty()) + << "Empty segment manager must produce zero StandbySegmentInfo " + "entries"; + + const auto& entry = snapshot->value().metadata.front(); + EXPECT_EQ(entry.key, kDefaultTestObjectKey); + EXPECT_EQ(entry.metadata.client_id, (UUID{1, 2})); + EXPECT_EQ(entry.metadata.size, kDefaultTestObjectSize); + EXPECT_EQ(entry.metadata.last_sequence_id, descriptor_.last_included_seq); + ASSERT_EQ(entry.metadata.replicas.size(), 1u); + + const auto& replica = entry.metadata.replicas.front(); EXPECT_EQ(replica.status, ReplicaStatus::COMPLETE); ASSERT_TRUE(replica.is_disk_replica()); EXPECT_EQ(replica.get_disk_descriptor().file_path, @@ -184,6 +197,12 @@ TEST_P(CatalogBackedSnapshotProviderTest, LoadLatestSnapshotWithGroupId) { ExpectLoadsDefaultObject(); } +TEST_P(CatalogBackedSnapshotProviderTest, + LoadLatestSnapshotIgnoresObjectChecksum) { + PublishSnapshotPayload(SnapshotMetadataFormat::kWithObjectChecksum); + ExpectLoadsDefaultObject(); +} + TEST_P(CatalogBackedSnapshotProviderTest, RejectsOverflowingReplicaCount) { // A near-UINT32_MAX replica_count must not wrap the format-detection // arithmetic into a valid-looking total and slip an out-of-bounds index diff --git a/mooncake-store/tests/ha/snapshot/master_service_promotion_test_for_snapshot.cpp b/mooncake-store/tests/ha/snapshot/master_service_promotion_test_for_snapshot.cpp index 20ca542b96..74189040b3 100644 --- a/mooncake-store/tests/ha/snapshot/master_service_promotion_test_for_snapshot.cpp +++ b/mooncake-store/tests/ha/snapshot/master_service_promotion_test_for_snapshot.cpp @@ -71,7 +71,9 @@ class MasterServicePromotionSnapshotTest int64_t size, const std::string& transport_endpoint) { std::vector tasks{ - OffloadTaskItem{.tenant_id = "default", .key = key, .size = size}}; + OffloadTaskItem{.tenant_id = TenantId::Default().value(), + .key = key, + .size = size}}; StorageObjectMetadata sm; sm.bucket_id = 0; sm.offset = 0; @@ -96,7 +98,7 @@ TEST_F(MasterServicePromotionSnapshotTest, LocalDiskReplicaRoundTrip) { ASSERT_TRUE( InjectLocalDiskReplica(client_id, "k_cold", 1024, "seg_a_endpoint")); - auto before = service_->GetReplicaList("k_cold", "default"); + auto before = service_->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(before.has_value()); ASSERT_EQ(before->replicas.size(), 1u); EXPECT_TRUE(before->replicas[0].is_local_disk_replica()); @@ -112,17 +114,19 @@ TEST_F(MasterServicePromotionSnapshotTest, MixedMemoryAndLocalDiskRoundTrip) { // Put a MEMORY replica. ReplicateConfig rc; rc.replica_num = 1; - ASSERT_TRUE(service_->PutStart(client_id, "k_mixed", "default", 1024, rc) - .has_value()); ASSERT_TRUE( - service_->PutEnd(client_id, "k_mixed", "default", ReplicaType::MEMORY) + service_->PutStart(client_id, "k_mixed", TenantId::Default(), 1024, rc) .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, "k_mixed", TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); // Add LOCAL_DISK alongside. ASSERT_TRUE( InjectLocalDiskReplica(client_id, "k_mixed", 1024, "seg_a_endpoint")); - auto descs = service_->GetReplicaList("k_mixed", "default"); + auto descs = service_->GetReplicaList("k_mixed", TenantId::Default()); ASSERT_TRUE(descs.has_value()); EXPECT_EQ(descs->replicas.size(), 2u); } @@ -153,8 +157,8 @@ TEST_F(MasterServicePromotionSnapshotTest, MultipleLocalDiskHoldersRoundTrip) { ASSERT_TRUE( InjectLocalDiskReplica(client_b, "k_b", 2048, "seg_b_endpoint")); - auto a = service_->GetReplicaList("k_a", "default"); - auto b = service_->GetReplicaList("k_b", "default"); + auto a = service_->GetReplicaList("k_a", TenantId::Default()); + auto b = service_->GetReplicaList("k_b", TenantId::Default()); ASSERT_TRUE(a.has_value()); ASSERT_TRUE(b.has_value()); EXPECT_EQ(a->replicas[0].get_local_disk_descriptor().client_id, client_a); @@ -177,7 +181,7 @@ TEST_F(MasterServicePromotionSnapshotTest, InFlightPromotionTaskSnapshotSafe) { // Trigger promotion gate to enqueue a task. This pins the source // replica's refcnt and adds a per-shard PromotionTask plus a per- // segment promotion_objects entry. - auto get = service_->GetReplicaList("k_cold", "default"); + auto get = service_->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(get.has_value()); // The visible replica list should still expose exactly one LOCAL_DISK diff --git a/mooncake-store/tests/ha/snapshot/master_service_ssd_test_for_snapshot.cpp b/mooncake-store/tests/ha/snapshot/master_service_ssd_test_for_snapshot.cpp index 73c6a1734a..0a1bfbb6d8 100644 --- a/mooncake-store/tests/ha/snapshot/master_service_ssd_test_for_snapshot.cpp +++ b/mooncake-store/tests/ha/snapshot/master_service_ssd_test_for_snapshot.cpp @@ -56,8 +56,8 @@ TEST_F(MasterServiceSSDSnapshotTest, PutEndBothReplica) { ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto replicas = put_start_result.value(); ASSERT_EQ(2, replicas.size()); @@ -70,17 +70,20 @@ TEST_F(MasterServiceSSDSnapshotTest, PutEndBothReplica) { EXPECT_TRUE(has_mem); EXPECT_TRUE(has_disk); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); // PutEnd for both memory and disk - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) - .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(2, get_result.value().replicas.size()); @@ -108,13 +111,17 @@ TEST_F(MasterServiceSSDSnapshotTest, RestorePreservesCacheTotalMetrics) { ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value()); std::string key = "restore_cache_total_metric_key"; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1}) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), 1024, + {.replica_num = 1}) .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); auto& metrics = MasterMetricManager::instance(); using CacheHitStat = MasterMetricManager::CacheHitStat; @@ -143,7 +150,8 @@ TEST_F(MasterServiceSSDSnapshotTest, RestorePreservesCacheTotalMetrics) { EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], 1); EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], 1); - auto get_result = restored_service->GetReplicaList(key, "default"); + auto get_result = + restored_service->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(2, get_result.value().replicas.size()); } @@ -170,22 +178,26 @@ TEST_F(MasterServiceSSDSnapshotTest, PutRevokeDiskReplica) { ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); ASSERT_TRUE(get_result.value().replicas[0].is_memory_replica()); EXPECT_TRUE( - service_->PutRevoke(client_id, key, "default", ReplicaType::DISK) + service_ + ->PutRevoke(client_id, key, TenantId::Default(), ReplicaType::DISK) .has_value()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); ASSERT_TRUE(get_result.value().replicas[0].is_memory_replica()); @@ -213,20 +225,23 @@ TEST_F(MasterServiceSSDSnapshotTest, PutRevokeMemoryReplica) { ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) - .has_value()); - EXPECT_TRUE( - service_->PutRevoke(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) + .has_value()); + EXPECT_TRUE(service_ + ->PutRevoke(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) - .has_value()); - get_result = service_->GetReplicaList(key, "default"); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); + get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); ASSERT_TRUE(get_result.value().replicas[0].is_disk_replica()); @@ -254,21 +269,24 @@ TEST_F(MasterServiceSSDSnapshotTest, PutRevokeBothReplica) { ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) - .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) + .has_value()); EXPECT_TRUE( - service_->PutRevoke(client_id, key, "default", ReplicaType::DISK) + service_ + ->PutRevoke(client_id, key, TenantId::Default(), ReplicaType::DISK) .has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); - EXPECT_TRUE( - service_->PutRevoke(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - get_result = service_->GetReplicaList(key, "default"); + EXPECT_TRUE(service_ + ->PutRevoke(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); + get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -295,23 +313,34 @@ TEST_F(MasterServiceSSDSnapshotTest, RemoveKey) { ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); - EXPECT_TRUE(service_->Remove(key, "default").has_value()); + EXPECT_TRUE(service_->Remove(key, TenantId::Default()).has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } TEST_F(MasterServiceSSDSnapshotTest, EvictObject) { - CreateMasterServiceWithSSDFeat("/mnt/ssd"); + // Keep the lease expiry wait below the client live TTL. This test verifies + // eviction and snapshot restore, not the process-wide default lease TTL. + constexpr uint64_t kv_lease_ttl = 2000; + auto service_config = MasterServiceConfig::builder() + .set_root_fs_dir("/mnt/ssd") + .set_default_kv_lease_ttl(kv_lease_ttl) + .build(); + CreateMasterServiceWithSSDFeatAndConfig(service_config); // Mount a segment that can hold about 1024 * 16 objects. // As the eviction is processed separately for each shard, // we need to fill each shard with enough objects to thoroughly @@ -337,13 +366,13 @@ TEST_F(MasterServiceSSDSnapshotTest, EvictObject) { uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { auto put_end_mem_result = service_->PutEnd( - client_id, key, "default", ReplicaType::MEMORY); - auto put_end_disk_result = - service_->PutEnd(client_id, key, "default", ReplicaType::DISK); + client_id, key, TenantId::Default(), ReplicaType::MEMORY); + auto put_end_disk_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::DISK); ASSERT_TRUE(put_end_mem_result.has_value()); ASSERT_TRUE(put_end_disk_result.has_value()); success_puts++; @@ -358,15 +387,14 @@ TEST_F(MasterServiceSSDSnapshotTest, EvictObject) { int success_gets = 0; for (int i = 0; i < 1024 * 16 + 50; ++i) { std::string key = "test_key" + std::to_string(i); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); if (get_result.has_value()) { success_gets++; } } ASSERT_GT(success_gets, 1024 * 16); - std::this_thread::sleep_for( - std::chrono::milliseconds(DEFAULT_DEFAULT_KV_LEASE_TTL)); + std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); // service_->RemoveAll(); } @@ -408,8 +436,8 @@ TEST_F(MasterServiceSSDSnapshotTest, PutStartExpires) { : ReplicaType::MEMORY; // Put key, should success. - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); auto replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -419,7 +447,7 @@ TEST_F(MasterServiceSSDSnapshotTest, PutStartExpires) { // Complete the reserved replica. auto put_end_result = - service_->PutEnd(client_id, key, "default", reserve_type); + service_->PutEnd(client_id, key, TenantId::Default(), reserve_type); EXPECT_TRUE(put_end_result.has_value()); // Wait for a while until the put-start expired. @@ -429,15 +457,16 @@ TEST_F(MasterServiceSSDSnapshotTest, PutStartExpires) { auto result = service_->Ping(client_id); EXPECT_TRUE(result.has_value()); // Protect the key from eviction. - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = + service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); std::this_thread::sleep_for(std::chrono::seconds(1)); } // Put key again, should fail because the object has had an completed // replica. - put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS); @@ -448,18 +477,20 @@ TEST_F(MasterServiceSSDSnapshotTest, PutStartExpires) { auto result = service_->Ping(client_id); EXPECT_TRUE(result.has_value()); // Protect the key from eviction. - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = + service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); std::this_thread::sleep_for(std::chrono::seconds(1)); } - // Try PutEnd the discarded replica. + // PutEnd must reject a replica discarded after the write expired. put_end_result = - service_->PutEnd(client_id, key, "default", discard_type); - EXPECT_TRUE(put_end_result.has_value()); + service_->PutEnd(client_id, key, TenantId::Default(), discard_type); + ASSERT_FALSE(put_end_result.has_value()); + EXPECT_EQ(put_end_result.error(), ErrorCode::INVALID_WRITE); // Check that the key has only one replica. - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(get_result.value().replicas.size(), 1); if (reserve_type == ReplicaType::MEMORY) { diff --git a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot.cpp b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot.cpp index 9ca98ab635..4231a9ef99 100644 --- a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot.cpp +++ b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot.cpp @@ -36,14 +36,14 @@ std::string GenerateKeyForSegment(const UUID& client_id, std::vector replica_list; // Check if the key already exists. - auto exist_result = service->ExistKey(key, "default"); + auto exist_result = service->ExistKey(key, TenantId::Default()); if (exist_result.has_value() && exist_result.value()) { continue; // Retry if the key already exists } // Attempt to put the key. - auto put_result = service->PutStart(client_id, key, "default", {1024}, - {.replica_num = 1}); + auto put_result = service->PutStart(client_id, key, TenantId::Default(), + {1024}, {.replica_num = 1}); if (put_result.has_value()) { replica_list = std::move(put_result.value()); } @@ -57,8 +57,8 @@ std::string GenerateKeyForSegment(const UUID& client_id, throw std::runtime_error("PutStart failed with code: " + std::to_string(static_cast(code))); } - auto put_end_result = - service->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); if (!put_end_result.has_value()) { throw std::runtime_error("PutEnd failed"); } @@ -68,7 +68,7 @@ std::string GenerateKeyForSegment(const UUID& client_id, return key; } // Clean up failed attempt - auto remove_result = service->Remove(key, "default"); + auto remove_result = service->Remove(key, TenantId::Default()); if (!remove_result.has_value()) { // Ignore cleanup failure } @@ -207,13 +207,14 @@ TEST_F(MasterServiceSnapshotTest, PutStartInvalidParams) { // Test invalid replica_num config.replica_num = 0; auto put_result1 = - service_->PutStart(client_id, key, "default", 1024, config); + service_->PutStart(client_id, key, TenantId::Default(), 1024, config); EXPECT_FALSE(put_result1.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result1.error()); // Test zero slice_length config.replica_num = 1; - auto put_result2 = service_->PutStart(client_id, key, "default", 0, config); + auto put_result2 = + service_->PutStart(client_id, key, TenantId::Default(), 0, config); EXPECT_FALSE(put_result2.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result2.error()); } @@ -231,40 +232,41 @@ TEST_F(MasterServiceSnapshotTest, PutStartEndFlow) { ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_FALSE(replica_list.empty()); EXPECT_EQ(ReplicaStatus::PROCESSING, replica_list[0].status); // During put, Get/Remove should fail - auto get_replica_result = service_->GetReplicaList(key, "default"); + auto get_replica_result = + service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_replica_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_replica_result.error()); - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error()); // PutEnd should fail if the client_id does not match. - auto put_end_fail_result = service_->PutEnd(invalid_client_id, key, - "default", ReplicaType::MEMORY); + auto put_end_fail_result = service_->PutEnd( + invalid_client_id, key, TenantId::Default(), ReplicaType::MEMORY); EXPECT_FALSE(put_end_fail_result.has_value()); EXPECT_EQ(put_end_fail_result.error(), ErrorCode::ILLEGAL_CLIENT); // PutRevoke should fail if the client_id does not match. auto put_revoke_fail_result = service_->PutRevoke( - invalid_client_id, key, "default", ReplicaType::MEMORY); + invalid_client_id, key, TenantId::Default(), ReplicaType::MEMORY); EXPECT_FALSE(put_revoke_fail_result.has_value()); EXPECT_EQ(put_revoke_fail_result.error(), ErrorCode::ILLEGAL_CLIENT); // Test PutEnd - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Verify replica list after PutEnd - auto final_get_result = service_->GetReplicaList(key, "default"); + auto final_get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(final_get_result.has_value()); replica_list = final_get_result.value().replicas; EXPECT_EQ(1, replica_list.size()); @@ -293,25 +295,25 @@ TEST_F(MasterServiceSnapshotTest, RandomPutStartEndFlow) { std::uniform_int_distribution<> dis(1, 5); int random_number = dis(gen); config.replica_num = random_number; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_FALSE(replica_list.empty()); EXPECT_EQ(ReplicaStatus::PROCESSING, replica_list[0].status); // During put, Get/Remove should fail - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error()); // Test PutEnd - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Verify replica list after PutEnd - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result2.has_value()); replica_list = get_result2.value().replicas; EXPECT_EQ(random_number, replica_list.size()); @@ -328,7 +330,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegex) { service_.reset(new MasterService(service_config)); const UUID client_id = generate_uuid(); // Test getting non-existent key - auto get_result = service_->GetReplicaList(".*non_existent.*", "default"); + auto get_result = + service_->GetReplicaList(".*non_existent.*", TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); @@ -340,20 +343,21 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegex) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); } // wait for all the lease to expire std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); // Test getting existing key - auto get_result2 = service_->GetReplicaListByRegex("^test_key", "default"); + auto get_result2 = + service_->GetReplicaListByRegex("^test_key", TenantId::Default()); EXPECT_TRUE(get_result2.has_value()); auto replica_list_local = get_result2.value(); EXPECT_EQ(10, replica_list_local.size()); @@ -365,15 +369,15 @@ void put_object(MasterService& service, const UUID& client_id, uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service.PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service.PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()) << "Failed to PutStart for key: " << key; - auto put_end_result = - service.PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service.PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()) << "Failed to PutEnd for key: " << key; - auto exist_result = service.ExistKey(key, "default"); + auto exist_result = service.ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()) << "Key does not exist after put: " << key; } @@ -417,7 +421,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) { // Test 3.1: Simple prefix matching { - auto result = service_->GetReplicaListByRegex("^test_key_", "default"); + auto result = + service_->GetReplicaListByRegex("^test_key_", TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 3); // Matches test_key_01, test_key_02, test_key_10 @@ -425,8 +430,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) { // Test 3.2: Matching with a wildcard for any number { - auto result = - service_->GetReplicaListByRegex("^test_key_\\d+$", "default"); + auto result = service_->GetReplicaListByRegex("^test_key_\\d+$", + TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 3); } @@ -435,7 +440,7 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) { { // Matches "data_part_1_chunk_a" and "data_part_2_chunk_b" auto result = service_->GetReplicaListByRegex("^data_part_\\d_chunk_.$", - "default"); + TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 2); } @@ -443,7 +448,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) { // Test 3.4: Matching keys containing a specific substring { // Matches all keys with "key" in them - auto result = service_->GetReplicaListByRegex("key", "default"); + auto result = + service_->GetReplicaListByRegex("key", TenantId::Default()); ASSERT_TRUE(result.has_value()); // Expected: test_key_01, test_key_02, test_key_10, // prod_key_alpha, prod_key_beta, @@ -455,7 +461,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) { // Test 3.5: Matching based on file-like paths { // Match all .log files - auto result = service_->GetReplicaListByRegex("\\.log$", "default"); + auto result = + service_->GetReplicaListByRegex("\\.log$", TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 1); EXPECT_EQ(result.value().begin()->first, "logs/app-2025-08-13.log"); @@ -464,8 +471,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) { // Test 3.6: OR condition using | { // Match keys starting with "prod" OR ending with "json" - auto result = - service_->GetReplicaListByRegex("^prod|\\.json$", "default"); + auto result = service_->GetReplicaListByRegex("^prod|\\.json$", + TenantId::Default()); ASSERT_TRUE(result.has_value()); // Expected: prod_key_alpha, prod_key_beta, config/user/settings.json EXPECT_EQ(result.value().size(), 3); @@ -473,8 +480,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) { // Test 3.7: Regex that should not match anything { - auto result = - service_->GetReplicaListByRegex("^non_existent_prefix_", "default"); + auto result = service_->GetReplicaListByRegex("^non_existent_prefix_", + TenantId::Default()); // This should succeed but return an empty map. ASSERT_TRUE(result.has_value()); EXPECT_TRUE(result.value().empty()); @@ -482,7 +489,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) { // Test 3.8: Exact match regex { - auto result = service_->GetReplicaListByRegex("^short$", "default"); + auto result = + service_->GetReplicaListByRegex("^short$", TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 1); EXPECT_EQ(result.value().begin()->first, "short"); @@ -491,7 +499,7 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaListByRegexComplex) { // Test 3.9: Initial test for non-existent key (as a sanity check) { auto get_result = service_->GetReplicaListByRegex( - ".*absolutely_non_existent.*", "default"); + ".*absolutely_non_existent.*", TenantId::Default()); // Depending on implementation, this could return an empty map or an // error. Let's assume it returns an empty map for a valid regex with no // matches. @@ -504,7 +512,8 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaList) { service_.reset(new MasterService()); const UUID client_id = generate_uuid(); // Test getting non-existent key - auto get_result = service_->GetReplicaList("non_existent", "default"); + auto get_result = + service_->GetReplicaList("non_existent", TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); @@ -514,15 +523,15 @@ TEST_F(MasterServiceSnapshotTest, GetReplicaList) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test getting existing key - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result2.has_value()); auto replica_list_local = get_result2.value().replicas; EXPECT_FALSE(replica_list_local.empty()); @@ -537,24 +546,24 @@ TEST_F(MasterServiceSnapshotTest, RemoveObject) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test removing the object - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result.has_value()); // Verify object is removed - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); // Test removing non-existent object - auto remove_result2 = service_->Remove("non_existent", "default"); + auto remove_result2 = service_->Remove("non_existent", TenantId::Default()); EXPECT_FALSE(remove_result2.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, remove_result2.error()); } @@ -572,19 +581,19 @@ TEST_F(MasterServiceSnapshotTest, RandomRemoveObject) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test removing the object - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result.has_value()); // Verify object is removed - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -604,24 +613,24 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegex) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); } // wait for all the lease to expire std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto res = service_->RemoveByRegex("^test_key", "default"); + auto res = service_->RemoveByRegex("^test_key", TenantId::Default()); ASSERT_TRUE(res.has_value()); ASSERT_EQ(10, res.value()); times = 10; while (times--) { std::string key = "test_key" + std::to_string(times); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()); } @@ -668,7 +677,8 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) { populate_store(); // Action: Remove keys starting with "test_key_" - auto remove_result = service_->RemoveByRegex("^test_key_", "default"); + auto remove_result = + service_->RemoveByRegex("^test_key_", TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); EXPECT_EQ(remove_result.value(), 3); // Should remove 3 keys @@ -676,7 +686,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) { std::vector deleted_keys = {"test_key_01", "test_key_02", "test_key_10"}; for (const auto& key : deleted_keys) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); EXPECT_FALSE(exist_result.value()) << "Key " << key << " should have been deleted."; @@ -685,7 +695,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) { std::vector remaining_keys = { "prod_key_alpha", "short", "test-key-extra"}; // Sample a few for (const auto& key : remaining_keys) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); EXPECT_TRUE(exist_result.value()) << "Key " << key << " should NOT have been deleted."; @@ -708,12 +718,13 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) { size_t total_keys = 13; // Count from the keys_to_put vector // Action: Remove all keys - auto remove_result = service_->RemoveByRegex(".*", "default"); + auto remove_result = service_->RemoveByRegex(".*", TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); EXPECT_EQ(remove_result.value(), total_keys); // Verification: Check that no keys remain - auto get_all_result = service_->GetReplicaListByRegex(".*", "default"); + auto get_all_result = + service_->GetReplicaListByRegex(".*", TenantId::Default()); ASSERT_TRUE(get_all_result.has_value()); EXPECT_TRUE(get_all_result.value().empty()); } @@ -733,13 +744,14 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) { size_t total_keys_before_remove = 13; // Action: Attempt to remove using a pattern that matches nothing - auto remove_result = - service_->RemoveByRegex("^nonexistent-pattern-", "default"); + auto remove_result = service_->RemoveByRegex("^nonexistent-pattern-", + TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); EXPECT_EQ(remove_result.value(), 0); // Should remove 0 keys // Verification: Check that all keys still exist - auto get_all_result = service_->GetReplicaListByRegex(".*", "default"); + auto get_all_result = + service_->GetReplicaListByRegex(".*", TenantId::Default()); ASSERT_TRUE(get_all_result.has_value()); EXPECT_EQ(get_all_result.value().size(), total_keys_before_remove); } @@ -757,7 +769,8 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) { populate_store(); // Action: Remove all keys that contain a slash '/' OR end with a number - auto remove_result = service_->RemoveByRegex("/|\\d$", "default"); + auto remove_result = + service_->RemoveByRegex("/|\\d$", TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); // Matches: "config/user/settings.json", "logs/app-2025-08-13.log", // "test_key_01", "test_key_02", "test_key_10" @@ -782,7 +795,8 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) { populate_store(); // Action: Remove all keys that contain "chunk" OR "config" - auto remove_result = service_->RemoveByRegex("chunk|config", "default"); + auto remove_result = + service_->RemoveByRegex("chunk|config", TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); // Matches: "data_part_1_chunk_a", "data_part_2_chunk_b", // "config/user/settings.json" @@ -790,17 +804,17 @@ TEST_F(MasterServiceSnapshotTest, RemoveByRegexComplex) { // Verification auto exist_result_chunk = - service_->ExistKey("data_part_1_chunk_a", "default"); + service_->ExistKey("data_part_1_chunk_a", TenantId::Default()); ASSERT_TRUE(exist_result_chunk.has_value()); EXPECT_FALSE(exist_result_chunk.value()); - auto exist_result_config = - service_->ExistKey("config/user/settings.json", "default"); + auto exist_result_config = service_->ExistKey( + "config/user/settings.json", TenantId::Default()); ASSERT_TRUE(exist_result_config.has_value()); EXPECT_FALSE(exist_result_config.value()); auto exist_result_untouched = - service_->ExistKey("prod_key_alpha", "default"); + service_->ExistKey("prod_key_alpha", TenantId::Default()); ASSERT_TRUE(exist_result_untouched.has_value()); EXPECT_TRUE(exist_result_untouched.value()); } @@ -820,13 +834,13 @@ TEST_F(MasterServiceSnapshotTest, RemoveAll) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); } // wait for all the lease to expire @@ -835,7 +849,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveAll) { // before TearDown snapshot verification ASSERT_EQ(10, // service_->RemoveAll()); times = 10; while (times--) { // std::string key = "test_key" + std::to_string(times); - // auto exist_result = service_->ExistKey(key, "default"); + // auto exist_result = service_->ExistKey(key, TenantId::Default()); // ASSERT_TRUE(exist_result.has_value()); // ASSERT_FALSE(exist_result.value()); // } @@ -869,8 +883,8 @@ TEST_F(MasterServiceSnapshotTest, SingleSliceMultiReplicaFlow) { std::vector replica_list; // Test PutStart with multiple slices and replicas - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); @@ -886,17 +900,17 @@ TEST_F(MasterServiceSnapshotTest, SingleSliceMultiReplicaFlow) { } // Test GetReplicaList during processing (should fail) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); // Complete the put operation - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test GetReplicaList after completion - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result2.has_value()); auto retrieved_replicas = get_result2.value().replicas; ASSERT_EQ(num_replicas, retrieved_replicas.size()); @@ -929,15 +943,15 @@ TEST_F(MasterServiceSnapshotTest, CleanupStaleHandlesTest) { config.replica_num = 1; // One replica // Create the object - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Verify object exists - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); auto retrieved_replicas = get_result.value().replicas; ASSERT_EQ(1, retrieved_replicas.size()); @@ -948,7 +962,7 @@ TEST_F(MasterServiceSnapshotTest, CleanupStaleHandlesTest) { // Try to get the object - it should be automatically removed since the // replica is invalid - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result2.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result2.error()); @@ -958,15 +972,15 @@ TEST_F(MasterServiceSnapshotTest, CleanupStaleHandlesTest) { // Create another object std::string key2 = "another_segment_object"; - auto put_start_result2 = - service_->PutStart(client_id, key2, "default", slice_length, config); + auto put_start_result2 = service_->PutStart( + client_id, key2, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result2.has_value()); - auto put_end_result2 = - service_->PutEnd(client_id, key2, "default", ReplicaType::MEMORY); + auto put_end_result2 = service_->PutEnd( + client_id, key2, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result2.has_value()); // Verify we can get it - auto get_result3 = service_->GetReplicaList(key2, "default"); + auto get_result3 = service_->GetReplicaList(key2, TenantId::Default()); ASSERT_TRUE(get_result3.has_value()); // Unmount the segment @@ -974,7 +988,7 @@ TEST_F(MasterServiceSnapshotTest, CleanupStaleHandlesTest) { ASSERT_TRUE(unmount_result2.has_value()); // Try to remove the object that should already be cleaned up - auto remove_result = service_->Remove(key2, "default"); + auto remove_result = service_->Remove(key2, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, remove_result.error()); } @@ -1007,10 +1021,11 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentWriteAndRemoveAll) { std::vector replica_list; auto put_start_result = service_->PutStart( - client_id, key, "default", slice_length, config); + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { - auto put_end_result = service_->PutEnd( - client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = + service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); if (put_end_result.has_value()) { success_writes++; } @@ -1076,11 +1091,11 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentReadAndRemoveAll) { ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } @@ -1093,7 +1108,8 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentReadAndRemoveAll) { readers.emplace_back([&]() { for (int j = 0; j < num_objects; ++j) { std::string key = "pre_key_" + std::to_string(j); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = + service_->GetReplicaList(key, TenantId::Default()); if (get_result.has_value()) { success_reads++; } @@ -1138,7 +1154,7 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentReadAndRemoveAll) { // // Verify all objects were removed // for (int i = 0; i < num_objects; ++i) { // std::string key = "pre_key_" + std::to_string(i); - // auto get_result = service_->GetReplicaList(key, "default"); + // auto get_result = service_->GetReplicaList(key, TenantId::Default()); // EXPECT_FALSE(get_result.has_value()); // EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); // } @@ -1161,11 +1177,11 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentRemoveAllOperations) { ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } @@ -1193,7 +1209,7 @@ TEST_F(MasterServiceSnapshotTest, ConcurrentRemoveAllOperations) { // // Verify all objects were removed // for (int i = 0; i < num_objects; ++i) { // std::string key = "pre_key_" + std::to_string(i); - // auto get_result = service_->GetReplicaList(key, "default"); + // auto get_result = service_->GetReplicaList(key, TenantId::Default()); // EXPECT_FALSE(get_result.has_value()); // EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); // } @@ -1230,23 +1246,23 @@ TEST_F(MasterServiceSnapshotTest, UnmountSegmentImmediateCleanup) { // Umount will remove all objects in the segment, include the key1 ASSERT_EQ(1, service_->GetKeyCount()); // Verify objects in segment1 is gone - auto get_result1 = service_->GetReplicaList(key1, "default"); + auto get_result1 = service_->GetReplicaList(key1, TenantId::Default()); ASSERT_FALSE(get_result1.has_value()); ASSERT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result1.error()); // Verify objects in segment2 is still there - auto get_result2 = service_->GetReplicaList(key2, "default"); + auto get_result2 = service_->GetReplicaList(key2, TenantId::Default()); ASSERT_TRUE(get_result2.has_value()); // Verify put key1 will put into segment2 rather than segment1 - auto put_start_result = - service_->PutStart(client_id, key1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key1, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); - auto put_end_result = - service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key1, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto get_result3 = service_->GetReplicaList(key1, "default"); + auto get_result3 = service_->GetReplicaList(key1, TenantId::Default()); ASSERT_TRUE(get_result3.has_value()); auto retrieved = get_result3.value(); ASSERT_EQ(replica_list[0] @@ -1278,15 +1294,17 @@ TEST_F(MasterServiceSnapshotTest, ReadableAfterPartialUnmountWithReplication) { ReplicateConfig config; config.replica_num = 2; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); ASSERT_EQ(2u, put_start_result->size()); - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); // Verify two replicas exist and they are on distinct segments - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); auto replicas = get_result.value().replicas; ASSERT_EQ(2u, replicas.size()); @@ -1304,7 +1322,7 @@ TEST_F(MasterServiceSnapshotTest, ReadableAfterPartialUnmountWithReplication) { ASSERT_TRUE(service_->UnmountSegment(segment1.id, client_id).has_value()); // Key should still be readable via the remaining replica - auto get_after_unmount = service_->GetReplicaList(key, "default"); + auto get_after_unmount = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_after_unmount.has_value()) << "Object should remain accessible with surviving replica"; } @@ -1354,7 +1372,7 @@ TEST_F(MasterServiceSnapshotTest, UnmountSegmentPerformance) { // Verify all keys are gone for (const auto& key : keys) { - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -1384,77 +1402,77 @@ TEST_F(MasterServiceSnapshotTest, RemoveLeasedObject) { config.replica_num = 1; // Verify lease is granted on ExistsKey - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result.error()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto remove_result2 = service_->Remove(key, "default"); + auto remove_result2 = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result2.has_value()); // Verify lease is extended on successive ExistsKey - auto put_start_result2 = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result2 = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result2.has_value()); - auto put_end_result2 = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result2 = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result2.has_value()); - auto exist_result2 = service_->ExistKey(key, "default"); + auto exist_result2 = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result2.has_value()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto exist_result3 = service_->ExistKey(key, "default"); + auto exist_result3 = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result3.has_value()); - auto remove_result3 = service_->Remove(key, "default"); + auto remove_result3 = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result3.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result3.error()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto remove_result4 = service_->Remove(key, "default"); + auto remove_result4 = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result4.has_value()); // Verify lease is granted on GetReplicaList - auto put_start_result3 = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result3 = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result3.has_value()); - auto put_end_result3 = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result3 = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result3.has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); - auto remove_result5 = service_->Remove(key, "default"); + auto remove_result5 = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result5.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result5.error()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto remove_result6 = service_->Remove(key, "default"); + auto remove_result6 = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result6.has_value()); // Verify lease is extended on successive GetReplicaList - auto put_start_result4 = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result4 = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result4.has_value()); - auto put_end_result4 = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result4 = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result4.has_value()); - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result2.has_value()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto get_result3 = service_->GetReplicaList(key, "default"); + auto get_result3 = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result3.has_value()); - auto remove_result7 = service_->Remove(key, "default"); + auto remove_result7 = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result7.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result7.error()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto remove_result8 = service_->Remove(key, "default"); + auto remove_result8 = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result8.has_value()); // Verify object is removed - auto get_result4 = service_->GetReplicaList(key, "default"); + auto get_result4 = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result4.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result4.error()); } @@ -1472,14 +1490,14 @@ TEST_F(MasterServiceSnapshotTest, RemoveAllLeasedObject) { uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); if (i >= 5) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); } } @@ -1487,7 +1505,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveAllLeasedObject) { // before TearDown snapshot verification ASSERT_EQ(5, // service_->RemoveAll()); for (int i = 0; i < 5; ++i) { // std::string key = "test_key" + std::to_string(i); - // auto exist_result = service_->ExistKey(key, "default"); + // auto exist_result = service_->ExistKey(key, TenantId::Default()); // ASSERT_FALSE(exist_result.value()); // } // // wait for all the lease to expire @@ -1495,7 +1513,7 @@ TEST_F(MasterServiceSnapshotTest, RemoveAllLeasedObject) { // ASSERT_EQ(5, service_->RemoveAll()); // for (int i = 5; i < 10; ++i) { // std::string key = "test_key" + std::to_string(i); - // auto exist_result = service_->ExistKey(key, "default"); + // auto exist_result = service_->ExistKey(key, TenantId::Default()); // ASSERT_FALSE(exist_result.value()); // } } @@ -1525,11 +1543,11 @@ TEST_F(MasterServiceSnapshotTest, EvictObject) { uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { - auto put_end_result = service_->PutEnd(client_id, key, "default", - ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); success_puts++; } else { @@ -1570,14 +1588,15 @@ TEST_F(MasterServiceSnapshotTest, TryEvictLeasedObject) { uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { - auto put_end_result = service_->PutEnd(client_id, key, "default", - ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // the object is leased - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = + service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); leased_keys.push_back(key); success_puts++; @@ -1591,7 +1610,7 @@ TEST_F(MasterServiceSnapshotTest, TryEvictLeasedObject) { std::this_thread::sleep_for(std::chrono::milliseconds(50)); // All leased objects should be accessible for (const auto& key : leased_keys) { - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); } // [Commented for snapshot test] The following RemoveAll would clear data @@ -1626,21 +1645,24 @@ TEST_F(MasterServiceSnapshotTest, RemoveSoftPinObject) { config.with_soft_pin = true; // Verify soft pin does not block remove + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) + .has_value()); ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) .has_value()); - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->Remove(key, "default").has_value()); + EXPECT_TRUE(service_->Remove(key, TenantId::Default()).has_value()); // [Commented for snapshot test] The following RemoveAll would clear data // before TearDown snapshot verification // // Verify soft pin does not block RemoveAll // ASSERT_TRUE( - // service_->PutStart(client_id, key, "default", slice_length, // - // config).has_value()); + // service_->PutStart(client_id, key, TenantId::Default(), slice_length, + // // config).has_value()); // ASSERT_TRUE( - // service_->PutEnd(client_id, key, "default", + // service_->PutEnd(client_id, key, TenantId::Default(), // ReplicaType::MEMORY).has_value()); // EXPECT_EQ(1, service_->RemoveAll()); } @@ -1679,13 +1701,13 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotEvictedBeforeOtherObjects) { soft_pin_config.with_soft_pin = true; ASSERT_TRUE(service_ - ->PutStart(client_id, pin_key, "default", + ->PutStart(client_id, pin_key, TenantId::Default(), slice_length, soft_pin_config) .has_value()); - ASSERT_TRUE( - service_ - ->PutEnd(client_id, pin_key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, pin_key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); } // Fill the segment to trigger eviction @@ -1696,12 +1718,13 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotEvictedBeforeOtherObjects) { ReplicateConfig config; config.replica_num = 1; if (service_ - ->PutStart(client_id, key, "default", slice_length, config) + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) .has_value()) { - ASSERT_TRUE( - service_ - ->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); } else { failed_puts++; } @@ -1713,8 +1736,8 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotEvictedBeforeOtherObjects) { // pin_key should still be accessible for (int i = 0; i < 2; i++) { std::string pin_key = "pin_key" + std::to_string(i); - ASSERT_TRUE( - service_->GetReplicaList(pin_key, "default").has_value()); + ASSERT_TRUE(service_->GetReplicaList(pin_key, TenantId::Default()) + .has_value()); } // wait for the lease to expire @@ -1757,11 +1780,14 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsCanBeEvicted) { ReplicateConfig config; config.replica_num = 1; config.with_soft_pin = true; - if (service_->PutStart(client_id, key, "default", slice_length, config) + if (service_ + ->PutStart(client_id, key, TenantId::Default(), slice_length, + config) .has_value()) { - ASSERT_TRUE( - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); success_puts++; } else { // wait for eviction to work @@ -1818,12 +1844,13 @@ TEST_F(MasterServiceSnapshotTest, SoftPinExtendedOnGet) { soft_pin_config.replica_num = 1; soft_pin_config.with_soft_pin = true; - ASSERT_TRUE(service_->PutStart(client_id, pin_key, "default", - slice_length, soft_pin_config)); - ASSERT_TRUE( - service_ - ->PutEnd(client_id, pin_key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_->PutStart(client_id, pin_key, + TenantId::Default(), slice_length, + soft_pin_config)); + ASSERT_TRUE(service_ + ->PutEnd(client_id, pin_key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); } // Wait for the soft pin to expire @@ -1832,8 +1859,8 @@ TEST_F(MasterServiceSnapshotTest, SoftPinExtendedOnGet) { // Get the pin_key to extend the soft pin for (int i = 0; i < 2; i++) { std::string pin_key = "pin_key" + std::to_string(i); - ASSERT_TRUE( - service_->GetReplicaList(pin_key, "default").has_value()); + ASSERT_TRUE(service_->GetReplicaList(pin_key, TenantId::Default()) + .has_value()); } // Fill the segment to trigger eviction @@ -1844,12 +1871,13 @@ TEST_F(MasterServiceSnapshotTest, SoftPinExtendedOnGet) { ReplicateConfig config; config.replica_num = 1; if (service_ - ->PutStart(client_id, key, "default", slice_length, config) + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) .has_value()) { - ASSERT_TRUE( - service_ - ->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); } else { failed_puts++; } @@ -1862,8 +1890,8 @@ TEST_F(MasterServiceSnapshotTest, SoftPinExtendedOnGet) { // pin_key should still be accessible for (int i = 0; i < 2; i++) { std::string pin_key = "pin_key" + std::to_string(i); - ASSERT_TRUE( - service_->GetReplicaList(pin_key, "default").has_value()); + ASSERT_TRUE(service_->GetReplicaList(pin_key, TenantId::Default()) + .has_value()); } // [Commented for snapshot test] The following RemoveAll would clear // data before TearDown snapshot verification Only remove all objects @@ -1907,11 +1935,14 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotAllowEvict) { ReplicateConfig config; config.replica_num = 1; config.with_soft_pin = true; - if (service_->PutStart(client_id, key, "default", slice_length, config) + if (service_ + ->PutStart(client_id, key, TenantId::Default(), slice_length, + config) .has_value()) { - ASSERT_TRUE( - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); success_keys.push_back(key); } else { // wait for eviction to work @@ -1921,7 +1952,8 @@ TEST_F(MasterServiceSnapshotTest, SoftPinObjectsNotAllowEvict) { ASSERT_LE(success_keys.size(), 17); // All soft pinned objects should be accessible for (const auto& key : success_keys) { - ASSERT_TRUE(service_->GetReplicaList(key, "default").has_value()); + ASSERT_TRUE( + service_->GetReplicaList(key, TenantId::Default()).has_value()); } // [Commented for snapshot test] The following RemoveAll would clear data // before TearDown snapshot verification @@ -1948,8 +1980,8 @@ TEST_F(MasterServiceSnapshotTest, ReplicaSegmentsAreUnique) { ReplicateConfig config; config.replica_num = 10; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto replica_list_local = put_start_result.value(); ASSERT_EQ(config.replica_num, replica_list_local.size()); @@ -1965,8 +1997,10 @@ TEST_F(MasterServiceSnapshotTest, ReplicaSegmentsAreUnique) { EXPECT_EQ(segment_names.size(), config.replica_num) << "Duplicate segment found"; - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); } TEST_F(MasterServiceSnapshotTest, ReplicationFactorTwoWithSingleSegment) { @@ -1986,8 +2020,8 @@ TEST_F(MasterServiceSnapshotTest, ReplicationFactorTwoWithSingleSegment) { ReplicateConfig config; config.replica_num = 2; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto replicas = put_start_result.value(); @@ -2020,22 +2054,23 @@ TEST_F(MasterServiceSnapshotTest, BatchExistKeyTest) { config.replica_num = 1; uint64_t slice_length = value_size; auto put_start_result = service_->PutStart( - client_id, test_keys[i], "default", slice_length, config); + client_id, test_keys[i], TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = service_->PutEnd(client_id, test_keys[i], - "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, test_keys[i], TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } // Test individual ExistKey calls to verify the underlying functionality for (int i = 0; i < test_object_num; ++i) { - auto exist_result = service_->ExistKey(test_keys[i], "default"); + auto exist_result = + service_->ExistKey(test_keys[i], TenantId::Default()); EXPECT_TRUE(exist_result.value()); } // Tets batch test_keys.push_back("non_existent_key"); - auto exist_resp = service_->BatchExistKey(test_keys, "default"); + auto exist_resp = service_->BatchExistKey(test_keys, TenantId::Default()); for (int i = 0; i < test_object_num; ++i) { ASSERT_TRUE(exist_resp[i].value()); } @@ -2329,8 +2364,8 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) { config.replica_num = kReplicaCnt; // Put key_1, should success. - auto put_start_result = - service_->PutStart(client_id, key_1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key_1, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -2339,8 +2374,8 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) { } // Put key_1 again, should fail because the key exists. - put_start_result = - service_->PutStart(client_id, key_1, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_1, TenantId::Default(), + slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS); @@ -2355,8 +2390,8 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) { // Put key_1 again, should success because the old one has expired and will // be discarded by this put. - put_start_result = - service_->PutStart(client_id, key_1, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_1, TenantId::Default(), + slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -2365,18 +2400,18 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) { } // Complete key_1. - auto put_end_result = - service_->PutEnd(client_id, key_1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key_1, TenantId::Default(), ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Protect key_1 from eviction. - auto get_result = service_->GetReplicaList(key_1, "default"); + auto get_result = service_->GetReplicaList(key_1, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); // Put key_2, should fail because the key_1 occupied 12MB (6MB processing, // 6MB discarded but not yet released) on each segment. - put_start_result = - service_->PutStart(client_id, key_2, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_2, TenantId::Default(), + slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE); @@ -2389,15 +2424,15 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) { EXPECT_TRUE(result.has_value()); } // Protect key_1 from eviction. - auto get_result = service_->GetReplicaList(key_1, "default"); + auto get_result = service_->GetReplicaList(key_1, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); std::this_thread::sleep_for(std::chrono::seconds(1)); } // Put key_2 again, should success because the discarded replica has been // released. - put_start_result = - service_->PutStart(client_id, key_2, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_2, TenantId::Default(), + slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -2412,15 +2447,15 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) { EXPECT_TRUE(result.has_value()); } // Protect key_1 from eviction. - auto get_result = service_->GetReplicaList(key_1, "default"); + auto get_result = service_->GetReplicaList(key_1, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); std::this_thread::sleep_for(std::chrono::seconds(1)); } // Put key_2 again, should fail because eviction has not been triggered. And // this PutStart should trigger the eviction. - put_start_result = - service_->PutStart(client_id, key_2, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_2, TenantId::Default(), + slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE); @@ -2429,8 +2464,8 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) { // Put key_2 again, should success because the previous one has been // discarded and released. - put_start_result = - service_->PutStart(client_id, key_2, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_2, TenantId::Default(), + slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -2439,8 +2474,8 @@ TEST_F(MasterServiceSnapshotTest, PutStartExpiringTest) { } // Complete key_2. - put_end_result = - service_->PutEnd(client_id, key_2, "default", ReplicaType::MEMORY); + put_end_result = service_->PutEnd(client_id, key_2, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); } @@ -2556,17 +2591,17 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearAllSegments) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } // Verify objects exist for (const auto& key : keys) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()); } @@ -2583,7 +2618,7 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearAllSegments) { // Verify objects are removed for (const auto& key : keys) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()) << "Key " << key << " should be removed"; @@ -2613,11 +2648,11 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearSpecificSegment) { config.replica_num = 1; config.preferred_segment = segment_name; // Ensure object is placed on segment1 - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // 4. Wait for lease to expire and verify it's actually expired @@ -2655,7 +2690,7 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearSpecificSegment) { const auto& cleared_keys = clear_result.value(); ASSERT_EQ(1u, cleared_keys.size()) << "Key should be cleared"; - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()) << "Key should be removed after being cleared."; @@ -2675,15 +2710,15 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithLeaseActive) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Grant a lease by calling GetReplicaList (similar to normal usage) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); // Try to clear immediately (lease should still be active) @@ -2697,7 +2732,7 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithLeaseActive) { << "No keys should be cleared when lease is active"; // Verify object still exists - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()) << "Key should still exist"; } @@ -2717,11 +2752,11 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithDifferentClientId) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id1, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id1, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id1, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id1, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Wait for lease to expire @@ -2738,7 +2773,7 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithDifferentClientId) { << "No keys should be cleared for different client_id"; // Verify object still exists - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()) << "Key should still exist"; } @@ -2793,11 +2828,11 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearWithEmptyStringKeys) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = service_->PutStart(client_id, valid_key, "default", - value_length, config); + auto put_start_result = service_->PutStart( + client_id, valid_key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, valid_key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, valid_key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Wait for lease to expire @@ -2834,26 +2869,26 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearMixedScenario) { config.replica_num = 1; // Create key1 and key2 with client_id1 - auto put_start1 = - service_->PutStart(client_id1, key1, "default", value_length, config); + auto put_start1 = service_->PutStart(client_id1, key1, TenantId::Default(), + value_length, config); ASSERT_TRUE(put_start1.has_value()); - auto put_end1 = - service_->PutEnd(client_id1, key1, "default", ReplicaType::MEMORY); + auto put_end1 = service_->PutEnd(client_id1, key1, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end1.has_value()); - auto put_start2 = - service_->PutStart(client_id1, key2, "default", value_length, config); + auto put_start2 = service_->PutStart(client_id1, key2, TenantId::Default(), + value_length, config); ASSERT_TRUE(put_start2.has_value()); - auto put_end2 = - service_->PutEnd(client_id1, key2, "default", ReplicaType::MEMORY); + auto put_end2 = service_->PutEnd(client_id1, key2, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end2.has_value()); // Create key3 with client_id2 - auto put_start3 = - service_->PutStart(client_id2, key3, "default", value_length, config); + auto put_start3 = service_->PutStart(client_id2, key3, TenantId::Default(), + value_length, config); ASSERT_TRUE(put_start3.has_value()); - auto put_end3 = - service_->PutEnd(client_id2, key3, "default", ReplicaType::MEMORY); + auto put_end3 = service_->PutEnd(client_id2, key3, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end3.has_value()); // Wait for lease to expire @@ -2871,16 +2906,16 @@ TEST_F(MasterServiceSnapshotTest, BatchReplicaClearMixedScenario) { << "Only keys belonging to client_id1 should be cleared"; // Verify key1 and key2 are cleared - auto exist1 = service_->ExistKey(key1, "default"); + auto exist1 = service_->ExistKey(key1, TenantId::Default()); ASSERT_TRUE(exist1.has_value()); ASSERT_FALSE(exist1.value()) << "key1 should be cleared"; - auto exist2 = service_->ExistKey(key2, "default"); + auto exist2 = service_->ExistKey(key2, TenantId::Default()); ASSERT_TRUE(exist2.has_value()); ASSERT_FALSE(exist2.value()) << "key2 should be cleared"; // Verify key3 still exists (different client_id) - auto exist3 = service_->ExistKey(key3, "default"); + auto exist3 = service_->ExistKey(key3, TenantId::Default()); ASSERT_TRUE(exist3.has_value()); ASSERT_TRUE(exist3.value()) << "key3 should still exist (different client_id)"; @@ -2915,16 +2950,16 @@ TEST_F(MasterServiceSnapshotTest, CreateCopyTaskTest) { ReplicateConfig config; config.replica_num = 1; config.preferred_segment = "segment_0"; - auto put_start_result = - service_->PutStart(client_id, key1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key1, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key1, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Copy key1 to "segment_1" and "segment_2" - auto copy_result = - service_->CreateCopyTask(key1, "default", {"segment_1", "segment_2"}); + auto copy_result = service_->CreateCopyTask(key1, TenantId::Default(), + {"segment_1", "segment_2"}); EXPECT_TRUE(copy_result.has_value()); // verify the copy task is created and assigned to the client who executed @@ -2935,19 +2970,19 @@ TEST_F(MasterServiceSnapshotTest, CreateCopyTaskTest) { EXPECT_EQ(contexts[0].client_id, task.value().assigned_client); // Copy with empty targets should fail - auto copy_result1 = service_->CreateCopyTask(key1, "default", {}); + auto copy_result1 = service_->CreateCopyTask(key1, TenantId::Default(), {}); EXPECT_FALSE(copy_result1.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_result1.error()); // Copy not exist key should fail - auto copy_result2 = - service_->CreateCopyTask("not_exist_key", "default", {"segment_1"}); + auto copy_result2 = service_->CreateCopyTask( + "not_exist_key", TenantId::Default(), {"segment_1"}); EXPECT_FALSE(copy_result2.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_result2.error()); // Copy to segment that not mounted should fail - auto copy_result3 = - service_->CreateCopyTask(key1, "default", {"not_mounted_segment"}); + auto copy_result3 = service_->CreateCopyTask(key1, TenantId::Default(), + {"not_mounted_segment"}); EXPECT_FALSE(copy_result3.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_result3.error()); } @@ -2981,16 +3016,16 @@ TEST_F(MasterServiceSnapshotTest, CreateMoveTaskTest) { ReplicateConfig config; config.replica_num = 1; config.preferred_segment = "segment_0"; - auto put_start_result = - service_->PutStart(client_id, key1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key1, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key1, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Move key1 from "segment_0" to "segment_1" - auto move_result = - service_->CreateMoveTask(key1, "default", "segment_0", "segment_1"); + auto move_result = service_->CreateMoveTask(key1, TenantId::Default(), + "segment_0", "segment_1"); EXPECT_TRUE(move_result.has_value()); // Verify the move task is created and assigned to the client owning the @@ -3001,32 +3036,32 @@ TEST_F(MasterServiceSnapshotTest, CreateMoveTaskTest) { EXPECT_EQ(contexts[0].client_id, task.value().assigned_client); // Move non-existent key should fail - auto move_result1 = service_->CreateMoveTask("not_exist_key", "default", - "segment_0", "segment_1"); + auto move_result1 = service_->CreateMoveTask( + "not_exist_key", TenantId::Default(), "segment_0", "segment_1"); EXPECT_FALSE(move_result1.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_result1.error()); // Move to segment that is same as source should fail - auto move_result_same = - service_->CreateMoveTask(key1, "default", "segment_1", "segment_1"); + auto move_result_same = service_->CreateMoveTask(key1, TenantId::Default(), + "segment_1", "segment_1"); EXPECT_FALSE(move_result_same.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result_same.error()); // Move to segment that is not mounted should fail - auto move_result2 = service_->CreateMoveTask(key1, "default", "segment_0", - "not_mounted_segment"); + auto move_result2 = service_->CreateMoveTask( + key1, TenantId::Default(), "segment_0", "not_mounted_segment"); EXPECT_FALSE(move_result2.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result2.error()); // Move from segment that does not have the replica should fail - auto move_result3 = - service_->CreateMoveTask(key1, "default", "segment_2", "segment_1"); + auto move_result3 = service_->CreateMoveTask(key1, TenantId::Default(), + "segment_2", "segment_1"); EXPECT_FALSE(move_result3.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result3.error()); // Move from segment that is not mounted should fail auto move_result4 = service_->CreateMoveTask( - key1, "default", "not_mounted_segment", "segment_1"); + key1, TenantId::Default(), "not_mounted_segment", "segment_1"); EXPECT_FALSE(move_result4.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result4.error()); } @@ -3060,16 +3095,16 @@ TEST_F(MasterServiceSnapshotTest, QueryTaskTest) { ReplicateConfig config; config.replica_num = 1; config.preferred_segment = "segment_0"; - auto put_start_result = - service_->PutStart(client_id, key1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key1, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key1, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Move key1 from "segment_0" to "segment_1" - auto move_result = - service_->CreateMoveTask(key1, "default", "segment_0", "segment_1"); + auto move_result = service_->CreateMoveTask(key1, TenantId::Default(), + "segment_0", "segment_1"); EXPECT_TRUE(move_result.has_value()); // Query non-existent task should fail @@ -3111,20 +3146,22 @@ TEST_F(MasterServiceSnapshotTest, config.preferred_segment = "segment_0"; ASSERT_TRUE(service_ - ->PutStart(put_client_id, key, "default", + ->PutStart(put_client_id, key, TenantId::Default(), /*slice_length=*/1024, config) .has_value()); - ASSERT_TRUE( - service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(put_client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); // Create two tasks; both should be assigned to the client owning source // segment_0. - auto copy_task_id = service_->CreateCopyTask(key, "default", {"segment_1"}); + auto copy_task_id = + service_->CreateCopyTask(key, TenantId::Default(), {"segment_1"}); ASSERT_TRUE(copy_task_id.has_value()); - auto move_task_id = - service_->CreateMoveTask(key, "default", "segment_0", "segment_1"); + auto move_task_id = service_->CreateMoveTask(key, TenantId::Default(), + "segment_0", "segment_1"); ASSERT_TRUE(move_task_id.has_value()); // Fetch from client_0 should get both tasks (order not guaranteed). @@ -3174,17 +3211,18 @@ TEST_F(MasterServiceSnapshotTest, FetchTasksRespectsBatchSize) { config.preferred_segment = "segment_0"; ASSERT_TRUE(service_ - ->PutStart(put_client_id, key, "default", + ->PutStart(put_client_id, key, TenantId::Default(), /*slice_length=*/1024, config) .has_value()); - ASSERT_TRUE( - service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(put_client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); - auto t1 = service_->CreateCopyTask(key, "default", {"segment_1"}); + auto t1 = service_->CreateCopyTask(key, TenantId::Default(), {"segment_1"}); ASSERT_TRUE(t1.has_value()); - auto t2 = - service_->CreateMoveTask(key, "default", "segment_0", "segment_1"); + auto t2 = service_->CreateMoveTask(key, TenantId::Default(), "segment_0", + "segment_1"); ASSERT_TRUE(t2.has_value()); auto fetch_first = service_->FetchTasks(ctx0.client_id, /*batch_size=*/1); @@ -3226,15 +3264,17 @@ TEST_F(MasterServiceSnapshotTest, UpdateTaskSuccessFlow) { config.preferred_segment = "segment_0"; ASSERT_TRUE(service_ - ->PutStart(put_client_id, key, "default", + ->PutStart(put_client_id, key, TenantId::Default(), /*slice_length=*/1024, config) .has_value()); - ASSERT_TRUE( - service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(put_client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); // Create a task assigned to client owning segment_0. - auto task_id_res = service_->CreateCopyTask(key, "default", {"segment_1"}); + auto task_id_res = + service_->CreateCopyTask(key, TenantId::Default(), {"segment_1"}); ASSERT_TRUE(task_id_res.has_value()); const UUID task_id = task_id_res.value(); @@ -3284,15 +3324,16 @@ TEST_F(MasterServiceSnapshotTest, UpdateTaskRejectsWrongClient) { config.preferred_segment = "segment_0"; ASSERT_TRUE(service_ - ->PutStart(put_client_id, key, "default", + ->PutStart(put_client_id, key, TenantId::Default(), /*slice_length=*/1024, config) .has_value()); - ASSERT_TRUE( - service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(put_client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); - auto task_id_res = - service_->CreateMoveTask(key, "default", "segment_0", "segment_1"); + auto task_id_res = service_->CreateMoveTask(key, TenantId::Default(), + "segment_0", "segment_1"); ASSERT_TRUE(task_id_res.has_value()); const UUID task_id = task_id_res.value(); @@ -3349,8 +3390,9 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) { UUID client_id = generate_uuid(); // Test Case 1: CopyStart a non-existent key, should fail. - auto copy_result = service_->CopyStart( - client_id, "non_existent_key", "default", "segment_1", {"segment_2"}); + auto copy_result = + service_->CopyStart(client_id, "non_existent_key", TenantId::Default(), + "segment_1", {"segment_2"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_result.error()); @@ -3362,25 +3404,25 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); // Test Case 2: CopyStart to segment_2 and segment_3, should fail because // the only replica is not completed. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_2", "segment_3"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_2", "segment_3"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, copy_result.error()); // PutEnd the object. - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 3: CopyStart to segment_2 and segment_3, should success. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_2", "segment_3"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_2", "segment_3"}); EXPECT_TRUE(copy_result.has_value()); auto copy_response = copy_result.value(); EXPECT_EQ("segment_1", copy_response.source.get_memory_descriptor() @@ -3388,42 +3430,44 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) { EXPECT_EQ(2, copy_response.targets.size()); // Test Case 4: Try remove the object, should fail because it is copying. - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error()); // Test Case 5: CopyStart to segment_4, should fail because there is an // ongoing copy task. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_4"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_4"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, copy_result.error()); // Test Case 6: CopyEnd, should success and the object now has 3 replicas. - auto copy_end_result = service_->CopyEnd(client_id, key, "default"); + auto copy_end_result = + service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_end_result.has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(3, get_result.value().replicas.size()); // Test Case 7: Copy from a non-existent replica to segment_3 and // segment_4, should fail. copy_result = - service_->CopyStart(client_id, key, "default", "non_existent_segment", - {"segment_3", "segment_4"}); + service_->CopyStart(client_id, key, TenantId::Default(), + "non_existent_segment", {"segment_3", "segment_4"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, copy_result.error()); // Test Case 8: Copy to segment_4 and a non-existent segment, should fail. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_4", "non_existent_segment"}); + copy_result = + service_->CopyStart(client_id, key, TenantId::Default(), "segment_1", + {"segment_4", "non_existent_segment"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::SEGMENT_NOT_FOUND, copy_result.error()); // Test Case 9: Copy to segment_3 and segment_4, should skip segment_3 and // successfully copy to segment_4. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_3", "segment_4"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_3", "segment_4"}); EXPECT_TRUE(copy_result.has_value()); copy_response = copy_result.value(); EXPECT_EQ("segment_1", copy_response.source.get_memory_descriptor() @@ -3435,16 +3479,16 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) { .buffer_descriptor.transport_endpoint_); // End the copy operation to clean up state - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_end_result.has_value()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(4, get_result.value().replicas.size()); // Test Case 10: Copy to segment_4 again, should skip because it's already // used. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_4"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_4"}); EXPECT_TRUE(copy_result.has_value()); copy_response = copy_result.value(); EXPECT_EQ("segment_1", copy_response.source.get_memory_descriptor() @@ -3457,19 +3501,19 @@ TEST_F(MasterServiceSnapshotTest, CopyStart) { std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2)); // Test Case 11: Try remove the object, should fail because it is copying. - remove_result = service_->Remove(key, "default"); + remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, remove_result.error()); // Clean up the copy operation - // copy_end_result = service_->CopyEnd(client_id, key, "default"); + // copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); // EXPECT_TRUE(copy_end_result.has_value()); // Wait for the lease to expire // std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2)); // Test Case 12: Try remove the object, should success. - // remove_result = service_->Remove(key, "default"); + // remove_result = service_->Remove(key, TenantId::Default()); // EXPECT_TRUE(remove_result.has_value()); } @@ -3490,7 +3534,7 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) { // Test Case 1: CopyEnd a non-existent key, should fail. auto copy_end_result = - service_->CopyEnd(client_id, "non_existent_key", "default"); + service_->CopyEnd(client_id, "non_existent_key", TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_end_result.error()); @@ -3501,46 +3545,48 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 2: CopyEnd the object, should fail because there is no ongoing // copy task. - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, copy_end_result.error()); // CopyStart the object to segment_2 - auto copy_start_result = service_->CopyStart(client_id, key, "default", - "segment_1", {"segment_2"}); + auto copy_start_result = service_->CopyStart( + client_id, key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_start_result.has_value()); // Test Case 3: CopyEnd with an invalid client id, should fail. - copy_end_result = service_->CopyEnd(invalid_client_id, key, "default"); + copy_end_result = + service_->CopyEnd(invalid_client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, copy_end_result.error()); // Test Case 4: MoveEnd the object, should fail because the ongoing task is // Copy. - auto move_end_result = service_->MoveEnd(client_id, key, "default"); + auto move_end_result = + service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_end_result.error()); // Test Case 5: CopyEnd, should success. - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_end_result.has_value()); // Verify we now have 2 replicas - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(2, get_result.value().replicas.size()); // CopyStart the object from segment_1 to segment_3, then unmount segment_1 - copy_start_result = service_->CopyStart(client_id, key, "default", + copy_start_result = service_->CopyStart(client_id, key, TenantId::Default(), "segment_1", {"segment_3"}); ASSERT_TRUE(copy_start_result.has_value()); @@ -3551,10 +3597,10 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) { // Test Case 6: CopyEnd, should fail because the source is gone, the object // should have only 1 replica from segment_2. - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, copy_end_result.error()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); auto& replicas = get_result.value().replicas; EXPECT_EQ(1, replicas.size()); @@ -3563,7 +3609,7 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) { .buffer_descriptor.transport_endpoint_); // CopyStart the object from segment_2 to segment_3, then unmount segment_3 - copy_start_result = service_->CopyStart(client_id, key, "default", + copy_start_result = service_->CopyStart(client_id, key, TenantId::Default(), "segment_2", {"segment_3"}); ASSERT_TRUE(copy_start_result.has_value()); @@ -3574,10 +3620,10 @@ TEST_F(MasterServiceSnapshotTest, CopyEnd) { // Test Case 7: CopyEnd, should fail because the target is gone, the object // should have only 1 replica from segment_2. - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, copy_end_result.error()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); replicas = get_result.value().replicas; EXPECT_EQ(1, replicas.size()); @@ -3600,8 +3646,8 @@ TEST_F(MasterServiceSnapshotTest, CopyRevoke) { UUID invalid_client_id = generate_uuid(); // Test Case 1: CopyRevoke a non-existent key, should fail. - auto copy_revoke_result = - service_->CopyRevoke(client_id, "non_existent_key", "default"); + auto copy_revoke_result = service_->CopyRevoke( + client_id, "non_existent_key", TenantId::Default()); EXPECT_FALSE(copy_revoke_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_revoke_result.error()); @@ -3612,49 +3658,52 @@ TEST_F(MasterServiceSnapshotTest, CopyRevoke) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 2: CopyRevoke the object, should fail because there is no // ongoing copy task. - copy_revoke_result = service_->CopyRevoke(client_id, key, "default"); + copy_revoke_result = + service_->CopyRevoke(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_revoke_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, copy_revoke_result.error()); // CopyStart the object to segment_2 - auto copy_start_result = service_->CopyStart(client_id, key, "default", - "segment_1", {"segment_2"}); + auto copy_start_result = service_->CopyStart( + client_id, key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_start_result.has_value()); // Test Case 3: CopyRevoke with an invalid client id, should fail. copy_revoke_result = - service_->CopyRevoke(invalid_client_id, key, "default"); + service_->CopyRevoke(invalid_client_id, key, TenantId::Default()); EXPECT_FALSE(copy_revoke_result.has_value()); EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, copy_revoke_result.error()); // Test Case 4: MoveRevoke the object, should fail because the ongoing task // is Copy. - auto move_revoke_result = service_->MoveRevoke(client_id, key, "default"); + auto move_revoke_result = + service_->MoveRevoke(client_id, key, TenantId::Default()); EXPECT_FALSE(move_revoke_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_revoke_result.error()); // Test Case 5: CopyRevoke, should success. - copy_revoke_result = service_->CopyRevoke(client_id, key, "default"); + copy_revoke_result = + service_->CopyRevoke(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_revoke_result.has_value()); // Verify we still have 1 replica (the copy was revoked) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); // CopyStart the object from segment_1 to segment_2 again, then unmount // segment_1 - copy_start_result = service_->CopyStart(client_id, key, "default", + copy_start_result = service_->CopyStart(client_id, key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_start_result.has_value()); @@ -3665,11 +3714,12 @@ TEST_F(MasterServiceSnapshotTest, CopyRevoke) { // Test Case 6: CopyRevoke, should success even though the source is gone, // the object should be erased too. - copy_revoke_result = service_->CopyRevoke(client_id, key, "default"); + copy_revoke_result = + service_->CopyRevoke(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_revoke_result.has_value()); // Verify the object has been removed. - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); } @@ -3688,7 +3738,7 @@ TEST_F(MasterServiceSnapshotTest, MoveEnd) { // Test Case 1: MoveEnd a non-existent key, should fail. auto move_end_result = - service_->MoveEnd(client_id, "non_existent_key", "default"); + service_->MoveEnd(client_id, "non_existent_key", TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_end_result.error()); @@ -3699,47 +3749,49 @@ TEST_F(MasterServiceSnapshotTest, MoveEnd) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 2: MoveEnd the object, should fail because there is no ongoing // move task. - move_end_result = service_->MoveEnd(client_id, key, "default"); + move_end_result = service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, move_end_result.error()); // MoveStart the object to segment_2 - auto move_start_result = service_->MoveStart(client_id, key, "default", - "segment_1", "segment_2"); + auto move_start_result = service_->MoveStart( + client_id, key, TenantId::Default(), "segment_1", "segment_2"); ASSERT_TRUE(move_start_result.has_value()); // Test Case 3: MoveEnd with an invalid client id, should fail. - move_end_result = service_->MoveEnd(invalid_client_id, key, "default"); + move_end_result = + service_->MoveEnd(invalid_client_id, key, TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, move_end_result.error()); // Test Case 4: CopyEnd the object, should fail because the ongoing task is // Move. - auto copy_end_result = service_->CopyEnd(client_id, key, "default"); + auto copy_end_result = + service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_end_result.error()); // Test Case 5: MoveEnd, should success. - move_end_result = service_->MoveEnd(client_id, key, "default"); + move_end_result = service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(move_end_result.has_value()); // Verify we still have 1 replica (the move was successful) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); // MoveStart the object from segment_2 to segment_1 again, then unmount // segment_2 - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_2", "segment_1"); ASSERT_TRUE(move_start_result.has_value()); @@ -3749,7 +3801,7 @@ TEST_F(MasterServiceSnapshotTest, MoveEnd) { ASSERT_TRUE(unmount_result.has_value()); // Test Case 6: MoveEnd, should fail because the source is gone. - move_end_result = service_->MoveEnd(client_id, key, "default"); + move_end_result = service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, move_end_result.error()); } @@ -3767,8 +3819,8 @@ TEST_F(MasterServiceSnapshotTest, MoveRevoke) { UUID invalid_client_id = generate_uuid(); // Test Case 1: MoveRevoke a non-existent key, should fail. - auto move_revoke_result = - service_->MoveRevoke(client_id, "non_existent_key", "default"); + auto move_revoke_result = service_->MoveRevoke( + client_id, "non_existent_key", TenantId::Default()); EXPECT_FALSE(move_revoke_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_revoke_result.error()); @@ -3779,43 +3831,46 @@ TEST_F(MasterServiceSnapshotTest, MoveRevoke) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 2: MoveRevoke the object, should fail because there is no // ongoing move task. - move_revoke_result = service_->MoveRevoke(client_id, key, "default"); + move_revoke_result = + service_->MoveRevoke(client_id, key, TenantId::Default()); EXPECT_FALSE(move_revoke_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, move_revoke_result.error()); // MoveStart the object from segment_1 to segment_2 - auto move_start_result = service_->MoveStart(client_id, key, "default", - "segment_1", "segment_2"); + auto move_start_result = service_->MoveStart( + client_id, key, TenantId::Default(), "segment_1", "segment_2"); ASSERT_TRUE(move_start_result.has_value()); // Test Case 3: MoveRevoke with an invalid client id, should fail. move_revoke_result = - service_->MoveRevoke(invalid_client_id, key, "default"); + service_->MoveRevoke(invalid_client_id, key, TenantId::Default()); EXPECT_FALSE(move_revoke_result.has_value()); EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, move_revoke_result.error()); // Test Case 4: CopyRevoke the object, should fail because the ongoing task // is Move. - auto copy_revoke_result = service_->CopyRevoke(client_id, key, "default"); + auto copy_revoke_result = + service_->CopyRevoke(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_revoke_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_revoke_result.error()); // Test Case 5: MoveRevoke, should succeed. - move_revoke_result = service_->MoveRevoke(client_id, key, "default"); + move_revoke_result = + service_->MoveRevoke(client_id, key, TenantId::Default()); EXPECT_TRUE(move_revoke_result.has_value()); // Verify we still have 1 replica (the move was revoked) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); auto& replicas = get_result.value().replicas; EXPECT_EQ(1, replicas.size()); @@ -3825,7 +3880,7 @@ TEST_F(MasterServiceSnapshotTest, MoveRevoke) { // MoveStart the object from segment_1 to segment_2 again, then unmount // segment_1 - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_2"); ASSERT_TRUE(move_start_result.has_value()); @@ -3835,11 +3890,12 @@ TEST_F(MasterServiceSnapshotTest, MoveRevoke) { ASSERT_TRUE(unmount_result.has_value()); // Test Case 6: MoveRevoke, should succeed even though the source is gone. - move_revoke_result = service_->MoveRevoke(client_id, key, "default"); + move_revoke_result = + service_->MoveRevoke(client_id, key, TenantId::Default()); EXPECT_TRUE(move_revoke_result.has_value()); // The object should be erased as there is no replica left. - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); } @@ -3862,8 +3918,9 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) { UUID client_id = generate_uuid(); // Test Case 1: MoveStart a non-existent key, should fail. - auto move_start_result = service_->MoveStart( - client_id, "non_existent_key", "default", "segment_1", "segment_2"); + auto move_start_result = + service_->MoveStart(client_id, "non_existent_key", TenantId::Default(), + "segment_1", "segment_2"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_start_result.error()); @@ -3874,37 +3931,38 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); // Test Case 2: MoveStart the object, should fail because the only replica // is not completed. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_2"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, move_start_result.error()); // PutEnd the object. - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Copy the object to segment_3. - auto copy_start_result = service_->CopyStart(client_id, key, "default", - "segment_1", {"segment_3"}); + auto copy_start_result = service_->CopyStart( + client_id, key, TenantId::Default(), "segment_1", {"segment_3"}); ASSERT_TRUE(copy_start_result.has_value()); - auto copy_end_result = service_->CopyEnd(client_id, key, "default"); + auto copy_end_result = + service_->CopyEnd(client_id, key, TenantId::Default()); ASSERT_TRUE(copy_end_result.has_value()); // Test Case 3: MoveStart with source and target be the same, should fail. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_1"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(move_start_result.error(), ErrorCode::INVALID_PARAMS); // Test Case 4: MoveStart to segment_2, should succeed. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_2"); EXPECT_TRUE(move_start_result.has_value()); auto move_response = move_start_result.value(); @@ -3916,13 +3974,13 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) { .buffer_descriptor.transport_endpoint_); // Test Case 5: Try remove the object, should fail because it is moving. - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error()); // Test Case 6: MoveStart again, should fail because there is an ongoing // move task. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_3"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, @@ -3930,29 +3988,32 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) { // Test Case 7: MoveEnd, should succeed and the object now has 2 replicas // from segment_2 and segment_3 - auto move_end_result = service_->MoveEnd(client_id, key, "default"); + auto move_end_result = + service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(move_end_result.has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); auto& replicas = get_result.value().replicas; EXPECT_EQ(2, replicas.size()); // Test Case 8: Move from a non-existent replica to segment_1, should fail. - move_start_result = service_->MoveStart( - client_id, key, "default", "non_existent_segment", "segment_1"); + move_start_result = + service_->MoveStart(client_id, key, TenantId::Default(), + "non_existent_segment", "segment_1"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, move_start_result.error()); // Test Case 8.5: Move to a non-existent target segment, should fail. - move_start_result = service_->MoveStart( - client_id, key, "default", "segment_2", "non_existent_segment"); + move_start_result = + service_->MoveStart(client_id, key, TenantId::Default(), "segment_2", + "non_existent_segment"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::SEGMENT_NOT_FOUND, move_start_result.error()); // Test Case 9: Move to an already existing segment, should succeed but // return nullopt. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_2", "segment_3"); EXPECT_TRUE(move_start_result.has_value()); move_response = move_start_result.value(); @@ -3962,16 +4023,16 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) { // Test Case 10: Try remove the object, should fail because it is moving. std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2)); - remove_result = service_->Remove(key, "default"); + remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, remove_result.error()); // End the move. - move_end_result = service_->MoveEnd(client_id, key, "default"); + move_end_result = service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(move_end_result.has_value()); // Now the object should have only 1 replica on segment_3. - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); replicas = get_result.value().replicas; EXPECT_EQ(1, replicas.size()); @@ -3981,7 +4042,7 @@ TEST_F(MasterServiceSnapshotTest, MoveStart) { // Test Case 11: Try remove the object, should succeed after lease expires. std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2)); - remove_result = service_->Remove(key, "default"); + remove_result = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result.has_value()); } @@ -4013,38 +4074,38 @@ TEST_F(MasterServiceSnapshotTest, ProtectCopyMoveSourceFromEviction) { config.preferred_segment = "segment_1"; // Put two objects for move and copy tests. - auto put_start_result = service_->PutStart(client_id, copy_key, "default", - slice_length, config); + auto put_start_result = service_->PutStart( + client_id, copy_key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, copy_key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, copy_key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - put_start_result = service_->PutStart(client_id, move_key, "default", - slice_length, config); + put_start_result = service_->PutStart( + client_id, move_key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - put_end_result = - service_->PutEnd(client_id, move_key, "default", ReplicaType::MEMORY); + put_end_result = service_->PutEnd(client_id, move_key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Start copy and move operations. - auto copy_start_result = service_->CopyStart(client_id, copy_key, "default", - "segment_1", {"segment_2"}); + auto copy_start_result = service_->CopyStart( + client_id, copy_key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_start_result.has_value()); - auto move_start_result = service_->MoveStart(client_id, move_key, "default", - "segment_1", "segment_2"); + auto move_start_result = service_->MoveStart( + client_id, move_key, TenantId::Default(), "segment_1", "segment_2"); ASSERT_TRUE(move_start_result.has_value()); // Put more objects to trigger eviction. Do not prefer any segments. config.preferred_segment = ""; for (size_t i = 0; i < 128 * (kSegmentSize * 2 / slice_length); ++i) { std::string key = "test_key_" + std::to_string(i); - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { - auto put_end_result = service_->PutEnd(client_id, key, "default", - ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } else { // wait for eviction to work @@ -4058,10 +4119,12 @@ TEST_F(MasterServiceSnapshotTest, ProtectCopyMoveSourceFromEviction) { ASSERT_TRUE(remove_all_result > 0); // Try end copy and move operations, should success. - auto copy_end_result = service_->CopyEnd(client_id, copy_key, "default"); + auto copy_end_result = + service_->CopyEnd(client_id, copy_key, TenantId::Default()); EXPECT_TRUE(copy_end_result.has_value()); - auto move_end_result = service_->MoveEnd(client_id, move_key, "default"); + auto move_end_result = + service_->MoveEnd(client_id, move_key, TenantId::Default()); EXPECT_TRUE(move_end_result.has_value()); } diff --git a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h index 5d48eb25e2..5486e9daf0 100644 --- a/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h +++ b/mooncake-store/tests/ha/snapshot/master_service_test_for_snapshot_base.h @@ -1,6 +1,7 @@ #pragma once #include "master_service.h" +#include "master_snapshot_manager.h" #include "master_metric_manager.h" #include "segment.h" #include "ha/snapshot/catalog/snapshot_catalog_store.h" @@ -159,11 +160,43 @@ class MasterServiceSnapshotTestBase : public ::testing::Test { // ==================== Snapshot Helper Methods ==================== - // Wrapper method: Call MasterService's private method PersistState - // This class is a friend of MasterService, so it can access private members + // Wrapper method: Call MasterSnapshotManager's PersistState through + // MasterService This class is a friend of MasterService, so it can access + // private members static tl::expected CallPersistState( MasterService* service, const std::string& snapshot_id) { - return service->PersistState(snapshot_id); + // If snapshot_manager_ exists, use it; otherwise create a temporary one + if (service->snapshot_manager_) { + return service->snapshot_manager_->PersistState(snapshot_id); + } + + // For tests that don't have snapshot_manager_ initialized, + // we need to access the old implementation or create a temporary + // manager This is a temporary compatibility layer for tests + EnsureSnapshotStores(service); + + MasterSnapshotManagerOptions options; + options.enable_snapshot = true; + options.snapshot_interval_seconds = 300; + options.snapshot_child_timeout_seconds = 300; + options.snapshot_retention_count = 3; + options.snapshot_backup_dir = ""; + options.use_snapshot_backup_dir = false; + options.snapshot_catalog_store_type = + service->snapshot_catalog_store_type_; + options.snapshot_catalog_store_connstring = + service->snapshot_catalog_store_connstring_; + options.ha_backend_type = service->ha_backend_type_; + options.ha_backend_connstring = service->ha_backend_connstring_; + options.cluster_id = service->cluster_id_; + options.enable_ha = service->enable_ha_; + + auto temp_manager = std::make_unique( + service, options, service->snapshot_mutex_, + service->snapshot_object_store_.get(), + service->snapshot_catalog_store_.get()); + + return temp_manager->PersistState(snapshot_id); } static void EnsureSnapshotStores(MasterService* service) { @@ -220,7 +253,7 @@ class MasterServiceSnapshotTestBase : public ::testing::Test { ServiceStateSnapshot state; // === Basic State === - auto keys_result = service->GetAllKeys("default"); + auto keys_result = service->GetAllKeys(TenantId::Default()); if (keys_result.has_value()) { state.all_keys = std::move(keys_result.value()); std::sort(state.all_keys.begin(), state.all_keys.end()); @@ -233,7 +266,8 @@ class MasterServiceSnapshotTestBase : public ::testing::Test { } for (const auto& key : state.all_keys) { - auto replica_result = service->GetReplicaList(key, "default"); + auto replica_result = + service->GetReplicaList(key, TenantId::Default()); if (replica_result.has_value()) { state.replica_lists[key] = std::move(replica_result.value()); } @@ -545,10 +579,10 @@ class MasterServiceSnapshotTestBase : public ::testing::Test { const std::string key = "snapshot_putstart_consistency_key"; const uint64_t slice_length = 1024; - auto before = - original->PutStart(client_id, key, "default", slice_length, config); - auto after = - restored->PutStart(client_id, key, "default", slice_length, config); + auto before = original->PutStart(client_id, key, TenantId::Default(), + slice_length, config); + auto after = restored->PutStart(client_id, key, TenantId::Default(), + slice_length, config); ASSERT_EQ(before.has_value(), after.has_value()) << "PutStart has_value mismatch between original and restored " diff --git a/mooncake-store/tests/ha/snapshot/master_snapshot_codec_test.cpp b/mooncake-store/tests/ha/snapshot/master_snapshot_codec_test.cpp new file mode 100644 index 0000000000..c83c7e4f03 --- /dev/null +++ b/mooncake-store/tests/ha/snapshot/master_snapshot_codec_test.cpp @@ -0,0 +1,201 @@ +#include + +#include +#include +#include + +#include + +#include "ha/snapshot/master_snapshot_codec.h" +#include "master_config.h" +#include "master_service.h" +#include "segment.h" +#include "task_manager.h" +#include "tenant_id.h" +#include "utils/zstd_util.h" + +namespace mooncake::ha { + +class MasterSnapshotCodecTest : public ::testing::Test { + protected: + void SetUp() override { master_service_ = MakeMasterService(); } + + void TearDown() override { master_service_.reset(); } + + static std::unique_ptr MakeMasterService() { + MasterServiceConfig config; + config.default_kv_lease_ttl = 10000; + config.eviction_ratio = 0.1; + return std::make_unique(config); + } + + // The fixture is befriended by MasterService, so private state access is + // funneled through this helper (friendship is not inherited by the + // TEST_F-generated subclasses). + static MasterSnapshotStateView MakeStateView(MasterService& service) { + return MasterSnapshotStateView(service, service.segment_manager_, + service.nof_segment_manager_, + service.task_manager_); + } + + std::unique_ptr master_service_; +}; + +TEST_F(MasterSnapshotCodecTest, EncodeManifestPreservesSnapshotId) { + std::vector bytes = MasterSnapshotCodec::EncodeManifest( + MasterSnapshotCodec::kSerializerType, + MasterSnapshotCodec::kSerializerVersion, "snapshot-000042"); + std::string manifest(bytes.begin(), bytes.end()); + EXPECT_EQ(manifest, "messagepack|1.0.0|snapshot-000042"); +} + +TEST_F(MasterSnapshotCodecTest, EncodeDecodeRoundTrip) { + MasterSnapshotCodec codec; + + MasterSnapshotStateView state_view = MakeStateView(*master_service_); + + auto encode_result = codec.Encode(state_view); + ASSERT_TRUE(encode_result.has_value()) + << "Encode failed: " << encode_result.error().message; + + const MasterSnapshotPayloads& payloads = encode_result.value(); + + // All three payload buffers must be produced. + EXPECT_FALSE(payloads.metadata.empty()); + EXPECT_FALSE(payloads.segments.empty()); + EXPECT_FALSE(payloads.task_manager.empty()); + + // Decode into a fresh service. + auto target_service = MakeMasterService(); + auto decode_result = codec.Decode(target_service.get(), payloads); + ASSERT_TRUE(decode_result.has_value()) + << "Decode failed: " << decode_result.error().message; +} + +TEST_F(MasterSnapshotCodecTest, EncodeDecodeRoundTripWithMemoryReplica) { + // Mount a segment and store an object backed by a MEMORY replica. On + // decode, the segment/allocator must be restored before the metadata, + // otherwise deserializing the replica fails with SEGMENT_NOT_FOUND because + // GetMountedSegment() cannot find its backing segment. + constexpr size_t kSegmentBase = 0x300000000; + constexpr size_t kSegmentSize = 1024 * 1024 * 16; // 16MB + const std::string kKey = "memory_replica_key"; + const TenantId& kTenant = TenantId::Default(); + + Segment segment; + segment.id = generate_uuid(); + segment.name = "codec_test_segment"; + segment.base = kSegmentBase; + segment.size = kSegmentSize; + segment.te_endpoint = segment.name; + + UUID client_id = generate_uuid(); + auto mount_result = master_service_->MountSegment(segment, client_id); + ASSERT_TRUE(mount_result.has_value()); + + auto put_start = master_service_->PutStart( + client_id, kKey, kTenant, + /*slice_length=*/1024, ReplicateConfig{.replica_num = 1}); + ASSERT_TRUE(put_start.has_value()) + << "PutStart failed: " << static_cast(put_start.error()); + auto put_end = + master_service_->PutEnd(client_id, kKey, kTenant, ReplicaType::MEMORY); + ASSERT_TRUE(put_end.has_value()) + << "PutEnd failed: " << static_cast(put_end.error()); + + MasterSnapshotCodec codec; + MasterSnapshotStateView state_view = MakeStateView(*master_service_); + + auto encode_result = codec.Encode(state_view); + ASSERT_TRUE(encode_result.has_value()) + << "Encode failed: " << encode_result.error().message; + EXPECT_FALSE(encode_result.value().segments.empty()); + + // Decode into a fresh service. This exercises the segments-before-metadata + // restore order. + auto target_service = MakeMasterService(); + auto decode_result = + codec.Decode(target_service.get(), encode_result.value()); + ASSERT_TRUE(decode_result.has_value()) + << "Decode failed: " << decode_result.error().message; + + // The MEMORY replica must be fully restored and queryable. + auto get_result = target_service->GetReplicaList(kKey, kTenant); + ASSERT_TRUE(get_result.has_value()) + << "GetReplicaList failed: " << static_cast(get_result.error()); + EXPECT_EQ(get_result.value().replicas.size(), 1u); +} + +TEST_F(MasterSnapshotCodecTest, DecodeWithCorruptPayloadFails) { + MasterSnapshotCodec codec; + + MasterSnapshotPayloads corrupt; + corrupt.metadata = std::vector{1, 2, 3}; + corrupt.segments = std::vector{4, 5, 6}; + corrupt.task_manager = std::vector{7, 8, 9}; + + auto decode_result = codec.Decode(master_service_.get(), corrupt); + EXPECT_FALSE(decode_result.has_value()); + EXPECT_EQ(decode_result.error().code, ErrorCode::DESERIALIZE_FAIL); +} + +TEST_F(MasterSnapshotCodecTest, DecodeWithNullService) { + MasterSnapshotCodec codec; + + MasterSnapshotPayloads payloads; + auto decode_result = codec.Decode(nullptr, payloads); + EXPECT_FALSE(decode_result.has_value()); + EXPECT_EQ(decode_result.error().code, ErrorCode::INVALID_PARAMS); +} + +// Regression test: a structurally valid MessagePack task-manager payload whose +// task id field has the wrong type used to throw msgpack::type_error out of +// TaskManagerSerializer::Deserialize() (the arr[0].as() call sits +// outside the field-conversion try block). Since RestoreState() no longer +// wraps each candidate in a try/catch, an escaping exception here would abort +// restore and prevent fallback to an older healthy snapshot. Decode() must +// convert it into a SerializationError instead of throwing. +TEST_F(MasterSnapshotCodecTest, DecodeWithInvalidTaskFieldTypeReturnsError) { + MasterSnapshotCodec codec; + + // Start from a valid encoded snapshot so the segments and metadata payloads + // decode cleanly; we only want to corrupt the task-manager payload. + MasterSnapshotStateView state_view = MakeStateView(*master_service_); + auto encode_result = codec.Encode(state_view); + ASSERT_TRUE(encode_result.has_value()) + << "Encode failed: " << encode_result.error().message; + MasterSnapshotPayloads payloads = std::move(encode_result.value()); + + // Build a structurally valid MessagePack task-manager payload: an outer + // array of one task, the task itself a valid array with the expected field + // count, but the id field (index 0, expected string) is an integer. This + // unpacks cleanly and only fails at the arr[0].as() step, + // which used to throw msgpack::type_error out of Deserialize(). + constexpr size_t kTaskSerializedFields = 8; // must match the serializer + msgpack::sbuffer sbuf; + msgpack::packer packer(&sbuf); + packer.pack_array(1); // one task + packer.pack_array(kTaskSerializedFields); + packer.pack(static_cast(12345)); // id: wrong type (int, not str) + packer.pack(static_cast(0)); // type + packer.pack(static_cast(0)); // status + packer.pack(std::string("payload")); // payload + packer.pack(static_cast(0)); // created_at + packer.pack(static_cast(0)); // last_updated_at + packer.pack(std::string("message")); // message + packer.pack(std::string("assigned")); // assigned_client + + payloads.task_manager = zstd_compress( + reinterpret_cast(sbuf.data()), sbuf.size(), 3); + + // Decode a fresh service. It must not throw; it must report a serialization + // error so RestoreState() can fall back to another candidate snapshot. + auto target_service = MakeMasterService(); + tl::expected decode_result; + ASSERT_NO_THROW( + { decode_result = codec.Decode(target_service.get(), payloads); }); + EXPECT_FALSE(decode_result.has_value()); + EXPECT_EQ(decode_result.error().code, ErrorCode::DESERIALIZE_FAIL); +} + +} // namespace mooncake::ha diff --git a/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp b/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp index 5c8f6a5a24..0d9b58769b 100644 --- a/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp +++ b/mooncake-store/tests/ha/snapshot/snapshot_child_process_test.cpp @@ -1,11 +1,14 @@ #include "master_service.h" +#include "master_snapshot_manager.h" +#include "master_snapshot_repository.h" #include "master_metric_manager.h" #include "ha/snapshot/catalog/snapshot_catalog_store.h" #include "ha/snapshot/object/snapshot_object_store.h" #include "ha/snapshot/snapshot_test_utils.h" #ifdef STORE_USE_ETCD #include "etcd_helper.h" -#include "ha/oplog/etcd_oplog_store.h" +#include "ha/kv/etcd_ha_kv_backend.h" +#include "ha/oplog/oplog_batch_storage.h" #endif #include @@ -98,6 +101,28 @@ class SnapshotChildProcessTest : public ::testing::Test { .set_enable_snapshot(false) .set_enable_snapshot_restore(true) .set_enable_ha(true) + .set_enable_oplog(true) + .set_ha_backend_type("etcd") + .set_ha_backend_connstring(etcd_endpoints) + .set_cluster_id(cluster_id) + .set_snapshot_backup_dir(tmp_dir() + "/backup") + .set_snapshot_interval_seconds(100) + .set_snapshot_child_timeout_seconds(60) + .set_snapshot_retention_count(3) + .set_snapshot_object_store_type("local") + .set_view_version(view_version) + .build(); + service_ = std::make_unique(config); + } + + void CreateBatchEtcdHASnapshotService(const std::string& cluster_id, + const std::string& etcd_endpoints, + ViewVersionId view_version) { + auto config = MasterServiceConfigBuilder() + .set_enable_snapshot(false) + .set_enable_snapshot_restore(true) + .set_enable_ha(true) + .set_enable_oplog(true) .set_ha_backend_type("etcd") .set_ha_backend_connstring(etcd_endpoints) .set_cluster_id(cluster_id) @@ -115,28 +140,58 @@ class SnapshotChildProcessTest : public ::testing::Test { // Helper wrappers for private methods (friend access) std::string CallFormatTimestamp( const std::chrono::system_clock::time_point& tp) { - return service_->FormatTimestamp(tp); + // FormatTimestamp is now in MasterSnapshotManager + if (service_->snapshot_manager_) { + return service_->snapshot_manager_->FormatTimestamp(tp); + } + // Fallback for tests without snapshot_manager_ + auto temp_manager = CreateTempSnapshotManager(); + return temp_manager->FormatTimestamp(tp); } void CallHandleChildExit(pid_t pid, int status, const std::string& snapshot_id) { - service_->HandleChildExit(pid, status, snapshot_id); + if (service_->snapshot_manager_) { + service_->snapshot_manager_->HandleChildExit(pid, status, + snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + temp_manager->HandleChildExit(pid, status, snapshot_id); + } } void CallHandleChildTimeout(pid_t pid, const std::string& snapshot_id) { - service_->HandleChildTimeout(pid, snapshot_id); + if (service_->snapshot_manager_) { + service_->snapshot_manager_->HandleChildTimeout(pid, snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + temp_manager->HandleChildTimeout(pid, snapshot_id); + } } void CallCleanupOldSnapshot(int keep_count, const std::string& snapshot_id) { - service_->CleanupOldSnapshot(keep_count, snapshot_id); + if (service_->snapshot_manager_) { + service_->snapshot_manager_->repository_->CleanupOldSnapshots( + keep_count, snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + temp_manager->repository_->CleanupOldSnapshots(keep_count, + snapshot_id); + } } tl::expected CallUploadSnapshotPayloadFile( const std::vector& data, const std::string& path, const std::string& local_filename, const std::string& snapshot_id) { - return service_->UploadSnapshotPayloadFile(data, path, local_filename, - snapshot_id); + if (service_->snapshot_manager_) { + return service_->snapshot_manager_->repository_->UploadPayloadFile( + data, path, local_filename, snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + return temp_manager->repository_->UploadPayloadFile( + data, path, local_filename, snapshot_id); + } } SnapshotObjectStore* GetSnapshotObjectStore() { @@ -149,12 +204,22 @@ class SnapshotChildProcessTest : public ::testing::Test { tl::expected CallPersistState( const std::string& snapshot_id) { - return service_->PersistState(snapshot_id); + if (service_->snapshot_manager_) { + return service_->snapshot_manager_->PersistState(snapshot_id); + } else { + auto temp_manager = CreateTempSnapshotManager(); + return temp_manager->PersistState(snapshot_id); + } } tl::expected CallPersistState( const ha::SnapshotDescriptor& descriptor) { - return service_->PersistState(descriptor); + if (service_->snapshot_manager_) { + return service_->snapshot_manager_->PersistState(descriptor); + } else { + auto temp_manager = CreateTempSnapshotManager(); + return temp_manager->PersistState(descriptor); + } } bool GetUseSnapshotBackupDir() { @@ -166,7 +231,7 @@ class SnapshotChildProcessTest : public ::testing::Test { size_t shard_idx = svc->getShardIndex(key); auto& shard = svc->metadata_shards_[shard_idx]; SharedMutexLocker lock(&shard.mutex, shared_lock_t{}); - auto tenant_it = shard.tenants.find("default"); + auto tenant_it = shard.tenants.find(TenantId::Default()); return tenant_it != shard.tenants.end() && tenant_it->second.metadata.find(key) != tenant_it->second.metadata.end(); @@ -206,6 +271,47 @@ class SnapshotChildProcessTest : public ::testing::Test { return key + "_group"; } + private: + // Helper to create a temporary snapshot manager for tests + std::unique_ptr CreateTempSnapshotManager() { + EnsureSnapshotStores(); + + MasterSnapshotManagerOptions options; + options.enable_snapshot = true; + options.snapshot_interval_seconds = + service_->snapshot_interval_seconds_; + options.snapshot_child_timeout_seconds = + service_->snapshot_child_timeout_seconds_; + options.snapshot_retention_count = service_->snapshot_retention_count_; + options.snapshot_backup_dir = service_->snapshot_backup_dir_; + options.use_snapshot_backup_dir = service_->use_snapshot_backup_dir_; + options.snapshot_catalog_store_type = + service_->snapshot_catalog_store_type_; + options.snapshot_catalog_store_connstring = + service_->snapshot_catalog_store_connstring_; + options.ha_backend_type = service_->ha_backend_type_; + options.ha_backend_connstring = service_->ha_backend_connstring_; + options.cluster_id = service_->cluster_id_; + options.enable_ha = service_->enable_ha_; + + return std::make_unique( + service_.get(), options, service_->snapshot_mutex_, + service_->snapshot_object_store_.get(), + service_->snapshot_catalog_store_.get()); + } + + void EnsureSnapshotStores() { + if (!service_->snapshot_object_store_) { + service_->snapshot_object_store_ = SnapshotObjectStore::Create( + SnapshotObjectStoreType::LOCAL_FILE); + } + if (!service_->snapshot_catalog_store_ && + service_->snapshot_object_store_) { + service_->snapshot_catalog_store_ = + service_->CreateSnapshotCatalogStore(); + } + } + private: std::string tmp_dir_; }; @@ -360,6 +466,33 @@ TEST_F(SnapshotChildProcessTest, UploadSnapshotPayloadFile_Success) { // ========== Auto Snapshot Thread ========== +TEST_F(SnapshotChildProcessTest, DestructorInterruptsSnapshotThreadSleep) { + auto config = MasterServiceConfigBuilder() + .set_enable_snapshot(true) + .set_enable_snapshot_restore(false) + .set_snapshot_backup_dir(tmp_dir() + "/backup") + .set_snapshot_interval_seconds(30) + .set_snapshot_child_timeout_seconds(60) + .set_snapshot_retention_count(3) + .set_snapshot_object_store_type("local") + .build(); + auto auto_service = std::make_unique(config); + + // Let the background thread enter its snapshot interval wait. + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + const auto destroy_start = std::chrono::steady_clock::now(); + auto_service.reset(); + const auto elapsed_ms = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - destroy_start) + .count(); + + EXPECT_LT(elapsed_ms, 5000) + << "MasterService shutdown should not wait for the full snapshot " + "interval"; +} + TEST_F(SnapshotChildProcessTest, AutoSnapshot_GeneratesFiles) { // Create a service with snapshot enabled and short interval auto config = MasterServiceConfigBuilder() @@ -482,12 +615,14 @@ TEST_F(SnapshotChildProcessTest, RestoreRebuildsGroupedObjectRouting) { replicate_config.group_ids = std::vector{ FindGroupIdOnDifferentShard(service_.get(), key)}; - auto put_start = - service_->PutStart(client_id, key, "default", 1024, replicate_config); + auto put_start = service_->PutStart(client_id, key, TenantId::Default(), + 1024, replicate_config); ASSERT_TRUE(put_start.has_value()) << toString(put_start.error()); - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - ASSERT_TRUE(service_->ExistKey(key, "default").value_or(false)); + ASSERT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + ASSERT_TRUE(service_->ExistKey(key, TenantId::Default()).value_or(false)); auto persist_result = CallPersistState("20240701_130000_000"); ASSERT_TRUE(persist_result.has_value()) @@ -496,11 +631,62 @@ TEST_F(SnapshotChildProcessTest, RestoreRebuildsGroupedObjectRouting) { service_.reset(); service_ = std::make_unique(make_config()); - auto restored_replicas = service_->GetReplicaList(key, "default"); + auto restored_replicas = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(restored_replicas.has_value()) << "Grouped key should remain reachable by key after restore"; - ASSERT_TRUE(service_->Remove(key, "default", /*force=*/true).has_value()); - EXPECT_FALSE(service_->ExistKey(key, "default").value_or(true)); + ASSERT_TRUE( + service_->Remove(key, TenantId::Default(), /*force=*/true).has_value()); + EXPECT_FALSE(service_->ExistKey(key, TenantId::Default()).value_or(true)); +} + +TEST_F(SnapshotChildProcessTest, RestorePreservesObjectChecksum) { + auto make_config = [this]() { + return MasterServiceConfigBuilder() + .set_enable_snapshot(false) + .set_enable_snapshot_restore(true) + .set_snapshot_backup_dir(tmp_dir() + "/backup") + .set_snapshot_interval_seconds(100) + .set_snapshot_child_timeout_seconds(60) + .set_snapshot_retention_count(3) + .set_snapshot_object_store_type("local") + .set_default_kv_lease_ttl(600000) + .build(); + }; + service_ = std::make_unique(make_config()); + + Segment segment; + segment.id = generate_uuid(); + segment.name = "checksum_snapshot_segment"; + segment.base = 0x320000000; + segment.size = 1024 * 1024 * 16; + segment.te_endpoint = segment.name; + const UUID client_id = generate_uuid(); + ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value()); + + constexpr uint64_t kChecksum = 0x123456789ABCDEF0ULL; + const std::string key = "snapshot_object_checksum_key"; + ReplicateConfig replicate_config; + replicate_config.replica_num = 1; + auto put_start = service_->PutStart(client_id, key, TenantId::Default(), + 1024, replicate_config); + ASSERT_TRUE(put_start.has_value()) << toString(put_start.error()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, ObjectMeta{key, kChecksum}, + TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + ASSERT_TRUE(service_->ExistKey(key, TenantId::Default()).value_or(false)); + + auto persist_result = CallPersistState("20240701_140000_000"); + ASSERT_TRUE(persist_result.has_value()) + << "PersistState failed: " << persist_result.error().message; + + service_.reset(); + service_ = std::make_unique(make_config()); + + auto restored = service_->GetReplicaList(key, TenantId::Default()); + ASSERT_TRUE(restored.has_value()); + ASSERT_TRUE(restored->object_checksum.has_value()); + EXPECT_EQ(*restored->object_checksum, kChecksum); } TEST_F(SnapshotChildProcessTest, @@ -556,6 +742,57 @@ TEST_F(SnapshotChildProcessTest, EXPECT_FALSE(ObjectIsGroupedInMetadata(key, shard_idx)); } +TEST_F(SnapshotChildProcessTest, DeserializeMetadataSkipsInvalidClientId) { + CreateDefaultService(); + const std::string key = "invalid_client_id_snapshot_key"; + const uint32_t shard_idx = GetShardIndexForTest(key); + + msgpack::sbuffer shard_buffer; + MsgpackPacker shard_packer(&shard_buffer); + shard_packer.pack_map(1); + shard_packer.pack(std::string("metadata")); + shard_packer.pack_array(1); + shard_packer.pack_array(2); + shard_packer.pack(key); + + shard_packer.pack_array(8); + shard_packer.pack(std::string("not-a-uuid")); + shard_packer.pack(kDefaultTestPutStartTimeMs); + shard_packer.pack(kDefaultTestObjectSize); + shard_packer.pack(kDefaultTestLeaseTimeoutMs); + shard_packer.pack(false); + shard_packer.pack(uint64_t{0}); + shard_packer.pack(uint32_t{1}); + PackDiskReplica(shard_packer, kDefaultTestDiskFilePath, + kDefaultTestObjectSize); + + auto compressed_shard = + zstd_compress(reinterpret_cast(shard_buffer.data()), + shard_buffer.size(), 3); + + msgpack::sbuffer root_buffer; + MsgpackPacker root_packer(&root_buffer); + root_packer.pack_map(3); + root_packer.pack(std::string("shards")); + root_packer.pack_map(1); + root_packer.pack(shard_idx); + root_packer.pack_bin(compressed_shard.size()); + root_packer.pack_bin_body( + reinterpret_cast(compressed_shard.data()), + compressed_shard.size()); + root_packer.pack(std::string("discarded_replicas")); + root_packer.pack_array(0); + root_packer.pack(std::string("replica_next_id")); + root_packer.pack(uint64_t{10}); + + auto deserialize_result = + DeserializeMetadataForTest(ToByteVector(root_buffer)); + ASSERT_TRUE(deserialize_result.has_value()) + << deserialize_result.error().message; + + EXPECT_FALSE(KeyExistsInMetadata(service_.get(), key)); +} + TEST_F(SnapshotChildProcessTest, LegacyEtcdConnstringFallbackIsPreserved) { MasterConfig legacy_config; legacy_config.enable_ha = true; @@ -574,7 +811,7 @@ TEST_F(SnapshotChildProcessTest, LegacyEtcdConnstringFallbackIsPreserved) { #ifdef STORE_USE_ETCD TEST_F(SnapshotChildProcessTest, - PersistState_UsesEtcdOplogBoundaryInSnapshotDescriptor) { + PersistState_UsesBatchDurablePrefixBoundaryInSnapshotDescriptor) { const std::string etcd_endpoints = "127.0.0.1:2379"; auto connect_err = EtcdHelper::ConnectToEtcdStoreClient(etcd_endpoints.c_str()); @@ -583,25 +820,30 @@ TEST_F(SnapshotChildProcessTest, } const std::string cluster_id = - "snapshot-descriptor-" + UuidToString(generate_uuid()); + "snapshot-batch-boundary-" + UuidToString(generate_uuid()); constexpr ViewVersionId kViewVersion = 19; - constexpr uint64_t kLatestSequenceId = 123; - const std::string snapshot_id = "20240601_120000_456"; + constexpr uint64_t kBatchLatest = 1; - EtcdOpLogStore oplog_store(cluster_id); - auto init_err = oplog_store.Init(); + auto backend = std::make_shared(); + OpLogBatchStorage storage(cluster_id, *backend); + DurablePrefix prefix; + auto init_err = storage.InitDurablePrefix(prefix); if (init_err != ErrorCode::OK) { - GTEST_SKIP() << "failed to initialize etcd oplog store: " + GTEST_SKIP() << "failed to initialize batch durable prefix: " << toString(init_err); } - - auto update_err = oplog_store.UpdateLatestSequenceId(kLatestSequenceId); - if (update_err != ErrorCode::OK) { - GTEST_SKIP() << "failed to update etcd latest sequence id: " - << toString(update_err); - } - - CreateEtcdHASnapshotService(cluster_id, etcd_endpoints, kViewVersion); + OpLogBatchRecord batch{.batch_id = 1, + .first_seq = kBatchLatest, + .last_seq = kBatchLatest, + .entries = {{.sequence_id = kBatchLatest, + .op_type = OpType::PUT_END, + .object_key = "snapshot-boundary", + .payload = {}}}}; + ASSERT_EQ(ErrorCode::OK, storage.WriteBatchAndAdvancePrefix(batch, prefix)); + + const std::string snapshot_id = "20240601_120000_789"; + CreateBatchEtcdHASnapshotService(cluster_id, etcd_endpoints, kViewVersion); + ASSERT_EQ(ErrorCode::OK, service_->SetBatchOpLogBackendForTesting(backend)); auto persist_result = CallPersistState(snapshot_id); ASSERT_TRUE(persist_result.has_value()) << "PersistState failed: " << persist_result.error().message; @@ -612,9 +854,8 @@ TEST_F(SnapshotChildProcessTest, ASSERT_TRUE(latest.has_value()); ASSERT_TRUE(latest->has_value()); - EXPECT_EQ(latest->value().last_included_seq, kLatestSequenceId); + EXPECT_EQ(latest->value().last_included_seq, kBatchLatest); EXPECT_EQ(latest->value().producer_view_version, kViewVersion); - EXPECT_GT(latest->value().created_at_ms, 0); } #endif @@ -712,14 +953,15 @@ TEST_F(SnapshotChildProcessTest, << "MountSegment failed"; const std::string key1 = "restore_fallback_key_1"; - auto put1 = service_->PutStart(client_id, key1, "default", {1024}, + auto put1 = service_->PutStart(client_id, key1, TenantId::Default(), {1024}, {.replica_num = 1}); ASSERT_TRUE(put1.has_value()) << "PutStart for key1 failed"; ASSERT_TRUE( - service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY) + service_ + ->PutEnd(client_id, key1, TenantId::Default(), ReplicaType::MEMORY) .has_value()) << "PutEnd for key1 failed"; - EXPECT_TRUE(service_->ExistKey(key1, "default").value_or(false)) + EXPECT_TRUE(service_->ExistKey(key1, TenantId::Default()).value_or(false)) << "ExistKey should refresh lease for key1 before snapshot1"; const std::string snapshot_id1 = "20240702_120000_000"; @@ -729,14 +971,15 @@ TEST_F(SnapshotChildProcessTest, << persist_result.error().message; const std::string key2 = "restore_fallback_key_2"; - auto put2 = service_->PutStart(client_id, key2, "default", {1024}, + auto put2 = service_->PutStart(client_id, key2, TenantId::Default(), {1024}, {.replica_num = 1}); ASSERT_TRUE(put2.has_value()) << "PutStart for key2 failed"; ASSERT_TRUE( - service_->PutEnd(client_id, key2, "default", ReplicaType::MEMORY) + service_ + ->PutEnd(client_id, key2, TenantId::Default(), ReplicaType::MEMORY) .has_value()) << "PutEnd for key2 failed"; - EXPECT_TRUE(service_->ExistKey(key2, "default").value_or(false)) + EXPECT_TRUE(service_->ExistKey(key2, TenantId::Default()).value_or(false)) << "ExistKey should refresh lease for key2 before snapshot2"; const std::string snapshot_id2 = "20240702_120500_000"; @@ -767,9 +1010,11 @@ TEST_F(SnapshotChildProcessTest, .build(); auto restored_service = std::make_unique(restore_config); - EXPECT_TRUE(restored_service->ExistKey(key1, "default").value_or(false)) + EXPECT_TRUE( + restored_service->ExistKey(key1, TenantId::Default()).value_or(false)) << "Restore should fall back to the previous healthy snapshot"; - EXPECT_FALSE(restored_service->ExistKey(key2, "default").value_or(false)) + EXPECT_FALSE( + restored_service->ExistKey(key2, TenantId::Default()).value_or(false)) << "Corrupted latest snapshot must not be partially restored"; restored_service.reset(); @@ -840,22 +1085,23 @@ TEST_F(SnapshotChildProcessTest, RestoreCleansNonCompleteReplica) { // Add a complete object (clean data) std::string clean_key = "clean_object"; - auto put_result = service_->PutStart(client_id, clean_key, "default", - {1024}, {.replica_num = 1}); + auto put_result = service_->PutStart( + client_id, clean_key, TenantId::Default(), {1024}, {.replica_num = 1}); ASSERT_TRUE(put_result.has_value()) << "PutStart clean failed"; - auto put_end_result = - service_->PutEnd(client_id, clean_key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, clean_key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()) << "PutEnd clean failed"; // Add an incomplete object (PutStart without PutEnd -> non-COMPLETE) std::string dirty_key = "dirty_incomplete"; - auto put_dirty_result = service_->PutStart(client_id, dirty_key, "default", - {1024}, {.replica_num = 1}); + auto put_dirty_result = service_->PutStart( + client_id, dirty_key, TenantId::Default(), {1024}, {.replica_num = 1}); ASSERT_TRUE(put_dirty_result.has_value()) << "PutStart dirty failed"; // Intentionally NO PutEnd -> replica stays in PENDING status // Verify both keys exist in metadata before snapshot - EXPECT_TRUE(service_->ExistKey(clean_key, "default").value_or(false)); + EXPECT_TRUE( + service_->ExistKey(clean_key, TenantId::Default()).value_or(false)); EXPECT_TRUE(KeyExistsInMetadata(service_.get(), dirty_key)) << "Dirty key should exist in raw metadata after PutStart"; @@ -878,8 +1124,8 @@ TEST_F(SnapshotChildProcessTest, RestoreCleansNonCompleteReplica) { auto restored_service = std::make_unique(restore_config); // Step 4: Verify non-COMPLETE replica was cleaned, complete one remains - EXPECT_TRUE( - restored_service->ExistKey(clean_key, "default").value_or(false)) + EXPECT_TRUE(restored_service->ExistKey(clean_key, TenantId::Default()) + .value_or(false)) << "Complete object should survive restore"; EXPECT_FALSE(KeyExistsInMetadata(restored_service.get(), dirty_key)) << "Non-COMPLETE object should be cleaned from metadata during restore"; @@ -915,25 +1161,29 @@ TEST_F(SnapshotChildProcessTest, RestoreCleansExpiredLease) { // Add two complete objects via PutStart + PutEnd // Note: PutEnd calls GrantLease(0, ...) so lease is immediately expired std::string expired_key = "expired_lease_object"; - auto put_exp = service_->PutStart(client_id, expired_key, "default", {1024}, - {.replica_num = 1}); + auto put_exp = + service_->PutStart(client_id, expired_key, TenantId::Default(), {1024}, + {.replica_num = 1}); ASSERT_TRUE(put_exp.has_value()) << "PutStart expired failed"; - ASSERT_TRUE( - service_->PutEnd(client_id, expired_key, "default", ReplicaType::MEMORY) - .has_value()) + ASSERT_TRUE(service_ + ->PutEnd(client_id, expired_key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()) << "PutEnd expired failed"; std::string normal_key = "normal_lease_object"; - auto put_norm = service_->PutStart(client_id, normal_key, "default", {1024}, - {.replica_num = 1}); + auto put_norm = service_->PutStart( + client_id, normal_key, TenantId::Default(), {1024}, {.replica_num = 1}); ASSERT_TRUE(put_norm.has_value()) << "PutStart normal failed"; - ASSERT_TRUE( - service_->PutEnd(client_id, normal_key, "default", ReplicaType::MEMORY) - .has_value()) + ASSERT_TRUE(service_ + ->PutEnd(client_id, normal_key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()) << "PutEnd normal failed"; // ExistKey grants a fresh lease (now + 600s) to normal_key - EXPECT_TRUE(service_->ExistKey(normal_key, "default").value_or(false)); + EXPECT_TRUE( + service_->ExistKey(normal_key, TenantId::Default()).value_or(false)); // Do NOT call ExistKey on expired_key, its lease stays expired from PutEnd // Step 2: Persist state @@ -955,8 +1205,8 @@ TEST_F(SnapshotChildProcessTest, RestoreCleansExpiredLease) { auto restored_service = std::make_unique(restore_config); // Step 4: Verify normal data retained, expired-lease data cleaned - EXPECT_TRUE( - restored_service->ExistKey(normal_key, "default").value_or(false)) + EXPECT_TRUE(restored_service->ExistKey(normal_key, TenantId::Default()) + .value_or(false)) << "Normal object with valid lease should survive restore"; EXPECT_FALSE(KeyExistsInMetadata(restored_service.get(), expired_key)) << "Lease-expired object should be cleaned during restore"; diff --git a/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h b/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h index a160c1433a..d6f0c28ef8 100644 --- a/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h +++ b/mooncake-store/tests/ha/snapshot/snapshot_test_utils.h @@ -19,7 +19,7 @@ #include "master_config.h" #include "replica.h" #include "segment.h" -#include "serialize/serializer.hpp" +#include "serialize/serializer.h" #include "types.h" #include "utils/zstd_util.h" @@ -53,13 +53,15 @@ struct CatalogBackendParam { // kDataTypeAndHardPinned: // 9 + replica_count, data_type plus trailing hard_pinned // kWithGroupId: 10 + replica_count, data_type + hard_pinned + group_id -// (the current writer format) +// kWithObjectChecksum: +// 11 + replica_count, current writer fields + checksum enum class SnapshotMetadataFormat { kLegacy, kDataTypeOnly, kHardPinnedOnly, kDataTypeAndHardPinned, kWithGroupId, + kWithObjectChecksum, }; class ScopedEnvVar { @@ -167,8 +169,9 @@ inline std::vector WrapShardIntoMetadataRoot( return ToByteVector(root_buffer); } -inline std::vector BuildMetadataPayload( - const UUID& client_id, std::string_view object_key = kDefaultTestObjectKey, +inline std::vector BuildMetadataPayloadWithClientIdString( + std::string_view client_id, + std::string_view object_key = kDefaultTestObjectKey, std::string_view disk_file_path = kDefaultTestDiskFilePath, uint64_t object_size = kDefaultTestObjectSize, uint64_t put_start_time_ms = kDefaultTestPutStartTimeMs, @@ -177,18 +180,24 @@ inline std::vector BuildMetadataPayload( const bool include_data_type = format == SnapshotMetadataFormat::kDataTypeOnly || format == SnapshotMetadataFormat::kDataTypeAndHardPinned || - format == SnapshotMetadataFormat::kWithGroupId; + format == SnapshotMetadataFormat::kWithGroupId || + format == SnapshotMetadataFormat::kWithObjectChecksum; const bool include_hard_pinned = format == SnapshotMetadataFormat::kHardPinnedOnly || format == SnapshotMetadataFormat::kDataTypeAndHardPinned || - format == SnapshotMetadataFormat::kWithGroupId; + format == SnapshotMetadataFormat::kWithGroupId || + format == SnapshotMetadataFormat::kWithObjectChecksum; const bool include_group_id = - format == SnapshotMetadataFormat::kWithGroupId; + format == SnapshotMetadataFormat::kWithGroupId || + format == SnapshotMetadataFormat::kWithObjectChecksum; + const bool include_object_checksum = + format == SnapshotMetadataFormat::kWithObjectChecksum; constexpr uint32_t kReplicaCount = 1; // 7 leading fields + replicas + optional data_type/hard_pinned/group_id. const size_t array_size = 7 + kReplicaCount + (include_data_type ? 1 : 0) + (include_hard_pinned ? 1 : 0) + - (include_group_id ? 1 : 0); + (include_group_id ? 1 : 0) + + (include_object_checksum ? 1 : 0); msgpack::sbuffer shard_buffer; MsgpackPacker shard_packer(&shard_buffer); @@ -199,7 +208,7 @@ inline std::vector BuildMetadataPayload( shard_packer.pack(std::string(object_key)); shard_packer.pack_array(array_size); - shard_packer.pack(UuidToString(client_id)); + shard_packer.pack(std::string(client_id)); shard_packer.pack(put_start_time_ms); shard_packer.pack(object_size); shard_packer.pack(lease_timeout_ms); @@ -216,10 +225,25 @@ inline std::vector BuildMetadataPayload( if (include_group_id) { shard_packer.pack(std::string("test-group")); } + if (include_object_checksum) { + shard_packer.pack(uint64_t{0x123456789ABCDEF0ULL}); + } return WrapShardIntoMetadataRoot(shard_buffer); } +inline std::vector BuildMetadataPayload( + const UUID& client_id, std::string_view object_key = kDefaultTestObjectKey, + std::string_view disk_file_path = kDefaultTestDiskFilePath, + uint64_t object_size = kDefaultTestObjectSize, + uint64_t put_start_time_ms = kDefaultTestPutStartTimeMs, + uint64_t lease_timeout_ms = kDefaultTestLeaseTimeoutMs, + SnapshotMetadataFormat format = SnapshotMetadataFormat::kLegacy) { + return BuildMetadataPayloadWithClientIdString( + UuidToString(client_id), object_key, disk_file_path, object_size, + put_start_time_ms, lease_timeout_ms, format); +} + // Builds a metadata payload whose declared replica_count field is set to // `declared_replica_count` while no replicas are actually packed (the entry // array stays at the 7 leading fields). Used to verify the deserializer diff --git a/mooncake-store/tests/ha/standby/ha_metric_manager_test.cpp b/mooncake-store/tests/ha/standby/ha_metric_manager_test.cpp index eb8fe526cd..7034db0979 100644 --- a/mooncake-store/tests/ha/standby/ha_metric_manager_test.cpp +++ b/mooncake-store/tests/ha/standby/ha_metric_manager_test.cpp @@ -107,6 +107,58 @@ TEST_F(HAMetricManagerTest, TestRecordOpLogApplyLatency) { SUCCEED(); } +#ifdef MOONCAKE_ENABLE_OPLOG_PERF_METRICS +TEST_F(HAMetricManagerTest, TestBatchRecordMetrics) { + const auto batches_before = M().get_batch_record_durable_batches_total(); + const auto entries_before = M().get_batch_record_durable_entries_total(); + const auto retries_before = M().get_batch_record_retries_total(); + + M().inc_batch_record_durable_batches(); + M().inc_batch_record_durable_entries(8); + M().inc_batch_record_retries(); + M().set_batch_record_committed_queue_depth(3); + M().set_batch_record_callback_queue_depth(2); + M().set_batch_record_last_batch_id(9); + M().set_batch_record_durable_sequence(72); + M().observe_batch_record_batch_entries(8); + M().observe_batch_record_batch_bytes(1024); + M().observe_batch_record_txn_latency_us(500); + M().observe_batch_record_commit_to_durable_us(700); + M().observe_batch_record_callback_latency_us(20); + + EXPECT_EQ(batches_before + 1, M().get_batch_record_durable_batches_total()); + EXPECT_EQ(entries_before + 8, M().get_batch_record_durable_entries_total()); + EXPECT_EQ(retries_before + 1, M().get_batch_record_retries_total()); + EXPECT_EQ(3, M().get_batch_record_committed_queue_depth()); + EXPECT_EQ(2, M().get_batch_record_callback_queue_depth()); + EXPECT_EQ(9, M().get_batch_record_last_batch_id()); + EXPECT_EQ(72, M().get_batch_record_durable_sequence()); + + const std::string text = M().serialize_metrics(); + for (const char* name : { + "ha_batch_record_durable_batches_total", + "ha_batch_record_durable_entries_total", + "ha_batch_record_retry_total", + "ha_batch_record_committed_queue_depth", + "ha_batch_record_callback_queue_depth", + "ha_batch_record_last_batch_id", + "ha_batch_record_durable_sequence", + "ha_batch_record_batch_entries", + "ha_batch_record_batch_bytes", + "ha_batch_record_txn_latency_us", + "ha_batch_record_commit_to_durable_us", + "ha_batch_record_callback_latency_us", + }) { + EXPECT_NE(std::string::npos, text.find(name)) << name; + } +} +#else +TEST_F(HAMetricManagerTest, BatchRecordMetricsAreAbsentWhenDisabled) { + EXPECT_EQ(std::string::npos, + M().serialize_metrics().find("ha_batch_record_")); +} +#endif + // ========== 7.1.2 Metric serialization tests ========== TEST_F(HAMetricManagerTest, TestSerializeMetrics) { diff --git a/mooncake-store/tests/ha/standby/hot_standby_service_test.cpp b/mooncake-store/tests/ha/standby/hot_standby_service_test.cpp index ff90b75a77..317a48d799 100644 --- a/mooncake-store/tests/ha/standby/hot_standby_service_test.cpp +++ b/mooncake-store/tests/ha/standby/hot_standby_service_test.cpp @@ -4,11 +4,23 @@ #include #include +#include #include #include +#include #include +#include + #include "master_service.h" +#include "ha/kv/ha_kv_backend.h" +#include "ha/oplog/oplog_batch_codec.h" +#include "ha/oplog/oplog_batch_storage.h" +#include "ha/oplog/oplog_types.h" +#ifdef STORE_USE_ETCD +#include "etcd_helper.h" +#include "ha/kv/etcd_ha_kv_backend.h" +#endif namespace mooncake::test { @@ -39,7 +51,7 @@ LoadedSnapshot MakeSnapshot(std::string snapshot_id, uint64_t seq_id, metadata.client_id = UUID{1, 2}; metadata.size = size; metadata.last_sequence_id = seq_id; - snapshot.metadata.emplace_back(std::move(key), metadata); + snapshot.metadata.emplace_back("default", std::move(key), metadata); return snapshot; } @@ -88,7 +100,7 @@ std::unique_ptr CreateSnapshotOnlyReadyStandby( metadata.client_id = UUID{1, 2}; metadata.size = 4096; metadata.last_sequence_id = 42; - snapshot.metadata.emplace_back("key-1", metadata); + snapshot.metadata.emplace_back("default", "key-1", metadata); service->SetSnapshotProvider(std::make_unique( std::optional(snapshot))); @@ -249,39 +261,45 @@ TEST_F(HotStandbyServiceTest, TestPromote_WhenReady) { EXPECT_EQ(42u, service_->GetLatestAppliedSequenceId()); } -TEST_F(HotStandbyServiceTest, TestPromote_FinalCatchUp) { -#ifdef STORE_USE_ETCD - GTEST_SKIP() << "Requires real etcd and OpLog data to exercise final " - "catch-up logic."; -#else - GTEST_SKIP() << "Requires an OpLog-following standby runtime."; -#endif -} +TEST_F(HotStandbyServiceTest, TestPromoteAndExportSnapshot_FinalCatchUp) { + // Setup: snapshot-only standby with baseline seq=10 + config_.enable_snapshot_bootstrap = true; + config_.enable_oplog_following = false; + service_ = std::make_unique(config_); -TEST_F(HotStandbyServiceTest, TestPromote_WithGaps) { -#ifdef STORE_USE_ETCD - GTEST_SKIP() - << "Requires real etcd and gaps in OpLog to validate gap resolution."; -#else - GTEST_SKIP() << "Requires an OpLog-following standby runtime."; -#endif -} + LoadedSnapshot snapshot; + snapshot.snapshot_id = "snap-001"; + snapshot.snapshot_sequence_id = 10; -TEST_F(HotStandbyServiceTest, TestPromote_Timeout) { -#ifdef STORE_USE_ETCD - GTEST_SKIP() - << "Requires real etcd and slow reads to trigger catch-up timeout."; -#else - GTEST_SKIP() << "Requires an OpLog-following standby runtime."; -#endif -} + StandbyObjectMetadata metadata; + metadata.client_id = UUID{1, 2}; + metadata.size = 4096; + metadata.last_sequence_id = 10; + snapshot.metadata.emplace_back("default", "key-1", metadata); -TEST_F(HotStandbyServiceTest, TestPromote_BatchLimit) { -#ifdef STORE_USE_ETCD - GTEST_SKIP() << "Requires real etcd and large OpLog to hit batch limit."; -#else - GTEST_SKIP() << "Requires an OpLog-following standby runtime."; -#endif + service_->SetSnapshotProvider(std::make_unique( + std::optional(snapshot))); + + ASSERT_EQ(ErrorCode::OK, service_->Start("", "", cluster_id_)); + EXPECT_EQ(StandbyState::WATCHING, service_->GetState()); + EXPECT_EQ(10u, service_->GetLatestAppliedSequenceId()); + + // Export before promotion: seq should be 10 + StandbySnapshot pre_snapshot; + EXPECT_TRUE(service_->ExportStandbySnapshot(pre_snapshot)); + EXPECT_EQ(10u, pre_snapshot.oplog_sequence_id); + + // Promote and export atomically + StandbySnapshot post_snapshot; + ErrorCode err = service_->PromoteAndExportSnapshot(post_snapshot); + EXPECT_EQ(ErrorCode::OK, err); + EXPECT_EQ(StandbyState::STOPPED, service_->GetState()); + + // After promotion, exported seq should still be 10 (no new OpLog in + // snapshot-only) + EXPECT_EQ(10u, post_snapshot.oplog_sequence_id); + ASSERT_EQ(1u, post_snapshot.objects.size()); + EXPECT_EQ("key-1", post_snapshot.objects[0].key); } // ========== 6.1.5 Warm start tests ========== @@ -340,11 +358,11 @@ TEST_F(HotStandbyServiceTest, TestStart_SnapshotOnlyWithSnapshot) { EXPECT_EQ(42u, status.primary_seq_id); EXPECT_TRUE(status.is_connected); - std::vector> exported; + std::vector exported; EXPECT_TRUE(service_->ExportMetadataSnapshot(exported)); ASSERT_EQ(1u, exported.size()); - EXPECT_EQ("key-1", exported[0].first); - EXPECT_EQ(4096u, exported[0].second.size); + EXPECT_EQ("key-1", exported[0].key); + EXPECT_EQ(4096u, exported[0].metadata.size); } TEST_F(HotStandbyServiceTest, @@ -372,12 +390,12 @@ TEST_F(HotStandbyServiceTest, EXPECT_EQ(84u, service_->GetLatestAppliedSequenceId()); EXPECT_EQ(1u, service_->GetMetadataCount()); - std::vector> exported; + std::vector exported; ASSERT_TRUE(service_->ExportMetadataSnapshot(exported)); ASSERT_EQ(1u, exported.size()); - EXPECT_EQ("key-new", exported[0].first); - EXPECT_EQ(8192u, exported[0].second.size); - EXPECT_EQ(84u, exported[0].second.last_sequence_id); + EXPECT_EQ("key-new", exported[0].key); + EXPECT_EQ(8192u, exported[0].metadata.size); + EXPECT_EQ(84u, exported[0].metadata.last_sequence_id); } TEST_F(HotStandbyServiceTest, TestStart_SnapshotOnlyWhenProviderFails) { @@ -400,7 +418,7 @@ TEST_F(HotStandbyServiceTest, TestGetMetadataCount) { } TEST_F(HotStandbyServiceTest, TestExportMetadataSnapshot) { - std::vector> snapshot; + std::vector snapshot; EXPECT_TRUE(service_->ExportMetadataSnapshot(snapshot)); EXPECT_TRUE(snapshot.empty()); } @@ -410,6 +428,45 @@ TEST_F(HotStandbyServiceTest, TestGetLatestAppliedSequenceId) { EXPECT_EQ(0u, seq); } +// ========== 6.1.6a ExportStandbySnapshot tests ========== + +TEST_F(HotStandbyServiceTest, TestExportStandbySnapshot_NotRunning) { + StandbySnapshot snapshot; + EXPECT_FALSE(service_->ExportStandbySnapshot(snapshot)); +} + +TEST_F(HotStandbyServiceTest, TestExportStandbySnapshot_SnapshotOnly) { + service_ = CreateSnapshotOnlyReadyStandby(config_, cluster_id_); + + StandbySnapshot snapshot; + EXPECT_TRUE(service_->ExportStandbySnapshot(snapshot)); + EXPECT_EQ(42u, snapshot.oplog_sequence_id); + ASSERT_EQ(1u, snapshot.objects.size()); + EXPECT_EQ("key-1", snapshot.objects[0].key); + EXPECT_EQ(4096u, snapshot.objects[0].metadata.size); + // Segment registry is empty because snapshot-only standby has no + // oplog_applier + EXPECT_TRUE(snapshot.segments.empty()); +} + +TEST_F(HotStandbyServiceTest, TestExportStandbySnapshot_Empty) { + config_.enable_snapshot_bootstrap = true; + config_.enable_oplog_following = false; + service_ = std::make_unique(config_); + + service_->SetSnapshotProvider(std::make_unique( + std::optional(LoadedSnapshot{}))); + + EXPECT_EQ(ErrorCode::OK, service_->Start("", "", cluster_id_)); + EXPECT_EQ(StandbyState::WATCHING, service_->GetState()); + + StandbySnapshot snapshot; + EXPECT_TRUE(service_->ExportStandbySnapshot(snapshot)); + EXPECT_EQ(0u, snapshot.oplog_sequence_id); + EXPECT_TRUE(snapshot.objects.empty()); + EXPECT_TRUE(snapshot.segments.empty()); +} + // ========== 6.1.7 Replication loop tests ========== TEST_F(HotStandbyServiceTest, TestReplicationLoop_UpdatesMetrics) { @@ -462,6 +519,426 @@ TEST_F(HotStandbyServiceTest, TestVerificationLoop_WhenDisabled) { #endif } +// ========== Issue 2 fail-closed catch-up ========== + +namespace { + +// Helper to create a valid OpLogEntry with checksum (mirrors the helper in +// oplog_applier_test.cpp). +OpLogEntry MakeEntry(uint64_t seq, OpType type, const std::string& key, + const std::string& payload) { + OpLogEntry e; + e.sequence_id = seq; + e.timestamp_ms = std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count(); + e.op_type = type; + e.object_key = key; + e.payload = payload; + // Compute checksum and prefix_hash using the production algorithm. + e.checksum = + static_cast(XXH32(payload.data(), payload.size(), 0)); + e.prefix_hash = + key.empty() ? 0 + : static_cast(XXH32(key.data(), key.size(), 0)); + return e; +} + +// Helper to create a valid struct_pack payload for PUT_END +std::string MakeValidPayload(uint64_t client_id_first = 1, + uint64_t client_id_second = 2, + uint64_t size = 1024) { + mooncake::MetadataPayload payload; + payload.client_id = {client_id_first, client_id_second}; + payload.size = size; + auto result = struct_pack::serialize(payload); + return std::string(result.begin(), result.end()); +} + +class FakeHaKvBackend : public HaKvBackend { + public: + ErrorCode Get(std::string_view key, std::string& value) override { + if (next_get_error_ != ErrorCode::OK) { + auto error = next_get_error_; + next_get_error_ = ErrorCode::OK; + return error; + } + if (get_error_ != ErrorCode::OK) { + return get_error_; + } + auto it = values_.find(std::string(key)); + if (it == values_.end()) { + return ErrorCode::ETCD_KEY_NOT_EXIST; + } + value = it->second; + return ErrorCode::OK; + } + ErrorCode Put(std::string_view key, std::string_view value) override { + values_[std::string(key)] = std::string(value); + return ErrorCode::OK; + } + ErrorCode Range(std::string_view begin_key, std::string_view end_key, + size_t limit, std::vector& kvs) override { + if (range_error_ != ErrorCode::OK) { + return range_error_; + } + kvs.clear(); + for (const auto& [key, value] : values_) { + if (key >= begin_key && key < end_key) { + kvs.push_back({.key = key, .value = value}); + if (kvs.size() >= limit) break; + } + } + return ErrorCode::OK; + } + bool SupportsTxn() const override { return true; } + ErrorCode Txn(const KvTxn&) override { return ErrorCode::OK; } + void SetGetError(ErrorCode err) { get_error_ = err; } + void SetRangeError(ErrorCode err) { range_error_ = err; } + void FailNextGet(ErrorCode err) { next_get_error_ = err; } + + private: + std::map values_; + ErrorCode get_error_{ErrorCode::OK}; + ErrorCode range_error_{ErrorCode::OK}; + ErrorCode next_get_error_{ErrorCode::OK}; +}; + +OpLogBatchRecord MakeBatch(uint64_t batch_id, uint64_t first_seq, + uint64_t last_seq) { + OpLogBatchRecord batch; + batch.batch_id = batch_id; + batch.first_seq = first_seq; + batch.last_seq = last_seq; + for (uint64_t seq = first_seq; seq <= last_seq; ++seq) { + batch.entries.push_back(MakeEntry(seq, OpType::PUT_END, + "batch_key_" + std::to_string(seq), + MakeValidPayload())); + } + return batch; +} + +} // namespace + +TEST_F(HotStandbyServiceTest, + BatchRecordStandbyPollsInjectedBackendWithoutNotifier) { + const std::string cluster_id = "batch-standby-injected-backend"; + auto batch_backend = std::make_shared(); + ASSERT_EQ(ErrorCode::OK, + batch_backend->Put( + BuildDurablePrefixKey(cluster_id), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + ASSERT_EQ(ErrorCode::OK, + batch_backend->Put(BuildBatchRecordKey(cluster_id, 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + + HotStandbyConfig config = config_; + config.enable_snapshot_bootstrap = false; + config.enable_oplog_following = true; + config.oplog_poll_interval_ms = 1; + + auto service = std::make_unique(config); + service->SetCatchUpBatchKvBackendForTesting(batch_backend); + ASSERT_EQ(ErrorCode::OK, + service->Start("primary_unused", "unused", cluster_id)); + + for (int i = 0; i < 100 && service->GetLatestAppliedSequenceId() < 1; ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + + EXPECT_EQ(StandbyState::WATCHING, service->GetState()); + EXPECT_EQ(1u, service->GetLatestAppliedSequenceId()); + service->Stop(); +} + +TEST_F(HotStandbyServiceTest, BatchRecordRetriesTransientBackendFailure) { + const std::string cluster_id = "batch-standby-transient-retry"; + auto batch_backend = std::make_shared(); + ASSERT_EQ(ErrorCode::OK, + batch_backend->Put( + BuildDurablePrefixKey(cluster_id), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + ASSERT_EQ(ErrorCode::OK, + batch_backend->Put(BuildBatchRecordKey(cluster_id, 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + batch_backend->FailNextGet(ErrorCode::ETCD_OPERATION_ERROR); + + HotStandbyConfig config = config_; + config.enable_snapshot_bootstrap = false; + config.oplog_poll_interval_ms = 1; + config.batch_oplog_retry_timeout_sec = 1; + + auto service = std::make_unique(config); + service->SetCatchUpBatchKvBackendForTesting(batch_backend); + ASSERT_EQ(ErrorCode::OK, + service->Start("primary_unused", "unused", cluster_id)); + + for (int i = 0; + i < 100 && (service->GetLatestAppliedSequenceId() < 1 || + service->GetSyncStatus().last_error != ErrorCode::OK); + ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + + EXPECT_EQ(StandbyState::WATCHING, service->GetState()); + EXPECT_EQ(1u, service->GetLatestAppliedSequenceId()); + EXPECT_EQ(ErrorCode::OK, service->GetSyncStatus().last_error); + service->Stop(); +} + +TEST_F(HotStandbyServiceTest, BatchRecordRetryTimeoutTransitionsToFailed) { + auto batch_backend = std::make_shared(); + batch_backend->SetGetError(ErrorCode::ETCD_OPERATION_ERROR); + + HotStandbyConfig config = config_; + config.enable_snapshot_bootstrap = false; + config.oplog_poll_interval_ms = 1; + config.batch_oplog_retry_timeout_sec = 0; + + auto service = std::make_unique(config); + service->SetCatchUpBatchKvBackendForTesting(batch_backend); + ASSERT_EQ(ErrorCode::OK, service->Start("primary_unused", "unused", + "batch-standby-timeout")); + + for (int i = 0; i < 100 && service->GetState() != StandbyState::FAILED; + ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + + EXPECT_EQ(StandbyState::FAILED, service->GetState()); + EXPECT_EQ(ErrorCode::ETCD_OPERATION_ERROR, + service->GetSyncStatus().last_error); + service->Stop(); +} + +TEST_F(HotStandbyServiceTest, BatchRecordCanRestartAfterRetryTimeout) { + auto batch_backend = std::make_shared(); + batch_backend->SetGetError(ErrorCode::ETCD_OPERATION_ERROR); + + HotStandbyConfig config = config_; + config.enable_snapshot_bootstrap = false; + config.oplog_poll_interval_ms = 1; + config.batch_oplog_retry_timeout_sec = 0; + + auto service = std::make_unique(config); + service->SetCatchUpBatchKvBackendForTesting(batch_backend); + ASSERT_EQ(ErrorCode::OK, service->Start("primary_unused", "unused", + "batch-standby-restart")); + for (int i = 0; i < 100 && service->GetState() != StandbyState::FAILED; + ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + ASSERT_EQ(StandbyState::FAILED, service->GetState()); + + batch_backend->SetGetError(ErrorCode::OK); + ASSERT_EQ(ErrorCode::OK, service->Start("primary_unused", "unused", + "batch-standby-restart")); + + EXPECT_EQ(StandbyState::WATCHING, service->GetState()); + EXPECT_EQ(ErrorCode::OK, service->GetSyncStatus().last_error); + service->Stop(); +} + +class PromotionCatchUpTest : public HotStandbyServiceTest { + protected: + void SetUp() override { + HotStandbyServiceTest::SetUp(); + config_.enable_oplog_following = true; + config_.enable_snapshot_bootstrap = false; + config_.oplog_poll_interval_ms = 1; + + batch_backend_ = std::make_shared(); + service_ = std::make_unique(config_); + service_->SetCatchUpBatchKvBackendForTesting(batch_backend_); + } + + std::shared_ptr batch_backend_; +}; + +#ifdef STORE_USE_ETCD +TEST_F(HotStandbyServiceTest, + BatchRecordStandbyPollsDurablePrefixInReplicationLoop) { + const std::string etcd_endpoints = "127.0.0.1:2379"; + auto connect_err = + EtcdHelper::ConnectToEtcdStoreClient(etcd_endpoints.c_str()); + if (connect_err != ErrorCode::OK) { + GTEST_SKIP() << "etcd is unavailable: " << toString(connect_err); + } + + const std::string cluster_id = + "batch-standby-loop-" + UuidToString(generate_uuid()); + EtcdHaKvBackend backend; + OpLogBatchStorage storage(cluster_id, backend); + DurablePrefix prefix; + auto init_err = storage.InitDurablePrefix(prefix); + if (init_err != ErrorCode::OK) { + GTEST_SKIP() << "etcd is unavailable: " << toString(init_err); + } + ASSERT_EQ(ErrorCode::OK, + storage.WriteBatchAndAdvancePrefix( + MakeBatch(prefix.batch_id + 1, prefix.last_seq + 1, + prefix.last_seq + 1), + prefix)); + + HotStandbyConfig config = config_; + config.enable_snapshot_bootstrap = false; + config.enable_oplog_following = true; + config.oplog_poll_interval_ms = 10; + + auto service = std::make_unique(config); + ASSERT_EQ(ErrorCode::OK, + service->Start("primary_unused", etcd_endpoints, cluster_id)); + + const uint64_t expected_seq = prefix.last_seq + 1; + constexpr int kMaxAttempts = 100; + for (int i = 0; i < kMaxAttempts && + service->GetLatestAppliedSequenceId() < expected_seq; + ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + + EXPECT_GE(service->GetLatestAppliedSequenceId(), expected_seq); + service->Stop(); +} +#endif + +TEST_F(PromotionCatchUpTest, UsesDurablePrefixLastSeqAsCatchUpTarget) { + ASSERT_EQ(ErrorCode::OK, + batch_backend_->Put( + BuildDurablePrefixKey(cluster_id_), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 2}))); + ASSERT_EQ(ErrorCode::OK, + batch_backend_->Put(BuildBatchRecordKey(cluster_id_, 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 2)))); + + auto err = service_->Start("", oplog_endpoints_, cluster_id_); + if (err != ErrorCode::OK) { + GTEST_SKIP() << "Service could not reach WATCHING state; " + "skipping promotion test"; + } + EXPECT_EQ(StandbyState::WATCHING, service_->GetState()); + + StandbySnapshot out; + ErrorCode promote_err = service_->PromoteAndExportSnapshot(out); + EXPECT_EQ(ErrorCode::OK, promote_err); + EXPECT_EQ(2u, out.oplog_sequence_id); +} + +TEST_F(PromotionCatchUpTest, RetriesTransientDurablePrefixReadFailure) { + auto batch_backend = std::make_shared(); + ASSERT_EQ(ErrorCode::OK, + batch_backend->Put( + BuildDurablePrefixKey(cluster_id_), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + ASSERT_EQ(ErrorCode::OK, + batch_backend->Put(BuildBatchRecordKey(cluster_id_, 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + service_->SetCatchUpBatchKvBackendForTesting(batch_backend); + + ASSERT_EQ(ErrorCode::OK, + service_->Start("", oplog_endpoints_, cluster_id_)); + batch_backend->FailNextGet(ErrorCode::ETCD_OPERATION_ERROR); + + StandbySnapshot out; + EXPECT_EQ(ErrorCode::OK, service_->PromoteAndExportSnapshot(out)); + EXPECT_EQ(1u, out.oplog_sequence_id); +} + +TEST_F(PromotionCatchUpTest, MissingDurablePrefixPromotesAtSequenceZero) { + ASSERT_EQ(ErrorCode::OK, + service_->Start("", oplog_endpoints_, cluster_id_)); + ASSERT_EQ(StandbyState::WATCHING, service_->GetState()); + StandbySnapshot out; + ASSERT_EQ(ErrorCode::OK, service_->PromoteAndExportSnapshot(out)); + EXPECT_EQ(0u, out.oplog_sequence_id); +} + +TEST_F(PromotionCatchUpTest, MissingDurablePrefixRejectsNonzeroSequence) { + config_.enable_snapshot_bootstrap = true; + service_ = std::make_unique(config_); + service_->SetCatchUpBatchKvBackendForTesting(batch_backend_); + service_->SetSnapshotProvider(std::make_unique( + std::optional(MakeSnapshot("baseline", 1, "key", 1)))); + + ASSERT_EQ(ErrorCode::OK, + service_->Start("", oplog_endpoints_, cluster_id_)); + ASSERT_EQ(StandbyState::WATCHING, service_->GetState()); + StandbySnapshot out; + EXPECT_EQ(ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, + service_->PromoteAndExportSnapshot(out)); + EXPECT_EQ(StandbyState::FAILED, service_->GetState()); +} + +TEST_F(PromotionCatchUpTest, CatchesUpPrefixThatAppearsBeforePromotion) { + ASSERT_EQ(ErrorCode::OK, + service_->Start("", oplog_endpoints_, cluster_id_)); + ASSERT_EQ(StandbyState::WATCHING, service_->GetState()); + ASSERT_EQ(ErrorCode::OK, + batch_backend_->Put( + BuildDurablePrefixKey(cluster_id_), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + ASSERT_EQ(ErrorCode::OK, + batch_backend_->Put(BuildBatchRecordKey(cluster_id_, 1), + EncodeOpLogBatchRecord(MakeBatch(1, 1, 1)))); + + StandbySnapshot out; + ASSERT_EQ(ErrorCode::OK, service_->PromoteAndExportSnapshot(out)); + EXPECT_EQ(1u, out.oplog_sequence_id); +} + +TEST_F(PromotionCatchUpTest, PaginatesBatchRecords) { + constexpr uint64_t kBatchCount = 1025; + auto batch_backend = std::make_shared(); + ASSERT_EQ( + ErrorCode::OK, + batch_backend->Put(BuildDurablePrefixKey(cluster_id_), + EncodeDurablePrefix({.batch_id = kBatchCount, + .last_seq = kBatchCount}))); + for (uint64_t id = 1; id <= kBatchCount; ++id) { + ASSERT_EQ( + ErrorCode::OK, + batch_backend->Put(BuildBatchRecordKey(cluster_id_, id), + EncodeOpLogBatchRecord(MakeBatch(id, id, id)))); + } + service_->SetCatchUpBatchKvBackendForTesting(batch_backend); + + auto err = service_->Start("", oplog_endpoints_, cluster_id_); + if (err != ErrorCode::OK) { + GTEST_SKIP() << "Service could not reach WATCHING state; " + "skipping promotion test"; + } + + StandbySnapshot out; + ASSERT_EQ(ErrorCode::OK, service_->PromoteAndExportSnapshot(out)); + EXPECT_EQ(kBatchCount, out.oplog_sequence_id); +} + +TEST_F(PromotionCatchUpTest, FailsPromotionWhenDurablePrefixUnreadable) { + ASSERT_EQ(ErrorCode::OK, + service_->Start("", oplog_endpoints_, cluster_id_)); + batch_backend_->SetGetError(ErrorCode::PERSISTENT_FAIL); + + StandbySnapshot out; + ErrorCode promote_err = service_->PromoteAndExportSnapshot(out); + EXPECT_EQ(ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, promote_err); + EXPECT_EQ(StandbyState::FAILED, service_->GetState()); +} + +TEST_F(PromotionCatchUpTest, FailsPromotionWhenTargetBatchUnreadable) { + ASSERT_EQ(ErrorCode::OK, + service_->Start("", oplog_endpoints_, cluster_id_)); + ASSERT_EQ(ErrorCode::OK, + batch_backend_->Put( + BuildDurablePrefixKey(cluster_id_), + EncodeDurablePrefix({.batch_id = 1, .last_seq = 1}))); + batch_backend_->SetRangeError(ErrorCode::PERSISTENT_FAIL); + + StandbySnapshot out; + ErrorCode promote_err = service_->PromoteAndExportSnapshot(out); + EXPECT_EQ(ErrorCode::INCOMPLETE_OPLOG_CATCH_UP, promote_err); + EXPECT_EQ(StandbyState::FAILED, service_->GetState()); +} + } // namespace mooncake::test int main(int argc, char** argv) { diff --git a/mooncake-store/tests/ha/standby/hot_standby_snapshot_bootstrap_test.cpp b/mooncake-store/tests/ha/standby/hot_standby_snapshot_bootstrap_test.cpp index 091d884651..c9a425b563 100644 --- a/mooncake-store/tests/ha/standby/hot_standby_snapshot_bootstrap_test.cpp +++ b/mooncake-store/tests/ha/standby/hot_standby_snapshot_bootstrap_test.cpp @@ -114,13 +114,13 @@ TEST_P(HotStandbySnapshotBootstrapTest, EXPECT_EQ(descriptor_.last_included_seq, status.applied_seq_id); EXPECT_EQ(descriptor_.last_included_seq, status.primary_seq_id); - std::vector> exported; + std::vector exported; ASSERT_TRUE(service.ExportMetadataSnapshot(exported)); ASSERT_EQ(1u, exported.size()); - EXPECT_EQ(kDefaultTestObjectKey, exported.front().first); - EXPECT_EQ(kDefaultTestObjectSize, exported.front().second.size); + EXPECT_EQ(kDefaultTestObjectKey, exported.front().key); + EXPECT_EQ(kDefaultTestObjectSize, exported.front().metadata.size); EXPECT_EQ(descriptor_.last_included_seq, - exported.front().second.last_sequence_id); + exported.front().metadata.last_sequence_id); } TEST_P(HotStandbySnapshotBootstrapTest, @@ -164,6 +164,91 @@ TEST(StandbyControllerTest, PromoteStandbyReturnsStartFailure) { EXPECT_EQ(ErrorCode::INVALID_PARAMS, controller->PromoteStandby()); } +TEST(StandbyControllerTest, OplogEnablementControlsReaderController) { + ha::HABackendSpec spec{ + .type = ha::HABackendType::ETCD, + .connstring = "http://localhost:2379", + .cluster_namespace = "oplog-reader-gate-test", + }; + MasterServiceSupervisorConfig config; + + auto disabled = ha::CreateStandbyController(spec, config); + ASSERT_NE(disabled, nullptr); + EXPECT_EQ(ErrorCode::OK, disabled->PromoteStandby()); + + config.enable_oplog = true; + auto enabled = ha::CreateStandbyController(spec, config); + ASSERT_NE(enabled, nullptr); + EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS, + enabled->PromoteStandby()); + + spec.type = ha::HABackendType::REDIS; + auto unsupported = ha::CreateStandbyController(spec, config); + ASSERT_NE(unsupported, nullptr); + EXPECT_EQ(ErrorCode::OK, unsupported->PromoteStandby()); +} + +TEST(StandbyControllerTest, + PromoteStandbyAndExport_ReturnsErrorWhenNotStarted) { + ha::HABackendSpec spec{ + .type = ha::HABackendType::UNKNOWN, + .connstring = "", + .cluster_namespace = "promotion-context-test", + }; + MasterServiceSupervisorConfig config; + config.cluster_id = "promotion-context-test"; + config.local_hostname = "127.0.0.1:50051"; + config.enable_snapshot_restore = true; + config.snapshot_object_store_type = "local"; + config.snapshot_catalog_store_type = "invalid"; + + auto controller = ha::CreateStandbyController(spec, config); + ASSERT_NE(controller, nullptr); + + // When never started + auto result = controller->PromoteStandbyAndExport(); + EXPECT_FALSE(result.has_value()); + EXPECT_EQ(ErrorCode::UNAVAILABLE_IN_CURRENT_STATUS, result.error()); + + // After failed start + EXPECT_EQ(ErrorCode::INVALID_PARAMS, + controller->StartStandby(std::nullopt)); + result = controller->PromoteStandbyAndExport(); + EXPECT_FALSE(result.has_value()); + EXPECT_EQ(ErrorCode::INVALID_PARAMS, result.error()); +} + +TEST_P(HotStandbySnapshotBootstrapTest, + PromoteStandbyAndExport_SuccessWithSnapshot) { + PublishSnapshot(); + + ha::HABackendSpec spec{ + .type = ha::HABackendType::UNKNOWN, + .connstring = "", + .cluster_namespace = cluster_id_, + }; + + auto config = MakeSnapshotProviderConfig(GetParam(), cluster_id_, + FLAGS_redis_endpoint); + config.enable_snapshot_restore = true; + config.local_hostname = "127.0.0.1:50051"; + + auto controller = ha::CreateStandbyController(spec, config); + ASSERT_NE(controller, nullptr); + + EXPECT_EQ(ErrorCode::OK, controller->StartStandby(std::nullopt)); + + auto result = controller->PromoteStandbyAndExport(); + ASSERT_TRUE(result.has_value()) << toString(result.error()); + + const auto& ctx = result.value(); + EXPECT_EQ(descriptor_.last_included_seq, ctx.applied_seq_id); + ASSERT_EQ(1u, ctx.objects.size()); + EXPECT_EQ(kDefaultTestObjectKey, ctx.objects[0].key); + EXPECT_EQ(kDefaultTestObjectSize, ctx.objects[0].metadata.size); + EXPECT_TRUE(ctx.segments.empty()); +} + INSTANTIATE_TEST_SUITE_P( SnapshotCatalogBackends, HotStandbySnapshotBootstrapTest, ::testing::ValuesIn(BuildCatalogBackendParams()), diff --git a/mooncake-store/tests/host_port_fix_test.cpp b/mooncake-store/tests/host_port_fix_test.cpp new file mode 100644 index 0000000000..6609ef9e7d --- /dev/null +++ b/mooncake-store/tests/host_port_fix_test.cpp @@ -0,0 +1,48 @@ +#include +#include +#include + +#include "common.h" +#include "utils.h" + +namespace mooncake { +namespace { + +// Regression: --host=ip:port should not break coro_rpc_server binding. +// When --host includes a port (e.g. 127.0.0.1:18007) for the TransferEngine +// data plane, the port must be stripped before the address is passed to +// coro_rpc_server, otherwise the server fails to start. + +class HostPortFixTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + FLAGS_logtostderr = true; + google::InitGoogleLogging("HostPortFixTest"); + } + static void TearDownTestSuite() { google::ShutdownGoogleLogging(); } +}; + +TEST_F(HostPortFixTest, StripsPortForRpcServer) { + auto rpc_bind_host = getHostNameWithoutPort("127.0.0.1:18007"); + EXPECT_EQ(rpc_bind_host, "127.0.0.1"); + + auto port = getFreeTcpPort(); + coro_rpc::coro_rpc_server server(/*thread_num=*/1, port, rpc_bind_host); + auto ec = server.async_start(); + EXPECT_FALSE(ec.hasResult()) << "Server should start with bare hostname"; + server.stop(); +} + +TEST_F(HostPortFixTest, BareHostPassesThrough) { + auto rpc_bind_host = getHostNameWithoutPort("127.0.0.1"); + EXPECT_EQ(rpc_bind_host, "127.0.0.1"); + + auto port = getFreeTcpPort(); + coro_rpc::coro_rpc_server server(/*thread_num=*/1, port, rpc_bind_host); + auto ec = server.async_start(); + EXPECT_FALSE(ec.hasResult()) << "Server should start with bare hostname"; + server.stop(); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/http_metadata_server_test.cpp b/mooncake-store/tests/http_metadata_server_test.cpp index 5ff426a592..26b3726b9f 100644 --- a/mooncake-store/tests/http_metadata_server_test.cpp +++ b/mooncake-store/tests/http_metadata_server_test.cpp @@ -96,4 +96,19 @@ TEST_F(HttpMetadataServerTest, RejectsChangedRpcMetaRepublish) { server.stop(); } +TEST_F(HttpMetadataServerTest, StartReportsBindFailure) { + int port = getFreeTcpPort(); + HttpMetadataServer first(static_cast(port), "127.0.0.1"); + ASSERT_TRUE(first.start()); + WaitUntilReady(port); + + // A second server cannot bind the already-taken port; start() must report + // the failure instead of claiming a healthy server that never came up. + HttpMetadataServer second(static_cast(port), "127.0.0.1"); + EXPECT_FALSE(second.start()); + EXPECT_FALSE(second.is_running()); + + first.stop(); +} + } // namespace mooncake::testing diff --git a/mooncake-store/tests/kv_event_publisher_test.cpp b/mooncake-store/tests/kv_event_publisher_test.cpp new file mode 100644 index 0000000000..9230e21aeb --- /dev/null +++ b/mooncake-store/tests/kv_event_publisher_test.cpp @@ -0,0 +1,218 @@ +#include + +#include +#include +#include +#include +#include +#include + +#include "kv_event/key_util.h" +#include "kv_event/kv_event_publisher.h" + +#if defined(MOONCAKE_ENABLE_KV_EVENTS) && MOONCAKE_ENABLE_KV_EVENTS +#include +#include +#endif + +namespace mooncake { +namespace { + +TEST(KvEventKeyUtilTest, ParseSeqHashFromObjectKey) { + EXPECT_EQ(ParseSeqHashFromObjectKey("12345"), 12345u); + EXPECT_EQ(ParseSeqHashFromObjectKey("0x2a"), 42u); + EXPECT_EQ(ParseSeqHashFromObjectKey("0XFF"), 255u); + EXPECT_FALSE(ParseSeqHashFromObjectKey("").has_value()); + EXPECT_FALSE(ParseSeqHashFromObjectKey("not-a-hash").has_value()); + EXPECT_FALSE(ParseSeqHashFromObjectKey("123abc").has_value()); + EXPECT_EQ(KvEventPublisher::ParseSeqHashFromObjectKey("99"), 99u); +} + +TEST(KvEventPublisherTest, DisabledPublisherIsNoop) { + KvEventConfig config; + config.enabled = false; + KvEventPublisher publisher(config); + EXPECT_FALSE(publisher.enabled()); + publisher.PublishStored("42", "cpu"); + publisher.PublishRemoved("42", "cpu"); + const auto stats = publisher.GetStats(); + EXPECT_EQ(stats.published_events, 0u); + EXPECT_EQ(stats.dropped_events, 0u); +} + +#if defined(MOONCAKE_ENABLE_KV_EVENTS) && MOONCAKE_ENABLE_KV_EVENTS + +namespace { + +std::string MakeIpcEndpoint() { + return "ipc:///tmp/kv_event_test_" + std::to_string(getpid()) + "_" + + std::to_string( + std::chrono::steady_clock::now().time_since_epoch().count()); +} + +bool ReceiveZmqMultipart(void* socket, std::vector& frames) { + frames.clear(); + while (true) { + zmq_msg_t msg; + if (zmq_msg_init(&msg) != 0) { + return false; + } + const int rc = zmq_msg_recv(&msg, socket, 0); + if (rc < 0) { + zmq_msg_close(&msg); + return false; + } + const char* data = static_cast(zmq_msg_data(&msg)); + frames.emplace_back(data, data + zmq_msg_size(&msg)); + const int more = zmq_msg_more(&msg); + zmq_msg_close(&msg); + if (!more) { + break; + } + } + return true; +} + +} // namespace + +TEST(KvEventPublisherTest, PublishesSglangObjectKeyOverZmq) { + const std::string endpoint = MakeIpcEndpoint(); + const std::string object_key = + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855_0_k"; + const std::string group_id = + "sglang-hicache:" + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"; + + KvEventConfig config; + config.enabled = true; + config.bind_endpoint = endpoint; + config.backend_id = "mooncake-test"; + config.emit_object_key = true; + config.emit_legacy_compat_fields = true; + config.queue_capacity = 64; + KvEventPublisher publisher(config); + ASSERT_TRUE(publisher.enabled()); + + void* ctx = zmq_ctx_new(); + ASSERT_NE(ctx, nullptr); + void* sub = zmq_socket(ctx, ZMQ_SUB); + ASSERT_NE(sub, nullptr); + ASSERT_EQ(zmq_connect(sub, endpoint.c_str()), 0); + ASSERT_EQ(zmq_setsockopt(sub, ZMQ_SUBSCRIBE, "", 0), 0); + + // Allow SUB connect before first publish propagates. + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + + publisher.PublishStored(object_key, "cpu", TenantId("tenant-a"), group_id); + + std::vector frames; + ASSERT_TRUE(ReceiveZmqMultipart(sub, frames)) << zmq_strerror(zmq_errno()); + ASSERT_EQ(frames.size(), 3u); + EXPECT_TRUE(frames[0].empty()); + ASSERT_EQ(frames[1].size(), sizeof(uint64_t)); + + const auto object_handle = + msgpack::unpack(frames[2].data(), frames[2].size()); + const auto& root = object_handle.get(); + ASSERT_EQ(root.type, msgpack::type::ARRAY); + ASSERT_EQ(root.via.array.size, 3u); + + const auto& events = root.via.array.ptr[1]; + ASSERT_EQ(events.type, msgpack::type::ARRAY); + ASSERT_EQ(events.via.array.size, 1u); + + const auto& event = events.via.array.ptr[0]; + ASSERT_EQ(event.type, msgpack::type::MAP); + + bool has_object_key = false; + bool has_group_id = false; + bool has_empty_seq_hashes = false; + std::string event_type; + std::string backend_id; + std::string tenant_id; + for (uint32_t i = 0; i < event.via.map.size; ++i) { + const auto& key = event.via.map.ptr[i].key; + const auto& val = event.via.map.ptr[i].val; + ASSERT_EQ(key.type, msgpack::type::STR); + const std::string field(key.via.str.ptr, key.via.str.size); + if (field == "object_key") { + ASSERT_EQ(val.type, msgpack::type::STR); + EXPECT_EQ(std::string(val.via.str.ptr, val.via.str.size), + object_key); + has_object_key = true; + } else if (field == "group_id") { + ASSERT_EQ(val.type, msgpack::type::STR); + EXPECT_EQ(std::string(val.via.str.ptr, val.via.str.size), group_id); + has_group_id = true; + } else if (field == "seq_hashes") { + ASSERT_EQ(val.type, msgpack::type::ARRAY); + EXPECT_EQ(val.via.array.size, 0u); + has_empty_seq_hashes = true; + } else if (field == "event_type") { + ASSERT_EQ(val.type, msgpack::type::STR); + event_type = std::string(val.via.str.ptr, val.via.str.size); + } else if (field == "backend_id") { + ASSERT_EQ(val.type, msgpack::type::STR); + backend_id = std::string(val.via.str.ptr, val.via.str.size); + } else if (field == "tenant_id") { + ASSERT_EQ(val.type, msgpack::type::STR); + tenant_id = std::string(val.via.str.ptr, val.via.str.size); + } + } + + EXPECT_EQ(event_type, "stored"); + EXPECT_EQ(backend_id, "mooncake-test"); + EXPECT_EQ(tenant_id, "tenant-a"); + EXPECT_TRUE(has_object_key); + EXPECT_TRUE(has_group_id); + EXPECT_TRUE(has_empty_seq_hashes); + + // Wait for async worker to finish publishing. + for (int i = 0; i < 50; ++i) { + const auto stats = publisher.GetStats(); + if (stats.published_events == 1 && stats.published_batches == 1) { + EXPECT_EQ(stats.dropped_events, 0u); + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + const auto stats = publisher.GetStats(); + EXPECT_EQ(stats.published_events, 1u); + EXPECT_EQ(stats.published_batches, 1u); + EXPECT_EQ(stats.dropped_events, 0u); + + zmq_close(sub); + zmq_ctx_destroy(ctx); +} + +TEST(KvEventPublisherTest, DropsOldestWhenQueueFull) { + const std::string endpoint = MakeIpcEndpoint(); + + KvEventConfig config; + config.enabled = true; + config.bind_endpoint = endpoint; + config.backend_id = "mooncake-test"; + config.emit_object_key = true; + config.queue_capacity = 2; + KvEventPublisher publisher(config); + ASSERT_TRUE(publisher.enabled()); + + for (int i = 0; i < 100; ++i) { + publisher.PublishStored(std::to_string(i), "cpu"); + } + + for (int i = 0; i < 50; ++i) { + const auto stats = publisher.GetStats(); + if (stats.dropped_events >= 1) { + EXPECT_GE(stats.dropped_events, 1u); + return; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + FAIL() << "expected dropped_events after queue overflow"; +} + +#endif // MOONCAKE_ENABLE_KV_EVENTS + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/localfs_hot_standby_integration_test.cpp b/mooncake-store/tests/localfs_hot_standby_integration_test.cpp deleted file mode 100644 index 7dd8d40336..0000000000 --- a/mooncake-store/tests/localfs_hot_standby_integration_test.cpp +++ /dev/null @@ -1,463 +0,0 @@ -// mooncake-store/tests/localfs_hot_standby_integration_test.cpp -// -// End-to-end integration tests for the HA replication flow using LocalFS -// backend. Mirrors the structure of hot_standby_integration_test.cpp but -// replaces etcd with a shared temp directory so the tests run without any -// external service. -// -// Test flow: -// Primary OpLogManager -> LocalFsOpLogStore (WRITER) -// -> shared filesystem directory <- -// HotStandbyService (PollingOpLogChangeNotifier -> OpLogReplicator -// -> OpLogApplier -> StandbyMetadataStore) - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "hot_standby_service.h" -#include "ha/oplog/oplog_manager.h" -#include "ha/oplog/oplog_store_factory.h" -#include "standby_state_machine.h" - -namespace mooncake { -namespace testing { - -// ============================================================ -// RAII helper to ensure HotStandbyService is always stopped -// ============================================================ - -class StandbyServiceGuard { - public: - explicit StandbyServiceGuard(HotStandbyService* service) - : service_(service) {} - ~StandbyServiceGuard() { - if (service_) { - service_->Stop(); - // LocalFS has no goroutines to drain; a short sleep suffices to - // let background polling threads exit cleanly. - std::this_thread::sleep_for(std::chrono::milliseconds(500)); - } - } - StandbyServiceGuard(const StandbyServiceGuard&) = delete; - StandbyServiceGuard& operator=(const StandbyServiceGuard&) = delete; - - private: - HotStandbyService* service_; -}; - -// ============================================================ -// Fixture -// ============================================================ - -class LocalFsHotStandbyIntegrationTest : public ::testing::Test { - protected: - void SetUp() override { - // Generate a unique temp directory per test - static std::atomic counter{0}; - test_dir_ = "/tmp/localfs_ha_test_" + std::to_string(getpid()) + "_" + - std::to_string(counter.fetch_add(1)); - cluster_id_ = "localfs_ha_cluster"; - poll_interval_ms_ = 100; // fast polling for tests - - std::filesystem::create_directories(test_dir_); - } - - void TearDown() override { - std::error_code ec; - std::filesystem::remove_all(test_dir_, ec); - if (ec) { - LOG(WARNING) << "Failed to remove test dir " << test_dir_ << ": " - << ec.message(); - } - } - - // -- Helpers -- - - std::unique_ptr CreatePrimaryOpLogManager() { - auto store = OpLogStoreFactory::Create( - OpLogStoreType::LOCAL_FS, cluster_id_, OpLogStoreRole::WRITER, - test_dir_, poll_interval_ms_); - if (!store) return nullptr; - auto mgr = std::make_unique(); - mgr->SetOpLogStore(std::shared_ptr(std::move(store))); - return mgr; - } - - HotStandbyConfig MakeHotStandbyConfig() { - HotStandbyConfig cfg; - cfg.enable_verification = false; - cfg.max_replication_lag_entries = 1000; - cfg.oplog_store_type = OpLogStoreType::LOCAL_FS; - cfg.oplog_store_root_dir = test_dir_; - cfg.oplog_poll_interval_ms = poll_interval_ms_; - return cfg; - } - - bool WaitForSync(HotStandbyService& standby, uint64_t target_seq, - int timeout_sec = 30) { - auto deadline = std::chrono::steady_clock::now() + - std::chrono::seconds(timeout_sec); - while (std::chrono::steady_clock::now() < deadline) { - auto status = standby.GetSyncStatus(); - LOG(INFO) << "Standby: state=" << StandbyStateToString(status.state) - << ", applied_seq_id=" << status.applied_seq_id - << ", primary_seq_id=" << status.primary_seq_id - << ", lag_entries=" << status.lag_entries; - if (status.state == StandbyState::WATCHING && - status.lag_entries == 0 && - status.applied_seq_id >= target_seq) { - return true; - } - std::this_thread::sleep_for(std::chrono::milliseconds(200)); - } - return false; - } - - std::string test_dir_; - std::string cluster_id_; - int poll_interval_ms_; -}; - -// ============================================================ -// Test cases -// ============================================================ - -TEST_F(LocalFsHotStandbyIntegrationTest, TestPrimaryStandbySync) { - // 1. Primary writes 10 entries (last one via AppendAndPersist to flush) - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - std::vector test_keys; - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - - for (int i = 0; i < 10; ++i) { - std::string key = "test_key_" + std::to_string(i); - if (i < 9) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()) - << "AppendAndPersist failed for last entry"; - } - test_keys.push_back(key); - } - - uint64_t last_seq_id = primary->GetLastSequenceId(); - LOG(INFO) << "Primary wrote " << last_seq_id << " OpLog entries"; - - // 2. Start the standby - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, - standby.Start("", /*oplog_endpoints=*/"", cluster_id_)); - - // 3. Wait for sync - ASSERT_TRUE(WaitForSync(standby, last_seq_id)) - << "Standby failed to sync within timeout"; - - // 4. Verify metadata snapshot - std::vector> snapshot; - ASSERT_TRUE(standby.ExportMetadataSnapshot(snapshot)); - LOG(INFO) << "Standby metadata snapshot size: " << snapshot.size(); - - std::set snapshot_keys; - for (const auto& kv : snapshot) { - snapshot_keys.insert(kv.first); - } - for (const auto& key : test_keys) { - EXPECT_NE(snapshot_keys.end(), snapshot_keys.find(key)) - << "Key " << key << " not found in Standby snapshot"; - } - EXPECT_GE(snapshot.size(), test_keys.size()); - - // 5. Verify sequence IDs - EXPECT_GE(standby.GetLatestAppliedSequenceId(), last_seq_id); -} - -TEST_F(LocalFsHotStandbyIntegrationTest, TestStandbyPromotion) { - // 1. Write entries - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - std::vector test_keys; - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - - for (int i = 0; i < 5; ++i) { - std::string key = "promote_test_key_" + std::to_string(i); - if (i < 4) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()); - } - test_keys.push_back(key); - } - - uint64_t last_seq_id = primary->GetLastSequenceId(); - - // 2. Start standby and wait for sync - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, standby.Start("", "", cluster_id_)); - - ASSERT_TRUE(WaitForSync(standby, last_seq_id)); - - // 3. Verify ready for promotion - ASSERT_TRUE(standby.IsReadyForPromotion()); - - // 4. Promote - EXPECT_EQ(ErrorCode::OK, standby.Promote()); - - // 5. Verify sequence ID - EXPECT_GE(standby.GetLatestAppliedSequenceId(), last_seq_id); - - // 6. Verify metadata preserved - std::vector> snapshot; - ASSERT_TRUE(standby.ExportMetadataSnapshot(snapshot)); - - std::set snapshot_keys; - for (const auto& kv : snapshot) { - snapshot_keys.insert(kv.first); - } - for (const auto& key : test_keys) { - EXPECT_NE(snapshot_keys.end(), snapshot_keys.find(key)) - << "Key " << key << " should be in snapshot after promotion"; - } -} - -TEST_F(LocalFsHotStandbyIntegrationTest, TestFailoverScenario) { - // 1. Primary writes data - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - std::vector test_keys; - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - - for (int i = 0; i < 10; ++i) { - std::string key = "failover_key_" + std::to_string(i); - if (i < 9) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()); - } - test_keys.push_back(key); - } - - uint64_t last_seq_id = primary->GetLastSequenceId(); - - // 2. Start standby - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, standby.Start("", "", cluster_id_)); - - ASSERT_TRUE(WaitForSync(standby, last_seq_id)); - - // 3. Simulate primary failure (destroy the OpLogManager) - primary.reset(); - - // 4. Verify standby data integrity - std::vector> snapshot; - ASSERT_TRUE(standby.ExportMetadataSnapshot(snapshot)); - - std::set snapshot_keys; - for (const auto& kv : snapshot) { - snapshot_keys.insert(kv.first); - } - for (const auto& key : test_keys) { - EXPECT_NE(snapshot_keys.end(), snapshot_keys.find(key)) - << "Key " << key << " should be in Standby after Primary failure"; - } - - // 5. Promote - ASSERT_TRUE(standby.IsReadyForPromotion()); - EXPECT_EQ(ErrorCode::OK, standby.Promote()); - - // 6. Verify state - EXPECT_GE(standby.GetLatestAppliedSequenceId(), last_seq_id); -} - -TEST_F(LocalFsHotStandbyIntegrationTest, TestDataConsistency) { - // 1. Mixed PUT and REMOVE operations - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - std::map expected_keys; // key -> should_exist - - // PUT 5 keys - for (int i = 0; i < 5; ++i) { - std::string key = "put_key_" + std::to_string(i); - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - primary->Append(OpType::PUT_END, key, payload); - expected_keys[key] = true; - } - - // REMOVE first 2 - for (int i = 0; i < 2; ++i) { - std::string key = "put_key_" + std::to_string(i); - primary->Append(OpType::REMOVE, key, ""); - expected_keys[key] = false; - } - - // PUT 3 more (last one sync) - for (int i = 5; i < 8; ++i) { - std::string key = "put_key_" + std::to_string(i); - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":2048,"replicas":[]})"; - if (i < 7) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()); - } - expected_keys[key] = true; - } - - uint64_t last_seq_id = primary->GetLastSequenceId(); - LOG(INFO) << "Primary wrote " << last_seq_id << " OpLog entries"; - - // 2. Start standby - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, standby.Start("", "", cluster_id_)); - - // 3. Wait for sync - ASSERT_TRUE(WaitForSync(standby, last_seq_id)); - - // 4. Verify consistency - std::vector> snapshot; - ASSERT_TRUE(standby.ExportMetadataSnapshot(snapshot)); - - std::set actual_keys; - for (const auto& kv : snapshot) { - actual_keys.insert(kv.first); - } - - for (const auto& kv : expected_keys) { - if (kv.second) { - EXPECT_NE(actual_keys.end(), actual_keys.find(kv.first)) - << "Key " << kv.first << " should exist but not found"; - } else { - EXPECT_EQ(actual_keys.end(), actual_keys.find(kv.first)) - << "Key " << kv.first << " should be removed but still exists"; - } - } - - size_t expected_count = 0; - for (const auto& kv : expected_keys) { - if (kv.second) expected_count++; - } - EXPECT_EQ(expected_count, actual_keys.size()); -} - -TEST_F(LocalFsHotStandbyIntegrationTest, TestHighThroughputSync) { - // 1. Create primary first (initializes directory structure) - auto primary = CreatePrimaryOpLogManager(); - ASSERT_NE(primary, nullptr); - - // 2. Start standby (will poll for new entries) - HotStandbyService standby(MakeHotStandbyConfig()); - StandbyServiceGuard guard(&standby); - - ASSERT_EQ(ErrorCode::OK, standby.Start("", "", cluster_id_)); - - // Wait for WATCHING state - { - auto deadline = - std::chrono::steady_clock::now() + std::chrono::seconds(10); - while (std::chrono::steady_clock::now() < deadline) { - if (standby.GetSyncStatus().state == StandbyState::WATCHING) { - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - } - } - - // 3. High-throughput writes (last one sync) - - const int num_writes = 100; - std::string payload = - R"({"client_id_first":1,"client_id_second":2,"size":1024,"replicas":[]})"; - - auto write_start = std::chrono::steady_clock::now(); - for (int i = 0; i < num_writes; ++i) { - std::string key = "throughput_key_" + std::to_string(i); - if (i < num_writes - 1) { - primary->Append(OpType::PUT_END, key, payload); - } else { - auto result = - primary->AppendAndPersist(OpType::PUT_END, key, payload); - ASSERT_TRUE(result.has_value()); - } - } - auto write_end = std::chrono::steady_clock::now(); - auto write_duration = std::chrono::duration_cast( - write_end - write_start); - - uint64_t last_seq_id = primary->GetLastSequenceId(); - LOG(INFO) << "Wrote " << num_writes << " entries in " - << write_duration.count() << "ms, last_seq_id=" << last_seq_id; - - // 3. Monitor lag - auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(30); - uint64_t max_lag = 0; - - while (std::chrono::steady_clock::now() < deadline) { - auto status = standby.GetSyncStatus(); - if (status.lag_entries > max_lag) { - max_lag = status.lag_entries; - } - LOG(INFO) << "Standby lag: " << status.lag_entries - << " entries, applied_seq_id=" << status.applied_seq_id - << ", primary_seq_id=" << status.primary_seq_id; - if (status.applied_seq_id >= last_seq_id && status.lag_entries == 0) { - break; - } - std::this_thread::sleep_for(std::chrono::milliseconds(200)); - } - - // 4. Verify - auto final_status = standby.GetSyncStatus(); - EXPECT_GE(final_status.applied_seq_id, last_seq_id) - << "Standby should have applied all entries"; - EXPECT_EQ(0u, final_status.lag_entries) - << "Standby lag should be zero after sync"; - - LOG(INFO) << "Max lag observed: " << max_lag << " entries"; -} - -} // namespace testing -} // namespace mooncake - -int main(int argc, char** argv) { - gflags::ParseCommandLineFlags(&argc, &argv, true); - google::InitGoogleLogging(argv[0]); - google::SetVLOGLevel("*", 1); - FLAGS_logtostderr = 1; - ::testing::InitGoogleTest(&argc, argv); - return RUN_ALL_TESTS(); -} diff --git a/mooncake-store/tests/master_admin_server_test.cpp b/mooncake-store/tests/master_admin_server_test.cpp index 92ba05604c..8c572a4199 100644 --- a/mooncake-store/tests/master_admin_server_test.cpp +++ b/mooncake-store/tests/master_admin_server_test.cpp @@ -3,6 +3,9 @@ #include #include +#include +#include +#include #include #include #include @@ -15,6 +18,7 @@ #include "master_admin_service.h" #include "master_config.h" #include "rpc_service.h" +#include "tenant_quota_policy_store.h" #include "types.h" #include "utils.h" @@ -82,6 +86,18 @@ struct HttpSegmentsDetailResponse { }; YLT_REFL(HttpSegmentsDetailResponse, total_segments); +std::string WriteTenantQuotaPolicyForTest( + const std::map& tenant_quotas) { + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = tenant_quotas; + auto path = std::filesystem::temp_directory_path() / + ("mooncake_admin_tenant_quota_" + + UuidToString(generate_uuid()) + ".yaml"); + std::ofstream out(path); + out << FormatTenantQuotaPolicyYaml(snapshot); + return path.string(); +} + } // namespace // ========================================================================= @@ -117,6 +133,21 @@ class MasterAdminServerTest : public ::testing::Test { coro_http::req_content_type::json); return {result.status, std::string(result.resp_body)}; } + + HttpResponse HttpPutJson(int port, const std::string& path, + const std::string& body) { + coro_http::coro_http_client client; + auto result = async_simple::coro::syncAwait(client.async_put( + BaseUrl(port) + path, body, coro_http::req_content_type::json)); + return {result.status, std::string(result.resp_body)}; + } + + HttpResponse HttpDelete(int port, const std::string& path) { + coro_http::coro_http_client client; + auto result = async_simple::coro::syncAwait(client.async_delete( + BaseUrl(port) + path, "", coro_http::req_content_type::json)); + return {result.status, std::string(result.resp_body)}; + } }; // ========================================================================= @@ -440,6 +471,10 @@ TEST_F(MasterAdminServerTest, ServiceEndpointsReturn503WhenServiceUnavailable) { EXPECT_EQ(seg_status.http_status, 503); EXPECT_NE(seg_status.body.find(unavailable_msg), std::string::npos); + auto tenant_quotas = HttpGet(port, "/api/v1/tenant_quotas"); + EXPECT_EQ(tenant_quotas.http_status, 503); + EXPECT_NE(tenant_quotas.body.find(unavailable_msg), std::string::npos); + auto drain_create = HttpPostJson(port, "/api/v1/drain_jobs", "{}"); EXPECT_EQ(drain_create.http_status, 503); @@ -457,6 +492,145 @@ TEST_F(MasterAdminServerTest, ServiceEndpointsReturn503WhenServiceUnavailable) { admin.Stop(); } +TEST_F(MasterAdminServerTest, TenantQuotaAdminLifecycleEndpoints) { + const std::string policy_path = WriteTenantQuotaPolicyForTest({}); + WrappedMasterServiceConfig svc_config; + svc_config.default_kv_lease_ttl = 5000; + svc_config.enable_metric_reporting = false; + svc_config.enable_multi_tenants = true; + svc_config.tenant_quota_connector_type = "file"; + svc_config.tenant_quota_connector_uri = policy_path; + auto service = std::make_shared(svc_config); + + Segment segment; + segment.id = generate_uuid(); + segment.name = "quota_admin_segment"; + segment.base = 0x600000000; + segment.size = 2000; + UUID client_id = generate_uuid(); + ASSERT_TRUE(service->MountSegment(segment, client_id).has_value()); + + int port = getFreeTcpPort(); + MasterAdminServer admin(static_cast(port), false); + ASSERT_TRUE(admin.Start()); + admin.SetRuntimeState(ha::MasterRuntimeState::kServing); + admin.SetServiceDelegate(service); + admin.SetServiceAvailable(true); + + auto upsert = HttpPutJson(port, "/api/v1/tenant_quotas?tenant_id=tenant-a", + "{\"requested_quota_bytes\":800}"); + EXPECT_EQ(upsert.http_status, 200); + EXPECT_NE(upsert.body.find("\"tenant_id\":\"tenant-a\""), + std::string::npos); + EXPECT_NE(upsert.body.find("\"requested_quota_bytes\":800"), + std::string::npos); + EXPECT_NE(upsert.body.find("\"effective_quota_bytes\":800"), + std::string::npos); + EXPECT_NE(upsert.body.find("\"has_explicit_policy\":true"), + std::string::npos); + + auto list = HttpGet(port, "/api/v1/tenant_quotas"); + EXPECT_EQ(list.http_status, 200); + EXPECT_NE(list.body.find("\"tenant_id\":\"tenant-a\""), std::string::npos); + + auto one = HttpGet(port, "/api/v1/tenant_quotas?tenant_id=tenant-a"); + EXPECT_EQ(one.http_status, 200); + EXPECT_NE(one.body.find("\"committed_count\":0"), std::string::npos); + EXPECT_NE(one.body.find("\"over_quota\":false"), std::string::npos); + + ReplicateConfig cfg; + cfg.replica_num = 1; + auto put = + service->PutStart(client_id, "quota_admin_key", 100, cfg, "tenant-a"); + ASSERT_TRUE(put.has_value()) << toString(put.error()); + ASSERT_TRUE(service + ->PutEnd(client_id, + ObjectMeta{"quota_admin_key", std::nullopt}, + ReplicaType::MEMORY, "tenant-a") + .has_value()); + + auto delete_non_empty = + HttpDelete(port, "/api/v1/tenant_quotas?tenant_id=tenant-a"); + EXPECT_EQ(delete_non_empty.http_status, 409); + EXPECT_NE(delete_non_empty.body.find("TENANT_NOT_EMPTY"), + std::string::npos); + + ASSERT_TRUE(service->Remove("quota_admin_key", /*force=*/true, "tenant-a") + .has_value()); + + auto deleted = HttpDelete(port, "/api/v1/tenant_quotas?tenant_id=tenant-a"); + EXPECT_EQ(deleted.http_status, 200); + + auto missing = HttpGet(port, "/api/v1/tenant_quotas?tenant_id=tenant-a"); + EXPECT_EQ(missing.http_status, 404); + + admin.Stop(); + std::filesystem::remove(policy_path); +} + +TEST_F(MasterAdminServerTest, TenantQuotaAdminValidationErrors) { + const std::string policy_path = WriteTenantQuotaPolicyForTest({}); + WrappedMasterServiceConfig svc_config; + svc_config.default_kv_lease_ttl = 5000; + svc_config.enable_metric_reporting = false; + svc_config.enable_multi_tenants = true; + svc_config.tenant_quota_connector_type = "file"; + svc_config.tenant_quota_connector_uri = policy_path; + auto service = std::make_shared(svc_config); + + int port = getFreeTcpPort(); + MasterAdminServer admin(static_cast(port), false); + ASSERT_TRUE(admin.Start()); + admin.SetRuntimeState(ha::MasterRuntimeState::kServing); + admin.SetServiceDelegate(service); + admin.SetServiceAvailable(true); + + auto missing_tenant = HttpPutJson(port, "/api/v1/tenant_quotas", + "{\"requested_quota_bytes\":100}"); + EXPECT_EQ(missing_tenant.http_status, 400); + + auto empty_tenant = HttpPutJson(port, "/api/v1/tenant_quotas?tenant_id=", + "{\"requested_quota_bytes\":100}"); + EXPECT_EQ(empty_tenant.http_status, 400); + + auto zero_explicit = + HttpPutJson(port, "/api/v1/tenant_quotas?tenant_id=tenant-a", + "{\"requested_quota_bytes\":0}"); + EXPECT_EQ(zero_explicit.http_status, 400); + + auto reserved_tenant = + HttpGet(port, "/api/v1/tenant_quotas?tenant_id=_system"); + EXPECT_EQ(reserved_tenant.http_status, 400); + + auto missing_query = + HttpGet(port, "/api/v1/tenant_quotas?tenant_id=missing"); + EXPECT_EQ(missing_query.http_status, 404); + + admin.Stop(); + std::filesystem::remove(policy_path); +} + +TEST_F(MasterAdminServerTest, TenantQuotaAdminDisabledModeReturns409) { + WrappedMasterServiceConfig svc_config; + svc_config.default_kv_lease_ttl = 5000; + svc_config.enable_metric_reporting = false; + svc_config.enable_multi_tenants = false; + auto service = std::make_shared(svc_config); + + int port = getFreeTcpPort(); + MasterAdminServer admin(static_cast(port), false); + ASSERT_TRUE(admin.Start()); + admin.SetRuntimeState(ha::MasterRuntimeState::kServing); + admin.SetServiceDelegate(service); + admin.SetServiceAvailable(true); + + auto list = HttpGet(port, "/api/v1/tenant_quotas"); + EXPECT_EQ(list.http_status, 409); + EXPECT_NE(list.body.find("UNAVAILABLE_IN_CURRENT_MODE"), std::string::npos); + + admin.Stop(); +} + // ========================================================================= // MasterAdminServerWithServiceTest — reuses a single server+service across // all tests via SetUpTestSuite / TearDownTestSuite for fast execution. @@ -488,7 +662,9 @@ class MasterAdminServerWithServiceTest : public ::testing::Test { cfg.replica_num = 1; auto ps = service_->PutStart(client_id, kDefaultKey, 1024, cfg); if (ps.has_value()) { - (void)service_->PutEnd(client_id, kDefaultKey, ReplicaType::MEMORY); + (void)service_->PutEnd(client_id, + ObjectMeta{kDefaultKey, std::nullopt}, + ReplicaType::MEMORY); } port_ = getFreeTcpPort(); @@ -557,7 +733,8 @@ TEST_F(MasterAdminServerWithServiceTest, GetAllKeysExcludesRemovedKey) { cfg.replica_num = 1; auto ps = service_->PutStart(client_id, key, 1024, cfg); if (ps.has_value()) { - (void)service_->PutEnd(client_id, key, ReplicaType::MEMORY); + (void)service_->PutEnd(client_id, ObjectMeta{key, std::nullopt}, + ReplicaType::MEMORY); } (void)service_->Remove(key, "default"); @@ -829,7 +1006,9 @@ TEST_F(MasterAdminServerWithServiceTest, BatchQueryKeysMultipleKeys) { cfg.replica_num = 1; auto ps = service_->PutStart(client_id, "second_key", 512, cfg); if (ps.has_value()) { - (void)service_->PutEnd(client_id, "second_key", ReplicaType::MEMORY); + (void)service_->PutEnd(client_id, + ObjectMeta{"second_key", std::nullopt}, + ReplicaType::MEMORY); } auto resp = HttpGet("/batch_query_keys?keys=" + std::string(kDefaultKey) + @@ -973,11 +1152,13 @@ TEST_F(MasterAdminServerTest, MultipleSegmentsAndKeys) { cfg.replica_num = 1; auto ps1 = service->PutStart(client_id, "key_one", 1024, cfg); if (ps1.has_value()) { - (void)service->PutEnd(client_id, "key_one", ReplicaType::MEMORY); + (void)service->PutEnd(client_id, ObjectMeta{"key_one", std::nullopt}, + ReplicaType::MEMORY); } auto ps2 = service->PutStart(client_id, "key_two", 2048, cfg); if (ps2.has_value()) { - (void)service->PutEnd(client_id, "key_two", ReplicaType::MEMORY); + (void)service->PutEnd(client_id, ObjectMeta{"key_two", std::nullopt}, + ReplicaType::MEMORY); } int port = getFreeTcpPort(); @@ -1033,6 +1214,55 @@ TEST_F(MasterAdminServerTest, MultipleSegmentsAndKeys) { admin.Stop(); } +// /batch_query_keys returns replica metadata for disk-based keys via the +// optional disk_values/local_disk_values/nof_values fields, while the existing +// values field stays present (empty array) for backward compatibility. +TEST_F(MasterAdminServerTest, BatchQueryKeysReturnsLocalDiskReplicaInfo) { + WrappedMasterServiceConfig svc_config; + svc_config.default_kv_lease_ttl = 5000; + svc_config.enable_metric_reporting = false; + svc_config.enable_offload = true; // required to mount local-disk segments + auto service = std::make_shared(svc_config); + + UUID client_id = generate_uuid(); + Segment segment; + segment.id = generate_uuid(); + segment.name = "ld_segment"; + segment.base = 0x600000000; + segment.size = 8 * 1024 * 1024; + ASSERT_TRUE(service->MountSegment(segment, client_id).has_value()); + ASSERT_TRUE(service->MountLocalDiskSegment(client_id, true).has_value()); + + const std::string key = "ld_only_key"; + const std::string endpoint = "127.0.0.1:9999"; + StorageObjectMetadata sm; + sm.bucket_id = 0; + sm.offset = 0; + sm.key_size = static_cast(key.size()); + sm.data_size = 2048; + sm.transport_endpoint = endpoint; + OffloadTaskItem task{.tenant_id = "default", .key = key, .size = 2048}; + ASSERT_TRUE( + service->NotifyOffloadSuccess(client_id, {task}, {sm}).has_value()); + + int port = getFreeTcpPort(); + MasterAdminServer admin(static_cast(port), false); + ASSERT_TRUE(admin.Start()); + admin.SetRuntimeState(ha::MasterRuntimeState::kServing); + admin.SetServiceDelegate(service); + admin.SetServiceAvailable(true); + + auto resp = HttpGet(port, "/batch_query_keys?keys=" + key); + EXPECT_EQ(resp.http_status, 200); + EXPECT_NE(resp.body.find("\"success\":true"), std::string::npos); + EXPECT_NE(resp.body.find("local_disk_values"), std::string::npos); + EXPECT_NE(resp.body.find(endpoint), std::string::npos); + // Backward compat: a disk-only key still reports an (empty) values array. + EXPECT_NE(resp.body.find("\"values\":[]"), std::string::npos); + + admin.Stop(); +} + } // namespace test } // namespace mooncake diff --git a/mooncake-store/tests/master_metrics_test.cpp b/mooncake-store/tests/master_metrics_test.cpp index 812d0ff37f..6f6889d59e 100644 --- a/mooncake-store/tests/master_metrics_test.cpp +++ b/mooncake-store/tests/master_metrics_test.cpp @@ -13,6 +13,7 @@ #include "utils.h" #include "master_admin_service.h" #include "master_service.h" +#include "segment.h" #include "rpc_service.h" #include "types.h" #include "master_config.h" @@ -64,6 +65,7 @@ TEST_F(MasterMetricsTest, InitialStatusTest) { ASSERT_EQ(metrics.get_put_start_requests(), 0); ASSERT_EQ(metrics.get_put_start_failures(), 0); ASSERT_EQ(metrics.get_put_start_alloc_failures(), 0); + ASSERT_EQ(metrics.get_put_start_partial_allocations(), 0); ASSERT_EQ(metrics.get_put_end_requests(), 0); ASSERT_EQ(metrics.get_put_end_failures(), 0); ASSERT_EQ(metrics.get_put_revoke_requests(), 0); @@ -201,7 +203,8 @@ TEST_F(MasterMetricsTest, BasicRequestTest) { value_length); ASSERT_EQ(metrics.get_put_start_requests(), 2); ASSERT_EQ(metrics.get_put_start_failures(), 0); - auto put_end_result = service_.PutEnd(client_id, key, ReplicaType::MEMORY); + auto put_end_result = service_.PutEnd( + client_id, ObjectMeta{key, std::nullopt}, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); ASSERT_EQ(metrics.get_key_count(), 1); ASSERT_EQ(metrics.get_allocated_mem_size(), value_length); @@ -237,7 +240,8 @@ TEST_F(MasterMetricsTest, BasicRequestTest) { auto put_start_result3 = service_.PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result3.has_value()); - auto put_end_result2 = service_.PutEnd(client_id, key, ReplicaType::MEMORY); + auto put_end_result2 = service_.PutEnd( + client_id, ObjectMeta{key, std::nullopt}, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result2.has_value()); ASSERT_EQ(metrics.get_key_count(), 1); ASSERT_EQ(1, service_.RemoveAll()); @@ -251,7 +255,8 @@ TEST_F(MasterMetricsTest, BasicRequestTest) { auto put_start_result4 = service_.PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result4.has_value()); - auto put_end_result3 = service_.PutEnd(client_id, key, ReplicaType::MEMORY); + auto put_end_result3 = service_.PutEnd( + client_id, ObjectMeta{key, std::nullopt}, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result3.has_value()); auto unmount_result = service_.UnmountSegment(segment_id, client_id); ASSERT_TRUE(unmount_result.has_value()); @@ -270,6 +275,84 @@ TEST_F(MasterMetricsTest, BasicRequestTest) { ASSERT_DOUBLE_EQ(metrics.get_segment_mem_used_ratio("xxxxxx_segment"), 0.0); } +TEST_F(MasterMetricsTest, ServiceTeardownReleasesSegmentCapacity) { + auto& metrics = MasterMetricManager::instance(); + const int64_t capacity_before = metrics.get_total_mem_capacity(); + + Segment segment; + segment.id = generate_uuid(); + segment.name = "teardown_test_segment"; + segment.base = 0x300000000; + segment.size = 1024 * 1024 * 16; + UUID client_id = generate_uuid(); + + { + WrappedMasterServiceConfig service_config; + service_config.default_kv_lease_ttl = 100; + service_config.enable_metric_reporting = false; + WrappedMasterService service(service_config); + ASSERT_TRUE(service.MountSegment(segment, client_id).has_value()); + ASSERT_EQ(metrics.get_total_mem_capacity(), + capacity_before + static_cast(segment.size)); + ASSERT_EQ(metrics.get_segment_total_mem_capacity(segment.name), + static_cast(segment.size)); + } + + // Destroying the service while the segment is still mounted (as happens + // when a master loses leadership) must release the segment's capacity + // contribution; MasterMetricManager outlives the service instance. + ASSERT_EQ(metrics.get_total_mem_capacity(), capacity_before); + ASSERT_EQ(metrics.get_segment_total_mem_capacity(segment.name), 0); +} + +TEST_F(MasterMetricsTest, SnapshotReaderTeardownKeepsCapacityIntact) { + auto& metrics = MasterMetricManager::instance(); + + // Mount a segment through the accounted path so the gauge is non-zero. + SegmentManager source_manager(BufferAllocatorType::OFFSET); + Segment segment; + segment.id = generate_uuid(); + segment.name = "snapshot_reader_segment"; + segment.base = 0x300000000; + segment.size = 1024 * 1024 * 16; + UUID client_id = generate_uuid(); + ASSERT_EQ( + source_manager.getSegmentAccess().MountSegment(segment, client_id), + ErrorCode::OK); + const int64_t capacity_after_mount = metrics.get_total_mem_capacity(); + ASSERT_EQ(metrics.get_segment_total_mem_capacity(segment.name), + static_cast(segment.size)); + + auto snapshot = SegmentSerializer(&source_manager).Serialize(); + ASSERT_TRUE(snapshot.has_value()); + + { + // Deserialize into a temporary reader (as + // CatalogBackedSnapshotProvider does). The reader's records never + // contributed to the capacity metrics, so destroying it must leave + // the gauges untouched. + SegmentManager reader(BufferAllocatorType::OFFSET); + SegmentSerializer reader_serializer(&reader); + ASSERT_TRUE( + reader_serializer.Deserialize(snapshot.value()).has_value()); + } + ASSERT_EQ(metrics.get_total_mem_capacity(), capacity_after_mount); + ASSERT_EQ(metrics.get_segment_total_mem_capacity(segment.name), + static_cast(segment.size)); + + // Unmount to restore the gauges for the other tests. + { + auto access = source_manager.getSegmentAccess(); + size_t dec_capacity = 0; + ASSERT_EQ(access.PrepareUnmountSegment(segment.id, dec_capacity), + ErrorCode::OK); + ASSERT_EQ( + access.CommitUnmountSegment(segment.id, client_id, dec_capacity), + ErrorCode::OK); + } + ASSERT_EQ(metrics.get_segment_total_mem_capacity(segment.name), 0); +} + TEST_F(MasterMetricsTest, CalcCacheStatsTest) { const uint64_t default_kv_lease_ttl = 100; auto& metrics = MasterMetricManager::instance(); @@ -360,7 +443,8 @@ TEST_F(MasterMetricsTest, CalcCacheStatsTest) { auto put_start_result1 = service_.PutStart(client_id, key, value_length, config); ASSERT_TRUE(put_start_result1.has_value()); - auto put_end_result1 = service_.PutEnd(client_id, key, ReplicaType::MEMORY); + auto put_end_result1 = service_.PutEnd( + client_id, ObjectMeta{key, std::nullopt}, ReplicaType::MEMORY); ASSERT_TRUE(put_end_result1.has_value()); auto stats_dict = metrics.calculate_cache_stats(); @@ -541,7 +625,12 @@ TEST_F(MasterMetricsTest, BatchRequestTest) { ASSERT_EQ(metrics.get_batch_get_replica_list_failed_items(), 3); // Test BatchPutEnd request - auto batch_put_end_result = service_.BatchPutEnd(client_id, keys); + std::vector object_metas; + object_metas.reserve(keys.size()); + for (const auto& key : keys) { + object_metas.emplace_back(ObjectMeta{key, std::nullopt}); + } + auto batch_put_end_result = service_.BatchPutEnd(client_id, object_metas); ASSERT_EQ(batch_put_end_result.size(), 3); ASSERT_EQ(metrics.get_batch_put_end_requests(), 1); ASSERT_EQ(metrics.get_batch_put_end_partial_successes(), 0); @@ -609,9 +698,10 @@ static std::string PutKeyAndOffload(MasterService& svc, const UUID& client_id, const std::string& key) { ReplicateConfig cfg; cfg.replica_num = 1; - auto put_start = svc.PutStart(client_id, key, "default", value_size, cfg); + auto put_start = + svc.PutStart(client_id, key, TenantId::Default(), value_size, cfg); if (!put_start) return ""; - svc.PutEnd(client_id, key, "default", ReplicaType::MEMORY); + svc.PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY); StorageObjectMetadata meta; meta.data_size = static_cast(value_size); @@ -656,7 +746,7 @@ TEST_F(MasterMetricsTest, LocalDiskReplicaAllocatedSize) { EXPECT_EQ(metrics.get_allocated_file_size(), baseline + kValueSize); // After removing the key the LocalDiskReplica is destroyed; gauge resets. - ASSERT_TRUE(svc.Remove(key, "default").has_value()); + ASSERT_TRUE(svc.Remove(key, TenantId::Default()).has_value()); EXPECT_EQ(metrics.get_allocated_file_size(), baseline); } @@ -754,7 +844,7 @@ TEST_F(MasterMetricsTest, SummaryUsesWindowRatesAndCumulativeEviction) { std::string::npos); EXPECT_NE( window_summary.find("Eviction: Success/Attempts=1/2, AllocFail=0, " - "keys=3, size=4.00 KB"), + "PartialAlloc=0, keys=3, size=4.00 KB"), std::string::npos); EXPECT_NE(window_summary.find("Mem Eviction: Success/Attempts=1/2, " "keys=3, size=4.00 KB"), @@ -767,7 +857,7 @@ TEST_F(MasterMetricsTest, SummaryUsesWindowRatesAndCumulativeEviction) { metrics.get_summary_string_and_update_snapshot(); EXPECT_NE( reported_summary.find("Eviction: Success/Attempts=1/2, AllocFail=0, " - "keys=3, size=4.00 KB"), + "PartialAlloc=0, keys=3, size=4.00 KB"), std::string::npos); std::this_thread::sleep_for(std::chrono::milliseconds(20)); @@ -775,7 +865,8 @@ TEST_F(MasterMetricsTest, SummaryUsesWindowRatesAndCumulativeEviction) { metrics.get_summary_string_and_update_snapshot(); EXPECT_NE(idle_summary.find("PutStart=0.00/0.00"), std::string::npos); EXPECT_NE(idle_summary.find("Eviction: Success/Attempts=1/2, " - "AllocFail=0, keys=3, size=4.00 KB"), + "AllocFail=0, PartialAlloc=0, keys=3, " + "size=4.00 KB"), std::string::npos); EXPECT_NE(idle_summary.find("Mem Eviction: Success/Attempts=1/2, " "keys=3, size=4.00 KB"), @@ -785,6 +876,151 @@ TEST_F(MasterMetricsTest, SummaryUsesWindowRatesAndCumulativeEviction) { std::string::npos); } +// Verify that the SSD Offload path (LOCAL_DISK replicas) is tracked +// consistently by both cache-total accounting and hit counters. +TEST_F(MasterMetricsTest, SsdOffloadCacheHitAndTotalConsistent) { + auto& metrics = MasterMetricManager::instance(); + using CacheHitStat = MasterMetricManager::CacheHitStat; + + WrappedMasterServiceConfig service_config; + service_config.default_kv_lease_ttl = 100; + service_config.enable_offload = true; + service_config.enable_metric_reporting = true; + WrappedMasterService service_(service_config); + + constexpr size_t kBufferAddress = 0x300000000; + constexpr size_t kSegmentSize = 1024 * 1024 * 16; + std::string segment_name = "test_segment"; + UUID segment_id = generate_uuid(); + Segment segment; + segment.id = segment_id; + segment.name = segment_name; + segment.base = kBufferAddress; + segment.size = kSegmentSize; + UUID client_id = generate_uuid(); + + std::string key = "ssd_offload_key"; + uint64_t value_length = 2048; + ReplicateConfig config; + config.replica_num = 1; + + // Record baselines (singleton counters are not reset between tests). + const auto base_stats = metrics.calculate_cache_stats(); + const int64_t base_mem_hit_nums = + static_cast(base_stats.at(CacheHitStat::MEMORY_HITS)); + const int64_t base_ssd_hit_nums = + static_cast(base_stats.at(CacheHitStat::SSD_HITS)); + const int64_t base_mem_total = + static_cast(base_stats.at(CacheHitStat::MEMORY_TOTAL)); + const int64_t base_ssd_total = + static_cast(base_stats.at(CacheHitStat::SSD_TOTAL)); + const int64_t base_mem_hit_bytes = metrics.get_mem_cache_hit_bytes(); + const int64_t base_file_hit_bytes = metrics.get_file_cache_hit_bytes(); + const int64_t base_mem_cache_nums = metrics.get_mem_cache_nums(); + const int64_t base_file_cache_nums = metrics.get_file_cache_nums(); + + // Step 1: Mount segment and create a completed MEMORY replica. + auto mount_result = service_.MountSegment(segment, client_id); + ASSERT_TRUE(mount_result.has_value()); + auto put_start_result = + service_.PutStart(client_id, key, value_length, config); + ASSERT_TRUE(put_start_result.has_value()); + auto put_end_result = service_.PutEnd( + client_id, ObjectMeta{key, std::nullopt}, ReplicaType::MEMORY); + ASSERT_TRUE(put_end_result.has_value()); + + // After PutEnd: MEMORY_TOTAL should increment by 1. + auto stats = metrics.calculate_cache_stats(); + ASSERT_EQ(static_cast(stats.at(CacheHitStat::MEMORY_TOTAL)), + base_mem_total + 1); + ASSERT_EQ(static_cast(stats.at(CacheHitStat::SSD_TOTAL)), + base_ssd_total); + ASSERT_EQ(metrics.get_mem_cache_nums(), base_mem_cache_nums + 1); + ASSERT_EQ(metrics.get_file_cache_nums(), base_file_cache_nums); + + // Step 2: Mount local disk segment and add LOCAL_DISK replica via + // NotifyOffloadSuccess. + auto mount_disk_result = + service_.MountLocalDiskSegment(client_id, /*enable_offloading=*/true); + ASSERT_TRUE(mount_disk_result.has_value()); + + OffloadTaskItem task{.tenant_id = "default", + .key = key, + .size = static_cast(value_length)}; + StorageObjectMetadata obj_meta{ + .bucket_id = 0, + .offset = 0, + .key_size = 0, + .data_size = static_cast(value_length), + .transport_endpoint = "tcp://127.0.0.1:9999"}; + auto offload_result = + service_.NotifyOffloadSuccess(client_id, {task}, {obj_meta}); + ASSERT_TRUE(offload_result.has_value()); + + // After offload: SSD_TOTAL should increment by 1, MEMORY_TOTAL unchanged. + stats = metrics.calculate_cache_stats(); + ASSERT_EQ(static_cast(stats.at(CacheHitStat::MEMORY_TOTAL)), + base_mem_total + 1); + ASSERT_EQ(static_cast(stats.at(CacheHitStat::SSD_TOTAL)), + base_ssd_total + 1); + ASSERT_EQ(metrics.get_mem_cache_nums(), base_mem_cache_nums + 1); + ASSERT_EQ(metrics.get_file_cache_nums(), base_file_cache_nums + 1); + + // Step 3: Remove the object (removes both MEMORY and LOCAL_DISK replicas). + // Wait for lease expiry so Remove succeeds. + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + auto evict_result = service_.Remove(key, "default"); + ASSERT_TRUE(evict_result.has_value()); + + // After Remove: both totals should return to baseline. + stats = metrics.calculate_cache_stats(); + ASSERT_EQ(static_cast(stats.at(CacheHitStat::MEMORY_TOTAL)), + base_mem_total); + ASSERT_EQ(static_cast(stats.at(CacheHitStat::SSD_TOTAL)), + base_ssd_total); + ASSERT_EQ(metrics.get_mem_cache_nums(), base_mem_cache_nums); + ASSERT_EQ(metrics.get_file_cache_nums(), base_file_cache_nums); + + // Step 4: Test SSD hit path by creating an object with only a LOCAL_DISK + // replica. Use NotifyOffloadSuccess with a key that has no existing + // metadata — AddReplica will create the metadata and add LOCAL_DISK. + std::string ssd_only_key = "ssd_only_key"; + OffloadTaskItem task2{.tenant_id = "default", + .key = ssd_only_key, + .size = static_cast(value_length)}; + StorageObjectMetadata obj_meta2{ + .bucket_id = 0, + .offset = 0, + .key_size = 0, + .data_size = static_cast(value_length), + .transport_endpoint = "tcp://127.0.0.1:9998"}; + auto offload_result2 = + service_.NotifyOffloadSuccess(client_id, {task2}, {obj_meta2}); + ASSERT_TRUE(offload_result2.has_value()); + + // Verify SSD_TOTAL incremented for the new key. + stats = metrics.calculate_cache_stats(); + ASSERT_EQ(static_cast(stats.at(CacheHitStat::SSD_TOTAL)), + base_ssd_total + 1); + + // GetReplicaList should hit LOCAL_DISK, incrementing SSD hit counters. + auto get_result = service_.GetReplicaList(ssd_only_key, "default"); + ASSERT_TRUE(get_result.has_value()); + + stats = metrics.calculate_cache_stats(); + ASSERT_EQ(static_cast(stats.at(CacheHitStat::MEMORY_HITS)), + base_mem_hit_nums); + ASSERT_EQ(static_cast(stats.at(CacheHitStat::SSD_HITS)), + base_ssd_hit_nums + 1); + ASSERT_EQ(metrics.get_mem_cache_hit_bytes(), base_mem_hit_bytes); + ASSERT_EQ(metrics.get_file_cache_hit_bytes(), + base_file_hit_bytes + value_length); + + // Clean up. + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + service_.Remove(ssd_only_key, "default"); +} + } // namespace mooncake::test int main(int argc, char** argv) { diff --git a/mooncake-store/tests/master_scenario.cpp b/mooncake-store/tests/master_scenario.cpp new file mode 100644 index 0000000000..ccf4187bf3 --- /dev/null +++ b/mooncake-store/tests/master_scenario.cpp @@ -0,0 +1,252 @@ +#include "master_scenario.h" + +#include + +#include +#include + +#include "types.h" + +namespace mooncake::test { +namespace { + +constexpr uint64_t kFnvOffset = 14695981039346656037ULL; +constexpr uint64_t kFnvPrime = 1099511628211ULL; + +uint64_t StableHash(std::string_view kind, std::string_view name, + uint64_t seed) { + uint64_t hash = seed; + for (const char value : kind) { + hash = (hash ^ static_cast(value)) * kFnvPrime; + } + for (const char value : name) { + hash = (hash ^ static_cast(value)) * kFnvPrime; + } + return hash; +} + +UUID StableUuid(std::string_view kind, std::string_view name) { + return {StableHash(kind, name, kFnvOffset), + StableHash(kind, name, kFnvOffset ^ 0x9e3779b97f4a7c15ULL)}; +} + +} // namespace + +MemoryNodeSpec MemoryNode(std::string name) { + return {.name = std::move(name)}; +} + +PutStartAction<> PutStart(std::string key, uint64_t size) { + return PutStartAction<>(std::move(key), size); +} + +PutEndAction PutEnd(std::string key) { return {.key = std::move(key)}; } + +PutRevokeAction PutRevoke(std::string key) { return {.key = std::move(key)}; } + +RemoveAction Remove(std::string key) { return {.key = std::move(key)}; } + +ObjectSpec<> Object(std::string key) { return ObjectSpec<>(std::move(key)); } + +MasterScenario::MasterScenario(std::string name) : name_(std::move(name)) {} + +MasterScenario::~MasterScenario() = default; + +MasterScenario& MasterScenario::Given(MemoryNodeSpec node) { + if (declarations_frozen_) { + Fail("MemoryNode declarations must precede actions and assertions"); + return *this; + } + if (node.name.empty()) { + Fail("MemoryNode requires a name"); + return *this; + } + if (node.capacity == 0) { + Fail("MemoryNode " + node.name + " requires non-zero capacity"); + return *this; + } + const auto duplicate = std::find_if( + nodes_.begin(), nodes_.end(), + [&](const auto& existing) { return existing.name == node.name; }); + if (duplicate != nodes_.end()) { + Fail("duplicate MemoryNode " + node.name); + return *this; + } + nodes_.push_back(std::move(node)); + return *this; +} + +MasterScenario& MasterScenario::WhenPutStart(PutStartActionData action) { + if (!EnsureService()) { + return *this; + } + + ReplicateConfig config; + config.replica_num = 1; + const auto result = + service_->PutStart(ActorId(action.actor), action.key, + TenantId::Default(), action.size, config); + ValidateActionResult("PutStart(" + action.key + ")", action.expected_error, + result.has_value(), + result ? ErrorCode::OK : result.error()); + if (!result) { + return *this; + } + if (action.expected_replica_count.has_value() && + result->size() != *action.expected_replica_count) { + Fail("PutStart(" + action.key + ") returned " + + std::to_string(result->size()) + " replicas; expected " + + std::to_string(*action.expected_replica_count)); + } + if (action.expected_replica_status.has_value() && + std::any_of(result->begin(), result->end(), [&](const auto& replica) { + return replica.status != *action.expected_replica_status; + })) { + Fail("PutStart(" + action.key + ") replica status mismatch"); + } + return *this; +} + +MasterScenario& MasterScenario::When(PutEndAction action) { + if (!EnsureService()) { + return *this; + } + + const auto result = + service_->PutEnd(ActorId(action.actor), action.key, TenantId::Default(), + ReplicaType::MEMORY); + ValidateActionResult("PutEnd(" + action.key + ")", action.expected_error, + result.has_value(), + result ? ErrorCode::OK : result.error()); + return *this; +} + +MasterScenario& MasterScenario::When(PutRevokeAction action) { + if (!EnsureService()) { + return *this; + } + + const auto result = + service_->PutRevoke(ActorId(action.actor), action.key, + TenantId::Default(), ReplicaType::MEMORY); + ValidateActionResult("PutRevoke(" + action.key + ")", action.expected_error, + result.has_value(), + result ? ErrorCode::OK : result.error()); + return *this; +} + +MasterScenario& MasterScenario::When(RemoveAction action) { + if (!EnsureService()) { + return *this; + } + + const auto result = service_->Remove(action.key, TenantId::Default()); + ValidateActionResult("Remove(" + action.key + ")", action.expected_error, + result.has_value(), + result ? ErrorCode::OK : result.error()); + return *this; +} + +MasterScenario& MasterScenario::ThenObject(ObjectSpecData object, + ObjectExpectation expectation) { + if (!EnsureService()) { + return *this; + } + + const auto result = + service_->GetReplicaList(object.key, TenantId::Default()); + if (expectation == ObjectExpectation::NOT_READY) { + if (result || result.error() != ErrorCode::REPLICA_IS_NOT_READY) { + Fail("Object(" + object.key + ") was expected to be not ready"); + } + return *this; + } + if (!result) { + Fail("Object(" + object.key + + ") is not readable: " + toString(result.error())); + return *this; + } + if (result->replicas.empty()) { + Fail("Object(" + object.key + ") has no readable replicas"); + } + if (object.expected_replica_count.has_value() && + result->replicas.size() != *object.expected_replica_count) { + Fail("Object(" + object.key + ") has " + + std::to_string(result->replicas.size()) + " replicas; expected " + + std::to_string(*object.expected_replica_count)); + } + if (object.expected_complete_replica_count.has_value()) { + const size_t complete = + std::count_if(result->replicas.begin(), result->replicas.end(), + [](const auto& replica) { + return replica.status == ReplicaStatus::COMPLETE; + }); + if (complete != *object.expected_complete_replica_count) { + Fail("Object(" + object.key + ") has " + std::to_string(complete) + + " complete replicas; expected " + + std::to_string(*object.expected_complete_replica_count)); + } + } + return *this; +} + +bool MasterScenario::EnsureService() { + if (service_) { + return true; + } + declarations_frozen_ = true; + if (nodes_.empty()) { + Fail("scenario requires at least one MemoryNode"); + return false; + } + + service_ = std::make_unique(); + for (const auto& node : nodes_) { + Segment segment; + segment.id = StableUuid("segment", node.name); + segment.name = node.name; + segment.base = next_segment_base_; + segment.size = node.capacity; + segment.te_endpoint = node.name; + next_segment_base_ += node.capacity + 4096; + + const auto result = service_->MountSegment(segment, ActorId(node.name)); + if (!result) { + Fail("failed to mount MemoryNode " + node.name + ": " + + toString(result.error())); + service_.reset(); + return false; + } + } + return true; +} + +UUID MasterScenario::ActorId(std::string_view actor) { + const std::string name(actor); + const auto result = actor_ids_.emplace(name, StableUuid("actor", name)); + return result.first->second; +} + +void MasterScenario::ValidateActionResult( + std::string_view action, const std::optional& expected_error, + bool succeeded, ErrorCode error) { + if (!expected_error.has_value()) { + if (!succeeded) { + Fail(std::string(action) + " failed: " + toString(error)); + } + return; + } + if (succeeded) { + Fail(std::string(action) + " succeeded; expected " + + toString(*expected_error)); + } else if (error != *expected_error) { + Fail(std::string(action) + " failed with " + toString(error) + + "; expected " + toString(*expected_error)); + } +} + +void MasterScenario::Fail(std::string message) const { + ADD_FAILURE() << "MasterScenario[" << name_ << "]: " << message; +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/master_scenario.h b/mooncake-store/tests/master_scenario.h new file mode 100644 index 0000000000..850fa796a9 --- /dev/null +++ b/mooncake-store/tests/master_scenario.h @@ -0,0 +1,241 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "master_service.h" + +namespace mooncake::test { + +constexpr uint64_t operator""_KB(unsigned long long value) { + return value * 1024; +} + +struct MemoryNodeSpec { + std::string name; + uint64_t capacity{16 * 1024 * 1024}; + + MemoryNodeSpec& Capacity(uint64_t value) { + capacity = value; + return *this; + } +}; + +MemoryNodeSpec MemoryNode(std::string name); + +enum class PutStartExpectation { + UNSPECIFIED, + SUCCESS, + ERROR, +}; + +struct PutStartActionData { + std::string key; + uint64_t size; + std::string actor{"default"}; + std::optional expected_error{}; + std::optional expected_replica_count{}; + std::optional expected_replica_status{}; +}; + +template +struct PutStartAction : PutStartActionData { + PutStartAction(std::string key, uint64_t size) + requires(expectation == PutStartExpectation::UNSPECIFIED) + : PutStartActionData{.key = std::move(key), .size = size} {} + + PutStartAction& By(std::string value) { + actor = std::move(value); + return *this; + } + + auto ExpectError(ErrorCode value) const + requires(expectation == PutStartExpectation::UNSPECIFIED) + { + PutStartAction action(*this); + action.expected_error = value; + return action; + } + + auto ExpectReplicas(size_t value) const + requires(expectation != PutStartExpectation::ERROR) + { + PutStartAction action(*this); + action.expected_replica_count = value; + return action; + } + + auto ExpectStatus(ReplicaStatus value) const + requires(expectation != PutStartExpectation::ERROR) + { + PutStartAction action(*this); + action.expected_replica_status = value; + return action; + } + + private: + template + friend struct PutStartAction; + + template + PutStartAction(const PutStartAction& action) + : PutStartActionData(action) {} +}; + +PutStartAction<> PutStart(std::string key, uint64_t size); + +struct PutEndAction { + std::string key; + std::string actor{"default"}; + std::optional expected_error{}; + + PutEndAction& By(std::string value) { + actor = std::move(value); + return *this; + } + + PutEndAction& ExpectError(ErrorCode value) { + expected_error = value; + return *this; + } +}; + +PutEndAction PutEnd(std::string key); + +struct PutRevokeAction { + std::string key; + std::string actor{"default"}; + std::optional expected_error{}; + + PutRevokeAction& By(std::string value) { + actor = std::move(value); + return *this; + } + + PutRevokeAction& ExpectError(ErrorCode value) { + expected_error = value; + return *this; + } +}; + +PutRevokeAction PutRevoke(std::string key); + +struct RemoveAction { + std::string key; + std::optional expected_error{}; + + RemoveAction& ExpectError(ErrorCode value) { + expected_error = value; + return *this; + } +}; + +RemoveAction Remove(std::string key); + +enum class ObjectExpectation { + UNSPECIFIED, + READABLE, + NOT_READY, +}; + +struct ObjectSpecData { + std::string key; + std::optional expected_replica_count{}; + std::optional expected_complete_replica_count{}; +}; + +template +struct ObjectSpec : ObjectSpecData { + explicit ObjectSpec(std::string key) + requires(expectation == ObjectExpectation::UNSPECIFIED) + : ObjectSpecData{.key = std::move(key)} {} + + auto IsReadable() const + requires(expectation != ObjectExpectation::NOT_READY) + { + return ObjectSpec(*this); + } + + auto IsNotReady() const + requires(expectation == ObjectExpectation::UNSPECIFIED) + { + return ObjectSpec(*this); + } + + auto HasReplicas(size_t value) const + requires(expectation != ObjectExpectation::NOT_READY) + { + ObjectSpec object(*this); + object.expected_replica_count = value; + return object; + } + + auto HasCompleteReplicas(size_t value) const + requires(expectation != ObjectExpectation::NOT_READY) + { + ObjectSpec object(*this); + object.expected_complete_replica_count = value; + return object; + } + + private: + template + friend struct ObjectSpec; + + template + ObjectSpec(const ObjectSpec& object) : ObjectSpecData(object) {} +}; + +ObjectSpec<> Object(std::string key); + +class MasterScenario { + public: + explicit MasterScenario(std::string name); + ~MasterScenario(); + + MasterScenario(const MasterScenario&) = delete; + MasterScenario& operator=(const MasterScenario&) = delete; + + MasterScenario& Given(MemoryNodeSpec node); + template + MasterScenario& When(PutStartAction action) { + return WhenPutStart(std::move(action)); + } + + MasterScenario& When(PutEndAction action); + MasterScenario& When(PutRevokeAction action); + MasterScenario& When(RemoveAction action); + + template + requires(expectation != ObjectExpectation::UNSPECIFIED) + MasterScenario& Then(ObjectSpec object) { + return ThenObject(std::move(object), expectation); + } + + private: + MasterScenario& WhenPutStart(PutStartActionData action); + MasterScenario& ThenObject(ObjectSpecData object, + ObjectExpectation expectation); + bool EnsureService(); + UUID ActorId(std::string_view actor); + void ValidateActionResult(std::string_view action, + const std::optional& expected_error, + bool succeeded, ErrorCode error); + void Fail(std::string message) const; + + std::string name_; + bool declarations_frozen_{false}; + uintptr_t next_segment_base_{0x300000000}; + std::vector nodes_; + std::unique_ptr service_; + std::unordered_map actor_ids_; +}; + +} // namespace mooncake::test diff --git a/mooncake-store/tests/master_scenario_test.cpp b/mooncake-store/tests/master_scenario_test.cpp new file mode 100644 index 0000000000..bea24d6761 --- /dev/null +++ b/mooncake-store/tests/master_scenario_test.cpp @@ -0,0 +1,120 @@ +#include "master_scenario.h" + +#include +#include + +namespace mooncake::test { +namespace { + +template +concept SupportsExpectError = + requires(T value) { value.ExpectError(ErrorCode::INTERNAL_ERROR); }; + +template +concept SupportsExpectReplicas = requires(T value) { value.ExpectReplicas(1); }; + +template +concept SupportsExpectStatus = + requires(T value) { value.ExpectStatus(ReplicaStatus::COMPLETE); }; + +template +concept SupportsIsReadable = requires(T value) { value.IsReadable(); }; + +template +concept SupportsIsNotReady = requires(T value) { value.IsNotReady(); }; + +template +concept SupportsHasReplicas = requires(T value) { value.HasReplicas(1); }; + +template +concept SupportsThen = + requires(MasterScenario& scenario, T value) { scenario.Then(value); }; + +using ErrorExpectedPutStart = + decltype(PutStart("compile-time", 1_KB) + .ExpectError(ErrorCode::INTERNAL_ERROR)); +using SuccessExpectedPutStart = + decltype(PutStart("compile-time", 1_KB).ExpectReplicas(1)); +using UnspecifiedObject = decltype(Object("compile-time")); +using NotReadyObject = decltype(Object("compile-time").IsNotReady()); +using ReadableObject = decltype(Object("compile-time").HasReplicas(1)); + +static_assert(!SupportsExpectReplicas); +static_assert(!SupportsExpectStatus); +static_assert(!SupportsExpectError); +static_assert(!SupportsIsReadable); +static_assert(!SupportsHasReplicas); +static_assert(!SupportsIsNotReady); +static_assert(!SupportsThen); +static_assert(SupportsThen); +static_assert(SupportsThen); + +} // namespace + +TEST(MasterScenarioContractTest, ReportsUnexpectedActionError) { + EXPECT_NONFATAL_FAILURE(MasterScenario("unexpected action error") + .Given(MemoryNode("memory")) + .When(PutEnd("missing")), + "PutEnd(missing) failed: OBJECT_NOT_FOUND"); +} + +TEST(MasterScenarioContractTest, ReportsUnexpectedActionSuccess) { + EXPECT_NONFATAL_FAILURE( + MasterScenario("unexpected action success") + .Given(MemoryNode("memory")) + .When(PutStart("key", 1_KB) + .ExpectError(ErrorCode::OBJECT_ALREADY_EXISTS)), + "PutStart(key) succeeded; expected OBJECT_ALREADY_EXISTS"); +} + +TEST(MasterScenarioContractTest, ReportsWrongErrorCode) { + EXPECT_NONFATAL_FAILURE( + MasterScenario("wrong error code") + .Given(MemoryNode("memory")) + .When(PutEnd("missing").ExpectError(ErrorCode::ILLEGAL_CLIENT)), + "PutEnd(missing) failed with OBJECT_NOT_FOUND; expected " + "ILLEGAL_CLIENT"); +} + +TEST(MasterScenarioContractTest, ReportsPutStartReplicaCountMismatch) { + EXPECT_NONFATAL_FAILURE(MasterScenario("put start replica count mismatch") + .Given(MemoryNode("memory")) + .When(PutStart("key", 1_KB).ExpectReplicas(2)), + "PutStart(key) returned 1 replicas; expected 2"); +} + +TEST(MasterScenarioContractTest, ReportsPutStartReplicaStatusMismatch) { + EXPECT_NONFATAL_FAILURE( + MasterScenario("put start replica status mismatch") + .Given(MemoryNode("memory")) + .When(PutStart("key", 1_KB).ExpectStatus(ReplicaStatus::COMPLETE)), + "PutStart(key) replica status mismatch"); +} + +TEST(MasterScenarioContractTest, ReportsUnreadableObject) { + EXPECT_NONFATAL_FAILURE( + MasterScenario("unreadable object") + .Given(MemoryNode("memory")) + .Then(Object("missing").IsReadable()), + "Object(missing) is not readable: OBJECT_NOT_FOUND"); +} + +TEST(MasterScenarioContractTest, ReportsObjectReplicaCountMismatch) { + EXPECT_NONFATAL_FAILURE(MasterScenario("object replica count mismatch") + .Given(MemoryNode("memory")) + .When(PutStart("key", 1_KB)) + .When(PutEnd("key")) + .Then(Object("key").HasReplicas(2)), + "Object(key) has 1 replicas; expected 2"); +} + +TEST(MasterScenarioContractTest, ReportsCompleteReplicaCountMismatch) { + EXPECT_NONFATAL_FAILURE(MasterScenario("complete replica count mismatch") + .Given(MemoryNode("memory")) + .When(PutStart("key", 1_KB)) + .When(PutEnd("key")) + .Then(Object("key").HasCompleteReplicas(0)), + "Object(key) has 1 complete replicas; expected 0"); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/master_service_config_test.cpp b/mooncake-store/tests/master_service_config_test.cpp new file mode 100644 index 0000000000..4b69bcf5f0 --- /dev/null +++ b/mooncake-store/tests/master_service_config_test.cpp @@ -0,0 +1,49 @@ +#include "master_config.h" + +#include + +namespace mooncake::test { + +TEST(MasterServiceConfigTest, OplogBatchMaxEntriesDefaultsTo1024) { + MasterConfig master_config; + EXPECT_EQ(1024u, master_config.oplog_batch_max_entries); + + MasterServiceConfig service_config; + EXPECT_EQ(1024u, service_config.oplog_batch_max_entries); +} + +TEST(MasterServiceConfigTest, OplogIsDisabledByDefault) { + MasterConfig master_config; + EXPECT_FALSE(master_config.enable_oplog); + + MasterServiceConfig service_config; + EXPECT_FALSE(service_config.enable_oplog); +} + +TEST(MasterServiceConfigTest, OplogBuilderOverrideIsRespected) { + auto config = MasterServiceConfig::builder().set_enable_oplog(true).build(); + + EXPECT_TRUE(config.enable_oplog); +} + +TEST(MasterServiceConfigTest, OplogEnablementPropagatesToServingConfig) { + MasterConfig master_config{}; + master_config.enable_oplog = true; + MasterServiceSupervisorConfig supervisor_config(master_config); + + WrappedMasterServiceConfig wrapped_config(supervisor_config, 1); + MasterServiceConfig service_config(wrapped_config); + + EXPECT_TRUE(supervisor_config.enable_oplog); + EXPECT_TRUE(wrapped_config.enable_oplog); + EXPECT_TRUE(service_config.enable_oplog); +} + +TEST(MasterServiceConfigTest, OplogBatchMaxEntriesBuilderOverrideRespected) { + auto config = + MasterServiceConfig::builder().set_oplog_batch_max_entries(17).build(); + + EXPECT_EQ(17u, config.oplog_batch_max_entries); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/master_service_processing_key_double_erase_test.cpp b/mooncake-store/tests/master_service_processing_key_double_erase_test.cpp new file mode 100644 index 0000000000..186e260bfa --- /dev/null +++ b/mooncake-store/tests/master_service_processing_key_double_erase_test.cpp @@ -0,0 +1,202 @@ +// Reproduction/regression test for the MetadataAccessorRW double-erase +// use-after-free (prod incident 2026-08-03: mooncake_master segfaulted every +// ~5 minutes after a snapshot restore, always at the same instruction inside +// std::unordered_set::erase(const_iterator) — the bucket-chain +// walk dereferencing the chain-end nullptr). +// +// The bug — MasterService::MetadataAccessorRW constructor +// (mooncake-store/include/master_service.h): +// +// if (!it_->second.IsValid()) { +// const bool had_processing = +// processing_it_ != tenant_state_->processing_keys.end(); +// this->Erase(); // -> EraseMetadata(), which already does +// // processing_keys.erase(key) +// // (master_service.cpp), freeing the +// // node processing_it_ points to +// if (tenant_state_ != nullptr && had_processing) { +// this->EraseFromProcessing(); // -> processing_keys.erase( +// } // processing_it_) +// } // STALE ITERATOR! +// +// Erase() already removes the key from processing_keys (by key); the follow-up +// EraseFromProcessing() erases the SAME node again via the now-dangling +// iterator. libstdc++ re-reads the cached hash from the freed node, walks the +// bucket chain looking for it by address, runs off the end of the chain and +// dereferences nullptr -> SIGSEGV (fault address 0x0, exactly as observed in +// the prod kernel logs). +// +// Production trigger chain reproduced here (public API + one friend hook): +// 1. MountSegment (a "ghost" client mounts a segment) +// 2. PutStart without PutEnd (key stays in processing_keys with an +// incomplete replica on that segment) +// 3. PrepareUnmountSegment (ghost client expires; the replica's +// allocator weak_ptr expires, so +// has_invalid_mem_handle() == true) +// 4. PutEnd (or any op constructing (ctor cleanup: erases invalid replicas +// MetadataAccessorRW for the key) -> !IsValid() -> Erase() + +// EraseFromProcessing() double-erase) +// +// Two scenario details matter for a deterministic repro: +// * Step 3 must NOT use MasterService::UnmountSegment: it internally runs +// ClearInvalidHandles(), which would erase the crafted object through the +// SAFE path (EraseMetadata) before step 4 can hit the buggy accessor path. +// Production had the same window: the expiry thread unmounts the ghost +// segment and only THEN slowly sweeps 23M keys in ClearInvalidHandles — +// any RPC landing in that window hits the buggy accessor cleanup first. +// * A second live key ON THE SAME METADATA SHARD must keep the TenantState +// non-empty. Otherwise MaybeEraseEmptyTenant() erases the tenant and +// nulls tenant_state_, masking the bug (the buggy branch is guarded by +// tenant_state_ != nullptr). Production tenants hold millions of keys, +// so the buggy branch always executed. +// +// On the buggy code step 4 segfaults; the forked-child assertion below turns +// that into a clean test failure. After the fix the child exits 0 (PutEnd +// simply reports OBJECT_NOT_FOUND) and the test passes. + +#include "master_service.h" + +#include +#include + +#include +#include +#include +#include +#include + +namespace mooncake::test { + +class MasterServiceProcessingKeyDoubleEraseTest : public ::testing::Test { + protected: + void SetUp() override { + google::InitGoogleLogging("MasterServiceProcessingKeyDoubleEraseTest"); + FLAGS_logtostderr = true; + } + + void TearDown() override { google::ShutdownGoogleLogging(); } + + static constexpr size_t kSegmentBase = 0x300000000; + static constexpr size_t kSegmentSize = 16 * 1024 * 1024; + + // Exit codes used by the child to report how far it got. + static constexpr int kExitOk = 0; // reached past the trigger + static constexpr int kExitMountFailed = 2; // scenario setup broken + static constexpr int kExitPutStartFailed = 3; // scenario setup broken + static constexpr int kExitUnmountFailed = 4; // scenario setup broken + + // Friend access: find a key that routes to the SAME metadata shard as + // `key` (getMetadataShardIndex hashes tenant+key, so a naive second key + // lands in a different shard's TenantState and cannot keep THIS shard's + // tenant non-empty). + std::string FindKeyOnSameShard(MasterService& service, + const std::string& key) { + const size_t target = + service.getMetadataShardIndex(TenantId::Default(), key); + for (int i = 0; i < 100000; ++i) { + std::string candidate = key + "_keepalive_" + std::to_string(i); + if (service.getMetadataShardIndex(TenantId::Default(), candidate) == + target) { + return candidate; + } + } + return key + "_keepalive_fallback"; + } + + // Builds the incident state and fires the trigger. Only returns on + // fixed code; on buggy code it dies with SIGSEGV inside the + // MetadataAccessorRW constructor invoked by PutEnd. + void RunIncidentScenario() { + MasterService service(MasterServiceConfig::builder().build()); + + // 1. Ghost client mounts a segment. + Segment segment; + segment.id = generate_uuid(); + segment.name = "ghost_segment"; + segment.base = kSegmentBase; + segment.size = kSegmentSize; + segment.te_endpoint = segment.name; + const UUID client_id = generate_uuid(); + if (!service.MountSegment(segment, client_id).has_value()) { + ::_exit(kExitMountFailed); + } + + // 2. PutStart a key onto the segment and never complete it — the key + // stays in TenantState::processing_keys (client "died" mid-put). + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segment = segment.name; + const std::string key = "orphan_processing_key"; + if (!service.PutStart(client_id, key, TenantId::Default(), 1024, config) + .has_value()) { + ::_exit(kExitPutStartFailed); + } + + // 2b. A second, completed key on the SAME shard keeps the TenantState + // non-empty in step 4 (see file header for why this is required). + const std::string keepalive_key = FindKeyOnSameShard(service, key); + if (!service + .PutStart(client_id, keepalive_key, TenantId::Default(), 1024, + config) + .has_value() || + !service + .PutEnd(client_id, keepalive_key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()) { + ::_exit(kExitPutStartFailed); + } + + // 3. Ghost client expires: the segment allocator is destroyed, + // invalidating the replica's memory handle (weak_ptr expires). + // No ClearInvalidHandles sweep here (see file header). + size_t metrics_dec_capacity = 0; + { + auto segment_access = service.segment_manager_.getSegmentAccess(); + if (segment_access.PrepareUnmountSegment( + segment.id, metrics_dec_capacity) != ErrorCode::OK) { + ::_exit(kExitUnmountFailed); + } + } + + // 4. Trigger: PutEnd constructs MetadataAccessorRW(service, key). + // The ctor erases the invalid replica, finds !IsValid(), calls + // Erase() (frees the processing_keys node) and then + // EraseFromProcessing() with the stale processing_it_ iterator. + (void)service.PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); + ::_exit(kExitOk); // only reachable on fixed code + } +}; + +// Regression assertion: the incident scenario must complete without crashing. +// On the current buggy code the forked child dies with SIGSEGV (this is the +// reproduction); after the fix it exits kExitOk and the test passes. +TEST_F(MasterServiceProcessingKeyDoubleEraseTest, + AccessorCleanupAfterSegmentUnmountDoesNotCrash) { + ::fflush(nullptr); + pid_t pid = ::fork(); + ASSERT_NE(pid, -1) << "fork failed: " << strerror(errno); + if (pid == 0) { + RunIncidentScenario(); + ::_exit(kExitOk); // unreachable (RunIncidentScenario exits itself) + } + + int status = 0; + ASSERT_EQ(::waitpid(pid, &status, 0), pid); + + if (WIFSIGNALED(status)) { + FAIL() << "MetadataAccessorRW double-erase reproduced: child died " + "with signal " + << WTERMSIG(status) + << (WTERMSIG(status) == SIGSEGV ? " (SIGSEGV)" : "") + << ". Erase() -> EraseMetadata() already erases the key from " + "processing_keys; the subsequent EraseFromProcessing() " + "re-erases it via the stale processing_it_ iterator."; + } + ASSERT_TRUE(WIFEXITED(status)) << "child did not exit normally"; + EXPECT_EQ(WEXITSTATUS(status), kExitOk) + << "scenario setup failed (exit " << WEXITSTATUS(status) + << ": 2=MountSegment, 3=PutStart, 4=UnmountSegment)"; +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/master_service_scenario_test.cpp b/mooncake-store/tests/master_service_scenario_test.cpp new file mode 100644 index 0000000000..ba342bf151 --- /dev/null +++ b/mooncake-store/tests/master_service_scenario_test.cpp @@ -0,0 +1,29 @@ +#include "master_scenario.h" + +#include + +namespace mooncake::test { + +TEST(MasterServiceTest, PutStartEndFlow) { + MasterScenario("put start/end flow") + .Given(MemoryNode("memory")) + .When(PutStart("test_key", 1_KB) + .By("writer") + .ExpectReplicas(1) + .ExpectStatus(ReplicaStatus::PROCESSING)) + .Then(Object("test_key").IsNotReady()) + .When(Remove("test_key").ExpectError(ErrorCode::REPLICA_IS_NOT_READY)) + .When(PutEnd("test_key") + .By("other-writer") + .ExpectError(ErrorCode::ILLEGAL_CLIENT)) + .When(PutRevoke("test_key") + .By("other-writer") + .ExpectError(ErrorCode::ILLEGAL_CLIENT)) + .When(PutEnd("test_key").By("writer")) + .Then(Object("test_key") + .IsReadable() + .HasReplicas(1) + .HasCompleteReplicas(1)); +} + +} // namespace mooncake::test diff --git a/mooncake-store/tests/master_service_ssd_test.cpp b/mooncake-store/tests/master_service_ssd_test.cpp index 86b1b5c831..fcde4b4768 100644 --- a/mooncake-store/tests/master_service_ssd_test.cpp +++ b/mooncake-store/tests/master_service_ssd_test.cpp @@ -28,6 +28,71 @@ class MasterServiceSSDTest : public ::testing::Test { void TearDown() override { google::ShutdownGoogleLogging(); } }; +std::unique_ptr CreateSsdAwareOffloadService() { + MasterServiceConfig config; + config.enable_offload = true; + config.default_kv_lease_ttl = 0; + config.allocation_strategy_type = + AllocationStrategyType::SSD_FREE_RATIO_FIRST; + return std::make_unique(config); +} + +void MountMemoryAndLocalDisk(MasterService& service, const UUID& client_id, + const std::string& segment_name, + size_t base_addr) { + Segment segment; + segment.id = generate_uuid(); + segment.name = segment_name; + segment.base = base_addr; + segment.size = 64 * 1024 * 1024; + segment.te_endpoint = segment.name; + + ASSERT_TRUE(service.MountSegment(segment, client_id).has_value()); + ASSERT_TRUE(service.MountLocalDiskSegment(client_id, true).has_value()); + ASSERT_TRUE(service.ReportSsdCapacity(client_id, 1000).has_value()); +} + +void PutAndOffload(MasterService& service, const UUID& client_id, + const std::string& key, int64_t object_size, + const std::string& local_disk_endpoint) { + ReplicateConfig config; + config.replica_num = 1; + + ASSERT_TRUE( + service + .PutStart(client_id, key, TenantId::Default(), object_size, config) + .has_value()); + ASSERT_TRUE( + service.PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + + StorageObjectMetadata metadata; + metadata.data_size = object_size; + metadata.transport_endpoint = local_disk_endpoint; + OffloadTaskItem task{.tenant_id = TenantId::Default().value(), + .key = key, + .size = object_size}; + ASSERT_TRUE(service.NotifyOffloadSuccess(client_id, {task}, {metadata}) + .has_value()); +} + +void ExpectNextAllocationOnSegment(MasterService& service, + const UUID& client_id, + const std::string& key, + const std::string& expected_segment) { + ReplicateConfig config; + config.replica_num = 1; + auto result = + service.PutStart(client_id, key, TenantId::Default(), 64, config); + ASSERT_TRUE(result.has_value()); + ASSERT_EQ(result->size(), 1u); + ASSERT_TRUE((*result)[0].is_memory_replica()); + EXPECT_EQ((*result)[0] + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + expected_segment); +} + TEST_F(MasterServiceSSDTest, PutEndBothReplica) { auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd"); @@ -50,8 +115,8 @@ TEST_F(MasterServiceSSDTest, PutEndBothReplica) { ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto replicas = put_start_result.value(); ASSERT_EQ(2, replicas.size()); @@ -64,17 +129,20 @@ TEST_F(MasterServiceSSDTest, PutEndBothReplica) { EXPECT_TRUE(has_mem); EXPECT_TRUE(has_disk); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); // PutEnd for both memory and disk - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) - .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(2, get_result.value().replicas.size()); @@ -105,22 +173,26 @@ TEST_F(MasterServiceSSDTest, PutRevokeDiskReplica) { ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); ASSERT_TRUE(get_result.value().replicas[0].is_memory_replica()); EXPECT_TRUE( - service_->PutRevoke(client_id, key, "default", ReplicaType::DISK) + service_ + ->PutRevoke(client_id, key, TenantId::Default(), ReplicaType::DISK) .has_value()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); ASSERT_TRUE(get_result.value().replicas[0].is_memory_replica()); @@ -147,25 +219,30 @@ TEST_F(MasterServiceSSDTest, PutRevokeProcessingDiskKeepsSsdTotal) { ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value()); std::string key = "revoke_processing_disk_metric_key"; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1}) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), 1024, + {.replica_num = 1}) .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); auto stats = metrics.calculate_cache_stats(); EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total + 1); EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], base_ssd_total); EXPECT_TRUE( - service_->PutRevoke(client_id, key, "default", ReplicaType::DISK) + service_ + ->PutRevoke(client_id, key, TenantId::Default(), ReplicaType::DISK) .has_value()); stats = metrics.calculate_cache_stats(); EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total + 1); EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], base_ssd_total); - ASSERT_TRUE(service_->Remove(key, "default", /*force=*/true).has_value()); + ASSERT_TRUE( + service_->Remove(key, TenantId::Default(), /*force=*/true).has_value()); stats = metrics.calculate_cache_stats(); EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total); EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], base_ssd_total); @@ -193,20 +270,23 @@ TEST_F(MasterServiceSSDTest, PutRevokeMemoryReplica) { ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) - .has_value()); - EXPECT_TRUE( - service_->PutRevoke(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) + .has_value()); + EXPECT_TRUE(service_ + ->PutRevoke(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) - .has_value()); - get_result = service_->GetReplicaList(key, "default"); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); + get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); ASSERT_TRUE(get_result.value().replicas[0].is_disk_replica()); @@ -234,21 +314,24 @@ TEST_F(MasterServiceSSDTest, PutRevokeBothReplica) { ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) - .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) + .has_value()); EXPECT_TRUE( - service_->PutRevoke(client_id, key, "default", ReplicaType::DISK) + service_ + ->PutRevoke(client_id, key, TenantId::Default(), ReplicaType::DISK) .has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); - EXPECT_TRUE( - service_->PutRevoke(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - get_result = service_->GetReplicaList(key, "default"); + EXPECT_TRUE(service_ + ->PutRevoke(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); + get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -275,17 +358,21 @@ TEST_F(MasterServiceSSDTest, RemoveKey) { ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); - EXPECT_TRUE(service_->Remove(key, "default").has_value()); + EXPECT_TRUE(service_->Remove(key, TenantId::Default()).has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -317,13 +404,13 @@ TEST_F(MasterServiceSSDTest, EvictObject) { uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { auto put_end_mem_result = service_->PutEnd( - client_id, key, "default", ReplicaType::MEMORY); - auto put_end_disk_result = - service_->PutEnd(client_id, key, "default", ReplicaType::DISK); + client_id, key, TenantId::Default(), ReplicaType::MEMORY); + auto put_end_disk_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::DISK); ASSERT_TRUE(put_end_mem_result.has_value()); ASSERT_TRUE(put_end_disk_result.has_value()); success_puts++; @@ -338,7 +425,7 @@ TEST_F(MasterServiceSSDTest, EvictObject) { int success_gets = 0; for (int i = 0; i < 1024 * 16 + 50; ++i) { std::string key = "test_key" + std::to_string(i); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); if (get_result.has_value()) { success_gets++; } @@ -388,8 +475,8 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) { : ReplicaType::MEMORY; // Put key, should success. - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); auto replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -399,7 +486,7 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) { // Complete the reserved replica. auto put_end_result = - service_->PutEnd(client_id, key, "default", reserve_type); + service_->PutEnd(client_id, key, TenantId::Default(), reserve_type); EXPECT_TRUE(put_end_result.has_value()); // Wait for a while until the put-start expired. @@ -409,15 +496,16 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) { auto result = service_->Ping(client_id); EXPECT_TRUE(result.has_value()); // Protect the key from eviction. - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = + service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); std::this_thread::sleep_for(std::chrono::seconds(1)); } // Put key again, should fail because the object has had an completed // replica. - put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS); @@ -428,18 +516,20 @@ TEST_F(MasterServiceSSDTest, PutStartExpires) { auto result = service_->Ping(client_id); EXPECT_TRUE(result.has_value()); // Protect the key from eviction. - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = + service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); std::this_thread::sleep_for(std::chrono::seconds(1)); } - // Try PutEnd the discarded replica. + // PutEnd must reject a replica discarded after the write expired. put_end_result = - service_->PutEnd(client_id, key, "default", discard_type); - EXPECT_TRUE(put_end_result.has_value()); + service_->PutEnd(client_id, key, TenantId::Default(), discard_type); + ASSERT_FALSE(put_end_result.has_value()); + EXPECT_EQ(put_end_result.error(), ErrorCode::INVALID_WRITE); // Check that the key has only one replica. - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(get_result.value().replicas.size(), 1); if (reserve_type == ReplicaType::MEMORY) { @@ -478,28 +568,31 @@ TEST_F(MasterServiceSSDTest, EvictDiskReplica_RemovesDiskReplica) { ASSERT_TRUE(mount_result.has_value()); std::string key = "evict_disk_key"; - auto put_result = - service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1}); + auto put_result = service_->PutStart(client_id, key, TenantId::Default(), + 1024, {.replica_num = 1}); ASSERT_TRUE(put_result.has_value()); // Complete both replicas - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) - .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); // Verify we have 2 replicas (MEM + DISK) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(2, get_result.value().replicas.size()); // Evict disk replica - auto evict_result = service_->EvictDiskReplica(client_id, key, "default", - ReplicaType::DISK); + auto evict_result = service_->EvictDiskReplica( + client_id, key, TenantId::Default(), ReplicaType::DISK); ASSERT_TRUE(evict_result.has_value()); // Verify only memory replica remains - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); EXPECT_TRUE(get_result.value().replicas[0].is_memory_replica()); @@ -526,19 +619,24 @@ TEST_F(MasterServiceSSDTest, RemoveDecrementsCacheTotalMetrics) { ASSERT_TRUE(service_->MountSegment(segment, client_id).has_value()); std::string key = "remove_cache_total_metric_key"; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1}) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), 1024, + {.replica_num = 1}) .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); auto stats = metrics.calculate_cache_stats(); EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total + 1); EXPECT_EQ(stats[CacheHitStat::SSD_TOTAL], base_ssd_total + 1); - ASSERT_TRUE(service_->Remove(key, "default", /*force=*/true).has_value()); + ASSERT_TRUE( + service_->Remove(key, TenantId::Default(), /*force=*/true).has_value()); stats = metrics.calculate_cache_stats(); EXPECT_EQ(stats[CacheHitStat::MEMORY_TOTAL], base_memory_total); @@ -550,7 +648,7 @@ TEST_F(MasterServiceSSDTest, EvictDiskReplica_NonExistentKeyReturnsError) { UUID client_id = generate_uuid(); auto evict_result = service_->EvictDiskReplica( - client_id, "nonexistent_key", "default", ReplicaType::DISK); + client_id, "nonexistent_key", TenantId::Default(), ReplicaType::DISK); EXPECT_FALSE(evict_result.has_value()); EXPECT_EQ(evict_result.error(), ErrorCode::OBJECT_NOT_FOUND); } @@ -572,21 +670,275 @@ TEST_F(MasterServiceSSDTest, EvictDiskReplica_InvalidReplicaTypeReturnsError) { ASSERT_TRUE(mount_result.has_value()); std::string key = "evict_invalid_type_key"; - auto put_result = - service_->PutStart(client_id, key, "default", 1024, {.replica_num = 1}); + auto put_result = service_->PutStart(client_id, key, TenantId::Default(), + 1024, {.replica_num = 1}); ASSERT_TRUE(put_result.has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::DISK) - .has_value()); + EXPECT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE( + service_->PutEnd(client_id, key, TenantId::Default(), ReplicaType::DISK) + .has_value()); // Attempting to evict with MEMORY type should fail - auto evict_result = service_->EvictDiskReplica(client_id, key, "default", - ReplicaType::MEMORY); + auto evict_result = service_->EvictDiskReplica( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); EXPECT_FALSE(evict_result.has_value()); EXPECT_EQ(evict_result.error(), ErrorCode::INVALID_PARAMS); } +TEST_F(MasterServiceSSDTest, RemoveReleasesLocalDiskUsageTracking) { + auto service = CreateSsdAwareOffloadService(); + UUID client1 = generate_uuid(); + UUID client2 = generate_uuid(); + const std::string segment1 = "ssd_remove_segment_1"; + const std::string segment2 = "ssd_remove_segment_2"; + MountMemoryAndLocalDisk(*service, client1, segment1, 0x400000000); + MountMemoryAndLocalDisk(*service, client2, segment2, 0x500000000); + + PutAndOffload(*service, client1, "ssd_remove_released", 800, segment1); + PutAndOffload(*service, client2, "ssd_remove_baseline", 100, segment2); + + ASSERT_TRUE(service->Remove("ssd_remove_released", TenantId::Default()) + .has_value()); + + ExpectNextAllocationOnSegment(*service, client1, "ssd_remove_probe", + segment1); +} + +TEST_F(MasterServiceSSDTest, + BatchReplicaClearAllSegmentsReleasesLocalDiskUsageTracking) { + auto service = CreateSsdAwareOffloadService(); + UUID client1 = generate_uuid(); + UUID client2 = generate_uuid(); + const std::string segment1 = "ssd_clear_segment_1"; + const std::string segment2 = "ssd_clear_segment_2"; + MountMemoryAndLocalDisk(*service, client1, segment1, 0x600000000); + MountMemoryAndLocalDisk(*service, client2, segment2, 0x700000000); + + PutAndOffload(*service, client1, "ssd_clear_released", 800, segment1); + PutAndOffload(*service, client2, "ssd_clear_baseline", 100, segment2); + + auto clear_result = + service->BatchReplicaClear({"ssd_clear_released"}, client1, ""); + ASSERT_TRUE(clear_result.has_value()); + ASSERT_EQ(clear_result->size(), 1u); + EXPECT_EQ((*clear_result)[0], "ssd_clear_released"); + + ExpectNextAllocationOnSegment(*service, client1, "ssd_clear_probe", + segment1); +} + +// Test that after offloading more data to segment1, the next allocation prefers +// segment2 which has more SSD free space. +TEST_F(MasterServiceSSDTest, SsdFreeRatioFirstPrefersFresherSsdAfterOffload) { + auto service = CreateSsdAwareOffloadService(); + UUID client1 = generate_uuid(); + UUID client2 = generate_uuid(); + const std::string segment1 = "ssd_fresher_seg_1"; + const std::string segment2 = "ssd_fresher_seg_2"; + // Each segment reports total SSD capacity = 1000 bytes + MountMemoryAndLocalDisk(*service, client1, segment1, 0x800000000); + MountMemoryAndLocalDisk(*service, client2, segment2, 0x900000000); + + // Offload 800 bytes to segment1 → ssd_used[seg1]=800, free=20% + PutAndOffload(*service, client1, "ssd_fresher_heavy", 800, segment1); + // Offload 100 bytes to segment2 → ssd_used[seg2]=100, free=90% + PutAndOffload(*service, client2, "ssd_fresher_light", 100, segment2); + + // segment2 has higher SSD free ratio → allocation should prefer segment2 + ExpectNextAllocationOnSegment(*service, client2, "ssd_fresher_probe", + segment2); +} + +// Test that EvictDiskReplica decrements ssd_used_bytes so that the evicted +// segment becomes preferred again for the next allocation. +TEST_F(MasterServiceSSDTest, EvictDiskReplicaDecrementsLocalDiskUsageTracking) { + auto service = CreateSsdAwareOffloadService(); + UUID client1 = generate_uuid(); + UUID client2 = generate_uuid(); + const std::string segment1 = "ssd_evict_dec_seg_1"; + const std::string segment2 = "ssd_evict_dec_seg_2"; + MountMemoryAndLocalDisk(*service, client1, segment1, 0xa00000000); + MountMemoryAndLocalDisk(*service, client2, segment2, 0xb00000000); + + // Offload 800 bytes to segment1 → ssd_used[seg1]=800 (20% free) + PutAndOffload(*service, client1, "ssd_evict_dec_heavy", 800, segment1); + // Offload 100 bytes to segment2 → ssd_used[seg2]=100 (90% free) + PutAndOffload(*service, client2, "ssd_evict_dec_light", 100, segment2); + + // segment2 has more SSD free space → should be preferred + ExpectNextAllocationOnSegment(*service, client2, "ssd_evict_dec_probe1", + segment2); + + // Evict the LOCAL_DISK replica of the heavy object from segment1. + // NotifyOffloadSuccess creates a LOCAL_DISK replica (not DISK). + // This decrements ssd_used[seg1] by 800 → ssd_used[seg1]=0 (100% free) + auto evict_result = + service->EvictDiskReplica(client1, "ssd_evict_dec_heavy", + TenantId::Default(), ReplicaType::LOCAL_DISK); + ASSERT_TRUE(evict_result.has_value()); + + // After eviction: segment1 has 100% free, segment2 has 90% free + // → segment1 should now be preferred + ExpectNextAllocationOnSegment(*service, client1, "ssd_evict_dec_probe2", + segment1); +} + +// Evicting a LOCAL_DISK replica via EvictDiskReplica must decrement +// file_cache_nums_ even when the object still has a MEMORY replica (so +// accessor.Erase() does not run). Without SyncCacheTotalAccounting in the +// LOCAL_DISK eviction branch, the gauge would stay over-counted. +TEST_F(MasterServiceSSDTest, EvictDiskReplicaDecrementsFileCacheNums) { + auto& metrics = MasterMetricManager::instance(); + auto service = CreateSsdAwareOffloadService(); + UUID client_id = generate_uuid(); + const std::string segment = "ssd_evict_cache_total_segment"; + MountMemoryAndLocalDisk(*service, client_id, segment, 0xc00000000); + + const int64_t baseline = metrics.get_file_cache_nums(); + const int64_t baseline_mem = metrics.get_mem_cache_nums(); + + PutAndOffload(*service, client_id, "ssd_evict_cache_total_key", 128, + segment); + + // After offload: file_cache_nums_ increments by 1 (LOCAL_DISK replica), + // mem_cache_nums_ also increments by 1 (MEMORY replica from PutEnd). + EXPECT_EQ(metrics.get_file_cache_nums(), baseline + 1); + EXPECT_EQ(metrics.get_mem_cache_nums(), baseline_mem + 1); + + auto evict_result = + service->EvictDiskReplica(client_id, "ssd_evict_cache_total_key", + TenantId::Default(), ReplicaType::LOCAL_DISK); + ASSERT_TRUE(evict_result.has_value()); + + // After evicting LOCAL_DISK: file_cache_nums_ returns to baseline, + // mem_cache_nums_ unchanged (MEMORY replica still present). + EXPECT_EQ(metrics.get_file_cache_nums(), baseline); + EXPECT_EQ(metrics.get_mem_cache_nums(), baseline_mem + 1); +} + +// Real-path performance comparison: MasterService PutStart throughput for +// three configurations: +// (A) RANDOM, no offload — baseline, original behavior +// (B) RANDOM, with offload — isolates disk-replica creation overhead +// (C) SSD_FREE_RATIO_FIRST, with offload — adds SSD metrics lock + sorting +// +// Comparing A→B separates the cost of mounting LocalDisk segments. +// Comparing B→C isolates the pure SSD-ranking strategy overhead. +// +// Each round: PutStart → PutEnd(MEMORY) (timed) → Remove (not timed). +TEST_F(MasterServiceSSDTest, + SsdFreeRatioFirstVsRandomMasterServicePerformance) { + constexpr int kNumNodes = 32; + constexpr size_t kSegmentSize = 8 * 1024 * 1024; // 8 MiB each + constexpr size_t kSliceSize = 512; // 512 B – focus on strategy cost + constexpr int kWarmupRounds = 50; + constexpr int kBenchmarkRounds = 300; + + // Build a MasterService with kNumNodes segments. with_ssd=true also + // mounts LocalDisk and reports varied SSD capacity per node. + auto buildAndMount = + [&](AllocationStrategyType strategy, bool with_ssd, size_t base_start, + const std::string& tag) -> std::unique_ptr { + MasterServiceConfig config; + config.enable_offload = with_ssd; + config.default_kv_lease_ttl = 10000; + config.allocation_strategy_type = strategy; + auto svc = std::make_unique(config); + + for (int i = 0; i < kNumNodes; i++) { + UUID cid = generate_uuid(); + Segment seg; + seg.id = generate_uuid(); + seg.name = "ms_perf_" + std::to_string(i) + "_" + tag; + seg.base = base_start + static_cast(i) * kSegmentSize; + seg.size = kSegmentSize; + seg.te_endpoint = seg.name; + (void)svc->MountSegment(seg, cid); + if (with_ssd) { + (void)svc->MountLocalDiskSegment(cid, true); + // Vary total SSD capacity so nodes have distinct free ratios + (void)svc->ReportSsdCapacity( + cid, static_cast(1024 * 1024) * (i + 1)); + } + } + return svc; + }; + + // Measure kRounds of PutStart + PutEnd(MEMORY). Remove is called after + // timing to free allocator space without inflating the measurement. + auto runBenchmark = [&](MasterService& svc, const std::string& key_pfx, + int rounds) -> std::chrono::microseconds { + const UUID writer = generate_uuid(); + ReplicateConfig cfg; + cfg.replica_num = 1; + std::chrono::microseconds total{0}; + + for (int i = 0; i < rounds; i++) { + const std::string key = key_pfx + std::to_string(i); + auto t0 = std::chrono::steady_clock::now(); + (void)svc.PutStart(writer, key, TenantId::Default(), kSliceSize, + cfg); + (void)svc.PutEnd(writer, key, TenantId::Default(), + ReplicaType::MEMORY); + total += std::chrono::duration_cast( + std::chrono::steady_clock::now() - t0); + (void)svc.Remove(key, TenantId::Default(), /*force=*/true); + } + return total; + }; + + // (A) RANDOM, no offload – baseline + auto svc_a = buildAndMount(AllocationStrategyType::RANDOM, false, + 0xc00000000ULL, "A"); + (void)runBenchmark(*svc_a, "ms_A_wu_", kWarmupRounds); + auto elapsed_a = runBenchmark(*svc_a, "ms_A_bm_", kBenchmarkRounds); + + // (B) RANDOM, with offload – quantify disk-replica overhead alone + auto svc_b = buildAndMount(AllocationStrategyType::RANDOM, true, + 0xd00000000ULL, "B"); + (void)runBenchmark(*svc_b, "ms_B_wu_", kWarmupRounds); + auto elapsed_b = runBenchmark(*svc_b, "ms_B_bm_", kBenchmarkRounds); + + // (C) SSD_FREE_RATIO_FIRST, with offload – full new feature + auto svc_c = buildAndMount(AllocationStrategyType::SSD_FREE_RATIO_FIRST, + true, 0xe00000000ULL, "C"); + (void)runBenchmark(*svc_c, "ms_C_wu_", kWarmupRounds); + auto elapsed_c = runBenchmark(*svc_c, "ms_C_bm_", kBenchmarkRounds); + + auto us_per_op = [&](std::chrono::microseconds us) { + return static_cast(us.count()) / kBenchmarkRounds; + }; + double ratio_b_a = + static_cast(elapsed_b.count()) / elapsed_a.count(); + double ratio_c_b = + static_cast(elapsed_c.count()) / elapsed_b.count(); + double ratio_c_a = + static_cast(elapsed_c.count()) / elapsed_a.count(); + + std::cout + << "\n=== MasterService Real-Path Performance (PutStart+PutEnd) ===\n" + << "Nodes: " << kNumNodes << " | Slice: " << kSliceSize + << " B | Rounds: " << kBenchmarkRounds << "\n\n" + << " (A) RANDOM, offload=OFF (baseline): " << elapsed_a.count() + << " us | " << std::fixed << std::setprecision(3) + << us_per_op(elapsed_a) << " us/op\n" + << " (B) RANDOM, offload=ON (disk replica cost):" << elapsed_b.count() + << " us | " << us_per_op(elapsed_b) << " us/op [" + << std::setprecision(2) << ratio_b_a << "x vs A]\n" + << " (C) SSD_FREE_RATIO_FIRST, offload=ON: " << elapsed_c.count() + << " us | " << us_per_op(elapsed_c) << " us/op [" << ratio_c_b + << "x vs B]\n\n" + << " A→B disk-replica overhead: " << std::setprecision(1) + << (ratio_b_a - 1.0) * 100.0 << "%\n" + << " B→C SSD-ranking overhead: " << (ratio_c_b - 1.0) * 100.0 + << "%\n" + << " A→C total overhead vs origin:" << (ratio_c_a - 1.0) * 100.0 + << "%\n\n"; +} + } // namespace mooncake::test int main(int argc, char** argv) { diff --git a/mooncake-store/tests/master_service_tenant_quota_test.cpp b/mooncake-store/tests/master_service_tenant_quota_test.cpp index 9f6de362e5..5016af2b55 100644 --- a/mooncake-store/tests/master_service_tenant_quota_test.cpp +++ b/mooncake-store/tests/master_service_tenant_quota_test.cpp @@ -1,30 +1,150 @@ #include "master_service.h" -#include +#include +#include +#include +#include +#include +#include +#include #include +#include +#include #include +#include #include #include +#include +#include "allocation_strategy.h" +#include "tenant_quota_policy_store.h" #include "types.h" namespace mooncake::test { +class BlockingTenantQuotaPolicyStore final : public TenantQuotaPolicyStore { + public: + explicit BlockingTenantQuotaPolicyStore(TenantQuotaPolicySnapshot snapshot) + : snapshot_(std::move(snapshot)), + allow_save_(allow_save_promise_.get_future()) {} + + std::future SaveStarted() { + return save_started_promise_.get_future(); + } + + void AllowSave() { allow_save_promise_.set_value(); } + + tl::expected Load() override { + return snapshot_; + } + + tl::expected Save( + const TenantQuotaPolicySnapshot& snapshot) override { + snapshot_ = snapshot; + save_started_promise_.set_value(); + allow_save_.wait(); + return {}; + } + + private: + TenantQuotaPolicySnapshot snapshot_; + std::promise save_started_promise_; + std::promise allow_save_promise_; + std::future allow_save_; +}; + +#ifdef USE_NOF +class BlockingAllocationStrategy final : public AllocationStrategy { + public: + BlockingAllocationStrategy() + : allow_allocation_(allow_allocation_promise_.get_future()) {} + + std::future AllocationStarted() { + return allocation_started_promise_.get_future(); + } + + void AllowAllocation() { allow_allocation_promise_.set_value(); } + + tl::expected, ErrorCode> Allocate( + const AllocatorManager& allocator_manager, const size_t slice_length, + const size_t replica_num, + const std::vector& preferred_segments, + const std::set& excluded_segments, + const ReplicaType replica_type) override { + BlockOnce(); + return delegate_.Allocate(allocator_manager, slice_length, replica_num, + preferred_segments, excluded_segments, + replica_type); + } + + tl::expected, ErrorCode> Allocate( + const AllocatorManager& allocator_manager, const size_t slice_length, + const size_t replica_num, + const std::vector& preferred_segments, + const std::set& excluded_segments, + const ReplicaType replica_type, + const SsdMetricsProvider* ssd_provider) override { + (void)ssd_provider; + return Allocate(allocator_manager, slice_length, replica_num, + preferred_segments, excluded_segments, replica_type); + } + + tl::expected AllocateFrom( + const AllocatorManager& allocator_manager, const size_t slice_length, + const std::string& segment_name) override { + return delegate_.AllocateFrom(allocator_manager, slice_length, + segment_name); + } + + private: + void BlockOnce() { + bool expected = true; + if (block_next_allocation_.compare_exchange_strong(expected, false)) { + allocation_started_promise_.set_value(); + allow_allocation_.wait(); + } + } + + RandomAllocationStrategy delegate_; + std::atomic block_next_allocation_{true}; + std::promise allocation_started_promise_; + std::promise allow_allocation_promise_; + std::future allow_allocation_; +}; +#endif + class MasterServiceTenantQuotaTest : public ::testing::Test { protected: static constexpr size_t kSegmentBase = 0x500000000; - MasterServiceConfig MakeConfig(uint64_t default_quota, - uint64_t pool_capacity, - bool enable_quota = true, - std::string root_fs_dir = "") { - return MasterServiceConfig::builder() - .set_root_fs_dir(root_fs_dir) - .set_enable_tenant_quota(enable_quota) - .set_default_tenant_quota_bytes(default_quota) - .set_tenant_quota_pool_capacity_bytes(pool_capacity) - .build(); + std::string WritePolicyFile( + const std::map& tenant_quotas) { + TenantQuotaPolicySnapshot snapshot; + for (const auto& [tenant_id, quota] : tenant_quotas) { + snapshot.tenant_quotas.emplace(tenant_id.value(), quota); + } + auto path = + std::filesystem::temp_directory_path() / + ("mooncake_tenant_quota_test_" + std::to_string(::getpid()) + "_" + + std::to_string(next_policy_file_++) + ".yaml"); + std::ofstream out(path); + out << FormatTenantQuotaPolicyYaml(snapshot); + out.close(); + policy_files_.push_back(path.string()); + return path.string(); + } + + MasterServiceConfig MakeConfig( + const std::map& tenant_quotas, + bool enable_multi_tenants = true) { + auto builder = MasterServiceConfig::builder().set_enable_multi_tenants( + enable_multi_tenants); + if (enable_multi_tenants) { + builder.set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri(WritePolicyFile(tenant_quotas)); + } + return builder.build(); } UUID MountSegment(MasterService& service, size_t size = 4096, @@ -39,29 +159,36 @@ class MasterServiceTenantQuotaTest : public ::testing::Test { UUID client_id = generate_uuid(); auto result = service.MountSegment(segment, client_id); - EXPECT_TRUE(result.has_value()); + EXPECT_TRUE(result.has_value()) << toString(result.error()); return client_id; } +#ifdef USE_NOF + UUID MountNoFSegment(MasterService& service, size_t size = 4096, + std::string name = "quota_nof_segment") { + NoFSegment segment; + segment.id = generate_uuid(); + segment.name = std::move(name); + segment.base = kSegmentBase + next_segment_offset_; + segment.size = size; + segment.te_endpoint = segment.name; + next_segment_offset_ += size + 4096; + + UUID client_id = generate_uuid(); + auto result = service.MountNoFSegment(segment, client_id); + EXPECT_TRUE(result.has_value()) << toString(result.error()); + return client_id; + } +#endif + ReplicateConfig MemoryConfig() { ReplicateConfig config; config.replica_num = 1; return config; } - MasterServiceConfig MakeOffloadConfig(uint64_t default_quota, - uint64_t pool_capacity) { - auto config = MakeConfig(default_quota, pool_capacity); - config.enable_offload = true; - config.offload_on_evict = true; - config.promotion_on_hit = true; - config.promotion_admission_threshold = 1; - config.default_kv_lease_ttl = 0; - return config; - } - void PutComplete(MasterService& service, const UUID& client_id, - const std::string& key, const std::string& tenant_id, + const std::string& key, const TenantId& tenant_id, uint64_t size) { auto start = service.PutStart(client_id, key, tenant_id, size, MemoryConfig()); @@ -71,487 +198,709 @@ class MasterServiceTenantQuotaTest : public ::testing::Test { ASSERT_TRUE(end.has_value()) << toString(end.error()); } - void MountLocalDiskSegment(MasterService& service, const UUID& client_id) { - auto mount = service.MountLocalDiskSegment(client_id, true); - ASSERT_TRUE(mount.has_value()) << toString(mount.error()); + TenantQuotaSnapshot Snapshot(MasterService& service, + const TenantId& tenant_id) { + auto snapshot = service.GetTenantQuotaSnapshot(tenant_id); + EXPECT_TRUE(snapshot.has_value()); + return *snapshot; } - void InjectLocalDiskReplica(MasterService& service, const UUID& client_id, - const std::string& key, - const std::string& tenant_id, int64_t size, - const std::string& transport_endpoint) { - std::vector tasks{ - OffloadTaskItem{.tenant_id = tenant_id, .key = key, .size = size}}; - StorageObjectMetadata metadata; - metadata.bucket_id = 0; - metadata.offset = 0; - metadata.key_size = static_cast(key.size()); - metadata.data_size = size; - metadata.transport_endpoint = transport_endpoint; - std::vector metadatas{metadata}; - auto result = service.NotifyOffloadSuccess(client_id, tasks, metadatas); - ASSERT_TRUE(result.has_value()) << toString(result.error()); + void ReloadTenantQuotaPolicyFromStore(MasterService& service) { + service.LoadTenantQuotaPoliciesFromStoreOrThrow(); + service.RebuildTenantQuotaUsageFromMetadata(); } - TenantQuotaSnapshot Snapshot(MasterService& service, - const std::string& tenant_id) { - auto snapshot = service.GetTenantQuotaSnapshotForTesting(tenant_id); - EXPECT_TRUE(snapshot.has_value()); - return *snapshot; + void ReplaceTenantQuotaPolicyStore( + MasterService& service, std::unique_ptr store) { + service.tenant_quota_policy_store_ = std::move(store); } - tl::expected ReserveQuota(MasterService& service, - const std::string& tenant_id, - uint64_t bytes) { - return service.ReserveTenantQuota(tenant_id, bytes); + int64_t LocalDiskUsedBytes(MasterService& service, const UUID& client_id) { + auto access = service.segment_manager_.getLocalDiskSegmentAccess(); + auto& segments = access.getClientLocalDiskSegment(); + auto it = segments.find(client_id); + EXPECT_TRUE(it != segments.end()); + if (it == segments.end()) { + return -1; + } + return it->second->ssd_used_bytes.load(std::memory_order_relaxed); } - void CommitQuota(MasterService& service, const std::string& tenant_id, - uint64_t bytes) { - service.CommitTenantQuota(tenant_id, bytes); +#ifdef USE_NOF + void ReplaceAllocationStrategy( + MasterService& service, std::shared_ptr strategy) { + service.allocation_strategy_ = std::move(strategy); } +#endif - void AbortQuota(MasterService& service, const std::string& tenant_id, - uint64_t bytes) { - service.AbortTenantQuota(tenant_id, bytes); + tl::expected ReserveTenantQuotaForTest( + MasterService& service, const TenantId& tenant_id, uint64_t bytes) { + return service.ReserveTenantQuota(tenant_id, bytes); } - void ReleaseQuota(MasterService& service, const std::string& tenant_id, - uint64_t bytes) { - service.ReleaseTenantQuota(tenant_id, bytes); + std::unique_lock LockSnapshotForTest( + MasterService& service) { + return std::unique_lock(service.snapshot_mutex_); } - void ReleaseQuotaPartial(MasterService& service, - const std::string& tenant_id, uint64_t bytes) { - service.ReleaseTenantQuotaPartial(tenant_id, bytes); + std::unique_lock LockTenantQuotaRecomputeForTest( + MasterService& service) { + return std::unique_lock( + service.tenant_quota_recompute_mutex_); } - void RecomputeTenantQuotas(MasterService& service) { + ErrorCode MountSegmentWithoutQuotaRecomputeForTest(MasterService& service, + size_t size, + std::string name) { + Segment segment; + segment.id = generate_uuid(); + segment.name = std::move(name); + segment.base = kSegmentBase + next_segment_offset_; + segment.size = size; + segment.te_endpoint = segment.name; + next_segment_offset_ += size + 4096; + + auto segment_access = service.segment_manager_.getSegmentAccess(); + return segment_access.MountSegment(segment, generate_uuid()); + } + + void RecomputeTenantEffectiveQuotasForTest(MasterService& service) { service.RecomputeTenantEffectiveQuotas(); } - void BatchEvict(MasterService& service) { - service.BatchEvict(/*evict_ratio_target=*/1.0, - /*evict_ratio_lowerbound=*/1.0); - } - - void SetExplicitTenantPolicy(MasterService& service, - const std::string& tenant_id, - uint64_t requested_quota_bytes) { - const auto normalized_tenant = NormalizeTenantId(tenant_id); - const auto shard_idx = - service.getTenantQuotaShardIndex(normalized_tenant); - auto& shard = service.tenant_quota_shards_[shard_idx]; - std::lock_guard lock(shard.mutex); - auto& state = shard.tenants[normalized_tenant]; - state.requested_quota_bytes = requested_quota_bytes; - state.effective_quota_bytes = requested_quota_bytes; - state.has_explicit_policy = true; - state.over_quota = false; - } - - void SetEmptyInheritedTenantState(MasterService& service, - const std::string& tenant_id) { - const auto normalized_tenant = NormalizeTenantId(tenant_id); - const auto shard_idx = - service.getTenantQuotaShardIndex(normalized_tenant); - auto& shard = service.tenant_quota_shards_[shard_idx]; - std::lock_guard lock(shard.mutex); - auto& state = shard.tenants[normalized_tenant]; - state.requested_quota_bytes = service.default_tenant_quota_bytes_; - state.effective_quota_bytes = service.default_tenant_quota_bytes_; - state.has_explicit_policy = false; - state.over_quota = false; - } - - void ExpectSameAccounting(const TenantQuotaSnapshot& before, - const TenantQuotaSnapshot& after) { - EXPECT_EQ(after.requested_quota_bytes, before.requested_quota_bytes); - EXPECT_EQ(after.effective_quota_bytes, before.effective_quota_bytes); - EXPECT_EQ(after.used_bytes, before.used_bytes); - EXPECT_EQ(after.reserved_bytes, before.reserved_bytes); - EXPECT_EQ(after.committed_count, before.committed_count); - EXPECT_EQ(after.has_explicit_policy, before.has_explicit_policy); - EXPECT_EQ(after.over_quota, before.over_quota); + bool WaitForTenantQuotaPolicyMutexContention(MasterService& service) { + for (int i = 0; i < 500; ++i) { + if (!service.tenant_quota_policy_mutex_.try_lock()) { + return true; + } + service.tenant_quota_policy_mutex_.unlock(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + return false; + } + + void TearDown() override { + for (const auto& path : policy_files_) { + std::error_code ec; + std::filesystem::remove(path, ec); + } } size_t next_segment_offset_ = 0; + size_t next_policy_file_ = 0; + std::vector policy_files_; }; -TEST_F(MasterServiceTenantQuotaTest, DisabledPreservesLegacyPutRemove) { - MasterService service(MakeConfig(/*default_quota=*/128, - /*pool_capacity=*/128, - /*enable_quota=*/false)); - UUID client_id = MountSegment(service); - - PutComplete(service, client_id, "large", "tenant-a", 512); - EXPECT_TRUE( - service.Remove("large", "tenant-a", /*force=*/true).has_value()); +TEST_F(MasterServiceTenantQuotaTest, + SingleTenantModeCollapsesTenantsAndDisablesQuota) { + MasterService service(MakeConfig({}, /*enable_multi_tenants=*/false)); + UUID client_id = MountSegment(service, /*size=*/1024); + + PutComplete(service, client_id, "shared-key", TenantId("tenant-a"), 800); + + EXPECT_TRUE(service.ExistKey("shared-key", TenantId("tenant-b")).value()); + auto duplicate = service.PutStart(client_id, "shared-key", + TenantId("tenant-b"), 1, MemoryConfig()); + ASSERT_FALSE(duplicate.has_value()); + EXPECT_EQ(duplicate.error(), ErrorCode::OBJECT_ALREADY_EXISTS); + EXPECT_TRUE(service + .Remove("shared-key", TenantId("tenant-b"), + /*force=*/true) + .has_value()); EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); + service.GetTenantQuotaSnapshot(TenantId("tenant-a")).has_value()); } -TEST_F(MasterServiceTenantQuotaTest, SameKeyDifferentTenantsIndependent) { - MasterService service(MakeConfig(/*default_quota=*/1000, - /*pool_capacity=*/2000)); +TEST_F(MasterServiceTenantQuotaTest, + MultiTenantModeRejectsUnregisteredAndImplicitDefaultWrites) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 1000}})); UUID client_id = MountSegment(service); - PutComplete(service, client_id, "shared-key", "tenant-a", 800); - PutComplete(service, client_id, "shared-key", "tenant-b", 800); + auto missing = service.PutStart(client_id, "missing", TenantId("tenant-b"), + 10, MemoryConfig()); + ASSERT_FALSE(missing.has_value()); + EXPECT_EQ(missing.error(), ErrorCode::TENANT_NOT_REGISTERED); - EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 800); - EXPECT_EQ(Snapshot(service, "tenant-b").used_bytes, 800); -} + auto implicit_default = service.PutStart( + client_id, "default-key", TenantId::Default(), 10, MemoryConfig()); + ASSERT_FALSE(implicit_default.has_value()); + EXPECT_EQ(implicit_default.error(), ErrorCode::TENANT_NOT_REGISTERED); -TEST_F(MasterServiceTenantQuotaTest, SameTenantSharesQuotaAcrossKeys) { - MasterService service(MakeConfig(/*default_quota=*/1000, - /*pool_capacity=*/1000)); - UUID client_id = MountSegment(service); + const std::string control_tenant("tenant\0bad", 10); + EXPECT_FALSE(TenantId(control_tenant).IsValid()); - PutComplete(service, client_id, "key-a", "tenant-a", 600); - auto over_quota = - service.PutStart(client_id, "key-b", "tenant-a", 500, MemoryConfig()); + auto register_default = + service.UpsertTenantQuotaPolicy(TenantId::Default(), 100); + ASSERT_TRUE(register_default.has_value()) + << toString(register_default.error()); + PutComplete(service, client_id, "registered-default", TenantId::Default(), + 10); - ASSERT_FALSE(over_quota.has_value()); - EXPECT_EQ(over_quota.error(), ErrorCode::TENANT_QUOTA_EXCEEDED); - EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 600); - EXPECT_EQ(Snapshot(service, "tenant-a").reserved_bytes, 0); + PutComplete(service, client_id, "ok", TenantId("tenant-a"), 10); } -TEST_F(MasterServiceTenantQuotaTest, FirstTenantPutStartUsesPoolCapacity) { - MasterService service(MakeConfig(/*default_quota=*/1000, - /*pool_capacity=*/100)); +TEST_F(MasterServiceTenantQuotaTest, + MultiTenantModeRejectsUnregisteredOffloadSuccess) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 1000}})); UUID client_id = MountSegment(service); - auto over_quota = - service.PutStart(client_id, "large", "tenant-a", 800, MemoryConfig()); + StorageObjectMetadata metadata; + metadata.data_size = 128; + metadata.transport_endpoint = "disk-endpoint"; + std::vector tasks{ + OffloadTaskItem{.tenant_id = "tenant-b", .key = "ghost", .size = 128}}; - ASSERT_FALSE(over_quota.has_value()); - EXPECT_EQ(over_quota.error(), ErrorCode::TENANT_QUOTA_EXCEEDED); - EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); - - auto start = - service.PutStart(client_id, "small", "tenant-a", 80, MemoryConfig()); + auto result = service.NotifyOffloadSuccess(client_id, tasks, {metadata}); - ASSERT_TRUE(start.has_value()) << toString(start.error()); - auto snapshot = Snapshot(service, "tenant-a"); - EXPECT_EQ(snapshot.requested_quota_bytes, 1000); - EXPECT_EQ(snapshot.effective_quota_bytes, 100); - EXPECT_EQ(snapshot.used_bytes, 0); - EXPECT_EQ(snapshot.reserved_bytes, 80); - AbortQuota(service, "tenant-a", 80); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::TENANT_NOT_REGISTERED); + auto missing = service.ExistKey("ghost", TenantId("tenant-b")); + ASSERT_TRUE(missing.has_value()) << toString(missing.error()); + EXPECT_FALSE(missing.value()); } TEST_F(MasterServiceTenantQuotaTest, - ZeroDefaultQuotaInitializesNewTenantAsUnlimited) { - MasterService service(MakeConfig(/*default_quota=*/0, - /*pool_capacity=*/1)); - - auto reserve = ReserveQuota(service, "tenant-a", 4096); - - ASSERT_TRUE(reserve.has_value()) << toString(reserve.error()); - auto snapshot = Snapshot(service, "tenant-a"); - EXPECT_EQ(snapshot.requested_quota_bytes, 0); - EXPECT_EQ(snapshot.effective_quota_bytes, - std::numeric_limits::max()); - EXPECT_EQ(snapshot.reserved_bytes, 4096); -} + MultiTenantModeAllowsRegisteredOffloadSuccess) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 1000}})); + UUID client_id = MountSegment(service); -TEST_F(MasterServiceTenantQuotaTest, - OverQuotaReserveDoesNotChangeReservedBytes) { - MasterService service(MakeConfig(/*default_quota=*/100, - /*pool_capacity=*/100)); - ASSERT_TRUE(ReserveQuota(service, "tenant-a", 80).has_value()); - auto before = Snapshot(service, "tenant-a"); + StorageObjectMetadata metadata; + metadata.data_size = 128; + metadata.transport_endpoint = "disk-endpoint"; + std::vector tasks{ + OffloadTaskItem{.tenant_id = "tenant-a", .key = "cold", .size = 128}}; - auto reserve = ReserveQuota(service, "tenant-a", 30); + auto result = service.NotifyOffloadSuccess(client_id, tasks, {metadata}); - ASSERT_FALSE(reserve.has_value()); - EXPECT_EQ(reserve.error(), ErrorCode::TENANT_QUOTA_EXCEEDED); - ExpectSameAccounting(before, Snapshot(service, "tenant-a")); + ASSERT_TRUE(result.has_value()) << toString(result.error()); + auto exists = service.ExistKey("cold", TenantId("tenant-a")); + ASSERT_TRUE(exists.has_value()) << toString(exists.error()); + EXPECT_TRUE(exists.value()); + EXPECT_EQ(Snapshot(service, TenantId("tenant-a")).used_bytes, 0); } TEST_F(MasterServiceTenantQuotaTest, - FirstOverQuotaReserveDoesNotCreateTenantState) { - MasterService service(MakeConfig(/*default_quota=*/100, - /*pool_capacity=*/100)); - - auto reserve = ReserveQuota(service, "tenant-a", 101); - - ASSERT_FALSE(reserve.has_value()); - EXPECT_EQ(reserve.error(), ErrorCode::TENANT_QUOTA_EXCEEDED); - EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); -} + ConnectorPolicyReloadKeepsLocalDiskOnlyOrphanVisible) { + const std::string initial_policy = WritePolicyFile( + {{TenantId("tenant-a"), 1000}, {TenantId("tenant-b"), 1000}}); + auto config = MasterServiceConfig::builder() + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri(initial_policy) + .build(); + MasterService service(config); + UUID client_id = MountSegment(service); -TEST_F(MasterServiceTenantQuotaTest, RecomputePrunesEmptyInheritedTenantState) { - MasterService service(MakeConfig(/*default_quota=*/100, - /*pool_capacity=*/1000)); - SetEmptyInheritedTenantState(service, "tenant-a"); + StorageObjectMetadata metadata; + metadata.data_size = 128; + metadata.transport_endpoint = "disk-endpoint"; + std::vector tasks{ + OffloadTaskItem{.tenant_id = "tenant-b", .key = "cold", .size = 128}}; ASSERT_TRUE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); - SetExplicitTenantPolicy(service, "tenant-explicit", 200); + service.NotifyOffloadSuccess(client_id, tasks, {metadata}).has_value()); + + { + std::ofstream out(initial_policy); + TenantQuotaPolicySnapshot replacement; + replacement.tenant_quotas = {{"tenant-a", 1000}}; + out << FormatTenantQuotaPolicyYaml(replacement); + } + ReloadTenantQuotaPolicyFromStore(service); - RecomputeTenantQuotas(service); + auto orphan = Snapshot(service, TenantId("tenant-b")); + EXPECT_FALSE(orphan.has_explicit_policy); + EXPECT_EQ(orphan.used_bytes, 0); + EXPECT_EQ(orphan.committed_count, 0); + EXPECT_EQ(orphan.metadata_object_count, 1); + EXPECT_TRUE(orphan.over_quota); + EXPECT_TRUE(service.Remove("cold", TenantId("tenant-b"), /*force=*/true) + .has_value()); EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); - auto explicit_snapshot = Snapshot(service, "tenant-explicit"); - EXPECT_TRUE(explicit_snapshot.has_explicit_policy); - EXPECT_EQ(explicit_snapshot.requested_quota_bytes, 200); + service.GetTenantQuotaSnapshot(TenantId("tenant-b")).has_value()); } TEST_F(MasterServiceTenantQuotaTest, - AbortPrunesInactiveInheritedTenantAndRecomputesQuotas) { - MasterService service(MakeConfig(/*default_quota=*/1000, - /*pool_capacity=*/100)); + NotifyOffloadSuccessCompletesExistingOrphanObject) { + const std::string initial_policy = WritePolicyFile( + {{TenantId("tenant-a"), 1000}, {TenantId("tenant-b"), 1000}}); + auto config = MasterServiceConfig::builder() + .set_enable_multi_tenants(true) + .set_enable_offload(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri(initial_policy) + .build(); + MasterService service(config); + UUID client_id = MountSegment(service); + ASSERT_TRUE(service.MountLocalDiskSegment(client_id, true).has_value()); + PutComplete(service, client_id, "warming", TenantId("tenant-b"), 128); + + { + std::ofstream out(initial_policy); + TenantQuotaPolicySnapshot replacement; + replacement.tenant_quotas = {{"tenant-a", 1000}}; + out << FormatTenantQuotaPolicyYaml(replacement); + } + ReloadTenantQuotaPolicyFromStore(service); + EXPECT_FALSE(Snapshot(service, TenantId("tenant-b")).has_explicit_policy); - ASSERT_TRUE(ReserveQuota(service, "tenant-a", 50).has_value()); - ASSERT_TRUE(ReserveQuota(service, "tenant-b", 40).has_value()); - EXPECT_EQ(Snapshot(service, "tenant-b").effective_quota_bytes, 50); + StorageObjectMetadata metadata; + metadata.data_size = 128; + metadata.transport_endpoint = "disk-endpoint"; + std::vector tasks{OffloadTaskItem{ + .tenant_id = "tenant-b", .key = "warming", .size = 128}}; - AbortQuota(service, "tenant-a", 50); + auto result = service.NotifyOffloadSuccess(client_id, tasks, {metadata}); - EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); - EXPECT_EQ(Snapshot(service, "tenant-b").effective_quota_bytes, 100); - ASSERT_TRUE(ReserveQuota(service, "tenant-b", 60).has_value()); - EXPECT_EQ(Snapshot(service, "tenant-b").reserved_bytes, 100); - AbortQuota(service, "tenant-b", 100); + ASSERT_TRUE(result.has_value()) << toString(result.error()); + auto replicas = service.GetReplicaList("warming", TenantId("tenant-b")); + ASSERT_TRUE(replicas.has_value()) << toString(replicas.error()); + EXPECT_TRUE(std::any_of(replicas->replicas.begin(), + replicas->replicas.end(), + [](const Replica::Descriptor& replica) { + return replica.is_local_disk_replica(); + })); } TEST_F(MasterServiceTenantQuotaTest, - FullReleasePrunesInactiveInheritedTenantAndRecomputesQuotas) { - MasterService service(MakeConfig(/*default_quota=*/1000, - /*pool_capacity=*/100)); + NotifyOffloadSuccessRejectsOrphanObjectWithoutOffloadTask) { + const std::string initial_policy = WritePolicyFile( + {{TenantId("tenant-a"), 1000}, {TenantId("tenant-b"), 1000}}); + auto config = MasterServiceConfig::builder() + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri(initial_policy) + .build(); + MasterService service(config); UUID client_id = MountSegment(service); - PutComplete(service, client_id, "key-a", "tenant-a", 50); - PutComplete(service, client_id, "key-b", "tenant-b", 40); - EXPECT_EQ(Snapshot(service, "tenant-b").effective_quota_bytes, 50); - - ASSERT_TRUE( - service.Remove("key-a", "tenant-a", /*force=*/true).has_value()); + PutComplete(service, client_id, "warming", TenantId("tenant-b"), 128); - EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); - EXPECT_EQ(Snapshot(service, "tenant-b").effective_quota_bytes, 100); - auto reserve = service.PutStart(client_id, "key-b-extra", "tenant-b", 60, - MemoryConfig()); - ASSERT_TRUE(reserve.has_value()) << toString(reserve.error()); - EXPECT_EQ(Snapshot(service, "tenant-b").reserved_bytes, 60); - AbortQuota(service, "tenant-b", 60); -} + { + std::ofstream out(initial_policy); + TenantQuotaPolicySnapshot replacement; + replacement.tenant_quotas = {{"tenant-a", 1000}}; + out << FormatTenantQuotaPolicyYaml(replacement); + } + ReloadTenantQuotaPolicyFromStore(service); + EXPECT_FALSE(Snapshot(service, TenantId("tenant-b")).has_explicit_policy); -TEST_F(MasterServiceTenantQuotaTest, CommitMismatchDoesNotMutateAccounting) { - MasterService service(MakeConfig(/*default_quota=*/100, - /*pool_capacity=*/100)); - ASSERT_TRUE(ReserveQuota(service, "tenant-a", 40).has_value()); - auto before = Snapshot(service, "tenant-a"); + StorageObjectMetadata metadata; + metadata.data_size = 128; + metadata.transport_endpoint = "disk-endpoint"; + std::vector tasks{OffloadTaskItem{ + .tenant_id = "tenant-b", .key = "warming", .size = 128}}; - CommitQuota(service, "tenant-a", 50); + auto result = service.NotifyOffloadSuccess(client_id, tasks, {metadata}); - ExpectSameAccounting(before, Snapshot(service, "tenant-a")); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::TENANT_NOT_REGISTERED); } -TEST_F(MasterServiceTenantQuotaTest, AbortMismatchDoesNotMutateAccounting) { - MasterService service(MakeConfig(/*default_quota=*/100, - /*pool_capacity=*/100)); - ASSERT_TRUE(ReserveQuota(service, "tenant-a", 40).has_value()); - auto before = Snapshot(service, "tenant-a"); +TEST_F(MasterServiceTenantQuotaTest, + NotifyOffloadSuccessDoesNotCountAddReplicaUpdateAsNewDiskUsage) { + const std::string policy = WritePolicyFile({{TenantId("tenant-a"), 1000}}); + auto config = MasterServiceConfig::builder() + .set_enable_multi_tenants(true) + .set_enable_offload(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri(policy) + .build(); + MasterService service(config); + UUID client_a = MountSegment(service, 4096, "quota_segment_a"); + UUID client_b = MountSegment(service, 4096, "quota_segment_b"); + ASSERT_TRUE(service.MountLocalDiskSegment(client_a, true).has_value()); + ASSERT_TRUE(service.MountLocalDiskSegment(client_b, true).has_value()); + + StorageObjectMetadata first_metadata; + first_metadata.data_size = 128; + first_metadata.transport_endpoint = "disk-endpoint-a"; + std::vector tasks{ + OffloadTaskItem{.tenant_id = "tenant-a", .key = "cold", .size = 128}}; + ASSERT_TRUE(service.NotifyOffloadSuccess(client_a, tasks, {first_metadata}) + .has_value()); + EXPECT_EQ(LocalDiskUsedBytes(service, client_a), 128); + EXPECT_EQ(LocalDiskUsedBytes(service, client_b), 0); + + StorageObjectMetadata second_metadata; + second_metadata.data_size = 128; + second_metadata.transport_endpoint = "disk-endpoint-b"; + auto result = + service.NotifyOffloadSuccess(client_b, tasks, {second_metadata}); + + ASSERT_TRUE(result.has_value()) << toString(result.error()); + EXPECT_EQ(LocalDiskUsedBytes(service, client_a), 128); + EXPECT_EQ(LocalDiskUsedBytes(service, client_b), 0); +} - AbortQuota(service, "tenant-a", 50); +TEST_F(MasterServiceTenantQuotaTest, + RegisteredTenantQuotaAdmissionDoesNotCreateImplicitTenants) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 100}})); + UUID client_id = MountSegment(service); - ExpectSameAccounting(before, Snapshot(service, "tenant-a")); + auto hard_pinned = MemoryConfig(); + hard_pinned.with_hard_pin = true; + auto first = service.PutStart(client_id, "key-a", TenantId("tenant-a"), 80, + hard_pinned); + ASSERT_TRUE(first.has_value()) << toString(first.error()); + ASSERT_TRUE(service + .PutEnd(client_id, "key-a", TenantId("tenant-a"), + ReplicaType::MEMORY) + .has_value()); + + auto over = service.PutStart(client_id, "key-b", TenantId("tenant-a"), 30, + MemoryConfig()); + + ASSERT_FALSE(over.has_value()); + EXPECT_EQ(over.error(), ErrorCode::TENANT_QUOTA_EXCEEDED); + EXPECT_EQ(Snapshot(service, TenantId("tenant-a")).used_bytes, 80); + EXPECT_FALSE( + service.GetTenantQuotaSnapshot(TenantId("tenant-b")).has_value()); } -TEST_F(MasterServiceTenantQuotaTest, ReleaseMismatchDoesNotMutateAccounting) { - MasterService service(MakeConfig(/*default_quota=*/100, - /*pool_capacity=*/100)); - ASSERT_TRUE(ReserveQuota(service, "tenant-a", 40).has_value()); - CommitQuota(service, "tenant-a", 40); - auto before = Snapshot(service, "tenant-a"); +TEST_F(MasterServiceTenantQuotaTest, CopyStartRequiresQuotaForNewReplica) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 150}})); + UUID client_id = MountSegment(service, /*size=*/1024, "segment-a"); + MountSegment(service, /*size=*/1024, "segment-b"); - ReleaseQuota(service, "tenant-a", 50); - - ExpectSameAccounting(before, Snapshot(service, "tenant-a")); + ReplicateConfig config = MemoryConfig(); + config.preferred_segment = "segment-a"; + auto put_start = + service.PutStart(client_id, "key", TenantId("tenant-a"), 100, config); + ASSERT_TRUE(put_start.has_value()) << toString(put_start.error()); + ASSERT_TRUE( + service + .PutEnd(client_id, "key", TenantId("tenant-a"), ReplicaType::MEMORY) + .has_value()); + + auto copy = service.CopyStart(client_id, "key", TenantId("tenant-a"), + "segment-a", {"segment-b"}); + + ASSERT_FALSE(copy.has_value()); + EXPECT_EQ(copy.error(), ErrorCode::TENANT_QUOTA_EXCEEDED); + auto snapshot = Snapshot(service, TenantId("tenant-a")); + EXPECT_EQ(snapshot.used_bytes, 100); + EXPECT_EQ(snapshot.reserved_bytes, 0); + EXPECT_EQ(snapshot.committed_count, 1); } TEST_F(MasterServiceTenantQuotaTest, - ReleasePartialMismatchDoesNotMutateAccounting) { - MasterService service(MakeConfig(/*default_quota=*/100, - /*pool_capacity=*/100)); - ASSERT_TRUE(ReserveQuota(service, "tenant-a", 40).has_value()); - CommitQuota(service, "tenant-a", 40); - auto before = Snapshot(service, "tenant-a"); + CopyEndCommitsAdditionalReplicaWithoutExtraObjectCount) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 300}})); + UUID client_id = MountSegment(service, /*size=*/1024, "segment-a"); + MountSegment(service, /*size=*/1024, "segment-b"); + + ReplicateConfig config = MemoryConfig(); + config.preferred_segment = "segment-a"; + auto put_start = + service.PutStart(client_id, "key", TenantId("tenant-a"), 100, config); + ASSERT_TRUE(put_start.has_value()) << toString(put_start.error()); + ASSERT_TRUE( + service + .PutEnd(client_id, "key", TenantId("tenant-a"), ReplicaType::MEMORY) + .has_value()); - ReleaseQuotaPartial(service, "tenant-a", 50); + auto copy = service.CopyStart(client_id, "key", TenantId("tenant-a"), + "segment-a", {"segment-b"}); + ASSERT_TRUE(copy.has_value()) << toString(copy.error()); + auto in_flight = Snapshot(service, TenantId("tenant-a")); + EXPECT_EQ(in_flight.used_bytes, 100); + EXPECT_EQ(in_flight.reserved_bytes, 100); - ExpectSameAccounting(before, Snapshot(service, "tenant-a")); + ASSERT_TRUE( + service.CopyEnd(client_id, "key", TenantId("tenant-a")).has_value()); + auto completed = Snapshot(service, TenantId("tenant-a")); + EXPECT_EQ(completed.used_bytes, 200); + EXPECT_EQ(completed.reserved_bytes, 0); + EXPECT_EQ(completed.committed_count, 1); + EXPECT_EQ(completed.metadata_object_count, 1); } TEST_F(MasterServiceTenantQuotaTest, - PhysicalAllocationFailureKeepsAllocatorError) { - MasterService service(MakeConfig(/*default_quota=*/4096, - /*pool_capacity=*/4096)); - UUID client_id = MountSegment(service, /*size=*/512); - - auto result = service.PutStart(client_id, "too-large", "tenant-a", 1024, - MemoryConfig()); - - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error(), ErrorCode::NO_AVAILABLE_HANDLE); - EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); + MoveStartRequiresQuotaForTemporaryReplica) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 150}})); + UUID client_id = MountSegment(service, /*size=*/1024, "segment-a"); + MountSegment(service, /*size=*/1024, "segment-b"); + + ReplicateConfig config = MemoryConfig(); + config.preferred_segment = "segment-a"; + auto put_start = + service.PutStart(client_id, "key", TenantId("tenant-a"), 100, config); + ASSERT_TRUE(put_start.has_value()) << toString(put_start.error()); + ASSERT_TRUE( + service + .PutEnd(client_id, "key", TenantId("tenant-a"), ReplicaType::MEMORY) + .has_value()); + + auto move = service.MoveStart(client_id, "key", TenantId("tenant-a"), + "segment-a", "segment-b"); + + ASSERT_FALSE(move.has_value()); + EXPECT_EQ(move.error(), ErrorCode::TENANT_QUOTA_EXCEEDED); + auto snapshot = Snapshot(service, TenantId("tenant-a")); + EXPECT_EQ(snapshot.used_bytes, 100); + EXPECT_EQ(snapshot.reserved_bytes, 0); } -TEST_F(MasterServiceTenantQuotaTest, RemoveReleasesCommittedCharge) { - MasterService service(MakeConfig(/*default_quota=*/1000, - /*pool_capacity=*/1000)); +TEST_F(MasterServiceTenantQuotaTest, DeletePolicyRequiresTenantWithoutObjects) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 1000}})); UUID client_id = MountSegment(service); - - PutComplete(service, client_id, "key", "tenant-a", 400); - ASSERT_EQ(Snapshot(service, "tenant-a").used_bytes, 400); - - ASSERT_TRUE(service.Remove("key", "tenant-a", /*force=*/true).has_value()); - EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); + PutComplete(service, client_id, "key", TenantId("tenant-a"), 100); + + auto delete_non_empty = + service.DeleteTenantQuotaPolicy(TenantId("tenant-a")); + ASSERT_FALSE(delete_non_empty.has_value()); + EXPECT_EQ(delete_non_empty.error(), ErrorCode::TENANT_NOT_EMPTY); + + auto upsert = service.UpsertTenantQuotaPolicy(TenantId("tenant-b"), 100); + ASSERT_TRUE(upsert.has_value()) << toString(upsert.error()); + auto delete_empty = service.DeleteTenantQuotaPolicy(TenantId("tenant-b")); + ASSERT_TRUE(delete_empty.has_value()) << toString(delete_empty.error()); + EXPECT_FALSE(delete_empty.value().has_value()); } TEST_F(MasterServiceTenantQuotaTest, - BatchEvictReleasesEvictedMemoryReplicaCharge) { - MasterService service(MakeOffloadConfig(/*default_quota=*/1000, - /*pool_capacity=*/1000)); - UUID client_id = - MountSegment(service, /*size=*/4096, "quota_evict_segment"); - MountLocalDiskSegment(service, client_id); + DeletePolicyBlocksValidatedReservationsBeforeConnectorSave) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 1000}})); + MountSegment(service); + + TenantQuotaPolicySnapshot current_policy; + current_policy.tenant_quotas = {{"tenant-a", 1000}}; + auto blocking_store = + std::make_unique(current_policy); + auto* blocking_store_ptr = blocking_store.get(); + auto save_started = blocking_store_ptr->SaveStarted(); + ReplaceTenantQuotaPolicyStore(service, std::move(blocking_store)); + + using DeleteResult = + tl::expected, ErrorCode>; + std::optional delete_result; + std::thread delete_thread([&] { + delete_result.emplace( + service.DeleteTenantQuotaPolicy(TenantId("tenant-a"))); + }); + + if (save_started.wait_for(std::chrono::seconds(5)) != + std::future_status::ready) { + blocking_store_ptr->AllowSave(); + delete_thread.join(); + FAIL() << "timed out waiting for connector save"; + } - PutComplete(service, client_id, "key", "tenant-a", 400); - InjectLocalDiskReplica(service, client_id, "key", "tenant-a", 400, - "quota_evict_segment"); - ASSERT_EQ(Snapshot(service, "tenant-a").used_bytes, 400); + auto reserve = ReserveTenantQuotaForTest(service, TenantId("tenant-a"), 1); + EXPECT_FALSE(reserve.has_value()); + EXPECT_EQ(reserve.error(), ErrorCode::TENANT_NOT_REGISTERED); - BatchEvict(service); + auto zero_byte_reserve = + ReserveTenantQuotaForTest(service, TenantId("tenant-a"), 0); + EXPECT_FALSE(zero_byte_reserve.has_value()); + EXPECT_EQ(zero_byte_reserve.error(), ErrorCode::TENANT_NOT_REGISTERED); + blocking_store_ptr->AllowSave(); + delete_thread.join(); + + ASSERT_TRUE(delete_result.has_value()); + ASSERT_TRUE(delete_result->has_value()) << toString(delete_result->error()); + EXPECT_FALSE(delete_result->value().has_value()); EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); - auto reserve = service.PutStart(client_id, "after-evict", "tenant-a", 1000, - MemoryConfig()); - ASSERT_TRUE(reserve.has_value()) << toString(reserve.error()); - EXPECT_EQ(Snapshot(service, "tenant-a").reserved_bytes, 1000); - auto revoke = service.PutRevoke(client_id, "after-evict", "tenant-a", - ReplicaType::MEMORY); - ASSERT_TRUE(revoke.has_value()) << toString(revoke.error()); + service.GetTenantQuotaSnapshot(TenantId("tenant-a")).has_value()); } -TEST_F(MasterServiceTenantQuotaTest, PromotionSuccessCommitsTenantQuota) { - MasterService service(MakeOffloadConfig(/*default_quota=*/1000, - /*pool_capacity=*/1000)); - UUID client_id = - MountSegment(service, /*size=*/4096, "quota_promotion_segment"); - MountLocalDiskSegment(service, client_id); - InjectLocalDiskReplica(service, client_id, "cold", "tenant-a", 400, - "quota_promotion_segment"); +TEST_F(MasterServiceTenantQuotaTest, + DeletePolicyWaitsForInFlightAddReplicaBeforeEmptyCheck) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 1000}})); + UUID client_id = MountSegment(service); - auto replicas = service.GetReplicaList("cold", "tenant-a"); - ASSERT_TRUE(replicas.has_value()) << toString(replicas.error()); - auto alloc = service.PromotionAllocStart(client_id, "cold", "tenant-a", 400, - {"quota_promotion_segment"}); - ASSERT_TRUE(alloc.has_value()) << toString(alloc.error()); - EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 0); - EXPECT_EQ(Snapshot(service, "tenant-a").reserved_bytes, 400); - - auto notify = service.NotifyPromotionSuccess(client_id, "cold", "tenant-a"); - ASSERT_TRUE(notify.has_value()) << toString(notify.error()); - - EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 400); - EXPECT_EQ(Snapshot(service, "tenant-a").reserved_bytes, 0); - auto over_quota = service.PutStart(client_id, "too-much", "tenant-a", 700, - MemoryConfig()); - ASSERT_FALSE(over_quota.has_value()); - EXPECT_EQ(over_quota.error(), ErrorCode::TENANT_QUOTA_EXCEEDED); + TenantQuotaPolicySnapshot current_policy; + current_policy.tenant_quotas = {{"tenant-a", 1000}}; + auto blocking_store = + std::make_unique(current_policy); + auto* blocking_store_ptr = blocking_store.get(); + auto save_started = blocking_store_ptr->SaveStarted(); + ReplaceTenantQuotaPolicyStore(service, std::move(blocking_store)); + + auto snapshot_lock = LockSnapshotForTest(service); + std::optional> add_result; + std::thread add_thread([&] { + Replica replica(client_id, 128, "disk-endpoint", + ReplicaStatus::COMPLETE); + add_result.emplace(service.AddReplica(client_id, "cold", + TenantId("tenant-a"), replica)); + }); + + if (!WaitForTenantQuotaPolicyMutexContention(service)) { + snapshot_lock.unlock(); + add_thread.join(); + FAIL() << "timed out waiting for AddReplica to enter tenant policy " + "critical section"; + } + + using DeleteResult = + tl::expected, ErrorCode>; + std::optional delete_result; + std::thread delete_thread([&] { + delete_result.emplace( + service.DeleteTenantQuotaPolicy(TenantId("tenant-a"))); + }); + + const auto premature_save = + save_started.wait_for(std::chrono::milliseconds(200)); + if (premature_save == std::future_status::ready) { + blocking_store_ptr->AllowSave(); + } + snapshot_lock.unlock(); + add_thread.join(); + delete_thread.join(); + + ASSERT_EQ(premature_save, std::future_status::timeout) + << "tenant deletion reached connector save before in-flight " + "AddReplica completed"; + ASSERT_TRUE(add_result.has_value()); + ASSERT_TRUE(add_result->has_value()) << toString(add_result->error()); + ASSERT_TRUE(delete_result.has_value()); + ASSERT_FALSE(delete_result->has_value()); + EXPECT_EQ(delete_result->error(), ErrorCode::TENANT_NOT_EMPTY); + auto exists = service.ExistKey("cold", TenantId("tenant-a")); + ASSERT_TRUE(exists.has_value()) << toString(exists.error()); + EXPECT_TRUE(exists.value()); } -TEST_F(MasterServiceTenantQuotaTest, PromotionAllocStartRejectsOverQuota) { - MasterService service(MakeOffloadConfig(/*default_quota=*/300, - /*pool_capacity=*/300)); - UUID client_id = - MountSegment(service, /*size=*/4096, "quota_promotion_reject_segment"); - MountLocalDiskSegment(service, client_id); - InjectLocalDiskReplica(service, client_id, "cold", "tenant-a", 400, - "quota_promotion_reject_segment"); +#ifdef USE_NOF +TEST_F(MasterServiceTenantQuotaTest, + DeletePolicyWaitsForZeroChargePutStartMetadataCreate) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 1000}})); + UUID client_id = MountNoFSegment(service); + + auto blocking_strategy = std::make_shared(); + auto* blocking_strategy_ptr = blocking_strategy.get(); + auto allocation_started = blocking_strategy_ptr->AllocationStarted(); + ReplaceAllocationStrategy(service, std::move(blocking_strategy)); + + ReplicateConfig config; + config.replica_num = 0; + config.nof_replica_num = 1; + + std::optional, ErrorCode>> + put_result; + std::thread put_thread([&] { + put_result.emplace(service.PutStart(client_id, "nof-key", + TenantId("tenant-a"), 128, config)); + }); + + if (allocation_started.wait_for(std::chrono::seconds(5)) != + std::future_status::ready) { + blocking_strategy_ptr->AllowAllocation(); + put_thread.join(); + FAIL() << "timed out waiting for PutStart allocation"; + } - auto replicas = service.GetReplicaList("cold", "tenant-a"); - ASSERT_TRUE(replicas.has_value()) << toString(replicas.error()); - auto alloc = service.PromotionAllocStart( - client_id, "cold", "tenant-a", 400, {"quota_promotion_reject_segment"}); + using DeleteResult = + tl::expected, ErrorCode>; + std::optional delete_result; + std::thread delete_thread([&] { + delete_result.emplace( + service.DeleteTenantQuotaPolicy(TenantId("tenant-a"))); + }); - ASSERT_FALSE(alloc.has_value()); - EXPECT_EQ(alloc.error(), ErrorCode::TENANT_QUOTA_EXCEEDED); - EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); + ASSERT_TRUE(WaitForTenantQuotaPolicyMutexContention(service)) + << "DeleteTenantQuotaPolicy did not wait for zero-charge PutStart"; + + blocking_strategy_ptr->AllowAllocation(); + put_thread.join(); + delete_thread.join(); + + ASSERT_TRUE(put_result.has_value()); + ASSERT_TRUE(put_result->has_value()) << toString(put_result->error()); + ASSERT_TRUE(delete_result.has_value()); + ASSERT_FALSE(delete_result->has_value()); + EXPECT_EQ(delete_result->error(), ErrorCode::TENANT_NOT_EMPTY); + + auto snapshot = Snapshot(service, TenantId("tenant-a")); + EXPECT_EQ(snapshot.used_bytes, 0); + EXPECT_EQ(snapshot.reserved_bytes, 0); + EXPECT_EQ(snapshot.metadata_object_count, 1); } +#endif TEST_F(MasterServiceTenantQuotaTest, - DiskPutEndBeforeMemoryKeepsQuotaReservation) { + EffectiveQuotaUsesOnlyExplicitPolicyAndScalesProportionally) { MasterService service( - MakeConfig(/*default_quota=*/1000, /*pool_capacity=*/1000, - /*enable_quota=*/true, /*root_fs_dir=*/"/tmp/mooncake")); - UUID client_id = MountSegment(service); + MakeConfig({{TenantId("tenant-a"), 200}, {TenantId("tenant-b"), 400}})); + MountSegment(service, /*size=*/300); - auto start = - service.PutStart(client_id, "key", "tenant-a", 400, MemoryConfig()); - ASSERT_TRUE(start.has_value()) << toString(start.error()); - EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 0); - EXPECT_EQ(Snapshot(service, "tenant-a").reserved_bytes, 400); - - auto disk_end = - service.PutEnd(client_id, "key", "tenant-a", ReplicaType::DISK); - ASSERT_TRUE(disk_end.has_value()) << toString(disk_end.error()); - EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 0); - EXPECT_EQ(Snapshot(service, "tenant-a").reserved_bytes, 400); - - auto memory_end = - service.PutEnd(client_id, "key", "tenant-a", ReplicaType::MEMORY); - ASSERT_TRUE(memory_end.has_value()) << toString(memory_end.error()); - EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 400); - EXPECT_EQ(Snapshot(service, "tenant-a").reserved_bytes, 0); + EXPECT_EQ(Snapshot(service, TenantId("tenant-a")).effective_quota_bytes, + 100); + EXPECT_EQ(Snapshot(service, TenantId("tenant-b")).effective_quota_bytes, + 200); } -TEST_F(MasterServiceTenantQuotaTest, ChangedSizeUpsertSuccessSwapsCharge) { - MasterService service(MakeConfig(/*default_quota=*/1000, - /*pool_capacity=*/1000)); - UUID client_id = MountSegment(service); - - PutComplete(service, client_id, "key", "tenant-a", 400); - auto start = - service.UpsertStart(client_id, "key", "tenant-a", 600, MemoryConfig()); - ASSERT_TRUE(start.has_value()) << toString(start.error()); - EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 400); - EXPECT_EQ(Snapshot(service, "tenant-a").reserved_bytes, 600); - - auto end = - service.UpsertEnd(client_id, "key", "tenant-a", ReplicaType::MEMORY); - ASSERT_TRUE(end.has_value()) << toString(end.error()); - EXPECT_EQ(Snapshot(service, "tenant-a").used_bytes, 600); - EXPECT_EQ(Snapshot(service, "tenant-a").reserved_bytes, 0); +TEST_F(MasterServiceTenantQuotaTest, + CapacityIsSampledInsideQuotaRecomputeCoordination) { + MasterService service(MakeConfig({{TenantId("tenant-a"), 1000}})); + auto recompute_lock = LockTenantQuotaRecomputeForTest(service); + ASSERT_EQ(MountSegmentWithoutQuotaRecomputeForTest(service, /*size=*/100, + "capacity-a"), + ErrorCode::OK); + + std::promise recompute_started; + auto recompute_started_future = recompute_started.get_future(); + auto recompute = std::async(std::launch::async, [&] { + recompute_started.set_value(); + RecomputeTenantEffectiveQuotasForTest(service); + }); + recompute_started_future.wait(); + + EXPECT_EQ(recompute.wait_for(std::chrono::milliseconds(100)), + std::future_status::timeout); + const auto second_mount_result = MountSegmentWithoutQuotaRecomputeForTest( + service, /*size=*/50, "capacity-b"); + recompute_lock.unlock(); + + ASSERT_EQ(second_mount_result, ErrorCode::OK); + ASSERT_EQ(recompute.wait_for(std::chrono::seconds(5)), + std::future_status::ready); + recompute.get(); + EXPECT_EQ(Snapshot(service, TenantId("tenant-a")).effective_quota_bytes, + 150); } -TEST_F(MasterServiceTenantQuotaTest, ChangedSizeUpsertRevokeReleasesOldAndNew) { - MasterService service(MakeConfig(/*default_quota=*/1000, - /*pool_capacity=*/1000)); +TEST_F(MasterServiceTenantQuotaTest, + ConnectorPolicyReloadCreatesOrphanStateAndAllowsCleanup) { + const std::string initial_policy = WritePolicyFile( + {{TenantId("tenant-a"), 1000}, {TenantId("tenant-b"), 1000}}); + auto config = MasterServiceConfig::builder() + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri(initial_policy) + .build(); + MasterService service(config); UUID client_id = MountSegment(service); + PutComplete(service, client_id, "orphan-key", TenantId("tenant-b"), 100); - PutComplete(service, client_id, "key", "tenant-a", 400); - auto start = - service.UpsertStart(client_id, "key", "tenant-a", 600, MemoryConfig()); - ASSERT_TRUE(start.has_value()) << toString(start.error()); + { + std::ofstream out(initial_policy); + TenantQuotaPolicySnapshot replacement; + replacement.tenant_quotas = {{"tenant-a", 1000}}; + out << FormatTenantQuotaPolicyYaml(replacement); + } + ReloadTenantQuotaPolicyFromStore(service); - auto revoke = - service.UpsertRevoke(client_id, "key", "tenant-a", ReplicaType::MEMORY); - ASSERT_TRUE(revoke.has_value()) << toString(revoke.error()); - EXPECT_FALSE( - service.GetTenantQuotaSnapshotForTesting("tenant-a").has_value()); + auto orphan = Snapshot(service, TenantId("tenant-b")); + EXPECT_FALSE(orphan.has_explicit_policy); + EXPECT_EQ(orphan.requested_quota_bytes, 0); + EXPECT_EQ(orphan.effective_quota_bytes, 0); + EXPECT_TRUE(orphan.over_quota); + + EXPECT_TRUE( + service.GetReplicaList("orphan-key", TenantId("tenant-b")).has_value()); + auto write = service.PutStart(client_id, "new-key", TenantId("tenant-b"), 1, + MemoryConfig()); + ASSERT_FALSE(write.has_value()); + EXPECT_EQ(write.error(), ErrorCode::TENANT_NOT_REGISTERED); + + EXPECT_TRUE(service + .Remove("orphan-key", TenantId("tenant-b"), + /*force=*/true) + .has_value()); } } // namespace mooncake::test diff --git a/mooncake-store/tests/master_service_test.cpp b/mooncake-store/tests/master_service_test.cpp index ad1c55e068..f8c453cbdc 100644 --- a/mooncake-store/tests/master_service_test.cpp +++ b/mooncake-store/tests/master_service_test.cpp @@ -8,18 +8,69 @@ #include #include #include +#include +#include +#include #include +#include #include +#include #include +#include #include +#include #include #include -#include +#include + +#include "tenant_quota_policy_store.h" #include "types.h" +#include "utils.h" namespace mooncake::test { +std::vector MakeObjectMetas(const std::vector& keys) { + std::vector object_metas; + object_metas.reserve(keys.size()); + for (const auto& key : keys) { + object_metas.emplace_back(ObjectMeta{key, std::nullopt}); + } + return object_metas; +} + +class ScopedEnvVar { + public: + explicit ScopedEnvVar(const char* name) : name_(name) { + Capture(); + ::unsetenv(name_.c_str()); + } + + ScopedEnvVar(const char* name, const char* value) : name_(name) { + Capture(); + ::setenv(name_.c_str(), value, 1); + } + + ~ScopedEnvVar() { + if (previous_value_.has_value()) { + ::setenv(name_.c_str(), previous_value_->c_str(), 1); + } else { + ::unsetenv(name_.c_str()); + } + } + + private: + void Capture() { + const char* value = ::getenv(name_.c_str()); + if (value != nullptr) { + previous_value_ = value; + } + } + + std::string name_; + std::optional previous_value_; +}; + class MasterServiceTest : public ::testing::Test { protected: void SetUp() override { @@ -34,16 +85,64 @@ class MasterServiceTest : public ::testing::Test { static constexpr size_t kDefaultSegmentBase = 0x300000000; static constexpr size_t kDefaultSegmentSize = 1024 * 1024 * 16; + static constexpr uint64_t kStrictTenantQuotaBytes = 4 * 1024 * 1024; + + std::string WriteTenantPolicyFile( + const std::map& tenant_quotas) { + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = tenant_quotas; + auto path = + std::filesystem::temp_directory_path() / + ("mooncake_master_service_test_" + std::to_string(::getpid()) + + "_" + std::to_string(next_policy_file_++) + ".yaml"); + std::ofstream out(path); + out << FormatTenantQuotaPolicyYaml(snapshot); + out.close(); + policy_files_.push_back(path.string()); + return path.string(); + } + + MasterServiceConfig MakeStrictTenantConfig( + const std::vector& tenants) { + std::map tenant_quotas; + for (const auto& tenant : tenants) { + tenant_quotas.emplace(tenant, kStrictTenantQuotaBytes); + } + return MasterServiceConfig::builder() + .set_enable_multi_tenants(true) + .set_tenant_quota_connector_type("file") + .set_tenant_quota_connector_uri( + WriteTenantPolicyFile(tenant_quotas)) + .build(); + } + + WrappedMasterServiceConfig MakeStrictWrappedConfig( + const std::vector& tenants) { + WrappedMasterServiceConfig config; + config.default_kv_lease_ttl = 100; + config.enable_metric_reporting = false; + config.enable_multi_tenants = true; + config.tenant_quota_connector_type = "file"; + std::map tenant_quotas; + for (const auto& tenant : tenants) { + tenant_quotas.emplace(tenant, kStrictTenantQuotaBytes); + } + config.tenant_quota_connector_uri = + WriteTenantPolicyFile(tenant_quotas); + return config; + } Segment MakeSegment(std::string name = "test_segment", size_t base = kDefaultSegmentBase, - size_t size = kDefaultSegmentSize) const { + size_t size = kDefaultSegmentSize, + std::string host_id = "") const { Segment segment; segment.id = generate_uuid(); segment.name = std::move(name); segment.base = base; segment.size = size; segment.te_endpoint = segment.name; + segment.host_id = std::move(host_id); return segment; } @@ -65,9 +164,10 @@ class MasterServiceTest : public ::testing::Test { MountedSegmentContext PrepareSimpleSegment( MasterService& service, std::string name = "test_segment", - size_t base = kDefaultSegmentBase, - size_t size = kDefaultSegmentSize) const { - Segment segment = MakeSegment(std::move(name), base, size); + size_t base = kDefaultSegmentBase, size_t size = kDefaultSegmentSize, + std::string host_id = "") const { + Segment segment = + MakeSegment(std::move(name), base, size, std::move(host_id)); UUID client_id = generate_uuid(); auto mount_result = service.MountSegment(segment, client_id); EXPECT_TRUE(mount_result.has_value()); @@ -86,12 +186,13 @@ class MasterServiceTest : public ::testing::Test { config.replica_num = 1; config.preferred_segment = segment_name; - auto put_start = - service.PutStart(client_id, key, "default", slice_length, config); + auto put_start = service.PutStart(client_id, key, TenantId::Default(), + slice_length, config); EXPECT_TRUE(put_start.has_value()); - EXPECT_TRUE( - service.PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + EXPECT_TRUE(service + .PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); return key; } @@ -114,19 +215,19 @@ class MasterServiceTest : public ::testing::Test { const std::string& key, const ReplicateConfig& config, uint64_t slice_length = 1024) const { - auto put_start = - service.PutStart(client_id, key, "default", slice_length, config); + auto put_start = service.PutStart(client_id, key, TenantId::Default(), + slice_length, config); ASSERT_TRUE(put_start.has_value()) << "PutStart failed for key=" << key << ", error=" << toString(put_start.error()); - ASSERT_TRUE( - service.PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service + .PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); } void PutCompletedObject(MasterService& service, const UUID& client_id, - const std::string& key, - const std::string& tenant_id, + const std::string& key, const TenantId& tenant_id, const ReplicateConfig& config, uint64_t slice_length = 1024) const { auto put_start = @@ -156,13 +257,17 @@ class MasterServiceTest : public ::testing::Test { ReplicaMovePayload payload; struct_json::from_json(payload, assignment.payload); + const TenantId tenant_id(payload.tenant_id); + EXPECT_TRUE(tenant_id.IsValid()); + if (!tenant_id.IsValid()) { + return false; + } auto move_start = - service.MoveStart(client_id, payload.key, payload.tenant_id, + service.MoveStart(client_id, payload.key, tenant_id, payload.source, payload.target); EXPECT_TRUE(move_start.has_value()); EXPECT_TRUE( - service.MoveEnd(client_id, payload.key, payload.tenant_id) - .has_value()); + service.MoveEnd(client_id, payload.key, tenant_id).has_value()); TaskCompleteRequest complete_request; complete_request.id = assignment.id; @@ -217,28 +322,72 @@ class MasterServiceTest : public ::testing::Test { } std::vector replica_list; + std::vector policy_files_; + size_t next_policy_file_ = 0; - void TearDown() override { google::ShutdownGoogleLogging(); } + void TearDown() override { + for (const auto& path : policy_files_) { + std::error_code ec; + std::filesystem::remove(path, ec); + } + google::ShutdownGoogleLogging(); + } }; +TEST_F(MasterServiceTest, ObjectChecksumIsStoredAndClearedByUpsert) { + MasterService service; + const auto segment = PrepareSimpleSegment(service, "checksum_segment"); + const std::string key = "checksum_key"; + ReplicateConfig config; + config.replica_num = 1; + config.preferred_segment = "checksum_segment"; + + ASSERT_TRUE( + service + .PutStart(segment.client_id, key, TenantId::Default(), 1024, config) + .has_value()); + constexpr uint64_t kChecksum = 0; + ASSERT_TRUE(service + .PutEnd(segment.client_id, ObjectMeta{key, kChecksum}, + TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + + auto query = service.GetReplicaList(key, TenantId::Default()); + ASSERT_TRUE(query.has_value()); + ASSERT_TRUE(query->object_checksum.has_value()); + EXPECT_EQ(*query->object_checksum, kChecksum); + + ASSERT_TRUE(service + .UpsertStart(segment.client_id, key, TenantId::Default(), + 1024, config) + .has_value()); + ASSERT_TRUE(service + .UpsertEnd(segment.client_id, ObjectMeta{key, std::nullopt}, + TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + query = service.GetReplicaList(key, TenantId::Default()); + ASSERT_TRUE(query.has_value()); + EXPECT_FALSE(query->object_checksum.has_value()); +} + TEST(TenantScopedStorageKeyTest, RoundTripsAndParsesLegacyKeys) { const auto scoped = - MakeTenantScopedStorageKey("tenant:with:colon", "path/key:with:colon"); + TenantId("tenant:with:colon").MakeScopedKey("path/key:with:colon"); EXPECT_NE(scoped.find('\0'), std::string::npos); - auto [tenant_id, key] = ParseTenantScopedStorageKey(scoped); - EXPECT_EQ(tenant_id, "tenant:with:colon"); + auto [tenant_id, key] = TenantId::ParseScopedKey(scoped); + EXPECT_EQ(tenant_id.value(), "tenant:with:colon"); EXPECT_EQ(key, "path/key:with:colon"); - auto [default_tenant, default_key] = ParseTenantScopedStorageKey("raw_key"); - EXPECT_EQ(default_tenant, "default"); + auto [default_tenant, default_key] = TenantId::ParseScopedKey("raw_key"); + EXPECT_EQ(default_tenant.value(), TenantId::kDefaultValue); EXPECT_EQ(default_key, "raw_key"); std::string legacy = "legacy_tenant"; legacy.push_back('\0'); legacy.append("legacy_key"); - auto [legacy_tenant, legacy_key] = ParseTenantScopedStorageKey(legacy); - EXPECT_EQ(legacy_tenant, "legacy_tenant"); + auto [legacy_tenant, legacy_key] = TenantId::ParseScopedKey(legacy); + EXPECT_EQ(legacy_tenant.value(), "legacy_tenant"); EXPECT_EQ(legacy_key, "legacy_key"); } @@ -252,14 +401,14 @@ std::string GenerateKeyForSegment(const UUID& client_id, std::vector replica_list; // Check if the key already exists. - auto exist_result = service->ExistKey(key, "default"); + auto exist_result = service->ExistKey(key, TenantId::Default()); if (exist_result.has_value() && exist_result.value()) { continue; // Retry if the key already exists } // Attempt to put the key. - auto put_result = service->PutStart(client_id, key, "default", {1024}, - {.replica_num = 1}); + auto put_result = service->PutStart(client_id, key, TenantId::Default(), + {1024}, {.replica_num = 1}); if (put_result.has_value()) { replica_list = std::move(put_result.value()); } @@ -273,8 +422,8 @@ std::string GenerateKeyForSegment(const UUID& client_id, throw std::runtime_error("PutStart failed with code: " + std::to_string(static_cast(code))); } - auto put_end_result = - service->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); if (!put_end_result.has_value()) { throw std::runtime_error("PutEnd failed"); } @@ -284,7 +433,7 @@ std::string GenerateKeyForSegment(const UUID& client_id, return key; } // Clean up failed attempt - auto remove_result = service->Remove(key, "default"); + auto remove_result = service->Remove(key, TenantId::Default()); if (!remove_result.has_value()) { // Ignore cleanup failure } @@ -499,13 +648,14 @@ TEST_F(MasterServiceTest, PutStartInvalidParams) { config.replica_num = 0; config.nof_replica_num = 0; auto put_result1 = - service_->PutStart(client_id, key, "default", 1024, config); + service_->PutStart(client_id, key, TenantId::Default(), 1024, config); EXPECT_FALSE(put_result1.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result1.error()); // Test zero slice_length config.replica_num = 1; - auto put_result2 = service_->PutStart(client_id, key, "default", 0, config); + auto put_result2 = + service_->PutStart(client_id, key, TenantId::Default(), 0, config); EXPECT_FALSE(put_result2.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result2.error()); @@ -513,7 +663,7 @@ TEST_F(MasterServiceTest, PutStartInvalidParams) { config.nof_replica_num = 1; config.prefer_alloc_in_same_node = true; auto put_result3 = - service_->PutStart(client_id, key, "default", 1024, config); + service_->PutStart(client_id, key, TenantId::Default(), 1024, config); EXPECT_FALSE(put_result3.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, put_result3.error()); } @@ -529,16 +679,16 @@ TEST_F(MasterServiceTest, PutEndAllCompletesMemoryAndNoFReplicas) { ReplicateConfig config; config.replica_num = 1; config.nof_replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, "test_key_all", "default", 1024, config); + auto put_start_result = service_->PutStart( + client_id, "test_key_all", TenantId::Default(), 1024, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = service_->PutEnd(client_id, "test_key_all", "default", - ReplicaType::ALL); + auto put_end_result = service_->PutEnd( + client_id, "test_key_all", TenantId::Default(), ReplicaType::ALL); ASSERT_TRUE(put_end_result.has_value()); auto get_replica_result = - service_->GetReplicaList("test_key_all", "default"); + service_->GetReplicaList("test_key_all", TenantId::Default()); ASSERT_TRUE(get_replica_result.has_value()); bool has_complete_memory = false; @@ -568,27 +718,27 @@ TEST_F(MasterServiceTest, PutEndMemoryDoesNotCompleteNoFReplica) { ReplicateConfig config; config.replica_num = 1; config.nof_replica_num = 1; - auto put_start_result = service_->PutStart(client_id, "test_key_split", - "default", 1024, config); + auto put_start_result = service_->PutStart( + client_id, "test_key_split", TenantId::Default(), 1024, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = service_->PutEnd(client_id, "test_key_split", - "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, "test_key_split", TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); auto get_replica_result = - service_->GetReplicaList("test_key_split", "default"); + service_->GetReplicaList("test_key_split", TenantId::Default()); ASSERT_TRUE(get_replica_result.has_value()); ASSERT_EQ(get_replica_result->replicas.size(), 1u); EXPECT_TRUE(get_replica_result->replicas[0].is_memory_replica()); EXPECT_EQ(get_replica_result->replicas[0].status, ReplicaStatus::COMPLETE); auto put_revoke_result = service_->PutRevoke( - client_id, "test_key_split", "default", ReplicaType::NOF_SSD); + client_id, "test_key_split", TenantId::Default(), ReplicaType::NOF_SSD); ASSERT_TRUE(put_revoke_result.has_value()); auto final_replica_result = - service_->GetReplicaList("test_key_split", "default"); + service_->GetReplicaList("test_key_split", TenantId::Default()); ASSERT_TRUE(final_replica_result.has_value()); ASSERT_EQ(final_replica_result->replicas.size(), 1u); EXPECT_TRUE(final_replica_result->replicas[0].is_memory_replica()); @@ -605,7 +755,7 @@ TEST_F(MasterServiceTest, PutStartOnePlusOneAllowsSingleAllocatedReplica) { config.replica_num = 1; config.nof_replica_num = 1; auto put_start_result = service_->PutStart( - client_id, "test_key_one_plus_one", "default", 1024, config); + client_id, "test_key_one_plus_one", TenantId::Default(), 1024, config); ASSERT_TRUE(put_start_result.has_value()); ASSERT_EQ(put_start_result->size(), 1u); EXPECT_TRUE(put_start_result->front().is_memory_replica()); @@ -621,26 +771,26 @@ TEST_F(MasterServiceTest, PutStartGroupIdsValidation) { config.replica_num = 1; config.group_ids = std::vector{}; - auto empty_group_ids = service_->PutStart(client_id, "empty_group_ids", - "default", 1024, config); + auto empty_group_ids = service_->PutStart( + client_id, "empty_group_ids", TenantId::Default(), 1024, config); EXPECT_FALSE(empty_group_ids.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, empty_group_ids.error()); config.group_ids = std::vector{"g0", "g1"}; auto too_many_group_ids = service_->PutStart( - client_id, "too_many_group_ids", "default", 1024, config); + client_id, "too_many_group_ids", TenantId::Default(), 1024, config); EXPECT_FALSE(too_many_group_ids.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, too_many_group_ids.error()); config.group_ids = std::vector{""}; auto ungrouped = service_->PutStart(client_id, "explicit_ungrouped", - "default", 1024, config); + TenantId::Default(), 1024, config); ASSERT_TRUE(ungrouped.has_value()); ASSERT_TRUE(service_ - ->PutEnd(client_id, "explicit_ungrouped", "default", - ReplicaType::MEMORY) + ->PutEnd(client_id, "explicit_ungrouped", + TenantId::Default(), ReplicaType::MEMORY) .has_value()); - auto exists = service_->ExistKey("explicit_ungrouped", "default"); + auto exists = service_->ExistKey("explicit_ungrouped", TenantId::Default()); ASSERT_TRUE(exists.has_value()); EXPECT_TRUE(exists.value()); } @@ -658,25 +808,27 @@ TEST_F(MasterServiceTest, GroupedObjectRoutesKeyLevelLookupAndRemove) { PutCompletedObject(*service_, client_id, key, config); - auto exists = service_->ExistKey(key, "default"); + auto exists = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exists.has_value()); EXPECT_TRUE(exists.value()); - EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value()); + EXPECT_TRUE(service_->GetReplicaList(key, TenantId::Default()).has_value()); - ASSERT_TRUE(service_->Remove(key, "default", /*force=*/true).has_value()); - auto exists_after_remove = service_->ExistKey(key, "default"); + ASSERT_TRUE( + service_->Remove(key, TenantId::Default(), /*force=*/true).has_value()); + auto exists_after_remove = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exists_after_remove.has_value()); EXPECT_FALSE(exists_after_remove.value()); } TEST_F(MasterServiceTest, GroupRoutingIsTenantScopedForSameUserKey) { - std::unique_ptr service_(new MasterService()); + const std::string key = "tenant_grouped_shared_user_key"; + const TenantId tenant_a("tenant_group_route_a"); + const TenantId tenant_b("tenant_group_route_b"); + auto service_ = std::make_unique( + MakeStrictTenantConfig({tenant_a.value(), tenant_b.value()})); [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_); const UUID client_id = generate_uuid(); - const std::string key = "tenant_grouped_shared_user_key"; - const std::string tenant_a = "tenant_group_route_a"; - const std::string tenant_b = "tenant_group_route_b"; const std::string group_a = FindGroupIdOnDifferentShard(key); std::string group_b; for (int i = 0; i < 10000; ++i) { @@ -713,6 +865,28 @@ TEST_F(MasterServiceTest, GroupRoutingIsTenantScopedForSameUserKey) { EXPECT_TRUE(service_->GetReplicaList(key, tenant_b).has_value()); } +TEST_F(MasterServiceTest, StandbySnapshotRestorePreservesTenantScopedKeys) { + const TenantId tenant_a("tenant_restore_a"); + const TenantId tenant_b("tenant_restore_b"); + MasterService service( + MakeStrictTenantConfig({tenant_a.value(), tenant_b.value()})); + const std::string key = "shared_restore_key"; + + Replica replica(generate_uuid(), 128, "local://standby", + ReplicaStatus::COMPLETE); + StandbyObjectMetadata metadata; + metadata.client_id = generate_uuid(); + metadata.size = 128; + metadata.replicas.push_back(replica.get_descriptor()); + + service.RestoreFromStandbySnapshot({{tenant_a.value(), key, metadata}}, + /*initial_oplog_sequence_id=*/0, {}); + + EXPECT_TRUE(service.ExistKey(key, tenant_a).value_or(false)); + EXPECT_FALSE(service.ExistKey(key, tenant_b).value_or(true)); + EXPECT_FALSE(service.ExistKey(key, TenantId::Default()).value_or(true)); +} + TEST_F(MasterServiceTest, BatchGetReplicaListPreservesOrderWithGroupedKeys) { std::unique_ptr service_(new MasterService()); [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_); @@ -741,13 +915,13 @@ TEST_F(MasterServiceTest, BatchGetReplicaListPreservesOrderWithGroupedKeys) { PutCompletedObject(*service_, client_id, grouped_key_b, grouped_config_b); ASSERT_TRUE(service_ - ->PutStart(client_id, pending_key, "default", 1024, - ungrouped_config) + ->PutStart(client_id, pending_key, TenantId::Default(), + 1024, ungrouped_config) .has_value()); const std::vector keys = { grouped_key_a, missing_key, ungrouped_key, grouped_key_b, pending_key}; - auto results = service_->BatchGetReplicaList(keys, "default"); + auto results = service_->BatchGetReplicaList(keys, TenantId::Default()); ASSERT_EQ(results.size(), keys.size()); ASSERT_TRUE(results[0].has_value()); @@ -763,14 +937,15 @@ TEST_F(MasterServiceTest, BatchGetReplicaListPreservesOrderWithGroupedKeys) { } TEST_F(MasterServiceTest, BatchGetReplicaListKeepsTenantIsolation) { - std::unique_ptr service_(new MasterService()); + const std::string key = "batch_get_tenant_shared_key"; + const TenantId tenant_a("batch_get_tenant_a"); + const TenantId tenant_b("batch_get_tenant_b"); + auto service_ = std::make_unique( + MakeStrictTenantConfig({std::string(TenantId::kDefaultValue), + tenant_a.value(), tenant_b.value()})); [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_); const UUID client_id = generate_uuid(); - const std::string key = "batch_get_tenant_shared_key"; - const std::string tenant_a = "batch_get_tenant_a"; - const std::string tenant_b = "batch_get_tenant_b"; - ReplicateConfig config_a; config_a.replica_num = 1; config_a.group_ids = @@ -789,7 +964,8 @@ TEST_F(MasterServiceTest, BatchGetReplicaListKeepsTenantIsolation) { auto tenant_a_results = service_->BatchGetReplicaList({key}, tenant_a); auto tenant_b_results = service_->BatchGetReplicaList({key}, tenant_b); - auto default_results = service_->BatchGetReplicaList({key}, "default"); + auto default_results = + service_->BatchGetReplicaList({key}, TenantId::Default()); ASSERT_EQ(tenant_a_results.size(), 1u); ASSERT_EQ(tenant_b_results.size(), 1u); @@ -801,10 +977,12 @@ TEST_F(MasterServiceTest, BatchGetReplicaListKeepsTenantIsolation) { } TEST_F(MasterServiceTest, GetAllKeysListsOnlyRequestedTenant) { - std::unique_ptr service_(new MasterService()); + const TenantId tenant_a("tenant_get_all_keys_a"); + auto service_ = std::make_unique(MakeStrictTenantConfig( + {std::string(TenantId::kDefaultValue), tenant_a.value()})); [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_); const UUID client_id = generate_uuid(); - const std::string tenant_a = "tenant_get_all_keys_a"; + const std::string shared_key = "shared_listing_key"; const std::string default_only_key = "default_listing_key"; const std::string tenant_only_key = "tenant_listing_key"; @@ -812,16 +990,19 @@ TEST_F(MasterServiceTest, GetAllKeysListsOnlyRequestedTenant) { ReplicateConfig config; config.replica_num = 1; ASSERT_TRUE( - service_->PutStart(client_id, shared_key, "default", 1024, config) - .has_value()); - ASSERT_TRUE( - service_->PutEnd(client_id, shared_key, "default", ReplicaType::MEMORY) - .has_value()); - ASSERT_TRUE( - service_->PutStart(client_id, default_only_key, "default", 1024, config) + service_ + ->PutStart(client_id, shared_key, TenantId::Default(), 1024, config) .has_value()); ASSERT_TRUE(service_ - ->PutEnd(client_id, default_only_key, "default", + ->PutEnd(client_id, shared_key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, default_only_key, TenantId::Default(), + 1024, config) + .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, default_only_key, TenantId::Default(), ReplicaType::MEMORY) .has_value()); ASSERT_TRUE( @@ -838,7 +1019,7 @@ TEST_F(MasterServiceTest, GetAllKeysListsOnlyRequestedTenant) { ->PutEnd(client_id, tenant_only_key, tenant_a, ReplicaType::MEMORY) .has_value()); - auto default_keys = service_->GetAllKeys("default"); + auto default_keys = service_->GetAllKeys(TenantId::Default()); ASSERT_TRUE(default_keys.has_value()); EXPECT_NE(std::find(default_keys->begin(), default_keys->end(), shared_key), default_keys->end()); @@ -868,7 +1049,7 @@ TEST_F(MasterServiceTest, const UUID client_id = generate_uuid(); const std::string key = "concurrent_grouped_ungrouped_first_create"; - const std::string tenant_id = "tenant_concurrent_first_create"; + const TenantId tenant_id("tenant_concurrent_first_create"); ReplicateConfig ungrouped_config; ungrouped_config.replica_num = 1; ReplicateConfig grouped_config; @@ -929,7 +1110,7 @@ TEST_F(MasterServiceTest, const UUID client_id = generate_uuid(); const std::string key = "concurrent_different_grouped_first_create"; - const std::string tenant_id = "tenant_concurrent_grouped_first_create"; + const TenantId tenant_id("tenant_concurrent_grouped_first_create"); const std::string group_a = FindGroupIdOnDifferentShard(key); std::string group_b; for (int i = 0; i < 10000; ++i) { @@ -1006,19 +1187,22 @@ TEST_F(MasterServiceTest, ExpiredGroupedPutCanBeReplacedByUngroupedPut) { grouped_config.group_ids = std::vector{FindGroupIdOnDifferentShard(key)}; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", 1024, grouped_config) - .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), 1024, + grouped_config) + .has_value()); std::this_thread::sleep_for(std::chrono::milliseconds(2)); ReplicateConfig ungrouped_config; ungrouped_config.replica_num = 1; - auto put_start = - service_->PutStart(client_id, key, "default", 1024, ungrouped_config); + auto put_start = service_->PutStart(client_id, key, TenantId::Default(), + 1024, ungrouped_config); ASSERT_TRUE(put_start.has_value()) << toString(put_start.error()); - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->ExistKey(key, "default").value_or(false)); + ASSERT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE(service_->ExistKey(key, TenantId::Default()).value_or(false)); } TEST_F(MasterServiceTest, BatchRemoveUnregistersGroupedRoute) { @@ -1033,15 +1217,15 @@ TEST_F(MasterServiceTest, BatchRemoveUnregistersGroupedRoute) { std::vector{FindGroupIdOnDifferentShard(key)}; PutCompletedObject(*service_, client_id, key, grouped_config); - auto remove_results = service_->BatchRemove(std::vector{key}, - "default", /*force=*/true); + auto remove_results = service_->BatchRemove( + std::vector{key}, TenantId::Default(), /*force=*/true); ASSERT_EQ(remove_results.size(), 1u); ASSERT_TRUE(remove_results[0].has_value()); ReplicateConfig ungrouped_config; ungrouped_config.replica_num = 1; PutCompletedObject(*service_, client_id, key, ungrouped_config); - EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value()); + EXPECT_TRUE(service_->GetReplicaList(key, TenantId::Default()).has_value()); } TEST_F(MasterServiceTest, RemoveByRegexUnregistersGroupedRoute) { @@ -1056,16 +1240,16 @@ TEST_F(MasterServiceTest, RemoveByRegexUnregistersGroupedRoute) { std::vector{FindGroupIdOnDifferentShard(key)}; PutCompletedObject(*service_, client_id, key, grouped_config); - auto removed = - service_->RemoveByRegex("^regex_remove_grouped_route$", "default", - /*force=*/true); + auto removed = service_->RemoveByRegex("^regex_remove_grouped_route$", + TenantId::Default(), + /*force=*/true); ASSERT_TRUE(removed.has_value()); EXPECT_EQ(removed.value(), 1); ReplicateConfig ungrouped_config; ungrouped_config.replica_num = 1; PutCompletedObject(*service_, client_id, key, ungrouped_config); - EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value()); + EXPECT_TRUE(service_->GetReplicaList(key, TenantId::Default()).has_value()); } TEST_F(MasterServiceTest, GroupedLeaseRefreshNearExpiryProtectsCurrentMembers) { @@ -1087,22 +1271,24 @@ TEST_F(MasterServiceTest, GroupedLeaseRefreshNearExpiryProtectsCurrentMembers) { PutCompletedObject(*service_, client_id, key_a, config_a); PutCompletedObject(*service_, client_id, key_b, config_b); - auto exists = service_->ExistKey(key_a, "default"); + auto exists = service_->ExistKey(key_a, TenantId::Default()); ASSERT_TRUE(exists.has_value()); ASSERT_TRUE(exists.value()); std::this_thread::sleep_for(std::chrono::milliseconds(120)); - exists = service_->ExistKey(key_a, "default"); + exists = service_->ExistKey(key_a, TenantId::Default()); ASSERT_TRUE(exists.has_value()); ASSERT_TRUE(exists.value()); std::this_thread::sleep_for(std::chrono::milliseconds(100)); - auto remove_group_peer = service_->Remove(key_b, "default"); + auto remove_group_peer = service_->Remove(key_b, TenantId::Default()); ASSERT_FALSE(remove_group_peer.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_group_peer.error()); - EXPECT_TRUE(service_->Remove(key_a, "default", /*force=*/true).has_value()); - EXPECT_TRUE(service_->Remove(key_b, "default", /*force=*/true).has_value()); + EXPECT_TRUE(service_->Remove(key_a, TenantId::Default(), /*force=*/true) + .has_value()); + EXPECT_TRUE(service_->Remove(key_b, TenantId::Default(), /*force=*/true) + .has_value()); } TEST_F(MasterServiceTest, @@ -1122,22 +1308,24 @@ TEST_F(MasterServiceTest, config.group_ids = std::vector{group_id}; PutCompletedObject(*service_, client_id, key_a, config); - ASSERT_TRUE(service_->ExistKey(key_a, "default").value_or(false)); + ASSERT_TRUE(service_->ExistKey(key_a, TenantId::Default()).value_or(false)); PutCompletedObject(*service_, client_id, key_b, config); std::this_thread::sleep_for(std::chrono::milliseconds(150)); - auto exists = service_->ExistKey(key_a, "default"); + auto exists = service_->ExistKey(key_a, TenantId::Default()); ASSERT_TRUE(exists.has_value()); ASSERT_TRUE(exists.value()); std::this_thread::sleep_for(std::chrono::milliseconds(390)); - auto remove_group_peer = service_->Remove(key_b, "default"); + auto remove_group_peer = service_->Remove(key_b, TenantId::Default()); ASSERT_FALSE(remove_group_peer.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_group_peer.error()); - EXPECT_TRUE(service_->Remove(key_a, "default", /*force=*/true).has_value()); - EXPECT_TRUE(service_->Remove(key_b, "default", /*force=*/true).has_value()); + EXPECT_TRUE(service_->Remove(key_a, TenantId::Default(), /*force=*/true) + .has_value()); + EXPECT_TRUE(service_->Remove(key_b, TenantId::Default(), /*force=*/true) + .has_value()); } TEST_F(MasterServiceTest, RemoveGroupedMemberPreservesOtherMembers) { @@ -1155,15 +1343,18 @@ TEST_F(MasterServiceTest, RemoveGroupedMemberPreservesOtherMembers) { PutCompletedObject(*service_, client_id, key_a, config); PutCompletedObject(*service_, client_id, key_b, config); - ASSERT_TRUE(service_->Remove(key_a, "default", /*force=*/true).has_value()); + ASSERT_TRUE(service_->Remove(key_a, TenantId::Default(), /*force=*/true) + .has_value()); - auto removed_exists = service_->ExistKey(key_a, "default"); + auto removed_exists = service_->ExistKey(key_a, TenantId::Default()); ASSERT_TRUE(removed_exists.has_value()); EXPECT_FALSE(removed_exists.value()); - EXPECT_TRUE(service_->GetReplicaList(key_b, "default").has_value()); + EXPECT_TRUE( + service_->GetReplicaList(key_b, TenantId::Default()).has_value()); - ASSERT_TRUE(service_->Remove(key_b, "default", /*force=*/true).has_value()); - auto group_empty = service_->ExistKey(key_b, "default"); + ASSERT_TRUE(service_->Remove(key_b, TenantId::Default(), /*force=*/true) + .has_value()); + auto group_empty = service_->ExistKey(key_b, TenantId::Default()); ASSERT_TRUE(group_empty.has_value()); EXPECT_FALSE(group_empty.value()); } @@ -1183,21 +1374,22 @@ TEST_F(MasterServiceTest, UpsertPreservesGroupMembership) { ReplicateConfig unset_group_config; unset_group_config.replica_num = 1; - auto preserve_result = service_->UpsertStart(client_id, key, "default", - 1024, unset_group_config); + auto preserve_result = service_->UpsertStart( + client_id, key, TenantId::Default(), 1024, unset_group_config); ASSERT_TRUE(preserve_result.has_value()) << "Unset group_ids should preserve existing group membership"; - ASSERT_TRUE( - service_->UpsertEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value()); + ASSERT_TRUE(service_ + ->UpsertEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE(service_->GetReplicaList(key, TenantId::Default()).has_value()); ReplicateConfig different_group_config; different_group_config.replica_num = 1; different_group_config.group_ids = std::vector{FindGroupIdOnDifferentShard(key + "_other")}; auto different_group_result = service_->UpsertStart( - client_id, key, "default", 1024, different_group_config); + client_id, key, TenantId::Default(), 1024, different_group_config); ASSERT_FALSE(different_group_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, different_group_result.error()); @@ -1205,7 +1397,7 @@ TEST_F(MasterServiceTest, UpsertPreservesGroupMembership) { explicit_ungrouped_config.replica_num = 1; explicit_ungrouped_config.group_ids = std::vector{""}; auto explicit_ungrouped_result = service_->UpsertStart( - client_id, key, "default", 1024, explicit_ungrouped_config); + client_id, key, TenantId::Default(), 1024, explicit_ungrouped_config); ASSERT_FALSE(explicit_ungrouped_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, explicit_ungrouped_result.error()); } @@ -1221,20 +1413,22 @@ TEST_F(MasterServiceTest, IncompleteGroupedUpsertCanBecomeUngrouped) { grouped_config.group_ids = std::vector{FindGroupIdOnDifferentShard(key)}; - ASSERT_TRUE( - service_->PutStart(client_id, key, "default", 1024, grouped_config) - .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), 1024, + grouped_config) + .has_value()); std::this_thread::sleep_for(std::chrono::milliseconds(2)); ReplicateConfig ungrouped_config; ungrouped_config.replica_num = 1; - auto upsert_start = service_->UpsertStart(client_id, key, "default", 1024, - ungrouped_config); + auto upsert_start = service_->UpsertStart( + client_id, key, TenantId::Default(), 1024, ungrouped_config); ASSERT_TRUE(upsert_start.has_value()) << toString(upsert_start.error()); - ASSERT_TRUE( - service_->UpsertEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->ExistKey(key, "default").value_or(false)); + ASSERT_TRUE(service_ + ->UpsertEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); + EXPECT_TRUE(service_->ExistKey(key, TenantId::Default()).value_or(false)); } TEST_F(MasterServiceTest, UpsertRejectsExistingUngroupedToGrouped) { @@ -1251,12 +1445,12 @@ TEST_F(MasterServiceTest, UpsertRejectsExistingUngroupedToGrouped) { grouped_config.replica_num = 1; grouped_config.group_ids = std::vector{FindGroupIdOnDifferentShard(key)}; - auto upsert_start = - service_->UpsertStart(client_id, key, "default", 2048, grouped_config); + auto upsert_start = service_->UpsertStart( + client_id, key, TenantId::Default(), 2048, grouped_config); ASSERT_FALSE(upsert_start.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, upsert_start.error()); - EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value()); + EXPECT_TRUE(service_->GetReplicaList(key, TenantId::Default()).has_value()); } TEST_F(MasterServiceTest, @@ -1287,16 +1481,18 @@ TEST_F(MasterServiceTest, ReplicateConfig trigger_config; trigger_config.replica_num = 1; - auto trigger_result = - service_->PutStart(client_id, "trigger_grouped_eviction", "default", - kObjectSize, trigger_config); + auto trigger_result = service_->PutStart( + client_id, "trigger_grouped_eviction", TenantId::Default(), + kObjectSize, trigger_config); ASSERT_FALSE(trigger_result.has_value()); EXPECT_EQ(ErrorCode::NO_AVAILABLE_HANDLE, trigger_result.error()); std::this_thread::sleep_for(std::chrono::milliseconds(200)); - EXPECT_FALSE(service_->ExistKey(evict_key_a, "default").value_or(true)); - EXPECT_FALSE(service_->ExistKey(evict_key_b, "default").value_or(true)); + EXPECT_FALSE(service_->ExistKey(evict_key_a, TenantId::Default()) + .value_or(true)); + EXPECT_FALSE(service_->ExistKey(evict_key_b, TenantId::Default()) + .value_or(true)); } { @@ -1318,24 +1514,24 @@ TEST_F(MasterServiceTest, PutCompletedObject(*service_, client_id, leased_key_b, leased_config, kObjectSize); - auto exists = service_->ExistKey(leased_key_a, "default"); + auto exists = service_->ExistKey(leased_key_a, TenantId::Default()); ASSERT_TRUE(exists.has_value()); ASSERT_TRUE(exists.value()); ReplicateConfig trigger_config; trigger_config.replica_num = 1; - auto trigger_result = - service_->PutStart(client_id, "trigger_leased_group_eviction", - "default", kObjectSize, trigger_config); + auto trigger_result = service_->PutStart( + client_id, "trigger_leased_group_eviction", TenantId::Default(), + kObjectSize, trigger_config); ASSERT_FALSE(trigger_result.has_value()); EXPECT_EQ(ErrorCode::NO_AVAILABLE_HANDLE, trigger_result.error()); std::this_thread::sleep_for(std::chrono::milliseconds(200)); - EXPECT_TRUE( - service_->GetReplicaList(leased_key_a, "default").has_value()); - EXPECT_TRUE( - service_->GetReplicaList(leased_key_b, "default").has_value()); + EXPECT_TRUE(service_->GetReplicaList(leased_key_a, TenantId::Default()) + .has_value()); + EXPECT_TRUE(service_->GetReplicaList(leased_key_b, TenantId::Default()) + .has_value()); } } @@ -1367,17 +1563,19 @@ TEST_F(MasterServiceTest, GroupedEvictionSkipsUnsafeMembersAndEvictsSafePeers) { trigger_config.replica_num = 1; auto trigger_result = service_->PutStart(client_id, "trigger_mixed_safety_group_eviction", - "default", kObjectSize, trigger_config); + TenantId::Default(), kObjectSize, trigger_config); ASSERT_FALSE(trigger_result.has_value()); EXPECT_EQ(ErrorCode::NO_AVAILABLE_HANDLE, trigger_result.error()); std::this_thread::sleep_for(std::chrono::milliseconds(200)); - EXPECT_FALSE(service_->ExistKey(safe_key, "default").value_or(true)); - EXPECT_TRUE( - service_->GetReplicaList(hard_pinned_key, "default").has_value()); - EXPECT_TRUE(service_->Remove(hard_pinned_key, "default", /*force=*/true) + EXPECT_FALSE( + service_->ExistKey(safe_key, TenantId::Default()).value_or(true)); + EXPECT_TRUE(service_->GetReplicaList(hard_pinned_key, TenantId::Default()) .has_value()); + EXPECT_TRUE( + service_->Remove(hard_pinned_key, TenantId::Default(), /*force=*/true) + .has_value()); } TEST_F(MasterServiceTest, BatchUpsertStartMixedGroupIdsPreservesOrder) { @@ -1398,27 +1596,29 @@ TEST_F(MasterServiceTest, BatchUpsertStartMixedGroupIdsPreservesOrder) { std::vector{FindGroupIdOnDifferentShard(keys[0]), "", FindGroupIdOnDifferentShard(keys[2])}; - auto results = - service_->BatchUpsertStart(client_id, keys, "default", sizes, config); + auto results = service_->BatchUpsertStart( + client_id, keys, TenantId::Default(), sizes, config); ASSERT_EQ(results.size(), keys.size()); for (const auto& result : results) { ASSERT_TRUE(result.has_value()); } - auto end_results = service_->BatchUpsertEnd(client_id, keys, "default"); + auto end_results = service_->BatchUpsertEnd( + client_id, MakeObjectMetas(keys), TenantId::Default()); ASSERT_EQ(end_results.size(), keys.size()); for (const auto& result : end_results) { ASSERT_TRUE(result.has_value()); } for (const auto& key : keys) { - EXPECT_TRUE(service_->GetReplicaList(key, "default").has_value()); + EXPECT_TRUE( + service_->GetReplicaList(key, TenantId::Default()).has_value()); } ReplicateConfig invalid_config = config; invalid_config.group_ids = std::vector{"only_one"}; auto invalid_results = service_->BatchUpsertStart( - client_id, keys, "default", sizes, invalid_config); + client_id, keys, TenantId::Default(), sizes, invalid_config); ASSERT_EQ(invalid_results.size(), keys.size()); for (const auto& result : invalid_results) { ASSERT_FALSE(result.has_value()); @@ -1455,14 +1655,16 @@ TEST_F(MasterServiceTest, WrappedBatchPutStartMixedGroupIdsPreservesOrder) { ASSERT_TRUE(result.has_value()) << toString(result.error()); } - auto end_results = service_.BatchPutEnd(client_id, keys); + auto end_results = service_.BatchPutEnd(client_id, MakeObjectMetas(keys)); ASSERT_EQ(end_results.size(), keys.size()); for (const auto& result : end_results) { ASSERT_TRUE(result.has_value()); } for (const auto& key : keys) { - EXPECT_TRUE(service_.GetReplicaList(key, "default").has_value()); + EXPECT_TRUE( + service_.GetReplicaList(key, std::string(TenantId::kDefaultValue)) + .has_value()); } ReplicateConfig invalid_config = config; @@ -1484,67 +1686,16 @@ TEST_F(MasterServiceTest, WrappedBatchPutStartMixedGroupIdsPreservesOrder) { } } -TEST_F(MasterServiceTest, PutStartEndFlow) { - std::unique_ptr service_(new MasterService()); - [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_); - const UUID client_id = generate_uuid(); - const UUID invalid_client_id = generate_uuid(); - ASSERT_NE(client_id, invalid_client_id); - - // Test PutStart - std::string key = "test_key"; - uint64_t value_length = 1024; - ReplicateConfig config; - config.replica_num = 1; - - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); - EXPECT_TRUE(put_start_result.has_value()); - replica_list = put_start_result.value(); - EXPECT_FALSE(replica_list.empty()); - EXPECT_EQ(ReplicaStatus::PROCESSING, replica_list[0].status); - - // During put, Get/Remove should fail - auto get_replica_result = service_->GetReplicaList(key, "default"); - EXPECT_FALSE(get_replica_result.has_value()); - EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_replica_result.error()); - auto remove_result = service_->Remove(key, "default"); - EXPECT_FALSE(remove_result.has_value()); - EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error()); - - // PutEnd should fail if the client_id does not match. - auto put_end_fail_result = service_->PutEnd(invalid_client_id, key, - "default", ReplicaType::MEMORY); - EXPECT_FALSE(put_end_fail_result.has_value()); - EXPECT_EQ(put_end_fail_result.error(), ErrorCode::ILLEGAL_CLIENT); - - // PutRevoke should fail if the client_id does not match. - auto put_revoke_fail_result = service_->PutRevoke( - invalid_client_id, key, "default", ReplicaType::MEMORY); - EXPECT_FALSE(put_revoke_fail_result.has_value()); - EXPECT_EQ(put_revoke_fail_result.error(), ErrorCode::ILLEGAL_CLIENT); - - // Test PutEnd - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); - EXPECT_TRUE(put_end_result.has_value()); - - // Verify replica list after PutEnd - auto final_get_result = service_->GetReplicaList(key, "default"); - EXPECT_TRUE(final_get_result.has_value()); - replica_list = final_get_result.value().replicas; - EXPECT_EQ(1, replica_list.size()); - EXPECT_EQ(ReplicaStatus::COMPLETE, replica_list[0].status); -} - TEST_F(MasterServiceTest, TenantPutGetRemoveIsolatesSameUserKey) { - std::unique_ptr service_(new MasterService()); + const std::string key = "shared_user_key"; + const TenantId tenant_a("tenant_a"); + const TenantId tenant_b("tenant_b"); + auto service_ = std::make_unique( + MakeStrictTenantConfig({std::string(TenantId::kDefaultValue), + tenant_a.value(), tenant_b.value()})); [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_); const UUID client_id = generate_uuid(); - const std::string key = "shared_user_key"; - const std::string tenant_a = "tenant_a"; - const std::string tenant_b = "tenant_b"; ReplicateConfig config; config.replica_num = 1; @@ -1557,8 +1708,9 @@ TEST_F(MasterServiceTest, TenantPutGetRemoveIsolatesSameUserKey) { ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_b, ReplicaType::MEMORY) .has_value()); - EXPECT_FALSE(service_->GetReplicaList(key, "default").has_value()); - EXPECT_FALSE(service_->ExistKey(key, "default").value()); + EXPECT_FALSE( + service_->GetReplicaList(key, TenantId::Default()).has_value()); + EXPECT_FALSE(service_->ExistKey(key, TenantId::Default()).value()); EXPECT_TRUE(service_->ExistKey(key, tenant_a).value()); EXPECT_TRUE(service_->ExistKey(key, tenant_b).value()); EXPECT_TRUE(service_->GetReplicaList(key, tenant_a).has_value()); @@ -1572,20 +1724,25 @@ TEST_F(MasterServiceTest, TenantPutGetRemoveIsolatesSameUserKey) { } TEST_F(MasterServiceTest, RegexOperationsAreTenantScoped) { - std::unique_ptr service_(new MasterService()); + const std::string key = "regex_shared_key"; + const TenantId tenant_a("tenant_regex_a"); + const TenantId tenant_b("tenant_regex_b"); + auto service_ = std::make_unique( + MakeStrictTenantConfig({std::string(TenantId::kDefaultValue), + tenant_a.value(), tenant_b.value()})); [[maybe_unused]] const auto context = PrepareSimpleSegment(*service_); const UUID client_id = generate_uuid(); - const std::string key = "regex_shared_key"; - const std::string tenant_a = "tenant_regex_a"; - const std::string tenant_b = "tenant_regex_b"; ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE(service_->PutStart(client_id, key, "default", 1024, config) - .has_value()); - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE( + service_->PutStart(client_id, key, TenantId::Default(), 1024, config) + .has_value()); + ASSERT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); ASSERT_TRUE( service_->PutStart(client_id, key, tenant_a, 1024, config).has_value()); ASSERT_TRUE(service_->PutEnd(client_id, key, tenant_a, ReplicaType::MEMORY) @@ -1596,15 +1753,16 @@ TEST_F(MasterServiceTest, RegexOperationsAreTenantScoped) { .has_value()); auto default_matches = - service_->GetReplicaListByRegex("^regex_shared", "default"); + service_->GetReplicaListByRegex("^regex_shared", TenantId::Default()); ASSERT_TRUE(default_matches.has_value()); EXPECT_EQ(default_matches->size(), 1); - auto remove_default = - service_->RemoveByRegex("^regex_shared", "default", /*force=*/true); + auto remove_default = service_->RemoveByRegex( + "^regex_shared", TenantId::Default(), /*force=*/true); ASSERT_TRUE(remove_default.has_value()); EXPECT_EQ(remove_default.value(), 1); - EXPECT_FALSE(service_->GetReplicaList(key, "default").has_value()); + EXPECT_FALSE( + service_->GetReplicaList(key, TenantId::Default()).has_value()); EXPECT_TRUE(service_->GetReplicaList(key, tenant_a).has_value()); EXPECT_TRUE(service_->GetReplicaList(key, tenant_b).has_value()); @@ -1617,15 +1775,16 @@ TEST_F(MasterServiceTest, RegexOperationsAreTenantScoped) { } TEST_F(MasterServiceTest, TenantBatchUpsertAndRevokeAreScoped) { - auto svc = std::make_unique(); - [[maybe_unused]] const auto context = PrepareSimpleSegment(*svc); - const UUID client_id = generate_uuid(); - const std::vector keys = {"tenant_batch_upsert_key_a", "tenant_batch_upsert_key_b"}; const std::vector sizes = {1024, 2048}; - const std::string tenant_a = "tenant_batch_upsert_a"; - const std::string tenant_b = "tenant_batch_upsert_b"; + const TenantId tenant_a("tenant_batch_upsert_a"); + const TenantId tenant_b("tenant_batch_upsert_b"); + auto svc = std::make_unique( + MakeStrictTenantConfig({std::string(TenantId::kDefaultValue), + tenant_a.value(), tenant_b.value()})); + [[maybe_unused]] const auto context = PrepareSimpleSegment(*svc); + const UUID client_id = generate_uuid(); ReplicateConfig config; config.replica_num = 1; @@ -1636,7 +1795,8 @@ TEST_F(MasterServiceTest, TenantBatchUpsertAndRevokeAreScoped) { for (const auto& result : tenant_a_results) { ASSERT_TRUE(result.has_value()); } - auto tenant_a_end = svc->BatchUpsertEnd(client_id, keys, tenant_a); + auto tenant_a_end = + svc->BatchUpsertEnd(client_id, MakeObjectMetas(keys), tenant_a); ASSERT_EQ(tenant_a_end.size(), keys.size()); for (const auto& result : tenant_a_end) { ASSERT_TRUE(result.has_value()); @@ -1648,14 +1808,15 @@ TEST_F(MasterServiceTest, TenantBatchUpsertAndRevokeAreScoped) { for (const auto& result : tenant_b_results) { ASSERT_TRUE(result.has_value()); } - auto tenant_b_end = svc->BatchUpsertEnd(client_id, keys, tenant_b); + auto tenant_b_end = + svc->BatchUpsertEnd(client_id, MakeObjectMetas(keys), tenant_b); ASSERT_EQ(tenant_b_end.size(), keys.size()); for (const auto& result : tenant_b_end) { ASSERT_TRUE(result.has_value()); } for (const auto& key : keys) { - EXPECT_FALSE(svc->GetReplicaList(key, "default").has_value()); + EXPECT_FALSE(svc->GetReplicaList(key, TenantId::Default()).has_value()); EXPECT_TRUE(svc->GetReplicaList(key, tenant_a).has_value()); EXPECT_TRUE(svc->GetReplicaList(key, tenant_b).has_value()); } @@ -1671,22 +1832,24 @@ TEST_F(MasterServiceTest, TenantBatchUpsertAndRevokeAreScoped) { } TEST_F(MasterServiceTest, TenantBatchRemoveAndRemoveAllAreScoped) { - auto svc = std::make_unique(); + const std::string shared_key = "tenant_batch_remove_shared_key"; + const TenantId tenant_a("tenant_batch_remove_a"); + const TenantId tenant_b("tenant_batch_remove_b"); + auto svc = std::make_unique( + MakeStrictTenantConfig({std::string(TenantId::kDefaultValue), + tenant_a.value(), tenant_b.value()})); [[maybe_unused]] const auto context = PrepareSimpleSegment(*svc); const UUID client_id = generate_uuid(); - const std::string shared_key = "tenant_batch_remove_shared_key"; - const std::string tenant_a = "tenant_batch_remove_a"; - const std::string tenant_b = "tenant_batch_remove_b"; - ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE(svc->PutStart(client_id, shared_key, "default", 1024, config) - .has_value()); ASSERT_TRUE( - svc->PutEnd(client_id, shared_key, "default", ReplicaType::MEMORY) + svc->PutStart(client_id, shared_key, TenantId::Default(), 1024, config) .has_value()); + ASSERT_TRUE(svc->PutEnd(client_id, shared_key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); ASSERT_TRUE(svc->PutStart(client_id, shared_key, tenant_a, 1024, config) .has_value()); ASSERT_TRUE( @@ -1702,33 +1865,38 @@ TEST_F(MasterServiceTest, TenantBatchRemoveAndRemoveAllAreScoped) { ASSERT_EQ(remove_a.size(), 1u); ASSERT_TRUE(remove_a[0].has_value()); EXPECT_FALSE(svc->GetReplicaList(shared_key, tenant_a).has_value()); - EXPECT_TRUE(svc->GetReplicaList(shared_key, "default").has_value()); + EXPECT_TRUE( + svc->GetReplicaList(shared_key, TenantId::Default()).has_value()); EXPECT_TRUE(svc->GetReplicaList(shared_key, tenant_b).has_value()); EXPECT_EQ(svc->RemoveAll(tenant_b, /*force=*/true), 1); EXPECT_FALSE(svc->GetReplicaList(shared_key, tenant_b).has_value()); - EXPECT_TRUE(svc->GetReplicaList(shared_key, "default").has_value()); + EXPECT_TRUE( + svc->GetReplicaList(shared_key, TenantId::Default()).has_value()); EXPECT_EQ(svc->RemoveAll(/*force=*/true), 1); - EXPECT_FALSE(svc->GetReplicaList(shared_key, "default").has_value()); + EXPECT_FALSE( + svc->GetReplicaList(shared_key, TenantId::Default()).has_value()); } TEST_F(MasterServiceTest, LegacyRemoveAllRemovesAllTenants) { - auto svc = std::make_unique(); + const std::string key = "legacy_remove_all_shared_key"; + const TenantId tenant_a("legacy_remove_all_a"); + const TenantId tenant_b("legacy_remove_all_b"); + auto svc = std::make_unique( + MakeStrictTenantConfig({std::string(TenantId::kDefaultValue), + tenant_a.value(), tenant_b.value()})); [[maybe_unused]] const auto context = PrepareSimpleSegment(*svc); const UUID client_id = generate_uuid(); - const std::string key = "legacy_remove_all_shared_key"; - const std::string tenant_a = "legacy_remove_all_a"; - const std::string tenant_b = "legacy_remove_all_b"; - ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - svc->PutStart(client_id, key, "default", 1024, config).has_value()); - ASSERT_TRUE(svc->PutEnd(client_id, key, "default", ReplicaType::MEMORY) + ASSERT_TRUE(svc->PutStart(client_id, key, TenantId::Default(), 1024, config) .has_value()); + ASSERT_TRUE( + svc->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); ASSERT_TRUE( svc->PutStart(client_id, key, tenant_a, 1024, config).has_value()); ASSERT_TRUE( @@ -1739,7 +1907,7 @@ TEST_F(MasterServiceTest, LegacyRemoveAllRemovesAllTenants) { svc->PutEnd(client_id, key, tenant_b, ReplicaType::MEMORY).has_value()); EXPECT_EQ(svc->RemoveAll(/*force=*/true), 3); - EXPECT_FALSE(svc->GetReplicaList(key, "default").has_value()); + EXPECT_FALSE(svc->GetReplicaList(key, TenantId::Default()).has_value()); EXPECT_FALSE(svc->GetReplicaList(key, tenant_a).has_value()); EXPECT_FALSE(svc->GetReplicaList(key, tenant_b).has_value()); EXPECT_EQ(svc->RemoveAll(/*force=*/true), 0); @@ -1769,8 +1937,8 @@ TEST_F(MasterServiceTest, PutWithPreferredSegment) { config.replica_num = 1; config.preferred_segment = preferred_segment; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(1, replica_list.size()); @@ -1780,8 +1948,8 @@ TEST_F(MasterServiceTest, PutWithPreferredSegment) { mem_desc.buffer_descriptor.transport_endpoint_); // Complete the Put operation - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); } @@ -1808,8 +1976,8 @@ TEST_F(MasterServiceTest, PutWithPreferredSegments) { config.replica_num = 2; config.preferred_segments = preferred_segments; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(2, replica_list.size()); @@ -1824,11 +1992,204 @@ TEST_F(MasterServiceTest, PutWithPreferredSegments) { EXPECT_TRUE(used_segments.find("segment_1") != used_segments.end()); // Complete the Put operation - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); } +TEST_F(MasterServiceTest, + ResolveMooncakeHostIdUsesLocalHostnameAndRejectsLoopback) { + { + EXPECT_EQ(ResolveMooncakeHostId("hostB:5000"), "hostB"); + EXPECT_EQ(ResolveMooncakeHostId("hostB:5001"), "hostB"); + EXPECT_EQ(ResolveMooncakeHostId("[2001:db8::1]:5000"), "2001:db8::1"); + EXPECT_TRUE(ResolveMooncakeHostId("localhost:5000").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("127.0.0.1:5000").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("0.0.0.0:5000").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("::1").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("[::1]:5000").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("::").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("[::]").empty()); + EXPECT_TRUE(ResolveMooncakeHostId("[::]:5000").empty()); + } +} + +TEST_F(MasterServiceTest, MasterConfigParsesLocalFirstStrategy) { + MasterConfig config{}; + config.allocation_strategy = "local_first"; + + WrappedMasterServiceConfig wrapped_config(config, 0); + MasterServiceConfig service_config(wrapped_config); + EXPECT_EQ(service_config.allocation_strategy_type, + AllocationStrategyType::LOCAL_FIRST); +} + +TEST_F(MasterServiceTest, LocalFirstPutPrefersWriterHost) { + auto service_config = + MasterServiceConfig::builder() + .set_allocation_strategy_type(AllocationStrategyType::LOCAL_FIRST) + .build(); + MasterService service(service_config); + const UUID writer_client_id = generate_uuid(); + + [[maybe_unused]] const auto host0 = PrepareSimpleSegment( + service, "segment_host0", 0x300000000, kDefaultSegmentSize, "host0"); + [[maybe_unused]] const auto host1 = PrepareSimpleSegment( + service, "segment_host1", 0x400000000, kDefaultSegmentSize, "host1"); + + ReplicateConfig config; + config.replica_num = 1; + config.host_id = "host1"; + + auto put_start = service.PutStart(writer_client_id, "local_first_key", + TenantId::Default(), 1024, config); + ASSERT_TRUE(put_start.has_value()); + ASSERT_EQ(put_start->size(), 1u); + EXPECT_EQ((*put_start)[0] + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + "segment_host1"); +} + +TEST_F(MasterServiceTest, LocalFirstPutFallsBackToNextOrderedHost) { + auto service_config = + MasterServiceConfig::builder() + .set_allocation_strategy_type(AllocationStrategyType::LOCAL_FIRST) + .build(); + MasterService service(service_config); + const UUID writer_client_id = generate_uuid(); + + [[maybe_unused]] const auto host0 = PrepareSimpleSegment( + service, "segment_host0", 0x300000000, kDefaultSegmentSize, "host0"); + [[maybe_unused]] const auto host2 = PrepareSimpleSegment( + service, "segment_host2", 0x400000000, kDefaultSegmentSize, "host2"); + + ReplicateConfig config; + config.replica_num = 1; + config.host_id = "host1"; + + auto put_start = service.PutStart(writer_client_id, "ordered_fallback_key", + TenantId::Default(), 1024, config); + ASSERT_TRUE(put_start.has_value()); + ASSERT_EQ(put_start->size(), 1u); + EXPECT_EQ((*put_start)[0] + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + "segment_host2"); +} + +TEST_F(MasterServiceTest, LocalFirstPutFallsBackWhenLocalSegmentIsFull) { + auto service_config = + MasterServiceConfig::builder() + .set_allocation_strategy_type(AllocationStrategyType::LOCAL_FIRST) + .build(); + MasterService service(service_config); + const UUID writer_client_id = generate_uuid(); + + [[maybe_unused]] const auto local = PrepareSimpleSegment( + service, "segment_host1", 0x300000000, 1024, "host1"); + [[maybe_unused]] const auto remote = PrepareSimpleSegment( + service, "segment_host2", 0x400000000, kDefaultSegmentSize, "host2"); + + ReplicateConfig config; + config.replica_num = 1; + config.host_id = "host1"; + + auto fill_start = service.PutStart(writer_client_id, "fill_local_segment", + TenantId::Default(), 1024, config); + ASSERT_TRUE(fill_start.has_value()); + ASSERT_EQ(fill_start->size(), 1u); + EXPECT_EQ((*fill_start)[0] + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + "segment_host1"); + ASSERT_TRUE(service + .PutEnd(writer_client_id, "fill_local_segment", + TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + + auto fallback_start = + service.PutStart(writer_client_id, "fallback_after_local_full", + TenantId::Default(), 1, config); + ASSERT_TRUE(fallback_start.has_value()); + ASSERT_EQ(fallback_start->size(), 1u); + EXPECT_EQ((*fallback_start)[0] + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + "segment_host2"); +} + +TEST_F(MasterServiceTest, ExplicitPreferredSegmentOverridesLocalFirst) { + auto service_config = + MasterServiceConfig::builder() + .set_allocation_strategy_type(AllocationStrategyType::LOCAL_FIRST) + .build(); + MasterService service(service_config); + const UUID writer_client_id = generate_uuid(); + + [[maybe_unused]] const auto host0 = PrepareSimpleSegment( + service, "segment_host0", 0x300000000, kDefaultSegmentSize, "host0"); + [[maybe_unused]] const auto host1 = PrepareSimpleSegment( + service, "segment_host1", 0x400000000, kDefaultSegmentSize, "host1"); + + ReplicateConfig config; + config.replica_num = 1; + config.host_id = "host1"; + config.preferred_segment = "segment_host0"; + + auto put_start = + service.PutStart(writer_client_id, "explicit_preferred_key", + TenantId::Default(), 1024, config); + ASSERT_TRUE(put_start.has_value()); + ASSERT_EQ(put_start->size(), 1u); + EXPECT_EQ((*put_start)[0] + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + "segment_host0"); +} + +TEST_F(MasterServiceTest, ExplicitPreferredSegmentFallsBackToLocalFirst) { + auto service_config = + MasterServiceConfig::builder() + .set_allocation_strategy_type(AllocationStrategyType::LOCAL_FIRST) + .build(); + MasterService service(service_config); + const UUID writer_client_id = generate_uuid(); + + [[maybe_unused]] const auto preferred = PrepareSimpleSegment( + service, "segment_host0", 0x300000000, 1024, "host0"); + [[maybe_unused]] const auto local = PrepareSimpleSegment( + service, "segment_host1", 0x400000000, kDefaultSegmentSize, "host1"); + + ReplicateConfig config; + config.replica_num = 1; + config.host_id = "host1"; + config.preferred_segment = "segment_host0"; + + auto fill_start = service.PutStart(writer_client_id, "fill_preferred", + TenantId::Default(), 1024, config); + ASSERT_TRUE(fill_start.has_value()); + ASSERT_EQ(fill_start->size(), 1u); + EXPECT_EQ((*fill_start)[0] + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + "segment_host0"); + ASSERT_TRUE(service + .PutEnd(writer_client_id, "fill_preferred", + TenantId::Default(), ReplicaType::MEMORY) + .has_value()); + + auto fallback_start = + service.PutStart(writer_client_id, "fallback_after_preferred_full", + TenantId::Default(), 1, config); + ASSERT_TRUE(fallback_start.has_value()); + ASSERT_EQ(fallback_start->size(), 1u); + EXPECT_EQ((*fallback_start)[0] + .get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + "segment_host1"); +} + TEST_F(MasterServiceTest, RandomPutStartEndFlow) { std::unique_ptr service_(new MasterService()); const UUID client_id = generate_uuid(); @@ -1851,25 +2212,25 @@ TEST_F(MasterServiceTest, RandomPutStartEndFlow) { std::uniform_int_distribution<> dis(1, 5); int random_number = dis(gen); config.replica_num = random_number; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_FALSE(replica_list.empty()); EXPECT_EQ(ReplicaStatus::PROCESSING, replica_list[0].status); // During put, Get/Remove should fail - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error()); // Test PutEnd - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Verify replica list after PutEnd - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result2.has_value()); replica_list = get_result2.value().replicas; EXPECT_EQ(random_number, replica_list.size()); @@ -1886,7 +2247,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegex) { std::unique_ptr service_(new MasterService(service_config)); const UUID client_id = generate_uuid(); // Test getting non-existent key - auto get_result = service_->GetReplicaList(".*non_existent.*", "default"); + auto get_result = + service_->GetReplicaList(".*non_existent.*", TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); @@ -1898,20 +2260,21 @@ TEST_F(MasterServiceTest, GetReplicaListByRegex) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); } // wait for all the lease to expire std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); // Test getting existing key - auto get_result2 = service_->GetReplicaListByRegex("^test_key", "default"); + auto get_result2 = + service_->GetReplicaListByRegex("^test_key", TenantId::Default()); EXPECT_TRUE(get_result2.has_value()); auto replica_list_local = get_result2.value(); EXPECT_EQ(10, replica_list_local.size()); @@ -1923,15 +2286,15 @@ void put_object(MasterService& service, const UUID& client_id, uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service.PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service.PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()) << "Failed to PutStart for key: " << key; - auto put_end_result = - service.PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service.PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()) << "Failed to PutEnd for key: " << key; - auto exist_result = service.ExistKey(key, "default"); + auto exist_result = service.ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()) << "Key does not exist after put: " << key; } @@ -1975,7 +2338,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { // Test 3.1: Simple prefix matching { - auto result = service_->GetReplicaListByRegex("^test_key_", "default"); + auto result = + service_->GetReplicaListByRegex("^test_key_", TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 3); // Matches test_key_01, test_key_02, test_key_10 @@ -1983,8 +2347,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { // Test 3.2: Matching with a wildcard for any number { - auto result = - service_->GetReplicaListByRegex("^test_key_\\d+$", "default"); + auto result = service_->GetReplicaListByRegex("^test_key_\\d+$", + TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 3); } @@ -1993,7 +2357,7 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { { // Matches "data_part_1_chunk_a" and "data_part_2_chunk_b" auto result = service_->GetReplicaListByRegex("^data_part_\\d_chunk_.$", - "default"); + TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 2); } @@ -2001,7 +2365,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { // Test 3.4: Matching keys containing a specific substring { // Matches all keys with "key" in them - auto result = service_->GetReplicaListByRegex("key", "default"); + auto result = + service_->GetReplicaListByRegex("key", TenantId::Default()); ASSERT_TRUE(result.has_value()); // Expected: test_key_01, test_key_02, test_key_10, // prod_key_alpha, prod_key_beta, @@ -2013,7 +2378,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { // Test 3.5: Matching based on file-like paths { // Match all .log files - auto result = service_->GetReplicaListByRegex("\\.log$", "default"); + auto result = + service_->GetReplicaListByRegex("\\.log$", TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 1); EXPECT_EQ(result.value().begin()->first, "logs/app-2025-08-13.log"); @@ -2022,8 +2388,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { // Test 3.6: OR condition using | { // Match keys starting with "prod" OR ending with "json" - auto result = - service_->GetReplicaListByRegex("^prod|\\.json$", "default"); + auto result = service_->GetReplicaListByRegex("^prod|\\.json$", + TenantId::Default()); ASSERT_TRUE(result.has_value()); // Expected: prod_key_alpha, prod_key_beta, config/user/settings.json EXPECT_EQ(result.value().size(), 3); @@ -2031,8 +2397,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { // Test 3.7: Regex that should not match anything { - auto result = - service_->GetReplicaListByRegex("^non_existent_prefix_", "default"); + auto result = service_->GetReplicaListByRegex("^non_existent_prefix_", + TenantId::Default()); // This should succeed but return an empty map. ASSERT_TRUE(result.has_value()); EXPECT_TRUE(result.value().empty()); @@ -2040,7 +2406,8 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { // Test 3.8: Exact match regex { - auto result = service_->GetReplicaListByRegex("^short$", "default"); + auto result = + service_->GetReplicaListByRegex("^short$", TenantId::Default()); ASSERT_TRUE(result.has_value()); EXPECT_EQ(result.value().size(), 1); EXPECT_EQ(result.value().begin()->first, "short"); @@ -2049,7 +2416,7 @@ TEST_F(MasterServiceTest, GetReplicaListByRegexComplex) { // Test 3.9: Initial test for non-existent key (as a sanity check) { auto get_result = service_->GetReplicaListByRegex( - ".*absolutely_non_existent.*", "default"); + ".*absolutely_non_existent.*", TenantId::Default()); // Depending on implementation, this could return an empty map or an // error. Let's assume it returns an empty map for a valid regex with no // matches. @@ -2062,7 +2429,8 @@ TEST_F(MasterServiceTest, GetReplicaList) { std::unique_ptr service_(new MasterService()); const UUID client_id = generate_uuid(); // Test getting non-existent key - auto get_result = service_->GetReplicaList("non_existent", "default"); + auto get_result = + service_->GetReplicaList("non_existent", TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); @@ -2072,15 +2440,15 @@ TEST_F(MasterServiceTest, GetReplicaList) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test getting existing key - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result2.has_value()); auto replica_list_local = get_result2.value().replicas; EXPECT_FALSE(replica_list_local.empty()); @@ -2095,24 +2463,24 @@ TEST_F(MasterServiceTest, RemoveObject) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test removing the object - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result.has_value()); // Verify object is removed - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); // Test removing non-existent object - auto remove_result2 = service_->Remove("non_existent", "default"); + auto remove_result2 = service_->Remove("non_existent", TenantId::Default()); EXPECT_FALSE(remove_result2.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, remove_result2.error()); } @@ -2130,19 +2498,19 @@ TEST_F(MasterServiceTest, RandomRemoveObject) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test removing the object - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result.has_value()); // Verify object is removed - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -2162,24 +2530,24 @@ TEST_F(MasterServiceTest, RemoveByRegex) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); } // wait for all the lease to expire std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto res = service_->RemoveByRegex("^test_key", "default"); + auto res = service_->RemoveByRegex("^test_key", TenantId::Default()); ASSERT_TRUE(res.has_value()); ASSERT_EQ(10, res.value()); times = 10; while (times--) { std::string key = "test_key" + std::to_string(times); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()); } @@ -2206,8 +2574,9 @@ TEST_F(MasterServiceTest, CopyStart) { UUID client_id = generate_uuid(); // Test Case 1: CopyStart a non-existent key, should fail. - auto copy_result = service_->CopyStart( - client_id, "non_existent_key", "default", "segment_1", {"segment_2"}); + auto copy_result = + service_->CopyStart(client_id, "non_existent_key", TenantId::Default(), + "segment_1", {"segment_2"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_result.error()); @@ -2219,25 +2588,25 @@ TEST_F(MasterServiceTest, CopyStart) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); // Test Case 2: CopyStart to segment_2 and segment_3, should fail because // the only replica is not completed. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_2", "segment_3"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_2", "segment_3"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, copy_result.error()); // PutEnd the object. - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 3: CopyStart to segment_2 and segment_3, should success. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_2", "segment_3"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_2", "segment_3"}); EXPECT_TRUE(copy_result.has_value()); auto copy_response = copy_result.value(); EXPECT_EQ("segment_1", copy_response.source.get_memory_descriptor() @@ -2245,42 +2614,44 @@ TEST_F(MasterServiceTest, CopyStart) { EXPECT_EQ(2, copy_response.targets.size()); // Test Case 4: Try remove the object, should fail because it is copying. - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error()); // Test Case 5: CopyStart to segment_4, should fail because there is an // ongoing copy task. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_4"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_4"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, copy_result.error()); // Test Case 6: CopyEnd, should success and the object now has 3 replicas. - auto copy_end_result = service_->CopyEnd(client_id, key, "default"); + auto copy_end_result = + service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_end_result.has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(3, get_result.value().replicas.size()); // Test Case 7: Copy from a non-existent replica to segment_3 and // segment_4, should fail. copy_result = - service_->CopyStart(client_id, key, "default", "non_existent_segment", - {"segment_3", "segment_4"}); + service_->CopyStart(client_id, key, TenantId::Default(), + "non_existent_segment", {"segment_3", "segment_4"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, copy_result.error()); // Test Case 8: Copy to segment_4 and a non-existent segment, should fail. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_4", "non_existent_segment"}); + copy_result = + service_->CopyStart(client_id, key, TenantId::Default(), "segment_1", + {"segment_4", "non_existent_segment"}); EXPECT_FALSE(copy_result.has_value()); EXPECT_EQ(ErrorCode::SEGMENT_NOT_FOUND, copy_result.error()); // Test Case 9: Copy to segment_3 and segment_4, should skip segment_3 and // successfully copy to segment_4. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_3", "segment_4"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_3", "segment_4"}); EXPECT_TRUE(copy_result.has_value()); copy_response = copy_result.value(); EXPECT_EQ("segment_1", copy_response.source.get_memory_descriptor() @@ -2292,16 +2663,16 @@ TEST_F(MasterServiceTest, CopyStart) { .buffer_descriptor.transport_endpoint_); // End the copy operation to clean up state - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_end_result.has_value()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(4, get_result.value().replicas.size()); // Test Case 10: Copy to segment_4 again, should skip because it's already // used. - copy_result = service_->CopyStart(client_id, key, "default", "segment_1", - {"segment_4"}); + copy_result = service_->CopyStart(client_id, key, TenantId::Default(), + "segment_1", {"segment_4"}); EXPECT_TRUE(copy_result.has_value()); copy_response = copy_result.value(); EXPECT_EQ("segment_1", copy_response.source.get_memory_descriptor() @@ -2314,19 +2685,19 @@ TEST_F(MasterServiceTest, CopyStart) { std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2)); // Test Case 11: Try remove the object, should fail because it is copying. - remove_result = service_->Remove(key, "default"); + remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, remove_result.error()); // Clean up the copy operation - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_end_result.has_value()); // Wait for the lease to expire std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2)); // Test Case 12: Try remove the object, should success. - remove_result = service_->Remove(key, "default"); + remove_result = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result.has_value()); } @@ -2347,7 +2718,7 @@ TEST_F(MasterServiceTest, CopyEnd) { // Test Case 1: CopyEnd a non-existent key, should fail. auto copy_end_result = - service_->CopyEnd(client_id, "non_existent_key", "default"); + service_->CopyEnd(client_id, "non_existent_key", TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_end_result.error()); @@ -2358,46 +2729,48 @@ TEST_F(MasterServiceTest, CopyEnd) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 2: CopyEnd the object, should fail because there is no ongoing // copy task. - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, copy_end_result.error()); // CopyStart the object to segment_2 - auto copy_start_result = service_->CopyStart(client_id, key, "default", - "segment_1", {"segment_2"}); + auto copy_start_result = service_->CopyStart( + client_id, key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_start_result.has_value()); // Test Case 3: CopyEnd with an invalid client id, should fail. - copy_end_result = service_->CopyEnd(invalid_client_id, key, "default"); + copy_end_result = + service_->CopyEnd(invalid_client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, copy_end_result.error()); // Test Case 4: MoveEnd the object, should fail because the ongoing task is // Copy. - auto move_end_result = service_->MoveEnd(client_id, key, "default"); + auto move_end_result = + service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_end_result.error()); // Test Case 5: CopyEnd, should success. - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_end_result.has_value()); // Verify we now have 2 replicas - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(2, get_result.value().replicas.size()); // CopyStart the object from segment_1 to segment_3, then unmount segment_1 - copy_start_result = service_->CopyStart(client_id, key, "default", + copy_start_result = service_->CopyStart(client_id, key, TenantId::Default(), "segment_1", {"segment_3"}); ASSERT_TRUE(copy_start_result.has_value()); @@ -2408,10 +2781,10 @@ TEST_F(MasterServiceTest, CopyEnd) { // Test Case 6: CopyEnd, should fail because the source is gone, the object // should have only 1 replica from segment_2. - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, copy_end_result.error()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); auto& replicas = get_result.value().replicas; EXPECT_EQ(1, replicas.size()); @@ -2420,7 +2793,7 @@ TEST_F(MasterServiceTest, CopyEnd) { .buffer_descriptor.transport_endpoint_); // CopyStart the object from segment_2 to segment_3, then unmount segment_3 - copy_start_result = service_->CopyStart(client_id, key, "default", + copy_start_result = service_->CopyStart(client_id, key, TenantId::Default(), "segment_2", {"segment_3"}); ASSERT_TRUE(copy_start_result.has_value()); @@ -2431,10 +2804,10 @@ TEST_F(MasterServiceTest, CopyEnd) { // Test Case 7: CopyEnd, should fail because the target is gone, the object // should have only 1 replica from segment_2. - copy_end_result = service_->CopyEnd(client_id, key, "default"); + copy_end_result = service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, copy_end_result.error()); - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); replicas = get_result.value().replicas; EXPECT_EQ(1, replicas.size()); @@ -2457,8 +2830,8 @@ TEST_F(MasterServiceTest, CopyRevoke) { UUID invalid_client_id = generate_uuid(); // Test Case 1: CopyRevoke a non-existent key, should fail. - auto copy_revoke_result = - service_->CopyRevoke(client_id, "non_existent_key", "default"); + auto copy_revoke_result = service_->CopyRevoke( + client_id, "non_existent_key", TenantId::Default()); EXPECT_FALSE(copy_revoke_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_revoke_result.error()); @@ -2469,49 +2842,52 @@ TEST_F(MasterServiceTest, CopyRevoke) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 2: CopyRevoke the object, should fail because there is no // ongoing copy task. - copy_revoke_result = service_->CopyRevoke(client_id, key, "default"); + copy_revoke_result = + service_->CopyRevoke(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_revoke_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, copy_revoke_result.error()); // CopyStart the object to segment_2 - auto copy_start_result = service_->CopyStart(client_id, key, "default", - "segment_1", {"segment_2"}); + auto copy_start_result = service_->CopyStart( + client_id, key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_start_result.has_value()); // Test Case 3: CopyRevoke with an invalid client id, should fail. copy_revoke_result = - service_->CopyRevoke(invalid_client_id, key, "default"); + service_->CopyRevoke(invalid_client_id, key, TenantId::Default()); EXPECT_FALSE(copy_revoke_result.has_value()); EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, copy_revoke_result.error()); // Test Case 4: MoveRevoke the object, should fail because the ongoing task // is Copy. - auto move_revoke_result = service_->MoveRevoke(client_id, key, "default"); + auto move_revoke_result = + service_->MoveRevoke(client_id, key, TenantId::Default()); EXPECT_FALSE(move_revoke_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_revoke_result.error()); // Test Case 5: CopyRevoke, should success. - copy_revoke_result = service_->CopyRevoke(client_id, key, "default"); + copy_revoke_result = + service_->CopyRevoke(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_revoke_result.has_value()); // Verify we still have 1 replica (the copy was revoked) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); // CopyStart the object from segment_1 to segment_2 again, then unmount // segment_1 - copy_start_result = service_->CopyStart(client_id, key, "default", + copy_start_result = service_->CopyStart(client_id, key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_start_result.has_value()); @@ -2522,11 +2898,12 @@ TEST_F(MasterServiceTest, CopyRevoke) { // Test Case 6: CopyRevoke, should success even though the source is gone, // the object should be erased too. - copy_revoke_result = service_->CopyRevoke(client_id, key, "default"); + copy_revoke_result = + service_->CopyRevoke(client_id, key, TenantId::Default()); EXPECT_TRUE(copy_revoke_result.has_value()); // Verify the object has been removed. - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); } @@ -2549,8 +2926,9 @@ TEST_F(MasterServiceTest, MoveStart) { UUID client_id = generate_uuid(); // Test Case 1: MoveStart a non-existent key, should fail. - auto move_start_result = service_->MoveStart( - client_id, "non_existent_key", "default", "segment_1", "segment_2"); + auto move_start_result = + service_->MoveStart(client_id, "non_existent_key", TenantId::Default(), + "segment_1", "segment_2"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_start_result.error()); @@ -2561,37 +2939,38 @@ TEST_F(MasterServiceTest, MoveStart) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); // Test Case 2: MoveStart the object, should fail because the only replica // is not completed. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_2"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, move_start_result.error()); // PutEnd the object. - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Copy the object to segment_3. - auto copy_start_result = service_->CopyStart(client_id, key, "default", - "segment_1", {"segment_3"}); + auto copy_start_result = service_->CopyStart( + client_id, key, TenantId::Default(), "segment_1", {"segment_3"}); ASSERT_TRUE(copy_start_result.has_value()); - auto copy_end_result = service_->CopyEnd(client_id, key, "default"); + auto copy_end_result = + service_->CopyEnd(client_id, key, TenantId::Default()); ASSERT_TRUE(copy_end_result.has_value()); // Test Case 3: MoveStart with source and target be the same, should fail. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_1"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(move_start_result.error(), ErrorCode::INVALID_PARAMS); // Test Case 4: MoveStart to segment_2, should succeed. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_2"); EXPECT_TRUE(move_start_result.has_value()); auto move_response = move_start_result.value(); @@ -2603,13 +2982,13 @@ TEST_F(MasterServiceTest, MoveStart) { .buffer_descriptor.transport_endpoint_); // Test Case 5: Try remove the object, should fail because it is moving. - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, remove_result.error()); // Test Case 6: MoveStart again, should fail because there is an ongoing // move task. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_3"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, @@ -2617,29 +2996,32 @@ TEST_F(MasterServiceTest, MoveStart) { // Test Case 7: MoveEnd, should succeed and the object now has 2 replicas // from segment_2 and segment_3 - auto move_end_result = service_->MoveEnd(client_id, key, "default"); + auto move_end_result = + service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(move_end_result.has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); auto& replicas = get_result.value().replicas; EXPECT_EQ(2, replicas.size()); // Test Case 8: Move from a non-existent replica to segment_1, should fail. - move_start_result = service_->MoveStart( - client_id, key, "default", "non_existent_segment", "segment_1"); + move_start_result = + service_->MoveStart(client_id, key, TenantId::Default(), + "non_existent_segment", "segment_1"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_NOT_FOUND, move_start_result.error()); // Test Case 8.5: Move to a non-existent target segment, should fail. - move_start_result = service_->MoveStart( - client_id, key, "default", "segment_2", "non_existent_segment"); + move_start_result = + service_->MoveStart(client_id, key, TenantId::Default(), "segment_2", + "non_existent_segment"); EXPECT_FALSE(move_start_result.has_value()); EXPECT_EQ(ErrorCode::SEGMENT_NOT_FOUND, move_start_result.error()); // Test Case 9: Move to an already existing segment, should succeed but // return nullopt. - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_2", "segment_3"); EXPECT_TRUE(move_start_result.has_value()); move_response = move_start_result.value(); @@ -2649,16 +3031,16 @@ TEST_F(MasterServiceTest, MoveStart) { // Test Case 10: Try remove the object, should fail because it is moving. std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2)); - remove_result = service_->Remove(key, "default"); + remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, remove_result.error()); // End the move. - move_end_result = service_->MoveEnd(client_id, key, "default"); + move_end_result = service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(move_end_result.has_value()); // Now the object should have only 1 replica on segment_3. - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); replicas = get_result.value().replicas; EXPECT_EQ(1, replicas.size()); @@ -2668,7 +3050,7 @@ TEST_F(MasterServiceTest, MoveStart) { // Test Case 11: Try remove the object, should succeed after lease expires. std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl * 2)); - remove_result = service_->Remove(key, "default"); + remove_result = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result.has_value()); } @@ -2687,7 +3069,7 @@ TEST_F(MasterServiceTest, MoveEnd) { // Test Case 1: MoveEnd a non-existent key, should fail. auto move_end_result = - service_->MoveEnd(client_id, "non_existent_key", "default"); + service_->MoveEnd(client_id, "non_existent_key", TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_end_result.error()); @@ -2698,47 +3080,49 @@ TEST_F(MasterServiceTest, MoveEnd) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 2: MoveEnd the object, should fail because there is no ongoing // move task. - move_end_result = service_->MoveEnd(client_id, key, "default"); + move_end_result = service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, move_end_result.error()); // MoveStart the object to segment_2 - auto move_start_result = service_->MoveStart(client_id, key, "default", - "segment_1", "segment_2"); + auto move_start_result = service_->MoveStart( + client_id, key, TenantId::Default(), "segment_1", "segment_2"); ASSERT_TRUE(move_start_result.has_value()); // Test Case 3: MoveEnd with an invalid client id, should fail. - move_end_result = service_->MoveEnd(invalid_client_id, key, "default"); + move_end_result = + service_->MoveEnd(invalid_client_id, key, TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, move_end_result.error()); // Test Case 4: CopyEnd the object, should fail because the ongoing task is // Move. - auto copy_end_result = service_->CopyEnd(client_id, key, "default"); + auto copy_end_result = + service_->CopyEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_end_result.error()); // Test Case 5: MoveEnd, should success. - move_end_result = service_->MoveEnd(client_id, key, "default"); + move_end_result = service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_TRUE(move_end_result.has_value()); // Verify we still have 1 replica (the move was successful) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); EXPECT_EQ(1, get_result.value().replicas.size()); // MoveStart the object from segment_2 to segment_1 again, then unmount // segment_2 - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_2", "segment_1"); ASSERT_TRUE(move_start_result.has_value()); @@ -2748,7 +3132,7 @@ TEST_F(MasterServiceTest, MoveEnd) { ASSERT_TRUE(unmount_result.has_value()); // Test Case 6: MoveEnd, should fail because the source is gone. - move_end_result = service_->MoveEnd(client_id, key, "default"); + move_end_result = service_->MoveEnd(client_id, key, TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_GONE, move_end_result.error()); } @@ -2766,8 +3150,8 @@ TEST_F(MasterServiceTest, MoveRevoke) { UUID invalid_client_id = generate_uuid(); // Test Case 1: MoveRevoke a non-existent key, should fail. - auto move_revoke_result = - service_->MoveRevoke(client_id, "non_existent_key", "default"); + auto move_revoke_result = service_->MoveRevoke( + client_id, "non_existent_key", TenantId::Default()); EXPECT_FALSE(move_revoke_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_revoke_result.error()); @@ -2778,43 +3162,46 @@ TEST_F(MasterServiceTest, MoveRevoke) { config.replica_num = 1; config.preferred_segment = "segment_1"; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test Case 2: MoveRevoke the object, should fail because there is no // ongoing move task. - move_revoke_result = service_->MoveRevoke(client_id, key, "default"); + move_revoke_result = + service_->MoveRevoke(client_id, key, TenantId::Default()); EXPECT_FALSE(move_revoke_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NO_REPLICATION_TASK, move_revoke_result.error()); // MoveStart the object from segment_1 to segment_2 - auto move_start_result = service_->MoveStart(client_id, key, "default", - "segment_1", "segment_2"); + auto move_start_result = service_->MoveStart( + client_id, key, TenantId::Default(), "segment_1", "segment_2"); ASSERT_TRUE(move_start_result.has_value()); // Test Case 3: MoveRevoke with an invalid client id, should fail. move_revoke_result = - service_->MoveRevoke(invalid_client_id, key, "default"); + service_->MoveRevoke(invalid_client_id, key, TenantId::Default()); EXPECT_FALSE(move_revoke_result.has_value()); EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, move_revoke_result.error()); // Test Case 4: CopyRevoke the object, should fail because the ongoing task // is Move. - auto copy_revoke_result = service_->CopyRevoke(client_id, key, "default"); + auto copy_revoke_result = + service_->CopyRevoke(client_id, key, TenantId::Default()); EXPECT_FALSE(copy_revoke_result.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_revoke_result.error()); // Test Case 5: MoveRevoke, should succeed. - move_revoke_result = service_->MoveRevoke(client_id, key, "default"); + move_revoke_result = + service_->MoveRevoke(client_id, key, TenantId::Default()); EXPECT_TRUE(move_revoke_result.has_value()); // Verify we still have 1 replica (the move was revoked) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); auto& replicas = get_result.value().replicas; EXPECT_EQ(1, replicas.size()); @@ -2824,7 +3211,7 @@ TEST_F(MasterServiceTest, MoveRevoke) { // MoveStart the object from segment_1 to segment_2 again, then unmount // segment_1 - move_start_result = service_->MoveStart(client_id, key, "default", + move_start_result = service_->MoveStart(client_id, key, TenantId::Default(), "segment_1", "segment_2"); ASSERT_TRUE(move_start_result.has_value()); @@ -2834,11 +3221,12 @@ TEST_F(MasterServiceTest, MoveRevoke) { ASSERT_TRUE(unmount_result.has_value()); // Test Case 6: MoveRevoke, should succeed even though the source is gone. - move_revoke_result = service_->MoveRevoke(client_id, key, "default"); + move_revoke_result = + service_->MoveRevoke(client_id, key, TenantId::Default()); EXPECT_TRUE(move_revoke_result.has_value()); // The object should be erased as there is no replica left. - get_result = service_->GetReplicaList(key, "default"); + get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); } @@ -2870,38 +3258,38 @@ TEST_F(MasterServiceTest, ProtectCopyMoveSourceFromEviction) { config.preferred_segment = "segment_1"; // Put two objects for move and copy tests. - auto put_start_result = service_->PutStart(client_id, copy_key, "default", - slice_length, config); + auto put_start_result = service_->PutStart( + client_id, copy_key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, copy_key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, copy_key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - put_start_result = service_->PutStart(client_id, move_key, "default", - slice_length, config); + put_start_result = service_->PutStart( + client_id, move_key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - put_end_result = - service_->PutEnd(client_id, move_key, "default", ReplicaType::MEMORY); + put_end_result = service_->PutEnd(client_id, move_key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Start copy and move operations. - auto copy_start_result = service_->CopyStart(client_id, copy_key, "default", - "segment_1", {"segment_2"}); + auto copy_start_result = service_->CopyStart( + client_id, copy_key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_start_result.has_value()); - auto move_start_result = service_->MoveStart(client_id, move_key, "default", - "segment_1", "segment_2"); + auto move_start_result = service_->MoveStart( + client_id, move_key, TenantId::Default(), "segment_1", "segment_2"); ASSERT_TRUE(move_start_result.has_value()); // Put more objects to trigger eviction. Do not prefer any segments. config.preferred_segment = ""; for (size_t i = 0; i < 128 * (kSegmentSize * 2 / slice_length); ++i) { std::string key = "test_key_" + std::to_string(i); - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { - auto put_end_result = service_->PutEnd(client_id, key, "default", - ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } else { // wait for eviction to work @@ -2915,10 +3303,12 @@ TEST_F(MasterServiceTest, ProtectCopyMoveSourceFromEviction) { ASSERT_TRUE(remove_all_result > 0); // Try end copy and move operations, should success. - auto copy_end_result = service_->CopyEnd(client_id, copy_key, "default"); + auto copy_end_result = + service_->CopyEnd(client_id, copy_key, TenantId::Default()); EXPECT_TRUE(copy_end_result.has_value()); - auto move_end_result = service_->MoveEnd(client_id, move_key, "default"); + auto move_end_result = + service_->MoveEnd(client_id, move_key, TenantId::Default()); EXPECT_TRUE(move_end_result.has_value()); } @@ -2955,27 +3345,27 @@ TEST_F(MasterServiceTest, DiscardTimeoutCopyMove) { config.preferred_segment = "segment_1"; // Put two objects for move and copy tests. - auto put_start_result = service_->PutStart(client_id, copy_key, "default", - slice_length, config); + auto put_start_result = service_->PutStart( + client_id, copy_key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, copy_key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, copy_key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - put_start_result = service_->PutStart(client_id, move_key, "default", - slice_length, config); + put_start_result = service_->PutStart( + client_id, move_key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - put_end_result = - service_->PutEnd(client_id, move_key, "default", ReplicaType::MEMORY); + put_end_result = service_->PutEnd(client_id, move_key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Start copy and move operations. - auto copy_start_result = service_->CopyStart(client_id, copy_key, "default", - "segment_1", {"segment_2"}); + auto copy_start_result = service_->CopyStart( + client_id, copy_key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_start_result.has_value()); - auto move_start_result = service_->MoveStart(client_id, move_key, "default", - "segment_1", "segment_2"); + auto move_start_result = service_->MoveStart( + client_id, move_key, TenantId::Default(), "segment_1", "segment_2"); ASSERT_TRUE(move_start_result.has_value()); // Wait for the operations timeout. @@ -2985,11 +3375,11 @@ TEST_F(MasterServiceTest, DiscardTimeoutCopyMove) { config.preferred_segment = ""; for (size_t i = 0; i < 128 * (kSegmentSize * 2 / slice_length); ++i) { std::string key = "test_key_" + std::to_string(i); - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { - auto put_end_result = service_->PutEnd(client_id, key, "default", - ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } else { // wait for eviction to work @@ -2999,11 +3389,13 @@ TEST_F(MasterServiceTest, DiscardTimeoutCopyMove) { // Try end copy and move operations, should fail because the objects are // evicted. - auto copy_end_result = service_->CopyEnd(client_id, copy_key, "default"); + auto copy_end_result = + service_->CopyEnd(client_id, copy_key, TenantId::Default()); EXPECT_FALSE(copy_end_result.has_value()); EXPECT_EQ(copy_end_result.error(), ErrorCode::OBJECT_NOT_FOUND); - auto move_end_result = service_->MoveEnd(client_id, move_key, "default"); + auto move_end_result = + service_->MoveEnd(client_id, move_key, TenantId::Default()); EXPECT_FALSE(move_end_result.has_value()); EXPECT_EQ(move_end_result.error(), ErrorCode::OBJECT_NOT_FOUND); } @@ -3049,7 +3441,8 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) { populate_store(); // Action: Remove keys starting with "test_key_" - auto remove_result = service_->RemoveByRegex("^test_key_", "default"); + auto remove_result = + service_->RemoveByRegex("^test_key_", TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); EXPECT_EQ(remove_result.value(), 3); // Should remove 3 keys @@ -3057,7 +3450,7 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) { std::vector deleted_keys = {"test_key_01", "test_key_02", "test_key_10"}; for (const auto& key : deleted_keys) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); EXPECT_FALSE(exist_result.value()) << "Key " << key << " should have been deleted."; @@ -3066,7 +3459,7 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) { std::vector remaining_keys = { "prod_key_alpha", "short", "test-key-extra"}; // Sample a few for (const auto& key : remaining_keys) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); EXPECT_TRUE(exist_result.value()) << "Key " << key << " should NOT have been deleted."; @@ -3089,12 +3482,13 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) { size_t total_keys = 13; // Count from the keys_to_put vector // Action: Remove all keys - auto remove_result = service_->RemoveByRegex(".*", "default"); + auto remove_result = service_->RemoveByRegex(".*", TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); EXPECT_EQ(remove_result.value(), total_keys); // Verification: Check that no keys remain - auto get_all_result = service_->GetReplicaListByRegex(".*", "default"); + auto get_all_result = + service_->GetReplicaListByRegex(".*", TenantId::Default()); ASSERT_TRUE(get_all_result.has_value()); EXPECT_TRUE(get_all_result.value().empty()); } @@ -3114,13 +3508,14 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) { size_t total_keys_before_remove = 13; // Action: Attempt to remove using a pattern that matches nothing - auto remove_result = - service_->RemoveByRegex("^nonexistent-pattern-", "default"); + auto remove_result = service_->RemoveByRegex("^nonexistent-pattern-", + TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); EXPECT_EQ(remove_result.value(), 0); // Should remove 0 keys // Verification: Check that all keys still exist - auto get_all_result = service_->GetReplicaListByRegex(".*", "default"); + auto get_all_result = + service_->GetReplicaListByRegex(".*", TenantId::Default()); ASSERT_TRUE(get_all_result.has_value()); EXPECT_EQ(get_all_result.value().size(), total_keys_before_remove); } @@ -3138,7 +3533,8 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) { populate_store(); // Action: Remove all keys that contain a slash '/' OR end with a number - auto remove_result = service_->RemoveByRegex("/|\\d$", "default"); + auto remove_result = + service_->RemoveByRegex("/|\\d$", TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); // Matches: "config/user/settings.json", "logs/app-2025-08-13.log", // "test_key_01", "test_key_02", "test_key_10" @@ -3163,7 +3559,8 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) { populate_store(); // Action: Remove all keys that contain "chunk" OR "config" - auto remove_result = service_->RemoveByRegex("chunk|config", "default"); + auto remove_result = + service_->RemoveByRegex("chunk|config", TenantId::Default()); ASSERT_TRUE(remove_result.has_value()); // Matches: "data_part_1_chunk_a", "data_part_2_chunk_b", // "config/user/settings.json" @@ -3171,17 +3568,17 @@ TEST_F(MasterServiceTest, RemoveByRegexComplex) { // Verification auto exist_result_chunk = - service_->ExistKey("data_part_1_chunk_a", "default"); + service_->ExistKey("data_part_1_chunk_a", TenantId::Default()); ASSERT_TRUE(exist_result_chunk.has_value()); EXPECT_FALSE(exist_result_chunk.value()); - auto exist_result_config = - service_->ExistKey("config/user/settings.json", "default"); + auto exist_result_config = service_->ExistKey( + "config/user/settings.json", TenantId::Default()); ASSERT_TRUE(exist_result_config.has_value()); EXPECT_FALSE(exist_result_config.value()); auto exist_result_untouched = - service_->ExistKey("prod_key_alpha", "default"); + service_->ExistKey("prod_key_alpha", TenantId::Default()); ASSERT_TRUE(exist_result_untouched.has_value()); EXPECT_TRUE(exist_result_untouched.value()); } @@ -3201,13 +3598,13 @@ TEST_F(MasterServiceTest, RemoveAll) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); } // wait for all the lease to expire @@ -3216,7 +3613,7 @@ TEST_F(MasterServiceTest, RemoveAll) { times = 10; while (times--) { std::string key = "test_key" + std::to_string(times); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()); } @@ -3250,8 +3647,8 @@ TEST_F(MasterServiceTest, SingleSliceMultiReplicaFlow) { std::vector replica_list; // Test PutStart with multiple slices and replicas - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); @@ -3267,17 +3664,17 @@ TEST_F(MasterServiceTest, SingleSliceMultiReplicaFlow) { } // Test GetReplicaList during processing (should fail) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); // Complete the put operation - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Test GetReplicaList after completion - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result2.has_value()); auto retrieved_replicas = get_result2.value().replicas; ASSERT_EQ(num_replicas, retrieved_replicas.size()); @@ -3310,15 +3707,15 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) { config.replica_num = 1; // One replica // Create the object - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Verify object exists - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); auto retrieved_replicas = get_result.value().replicas; ASSERT_EQ(1, retrieved_replicas.size()); @@ -3329,7 +3726,7 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) { // Try to get the object - it should be automatically removed since the // replica is invalid - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result2.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result2.error()); @@ -3339,15 +3736,15 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) { // Create another object std::string key2 = "another_segment_object"; - auto put_start_result2 = - service_->PutStart(client_id, key2, "default", slice_length, config); + auto put_start_result2 = service_->PutStart( + client_id, key2, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result2.has_value()); - auto put_end_result2 = - service_->PutEnd(client_id, key2, "default", ReplicaType::MEMORY); + auto put_end_result2 = service_->PutEnd( + client_id, key2, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result2.has_value()); // Verify we can get it - auto get_result3 = service_->GetReplicaList(key2, "default"); + auto get_result3 = service_->GetReplicaList(key2, TenantId::Default()); ASSERT_TRUE(get_result3.has_value()); // Unmount the segment @@ -3355,7 +3752,7 @@ TEST_F(MasterServiceTest, CleanupStaleHandlesTest) { ASSERT_TRUE(unmount_result2.has_value()); // Try to remove the object that should already be cleaned up - auto remove_result = service_->Remove(key2, "default"); + auto remove_result = service_->Remove(key2, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, remove_result.error()); } @@ -3388,10 +3785,11 @@ TEST_F(MasterServiceTest, ConcurrentWriteAndRemoveAll) { std::vector replica_list; auto put_start_result = service_->PutStart( - client_id, key, "default", slice_length, config); + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { - auto put_end_result = service_->PutEnd( - client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = + service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); if (put_end_result.has_value()) { success_writes++; } @@ -3456,11 +3854,11 @@ TEST_F(MasterServiceTest, ConcurrentReadAndRemoveAll) { ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } @@ -3473,7 +3871,8 @@ TEST_F(MasterServiceTest, ConcurrentReadAndRemoveAll) { readers.emplace_back([&]() { for (int j = 0; j < num_objects; ++j) { std::string key = "pre_key_" + std::to_string(j); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = + service_->GetReplicaList(key, TenantId::Default()); if (get_result.has_value()) { success_reads++; } @@ -3514,7 +3913,7 @@ TEST_F(MasterServiceTest, ConcurrentReadAndRemoveAll) { // Verify all objects were removed for (int i = 0; i < num_objects; ++i) { std::string key = "pre_key_" + std::to_string(i); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -3537,11 +3936,11 @@ TEST_F(MasterServiceTest, ConcurrentRemoveAllOperations) { ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } @@ -3568,7 +3967,7 @@ TEST_F(MasterServiceTest, ConcurrentRemoveAllOperations) { // Verify all objects were removed for (int i = 0; i < num_objects; ++i) { std::string key = "pre_key_" + std::to_string(i); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -3605,23 +4004,23 @@ TEST_F(MasterServiceTest, UnmountSegmentImmediateCleanup) { // Umount will remove all objects in the segment, include the key1 ASSERT_EQ(1, service_->GetKeyCount()); // Verify objects in segment1 is gone - auto get_result1 = service_->GetReplicaList(key1, "default"); + auto get_result1 = service_->GetReplicaList(key1, TenantId::Default()); ASSERT_FALSE(get_result1.has_value()); ASSERT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result1.error()); // Verify objects in segment2 is still there - auto get_result2 = service_->GetReplicaList(key2, "default"); + auto get_result2 = service_->GetReplicaList(key2, TenantId::Default()); ASSERT_TRUE(get_result2.has_value()); // Verify put key1 will put into segment2 rather than segment1 - auto put_start_result = - service_->PutStart(client_id, key1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key1, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); - auto put_end_result = - service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key1, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto get_result3 = service_->GetReplicaList(key1, "default"); + auto get_result3 = service_->GetReplicaList(key1, TenantId::Default()); ASSERT_TRUE(get_result3.has_value()); auto retrieved = get_result3.value(); ASSERT_EQ(replica_list[0] @@ -3653,15 +4052,17 @@ TEST_F(MasterServiceTest, ReadableAfterPartialUnmountWithReplication) { ReplicateConfig config; config.replica_num = 2; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); ASSERT_EQ(2u, put_start_result->size()); - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); // Verify two replicas exist and they are on distinct segments - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); auto replicas = get_result.value().replicas; ASSERT_EQ(2u, replicas.size()); @@ -3679,11 +4080,48 @@ TEST_F(MasterServiceTest, ReadableAfterPartialUnmountWithReplication) { ASSERT_TRUE(service_->UnmountSegment(segment1.id, client_id).has_value()); // Key should still be readable via the remaining replica - auto get_after_unmount = service_->GetReplicaList(key, "default"); + auto get_after_unmount = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_after_unmount.has_value()) << "Object should remain accessible with surviving replica"; } +TEST_F(MasterServiceTest, PutStartPartialAllocationIsObservable) { + std::unique_ptr service_(new MasterService()); + + // Mount two segments only + constexpr size_t buffer1 = 0x300000000; + constexpr size_t buffer2 = 0x400000000; + constexpr size_t segment_size = 1024 * 1024 * 64; // 64MB + + auto segment1 = MakeSegment("segment1", buffer1, segment_size); + auto segment2 = MakeSegment("segment2", buffer2, segment_size); + UUID client_id = generate_uuid(); + ASSERT_TRUE(service_->MountSegment(segment1, client_id).has_value()); + ASSERT_TRUE(service_->MountSegment(segment2, client_id).has_value()); + + auto& metrics = MasterMetricManager::instance(); + const int64_t partial_before = metrics.get_put_start_partial_allocations(); + + // Request more replicas than available segments: best-effort keeps the + // put successful but the degradation must be recorded. + ReplicateConfig config; + config.replica_num = 3; + auto put_start_result = service_->PutStart( + client_id, "partial_alloc_key", TenantId::Default(), 1024, config); + ASSERT_TRUE(put_start_result.has_value()); + ASSERT_EQ(2u, put_start_result->size()); + ASSERT_EQ(metrics.get_put_start_partial_allocations(), partial_before + 1); + + // A fully satisfied allocation must not be counted as partial. + ReplicateConfig full_config; + full_config.replica_num = 2; + auto full_result = service_->PutStart( + client_id, "full_alloc_key", TenantId::Default(), 1024, full_config); + ASSERT_TRUE(full_result.has_value()); + ASSERT_EQ(2u, full_result->size()); + ASSERT_EQ(metrics.get_put_start_partial_allocations(), partial_before + 1); +} + TEST_F(MasterServiceTest, UnmountSegmentPerformance) { std::unique_ptr service_(new MasterService()); constexpr size_t kBufferAddress = 0x300000000; @@ -3729,7 +4167,7 @@ TEST_F(MasterServiceTest, UnmountSegmentPerformance) { // Verify all keys are gone for (const auto& key : keys) { - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -3759,77 +4197,77 @@ TEST_F(MasterServiceTest, RemoveLeasedObject) { config.replica_num = 1; // Verify lease is granted on ExistsKey - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); - auto remove_result = service_->Remove(key, "default"); + auto remove_result = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result.error()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto remove_result2 = service_->Remove(key, "default"); + auto remove_result2 = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result2.has_value()); // Verify lease is extended on successive ExistsKey - auto put_start_result2 = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result2 = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result2.has_value()); - auto put_end_result2 = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result2 = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result2.has_value()); - auto exist_result2 = service_->ExistKey(key, "default"); + auto exist_result2 = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result2.has_value()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto exist_result3 = service_->ExistKey(key, "default"); + auto exist_result3 = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result3.has_value()); - auto remove_result3 = service_->Remove(key, "default"); + auto remove_result3 = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result3.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result3.error()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto remove_result4 = service_->Remove(key, "default"); + auto remove_result4 = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result4.has_value()); // Verify lease is granted on GetReplicaList - auto put_start_result3 = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result3 = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result3.has_value()); - auto put_end_result3 = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result3 = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result3.has_value()); - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); - auto remove_result5 = service_->Remove(key, "default"); + auto remove_result5 = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result5.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result5.error()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto remove_result6 = service_->Remove(key, "default"); + auto remove_result6 = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result6.has_value()); // Verify lease is extended on successive GetReplicaList - auto put_start_result4 = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result4 = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result4.has_value()); - auto put_end_result4 = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result4 = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result4.has_value()); - auto get_result2 = service_->GetReplicaList(key, "default"); + auto get_result2 = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result2.has_value()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto get_result3 = service_->GetReplicaList(key, "default"); + auto get_result3 = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result3.has_value()); - auto remove_result7 = service_->Remove(key, "default"); + auto remove_result7 = service_->Remove(key, TenantId::Default()); EXPECT_FALSE(remove_result7.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result7.error()); std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); - auto remove_result8 = service_->Remove(key, "default"); + auto remove_result8 = service_->Remove(key, TenantId::Default()); EXPECT_TRUE(remove_result8.has_value()); // Verify object is removed - auto get_result4 = service_->GetReplicaList(key, "default"); + auto get_result4 = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result4.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result4.error()); } @@ -3847,21 +4285,21 @@ TEST_F(MasterServiceTest, RemoveAllLeasedObject) { uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); if (i >= 5) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); } } ASSERT_EQ(5, service_->RemoveAll()); for (int i = 0; i < 5; ++i) { std::string key = "test_key" + std::to_string(i); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_FALSE(exist_result.value()); } // wait for all the lease to expire @@ -3869,7 +4307,7 @@ TEST_F(MasterServiceTest, RemoveAllLeasedObject) { ASSERT_EQ(5, service_->RemoveAll()); for (int i = 5; i < 10; ++i) { std::string key = "test_key" + std::to_string(i); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_FALSE(exist_result.value()); } } @@ -3899,11 +4337,11 @@ TEST_F(MasterServiceTest, EvictObject) { uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { - auto put_end_result = service_->PutEnd(client_id, key, "default", - ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); success_puts++; } else { @@ -3939,14 +4377,15 @@ TEST_F(MasterServiceTest, TryEvictLeasedObject) { uint64_t slice_length = object_size; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); if (put_start_result.has_value()) { - auto put_end_result = service_->PutEnd(client_id, key, "default", - ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // the object is leased - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = + service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); leased_keys.push_back(key); success_puts++; @@ -3960,7 +4399,7 @@ TEST_F(MasterServiceTest, TryEvictLeasedObject) { std::this_thread::sleep_for(std::chrono::milliseconds(50)); // All leased objects should be accessible for (const auto& key : leased_keys) { - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); } std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); @@ -3993,19 +4432,25 @@ TEST_F(MasterServiceTest, RemoveSoftPinObject) { config.with_soft_pin = true; // Verify soft pin does not block remove + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) + .has_value()); ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) .has_value()); - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); - EXPECT_TRUE(service_->Remove(key, "default").has_value()); + EXPECT_TRUE(service_->Remove(key, TenantId::Default()).has_value()); // Verify soft pin does not block RemoveAll + ASSERT_TRUE(service_ + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) + .has_value()); ASSERT_TRUE( - service_->PutStart(client_id, key, "default", slice_length, config) + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) .has_value()); - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); EXPECT_EQ(1, service_->RemoveAll()); } @@ -4043,13 +4488,13 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotEvictedBeforeOtherObjects) { soft_pin_config.with_soft_pin = true; ASSERT_TRUE(service_ - ->PutStart(client_id, pin_key, "default", + ->PutStart(client_id, pin_key, TenantId::Default(), slice_length, soft_pin_config) .has_value()); - ASSERT_TRUE( - service_ - ->PutEnd(client_id, pin_key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, pin_key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); } // Fill the segment to trigger eviction @@ -4060,12 +4505,13 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotEvictedBeforeOtherObjects) { ReplicateConfig config; config.replica_num = 1; if (service_ - ->PutStart(client_id, key, "default", slice_length, config) + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) .has_value()) { - ASSERT_TRUE( - service_ - ->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); } else { failed_puts++; } @@ -4077,8 +4523,8 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotEvictedBeforeOtherObjects) { // pin_key should still be accessible for (int i = 0; i < 2; i++) { std::string pin_key = "pin_key" + std::to_string(i); - ASSERT_TRUE( - service_->GetReplicaList(pin_key, "default").has_value()); + ASSERT_TRUE(service_->GetReplicaList(pin_key, TenantId::Default()) + .has_value()); } // wait for the lease to expire @@ -4117,11 +4563,14 @@ TEST_F(MasterServiceTest, SoftPinObjectsCanBeEvicted) { ReplicateConfig config; config.replica_num = 1; config.with_soft_pin = true; - if (service_->PutStart(client_id, key, "default", slice_length, config) + if (service_ + ->PutStart(client_id, key, TenantId::Default(), slice_length, + config) .has_value()) { - ASSERT_TRUE( - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); success_puts++; } else { // wait for eviction to work @@ -4170,12 +4619,13 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) { soft_pin_config.replica_num = 1; soft_pin_config.with_soft_pin = true; - ASSERT_TRUE(service_->PutStart(client_id, pin_key, "default", - slice_length, soft_pin_config)); - ASSERT_TRUE( - service_ - ->PutEnd(client_id, pin_key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_->PutStart(client_id, pin_key, + TenantId::Default(), slice_length, + soft_pin_config)); + ASSERT_TRUE(service_ + ->PutEnd(client_id, pin_key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); } // Wait for the soft pin to expire @@ -4184,8 +4634,8 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) { // Get the pin_key to extend the soft pin for (int i = 0; i < 2; i++) { std::string pin_key = "pin_key" + std::to_string(i); - ASSERT_TRUE( - service_->GetReplicaList(pin_key, "default").has_value()); + ASSERT_TRUE(service_->GetReplicaList(pin_key, TenantId::Default()) + .has_value()); } // Fill the segment to trigger eviction @@ -4196,12 +4646,13 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) { ReplicateConfig config; config.replica_num = 1; if (service_ - ->PutStart(client_id, key, "default", slice_length, config) + ->PutStart(client_id, key, TenantId::Default(), + slice_length, config) .has_value()) { - ASSERT_TRUE( - service_ - ->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); } else { failed_puts++; } @@ -4214,8 +4665,8 @@ TEST_F(MasterServiceTest, SoftPinExtendedOnGet) { // pin_key should still be accessible for (int i = 0; i < 2; i++) { std::string pin_key = "pin_key" + std::to_string(i); - ASSERT_TRUE( - service_->GetReplicaList(pin_key, "default").has_value()); + ASSERT_TRUE(service_->GetReplicaList(pin_key, TenantId::Default()) + .has_value()); } // wait for the lease to expire @@ -4256,11 +4707,14 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotAllowEvict) { ReplicateConfig config; config.replica_num = 1; config.with_soft_pin = true; - if (service_->PutStart(client_id, key, "default", slice_length, config) + if (service_ + ->PutStart(client_id, key, TenantId::Default(), slice_length, + config) .has_value()) { - ASSERT_TRUE( - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); success_keys.push_back(key); } else { // wait for eviction to work @@ -4270,7 +4724,8 @@ TEST_F(MasterServiceTest, SoftPinObjectsNotAllowEvict) { ASSERT_LE(success_keys.size(), 17); // All soft pinned objects should be accessible for (const auto& key : success_keys) { - ASSERT_TRUE(service_->GetReplicaList(key, "default").has_value()); + ASSERT_TRUE( + service_->GetReplicaList(key, TenantId::Default()).has_value()); } std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); service_->RemoveAll(); @@ -4295,8 +4750,8 @@ TEST_F(MasterServiceTest, ReplicaSegmentsAreUnique) { ReplicateConfig config; config.replica_num = 10; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto replica_list_local = put_start_result.value(); ASSERT_EQ(config.replica_num, replica_list_local.size()); @@ -4312,8 +4767,10 @@ TEST_F(MasterServiceTest, ReplicaSegmentsAreUnique) { EXPECT_EQ(segment_names.size(), config.replica_num) << "Duplicate segment found"; - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); } TEST_F(MasterServiceTest, ReplicationFactorTwoWithSingleSegment) { @@ -4333,8 +4790,8 @@ TEST_F(MasterServiceTest, ReplicationFactorTwoWithSingleSegment) { ReplicateConfig config; config.replica_num = 2; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); auto replicas = put_start_result.value(); @@ -4367,22 +4824,23 @@ TEST_F(MasterServiceTest, BatchExistKeyTest) { config.replica_num = 1; uint64_t slice_length = value_size; auto put_start_result = service_->PutStart( - client_id, test_keys[i], "default", slice_length, config); + client_id, test_keys[i], TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = service_->PutEnd(client_id, test_keys[i], - "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, test_keys[i], TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } // Test individual ExistKey calls to verify the underlying functionality for (int i = 0; i < test_object_num; ++i) { - auto exist_result = service_->ExistKey(test_keys[i], "default"); + auto exist_result = + service_->ExistKey(test_keys[i], TenantId::Default()); EXPECT_TRUE(exist_result.value()); } // Tets batch test_keys.push_back("non_existent_key"); - auto exist_resp = service_->BatchExistKey(test_keys, "default"); + auto exist_resp = service_->BatchExistKey(test_keys, TenantId::Default()); for (int i = 0; i < test_object_num; ++i) { ASSERT_TRUE(exist_resp[i].value()); } @@ -4413,16 +4871,17 @@ TEST_F(MasterServiceTest, BatchExistKeyGroupedAndIncompletePreservesOrder) { PutCompletedObject(*service_, client_id, completed_key, config); const std::string incomplete_key = "batch_incomplete_key"; - ASSERT_TRUE( - service_->PutStart(client_id, incomplete_key, "default", 1024, config) - .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, incomplete_key, TenantId::Default(), + 1024, config) + .has_value()); const std::string missing_key = "batch_missing_key"; std::vector keys = {grouped_key_a, completed_key, incomplete_key, missing_key, grouped_key_b}; - auto resp = service_->BatchExistKey(keys, "default"); + auto resp = service_->BatchExistKey(keys, TenantId::Default()); ASSERT_EQ(resp.size(), keys.size()); ASSERT_TRUE(resp[0].has_value()); ASSERT_TRUE(resp[1].has_value()); @@ -4437,7 +4896,9 @@ TEST_F(MasterServiceTest, BatchExistKeyGroupedAndIncompletePreservesOrder) { } TEST_F(MasterServiceTest, BatchExistKeyTenantAwarePreservesOrder) { - std::unique_ptr service_(new MasterService()); + const TenantId tenant_id("tenant_batch_exist"); + auto service_ = std::make_unique(MakeStrictTenantConfig( + {std::string(TenantId::kDefaultValue), tenant_id.value()})); const UUID client_id = generate_uuid(); constexpr size_t buffer = 0x300000000; @@ -4447,7 +4908,7 @@ TEST_F(MasterServiceTest, BatchExistKeyTenantAwarePreservesOrder) { ReplicateConfig config; config.replica_num = 1; - const std::string tenant_id = "tenant_batch_exist"; + const std::string tenant_only_key = "batch_tenant_only"; const std::string default_only_key = "batch_default_only"; const std::string incomplete_key = "batch_tenant_incomplete"; @@ -4472,16 +4933,17 @@ TEST_F(MasterServiceTest, BatchExistKeyTenantAwarePreservesOrder) { EXPECT_TRUE(tenant_resp[4].value()); std::vector default_keys = {tenant_only_key, default_only_key}; - auto default_resp = service_->BatchExistKey(default_keys, "default"); + auto default_resp = + service_->BatchExistKey(default_keys, TenantId::Default()); ASSERT_EQ(default_resp.size(), default_keys.size()); EXPECT_FALSE(default_resp[0].value()); EXPECT_TRUE(default_resp[1].value()); } TEST_F(MasterServiceTest, WrappedBatchExistKeyUsesTenantAwareBatchPath) { - WrappedMasterServiceConfig service_config; - service_config.default_kv_lease_ttl = 100; - service_config.enable_metric_reporting = false; + const TenantId tenant_id("wrapped_batch_exist_tenant"); + auto service_config = MakeStrictWrappedConfig( + {std::string(TenantId::kDefaultValue), tenant_id.value()}); WrappedMasterService service_(service_config); Segment segment = MakeSegment("wrapped_batch_exist_segment"); @@ -4490,7 +4952,6 @@ TEST_F(MasterServiceTest, WrappedBatchExistKeyUsesTenantAwareBatchPath) { ReplicateConfig config; config.replica_num = 1; - const std::string tenant_id = "wrapped_batch_exist_tenant"; const std::string tenant_key_a = "wrapped_batch_tenant_a"; const std::string tenant_key_b = "wrapped_batch_tenant_b"; const std::string default_only_key = "wrapped_batch_default_only"; @@ -4499,13 +4960,14 @@ TEST_F(MasterServiceTest, WrappedBatchExistKeyUsesTenantAwareBatchPath) { std::vector tenant_keys = {tenant_key_a, tenant_key_b}; std::vector tenant_sizes = {1024, 2048}; auto tenant_put_start = service_.BatchPutStart( - client_id, tenant_keys, tenant_sizes, config, tenant_id); + client_id, tenant_keys, tenant_sizes, config, tenant_id.value()); ASSERT_EQ(tenant_put_start.size(), tenant_keys.size()); for (const auto& result : tenant_put_start) { ASSERT_TRUE(result.has_value()) << toString(result.error()); } - auto tenant_put_end = service_.BatchPutEnd(client_id, tenant_keys, - ReplicaType::MEMORY, tenant_id); + auto tenant_put_end = + service_.BatchPutEnd(client_id, MakeObjectMetas(tenant_keys), + ReplicaType::MEMORY, tenant_id.value()); ASSERT_EQ(tenant_put_end.size(), tenant_keys.size()); for (const auto& result : tenant_put_end) { ASSERT_TRUE(result.has_value()) << toString(result.error()); @@ -4514,9 +4976,11 @@ TEST_F(MasterServiceTest, WrappedBatchExistKeyUsesTenantAwareBatchPath) { auto default_put_start = service_.PutStart(client_id, default_only_key, 1024, config); ASSERT_TRUE(default_put_start.has_value()); - ASSERT_TRUE( - service_.PutEnd(client_id, default_only_key, ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + .PutEnd(client_id, + ObjectMeta{default_only_key, std::nullopt}, + ReplicaType::MEMORY) + .has_value()); auto& metrics = MasterMetricManager::instance(); const auto base_requests = metrics.get_batch_exist_key_requests(); @@ -4527,7 +4991,7 @@ TEST_F(MasterServiceTest, WrappedBatchExistKeyUsesTenantAwareBatchPath) { std::vector lookup_keys = {tenant_key_a, default_only_key, missing_key, tenant_key_b}; - auto resp = service_.BatchExistKey(lookup_keys, tenant_id); + auto resp = service_.BatchExistKey(lookup_keys, tenant_id.value()); ASSERT_EQ(resp.size(), lookup_keys.size()); EXPECT_TRUE(resp[0].value()); EXPECT_FALSE(resp[1].value()); @@ -4542,6 +5006,77 @@ TEST_F(MasterServiceTest, WrappedBatchExistKeyUsesTenantAwareBatchPath) { EXPECT_EQ(base_failed_items, metrics.get_batch_exist_key_failed_items()); } +TEST_F(MasterServiceTest, WrappedWriteBoundaryRejectsInvalidTenantIds) { + WrappedMasterService service( + MakeStrictWrappedConfig({"registered-tenant"})); + ReplicateConfig config; + config.replica_num = 1; + const UUID client_id = generate_uuid(); + + auto empty = + service.PutStart(client_id, "empty-tenant-key", 1024, config, ""); + ASSERT_FALSE(empty.has_value()); + EXPECT_EQ(empty.error(), ErrorCode::TENANT_NOT_REGISTERED); + + const std::string control_tenant("tenant\0bad", 10); + auto invalid = service.PutStart(client_id, "invalid-tenant-key", 1024, + config, control_tenant); + ASSERT_FALSE(invalid.has_value()); + EXPECT_EQ(invalid.error(), ErrorCode::TENANT_NOT_REGISTERED); +} + +TEST_F(MasterServiceTest, WrappedRequestBoundaryRejectsInvalidTenantIds) { + WrappedMasterService service( + MakeStrictWrappedConfig({std::string(TenantId::kDefaultValue)})); + + auto invalid = service.GetReplicaList("missing-key", "_invalid-tenant"); + ASSERT_FALSE(invalid.has_value()); + EXPECT_EQ(invalid.error(), ErrorCode::INVALID_PARAMS); + + const std::string control_tenant("tenant\0bad", 10); + auto batch = + service.BatchGetReplicaList({"key-a", "key-b"}, control_tenant); + ASSERT_EQ(batch.size(), 2u); + for (const auto& result : batch) { + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); + } + + std::vector tasks = { + {.tenant_id = "_invalid-tenant", .key = "key", .size = 1}}; + std::vector metadatas = { + {.bucket_id = 0, + .offset = 0, + .key_size = 3, + .data_size = 1, + .transport_endpoint = "segment"}}; + auto offload = + service.NotifyOffloadSuccess(generate_uuid(), tasks, metadatas); + ASSERT_FALSE(offload.has_value()); + EXPECT_EQ(offload.error(), ErrorCode::INVALID_PARAMS); + + // RemoveAll has a legacy scalar return type and cannot carry ErrorCode. + EXPECT_EQ(service.RemoveAll(false, "_invalid-tenant"), 0); +} + +TEST_F(MasterServiceTest, WrappedRequestBoundaryPreservesTenantNormalization) { + WrappedMasterService multi_tenant_service( + MakeStrictWrappedConfig({std::string(TenantId::kDefaultValue)})); + auto empty = multi_tenant_service.ExistKey("missing-key", ""); + ASSERT_TRUE(empty.has_value()); + EXPECT_FALSE(empty.value()); + + WrappedMasterServiceConfig single_tenant_config; + single_tenant_config.default_kv_lease_ttl = 100; + single_tenant_config.enable_metric_reporting = false; + single_tenant_config.enable_multi_tenants = false; + WrappedMasterService single_tenant_service(single_tenant_config); + auto invalid = + single_tenant_service.ExistKey("missing-key", "_invalid-tenant"); + ASSERT_TRUE(invalid.has_value()); + EXPECT_FALSE(invalid.value()); +} + TEST_F(MasterServiceTest, BatchQueryIpTest) { std::unique_ptr service_(new MasterService()); const UUID client_id = generate_uuid(); @@ -4828,8 +5363,8 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { config.replica_num = kReplicaCnt; // Put key_1, should success. - auto put_start_result = - service_->PutStart(client_id, key_1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key_1, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -4838,8 +5373,8 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { } // Put key_1 again, should fail because the key exists. - put_start_result = - service_->PutStart(client_id, key_1, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_1, TenantId::Default(), + slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::OBJECT_ALREADY_EXISTS); @@ -4854,8 +5389,8 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { // Put key_1 again, should success because the old one has expired and will // be discarded by this put. - put_start_result = - service_->PutStart(client_id, key_1, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_1, TenantId::Default(), + slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -4864,18 +5399,18 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { } // Complete key_1. - auto put_end_result = - service_->PutEnd(client_id, key_1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key_1, TenantId::Default(), ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Protect key_1 from eviction. - auto get_result = service_->GetReplicaList(key_1, "default"); + auto get_result = service_->GetReplicaList(key_1, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); // Put key_2, should fail because the key_1 occupied 12MB (6MB processing, // 6MB discarded but not yet released) on each segment. - put_start_result = - service_->PutStart(client_id, key_2, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_2, TenantId::Default(), + slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE); @@ -4888,15 +5423,15 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { EXPECT_TRUE(result.has_value()); } // Protect key_1 from eviction. - auto get_result = service_->GetReplicaList(key_1, "default"); + auto get_result = service_->GetReplicaList(key_1, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); std::this_thread::sleep_for(std::chrono::seconds(1)); } // Put key_2 again, should success because the discarded replica has been // released. - put_start_result = - service_->PutStart(client_id, key_2, "default", slice_length, config); + put_start_result = service_->PutStart(client_id, key_2, TenantId::Default(), + slice_length, config); EXPECT_TRUE(put_start_result.has_value()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); @@ -4911,26 +5446,39 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { EXPECT_TRUE(result.has_value()); } // Protect key_1 from eviction. - auto get_result = service_->GetReplicaList(key_1, "default"); + auto get_result = service_->GetReplicaList(key_1, TenantId::Default()); EXPECT_TRUE(get_result.has_value()); std::this_thread::sleep_for(std::chrono::seconds(1)); } // Put key_2 again, should fail because eviction has not been triggered. And - // this PutStart should trigger the eviction. - put_start_result = - service_->PutStart(client_id, key_2, "default", slice_length, config); + // this PutStart should trigger the eviction. Only BatchEvict moves the + // eviction attempt counter, so take the baseline before the trigger: the + // eviction thread polls every 10 ms, and sampling after the failing + // PutStart could race a completed BatchEvict and wait for a second one + // that never comes. + const int64_t eviction_attempts_before = + MasterMetricManager::instance().get_mem_eviction_attempts(); + put_start_result = service_->PutStart(client_id, key_2, TenantId::Default(), + slice_length, config); EXPECT_FALSE(put_start_result.has_value()); EXPECT_EQ(put_start_result.error(), ErrorCode::NO_AVAILABLE_HANDLE); - // Wait a moment for the eviction to complete. - std::this_thread::sleep_for(std::chrono::milliseconds(50)); - - // Put key_2 again, should success because the previous one has been - // discarded and released. - put_start_result = - service_->PutStart(client_id, key_2, "default", slice_length, config); - EXPECT_TRUE(put_start_result.has_value()); + // The failed PutStart above sets need_mem_eviction_, and the eviction + // thread answers with an asynchronous BatchEvict. Polling PutStart for up + // to put_start_release_timeout_sec cannot tell that path apart from the + // periodic DiscardExpiredProcessingReplicas fallback, which releases the + // same replicas on the same 5 s scale and would pass the test without + // exercising the immediate eviction, and the periodic path never touches + // the attempt counter. + WaitUntil([&] { + return MasterMetricManager::instance().get_mem_eviction_attempts() > + eviction_attempts_before; + }); + put_start_result = service_->PutStart(client_id, key_2, TenantId::Default(), + slice_length, config); + ASSERT_TRUE(put_start_result.has_value()) + << toString(put_start_result.error()); replica_list = put_start_result.value(); EXPECT_EQ(replica_list.size(), kReplicaCnt); for (size_t i = 0; i < kReplicaCnt; i++) { @@ -4938,8 +5486,8 @@ TEST_F(MasterServiceTest, PutStartExpiringTest) { } // Complete key_2. - put_end_result = - service_->PutEnd(client_id, key_2, "default", ReplicaType::MEMORY); + put_end_result = service_->PutEnd(client_id, key_2, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); } @@ -5055,17 +5603,17 @@ TEST_F(MasterServiceTest, BatchReplicaClearAllSegments) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); } // Verify objects exist for (const auto& key : keys) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()); } @@ -5082,7 +5630,7 @@ TEST_F(MasterServiceTest, BatchReplicaClearAllSegments) { // Verify objects are removed for (const auto& key : keys) { - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()) << "Key " << key << " should be removed"; @@ -5112,11 +5660,11 @@ TEST_F(MasterServiceTest, BatchReplicaClearSpecificSegment) { config.replica_num = 1; config.preferred_segment = segment_name; // Ensure object is placed on segment1 - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // 4. Wait for lease to expire and verify it's actually expired @@ -5154,7 +5702,7 @@ TEST_F(MasterServiceTest, BatchReplicaClearSpecificSegment) { const auto& cleared_keys = clear_result.value(); ASSERT_EQ(1u, cleared_keys.size()) << "Key should be cleared"; - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()) << "Key should be removed after being cleared."; @@ -5174,15 +5722,15 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithLeaseActive) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Grant a lease by calling GetReplicaList (similar to normal usage) - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(get_result.has_value()); // Try to clear immediately (lease should still be active) @@ -5196,7 +5744,7 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithLeaseActive) { << "No keys should be cleared when lease is active"; // Verify object still exists - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()) << "Key should still exist"; } @@ -5216,11 +5764,11 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithDifferentClientId) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id1, key, "default", value_length, config); + auto put_start_result = service_->PutStart( + client_id1, key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id1, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id1, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Wait for lease to expire @@ -5237,7 +5785,7 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithDifferentClientId) { << "No keys should be cleared for different client_id"; // Verify object still exists - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()) << "Key should still exist"; } @@ -5292,11 +5840,11 @@ TEST_F(MasterServiceTest, BatchReplicaClearWithEmptyStringKeys) { uint64_t value_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = service_->PutStart(client_id, valid_key, "default", - value_length, config); + auto put_start_result = service_->PutStart( + client_id, valid_key, TenantId::Default(), value_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, valid_key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, valid_key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Wait for lease to expire @@ -5333,26 +5881,26 @@ TEST_F(MasterServiceTest, BatchReplicaClearMixedScenario) { config.replica_num = 1; // Create key1 and key2 with client_id1 - auto put_start1 = - service_->PutStart(client_id1, key1, "default", value_length, config); + auto put_start1 = service_->PutStart(client_id1, key1, TenantId::Default(), + value_length, config); ASSERT_TRUE(put_start1.has_value()); - auto put_end1 = - service_->PutEnd(client_id1, key1, "default", ReplicaType::MEMORY); + auto put_end1 = service_->PutEnd(client_id1, key1, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end1.has_value()); - auto put_start2 = - service_->PutStart(client_id1, key2, "default", value_length, config); + auto put_start2 = service_->PutStart(client_id1, key2, TenantId::Default(), + value_length, config); ASSERT_TRUE(put_start2.has_value()); - auto put_end2 = - service_->PutEnd(client_id1, key2, "default", ReplicaType::MEMORY); + auto put_end2 = service_->PutEnd(client_id1, key2, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end2.has_value()); // Create key3 with client_id2 - auto put_start3 = - service_->PutStart(client_id2, key3, "default", value_length, config); + auto put_start3 = service_->PutStart(client_id2, key3, TenantId::Default(), + value_length, config); ASSERT_TRUE(put_start3.has_value()); - auto put_end3 = - service_->PutEnd(client_id2, key3, "default", ReplicaType::MEMORY); + auto put_end3 = service_->PutEnd(client_id2, key3, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end3.has_value()); // Wait for lease to expire @@ -5370,16 +5918,16 @@ TEST_F(MasterServiceTest, BatchReplicaClearMixedScenario) { << "Only keys belonging to client_id1 should be cleared"; // Verify key1 and key2 are cleared - auto exist1 = service_->ExistKey(key1, "default"); + auto exist1 = service_->ExistKey(key1, TenantId::Default()); ASSERT_TRUE(exist1.has_value()); ASSERT_FALSE(exist1.value()) << "key1 should be cleared"; - auto exist2 = service_->ExistKey(key2, "default"); + auto exist2 = service_->ExistKey(key2, TenantId::Default()); ASSERT_TRUE(exist2.has_value()); ASSERT_FALSE(exist2.value()) << "key2 should be cleared"; // Verify key3 still exists (different client_id) - auto exist3 = service_->ExistKey(key3, "default"); + auto exist3 = service_->ExistKey(key3, TenantId::Default()); ASSERT_TRUE(exist3.has_value()); ASSERT_TRUE(exist3.value()) << "key3 should still exist (different client_id)"; @@ -5414,16 +5962,16 @@ TEST_F(MasterServiceTest, CreateCopyTaskTest) { ReplicateConfig config; config.replica_num = 1; config.preferred_segment = "segment_0"; - auto put_start_result = - service_->PutStart(client_id, key1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key1, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key1, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Copy key1 to "segment_1" and "segment_2" - auto copy_result = - service_->CreateCopyTask(key1, "default", {"segment_1", "segment_2"}); + auto copy_result = service_->CreateCopyTask(key1, TenantId::Default(), + {"segment_1", "segment_2"}); EXPECT_TRUE(copy_result.has_value()); // verify the copy task is created and assigned to the client who executed @@ -5434,19 +5982,19 @@ TEST_F(MasterServiceTest, CreateCopyTaskTest) { EXPECT_EQ(contexts[0].client_id, task.value().assigned_client); // Copy with empty targets should fail - auto copy_result1 = service_->CreateCopyTask(key1, "default", {}); + auto copy_result1 = service_->CreateCopyTask(key1, TenantId::Default(), {}); EXPECT_FALSE(copy_result1.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_result1.error()); // Copy not exist key should fail - auto copy_result2 = - service_->CreateCopyTask("not_exist_key", "default", {"segment_1"}); + auto copy_result2 = service_->CreateCopyTask( + "not_exist_key", TenantId::Default(), {"segment_1"}); EXPECT_FALSE(copy_result2.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, copy_result2.error()); // Copy to segment that not mounted should fail - auto copy_result3 = - service_->CreateCopyTask(key1, "default", {"not_mounted_segment"}); + auto copy_result3 = service_->CreateCopyTask(key1, TenantId::Default(), + {"not_mounted_segment"}); EXPECT_FALSE(copy_result3.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, copy_result3.error()); } @@ -5480,16 +6028,16 @@ TEST_F(MasterServiceTest, CreateMoveTaskTest) { ReplicateConfig config; config.replica_num = 1; config.preferred_segment = "segment_0"; - auto put_start_result = - service_->PutStart(client_id, key1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key1, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key1, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Move key1 from "segment_0" to "segment_1" - auto move_result = - service_->CreateMoveTask(key1, "default", "segment_0", "segment_1"); + auto move_result = service_->CreateMoveTask(key1, TenantId::Default(), + "segment_0", "segment_1"); EXPECT_TRUE(move_result.has_value()); // Verify the move task is created and assigned to the client owning the @@ -5500,32 +6048,32 @@ TEST_F(MasterServiceTest, CreateMoveTaskTest) { EXPECT_EQ(contexts[0].client_id, task.value().assigned_client); // Move non-existent key should fail - auto move_result1 = service_->CreateMoveTask("not_exist_key", "default", - "segment_0", "segment_1"); + auto move_result1 = service_->CreateMoveTask( + "not_exist_key", TenantId::Default(), "segment_0", "segment_1"); EXPECT_FALSE(move_result1.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, move_result1.error()); // Move to segment that is same as source should fail - auto move_result_same = - service_->CreateMoveTask(key1, "default", "segment_1", "segment_1"); + auto move_result_same = service_->CreateMoveTask(key1, TenantId::Default(), + "segment_1", "segment_1"); EXPECT_FALSE(move_result_same.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result_same.error()); // Move to segment that is not mounted should fail - auto move_result2 = service_->CreateMoveTask(key1, "default", "segment_0", - "not_mounted_segment"); + auto move_result2 = service_->CreateMoveTask( + key1, TenantId::Default(), "segment_0", "not_mounted_segment"); EXPECT_FALSE(move_result2.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result2.error()); // Move from segment that does not have the replica should fail - auto move_result3 = - service_->CreateMoveTask(key1, "default", "segment_2", "segment_1"); + auto move_result3 = service_->CreateMoveTask(key1, TenantId::Default(), + "segment_2", "segment_1"); EXPECT_FALSE(move_result3.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result3.error()); // Move from segment that is not mounted should fail auto move_result4 = service_->CreateMoveTask( - key1, "default", "not_mounted_segment", "segment_1"); + key1, TenantId::Default(), "not_mounted_segment", "segment_1"); EXPECT_FALSE(move_result4.has_value()); EXPECT_EQ(ErrorCode::INVALID_PARAMS, move_result4.error()); } @@ -5559,16 +6107,16 @@ TEST_F(MasterServiceTest, QueryTaskTest) { ReplicateConfig config; config.replica_num = 1; config.preferred_segment = "segment_0"; - auto put_start_result = - service_->PutStart(client_id, key1, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key1, TenantId::Default(), slice_length, config); EXPECT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key1, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key1, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_TRUE(put_end_result.has_value()); // Move key1 from "segment_0" to "segment_1" - auto move_result = - service_->CreateMoveTask(key1, "default", "segment_0", "segment_1"); + auto move_result = service_->CreateMoveTask(key1, TenantId::Default(), + "segment_0", "segment_1"); EXPECT_TRUE(move_result.has_value()); // Query non-existent task should fail @@ -5609,20 +6157,22 @@ TEST_F(MasterServiceTest, FetchTasksReturnsAssignedTasksOnlyAndDrainsQueue) { config.preferred_segment = "segment_0"; ASSERT_TRUE(service_ - ->PutStart(put_client_id, key, "default", + ->PutStart(put_client_id, key, TenantId::Default(), /*slice_length=*/1024, config) .has_value()); - ASSERT_TRUE( - service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(put_client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); // Create two tasks; both should be assigned to the client owning source // segment_0. - auto copy_task_id = service_->CreateCopyTask(key, "default", {"segment_1"}); + auto copy_task_id = + service_->CreateCopyTask(key, TenantId::Default(), {"segment_1"}); ASSERT_TRUE(copy_task_id.has_value()); - auto move_task_id = - service_->CreateMoveTask(key, "default", "segment_0", "segment_1"); + auto move_task_id = service_->CreateMoveTask(key, TenantId::Default(), + "segment_0", "segment_1"); ASSERT_TRUE(move_task_id.has_value()); // Fetch from client_0 should get both tasks (order not guaranteed). @@ -5657,7 +6207,9 @@ TEST_F(MasterServiceTest, FetchTasksReturnsAssignedTasksOnlyAndDrainsQueue) { } TEST_F(MasterServiceTest, TenantTasksCarryTenantInPayload) { - auto service = std::make_unique(); + const TenantId tenant_id("tenant_for_async_task"); + auto service = std::make_unique( + MakeStrictTenantConfig({tenant_id.value()})); const auto ctx0 = PrepareSimpleSegment(*service, "segment_0", 0x300000000, kDefaultSegmentSize); [[maybe_unused]] const auto ctx1 = PrepareSimpleSegment( @@ -5665,7 +6217,6 @@ TEST_F(MasterServiceTest, TenantTasksCarryTenantInPayload) { const UUID put_client_id = generate_uuid(); const std::string key = "tenant_task_key"; - const std::string tenant_id = "tenant_for_async_task"; ReplicateConfig config; config.replica_num = 1; @@ -5695,13 +6246,13 @@ TEST_F(MasterServiceTest, TenantTasksCarryTenantInPayload) { if (assignment.id == copy_task_id.value()) { ReplicaCopyPayload payload; struct_json::from_json(payload, assignment.payload); - EXPECT_EQ(payload.tenant_id, tenant_id); + EXPECT_EQ(payload.tenant_id, tenant_id.value()); EXPECT_EQ(payload.key, key); saw_copy = true; } else if (assignment.id == move_task_id.value()) { ReplicaMovePayload payload; struct_json::from_json(payload, assignment.payload); - EXPECT_EQ(payload.tenant_id, tenant_id); + EXPECT_EQ(payload.tenant_id, tenant_id.value()); EXPECT_EQ(payload.key, key); saw_move = true; } @@ -5715,7 +6266,7 @@ TEST_F(MasterServiceTest, LegacyTaskPayloadDefaultsTenant) { struct_json::from_json( copy_payload, R"({"key":"legacy_copy_key","source":"segment_0","targets":["segment_1"]})"); - EXPECT_EQ(copy_payload.tenant_id, "default"); + EXPECT_EQ(copy_payload.tenant_id, TenantId::kDefaultValue); EXPECT_EQ(copy_payload.key, "legacy_copy_key"); EXPECT_EQ(copy_payload.source, "segment_0"); ASSERT_EQ(copy_payload.targets.size(), 1u); @@ -5725,7 +6276,7 @@ TEST_F(MasterServiceTest, LegacyTaskPayloadDefaultsTenant) { struct_json::from_json( move_payload, R"({"key":"legacy_move_key","source":"segment_0","target":"segment_1"})"); - EXPECT_EQ(move_payload.tenant_id, "default"); + EXPECT_EQ(move_payload.tenant_id, TenantId::kDefaultValue); EXPECT_EQ(move_payload.key, "legacy_move_key"); EXPECT_EQ(move_payload.source, "segment_0"); EXPECT_EQ(move_payload.target, "segment_1"); @@ -5747,17 +6298,18 @@ TEST_F(MasterServiceTest, FetchTasksRespectsBatchSize) { config.preferred_segment = "segment_0"; ASSERT_TRUE(service_ - ->PutStart(put_client_id, key, "default", + ->PutStart(put_client_id, key, TenantId::Default(), /*slice_length=*/1024, config) .has_value()); - ASSERT_TRUE( - service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(put_client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); - auto t1 = service_->CreateCopyTask(key, "default", {"segment_1"}); + auto t1 = service_->CreateCopyTask(key, TenantId::Default(), {"segment_1"}); ASSERT_TRUE(t1.has_value()); - auto t2 = - service_->CreateMoveTask(key, "default", "segment_0", "segment_1"); + auto t2 = service_->CreateMoveTask(key, TenantId::Default(), "segment_0", + "segment_1"); ASSERT_TRUE(t2.has_value()); auto fetch_first = service_->FetchTasks(ctx0.client_id, /*batch_size=*/1); @@ -5799,15 +6351,17 @@ TEST_F(MasterServiceTest, UpdateTaskSuccessFlow) { config.preferred_segment = "segment_0"; ASSERT_TRUE(service_ - ->PutStart(put_client_id, key, "default", + ->PutStart(put_client_id, key, TenantId::Default(), /*slice_length=*/1024, config) .has_value()); - ASSERT_TRUE( - service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(put_client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); // Create a task assigned to client owning segment_0. - auto task_id_res = service_->CreateCopyTask(key, "default", {"segment_1"}); + auto task_id_res = + service_->CreateCopyTask(key, TenantId::Default(), {"segment_1"}); ASSERT_TRUE(task_id_res.has_value()); const UUID task_id = task_id_res.value(); @@ -5857,15 +6411,16 @@ TEST_F(MasterServiceTest, UpdateTaskRejectsWrongClient) { config.preferred_segment = "segment_0"; ASSERT_TRUE(service_ - ->PutStart(put_client_id, key, "default", + ->PutStart(put_client_id, key, TenantId::Default(), /*slice_length=*/1024, config) .has_value()); - ASSERT_TRUE( - service_->PutEnd(put_client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(put_client_id, key, TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); - auto task_id_res = - service_->CreateMoveTask(key, "default", "segment_0", "segment_1"); + auto task_id_res = service_->CreateMoveTask(key, TenantId::Default(), + "segment_0", "segment_1"); ASSERT_TRUE(task_id_res.has_value()); const UUID task_id = task_id_res.value(); @@ -5928,8 +6483,9 @@ TEST_F(MasterServiceTest, config.replica_num = 1; config.preferred_segment = "segment_0"; - auto put_result = service_->PutStart( - ctx0.client_id, "drain_skip_allocation_key", "default", 1024, config); + auto put_result = + service_->PutStart(ctx0.client_id, "drain_skip_allocation_key", + TenantId::Default(), 1024, config); ASSERT_TRUE(put_result.has_value()); ASSERT_EQ(put_result->size(), 1u); EXPECT_EQ(put_result->front() @@ -5938,7 +6494,7 @@ TEST_F(MasterServiceTest, "segment_1"); ASSERT_TRUE(service_ ->PutEnd(ctx0.client_id, "drain_skip_allocation_key", - "default", ReplicaType::MEMORY) + TenantId::Default(), ReplicaType::MEMORY) .has_value()); } @@ -5982,7 +6538,7 @@ TEST_F(MasterServiceTest, DrainJobSchedulesMoveTaskAndConvergesToDrained) { ASSERT_TRUE(segment_status.has_value()); EXPECT_EQ(segment_status.value(), SegmentStatus::DRAINED); - auto replicas = service_->GetReplicaList(key, "default"); + auto replicas = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(replicas.has_value()); std::unordered_set segment_names; for (const auto& replica : replicas->replicas) { @@ -6113,29 +6669,30 @@ TEST_F(MasterServiceTest, ForceRemoveLeasedObject) { uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Verify object exists - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()); // Normal remove should fail because object has active lease - auto remove_result_no_force = service_->Remove(key, "default", false); + auto remove_result_no_force = + service_->Remove(key, TenantId::Default(), false); EXPECT_FALSE(remove_result_no_force.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_LEASE, remove_result_no_force.error()); // Force remove should succeed even with active lease - auto remove_result_force = service_->Remove(key, "default", true); + auto remove_result_force = service_->Remove(key, TenantId::Default(), true); EXPECT_TRUE(remove_result_force.has_value()); // Verify object is removed - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error()); } @@ -6157,43 +6714,43 @@ TEST_F(MasterServiceTest, ForceRemoveByRegexLeasedObjects) { uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Grant lease by reading the object - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()); } // Normal RemoveByRegex should remove 0 because all objects have active // leases - auto remove_result_no_force = - service_->RemoveByRegex("^force_regex_key_", "default", false); + auto remove_result_no_force = service_->RemoveByRegex( + "^force_regex_key_", TenantId::Default(), false); ASSERT_TRUE(remove_result_no_force.has_value()); EXPECT_EQ(0, remove_result_no_force.value()); // All objects should still exist for (int i = 0; i < 5; ++i) { std::string key = "force_regex_key_" + std::to_string(i); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()); } // Force RemoveByRegex should remove all 5 objects auto remove_result_force = - service_->RemoveByRegex("^force_regex_key_", "default", true); + service_->RemoveByRegex("^force_regex_key_", TenantId::Default(), true); ASSERT_TRUE(remove_result_force.has_value()); EXPECT_EQ(5, remove_result_force.value()); // All objects should be removed for (int i = 0; i < 5; ++i) { std::string key = "force_regex_key_" + std::to_string(i); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()); } @@ -6216,14 +6773,14 @@ TEST_F(MasterServiceTest, ForceRemoveAllLeasedObjects) { uint64_t slice_length = 1024; ReplicateConfig config; config.replica_num = 1; - auto put_start_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_start_result = service_->PutStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(put_start_result.has_value()); - auto put_end_result = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end_result = service_->PutEnd( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(put_end_result.has_value()); // Grant lease by reading the object - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()); } @@ -6234,7 +6791,7 @@ TEST_F(MasterServiceTest, ForceRemoveAllLeasedObjects) { // All objects should still exist for (int i = 0; i < 10; ++i) { std::string key = "force_all_key_" + std::to_string(i); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_TRUE(exist_result.value()); } @@ -6245,7 +6802,7 @@ TEST_F(MasterServiceTest, ForceRemoveAllLeasedObjects) { // All objects should be removed for (int i = 0; i < 10; ++i) { std::string key = "force_all_key_" + std::to_string(i); - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()); } @@ -6263,25 +6820,25 @@ TEST_F(MasterServiceTest, UpsertNewKey) { ReplicateConfig config; config.replica_num = 1; - auto upsert_result = - service_->UpsertStart(client_id, key, "default", slice_length, config); + auto upsert_result = service_->UpsertStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(upsert_result.has_value()); auto replicas = upsert_result.value(); EXPECT_EQ(1, replicas.size()); EXPECT_EQ(ReplicaStatus::PROCESSING, replicas[0].status); // During upsert, GetReplicaList should return not ready - auto get_result = service_->GetReplicaList(key, "default"); + auto get_result = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_result.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error()); // UpsertEnd completes the operation - auto end_result = - service_->UpsertEnd(client_id, key, "default", ReplicaType::MEMORY); + auto end_result = service_->UpsertEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(end_result.has_value()); // Verify replica is COMPLETE - auto final_result = service_->GetReplicaList(key, "default"); + auto final_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(final_result.has_value()); EXPECT_EQ(1, final_result.value().replicas.size()); EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status); @@ -6299,18 +6856,18 @@ TEST_F(MasterServiceTest, UpsertSameSize) { config.replica_num = 1; // First: PutStart + PutEnd to create the object - auto put_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_result = service_->PutStart(client_id, key, TenantId::Default(), + slice_length, config); ASSERT_TRUE(put_result.has_value()); auto original_replicas = put_result.value(); - auto put_end = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()); // UpsertStart with same size — should reuse buffers const UUID new_client_id = generate_uuid(); - auto upsert_result = service_->UpsertStart(new_client_id, key, "default", - slice_length, config); + auto upsert_result = service_->UpsertStart( + new_client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(upsert_result.has_value()); auto upsert_replicas = upsert_result.value(); EXPECT_EQ(1, upsert_replicas.size()); @@ -6325,12 +6882,12 @@ TEST_F(MasterServiceTest, UpsertSameSize) { .buffer_descriptor.buffer_address_); // UpsertEnd with the new client_id - auto end_result = - service_->UpsertEnd(new_client_id, key, "default", ReplicaType::MEMORY); + auto end_result = service_->UpsertEnd( + new_client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(end_result.has_value()); // Verify replica is COMPLETE again - auto final_result = service_->GetReplicaList(key, "default"); + auto final_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(final_result.has_value()); EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status); } @@ -6348,27 +6905,27 @@ TEST_F(MasterServiceTest, UpsertSameSizeRefreshesMetadata) { config.replica_num = 1; // Create object with client_a - auto put_result = - service_->PutStart(client_id_a, key, "default", slice_length, config); + auto put_result = service_->PutStart(client_id_a, key, TenantId::Default(), + slice_length, config); ASSERT_TRUE(put_result.has_value()); - auto put_end = - service_->PutEnd(client_id_a, key, "default", ReplicaType::MEMORY); + auto put_end = service_->PutEnd(client_id_a, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()); // UpsertStart with client_b - auto upsert_result = service_->UpsertStart(client_id_b, key, "default", - slice_length, config); + auto upsert_result = service_->UpsertStart( + client_id_b, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(upsert_result.has_value()); // UpsertEnd with client_a should fail (client_id was refreshed to client_b) - auto end_fail = - service_->UpsertEnd(client_id_a, key, "default", ReplicaType::MEMORY); + auto end_fail = service_->UpsertEnd(client_id_a, key, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_FALSE(end_fail.has_value()); EXPECT_EQ(ErrorCode::ILLEGAL_CLIENT, end_fail.error()); // UpsertEnd with client_b should succeed - auto end_ok = - service_->UpsertEnd(client_id_b, key, "default", ReplicaType::MEMORY); + auto end_ok = service_->UpsertEnd(client_id_b, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(end_ok.has_value()); } @@ -6385,17 +6942,17 @@ TEST_F(MasterServiceTest, UpsertDifferentSize) { config.replica_num = 1; // Create object with original_size - auto put_result = - service_->PutStart(client_id, key, "default", original_size, config); + auto put_result = service_->PutStart(client_id, key, TenantId::Default(), + original_size, config); ASSERT_TRUE(put_result.has_value()); auto original_replicas = put_result.value(); - auto put_end = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()); // UpsertStart with different size - auto upsert_result = - service_->UpsertStart(client_id, key, "default", new_size, config); + auto upsert_result = service_->UpsertStart( + client_id, key, TenantId::Default(), new_size, config); ASSERT_TRUE(upsert_result.has_value()); auto new_replicas = upsert_result.value(); EXPECT_EQ(1, new_replicas.size()); @@ -6410,12 +6967,12 @@ TEST_F(MasterServiceTest, UpsertDifferentSize) { .buffer_descriptor.buffer_address_); // UpsertEnd - auto end_result = - service_->UpsertEnd(client_id, key, "default", ReplicaType::MEMORY); + auto end_result = service_->UpsertEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(end_result.has_value()); // Verify the object is complete - auto final_result = service_->GetReplicaList(key, "default"); + auto final_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(final_result.has_value()); EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status); } @@ -6441,21 +6998,21 @@ TEST_F(MasterServiceTest, UpsertConflictReplicationTask) { config.preferred_segment = "segment_1"; // Create object - auto put_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_result = service_->PutStart(client_id, key, TenantId::Default(), + slice_length, config); ASSERT_TRUE(put_result.has_value()); - auto put_end = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()); // Start a Copy - auto copy_result = service_->CopyStart(client_id, key, "default", + auto copy_result = service_->CopyStart(client_id, key, TenantId::Default(), "segment_1", {"segment_2"}); ASSERT_TRUE(copy_result.has_value()); // UpsertStart should fail with OBJECT_HAS_REPLICATION_TASK - auto upsert_result = - service_->UpsertStart(client_id, key, "default", slice_length, config); + auto upsert_result = service_->UpsertStart( + client_id, key, TenantId::Default(), slice_length, config); EXPECT_FALSE(upsert_result.has_value()); EXPECT_EQ(ErrorCode::OBJECT_HAS_REPLICATION_TASK, upsert_result.error()); } @@ -6474,30 +7031,30 @@ TEST_F(MasterServiceTest, UpsertPreemptsInProgressPut) { config.replica_num = 1; // Client A starts a Put but doesn't finish - auto put_result = - service_->PutStart(client_a, key, "default", slice_length, config); + auto put_result = service_->PutStart(client_a, key, TenantId::Default(), + slice_length, config); ASSERT_TRUE(put_result.has_value()); // Client B upserts the same key — should preempt client A - auto upsert_result = - service_->UpsertStart(client_b, key, "default", slice_length, config); + auto upsert_result = service_->UpsertStart( + client_b, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(upsert_result.has_value()); auto upsert_replicas = upsert_result.value(); EXPECT_EQ(1, upsert_replicas.size()); EXPECT_EQ(ReplicaStatus::PROCESSING, upsert_replicas[0].status); // Client A's PutEnd should fail - auto put_end_a = - service_->PutEnd(client_a, key, "default", ReplicaType::MEMORY); + auto put_end_a = service_->PutEnd(client_a, key, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_FALSE(put_end_a.has_value()); // Client B's UpsertEnd should succeed - auto upsert_end = - service_->UpsertEnd(client_b, key, "default", ReplicaType::MEMORY); + auto upsert_end = service_->UpsertEnd(client_b, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(upsert_end.has_value()); // Verify final state - auto final_result = service_->GetReplicaList(key, "default"); + auto final_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(final_result.has_value()); EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status); } @@ -6514,17 +7071,17 @@ TEST_F(MasterServiceTest, UpsertRevoke) { config.replica_num = 1; // UpsertStart (Case A — new key) - auto upsert_result = - service_->UpsertStart(client_id, key, "default", slice_length, config); + auto upsert_result = service_->UpsertStart( + client_id, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(upsert_result.has_value()); // UpsertRevoke - auto revoke_result = - service_->UpsertRevoke(client_id, key, "default", ReplicaType::MEMORY); + auto revoke_result = service_->UpsertRevoke( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(revoke_result.has_value()); // Key should be gone - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); EXPECT_FALSE(exist_result.value()); } @@ -6541,26 +7098,26 @@ TEST_F(MasterServiceTest, UpsertInPlaceThenRevoke) { config.replica_num = 1; // Create object first - auto put_result = - service_->PutStart(client_id, key, "default", slice_length, config); + auto put_result = service_->PutStart(client_id, key, TenantId::Default(), + slice_length, config); ASSERT_TRUE(put_result.has_value()); - auto put_end = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()); // UpsertStart in-place (same size) const UUID new_client = generate_uuid(); - auto upsert_result = - service_->UpsertStart(new_client, key, "default", slice_length, config); + auto upsert_result = service_->UpsertStart( + new_client, key, TenantId::Default(), slice_length, config); ASSERT_TRUE(upsert_result.has_value()); // UpsertRevoke — replicas are PROCESSING, should be erased - auto revoke_result = - service_->UpsertRevoke(new_client, key, "default", ReplicaType::MEMORY); + auto revoke_result = service_->UpsertRevoke( + new_client, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(revoke_result.has_value()); // Key should be gone (no valid replicas left) - auto exist_result = service_->ExistKey(key, "default"); + auto exist_result = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); EXPECT_FALSE(exist_result.value()); } @@ -6575,25 +7132,26 @@ TEST_F(MasterServiceTest, BatchUpsertStart) { config.replica_num = 1; // Create key_1 with size 1024 - auto put_result = - service_->PutStart(client_id, "key_1", "default", 1024, config); + auto put_result = service_->PutStart(client_id, "key_1", + TenantId::Default(), 1024, config); ASSERT_TRUE(put_result.has_value()); - auto put_end = - service_->PutEnd(client_id, "key_1", "default", ReplicaType::MEMORY); + auto put_end = service_->PutEnd(client_id, "key_1", TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()); // BatchUpsertStart: key_1 (same size), key_2 (new) std::vector keys = {"key_1", "key_2"}; std::vector slice_lengths = {1024, 2048}; - auto results = service_->BatchUpsertStart(client_id, keys, "default", - slice_lengths, config); + auto results = service_->BatchUpsertStart( + client_id, keys, TenantId::Default(), slice_lengths, config); ASSERT_EQ(2, results.size()); EXPECT_TRUE(results[0].has_value()); // key_1: Case B (in-place) EXPECT_TRUE(results[1].has_value()); // key_2: Case A (new) // Complete both - auto end_results = service_->BatchUpsertEnd(client_id, keys, "default"); + auto end_results = service_->BatchUpsertEnd( + client_id, MakeObjectMetas(keys), TenantId::Default()); ASSERT_EQ(2, end_results.size()); EXPECT_TRUE(end_results[0].has_value()); EXPECT_TRUE(end_results[1].has_value()); @@ -6615,42 +7173,42 @@ TEST_F(MasterServiceTest, UpsertPreemptsInProgressUpsert) { config.replica_num = 1; // Step 1: Create the object via Put - auto put_result = - service_->PutStart(client_a, key, "default", slice_length, config); + auto put_result = service_->PutStart(client_a, key, TenantId::Default(), + slice_length, config); ASSERT_TRUE(put_result.has_value()); - auto put_end = - service_->PutEnd(client_a, key, "default", ReplicaType::MEMORY); + auto put_end = service_->PutEnd(client_a, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()); // Step 2: Client B starts in-place upsert (Case B) — marks COMPLETE → // PROCESSING - auto upsert_b = - service_->UpsertStart(client_b, key, "default", slice_length, config); + auto upsert_b = service_->UpsertStart(client_b, key, TenantId::Default(), + slice_length, config); ASSERT_TRUE(upsert_b.has_value()); // Key should be unreadable now (all replicas are PROCESSING) - auto get_mid = service_->GetReplicaList(key, "default"); + auto get_mid = service_->GetReplicaList(key, TenantId::Default()); EXPECT_FALSE(get_mid.has_value()); EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_mid.error()); // Step 3: Client C upserts the same key — preempts Client B - auto upsert_c = - service_->UpsertStart(client_c, key, "default", slice_length, config); + auto upsert_c = service_->UpsertStart(client_c, key, TenantId::Default(), + slice_length, config); ASSERT_TRUE(upsert_c.has_value()); EXPECT_EQ(1, upsert_c.value().size()); // Step 4: Client B's UpsertEnd should fail (preempted) - auto end_b = - service_->UpsertEnd(client_b, key, "default", ReplicaType::MEMORY); + auto end_b = service_->UpsertEnd(client_b, key, TenantId::Default(), + ReplicaType::MEMORY); EXPECT_FALSE(end_b.has_value()); // Step 5: Client C's UpsertEnd should succeed - auto end_c = - service_->UpsertEnd(client_c, key, "default", ReplicaType::MEMORY); + auto end_c = service_->UpsertEnd(client_c, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(end_c.has_value()); // Final verification - auto final_result = service_->GetReplicaList(key, "default"); + auto final_result = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(final_result.has_value()); EXPECT_EQ(1, final_result.value().replicas.size()); EXPECT_EQ(ReplicaStatus::COMPLETE, final_result.value().replicas[0].status); @@ -6671,31 +7229,31 @@ TEST_F(MasterServiceTest, UpsertDifferentSizeThenRevoke) { config.replica_num = 1; // Create object with original size - auto put_result = - service_->PutStart(client_id, key, "default", original_size, config); + auto put_result = service_->PutStart(client_id, key, TenantId::Default(), + original_size, config); ASSERT_TRUE(put_result.has_value()); - auto put_end = - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end = service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()); // Verify the key exists - auto exist_before = service_->ExistKey(key, "default"); + auto exist_before = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_before.has_value()); EXPECT_TRUE(exist_before.value()); // UpsertStart with different size (Case C) — old replicas discarded, // new replicas allocated - auto upsert_result = - service_->UpsertStart(client_id, key, "default", new_size, config); + auto upsert_result = service_->UpsertStart( + client_id, key, TenantId::Default(), new_size, config); ASSERT_TRUE(upsert_result.has_value()); // Revoke — erase the newly allocated PROCESSING replicas - auto revoke_result = - service_->UpsertRevoke(client_id, key, "default", ReplicaType::MEMORY); + auto revoke_result = service_->UpsertRevoke( + client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(revoke_result.has_value()); // Key should be gone (old replicas in discarded, new replicas erased) - auto exist_after = service_->ExistKey(key, "default"); + auto exist_after = service_->ExistKey(key, TenantId::Default()); ASSERT_TRUE(exist_after.has_value()); EXPECT_FALSE(exist_after.value()); } @@ -6723,11 +7281,11 @@ TEST_F(MasterServiceTest, HardPinObjectNotEvicted) { ReplicateConfig config; config.replica_num = 1; config.with_hard_pin = true; - auto result = service_->PutStart(client_id, "pinned_model", "default", - value_size, config); + auto result = service_->PutStart( + client_id, "pinned_model", TenantId::Default(), value_size, config); ASSERT_TRUE(result.has_value()); ASSERT_TRUE(service_ - ->PutEnd(client_id, "pinned_model", "default", + ->PutEnd(client_id, "pinned_model", TenantId::Default(), ReplicaType::MEMORY) .has_value()); } @@ -6737,10 +7295,11 @@ TEST_F(MasterServiceTest, HardPinObjectNotEvicted) { std::string key = "filler_" + std::to_string(i); ReplicateConfig config; config.replica_num = 1; - auto result = - service_->PutStart(client_id, key, "default", value_size, config); + auto result = service_->PutStart(client_id, key, TenantId::Default(), + value_size, config); if (result.has_value()) { - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); } } @@ -6748,15 +7307,16 @@ TEST_F(MasterServiceTest, HardPinObjectNotEvicted) { std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl + 500)); // Hard-pinned object must still be there - auto get_result = service_->GetReplicaList("pinned_model", "default"); + auto get_result = + service_->GetReplicaList("pinned_model", TenantId::Default()); ASSERT_TRUE(get_result.has_value()) << "Hard-pinned object was evicted, but it should never be"; // Explicit Remove should still work on hard-pinned objects auto remove_result = - service_->Remove("pinned_model", "default", /*force=*/true); + service_->Remove("pinned_model", TenantId::Default(), /*force=*/true); ASSERT_TRUE(remove_result.has_value()); - auto exist_result = service_->ExistKey("pinned_model", "default"); + auto exist_result = service_->ExistKey("pinned_model", TenantId::Default()); ASSERT_TRUE(exist_result.has_value()); ASSERT_FALSE(exist_result.value()); @@ -6791,11 +7351,11 @@ TEST_F(MasterServiceTest, HardPinWithSoftPinEvictionOrder) { config.replica_num = 1; config.with_hard_pin = true; ASSERT_TRUE(service_ - ->PutStart(client_id, "hard_pinned", "default", - value_size, config) + ->PutStart(client_id, "hard_pinned", + TenantId::Default(), value_size, config) .has_value()); ASSERT_TRUE(service_ - ->PutEnd(client_id, "hard_pinned", "default", + ->PutEnd(client_id, "hard_pinned", TenantId::Default(), ReplicaType::MEMORY) .has_value()); } @@ -6806,11 +7366,11 @@ TEST_F(MasterServiceTest, HardPinWithSoftPinEvictionOrder) { config.replica_num = 1; config.with_soft_pin = true; ASSERT_TRUE(service_ - ->PutStart(client_id, "soft_pinned", "default", - value_size, config) + ->PutStart(client_id, "soft_pinned", + TenantId::Default(), value_size, config) .has_value()); ASSERT_TRUE(service_ - ->PutEnd(client_id, "soft_pinned", "default", + ->PutEnd(client_id, "soft_pinned", TenantId::Default(), ReplicaType::MEMORY) .has_value()); } @@ -6820,10 +7380,11 @@ TEST_F(MasterServiceTest, HardPinWithSoftPinEvictionOrder) { std::string key = "normal_" + std::to_string(i); ReplicateConfig config; config.replica_num = 1; - auto result = - service_->PutStart(client_id, key, "default", value_size, config); + auto result = service_->PutStart(client_id, key, TenantId::Default(), + value_size, config); if (result.has_value()) { - service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY); + service_->PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); } } @@ -6831,7 +7392,8 @@ TEST_F(MasterServiceTest, HardPinWithSoftPinEvictionOrder) { std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl + 500)); // Hard-pinned always survives - ASSERT_TRUE(service_->GetReplicaList("hard_pinned", "default").has_value()) + ASSERT_TRUE(service_->GetReplicaList("hard_pinned", TenantId::Default()) + .has_value()) << "Hard-pinned object was evicted"; std::this_thread::sleep_for(std::chrono::milliseconds(kv_lease_ttl)); @@ -6853,28 +7415,33 @@ TEST_F(MasterServiceTest, HardPinDefaultIsFalse) { // Put without hard_pin (default) ReplicateConfig config; config.replica_num = 1; - ASSERT_TRUE( - service_->PutStart(client_id, "normal_key", "default", 1024, config) - .has_value()); - ASSERT_TRUE( - service_ - ->PutEnd(client_id, "normal_key", "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, "normal_key", TenantId::Default(), + 1024, config) + .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, "normal_key", TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); // Put with hard_pin ReplicateConfig hp_config; hp_config.replica_num = 1; hp_config.with_hard_pin = true; - ASSERT_TRUE( - service_->PutStart(client_id, "hp_key", "default", 1024, hp_config) - .has_value()); - ASSERT_TRUE( - service_->PutEnd(client_id, "hp_key", "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE(service_ + ->PutStart(client_id, "hp_key", TenantId::Default(), 1024, + hp_config) + .has_value()); + ASSERT_TRUE(service_ + ->PutEnd(client_id, "hp_key", TenantId::Default(), + ReplicaType::MEMORY) + .has_value()); // Both should exist - ASSERT_TRUE(service_->GetReplicaList("normal_key", "default").has_value()); - ASSERT_TRUE(service_->GetReplicaList("hp_key", "default").has_value()); + ASSERT_TRUE(service_->GetReplicaList("normal_key", TenantId::Default()) + .has_value()); + ASSERT_TRUE( + service_->GetReplicaList("hp_key", TenantId::Default()).has_value()); service_->RemoveAll(); } @@ -7063,10 +7630,12 @@ TEST_F(MasterServiceTest, GracefulUnmountSegment_PreventAllocation) { config.preferred_segment = segment1.name; auto put_start = - service_->PutStart(client_id, key, "default", 1024, config); + service_->PutStart(client_id, key, TenantId::Default(), 1024, config); ASSERT_TRUE(put_start.has_value()); - ASSERT_TRUE(service_->PutEnd(client_id, key, "default", ReplicaType::MEMORY) - .has_value()); + ASSERT_TRUE( + service_ + ->PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY) + .has_value()); // Graceful unmount segment1 ASSERT_TRUE(service_->GracefulUnmountSegment(segment1.id, client_id, 1000) @@ -7079,7 +7648,7 @@ TEST_F(MasterServiceTest, GracefulUnmountSegment_PreventAllocation) { // Existing replicas on the graceful segment should remain readable during // the grace window. - auto existing_replicas = service_->GetReplicaList(key, "default"); + auto existing_replicas = service_->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(existing_replicas.has_value()); ASSERT_EQ(existing_replicas->replicas.size(), 1u); EXPECT_EQ(existing_replicas->replicas[0] @@ -7098,7 +7667,7 @@ TEST_F(MasterServiceTest, GracefulUnmountSegment_PreventAllocation) { config2.replica_num = 1; auto put_start2 = - service_->PutStart(client_id, key2, "default", 1024, config2); + service_->PutStart(client_id, key2, TenantId::Default(), 1024, config2); ASSERT_TRUE(put_start2.has_value()); auto replicas = put_start2.value(); ASSERT_EQ(replicas.size(), 1u); diff --git a/mooncake-store/tests/mmap_arena_fallback_test.cpp b/mooncake-store/tests/mmap_arena_fallback_test.cpp index 8e6289b11b..85f0fc7806 100644 --- a/mooncake-store/tests/mmap_arena_fallback_test.cpp +++ b/mooncake-store/tests/mmap_arena_fallback_test.cpp @@ -3,12 +3,15 @@ #include #include +#include +#include #include #include #include #include #include +#include #include "utils.h" @@ -42,10 +45,71 @@ class MmapArenaFallbackTest : public ::testing::Test { void TearDown() override { unsetenv("MC_DISABLE_MMAP_ARENA"); unsetenv("MC_MMAP_ARENA_POOL_SIZE"); + unsetenv("MC_STORE_HUGEPAGE_SIZE"); unsetenv("MC_STORE_USE_HUGEPAGE"); } }; +TEST_F(MmapArenaFallbackTest, PopulateHugetlbMappingUsesConfiguredPageStride) { + setenv("MC_STORE_USE_HUGEPAGE", "1", 1); + setenv("MC_STORE_HUGEPAGE_SIZE", "2MB", 1); + + constexpr size_t kPageCount = 3; + constexpr size_t kMapSize = kPageCount * SZ_2MB; + void* mapping = mmap(nullptr, kMapSize, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + ASSERT_NE(mapping, MAP_FAILED); + + auto* bytes = static_cast(mapping); + for (size_t page = 0; page < kPageCount; ++page) { + bytes[page * SZ_2MB] = 0xAB; + } + + populate_hugetlb_mapping(mapping, kMapSize); + + for (size_t page = 0; page < kPageCount; ++page) { + EXPECT_EQ(bytes[page * SZ_2MB], 0); + } + EXPECT_EQ(munmap(mapping, kMapSize), 0); +} + +TEST_F(MmapArenaFallbackTest, PopulateNumaHugetlbMappingTouchesEveryRegion) { + if (numa_available() < 0) { + GTEST_SKIP() << "NUMA is unavailable"; + } + + setenv("MC_STORE_USE_HUGEPAGE", "1", 1); + setenv("MC_STORE_HUGEPAGE_SIZE", "2MB", 1); + + std::vector numa_nodes; + for (int node = 0; node <= numa_max_node() && numa_nodes.size() < 2; + ++node) { + if (numa_bitmask_isbitset(numa_all_nodes_ptr, node)) { + numa_nodes.push_back(node); + } + } + ASSERT_FALSE(numa_nodes.empty()); + + constexpr size_t kPagesPerRegion = 2; + const size_t page_count = kPagesPerRegion * numa_nodes.size(); + const size_t map_size = page_count * SZ_2MB; + void* mapping = mmap(nullptr, map_size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + ASSERT_NE(mapping, MAP_FAILED); + + auto* bytes = static_cast(mapping); + for (size_t page = 0; page < page_count; ++page) { + bytes[page * SZ_2MB] = 0xAB; + } + + populate_hugetlb_numa_mapping(mapping, map_size, numa_nodes); + + for (size_t page = 0; page < page_count; ++page) { + EXPECT_EQ(bytes[page * SZ_2MB], 0); + } + EXPECT_EQ(munmap(mapping, map_size), 0); +} + TEST_F(MmapArenaFallbackTest, ArenaInitFailureIsStickyForProcessLifetime) { unsetenv("MC_DISABLE_MMAP_ARENA"); unsetenv("MC_STORE_USE_HUGEPAGE"); diff --git a/mooncake-store/tests/mmap_arena_test.cpp b/mooncake-store/tests/mmap_arena_test.cpp index 05425f0bd2..438aea2b14 100644 --- a/mooncake-store/tests/mmap_arena_test.cpp +++ b/mooncake-store/tests/mmap_arena_test.cpp @@ -651,10 +651,11 @@ TEST_F(MmapArenaTest, PagesArePhysicallyBackedAfterInit) { << "), falling back to read-verification"; // Read every page — if MAP_POPULATE didn't work, this would trigger // page faults (which is fine for CPU but would crash GPU DMA). - volatile char sink = 0; + char sum = 0; for (size_t off = 0; off < pool_size; off += sys_page_size) { - sink += static_cast(base)[off]; + sum += static_cast(base)[off]; } + volatile char sink = sum; (void)sink; // If we get here without SIGSEGV, at least CPU access works. // The real MAP_POPULATE guarantee is that DMA works too, which diff --git a/mooncake-store/tests/object_data_type_test.cpp b/mooncake-store/tests/object_data_type_test.cpp index 9b50dc4f8e..a14a72e5eb 100644 --- a/mooncake-store/tests/object_data_type_test.cpp +++ b/mooncake-store/tests/object_data_type_test.cpp @@ -110,13 +110,13 @@ TEST_F(ObjectDataTypeTest, PutStartWithDataType) { config.replica_num = 1; config.data_type = ObjectDataType::WEIGHT; - auto result = - service->PutStart(put_client, "key_weight", "default", 1024, config); + auto result = service->PutStart(put_client, "key_weight", + TenantId::Default(), 1024, config); ASSERT_TRUE(result.has_value()); EXPECT_FALSE(result.value().empty()); - auto end_result = service->PutEnd(put_client, "key_weight", "default", - ReplicaType::MEMORY); + auto end_result = service->PutEnd(put_client, "key_weight", + TenantId::Default(), ReplicaType::MEMORY); EXPECT_TRUE(end_result.has_value()); } @@ -133,8 +133,8 @@ TEST_F(ObjectDataTypeTest, PutStartDefaultDataType) { config.replica_num = 1; // data_type left as default (UNKNOWN) - auto result = - service->PutStart(put_client, "key_default", "default", 1024, config); + auto result = service->PutStart(put_client, "key_default", + TenantId::Default(), 1024, config); ASSERT_TRUE(result.has_value()); EXPECT_FALSE(result.value().empty()); } diff --git a/mooncake-store/tests/offload_on_evict_test.cpp b/mooncake-store/tests/offload_on_evict_test.cpp index fdb687a735..08810e6fdf 100644 --- a/mooncake-store/tests/offload_on_evict_test.cpp +++ b/mooncake-store/tests/offload_on_evict_test.cpp @@ -56,10 +56,10 @@ class OffloadOnEvictTest : public ::testing::Test { ReplicateConfig config; config.replica_num = 1; auto put_start = - service.PutStart(client_id, key, "default", size, config); + service.PutStart(client_id, key, TenantId::Default(), size, config); ASSERT_TRUE(put_start.has_value()) << "PutStart failed for key=" << key; - auto put_end = - service.PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end = service.PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()) << "PutEnd failed for key=" << key; } @@ -103,10 +103,10 @@ class OffloadOnEvictTest : public ::testing::Test { std::string key = key_prefix + std::to_string(i); ReplicateConfig config; config.replica_num = 1; - auto result = service.PutStart(client_id, key, "default", + auto result = service.PutStart(client_id, key, TenantId::Default(), object_size, config); if (result.has_value()) { - auto end = service.PutEnd(client_id, key, "default", + auto end = service.PutEnd(client_id, key, TenantId::Default(), ReplicaType::MEMORY); EXPECT_TRUE(end.has_value()); success_puts++; @@ -228,10 +228,10 @@ TEST_F(OffloadOnEvictTest, ComboB_EvictionTriggersOffload) { std::string key = "evict_b_" + std::to_string(i); ReplicateConfig config; config.replica_num = 1; - auto result = service->PutStart(ctx.client_id, key, "default", + auto result = service->PutStart(ctx.client_id, key, TenantId::Default(), object_size, config); if (result.has_value()) { - auto end = service->PutEnd(ctx.client_id, key, "default", + auto end = service->PutEnd(ctx.client_id, key, TenantId::Default(), ReplicaType::MEMORY); ASSERT_TRUE(end.has_value()); success_puts++; @@ -389,6 +389,41 @@ TEST_F(OffloadOnEvictTest, ComboD_EvictionWorks) { service->RemoveAll(); } +// Regression: EraseMetadata must drop the mirror entry in +// LocalDiskSegment::offloading_objects. Otherwise BatchRemove leaves a +// task-less key in the offload queue which the next +// OffloadObjectHeartbeat drains back to the client, producing an +// orphan bucket on SSD. +TEST_F(OffloadOnEvictTest, BatchRemoveDropsOffloadingObjectsMirror) { + MasterServiceConfig config; + config.enable_offload = true; + config.default_kv_lease_ttl = 0; // no lease so BatchRemove succeeds + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto ctx = + PrepareSegment(*service, "test_segment", kDefaultSegmentBase, seg_size); + auto mount_ld = service->MountLocalDiskSegment(ctx.client_id, true); + ASSERT_TRUE(mount_ld.has_value()); + + const std::vector keys = {"key_r1", "key_r2", "key_r3"}; + for (const auto& k : keys) { + PutObject(*service, ctx.client_id, k); + } + + // Remove before heartbeat drains the offload queue. + auto rm = service->BatchRemove(keys, TenantId::Default(), /*force=*/true); + for (const auto& r : rm) { + EXPECT_TRUE(r.has_value()); + } + + auto queued = DrainOffloadQueue(*service, ctx.client_id); + EXPECT_TRUE(queued.empty()) + << "OffloadObjectHeartbeat returned " << queued.size() + << " stale entries after BatchRemove; EraseMetadata failed to clean " + "offloading_objects."; +} + } // namespace mooncake::test int main(int argc, char** argv) { diff --git a/mooncake-store/tests/offset_allocator_test.cpp b/mooncake-store/tests/offset_allocator_test.cpp index d0e5519e44..abf976a2d5 100644 --- a/mooncake-store/tests/offset_allocator_test.cpp +++ b/mooncake-store/tests/offset_allocator_test.cpp @@ -1,10 +1,13 @@ -#include "offset_allocator/offset_allocator.hpp" +#include "offset_allocator/offset_allocator.h" #include "mutex.h" #include "serializer.h" #include "types.h" #include +#include +#include +#include #include #include #include @@ -283,6 +286,49 @@ class OffsetAllocatorTest : public ::testing::Test { void TearDown() override {} + void lockAllocator(const std::shared_ptr& allocator) { + allocator->m_mutex.lock(); + } + + void unlockAllocator(const std::shared_ptr& allocator) { + allocator->m_mutex.unlock(); + } + + bool isLargestFreeRegionHintTightened( + const std::shared_ptr& allocator) { + MutexLocker lock(&allocator->m_mutex); + return allocator->m_largest_free_region_tightened; + } + + bool canAllocateWithoutFastFail( + const std::shared_ptr& allocator, size_t size) { + MutexLocker lock(&allocator->m_mutex); + const uint64_t quantum = uint64_t{1} << allocator->m_multiplier_bits; + if (size == 0 || + size > std::numeric_limits::max() - (quantum - 1)) { + return false; + } + const uint64_t normalized_size = (size + quantum - 1) / quantum; + if (!allocator->m_allocator || + normalized_size > allocator->m_allocator->m_size) { + return false; + } + + auto allocation = allocator->m_allocator->allocate( + static_cast(normalized_size)); + if (allocation.isNoSpace()) { + return false; + } + allocator->m_allocator->free(allocation); + allocator->refreshLargestFreeRegion(); + return true; + } + + uint64_t multiplierBits( + const std::shared_ptr& allocator) const { + return allocator->m_multiplier_bits; + } + OffsetAllocationHandle copyHandleWithNewAllocator( const OffsetAllocationHandle& handle, const std::shared_ptr& new_allocator) { @@ -306,6 +352,15 @@ class OffsetAllocatorTest : public ::testing::Test { ASSERT_EQ(a->m_capacity, b->m_capacity); ASSERT_EQ(a->m_allocated_size, b->m_allocated_size); ASSERT_EQ(a->m_allocated_num, b->m_allocated_num); + const uint64_t actual_largest_a = + a->m_allocator->storageReport().largestFreeRegion + << a->m_multiplier_bits; + const uint64_t actual_largest_b = + b->m_allocator->storageReport().largestFreeRegion + << b->m_multiplier_bits; + ASSERT_EQ(actual_largest_a, actual_largest_b); + ASSERT_GE(a->getLargestFreeRegion(), actual_largest_a); + ASSERT_GE(b->getLargestFreeRegion(), actual_largest_b); // Compare __Allocator member variables ASSERT_EQ(a->m_allocator->m_size, b->m_allocator->m_size); @@ -516,6 +571,182 @@ TEST_F(OffsetAllocatorTest, AllocationFailure) { EXPECT_FALSE(handle.has_value()); } +TEST_F(OffsetAllocatorTest, FreeReturnsMergedRegionSize) { + constexpr uint32 ALLOCATOR_SIZE = 1024; + constexpr uint32 MAX_ALLOCS = 8; + constexpr uint32 BLOCK_SIZE = 128; + __Allocator allocator(ALLOCATOR_SIZE, MAX_ALLOCS, MAX_ALLOCS); + + const auto first = allocator.allocate(BLOCK_SIZE); + const auto second = allocator.allocate(BLOCK_SIZE); + const auto third = allocator.allocate(BLOCK_SIZE); + ASSERT_FALSE(first.isNoSpace()); + ASSERT_FALSE(second.isNoSpace()); + ASSERT_FALSE(third.isNoSpace()); + + EXPECT_EQ(allocator.free(first), BLOCK_SIZE); + EXPECT_EQ(allocator.free(third), ALLOCATOR_SIZE - 2 * BLOCK_SIZE); + EXPECT_EQ(allocator.free(second), ALLOCATOR_SIZE); +} + +TEST_F(OffsetAllocatorTest, FreeReturnsZeroWhenMetadataCapacityIsExhausted) { + constexpr uint32 ALLOCATOR_SIZE = 384; + constexpr uint32 MAX_ALLOCS = 3; + constexpr uint32 BLOCK_SIZE = 128; + __Allocator allocator(ALLOCATOR_SIZE, MAX_ALLOCS, MAX_ALLOCS); + + const auto first = allocator.allocate(BLOCK_SIZE); + const auto second = allocator.allocate(BLOCK_SIZE); + ASSERT_FALSE(first.isNoSpace()); + ASSERT_FALSE(second.isNoSpace()); + + EXPECT_EQ(allocator.free(first), 0); + EXPECT_EQ(allocator.storageReport().largestFreeRegion, 0); +} + +TEST_F(OffsetAllocatorTest, + LargestFreeRegionHintTightensOnFailureAndRaisesOnFree) { + constexpr uint32 ALLOCATOR_SIZE = 1024 * 1024; + constexpr uint32 MAX_ALLOCS = 1000; + auto allocator = + OffsetAllocator::create(0, ALLOCATOR_SIZE, MAX_ALLOCS, MAX_ALLOCS); + + const uint64_t initial_largest = allocator->getLargestFreeRegion(); + EXPECT_EQ(initial_largest, allocator->storageReport().largestFreeRegion); + + auto handle = allocator->allocate(ALLOCATOR_SIZE / 2); + ASSERT_TRUE(handle.has_value()); + const uint64_t actual_largest = + allocator->storageReport().largestFreeRegion; + ASSERT_LT(actual_largest, initial_largest); + + // A successful allocation can leave the hint conservatively high so the + // success path does not need to publish a new atomic value. + EXPECT_EQ(allocator->getLargestFreeRegion(), initial_largest); + + // The first request that passes the hint but fails in the allocator + // tightens it for subsequent requests. + EXPECT_FALSE(allocator->allocate(actual_largest + 1).has_value()); + EXPECT_EQ(allocator->getLargestFreeRegion(), actual_largest); + + handle.reset(); + EXPECT_EQ(allocator->getLargestFreeRegion(), + allocator->storageReport().largestFreeRegion); + EXPECT_EQ(allocator->getLargestFreeRegion(), ALLOCATOR_SIZE); +} + +TEST_F(OffsetAllocatorTest, HintMaintenanceActivatesOnlyAfterFailure) { + constexpr uint32 ALLOCATOR_SIZE = 1024 * 1024; + constexpr uint32 MAX_ALLOCS = 1000; + auto allocator = + OffsetAllocator::create(0, ALLOCATOR_SIZE, MAX_ALLOCS, MAX_ALLOCS); + + EXPECT_FALSE(isLargestFreeRegionHintTightened(allocator)); + auto handle = allocator->allocate(ALLOCATOR_SIZE / 2); + ASSERT_TRUE(handle.has_value()); + EXPECT_FALSE(isLargestFreeRegionHintTightened(allocator)); + + const uint64_t actual_largest = + allocator->storageReport().largestFreeRegion; + ASSERT_FALSE(allocator->allocate(actual_largest + 1).has_value()); + EXPECT_TRUE(isLargestFreeRegionHintTightened(allocator)); + + handle.reset(); + EXPECT_FALSE(isLargestFreeRegionHintTightened(allocator)); + EXPECT_EQ(allocator->getLargestFreeRegion(), ALLOCATOR_SIZE); +} + +TEST_F(OffsetAllocatorTest, + HintMaintenanceUsesScaledAllocatorCapacityWhenFullyFreed) { + const uint64_t internal_capacity = bin_sizes[NUM_BINS - 1]; + const uint64_t allocator_size = 2 * internal_capacity + 1; + constexpr uint32 MAX_ALLOCS = 1000; + auto allocator = + OffsetAllocator::create(0, allocator_size, MAX_ALLOCS, MAX_ALLOCS); + + const uint64_t maximum_free_region = allocator->getLargestFreeRegion(); + ASSERT_LT(maximum_free_region, allocator_size); + auto handle = allocator->allocate(maximum_free_region / 2); + ASSERT_TRUE(handle.has_value()); + + const uint64_t actual_largest = + allocator->storageReport().largestFreeRegion; + ASSERT_FALSE(allocator->allocate(actual_largest + 1).has_value()); + ASSERT_TRUE(isLargestFreeRegionHintTightened(allocator)); + + handle.reset(); + EXPECT_FALSE(isLargestFreeRegionHintTightened(allocator)); + EXPECT_EQ(allocator->getLargestFreeRegion(), maximum_free_region); +} + +TEST_F(OffsetAllocatorTest, ImpossibleAllocationDoesNotAcquireMutex) { + constexpr uint32 ALLOCATOR_SIZE = 1024 * 1024; + constexpr uint32 MAX_ALLOCS = 1000; + auto allocator = + OffsetAllocator::create(0, ALLOCATOR_SIZE, MAX_ALLOCS, MAX_ALLOCS); + auto full_handle = allocator->allocate(ALLOCATOR_SIZE); + ASSERT_TRUE(full_handle.has_value()); + + // The first failed request observes the exact allocator state and tightens + // the conservative hint. + ASSERT_FALSE(allocator->allocate(1).has_value()); + ASSERT_EQ(allocator->getLargestFreeRegion(), 0); + + lockAllocator(allocator); + std::promise started; + auto started_future = started.get_future(); + auto result = std::async(std::launch::async, [&] { + started.set_value(); + return allocator->allocate(1); + }); + started_future.wait(); + const auto status = result.wait_for(std::chrono::seconds(1)); + unlockAllocator(allocator); + + EXPECT_EQ(status, std::future_status::ready); + EXPECT_FALSE(result.get().has_value()); +} + +TEST_F(OffsetAllocatorTest, FastFailMatchesAllocatorAcrossBinBoundaries) { + constexpr uint32 MAX_ALLOCS = 8; + for (uint32 i = 16; i + 1 < NUM_BINS; ++i) { + const uint64_t capacity = static_cast(bin_sizes[i + 1]) - 1; + const uint64_t request_size = static_cast(bin_sizes[i]) + 1; + auto allocator = + OffsetAllocator::create(0, capacity, MAX_ALLOCS, MAX_ALLOCS); + + ASSERT_EQ(allocator->getLargestFreeRegion(), bin_sizes[i]) + << "bin index=" << i; + ASSERT_FALSE(canAllocateWithoutFastFail(allocator, request_size)) + << "bin index=" << i; + EXPECT_FALSE(allocator->allocate(request_size).has_value()) + << "bin index=" << i; + } +} + +TEST_F(OffsetAllocatorTest, FastFailMatchesAllocatorWithMultipliers) { + constexpr uint32 MAX_ALLOCS = 8; + constexpr uint32 PREVIOUS_BIN_INDEX = NUM_BINS - 2; + const uint64_t INTERNAL_CAPACITY = + static_cast(bin_sizes[NUM_BINS - 1]) - 1; + + for (uint64_t multiplier_bits = 1; multiplier_bits <= 4; + ++multiplier_bits) { + const uint64_t capacity = INTERNAL_CAPACITY << multiplier_bits; + auto allocator = + OffsetAllocator::create(0, capacity, MAX_ALLOCS, MAX_ALLOCS); + const uint64_t largest_free_region = + static_cast(bin_sizes[PREVIOUS_BIN_INDEX]) + << multiplier_bits; + const uint64_t request_size = largest_free_region + 1; + + ASSERT_EQ(multiplierBits(allocator), multiplier_bits); + ASSERT_EQ(allocator->getLargestFreeRegion(), largest_free_region); + ASSERT_FALSE(canAllocateWithoutFastFail(allocator, request_size)); + EXPECT_FALSE(allocator->allocate(request_size).has_value()); + } +} + // Test multiple allocations TEST_F(OffsetAllocatorTest, MultipleAllocations) { constexpr uint32 ALLOCATOR_SIZE = 1024 * 1024 * 1024; // 1GB @@ -1422,6 +1653,29 @@ TEST_F(OffsetAllocatorTest, SerializationOneElementAllocator) { testSerializeAllocator(alloc_a, handles); } +TEST_F(OffsetAllocatorTest, DeserializedHintRaisesWhenAllocationsAreFreed) { + constexpr size_t ALLOCATOR_SIZE = 1024; + constexpr size_t BLOCK_SIZE = ALLOCATOR_SIZE / 2; + constexpr uint32 MAX_ALLOCS = 1000; + auto original = + OffsetAllocator::create(0, ALLOCATOR_SIZE, MAX_ALLOCS, MAX_ALLOCS); + + auto handle = original->allocate(BLOCK_SIZE); + ASSERT_TRUE(handle.has_value()); + + std::vector buffer; + ASSERT_EQ(serialize_to(original, buffer), ErrorCode::OK); + auto restored = deserialize_from(buffer); + ASSERT_NE(restored, nullptr); + + std::optional restored_handle( + copyHandleWithNewAllocator(*handle, restored)); + restored_handle.reset(); + + EXPECT_EQ(restored->getLargestFreeRegion(), ALLOCATOR_SIZE); + EXPECT_TRUE(restored->allocate(ALLOCATOR_SIZE).has_value()); +} + TEST_F(OffsetAllocatorTest, SerializationRandomAllocatedAllocator) { // The size multiplier is larger than 1 when the allocator size is larger // than MAX_BIN_SIZE. diff --git a/mooncake-store/tests/promotion_on_hit_test.cpp b/mooncake-store/tests/promotion_on_hit_test.cpp index b030f4574f..dc14273626 100644 --- a/mooncake-store/tests/promotion_on_hit_test.cpp +++ b/mooncake-store/tests/promotion_on_hit_test.cpp @@ -9,11 +9,17 @@ #include #include +#include +#include +#include #include #include #include #include +#include + +#include "tenant_quota_policy_store.h" #include "types.h" namespace mooncake::test { @@ -32,7 +38,13 @@ class PromotionOnHitTest : public ::testing::Test { FLAGS_logtostderr = true; } - void TearDown() override { google::ShutdownGoogleLogging(); } + void TearDown() override { + for (const auto& path : policy_files_) { + std::error_code ec; + std::filesystem::remove(path, ec); + } + google::ShutdownGoogleLogging(); + } // Friend access to MasterService::promotion_admission_threshold_, which // is otherwise private. PromotionOnHitTest is friended; TEST_F-generated @@ -42,8 +54,77 @@ class PromotionOnHitTest : public ::testing::Test { return service->promotion_admission_threshold_; } + static size_t CountPromotionCandidatesForTesting(MasterService* service, + const TenantId& tenant) { + return service->CountCandidatesForTesting(tenant); + } + + static void ResetCandidateBackoffsForTesting(MasterService* service) { + service->ResetCandidateBackoffsForTesting(); + } + + static size_t RunPromotionCandidateRetryForTesting(MasterService* service) { + return service->RunPromotionCandidateRetryForTesting(); + } + + static size_t RunPromotionCandidateRetryForTesting(MasterService* service, + size_t shards_to_scan) { + return service->RunPromotionCandidateRetry(shards_to_scan); + } + + static void ClearCandidatesForReloadForTesting(MasterService* service) { + service->ClearCandidatesForReload(); + } + + static uint64_t GetPromotionCandidateCountForTesting( + MasterService* service) { + return service->promotion_candidate_count_.load( + std::memory_order_relaxed); + } + + static uint64_t GetPromotionInFlightForTesting(MasterService* service) { + return service->promotion_in_flight_.load(std::memory_order_relaxed); + } + + static bool PromotionAdmissionBlockedByPrimaryWriteForTesting( + MasterService* service, const TenantId& tenant_id, + const std::string& key) { + const auto result = + service->TryPushPromotionQueue(MasterService::ObjectIdentity{ + .tenant_id = tenant_id, .user_key = key}); + return result == MasterService::PromotionQueueResult::kAlreadyInFlight; + } + + static void MarkReplicaCompleteForTesting(MasterService* service, + const TenantId& tenant_id, + const std::string& key, + ReplicaID replica_id) { + MasterService::MetadataAccessorRW accessor( + service, MasterService::ObjectIdentity{.tenant_id = tenant_id, + .user_key = key}); + ASSERT_TRUE(accessor.Exists()); + auto* replica = accessor.Get().GetReplicaByID(replica_id); + ASSERT_NE(replica, nullptr); + replica->mark_complete(); + } + static constexpr size_t kDefaultSegmentBase = 0x300000000; + std::string WriteTenantQuotaPolicyFile( + const std::map& tenant_quotas) { + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = tenant_quotas; + auto path = + std::filesystem::temp_directory_path() / + ("mooncake_promotion_tenant_policy_" + std::to_string(::getpid()) + + "_" + std::to_string(next_policy_file_++) + ".yaml"); + std::ofstream out(path); + out << FormatTenantQuotaPolicyYaml(snapshot); + out.close(); + policy_files_.push_back(path.string()); + return path.string(); + } + Segment MakeSegment(std::string name, size_t base, size_t size) const { Segment segment; segment.id = generate_uuid(); @@ -80,20 +161,20 @@ class PromotionOnHitTest : public ::testing::Test { ReplicateConfig config; config.replica_num = 1; auto put_start = - service.PutStart(client_id, key, "default", size, config); + service.PutStart(client_id, key, TenantId::Default(), size, config); ASSERT_TRUE(put_start.has_value()) << "PutStart failed for key=" << key; - auto put_end = - service.PutEnd(client_id, key, "default", ReplicaType::MEMORY); + auto put_end = service.PutEnd(client_id, key, TenantId::Default(), + ReplicaType::MEMORY); ASSERT_TRUE(put_end.has_value()) << "PutEnd failed for key=" << key; } // Inject a synthetic LOCAL_DISK replica for `key` on `client_id`'s // segment via NotifyOffloadSuccess. Lets tests put a key into // LOCAL_DISK-only state without running the full offload pipeline. - bool InjectLocalDiskReplica(MasterService& service, const UUID& client_id, - const std::string& key, int64_t size, - const std::string& transport_endpoint, - const std::string& tenant_id = "default") { + bool InjectLocalDiskReplica( + MasterService& service, const UUID& client_id, const std::string& key, + int64_t size, const std::string& transport_endpoint, + std::string tenant_id = std::string(TenantId::kDefaultValue)) { std::vector tasks{ OffloadTaskItem{.tenant_id = tenant_id, .key = key, .size = size}}; StorageObjectMetadata sm; @@ -116,6 +197,9 @@ class PromotionOnHitTest : public ::testing::Test { EXPECT_TRUE(mount_ld.has_value()); return client_id; } + + std::vector policy_files_; + size_t next_policy_file_ = 0; }; // Sanity: with promotion disabled, no path mutates promotion_objects. @@ -134,7 +218,7 @@ TEST_F(PromotionOnHitTest, DefaultOffNoPromotion) { // GetReplicaList many times. With promotion_on_hit=false, nothing should // appear in promotion_objects regardless of access count. for (int i = 0; i < 5; ++i) { - auto resp = service->GetReplicaList("k1", "default"); + auto resp = service->GetReplicaList("k1", TenantId::Default()); ASSERT_TRUE(resp.has_value()); } @@ -161,7 +245,7 @@ TEST_F(PromotionOnHitTest, NoLocalDiskNoPromotion) { PutObject(*service, ctx.client_id, "k_mem_only"); for (int i = 0; i < 5; ++i) { - auto resp = service->GetReplicaList("k_mem_only", "default"); + auto resp = service->GetReplicaList("k_mem_only", TenantId::Default()); ASSERT_TRUE(resp.has_value()); } @@ -193,7 +277,7 @@ TEST_F(PromotionOnHitTest, MemoryReplicaPresentNoPromotion) { ctx.segment_name)); for (int i = 0; i < 5; ++i) { - auto resp = service->GetReplicaList("k_dual", "default"); + auto resp = service->GetReplicaList("k_dual", TenantId::Default()); ASSERT_TRUE(resp.has_value()); } @@ -229,13 +313,14 @@ TEST_F(PromotionOnHitTest, BatchGetReplicaListPromotesLocalDiskOnlyObject) { const int64_t admitted_pre = mm.get_promotion_admitted(); const int64_t in_flight_pre = mm.get_promotion_in_flight(); - auto single_result = service->GetReplicaList(single_key, "default"); + auto single_result = + service->GetReplicaList(single_key, TenantId::Default()); ASSERT_TRUE(single_result.has_value()); ASSERT_EQ(single_result->replicas.size(), 1u); EXPECT_TRUE(single_result->replicas[0].is_local_disk_replica()); auto batch_result = service->BatchGetReplicaList( - std::vector{batch_key}, "default"); + std::vector{batch_key}, TenantId::Default()); ASSERT_EQ(batch_result.size(), 1u); ASSERT_TRUE(batch_result[0].has_value()); ASSERT_EQ(batch_result[0]->replicas.size(), 1u); @@ -253,6 +338,91 @@ TEST_F(PromotionOnHitTest, BatchGetReplicaListPromotesLocalDiskOnlyObject) { service->RemoveAll(); } +// The read-only admin batch query must NOT trigger promotion-on-hit, even for a +// LOCAL_DISK-only key. Contrast with BatchGetReplicaListPromotesLocalDiskOnly +// Object above, where the client-facing BatchGetReplicaList admits the key for +// promotion. +TEST_F(PromotionOnHitTest, + BatchGetReplicaListForAdminDoesNotPromoteLocalDiskOnlyObject) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + config.promotion_max_per_heartbeat = 2; + config.default_kv_lease_ttl = 2000; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto ctx = PrepareSegment(*service, "test_segment_admin", + kDefaultSegmentBase, seg_size); + + const std::string key = "k_batch_admin_no_promote"; + ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, key, 1024, + ctx.segment_name)); + + auto& mm = MasterMetricManager::instance(); + const int64_t admitted_pre = mm.get_promotion_admitted(); + const int64_t in_flight_pre = mm.get_promotion_in_flight(); + + auto result = service->BatchGetReplicaListForAdmin( + std::vector{key}, TenantId::Default()); + ASSERT_EQ(result.size(), 1u); + ASSERT_TRUE(result[0].has_value()); + ASSERT_EQ(result[0]->replicas.size(), 1u); + EXPECT_TRUE(result[0]->replicas[0].is_local_disk_replica()); + + // No promotion admitted or enqueued by the read-only admin path. + EXPECT_EQ(mm.get_promotion_admitted() - admitted_pre, 0); + EXPECT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 0); + auto pending = service->PromotionObjectHeartbeat(ctx.client_id); + ASSERT_TRUE(pending.has_value()); + EXPECT_EQ(pending->size(), 0u); + EXPECT_EQ(CountPromotionTask(*pending, key), 0u); + + service->RemoveAll(); +} + +// The read-only admin batch query must NOT update the store-observed cache-hit +// counters. Contrast with the client-facing BatchGetReplicaList, which bumps +// the memory-cache-hit counter for a MEMORY replica. +TEST_F(PromotionOnHitTest, + BatchGetReplicaListForAdminDoesNotUpdateCacheHitMetrics) { + MasterServiceConfig config; + config.enable_offload = true; + config.default_kv_lease_ttl = 2000; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto ctx = PrepareSegment(*service, "metrics_segment", kDefaultSegmentBase, + seg_size); + + const std::string key = "k_admin_no_metric"; + PutObject(*service, ctx.client_id, key, 1024); + + using CacheHitStat = MasterMetricManager::CacheHitStat; + auto mem_hits = []() { + auto stats = MasterMetricManager::instance().calculate_cache_stats(); + return stats[CacheHitStat::MEMORY_HITS]; + }; + + const double before_admin = mem_hits(); + + // Read-only admin query must leave the memory-cache-hit counter untouched. + auto admin_result = service->BatchGetReplicaListForAdmin( + std::vector{key}, TenantId::Default()); + ASSERT_EQ(admin_result.size(), 1u); + ASSERT_TRUE(admin_result[0].has_value()); + EXPECT_EQ(mem_hits(), before_admin); + + // The client-facing path does bump it, proving the assertion above is + // meaningful rather than a counter that never moves. + (void)service->BatchGetReplicaList(std::vector{key}, + TenantId::Default()); + EXPECT_GT(mem_hits(), before_admin); + + service->RemoveAll(); +} + // PromotionObjectHeartbeat returns an empty task list when called against a // client that has no LocalDiskSegment registered. TEST_F(PromotionOnHitTest, HeartbeatReturnsErrorForUnknownClient) { @@ -275,11 +445,173 @@ TEST_F(PromotionOnHitTest, AllocStartUnknownKey) { auto service = std::make_unique(config); auto resp = service->PromotionAllocStart(generate_uuid(), "nonexistent", - "default", 1024, {}); + TenantId::Default(), 1024, {}); ASSERT_FALSE(resp.has_value()); EXPECT_EQ(resp.error(), ErrorCode::OBJECT_NOT_FOUND); } +TEST_F(PromotionOnHitTest, InvalidPrimaryEndCannotCompletePromotionReplica) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + Segment holder_segment = + MakeSegment("issue_3203_end", kDefaultSegmentBase, seg_size); + const UUID holder_id = generate_uuid(); + ASSERT_TRUE(service->MountSegment(holder_segment, holder_id).has_value()); + ASSERT_TRUE(service->MountLocalDiskSegment(holder_id, true).has_value()); + ASSERT_TRUE( + service->ReMountSegment({holder_segment}, holder_id).has_value()); + ASSERT_TRUE(InjectLocalDiskReplica(*service, holder_id, "k_cold", 1024, + holder_segment.name)); + + ReplicateConfig upsert_config; + upsert_config.replica_num = 1; + ASSERT_TRUE(service + ->UpsertStart(holder_id, "k_cold", TenantId::Default(), + 1024, upsert_config) + .has_value()); + ASSERT_TRUE(service + ->UpsertEnd(holder_id, "k_cold", TenantId::Default(), + ReplicaType::LOCAL_DISK) + .has_value()); + + auto read = service->GetReplicaList("k_cold", TenantId::Default()); + ASSERT_TRUE(read.has_value()); + auto alloc = service->PromotionAllocStart(holder_id, "k_cold", + TenantId::Default(), 1024, {}); + ASSERT_TRUE(alloc.has_value()); + + // Retrying the completed End is a no-op: it succeeds, and the promotion + // still owns its PROCESSING MEMORY replica and task. + ASSERT_TRUE(service + ->UpsertEnd(holder_id, "k_cold", TenantId::Default(), + ReplicaType::LOCAL_DISK) + .has_value()); + + // There was no primary PutStart. PutEnd must not complete the staged + // promotion replica even though the metadata client_id matches. + auto invalid_put_end = service->PutEnd( + holder_id, "k_cold", TenantId::Default(), ReplicaType::MEMORY); + ASSERT_FALSE(invalid_put_end.has_value()); + EXPECT_EQ(invalid_put_end.error(), ErrorCode::INVALID_WRITE); + + auto rejected_upsert = service->UpsertStart( + holder_id, "k_cold", TenantId::Default(), 1024, upsert_config); + ASSERT_FALSE(rejected_upsert.has_value()); + EXPECT_EQ(rejected_upsert.error(), ErrorCode::OBJECT_HAS_REPLICATION_TASK); + + // A caller must not be able to finalize after the rejected UpsertStart. + auto invalid_upsert_end = service->UpsertEnd( + holder_id, "k_cold", TenantId::Default(), ReplicaType::MEMORY); + ASSERT_FALSE(invalid_upsert_end.has_value()); + EXPECT_EQ(invalid_upsert_end.error(), ErrorCode::INVALID_WRITE); + + // Only the promotion completion path owns the staged replica. + auto notify = service->NotifyPromotionSuccess(holder_id, "k_cold", + TenantId::Default()); + ASSERT_TRUE(notify.has_value()); + + service->RemoveAll(); +} + +TEST_F(PromotionOnHitTest, PrimaryWriteBlocksPromotionAdmission) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto holder = PrepareSegment(*service, "issue_3203_primary", + kDefaultSegmentBase, seg_size); + ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_cold", + 1024, holder.segment_name)); + + ReplicateConfig upsert_config; + upsert_config.replica_num = 1; + auto upsert = service->UpsertStart( + holder.client_id, "k_cold", TenantId::Default(), 1024, upsert_config); + ASSERT_TRUE(upsert.has_value()); + + EXPECT_TRUE(PromotionAdmissionBlockedByPrimaryWriteForTesting( + service.get(), TenantId::Default(), "k_cold")); + EXPECT_EQ(GetPromotionInFlightForTesting(service.get()), 0u); + + auto end = service->UpsertEnd(holder.client_id, "k_cold", + TenantId::Default(), ReplicaType::LOCAL_DISK); + ASSERT_TRUE(end.has_value()); + + service->RemoveAll(); +} + +TEST_F(PromotionOnHitTest, StalePromotionReplicaCleanupErasesTask) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + Segment holder_segment = + MakeSegment("issue_3203_holder", kDefaultSegmentBase, seg_size); + const UUID holder_id = generate_uuid(); + ASSERT_TRUE(service->MountSegment(holder_segment, holder_id).has_value()); + ASSERT_TRUE(service->MountLocalDiskSegment(holder_id, true).has_value()); + ASSERT_TRUE( + service->ReMountSegment({holder_segment}, holder_id).has_value()); + + Segment target_segment = MakeSegment( + "issue_3203_target", kDefaultSegmentBase + seg_size, seg_size); + const UUID target_id = generate_uuid(); + ASSERT_TRUE(service->MountSegment(target_segment, target_id).has_value()); + + ASSERT_TRUE(InjectLocalDiskReplica(*service, holder_id, "k_cold", 1024, + holder_segment.name)); + auto read = service->GetReplicaList("k_cold", TenantId::Default()); + ASSERT_TRUE(read.has_value()); + auto alloc = service->PromotionAllocStart( + holder_id, "k_cold", TenantId::Default(), 1024, {target_segment.name}); + ASSERT_TRUE(alloc.has_value()); + ASSERT_EQ(alloc->memory_descriptor.get_memory_descriptor() + .buffer_descriptor.transport_endpoint_, + target_segment.name); + + // Recreate the legacy bad state: an invalid primary End used to mark the + // promotion-owned replica COMPLETE before stale-handle cleanup removed it. + MarkReplicaCompleteForTesting(service.get(), TenantId::Default(), "k_cold", + alloc->memory_descriptor.id); + + auto& metrics = MasterMetricManager::instance(); + const int64_t cancelled_before = metrics.get_promotion_cancelled(); + ASSERT_EQ(GetPromotionInFlightForTesting(service.get()), 1u); + + ASSERT_TRUE( + service->UnmountSegment(target_segment.id, target_id).has_value()); + + EXPECT_EQ(GetPromotionInFlightForTesting(service.get()), 0u); + EXPECT_EQ(metrics.get_promotion_cancelled() - cancelled_before, 1); + auto pending = service->PromotionObjectHeartbeat(holder_id); + ASSERT_TRUE(pending.has_value()); + EXPECT_TRUE(pending->empty()); + + auto notify = service->NotifyPromotionSuccess(holder_id, "k_cold", + TenantId::Default()); + ASSERT_FALSE(notify.has_value()); + EXPECT_EQ(notify.error(), ErrorCode::REPLICA_IS_NOT_READY); + + auto remaining = + service->GetReplicaListForAdmin("k_cold", TenantId::Default()); + ASSERT_TRUE(remaining.has_value()); + ASSERT_EQ(remaining->replicas.size(), 1u); + EXPECT_TRUE(remaining->replicas[0].is_local_disk_replica()); + + service->RemoveAll(); +} + // NotifyPromotionSuccess on a non-existent key returns OBJECT_NOT_FOUND. TEST_F(PromotionOnHitTest, NotifyUnknownKey) { MasterServiceConfig config; @@ -288,8 +620,8 @@ TEST_F(PromotionOnHitTest, NotifyUnknownKey) { auto service = std::make_unique(config); UUID client_id = generate_uuid(); - auto resp = - service->NotifyPromotionSuccess(client_id, "nonexistent", "default"); + auto resp = service->NotifyPromotionSuccess(client_id, "nonexistent", + TenantId::Default()); ASSERT_FALSE(resp.has_value()); EXPECT_EQ(resp.error(), ErrorCode::OBJECT_NOT_FOUND); } @@ -324,7 +656,7 @@ TEST_F(PromotionOnHitTest, RacingReadersDedup) { for (int t = 0; t < kThreads; ++t) { threads.emplace_back([&service]() { for (int j = 0; j < kReadsPerThread; ++j) { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); EXPECT_TRUE(r.has_value()); } }); @@ -374,7 +706,7 @@ TEST_F(PromotionOnHitTest, StalePromotionReaper) { // the per-shard PromotionTask intact (the heartbeat is best-effort GC, // not the authoritative state). { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); auto pending = service->PromotionObjectHeartbeat(ctx.client_id); ASSERT_TRUE(pending.has_value()); @@ -384,7 +716,7 @@ TEST_F(PromotionOnHitTest, StalePromotionReaper) { // Without reap, dedup blocks re-enqueue. Confirm: GetReplicaList again, // heartbeat must be empty because PromotionTask still pins the slot. { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); auto pending = service->PromotionObjectHeartbeat(ctx.client_id); ASSERT_TRUE(pending.has_value()); @@ -403,7 +735,7 @@ TEST_F(PromotionOnHitTest, StalePromotionReaper) { // Trigger #3: with the task reaped, dedup is unblocked and a fresh // GetReplicaList must enqueue again. { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); auto pending = service->PromotionObjectHeartbeat(ctx.client_id); ASSERT_TRUE(pending.has_value()); @@ -443,12 +775,12 @@ TEST_F(PromotionOnHitTest, RemoveDuringPromotion) { // Queue a promotion task. { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } // Lease must expire before we can call non-force Remove (or use force). - auto rm = service->Remove("k_cold", "default", /*force=*/true); + auto rm = service->Remove("k_cold", TenantId::Default(), /*force=*/true); // Remove returns REPLICA_IS_NOT_READY if any replica is non-COMPLETE. // The injected LOCAL_DISK replica is COMPLETE, so this should succeed. ASSERT_TRUE(rm.has_value()) @@ -460,8 +792,8 @@ TEST_F(PromotionOnHitTest, RemoveDuringPromotion) { // NotifyPromotionSuccess on the now-removed key must surface the missing // metadata cleanly, not crash. - auto notify = - service->NotifyPromotionSuccess(ctx.client_id, "k_cold", "default"); + auto notify = service->NotifyPromotionSuccess(ctx.client_id, "k_cold", + TenantId::Default()); ASSERT_FALSE(notify.has_value()); EXPECT_EQ(notify.error(), ErrorCode::OBJECT_NOT_FOUND); @@ -475,7 +807,7 @@ TEST_F(PromotionOnHitTest, RemoveDuringPromotion) { ASSERT_TRUE(InjectLocalDiskReplica(*service, ctx.client_id, "k_cold", 1024, ctx.segment_name)); { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); auto pending = service->PromotionObjectHeartbeat(ctx.client_id); ASSERT_TRUE(pending.has_value()); @@ -515,14 +847,14 @@ TEST_F(PromotionOnHitTest, MultiSegmentAllocPicksAvailableSegment) { // Seed the PromotionTask through the gate — PromotionAllocStart now // requires an in-flight task to exist (rejects orphaned-stage path). { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } // PromotionAllocStart — test the segment-selection logic on top of // the gate-seeded task. auto resp = service->PromotionAllocStart(holder_client_id, "k_cold", - "default", 1024, {}); + TenantId::Default(), 1024, {}); ASSERT_TRUE(resp.has_value()) << "PromotionAllocStart should succeed when any DRAM segment has " << "capacity; error=" << resp.error(); @@ -565,12 +897,13 @@ TEST_F(PromotionOnHitTest, MultiSegmentAllocRespectsPreferred) { // Seed the PromotionTask through the gate so AllocStart's // task-existence check passes. { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } - auto resp = service->PromotionAllocStart( - seg_a.client_id, "k_cold", "default", 1024, {seg_b.segment_name}); + auto resp = service->PromotionAllocStart(seg_a.client_id, "k_cold", + TenantId::Default(), 1024, + {seg_b.segment_name}); ASSERT_TRUE(resp.has_value()); const auto& mem_desc = resp.value().memory_descriptor; EXPECT_EQ( @@ -629,13 +962,13 @@ TEST_F(PromotionOnHitTest, QueueLimitRejectsBeyondCap) { const int64_t cap_rej_pre = mm.get_promotion_rejected_cap(); // First read on k1 enqueues a task in shard S. - auto r1 = service->GetReplicaList(k1, "default"); + auto r1 = service->GetReplicaList(k1, TenantId::Default()); ASSERT_TRUE(r1.has_value()); // Second read on k2 (same shard S, different key, so no dedup) must // be dropped by the cap gate: the cluster-wide in-flight counter is // already 1, which meets promotion_queue_limit_ = 1. - auto r2 = service->GetReplicaList(k2, "default"); + auto r2 = service->GetReplicaList(k2, TenantId::Default()); ASSERT_TRUE(r2.has_value()) << "read itself must still succeed; " << "queue gate is silent"; @@ -680,7 +1013,7 @@ TEST_F(PromotionOnHitTest, HeartbeatBoundedBatchPreservesLeftovers) { for (const auto& k : keys) { ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, k, 1024, seg.segment_name)); - auto r = service->GetReplicaList(k, "default"); + auto r = service->GetReplicaList(k, TenantId::Default()); ASSERT_TRUE(r.has_value()); } @@ -719,7 +1052,7 @@ TEST_F(PromotionOnHitTest, HeartbeatBoundedBatchPreservesLeftovers) { // (they're cleared by NotifyPromotionSuccess, not by Heartbeat), so the // source refcnts remain pinned until processed. for (const auto& k : keys) { - auto rl = service->GetReplicaList(k, "default"); + auto rl = service->GetReplicaList(k, TenantId::Default()); ASSERT_TRUE(rl.has_value()) << "key " << k << " should still exist"; } @@ -761,14 +1094,14 @@ TEST_F(PromotionOnHitTest, ReaperPopsStagedMemoryReplicaOnExpiry) { // Trigger the gate to enqueue a PromotionTask. { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } // Drive the AllocStart side so alloc_id != 0 — this is the exact // setup that produces an orphaned PROCESSING MEMORY replica if the // reaper does not pop it. auto alloc = service->PromotionAllocStart(ctx.client_id, "k_cold", - "default", 1024, {}); + TenantId::Default(), 1024, {}); ASSERT_TRUE(alloc.has_value()); // After AllocStart, the DRAM allocator must have committed bytes for @@ -806,8 +1139,8 @@ TEST_F(PromotionOnHitTest, ReaperPopsStagedMemoryReplicaOnExpiry) { // and must return REPLICA_IS_NOT_READY (the task entry is gone, so // the alloc_id lookup at the top of NotifyPromotionSuccess fails // fast). - auto notify = - service->NotifyPromotionSuccess(ctx.client_id, "k_cold", "default"); + auto notify = service->NotifyPromotionSuccess(ctx.client_id, "k_cold", + TenantId::Default()); ASSERT_FALSE(notify.has_value()); EXPECT_EQ(notify.error(), ErrorCode::REPLICA_IS_NOT_READY); @@ -855,12 +1188,12 @@ TEST_F(PromotionOnHitTest, QueueLimitRejectsCrossShard) { ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, k2, 1024, seg.segment_name)); - auto r1 = service->GetReplicaList(k1, "default"); + auto r1 = service->GetReplicaList(k1, TenantId::Default()); ASSERT_TRUE(r1.has_value()); // k2 lives in a different shard, but the global cap is already met // by k1's task — k2 must be rejected. - auto r2 = service->GetReplicaList(k2, "default"); + auto r2 = service->GetReplicaList(k2, TenantId::Default()); ASSERT_TRUE(r2.has_value()) << "read itself still succeeds"; auto heartbeat = service->PromotionObjectHeartbeat(seg.client_id); @@ -934,7 +1267,7 @@ TEST_F(PromotionOnHitTest, AllocStartResetsTaskDeadline) { // T=0 : admit. start_time = T=0. { - auto r = service->GetReplicaList("k_late", "default"); + auto r = service->GetReplicaList("k_late", TenantId::Default()); ASSERT_TRUE(r.has_value()); } @@ -944,7 +1277,7 @@ TEST_F(PromotionOnHitTest, AllocStartResetsTaskDeadline) { std::this_thread::sleep_for(std::chrono::milliseconds(1500)); auto alloc = service->PromotionAllocStart(ctx.client_id, "k_late", - "default", 1024, {}); + TenantId::Default(), 1024, {}); ASSERT_TRUE(alloc.has_value()) << "AllocStart must succeed before the queue-wait phase's TTL " << "expires (1.5s elapsed, TTL is 2s). If this fires the test " @@ -1014,7 +1347,7 @@ TEST_F(PromotionOnHitTest, NotifySuccessDecrementsCounter) { // Admit task 1. promotion_in_flight_ goes from 0 -> 1. { - auto r = service->GetReplicaList("k_first", "default"); + auto r = service->GetReplicaList("k_first", TenantId::Default()); ASSERT_TRUE(r.has_value()); } @@ -1023,11 +1356,11 @@ TEST_F(PromotionOnHitTest, NotifySuccessDecrementsCounter) { // COMPLETE, drops the source LOCAL_DISK refcnt, erases the task, and // decrements the counter. auto alloc = service->PromotionAllocStart(seg.client_id, "k_first", - "default", 1024, {}); + TenantId::Default(), 1024, {}); ASSERT_TRUE(alloc.has_value()) << "AllocStart should succeed; error=" << alloc.error(); - auto notify = - service->NotifyPromotionSuccess(seg.client_id, "k_first", "default"); + auto notify = service->NotifyPromotionSuccess(seg.client_id, "k_first", + TenantId::Default()); ASSERT_TRUE(notify.has_value()) << "NotifyPromotionSuccess happy path should succeed; if this " << "fires, AllocStart did not record alloc_id, or the staged " @@ -1038,7 +1371,7 @@ TEST_F(PromotionOnHitTest, NotifySuccessDecrementsCounter) { // still saturated at 1 and TryPushPromotionQueue silently drops this // attempt. { - auto r = service->GetReplicaList("k_second", "default"); + auto r = service->GetReplicaList("k_second", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto pending = service->PromotionObjectHeartbeat(seg.client_id); @@ -1083,7 +1416,7 @@ TEST_F(PromotionOnHitTest, AllocStartRejectsReapedTask) { // Admit the task. { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } @@ -1101,7 +1434,7 @@ TEST_F(PromotionOnHitTest, AllocStartRejectsReapedTask) { // Allocating would leave an orphaned PROCESSING MEMORY replica // attached to the object. auto alloc = service->PromotionAllocStart(ctx.client_id, "k_cold", - "default", 1024, {}); + TenantId::Default(), 1024, {}); ASSERT_FALSE(alloc.has_value()) << "AllocStart must reject when the task has been reaped — " << "otherwise the staged PROCESSING MEMORY replica is orphaned"; @@ -1140,19 +1473,19 @@ TEST_F(PromotionOnHitTest, NotifyRejectsNonHolder) { // Admit + stage. { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto alloc = service->PromotionAllocStart(holder.client_id, "k_cold", - "default", 1024, {}); + TenantId::Default(), 1024, {}); ASSERT_TRUE(alloc.has_value()); // An unrelated client tries to Notify. Must be rejected as // INVALID_PARAMS so the staged replica stays PROCESSING. UUID intruder_id = generate_uuid(); ASSERT_NE(intruder_id, holder.client_id); - auto bad_notify = - service->NotifyPromotionSuccess(intruder_id, "k_cold", "default"); + auto bad_notify = service->NotifyPromotionSuccess(intruder_id, "k_cold", + TenantId::Default()); ASSERT_FALSE(bad_notify.has_value()) << "Notify from a non-holder client must be rejected — otherwise " << "any client knowing the key can commit someone else's " @@ -1164,7 +1497,7 @@ TEST_F(PromotionOnHitTest, NotifyRejectsNonHolder) { // GetReplicaList filters PROCESSING replicas, and pre-AllocStart // there are no COMPLETE replicas to read for a LOCAL_DISK-only key // (only the LOCAL_DISK descriptor itself). - auto r_check = service->GetReplicaList("k_cold", "default"); + auto r_check = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r_check.has_value()); bool saw_complete_memory = false; for (const auto& d : r_check.value().replicas) { @@ -1177,8 +1510,8 @@ TEST_F(PromotionOnHitTest, NotifyRejectsNonHolder) { << "Rejected Notify must not have committed the staged replica"; // The legitimate holder must still be able to Notify successfully. - auto good_notify = - service->NotifyPromotionSuccess(holder.client_id, "k_cold", "default"); + auto good_notify = service->NotifyPromotionSuccess( + holder.client_id, "k_cold", TenantId::Default()); ASSERT_TRUE(good_notify.has_value()) << "Holder Notify on the same task must succeed after a rejected " << "intruder Notify — the task entry should be untouched by the " @@ -1224,11 +1557,11 @@ TEST_F(PromotionOnHitTest, NotifyFailureReleasesStateImmediately) { // Admit + stage k_a. promotion_in_flight_ goes 0 -> 1. { - auto r = service->GetReplicaList("k_a", "default"); + auto r = service->GetReplicaList("k_a", TenantId::Default()); ASSERT_TRUE(r.has_value()); } - auto alloc = - service->PromotionAllocStart(seg.client_id, "k_a", "default", 1024, {}); + auto alloc = service->PromotionAllocStart(seg.client_id, "k_a", + TenantId::Default(), 1024, {}); ASSERT_TRUE(alloc.has_value()); // The staged PROCESSING MEMORY buffer is allocated. @@ -1240,8 +1573,8 @@ TEST_F(PromotionOnHitTest, NotifyFailureReleasesStateImmediately) { // Holder reports failure (simulating SSD read error after AllocStart // succeeded). Master must immediately reap the staged replica and // decrement the slot counter. - auto failure = - service->NotifyPromotionFailure(seg.client_id, "k_a", "default"); + auto failure = service->NotifyPromotionFailure(seg.client_id, "k_a", + TenantId::Default()); ASSERT_TRUE(failure.has_value()) << "NotifyPromotionFailure on a valid in-flight task from the " << "legitimate holder must succeed; error=" << failure.error(); @@ -1263,7 +1596,7 @@ TEST_F(PromotionOnHitTest, NotifyFailureReleasesStateImmediately) { // succeed even though queue_limit=1. Without the failure-side // decrement the cap would stay saturated until reaper TTL. { - auto r = service->GetReplicaList("k_b", "default"); + auto r = service->GetReplicaList("k_b", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto heartbeat = service->PromotionObjectHeartbeat(seg.client_id); @@ -1276,8 +1609,8 @@ TEST_F(PromotionOnHitTest, NotifyFailureReleasesStateImmediately) { // Idempotency: repeated failure notification on the same key must be // safe (return OK without underflowing the counter). - auto failure_again = - service->NotifyPromotionFailure(seg.client_id, "k_a", "default"); + auto failure_again = service->NotifyPromotionFailure(seg.client_id, "k_a", + TenantId::Default()); EXPECT_TRUE(failure_again.has_value()) << "Repeated NotifyPromotionFailure should be idempotent (return " << "OK on already-released task), not error."; @@ -1306,18 +1639,18 @@ TEST_F(PromotionOnHitTest, NotifyFailureRejectsNonHolder) { 1024, holder.segment_name)); { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto alloc = service->PromotionAllocStart(holder.client_id, "k_cold", - "default", 1024, {}); + TenantId::Default(), 1024, {}); ASSERT_TRUE(alloc.has_value()); // Intruder calls Failure with the wrong client_id. UUID intruder_id = generate_uuid(); ASSERT_NE(intruder_id, holder.client_id); - auto bad_failure = - service->NotifyPromotionFailure(intruder_id, "k_cold", "default"); + auto bad_failure = service->NotifyPromotionFailure(intruder_id, "k_cold", + TenantId::Default()); ASSERT_FALSE(bad_failure.has_value()) << "Failure from a non-holder client must be rejected."; EXPECT_EQ(bad_failure.error(), ErrorCode::INVALID_PARAMS); @@ -1325,8 +1658,8 @@ TEST_F(PromotionOnHitTest, NotifyFailureRejectsNonHolder) { // The legitimate holder must still be able to either commit via // Notify-Success or release via Notify-Failure on the same task. Use // Notify-Failure here to exercise the surviving-task path. - auto good_failure = - service->NotifyPromotionFailure(holder.client_id, "k_cold", "default"); + auto good_failure = service->NotifyPromotionFailure( + holder.client_id, "k_cold", TenantId::Default()); ASSERT_TRUE(good_failure.has_value()) << "Holder Failure must succeed after a rejected intruder " << "Failure — the task entry should be untouched by the " @@ -1360,14 +1693,14 @@ TEST_F(PromotionOnHitTest, AllocStartRejectsNonHolder) { const size_t used_baseline = seg_baseline->first; { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } UUID intruder_id = generate_uuid(); ASSERT_NE(intruder_id, holder.client_id); - auto bad_alloc = service->PromotionAllocStart(intruder_id, "k_cold", - "default", 1024, {}); + auto bad_alloc = service->PromotionAllocStart( + intruder_id, "k_cold", TenantId::Default(), 1024, {}); ASSERT_FALSE(bad_alloc.has_value()) << "AllocStart from a non-holder client must be rejected — " << "otherwise an attacker that drained another's queue could " @@ -1382,8 +1715,8 @@ TEST_F(PromotionOnHitTest, AllocStartRejectsNonHolder) { // The legitimate holder must still be able to AllocStart (task // untouched by the rejection). - auto good_alloc = service->PromotionAllocStart(holder.client_id, "k_cold", - "default", 1024, {}); + auto good_alloc = service->PromotionAllocStart( + holder.client_id, "k_cold", TenantId::Default(), 1024, {}); ASSERT_TRUE(good_alloc.has_value()) << "Holder AllocStart on the same task must succeed after a " << "rejected intruder AllocStart; error=" << good_alloc.error(); @@ -1417,7 +1750,7 @@ TEST_F(PromotionOnHitTest, AllocStartRejectsSizeMismatch) { const size_t used_baseline = seg_baseline->first; { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } @@ -1425,7 +1758,7 @@ TEST_F(PromotionOnHitTest, AllocStartRejectsSizeMismatch) { for (uint64_t bad_size : {static_cast(kRealSize) / 2, static_cast(kRealSize) * 4}) { auto bad_alloc = service->PromotionAllocStart( - holder.client_id, "k_cold", "default", bad_size, {}); + holder.client_id, "k_cold", TenantId::Default(), bad_size, {}); ASSERT_FALSE(bad_alloc.has_value()) << "AllocStart with size=" << bad_size << " (task.object_size=" << kRealSize << ") must be rejected."; @@ -1441,9 +1774,9 @@ TEST_F(PromotionOnHitTest, AllocStartRejectsSizeMismatch) { // Correct size must still work — the task was not consumed by the // rejections. - auto good_alloc = - service->PromotionAllocStart(holder.client_id, "k_cold", "default", - static_cast(kRealSize), {}); + auto good_alloc = service->PromotionAllocStart( + holder.client_id, "k_cold", TenantId::Default(), + static_cast(kRealSize), {}); ASSERT_TRUE(good_alloc.has_value()) << "AllocStart with the correct size must succeed after rejected " << "size-mismatch attempts; error=" << good_alloc.error(); @@ -1488,7 +1821,7 @@ TEST_F(PromotionOnHitTest, ClientExpiryClearsPromotionTask) { // Admit the promotion. promotion_in_flight_ goes 0 -> 1. { - auto r = service->GetReplicaList("k_cold", "default"); + auto r = service->GetReplicaList("k_cold", TenantId::Default()); ASSERT_TRUE(r.has_value()); } @@ -1516,7 +1849,7 @@ TEST_F(PromotionOnHitTest, ClientExpiryClearsPromotionTask) { "k_other", 1024, second_holder.segment_name)); { - auto r = service->GetReplicaList("k_other", "default"); + auto r = service->GetReplicaList("k_other", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto pending_pre = @@ -1549,7 +1882,7 @@ TEST_F(PromotionOnHitTest, ClientExpiryClearsPromotionTask) { // on the second holder; with queue_limit=1 this can only succeed if // the slot was freed. { - auto r = service->GetReplicaList("k_other", "default"); + auto r = service->GetReplicaList("k_other", TenantId::Default()); ASSERT_TRUE(r.has_value()) << "GetReplicaList(k_other) failed with error=" << r.error(); } @@ -1631,7 +1964,7 @@ TEST_F(PromotionOnHitTest, RemoveErasesPromotionTask) { // Admit task 1. { - auto r = service->GetReplicaList("k_first", "default"); + auto r = service->GetReplicaList("k_first", TenantId::Default()); ASSERT_TRUE(r.has_value()); } { @@ -1643,7 +1976,7 @@ TEST_F(PromotionOnHitTest, RemoveErasesPromotionTask) { // Remove k_first with force=true. With the fix, this also wipes // k_first's promotion_tasks entry and decrements // promotion_in_flight_ back to 0. - auto rm = service->Remove("k_first", "default", /*force=*/true); + auto rm = service->Remove("k_first", TenantId::Default(), /*force=*/true); ASSERT_TRUE(rm.has_value()) << "Remove should succeed; error=" << rm.error(); @@ -1651,7 +1984,7 @@ TEST_F(PromotionOnHitTest, RemoveErasesPromotionTask) { // if the slot was freed by Remove. Without the in-flight cleanup, // it would stay pinned for the full 300s TTL. { - auto r = service->GetReplicaList("k_second", "default"); + auto r = service->GetReplicaList("k_second", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto pending_post = service->PromotionObjectHeartbeat(holder.client_id); @@ -1688,12 +2021,13 @@ TEST_F(PromotionOnHitTest, RemoveByRegexErasesPromotionTask) { // Admit task on regex_k1. { - auto r = service->GetReplicaList("regex_k1", "default"); + auto r = service->GetReplicaList("regex_k1", TenantId::Default()); ASSERT_TRUE(r.has_value()); } // RemoveByRegex matches regex_k1 only. - auto removed = service->RemoveByRegex("^regex_", "default", /*force=*/true); + auto removed = + service->RemoveByRegex("^regex_", TenantId::Default(), /*force=*/true); ASSERT_TRUE(removed.has_value()) << "RemoveByRegex should succeed; error=" << removed.error(); EXPECT_EQ(removed.value(), 1) << "exactly one key (regex_k1) should match"; @@ -1701,7 +2035,7 @@ TEST_F(PromotionOnHitTest, RemoveByRegexErasesPromotionTask) { // Slot must be free — admit on other_k2 (different shard or same, // doesn't matter because counter is global). { - auto r = service->GetReplicaList("other_k2", "default"); + auto r = service->GetReplicaList("other_k2", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto pending_post = service->PromotionObjectHeartbeat(holder.client_id); @@ -1739,7 +2073,7 @@ TEST_F(PromotionOnHitTest, RemoveAllErasesPromotionTask) { const int64_t cancelled_pre = mm.get_promotion_cancelled(); { - auto r = service->GetReplicaList("k_first", "default"); + auto r = service->GetReplicaList("k_first", TenantId::Default()); ASSERT_TRUE(r.has_value()); } { @@ -1758,7 +2092,7 @@ TEST_F(PromotionOnHitTest, RemoveAllErasesPromotionTask) { ASSERT_TRUE(InjectLocalDiskReplica(*service, holder.client_id, "k_second", 1024, holder.segment_name)); { - auto r = service->GetReplicaList("k_second", "default"); + auto r = service->GetReplicaList("k_second", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto pending_post = service->PromotionObjectHeartbeat(holder.client_id); @@ -1803,7 +2137,7 @@ TEST_F(PromotionOnHitTest, BatchRemoveErasesPromotionTask) { const int64_t cancelled_pre = mm.get_promotion_cancelled(); { - auto r = service->GetReplicaList("k_first", "default"); + auto r = service->GetReplicaList("k_first", TenantId::Default()); ASSERT_TRUE(r.has_value()); } { @@ -1812,7 +2146,8 @@ TEST_F(PromotionOnHitTest, BatchRemoveErasesPromotionTask) { EXPECT_EQ(CountPromotionTask(*pending, "k_first"), 1u); } - auto results = service->BatchRemove({"k_first"}, "default", /*force=*/true); + auto results = + service->BatchRemove({"k_first"}, TenantId::Default(), /*force=*/true); ASSERT_EQ(results.size(), 1u); EXPECT_TRUE(results[0].has_value()) << "BatchRemove should succeed; error=" << results[0].error(); @@ -1822,7 +2157,7 @@ TEST_F(PromotionOnHitTest, BatchRemoveErasesPromotionTask) { << "must bump promotion_cancelled_total."; { - auto r = service->GetReplicaList("k_second", "default"); + auto r = service->GetReplicaList("k_second", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto pending_post = service->PromotionObjectHeartbeat(holder.client_id); @@ -1859,7 +2194,7 @@ TEST_F(PromotionOnHitTest, BatchRemoveStaleHandleErasesPromotionTask) { const int64_t cancelled_pre = mm.get_promotion_cancelled(); { - auto r = service->GetReplicaList("k_first", "default"); + auto r = service->GetReplicaList("k_first", TenantId::Default()); ASSERT_TRUE(r.has_value()); } { @@ -1868,7 +2203,8 @@ TEST_F(PromotionOnHitTest, BatchRemoveStaleHandleErasesPromotionTask) { EXPECT_EQ(CountPromotionTask(*pending, "k_first"), 1u); } - auto results = service->BatchRemove({"k_first"}, "default", /*force=*/true); + auto results = + service->BatchRemove({"k_first"}, TenantId::Default(), /*force=*/true); ASSERT_EQ(results.size(), 1u); EXPECT_FALSE(results[0].has_value()) << "stale-handle path should report OBJECT_NOT_FOUND once the " @@ -1892,7 +2228,7 @@ TEST_F(PromotionOnHitTest, BatchRemoveStaleHandleErasesPromotionTask) { "k_second", 1024, second_holder.segment_name)); { - auto r = service->GetReplicaList("k_second", "default"); + auto r = service->GetReplicaList("k_second", TenantId::Default()); ASSERT_TRUE(r.has_value()); } auto pending_post = @@ -1933,18 +2269,18 @@ TEST_F(PromotionOnHitTest, MetricsFunnelTracksSuccessfulPromotion) { // Admit. { - auto r = service->GetReplicaList("k_hot", "default"); + auto r = service->GetReplicaList("k_hot", TenantId::Default()); ASSERT_TRUE(r.has_value()); } EXPECT_EQ(mm.get_promotion_admitted() - admitted_pre, 1); EXPECT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 1); // Drive AllocStart + NotifyPromotionSuccess. - auto alloc = service->PromotionAllocStart(seg.client_id, "k_hot", "default", - kObjBytes, {}); + auto alloc = service->PromotionAllocStart( + seg.client_id, "k_hot", TenantId::Default(), kObjBytes, {}); ASSERT_TRUE(alloc.has_value()); - auto notify = - service->NotifyPromotionSuccess(seg.client_id, "k_hot", "default"); + auto notify = service->NotifyPromotionSuccess(seg.client_id, "k_hot", + TenantId::Default()); ASSERT_TRUE(notify.has_value()); EXPECT_EQ(mm.get_promotion_completed() - completed_pre, 1); @@ -1979,25 +2315,25 @@ TEST_F(PromotionOnHitTest, MetricsRejectionCountersIncrementOnGateMiss) { // 1st Get on k_a: freq=1, threshold=2 → rejected on frequency. { - auto r = service->GetReplicaList("k_a", "default"); + auto r = service->GetReplicaList("k_a", TenantId::Default()); ASSERT_TRUE(r.has_value()); } EXPECT_EQ(mm.get_promotion_rejected_frequency() - freq_pre, 1); // 2nd Get on k_a: freq=2, admits. Now in-flight = 1 == limit. { - auto r = service->GetReplicaList("k_a", "default"); + auto r = service->GetReplicaList("k_a", TenantId::Default()); ASSERT_TRUE(r.has_value()); } // Get on k_b: freq=1, threshold=2 → rejected on frequency. { - auto r = service->GetReplicaList("k_b", "default"); + auto r = service->GetReplicaList("k_b", TenantId::Default()); ASSERT_TRUE(r.has_value()); } // Get on k_b again: freq=2, gets past frequency, but cap=1 // saturated → rejected on cap. { - auto r = service->GetReplicaList("k_b", "default"); + auto r = service->GetReplicaList("k_b", TenantId::Default()); ASSERT_TRUE(r.has_value()); } EXPECT_EQ(mm.get_promotion_rejected_cap() - cap_pre, 1); @@ -2024,7 +2360,7 @@ TEST_F(PromotionOnHitTest, MetricsRejectionCountersIncrementOnGateMiss) { const int64_t wm_pre = mm.get_promotion_rejected_watermark(); { - auto r = wm_service->GetReplicaList("k_w", "default"); + auto r = wm_service->GetReplicaList("k_w", TenantId::Default()); ASSERT_TRUE(r.has_value()); } EXPECT_EQ(mm.get_promotion_rejected_watermark() - wm_pre, 1); @@ -2035,21 +2371,26 @@ TEST_F(PromotionOnHitTest, MetricsRejectionCountersIncrementOnGateMiss) { TEST_F(PromotionOnHitTest, AdmissionFrequencyIsTenantScoped) { MasterServiceConfig config; config.enable_offload = true; + config.enable_multi_tenants = true; + config.tenant_quota_connector_type = "file"; config.promotion_on_hit = true; config.promotion_admission_threshold = 2; config.default_kv_lease_ttl = 2000; + const std::string key = "shared_hot_key"; + const TenantId tenant_a("tenant_promotion_a"); + const TenantId tenant_b("tenant_promotion_b"); + config.tenant_quota_connector_uri = + WriteTenantQuotaPolicyFile({{tenant_a.value(), 64 * 1024 * 1024}, + {tenant_b.value(), 64 * 1024 * 1024}}); auto service = std::make_unique(config); constexpr size_t seg_size = 1024 * 1024 * 16; auto seg = PrepareSegment(*service, "seg_tenant", kDefaultSegmentBase, seg_size); - const std::string key = "shared_hot_key"; - const std::string tenant_a = "tenant_promotion_a"; - const std::string tenant_b = "tenant_promotion_b"; ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, key, 1024, - seg.segment_name, tenant_a)); + seg.segment_name, tenant_a.value())); ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, key, 1024, - seg.segment_name, tenant_b)); + seg.segment_name, tenant_b.value())); { auto r = service->GetReplicaList(key, tenant_a); @@ -2073,7 +2414,7 @@ TEST_F(PromotionOnHitTest, AdmissionFrequencyIsTenantScoped) { service->PromotionObjectHeartbeat(seg.client_id); ASSERT_TRUE(pending_after_tenant_a_second.has_value()); ASSERT_EQ(pending_after_tenant_a_second->size(), 1u); - EXPECT_EQ((*pending_after_tenant_a_second)[0].tenant_id, tenant_a); + EXPECT_EQ((*pending_after_tenant_a_second)[0].tenant_id, tenant_a.value()); EXPECT_EQ((*pending_after_tenant_a_second)[0].key, key); service->RemoveAll(); @@ -2099,7 +2440,7 @@ TEST_F(PromotionOnHitTest, MaxPerHeartbeatKnobControlsBatchSize) { const auto key = "k_" + std::to_string(i); ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, key, 1024, seg.segment_name)); - auto r = service->GetReplicaList(key, "default"); + auto r = service->GetReplicaList(key, TenantId::Default()); ASSERT_TRUE(r.has_value()); } @@ -2140,7 +2481,7 @@ TEST_F(PromotionOnHitTest, MaxPerHeartbeatZeroClampsToOne) { ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k1", 1024, seg.segment_name)); { - auto r = service->GetReplicaList("k1", "default"); + auto r = service->GetReplicaList("k1", TenantId::Default()); ASSERT_TRUE(r.has_value()); } @@ -2178,13 +2519,13 @@ TEST_F(PromotionOnHitTest, MetricsRemoveMidPromotionCountsAsCancelled) { // Admit a promotion. in_flight goes from 0 to 1. { - auto r = service->GetReplicaList("k_drop", "default"); + auto r = service->GetReplicaList("k_drop", TenantId::Default()); ASSERT_TRUE(r.has_value()); } ASSERT_EQ(mm.get_promotion_admitted() - admitted_pre, 1); ASSERT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 1); - auto rm = service->Remove("k_drop", "default", /*force=*/true); + auto rm = service->Remove("k_drop", TenantId::Default(), /*force=*/true); ASSERT_TRUE(rm.has_value()) << "error=" << rm.error(); EXPECT_EQ(mm.get_promotion_in_flight() - in_flight_pre, 0); @@ -2193,6 +2534,316 @@ TEST_F(PromotionOnHitTest, MetricsRemoveMidPromotionCountsAsCancelled) { service->RemoveAll(); } +// --- Promotion retry candidate tests --- + +// A transient watermark rejection records a candidate; the retry scheduler +// can later see and process it. +TEST_F(PromotionOnHitTest, RetryCandidate_WatermarkRejectionRecordsCandidate) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + config.default_kv_lease_ttl = 2000; + config.eviction_high_watermark_ratio = 0.0; // all promotions rejected + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto seg = + PrepareSegment(*service, "retry_seg", kDefaultSegmentBase, seg_size); + // LOCAL_DISK-only key: inject without a prior PutStart. + ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_wm", 1024, + seg.segment_name)); + + auto& mm = MasterMetricManager::instance(); + const int64_t recorded_pre = mm.get_promotion_candidate_recorded(); + + // Get triggers frequency gate (threshold=1 → passes) then hits watermark=0. + { + auto r = service->GetReplicaList("k_wm", TenantId::Default()); + ASSERT_TRUE(r.has_value()); + } + + EXPECT_EQ(mm.get_promotion_candidate_recorded() - recorded_pre, 1) + << "Expected one candidate recorded on watermark rejection"; + EXPECT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 1u); + + // A second Get on the same key should update last_seen but not add a + // duplicate candidate. + { + auto r = service->GetReplicaList("k_wm", TenantId::Default()); + ASSERT_TRUE(r.has_value()); + } + EXPECT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 1u) + << "Duplicate candidate must not be created"; + + service->RemoveAll(); +} + +// A candidate recorded while the promotion queue is full is admitted once the +// active slot is released and the retry scanner reaches it. +TEST_F(PromotionOnHitTest, RetryCandidate_CapRejectedThenQueuedOnRetry) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + config.default_kv_lease_ttl = 2000; + config.promotion_queue_limit = 1; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto seg = PrepareSegment(*service, "retry_cap_seg", kDefaultSegmentBase, + seg_size); + ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_busy", 1024, + seg.segment_name)); + ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_retry", 1024, + seg.segment_name)); + + auto& mm = MasterMetricManager::instance(); + const int64_t candidate_admitted_pre = + mm.get_promotion_candidate_admitted(); + const int64_t promotion_admitted_pre = mm.get_promotion_admitted(); + + { + auto r = service->GetReplicaList("k_busy", TenantId::Default()); + ASSERT_TRUE(r.has_value()); + } + ASSERT_EQ(mm.get_promotion_admitted() - promotion_admitted_pre, 1); + ASSERT_EQ(GetPromotionInFlightForTesting(service.get()), 1u); + + { + auto r = service->GetReplicaList("k_retry", TenantId::Default()); + ASSERT_TRUE(r.has_value()); + } + ASSERT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 1u); + + auto failed = service->NotifyPromotionFailure(seg.client_id, "k_busy", + TenantId::Default()); + ASSERT_TRUE(failed.has_value()); + ASSERT_EQ(GetPromotionInFlightForTesting(service.get()), 0u); + + ResetCandidateBackoffsForTesting(service.get()); + EXPECT_EQ(RunPromotionCandidateRetryForTesting(service.get()), 1u); + + EXPECT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 0u); + EXPECT_EQ(GetPromotionInFlightForTesting(service.get()), 1u); + EXPECT_EQ(mm.get_promotion_candidate_admitted() - candidate_admitted_pre, + 1); + + auto heartbeat = service->PromotionObjectHeartbeat(seg.client_id); + ASSERT_TRUE(heartbeat.has_value()); + ASSERT_EQ(heartbeat->size(), 1u); + EXPECT_EQ(heartbeat->front().key, "k_retry"); + + service->RemoveAll(); +} + +TEST_F(PromotionOnHitTest, RetryCandidate_NoCandidatesOrNoShardBudgetNoops) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + config.default_kv_lease_ttl = 2000; + config.eviction_high_watermark_ratio = 0.0; + auto service = std::make_unique(config); + + EXPECT_EQ(RunPromotionCandidateRetryForTesting(service.get()), 0u); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto seg = PrepareSegment(*service, "retry_noop_seg", kDefaultSegmentBase, + seg_size); + ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_noop", 1024, + seg.segment_name)); + + { + auto r = service->GetReplicaList("k_noop", TenantId::Default()); + ASSERT_TRUE(r.has_value()); + } + ASSERT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 1u); + + EXPECT_EQ(RunPromotionCandidateRetryForTesting(service.get(), + /*shards_to_scan=*/0), + 0u); + EXPECT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 1u); + + service->RemoveAll(); +} + +// Candidate is cleaned up after kPromotionCandidateMaxRetries scans. +TEST_F(PromotionOnHitTest, RetryCandidate_ExhaustedAfterMaxRetries) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + config.default_kv_lease_ttl = 2000; + config.eviction_high_watermark_ratio = 0.0; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto seg = + PrepareSegment(*service, "exhaust_seg", kDefaultSegmentBase, seg_size); + ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_exhaust", + 1024, seg.segment_name)); + + // Record a candidate via a Get (watermark=0 rejects). + { + auto r = service->GetReplicaList("k_exhaust", TenantId::Default()); + ASSERT_TRUE(r.has_value()); + } + ASSERT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 1u); + + auto& mm = MasterMetricManager::instance(); + const int64_t expired_pre = mm.get_promotion_candidate_expired_evaluated(); + + // Drive retries; each scan increments retry_count until exhausted. + // Reset backoff timestamps before each scan so wall-clock time doesn't + // gate retries and the loop completes without sleeping. + for (int i = 0; i <= 10; ++i) { + ResetCandidateBackoffsForTesting(service.get()); + RunPromotionCandidateRetryForTesting(service.get()); + } + + EXPECT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 0u) + << "Candidate must be erased after max retries"; + EXPECT_GT(mm.get_promotion_candidate_expired_evaluated() - expired_pre, 0); + + service->RemoveAll(); +} + +// When an object is removed while a candidate is pending, the retry scan +// cleans up the candidate without promoting. +TEST_F(PromotionOnHitTest, RetryCandidate_ObjectRemovedMidRetry) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + config.default_kv_lease_ttl = 2000; + config.eviction_high_watermark_ratio = 0.0; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto seg = + PrepareSegment(*service, "rm_seg", kDefaultSegmentBase, seg_size); + ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_rm", 1024, + seg.segment_name)); + + auto& mm = MasterMetricManager::instance(); + const int64_t admitted_pre = mm.get_promotion_candidate_admitted(); + + // Record candidate. + { + auto r = service->GetReplicaList("k_rm", TenantId::Default()); + ASSERT_TRUE(r.has_value()); + } + ASSERT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 1u); + + // Remove the object while candidate is pending. + auto rm = service->Remove("k_rm", TenantId::Default(), /*force=*/true); + ASSERT_TRUE(rm.has_value()); + + // Retry scan should find object gone and erase candidate. + RunPromotionCandidateRetryForTesting(service.get()); + EXPECT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 0u); + + EXPECT_EQ(mm.get_promotion_candidate_admitted() - admitted_pre, 0); + + service->RemoveAll(); +} + +// Records multiple candidates and verifies per-key tracking is correct. +TEST_F(PromotionOnHitTest, RetryCandidate_MultipleKeysTracked) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + config.default_kv_lease_ttl = 2000; + config.eviction_high_watermark_ratio = 0.0; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 64; + auto seg = + PrepareSegment(*service, "multi_seg", kDefaultSegmentBase, seg_size); + + auto& mm = MasterMetricManager::instance(); + const int64_t recorded_pre = mm.get_promotion_candidate_recorded(); + const int64_t unevaluated_pre = + mm.get_promotion_candidate_expired_unevaluated(); + + constexpr int kKeys = 5; + for (int i = 0; i < kKeys; i++) { + std::string key = "k_multi_" + std::to_string(i); + ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, key, 512, + seg.segment_name)); + auto r = service->GetReplicaList(key, TenantId::Default()); + ASSERT_TRUE(r.has_value()); + } + + EXPECT_EQ(mm.get_promotion_candidate_recorded() - recorded_pre, kKeys); + EXPECT_EQ( + mm.get_promotion_candidate_expired_unevaluated() - unevaluated_pre, 0); + EXPECT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + static_cast(kKeys)); + + service->RemoveAll(); +} + +// ClearCandidatesForReload resets all candidate state and the global count. +TEST_F(PromotionOnHitTest, RetryCandidate_ClearOnReload) { + MasterServiceConfig config; + config.enable_offload = true; + config.promotion_on_hit = true; + config.promotion_admission_threshold = 1; + config.default_kv_lease_ttl = 2000; + config.eviction_high_watermark_ratio = 0.0; + auto service = std::make_unique(config); + + constexpr size_t seg_size = 1024 * 1024 * 16; + auto seg = + PrepareSegment(*service, "reload_seg", kDefaultSegmentBase, seg_size); + ASSERT_TRUE(InjectLocalDiskReplica(*service, seg.client_id, "k_reload", + 1024, seg.segment_name)); + + // Record a candidate. + { + auto r = service->GetReplicaList("k_reload", TenantId::Default()); + ASSERT_TRUE(r.has_value()); + } + ASSERT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 1u); + + // Simulate metadata reload. + ClearCandidatesForReloadForTesting(service.get()); + + EXPECT_EQ( + CountPromotionCandidatesForTesting(service.get(), TenantId::Default()), + 0u); + EXPECT_EQ(GetPromotionCandidateCountForTesting(service.get()), 0u); + EXPECT_EQ(GetPromotionInFlightForTesting(service.get()), 0u); + + service->RemoveAll(); +} + } // namespace mooncake::test int main(int argc, char** argv) { diff --git a/mooncake-store/tests/pybind_client_test.cpp b/mooncake-store/tests/pybind_client_test.cpp index 53492a5dd4..c08a48a8e7 100644 --- a/mooncake-store/tests/pybind_client_test.cpp +++ b/mooncake-store/tests/pybind_client_test.cpp @@ -992,6 +992,23 @@ TEST_F(RealClientTest, SetupWithConfigDictAllowsZeroSizes) { << "Setup should preserve zero-size pure client/server semantics"; } +TEST_F(RealClientTest, + ConfigDictGlobalSegmentSizeAboveMaxPassesSizeValidation) { + ConfigDict config = MakeConfigDict( + "localhost:17816", std::to_string(MAX_SEGMENT_SIZE + 1), "0"); + config[CONFIG_KEY_PROTOCOL] = "tcp"; + config[CONFIG_KEY_MASTER_SERVER_ADDR] = "127.0.0.1:1"; + + ::testing::internal::CaptureStderr(); + auto result = py_client_->setup_internal(config); + const std::string logs = ::testing::internal::GetCapturedStderr(); + + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(logs.find("Invalid global_segment_size"), std::string::npos) + << "global_segment_size above MAX_SEGMENT_SIZE should be accepted as " + "total capacity by ConfigDict validation"; +} + TEST_F(RealClientTest, ErrSetupWithInvalidConfigDictSize) { GLogMuter muter; ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build())) @@ -999,15 +1016,16 @@ TEST_F(RealClientTest, ErrSetupWithInvalidConfigDictSize) { master_address_ = master_.master_address(); struct InvalidSizeCase { - const char* local_hostname; - const char* global_segment_size; - const char* local_buffer_size; + std::string local_hostname; + std::string global_segment_size; + std::string local_buffer_size; }; const InvalidSizeCase invalid_size_cases[] = { {"localhost:17816", "50%", "16MB"}, {"localhost:17817", "16MB", "16XB"}, {"localhost:17818", "-5", "16MB"}, + {"localhost:17819", "0", std::to_string(MAX_SEGMENT_SIZE + 1)}, }; for (const auto& test_case : invalid_size_cases) { @@ -1794,6 +1812,65 @@ TEST_F(RealClientTest, ErrUnmountAndFreeInvalidSegmentIds) { << "Unmount-and-free of non-existent segment ids should fail"; } +#if defined(USE_SUNRISE) +TEST_F(RealClientTest, SetupWithConfigDictAcceptsSunriseLink) { + ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build())) + << "Failed to start in-proc master"; + master_address_ = master_.master_address(); + + ConfigDict config; + config[CONFIG_KEY_LOCAL_HOSTNAME] = "localhost:17813"; + config[CONFIG_KEY_METADATA_SERVER] = "P2PHANDSHAKE"; + config[CONFIG_KEY_GLOBAL_SEGMENT_SIZE] = std::to_string(16 * 1024 * 1024); + config[CONFIG_KEY_LOCAL_BUFFER_SIZE] = std::to_string(16 * 1024 * 1024); + config[CONFIG_KEY_PROTOCOL] = "sunrise_link"; + config[CONFIG_KEY_MASTER_SERVER_ADDR] = master_address_; + + auto result = py_client_->setup_internal(config); + ASSERT_TRUE(result.has_value()) + << "setup_internal(ConfigDict) should accept sunrise_link protocol"; +} + +TEST_F(RealClientTest, SetupAcceptsSunriseLink) { + ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build())); + master_address_ = master_.master_address(); + + int result = py_client_->setup_real("localhost:17813", "P2PHANDSHAKE", + 16 * 1024 * 1024, 16 * 1024 * 1024, + "sunrise_link", "", master_address_); + EXPECT_EQ(result, 0) << "sunrise_link should work with Classic TE"; +} + +TEST_F(RealClientTest, SunriseLinkHostPutGetRemove) { + ASSERT_TRUE(master_.Start(InProcMasterConfigBuilder().build())); + master_address_ = master_.master_address(); + + ASSERT_EQ(py_client_->setup_real("localhost:17813", "P2PHANDSHAKE", + 16 * 1024 * 1024, 16 * 1024 * 1024, + "sunrise_link", "", master_address_), + 0); + + const std::string key = "sunrise_link_host_key"; + const std::string value = "sunrise_link_host_value"; + std::span data_span(value.data(), value.size()); + ReplicateConfig config; + config.replica_num = 1; + + ASSERT_EQ(py_client_->put(key, data_span, config), 0); + + { + auto buffer_handle = py_client_->get_buffer(key); + ASSERT_NE(buffer_handle, nullptr); + std::string actual(static_cast(buffer_handle->ptr()), + buffer_handle->size()); + EXPECT_EQ(actual, value); + } + + EXPECT_EQ(py_client_->remove(key, true), 0); + EXPECT_EQ(py_client_->isExist(key), 0); +} +#endif + } // namespace testing } // namespace mooncake diff --git a/mooncake-store/tests/registered_pinned_memory_test.cpp b/mooncake-store/tests/registered_pinned_memory_test.cpp new file mode 100644 index 0000000000..04219025e3 --- /dev/null +++ b/mooncake-store/tests/registered_pinned_memory_test.cpp @@ -0,0 +1,133 @@ +#define MOONCAKE_STORE_TEST +#include "../src/registered_pinned_memory.h" + +#include + +#include + +namespace mooncake { +namespace { + +using Manager = RegisteredPinnedMemoryManager; +using UnregisterResult = Manager::UnregisterResult; + +struct FakePinState { + bool register_succeeds = true; + UnregisterResult unregister_result = UnregisterResult::kSuccess; + int register_calls = 0; + int unregister_calls = 0; +}; + +FakePinState& State() { + static FakePinState state; + return state; +} + +bool FakeRegister(void*, size_t, std::string* error_message) { + ++State().register_calls; + if (State().register_succeeds) return true; + if (error_message) *error_message = "fake register failure"; + return false; +} + +UnregisterResult FakeUnregister(void*, std::string* error_message) { + ++State().unregister_calls; + if (State().unregister_result == UnregisterResult::kError && + error_message) { + *error_message = "fake unregister failure"; + } + return State().unregister_result; +} + +class RegisteredPinnedMemoryManagerTest : public ::testing::Test { + protected: + void SetUp() override { State() = FakePinState(); } + + Manager MakeManager(size_t limit) { + return Manager({true, limit}, {FakeRegister, FakeUnregister}); + } + + std::shared_ptr Pin(Manager& manager, size_t offset, + size_t size) { + return manager.try_pin(buffer_.data() + offset, size, "segment"); + } + + void ExpectCalls(int register_calls, int unregister_calls) { + EXPECT_EQ(State().register_calls, register_calls); + EXPECT_EQ(State().unregister_calls, unregister_calls); + } + + std::array buffer_{}; +}; + +TEST_F(RegisteredPinnedMemoryManagerTest, QuotaRejectsAndReleaseRefunds) { + auto manager = MakeManager(64); + + auto first = Pin(manager, 0, 64); + ASSERT_NE(first, nullptr); + EXPECT_EQ(Pin(manager, 64, 1), nullptr); + ExpectCalls(1, 0); + + first.reset(); + ExpectCalls(1, 1); + + auto second = Pin(manager, 64, 64); + ASSERT_NE(second, nullptr); + ExpectCalls(2, 1); + + second.reset(); + ExpectCalls(2, 2); +} + +TEST_F(RegisteredPinnedMemoryManagerTest, OverlapAndDuplicateAreRejected) { + auto manager = MakeManager(128); + + auto first = Pin(manager, 16, 32); + ASSERT_NE(first, nullptr); + + EXPECT_EQ(Pin(manager, 16, 32), nullptr); + EXPECT_EQ(Pin(manager, 32, 16), nullptr); + + auto adjacent = Pin(manager, 48, 16); + ASSERT_NE(adjacent, nullptr); + ExpectCalls(2, 0); + + adjacent.reset(); + first.reset(); + ExpectCalls(2, 2); +} + +TEST_F(RegisteredPinnedMemoryManagerTest, RegisterFailureRefundsReservation) { + auto manager = MakeManager(32); + + State().register_succeeds = false; + EXPECT_EQ(Pin(manager, 0, 32), nullptr); + ExpectCalls(1, 0); + + State().register_succeeds = true; + auto retried = Pin(manager, 0, 32); + ASSERT_NE(retried, nullptr); + ExpectCalls(2, 0); + + retried.reset(); + ExpectCalls(2, 1); +} + +TEST_F(RegisteredPinnedMemoryManagerTest, UnregisterFailureRetainsReservation) { + auto manager = MakeManager(32); + + auto first = Pin(manager, 0, 32); + ASSERT_NE(first, nullptr); + + State().unregister_result = UnregisterResult::kError; + EXPECT_FALSE(first->release()); + ExpectCalls(1, 1); + + State().unregister_result = UnregisterResult::kSuccess; + EXPECT_EQ(Pin(manager, 0, 32), nullptr); + EXPECT_EQ(Pin(manager, 32, 32), nullptr); + ExpectCalls(1, 1); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/replica_selection_test.cpp b/mooncake-store/tests/replica_selection_test.cpp new file mode 100644 index 0000000000..388cede25a --- /dev/null +++ b/mooncake-store/tests/replica_selection_test.cpp @@ -0,0 +1,361 @@ +// Copyright 2025 Mooncake Authors +// +// Unit tests for replica_selection.h: verifies the base type+locality policy is +// unchanged, and that the opt-in remote-replica scoring picks a better remote +// MEMORY replica instead of the first one the master happened to return +// (issue #2516). + +#include "replica_selection.h" + +#include + +#include +#include +#include +#include +#include + +namespace mooncake { +namespace { + +// Build a COMPLETE MEMORY replica descriptor with the given endpoint/protocol. +Replica::Descriptor MakeMemory(const std::string& endpoint, + const std::string& protocol, + ReplicaStatus status = ReplicaStatus::COMPLETE) { + Replica::Descriptor d; + d.id = 0; + MemoryDescriptor mem; + mem.buffer_descriptor.size_ = 1024; + mem.buffer_descriptor.buffer_address_ = 0x1000; + mem.buffer_descriptor.protocol_ = protocol; + mem.buffer_descriptor.transport_endpoint_ = endpoint; + d.descriptor_variant = mem; + d.status = status; + return d; +} + +Replica::Descriptor MakeNoF(const std::string& endpoint, + ReplicaStatus status = ReplicaStatus::COMPLETE) { + Replica::Descriptor d; + d.id = 0; + NoFDescriptor nof; + nof.buffer_descriptor.size_ = 1024; + nof.buffer_descriptor.buffer_address_ = 0x2000; + nof.buffer_descriptor.protocol_ = "nvmeof"; + nof.buffer_descriptor.transport_endpoint_ = endpoint; + d.descriptor_variant = nof; + d.status = status; + return d; +} + +Replica::Descriptor MakeDisk(const std::string& path) { + Replica::Descriptor d; + d.id = 0; + d.descriptor_variant = DiskDescriptor{path, 1024}; + d.status = ReplicaStatus::COMPLETE; + return d; +} + +Replica::Descriptor MakeLocalDisk(const std::string& endpoint) { + Replica::Descriptor d; + d.id = 0; + LocalDiskDescriptor local_disk; + local_disk.object_size = 1024; + local_disk.transport_endpoint = endpoint; + d.descriptor_variant = local_disk; + d.status = ReplicaStatus::COMPLETE; + return d; +} + +// A test fixture that guarantees scoring state is reset between tests, since +// the enable flag / injected scorer are process-wide. +class ReplicaSelectionTest : public ::testing::Test { + protected: + void TearDown() override { SetRemoteReplicaScorer(nullptr); } +}; + +// --- Base policy (scoring off): behaviour must be unchanged -------------- + +TEST_F(ReplicaSelectionTest, LocalMemoryAlwaysWins) { + std::unordered_set local = {"nodeB"}; + std::vector reps = { + MakeMemory("nodeA", "rdma"), + MakeMemory("nodeB", "tcp"), // local, slower protocol + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ( + sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_, + "nodeB"); // locality beats protocol +} + +TEST_F(ReplicaSelectionTest, ScoringOffKeepsFirstRemoteMemory) { + // No scorer injected and env not set -> must return the FIRST remote + // MEMORY. + ASSERT_FALSE(RemoteReplicaScoringEnabled()); + std::unordered_set local; // nothing local + std::vector reps = { + MakeMemory("nodeA", "tcp"), // first + MakeMemory("nodeB", "rdma"), // "better" but must be ignored when off + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ( + sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_, + "nodeA"); +} + +// --- Opt-in scoring: pick the better remote replica --------------------- + +TEST_F(ReplicaSelectionTest, InjectedScorerPicksLowestScore) { + // Inject a scorer that prefers "nodeB" regardless of order. + SetRemoteReplicaScorer([](const Replica::Descriptor& r) { + const auto& ep = + r.get_memory_descriptor().buffer_descriptor.transport_endpoint_; + return ep == "nodeB" ? 0.0 : 10.0; + }); + ASSERT_TRUE(RemoteReplicaScoringEnabled()); + + std::unordered_set local; + std::vector reps = { + MakeMemory("nodeA", "rdma"), // first, but higher score + MakeMemory("nodeB", "rdma"), // lower score -> should win + MakeMemory("nodeC", "rdma"), + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ( + sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_, + "nodeB"); +} + +TEST_F(ReplicaSelectionTest, BuiltinScorerPrefersRdmaOverTcp) { + // Built-in scorer active via injected scorer? No — use it directly by + // enabling through injection of the built-in. Simulate env-on path by + // injecting the built-in function. + SetRemoteReplicaScorer(BuiltinRemoteReplicaScore); + ASSERT_TRUE(RemoteReplicaScoringEnabled()); + + std::unordered_set local; + std::vector reps = { + MakeMemory("nodeA", "tcp"), // first, but tcp + MakeMemory("nodeB", "rdma"), // rdma -> preferred + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ( + sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_, + "nodeB"); +} + +TEST_F(ReplicaSelectionTest, BuiltinScorerRanksUnknownAndNonMemory) { + EXPECT_DOUBLE_EQ(BuiltinRemoteReplicaScore(MakeMemory("nodeA", "ucx")), + 2.0); + EXPECT_DOUBLE_EQ(BuiltinRemoteReplicaScore(MakeNoF("nodeB")), 100.0); +} + +TEST_F(ReplicaSelectionTest, EnvironmentOptInUsesBuiltinScorer) { + const char* env = std::getenv("MC_STORE_REPLICA_SCORING"); + if (env == nullptr || std::string(env) != "1") { + GTEST_SKIP() << "covered by the env-enabled CTest process"; + } + + std::unordered_set local; + std::vector reps = { + MakeMemory("nodeA", "tcp"), + MakeMemory("nodeB", "rdma"), + }; + + EXPECT_TRUE(RemoteReplicaScoringEnabled()); + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ(sel->get_memory_descriptor().buffer_descriptor.protocol_, "rdma"); +} + +TEST_F(ReplicaSelectionTest, ScorerTieKeepsMasterOrder) { + // All equal score -> strictly-less comparison keeps the first one. + SetRemoteReplicaScorer([](const Replica::Descriptor&) { return 5.0; }); + std::unordered_set local; + std::vector reps = { + MakeMemory("nodeA", "rdma"), + MakeMemory("nodeB", "rdma"), + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ( + sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_, + "nodeA"); +} + +TEST_F(ReplicaSelectionTest, ScorerSkipsIncompleteReplicas) { + SetRemoteReplicaScorer([](const Replica::Descriptor& r) { + const auto& ep = + r.get_memory_descriptor().buffer_descriptor.transport_endpoint_; + return ep == "nodeB" ? 0.0 : 10.0; + }); + std::unordered_set local; + std::vector reps = { + MakeMemory("nodeA", "rdma"), + MakeMemory("nodeB", "rdma", + ReplicaStatus::PROCESSING), // best score + // but not ready + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ( + sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_, + "nodeA"); // nodeB skipped -> falls back to only COMPLETE one +} + +TEST_F(ReplicaSelectionTest, LocalStillWinsWhenScoringOn) { + SetRemoteReplicaScorer(BuiltinRemoteReplicaScore); + std::unordered_set local = {"nodeB"}; + std::vector reps = { + MakeMemory("nodeA", "rdma"), // remote, best protocol + MakeMemory("nodeB", "tcp"), // local -> must still win over scoring + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ( + sel->get_memory_descriptor().buffer_descriptor.transport_endpoint_, + "nodeB"); +} + +TEST_F(ReplicaSelectionTest, PickBestRemoteMemoryCanFindNoCandidate) { + SetRemoteReplicaScorer(BuiltinRemoteReplicaScore); + std::unordered_set local = {"nodeA"}; + std::vector reps = { + MakeMemory("nodeA", "rdma"), + MakeMemory("nodeB", "rdma", ReplicaStatus::PROCESSING), + MakeNoF("nodeC"), + }; + EXPECT_EQ(PickBestRemoteMemory(reps, local), nullptr); +} + +// --- Non-MEMORY fallbacks: preserve the complete historical policy ------- + +TEST_F(ReplicaSelectionTest, LocalNoFPrecedesRemoteNoF) { + std::unordered_set local = {"nodeB"}; + std::vector reps = { + MakeNoF("nodeA"), + MakeNoF("nodeB"), + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ(sel->get_nof_descriptor().buffer_descriptor.transport_endpoint_, + "nodeB"); +} + +TEST_F(ReplicaSelectionTest, RemoteNoFFallbackKeepsMasterOrder) { + std::unordered_set local; + std::vector reps = { + MakeNoF("nodeA"), + MakeNoF("nodeB"), + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_EQ(sel->get_nof_descriptor().buffer_descriptor.transport_endpoint_, + "nodeA"); +} + +TEST_F(ReplicaSelectionTest, LocalDiskPrecedesDisk) { + std::unordered_set local; + std::vector reps = { + MakeDisk("/remote/object"), + MakeLocalDisk("nodeA"), + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_TRUE(sel->is_local_disk_replica()); +} + +TEST_F(ReplicaSelectionTest, DiskIsLastCompleteFallback) { + std::unordered_set local; + std::vector reps = { + MakeDisk("/remote/object"), + MakeMemory("nodeA", "rdma", ReplicaStatus::FAILED), + }; + const auto* sel = SelectBestReplica(reps, local); + ASSERT_NE(sel, nullptr); + EXPECT_TRUE(sel->is_disk_replica()); +} + +TEST_F(ReplicaSelectionTest, NoCompleteReplicaReturnsNull) { + std::unordered_set local; + std::vector reps = { + MakeMemory("nodeA", "rdma", ReplicaStatus::PROCESSING), + MakeNoF("nodeB", ReplicaStatus::FAILED), + }; + EXPECT_EQ(SelectBestReplica(reps, local), nullptr); +} + +// --- Concurrency: verify no data race on SetRemoteReplicaScorer vs reads --- + +TEST_F(ReplicaSelectionTest, ConcurrentSetAndSelectIsRaceFree) { + constexpr int kIterations = 50000; + constexpr int kReaderThreads = 4; + + std::unordered_set local; + std::vector reps = { + MakeMemory("nodeA", "tcp"), + MakeMemory("nodeB", "rdma"), + MakeMemory("nodeC", "rdma"), + }; + + std::atomic ready{0}; + std::atomic start{false}; + std::atomic selection_failed{false}; + std::atomic read_count{0}; + + auto wait_for_start = [&] { + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + }; + + // Writer: repeatedly swap scorers while readers are active. + std::thread writer([&] { + wait_for_start(); + for (int i = 0; i < kIterations; ++i) { + if (i % 2 == 0) { + SetRemoteReplicaScorer([](const Replica::Descriptor& r) { + const auto& proto = + r.get_memory_descriptor().buffer_descriptor.protocol_; + return proto == "rdma" ? 0.0 : 10.0; + }); + } else { + SetRemoteReplicaScorer(nullptr); + } + } + }); + + // Readers: call SelectBestReplica (which reads the scorer) concurrently. + std::vector readers; + for (int t = 0; t < kReaderThreads; ++t) { + readers.emplace_back([&] { + wait_for_start(); + for (int i = 0; i < kIterations; ++i) { + const auto* sel = SelectBestReplica(reps, local); + if (sel == nullptr) { + selection_failed.store(true, std::memory_order_relaxed); + } + read_count.fetch_add(1, std::memory_order_relaxed); + } + }); + } + + while (ready.load(std::memory_order_acquire) < kReaderThreads + 1) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + + writer.join(); + for (auto& r : readers) r.join(); + + EXPECT_FALSE(selection_failed.load()); + EXPECT_EQ(read_count.load(), kIterations * kReaderThreads); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/runtime_accelerator_test.cpp b/mooncake-store/tests/runtime_accelerator_test.cpp new file mode 100644 index 0000000000..395084f4b5 --- /dev/null +++ b/mooncake-store/tests/runtime_accelerator_test.cpp @@ -0,0 +1,114 @@ +#include "device/runtime_accelerator.h" + +#include + +#include + +namespace mooncake::device { +namespace { + +class FakeAcceleratorDevice final : public AcceleratorDevice { + public: + FakeAcceleratorDevice(AcceleratorVendor vendor, const void* device_ptr, + int32_t device_id) + : vendor_(vendor), device_ptr_(device_ptr), device_id_(device_id) {} + + AcceleratorVendor Vendor() const override { return vendor_; } + + bool Available(bool ensure = false) const override { + (void)ensure; + return true; + } + + PointerInfo QueryPointer(const void* ptr) const override { + ++query_count_; + if (ptr != device_ptr_) { + return PointerInfo{.kind = MemoryKind::kHost, .device_id = -1}; + } + return PointerInfo{.kind = MemoryKind::kDevice, + .device_id = device_id_}; + } + + int32_t CurrentDeviceId() const override { return current_device_id_; } + + void SetContext(int32_t device_id) const override { + current_device_id_ = device_id; + } + + bool Copy(void* dst, const void* src, size_t size, + CopyDirection direction) const override { + last_direction_ = direction; + std::memcpy(dst, src, size); + return copy_succeeds_; + } + + PinnedHostBuffer AllocatePinnedHost(size_t size) const override { + (void)size; + return PinnedHostBuffer(); + } + + void set_copy_succeeds(bool succeeds) { copy_succeeds_ = succeeds; } + int32_t current_device_id() const { return current_device_id_; } + CopyDirection last_direction() const { return last_direction_; } + int query_count() const { return query_count_; } + + private: + AcceleratorVendor vendor_; + const void* device_ptr_; + int32_t device_id_; + mutable int32_t current_device_id_ = -1; + mutable CopyDirection last_direction_ = CopyDirection::kAuto; + mutable bool copy_succeeds_ = true; + mutable int query_count_ = 0; +}; + +TEST(RuntimeAcceleratorTest, FindDeviceForPointerReturnsMatchingDevice) { + char device_byte = 'd'; + FakeAcceleratorDevice device(AcceleratorVendor::kNvidia, &device_byte, 7); + RuntimeAccelerator runtime_accelerator({&device}); + + PointerInfo info; + auto* found = runtime_accelerator.FindDeviceForPointer(&device_byte, &info); + + EXPECT_EQ(found, &device); + EXPECT_EQ(info.kind, MemoryKind::kDevice); + EXPECT_EQ(info.device_id, 7); +} + +TEST(RuntimeAcceleratorTest, FindDeviceForPointerSkipsNullPointerQueries) { + char device_byte = 'd'; + FakeAcceleratorDevice device(AcceleratorVendor::kNvidia, &device_byte, 0); + RuntimeAccelerator runtime_accelerator({&device}); + + EXPECT_EQ(runtime_accelerator.FindDeviceForPointer(nullptr), nullptr); + EXPECT_EQ(device.query_count(), 0); +} + +TEST(RuntimeAcceleratorTest, CopyToHostUsesDeviceToHostCopy) { + char src[] = "abc"; + char dst[sizeof(src)] = {}; + FakeAcceleratorDevice device(AcceleratorVendor::kNvidia, src, 3); + RuntimeAccelerator runtime_accelerator({&device}); + + EXPECT_TRUE(runtime_accelerator.CopyToHost(dst, src, sizeof(src))); + + EXPECT_STREQ(dst, src); + EXPECT_EQ(device.current_device_id(), 3); + EXPECT_EQ(device.last_direction(), CopyDirection::kDeviceToHost); +} + +TEST(RuntimeAcceleratorTest, CopyFromHostUsesHostToDeviceCopy) { + char src[] = "abc"; + char dst[sizeof(src)] = {}; + FakeAcceleratorDevice device(AcceleratorVendor::kNvidia, dst, 4); + RuntimeAccelerator runtime_accelerator({&device}); + + EXPECT_TRUE(runtime_accelerator.CopyFromHost(dst, src, sizeof(src))); + + EXPECT_STREQ(dst, src); + EXPECT_EQ(device.current_device_id(), 4); + EXPECT_EQ(device.last_direction(), CopyDirection::kHostToDevice); +} + +} // namespace +} // namespace mooncake::device diff --git a/mooncake-store/tests/segment_test.cpp b/mooncake-store/tests/segment_test.cpp index 4fa0c62508..60b7362061 100644 --- a/mooncake-store/tests/segment_test.cpp +++ b/mooncake-store/tests/segment_test.cpp @@ -5,6 +5,11 @@ #include +#include +#include +#include +#include + namespace mooncake { // Test fixture for Segment tests @@ -346,6 +351,172 @@ TEST_F(SegmentTest, SegmentLifecycleStatusControlsAllocation) { EXPECT_TRUE(HasAllocatorForSegment(segment_manager, segment.id)); } +TEST_F(SegmentTest, HostOrderedSegmentsTracksMountStatusAndUnmount) { + SegmentManager segment_manager; + + Segment segment0; + segment0.id = generate_uuid(); + segment0.name = "host0_segment"; + segment0.size = 1024 * 1024 * 16; + segment0.base = 0x100000000; + segment0.host_id = "host0"; + + Segment segment1; + segment1.id = generate_uuid(); + segment1.name = "host1_segment"; + segment1.size = 1024 * 1024 * 16; + segment1.base = 0x200000000; + segment1.host_id = "host1"; + + UUID client_id = generate_uuid(); + + { + auto segment_access = segment_manager.getSegmentAccess(); + ASSERT_EQ(segment_access.MountSegment(segment0, client_id), + ErrorCode::OK); + ASSERT_EQ(segment_access.MountSegment(segment1, client_id), + ErrorCode::OK); + } + + { + auto allocator_access = segment_manager.getAllocatorAccess(); + auto ordered = + allocator_access.GetHostOrderedSegments("host1", "test_key"); + ASSERT_GE(ordered.size(), 2u); + EXPECT_EQ(ordered[0], segment1.name); + } + + { + auto segment_access = segment_manager.getSegmentAccess(); + ASSERT_EQ(segment_access.SetSegmentStatusByName( + segment1.name, SegmentStatus::DRAINING), + ErrorCode::OK); + } + + { + auto allocator_access = segment_manager.getAllocatorAccess(); + auto ordered = + allocator_access.GetHostOrderedSegments("host1", "test_key"); + ASSERT_EQ(ordered.size(), 1u); + EXPECT_EQ(ordered[0], segment0.name); + } + + { + auto segment_access = segment_manager.getSegmentAccess(); + ASSERT_EQ(segment_access.SetSegmentStatusByName(segment1.name, + SegmentStatus::OK), + ErrorCode::OK); + size_t metrics_dec_capacity = 0; + ASSERT_EQ(segment_access.PrepareUnmountSegment(segment1.id, + metrics_dec_capacity), + ErrorCode::OK); + ASSERT_EQ(segment_access.CommitUnmountSegment(segment1.id, client_id, + metrics_dec_capacity), + ErrorCode::OK); + } + + { + auto allocator_access = segment_manager.getAllocatorAccess(); + auto ordered = + allocator_access.GetHostOrderedSegments("host1", "test_key"); + ASSERT_EQ(ordered.size(), 1u); + EXPECT_EQ(ordered[0], segment0.name); + } +} + +TEST_F(SegmentTest, HostOrderedSegmentsKeepsNameUntilLastSameNameSegmentGone) { + SegmentManager segment_manager; + + Segment segment0; + segment0.id = generate_uuid(); + segment0.name = "shared_host_segment"; + segment0.size = 1024 * 1024 * 16; + segment0.base = 0x100000000; + segment0.host_id = "host1"; + + Segment segment1; + segment1.id = generate_uuid(); + segment1.name = segment0.name; + segment1.size = 1024 * 1024 * 16; + segment1.base = 0x200000000; + segment1.host_id = "host1"; + + UUID client_id = generate_uuid(); + { + auto segment_access = segment_manager.getSegmentAccess(); + ASSERT_EQ(segment_access.MountSegment(segment0, client_id), + ErrorCode::OK); + ASSERT_EQ(segment_access.MountSegment(segment1, client_id), + ErrorCode::OK); + } + + { + auto allocator_access = segment_manager.getAllocatorAccess(); + auto ordered = + allocator_access.GetHostOrderedSegments("host1", "test_key"); + ASSERT_EQ(ordered.size(), 1u); + EXPECT_EQ(ordered[0], segment0.name); + } + + { + auto segment_access = segment_manager.getSegmentAccess(); + size_t metrics_dec_capacity = 0; + ASSERT_EQ(segment_access.PrepareUnmountSegment(segment0.id, + metrics_dec_capacity), + ErrorCode::OK); + ASSERT_EQ(segment_access.CommitUnmountSegment(segment0.id, client_id, + metrics_dec_capacity), + ErrorCode::OK); + } + + { + auto allocator_access = segment_manager.getAllocatorAccess(); + auto ordered = + allocator_access.GetHostOrderedSegments("host1", "test_key"); + ASSERT_EQ(ordered.size(), 1u); + EXPECT_EQ(ordered[0], segment1.name); + } +} + +TEST_F(SegmentTest, HostOrderedSegmentsRotateWithinSameHostByKey) { + SegmentManager segment_manager; + + Segment segment_a; + segment_a.id = generate_uuid(); + segment_a.name = "host1_segment_a"; + segment_a.size = 1024 * 1024 * 16; + segment_a.base = 0x100000000; + segment_a.host_id = "host1"; + + Segment segment_b; + segment_b.id = generate_uuid(); + segment_b.name = "host1_segment_b"; + segment_b.size = 1024 * 1024 * 16; + segment_b.base = 0x200000000; + segment_b.host_id = "host1"; + + UUID client_id = generate_uuid(); + { + auto segment_access = segment_manager.getSegmentAccess(); + ASSERT_EQ(segment_access.MountSegment(segment_a, client_id), + ErrorCode::OK); + ASSERT_EQ(segment_access.MountSegment(segment_b, client_id), + ErrorCode::OK); + } + + const std::string key = "stable_rotation_key"; + std::vector sorted_segments = {segment_a.name, segment_b.name}; + std::sort(sorted_segments.begin(), sorted_segments.end()); + const size_t start = std::hash{}(key) % sorted_segments.size(); + + auto allocator_access = segment_manager.getAllocatorAccess(); + auto ordered = allocator_access.GetHostOrderedSegments("host1", key); + ASSERT_EQ(ordered.size(), 2u); + EXPECT_EQ(ordered[0], sorted_segments[start]); + EXPECT_EQ(ordered[1], + sorted_segments[(start + 1) % sorted_segments.size()]); +} + TEST_F(SegmentTest, PrepareUnmountDrainedSegment) { SegmentManager segment_manager; diff --git a/mooncake-store/tests/serializer_test.cpp b/mooncake-store/tests/serializer_test.cpp index 34dfb2ba49..538d2ee05c 100644 --- a/mooncake-store/tests/serializer_test.cpp +++ b/mooncake-store/tests/serializer_test.cpp @@ -2,6 +2,8 @@ #include #include "serializer.h" +#include "serialize/serializer.h" +#include "segment.h" namespace mooncake::test { @@ -165,9 +167,59 @@ TEST_F(SerializerTest, ExampleClassDeserializationWithException) { EXPECT_EQ(restored, nullptr); } +TEST_F(SerializerTest, MountedSegmentSerializationPreservesHostId) { + MountedSegment original; + original.segment.id = generate_uuid(); + original.segment.name = "segment_host1"; + original.segment.base = 0x300000000; + original.segment.size = 1024 * 1024; + original.segment.te_endpoint = "segment_host1"; + original.segment.host_id = "host1"; + original.status = SegmentStatus::OK; + + msgpack::sbuffer buffer; + MsgpackPacker packer(&buffer); + ASSERT_TRUE( + Serializer::serialize(original, packer).has_value()); + + auto object_handle = msgpack::unpack(buffer.data(), buffer.size()); + auto restored = + Serializer::deserialize(object_handle.get()); + ASSERT_TRUE(restored.has_value()); + EXPECT_EQ(restored->segment.id, original.segment.id); + EXPECT_EQ(restored->segment.name, original.segment.name); + EXPECT_EQ(restored->segment.host_id, original.segment.host_id); + EXPECT_EQ(restored->status, original.status); +} + +TEST_F(SerializerTest, MountedSegmentDeserializesLegacyFormatWithoutHostId) { + const UUID segment_id = generate_uuid(); + + msgpack::sbuffer buffer; + MsgpackPacker packer(&buffer); + packer.pack_array(8); + packer.pack(UuidToString(segment_id)); + packer.pack(std::string("legacy_segment")); + packer.pack(static_cast(0x300000000)); + packer.pack(static_cast(1024 * 1024)); + packer.pack(std::string("legacy_segment")); + packer.pack(static_cast(SegmentStatus::OK)); + packer.pack(false); + packer.pack_nil(); + + auto object_handle = msgpack::unpack(buffer.data(), buffer.size()); + auto restored = + Serializer::deserialize(object_handle.get()); + ASSERT_TRUE(restored.has_value()); + EXPECT_EQ(restored->segment.id, segment_id); + EXPECT_EQ(restored->segment.name, "legacy_segment"); + EXPECT_TRUE(restored->segment.host_id.empty()); + EXPECT_EQ(restored->status, SegmentStatus::OK); +} + } // namespace mooncake::test int main(int argc, char** argv) { ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); -} \ No newline at end of file +} diff --git a/mooncake-store/tests/storage_backend_test.cpp b/mooncake-store/tests/storage_backend_test.cpp index 1c66f376e4..1f3e1ae246 100644 --- a/mooncake-store/tests/storage_backend_test.cpp +++ b/mooncake-store/tests/storage_backend_test.cpp @@ -4,10 +4,14 @@ #include #include +#include #include +#include #include #include #include +#include +#include #include #include #include @@ -152,6 +156,85 @@ class StorageBackendTest : public ::testing::Test { } }; +// Regression tests for StorageBackend::Create validation (issue #3134): +// invalid configuration must be reported as INVALID_PARAMS instead of +// producing a nullptr that callers may dereference. +TEST_F(StorageBackendTest, CreateRejectsNonexistentRootDir) { + auto result = StorageBackend::Create(data_path + "/does/not/exist/12345", + "fsdir", true); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); +} + +TEST_F(StorageBackendTest, CreateRejectsRootDirThatIsAFile) { + std::string file_path = data_path + "/root_is_a_file"; + { + std::ofstream ofs(file_path); + ofs << "not a directory"; + } + auto result = StorageBackend::Create(file_path, "fsdir", true); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); +} + +TEST_F(StorageBackendTest, CreateRejectsEmptyFsdir) { + auto result = StorageBackend::Create(data_path, "", true); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), ErrorCode::INVALID_PARAMS); +} + +TEST_F(StorageBackendTest, CreateAcceptsValidConfig) { + auto result = StorageBackend::Create(data_path, "fsdir", true); + ASSERT_TRUE(result.has_value()); + EXPECT_NE(result.value(), nullptr); +} + +// Regression tests for StorageBackendAdaptor::Init validation (issue #3134 +// follow-up): obviously invalid configuration must be reported as +// INVALID_PARAMS, while a missing root directory is still auto-created. +TEST_F(StorageBackendTest, AdaptorInitRejectsStoragePathThatIsAFile) { + std::string file_path = data_path + "/storage_path_is_a_file"; + { + std::ofstream ofs(file_path); + ofs << "not a directory"; + } + FileStorageConfig cfg; + cfg.storage_filepath = file_path; + FilePerKeyConfig file_per_key_config; + file_per_key_config.fsdir = "file_per_key_dir"; + file_per_key_config.enable_eviction = false; + + StorageBackendAdaptor adaptor(cfg, file_per_key_config); + auto init_result = adaptor.Init(); + ASSERT_FALSE(init_result.has_value()); + EXPECT_EQ(init_result.error(), ErrorCode::INVALID_PARAMS); +} + +TEST_F(StorageBackendTest, AdaptorInitRejectsEmptyFsdir) { + FileStorageConfig cfg; + cfg.storage_filepath = data_path; + FilePerKeyConfig file_per_key_config; + file_per_key_config.fsdir = ""; + file_per_key_config.enable_eviction = false; + + StorageBackendAdaptor adaptor(cfg, file_per_key_config); + auto init_result = adaptor.Init(); + ASSERT_FALSE(init_result.has_value()); + EXPECT_EQ(init_result.error(), ErrorCode::INVALID_PARAMS); +} + +TEST_F(StorageBackendTest, AdaptorInitCreatesMissingRootDir) { + FileStorageConfig cfg; + cfg.storage_filepath = data_path + "/auto_created_root"; + FilePerKeyConfig file_per_key_config; + file_per_key_config.fsdir = "file_per_key_dir"; + file_per_key_config.enable_eviction = true; + + StorageBackendAdaptor adaptor(cfg, file_per_key_config); + ASSERT_TRUE(adaptor.Init().has_value()); + EXPECT_TRUE(fs::is_directory(cfg.storage_filepath)); +} + TEST_F(StorageBackendTest, StorageBackendAll) { std::shared_ptr client_buffer_allocator = std::make_shared(128 * 1024 * 1024); @@ -507,6 +590,132 @@ TEST_F(StorageBackendTest, OrphanedBucketFileCleanup) { ASSERT_TRUE(is_exist.value()); } +TEST_F(StorageBackendTest, MissingBucketDataFileCleanup) { + FileStorageConfig config; + config.storage_filepath = data_path; + BucketBackendConfig bucket_config; + + int64_t bucket_id = 0; + { + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::string value = "restart_data"; + std::unordered_map> batch; + batch.emplace("missing_data_key", + std::vector{Slice{value.data(), value.size()}}); + + auto result = storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(result.has_value()); + bucket_id = result.value(); + } + + const auto bucket_path = + fs::path(data_path) / (std::to_string(bucket_id) + ".bucket"); + const auto metadata_path = + fs::path(data_path) / (std::to_string(bucket_id) + ".meta"); + ASSERT_TRUE(fs::remove(bucket_path)); + ASSERT_TRUE(fs::exists(metadata_path)); + + BucketStorageBackend restarted_backend(config, bucket_config); + ASSERT_TRUE(restarted_backend.Init()); + EXPECT_FALSE(fs::exists(metadata_path)); + + auto exists = restarted_backend.IsExist("missing_data_key"); + ASSERT_TRUE(exists.has_value()); + EXPECT_FALSE(exists.value()); + + std::vector recovered_keys; + auto scan_result = + restarted_backend.ScanMeta([&](const std::vector& keys, + std::vector&) { + recovered_keys.insert(recovered_keys.end(), keys.begin(), + keys.end()); + return ErrorCode::OK; + }); + ASSERT_TRUE(scan_result.has_value()); + EXPECT_TRUE(recovered_keys.empty()); +} + +TEST_F(StorageBackendTest, BatchOffloadRollbackOnCompleteHandlerFailure) { + std::string test_dir = data_path + "/rollback_test"; + fs::create_directories(test_dir); + + FileStorageConfig config; + config.storage_filepath = test_dir; + BucketBackendConfig bucket_config; + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::shared_ptr client_buffer_allocator = + std::make_shared(128 * 1024 * 1024); + + // Prepare 3 keys to verify ALL keys are rolled back, not just one + std::unordered_map> batched_slices; + std::vector test_keys; + for (int i = 0; i < 3; ++i) { + std::string key = "rollback_test_key_" + std::to_string(i); + std::string data = "test_data_for_rollback_" + std::to_string(i); + void* buffer = client_buffer_allocator->allocate(data.size()); + memcpy(buffer, data.data(), data.size()); + batched_slices.emplace(key, + std::vector{Slice{buffer, data.size()}}); + test_keys.push_back(key); + } + + // Capture pre-offload state + auto pre_meta = storage_backend.GetStoreMetadata(); + ASSERT_TRUE(pre_meta.has_value()); + int64_t pre_total_size = pre_meta->total_size; + int64_t pre_total_keys = pre_meta->total_keys; + + // Trigger rollback via failing complete_handler + auto offload_res = storage_backend.BatchOffload( + batched_slices, [](const std::vector& keys, + std::vector& metadatas) { + return ErrorCode::INTERNAL_ERROR; + }); + + // Assertion 1: Offload returns the handler's error + EXPECT_FALSE(offload_res.has_value()); + EXPECT_EQ(offload_res.error(), ErrorCode::INTERNAL_ERROR); + + // Assertion 2: All keys removed from object_bucket_map_ + for (const auto& key : test_keys) { + auto exist_res = storage_backend.IsExist(key); + ASSERT_TRUE(exist_res.has_value()); + EXPECT_FALSE(exist_res.value()) + << "Key '" << key << "' should not exist after rollback"; + } + + // Assertion 3: total_size_ and total_keys restored to pre-offload values + auto post_meta = storage_backend.GetStoreMetadata(); + ASSERT_TRUE(post_meta.has_value()); + EXPECT_EQ(post_meta->total_size, pre_total_size) + << "total_size_ should be restored after rollback"; + EXPECT_EQ(post_meta->total_keys, pre_total_keys) + << "total_keys should be restored after rollback"; + + // Assertions 4 & 5: Bucket files deleted from disk. + // The bucket ID is generated from a timestamp-based BucketIdGenerator, so + // we can't hardcode it. Instead, scan the test directory to confirm no + // .bucket or .meta files remain after rollback. + int bucket_file_count = 0; + for (const auto& entry : fs::directory_iterator(test_dir)) { + if (entry.is_regular_file()) { + std::string ext = entry.path().extension().string(); + if (ext == ".bucket" || ext == ".meta") { + bucket_file_count++; + } + } + } + EXPECT_EQ(bucket_file_count, 0) + << "No bucket data or metadata files should remain after rollback"; +} + TEST_F(StorageBackendTest, AdaptorBatchOffloadAndBatchLoad) { FileStorageConfig cfg; @@ -2666,49 +2875,1854 @@ TEST_F(StorageBackendTest, StoreObjectEvictionWithEmptyKey) { EXPECT_TRUE(r3.value().empty()); } -TEST_F(StorageBackendTest, AdaptorBatchOffload_EvictionHandlerCalled) { - // Test that the eviction_handler callback in BatchOffload is correctly - // invoked when the underlying StorageBackend evicts files during - // StoreObject. We use a direct StorageBackend with a small quota to - // guarantee eviction, then verify via StorageBackendAdaptor that - // the handler fires. - // - // Since StorageBackendAdaptor::Init doesn't forward quota to the - // underlying StorageBackend, we test at the StoreObject level (already - // covered by StoreObjectReturnsEvictedKeys) and verify the BatchOffload - // handler wiring here with a mock-like capture. +TEST_F(StorageBackendTest, StoreObjectWatermarkEvictionReturnsEvictedKeys) { + std::string test_dir = data_path + "/watermark_evict_test"; + std::filesystem::create_directories(test_dir); + + StorageBackend backend(test_dir, "", true); + auto init_result = backend.Init(4096); + ASSERT_TRUE(init_result.has_value()); + + std::string data(1024, 'Z'); + ASSERT_TRUE(backend.StoreObject(test_dir + "/f1", data, "key_1")); + ASSERT_TRUE(backend.StoreObject(test_dir + "/f2", data, "key_2")); + ASSERT_TRUE(backend.StoreObject(test_dir + "/f3", data, "key_3")); + + auto evict_result = backend.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/0.70, /*low_watermark_ratio=*/0.40); + ASSERT_TRUE(evict_result.has_value()); + + const auto& evicted_keys = evict_result.value(); + ASSERT_EQ(evicted_keys.size(), 2); + EXPECT_EQ(evicted_keys[0], "key_1"); + EXPECT_EQ(evicted_keys[1], "key_2"); + EXPECT_FALSE(std::filesystem::exists(test_dir + "/f1")); + EXPECT_FALSE(std::filesystem::exists(test_dir + "/f2")); + EXPECT_TRUE(std::filesystem::exists(test_dir + "/f3")); + + auto second_evict = backend.EvictAboveDiskWatermark(0.70, 0.40); + ASSERT_TRUE(second_evict.has_value()); + EXPECT_TRUE(second_evict.value().empty()); +} - std::string test_dir = data_path + "/eviction_handler_test"; +TEST_F(StorageBackendTest, + StoreObjectWatermarkEvictionKeepsFilesWhenNotificationFails) { + std::string test_dir = data_path + "/watermark_notify_fail_test"; std::filesystem::create_directories(test_dir); - // Create backend with small quota (3072 bytes = room for ~3 files of 1024) StorageBackend backend(test_dir, "", true); - auto init_result = backend.Init(3072); + auto init_result = backend.Init(4096); ASSERT_TRUE(init_result.has_value()); - // Pre-fill with keyed files - std::string data(1024, 'A'); - auto r1 = backend.StoreObject(test_dir + "/f1", data, "key_1"); - ASSERT_TRUE(r1.has_value()); - auto r2 = backend.StoreObject(test_dir + "/f2", data, "key_2"); - ASSERT_TRUE(r2.has_value()); - auto r3 = backend.StoreObject(test_dir + "/f3", data, "key_3"); - ASSERT_TRUE(r3.has_value()); + std::string data(1024, 'N'); + ASSERT_TRUE(backend.StoreObject(test_dir + "/f1", data, "key_1")); + ASSERT_TRUE(backend.StoreObject(test_dir + "/f2", data, "key_2")); + ASSERT_TRUE(backend.StoreObject(test_dir + "/f3", data, "key_3")); + + std::vector notified_keys; + auto failed_evict = backend.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/0.70, /*low_watermark_ratio=*/0.40, + [&](const std::vector& evicted_keys) + -> tl::expected { + notified_keys = evicted_keys; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + }); + ASSERT_FALSE(failed_evict.has_value()); + EXPECT_EQ(failed_evict.error(), ErrorCode::INTERNAL_ERROR); + ASSERT_EQ(notified_keys.size(), 2); + EXPECT_EQ(notified_keys[0], "key_1"); + EXPECT_EQ(notified_keys[1], "key_2"); + EXPECT_TRUE(std::filesystem::exists(test_dir + "/f1")); + EXPECT_TRUE(std::filesystem::exists(test_dir + "/f2")); + EXPECT_TRUE(std::filesystem::exists(test_dir + "/f3")); + + auto successful_evict = backend.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/0.70, /*low_watermark_ratio=*/0.40, + [](const std::vector&) -> tl::expected { + return {}; + }); + ASSERT_TRUE(successful_evict.has_value()); + ASSERT_EQ(successful_evict.value().size(), 2); + EXPECT_EQ(successful_evict.value()[0], "key_1"); + EXPECT_EQ(successful_evict.value()[1], "key_2"); + EXPECT_FALSE(std::filesystem::exists(test_dir + "/f1")); + EXPECT_FALSE(std::filesystem::exists(test_dir + "/f2")); + EXPECT_TRUE(std::filesystem::exists(test_dir + "/f3")); +} + +TEST_F(StorageBackendTest, StoreObjectRejectsOverwriteDuringFailedEviction) { + std::string test_dir = data_path + "/eviction_overwrite_race_test"; + std::filesystem::create_directories(test_dir); + + constexpr size_t kUnit = 1024; + StorageBackend backend(test_dir, "", true); + ASSERT_TRUE(backend.Init(3 * kUnit)); + + const std::string overwritten_path = test_dir + "/overwritten"; + const std::string incoming_path = test_dir + "/incoming"; + const std::string old_value(kUnit, 'A'); + const std::string replacement_value(kUnit / 2, 'B'); + const std::string incoming_value(5 * kUnit / 2, 'C'); + ASSERT_TRUE(backend.StoreObject(overwritten_path, old_value, "old_key") + .has_value()); + + std::optional, ErrorCode>> + overwrite_result; + auto eviction_result = backend.StoreObject( + incoming_path, incoming_value, "incoming_key", + [&](const std::vector& keys) + -> tl::expected { + EXPECT_EQ(keys, std::vector{"old_key"}); + overwrite_result.emplace(backend.StoreObject( + overwritten_path, replacement_value, "new_key")); + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + }); + + ASSERT_TRUE(overwrite_result.has_value()); + ASSERT_FALSE(overwrite_result->has_value()); + EXPECT_EQ(overwrite_result->error(), ErrorCode::FILE_WRITE_FAIL); + ASSERT_FALSE(eviction_result.has_value()); + EXPECT_EQ(eviction_result.error(), ErrorCode::INTERNAL_ERROR); + + std::string loaded; + auto load_result = backend.LoadObject( + overwritten_path, loaded, static_cast(old_value.size())); + ASSERT_TRUE(load_result.has_value()); + EXPECT_EQ(loaded, old_value); + + ASSERT_TRUE( + backend.StoreObject(overwritten_path, replacement_value, "new_key") + .has_value()); + ASSERT_TRUE( + backend.StoreObject(incoming_path, incoming_value, "incoming_key") + .has_value()); + + loaded.clear(); + ASSERT_TRUE(backend + .LoadObject(overwritten_path, loaded, + static_cast(replacement_value.size())) + .has_value()); + EXPECT_EQ(loaded, replacement_value); +} + +TEST_F(StorageBackendTest, StoreObjectOverwriteReleasesPreviousReservation) { + std::string test_dir = data_path + "/overwrite_accounting_test"; + std::filesystem::create_directories(test_dir); + + constexpr size_t kUnit = 1024; + StorageBackend backend(test_dir, "", true); + ASSERT_TRUE(backend.Init(2 * kUnit)); + + const std::string overwritten_path = test_dir + "/overwritten"; + const std::string second_path = test_dir + "/second"; + const std::string old_value(kUnit, 'A'); + const std::string replacement_value(kUnit, 'B'); + const std::string second_value(kUnit, 'C'); + + ASSERT_TRUE(backend.StoreObject(overwritten_path, old_value, "old_key") + .has_value()); + ASSERT_TRUE( + backend + .StoreObject(overwritten_path, replacement_value, "replacement_key") + .has_value()); - // Now store one more, which should evict key_1 std::vector evicted_keys; - auto r4 = backend.StoreObject(test_dir + "/f4", data, "key_4"); - ASSERT_TRUE(r4.has_value()); - for (const auto& ek : r4.value()) { - evicted_keys.push_back(ek); + auto second_result = + backend.StoreObject(second_path, second_value, "second_key", + [&](const std::vector& keys) + -> tl::expected { + evicted_keys = keys; + return {}; + }); + + ASSERT_TRUE(second_result.has_value()); + EXPECT_TRUE(evicted_keys.empty()); + EXPECT_TRUE(std::filesystem::exists(overwritten_path)); + EXPECT_TRUE(std::filesystem::exists(second_path)); + + std::string loaded; + ASSERT_TRUE(backend + .LoadObject(overwritten_path, loaded, + static_cast(replacement_value.size())) + .has_value()); + EXPECT_EQ(loaded, replacement_value); +} + +TEST_F(StorageBackendTest, + AdaptorWatermarkEvictionNotifiesRecoveredKeysAfterRestart) { + FileStorageConfig cfg; + cfg.storage_filepath = data_path + "/"; + cfg.scanmeta_iterator_keys_limit = 16; + + FilePerKeyConfig file_per_key_config; + file_per_key_config.fsdir = "file_per_key_watermark_restart"; + file_per_key_config.enable_eviction = true; + + std::unordered_map test_data = { + {"restart_key_1", std::string(512, 'a')}, + {"restart_key_2", std::string(512, 'b')}, + {"restart_key_3", std::string(512, 'c')}, + }; + + { + StorageBackendAdaptor adaptor(cfg, file_per_key_config); + ASSERT_TRUE(adaptor.Init()); + + std::unordered_map> batch_object; + std::vector> write_buffers; + for (auto& [key, value] : test_data) { + auto buf = std::make_unique(value.size()); + std::memcpy(buf.get(), value.data(), value.size()); + batch_object.emplace( + key, std::vector{Slice{buf.get(), value.size()}}); + write_buffers.emplace_back(std::move(buf)); + } + + auto offload_res = adaptor.BatchOffload( + batch_object, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(offload_res); } - EXPECT_FALSE(evicted_keys.empty()) - << "Should have evicted at least one key"; - EXPECT_EQ(evicted_keys[0], "key_1") - << "FIFO eviction should evict key_1 first"; + StorageBackendAdaptor restart_adaptor(cfg, file_per_key_config); + ASSERT_TRUE(restart_adaptor.Init()); + auto scan_res = restart_adaptor.ScanMeta( + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(scan_res); + + std::vector notified_keys; + auto evict_result = restart_adaptor.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/1e-12, + /*low_watermark_ratio=*/0.5e-12, + [&](const std::vector& evicted_keys) + -> tl::expected { + notified_keys = evicted_keys; + return {}; + }); + ASSERT_TRUE(evict_result.has_value()); + + auto returned_keys = evict_result.value(); + std::sort(returned_keys.begin(), returned_keys.end()); + std::sort(notified_keys.begin(), notified_keys.end()); + std::vector expected_keys = {"restart_key_1", "restart_key_2", + "restart_key_3"}; + EXPECT_EQ(returned_keys, expected_keys); + EXPECT_EQ(notified_keys, expected_keys); } -//----------------------------------------------------------------------------- +TEST_F(StorageBackendTest, BucketWatermarkEvictionUsesHandlerAndKeepsNewest) { + FileStorageConfig config; + config.storage_filepath = data_path; + + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 10; + bucket_config.bucket_size_limit = 8 * 1024; + bucket_config.max_total_size = 30 * 1024; + bucket_config.eviction_policy = BucketEvictionPolicy::FIFO; + + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::vector> buffers; + for (int i = 0; i < 3; ++i) { + const std::string key = "bucket_key_" + std::to_string(i); + auto buffer = std::make_unique(6 * 1024); + std::memset(buffer.get(), static_cast('A' + i), 6 * 1024); + std::unordered_map> batch; + batch.emplace(key, std::vector{Slice{buffer.get(), 6 * 1024}}); + buffers.push_back(std::move(buffer)); + + auto result = storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(result.has_value()); + } + + std::vector notified_keys; + auto evict_result = storage_backend.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/0.50, /*low_watermark_ratio=*/0.25, + [&](const std::vector& evicted_keys) { + notified_keys.insert(notified_keys.end(), evicted_keys.begin(), + evicted_keys.end()); + return tl::expected{}; + }); + ASSERT_TRUE(evict_result.has_value()); + ASSERT_FALSE(evict_result.value().empty()); + EXPECT_EQ(notified_keys, evict_result.value()); + EXPECT_EQ(evict_result.value().front(), "bucket_key_0"); + + auto oldest_exists = storage_backend.IsExist("bucket_key_0"); + ASSERT_TRUE(oldest_exists.has_value()); + EXPECT_FALSE(oldest_exists.value()); + + auto newest_exists = storage_backend.IsExist("bucket_key_2"); + ASSERT_TRUE(newest_exists.has_value()); + EXPECT_TRUE(newest_exists.value()); +} + +TEST_F(StorageBackendTest, + BucketWatermarkEvictionDoesNotOverEvictForSharedDisk) { + std::error_code ec; + auto space_info = fs::space(data_path, ec); + ASSERT_FALSE(ec); + constexpr uint64_t kMinFreeSpace = 256ULL * 1024 * 1024; + if (space_info.available <= kMinFreeSpace) { + GTEST_SKIP() << "Need more than 256MB free space to isolate the " + "synthetic-size check"; + } + if (space_info.available > + static_cast(std::numeric_limits::max() / 2)) { + GTEST_SKIP() << "Filesystem is too large for this synthetic quota test"; + } + + FileStorageConfig config; + config.storage_filepath = data_path; + + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 10; + bucket_config.bucket_size_limit = 8 * 1024; + bucket_config.max_total_size = + static_cast(space_info.available * 2); + bucket_config.eviction_policy = BucketEvictionPolicy::FIFO; + + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::vector> buffers; + for (int i = 0; i < 3; ++i) { + const std::string key = "shared_disk_bucket_key_" + std::to_string(i); + auto buffer = std::make_unique(6 * 1024); + std::memset(buffer.get(), static_cast('A' + i), 6 * 1024); + std::unordered_map> batch; + batch.emplace(key, std::vector{Slice{buffer.get(), 6 * 1024}}); + buffers.push_back(std::move(buffer)); + + auto result = storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(result.has_value()); + } + + constexpr int64_t kHighWatermarkBytes = 14 * 1024; + constexpr int64_t kLowWatermarkBytes = 12 * 1024; + auto evict_result = storage_backend.EvictAboveDiskWatermark( + static_cast(kHighWatermarkBytes) / bucket_config.max_total_size, + static_cast(kLowWatermarkBytes) / bucket_config.max_total_size, + [](const std::vector&) { + return tl::expected{}; + }); + + ASSERT_TRUE(evict_result.has_value()); + EXPECT_FALSE(evict_result.value().empty()); + EXPECT_LT(evict_result.value().size(), 3); + + auto newest_exists = storage_backend.IsExist("shared_disk_bucket_key_2"); + ASSERT_TRUE(newest_exists.has_value()); + EXPECT_TRUE(newest_exists.value()); +} + +TEST_F(StorageBackendTest, + BucketPendingEvictionRejectsConcurrentDuplicateWrite) { + FileStorageConfig config; + config.storage_filepath = data_path; + + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 10; + bucket_config.bucket_size_limit = 8 * 1024; + bucket_config.max_total_size = 10 * 1024; + bucket_config.eviction_policy = BucketEvictionPolicy::FIFO; + + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::string old_value(6 * 1024, 'A'); + std::unordered_map> old_batch; + old_batch.emplace("old_key", std::vector{Slice{old_value.data(), + old_value.size()}}); + ASSERT_TRUE(storage_backend + .BatchOffload(old_batch, + [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }) + .has_value()); + + std::string incoming_value(6 * 1024, 'B'); + std::unordered_map> incoming_batch; + incoming_batch.emplace("incoming_key", + std::vector{Slice{incoming_value.data(), + incoming_value.size()}}); + std::string replacement_value(3 * 1024, 'C'); + std::unordered_map> replacement_batch; + replacement_batch.emplace( + "old_key", std::vector{Slice{replacement_value.data(), + replacement_value.size()}}); + std::optional> replacement_result; + auto failed_eviction_result = storage_backend.BatchOffload( + incoming_batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }, + [&](const std::vector& keys) + -> tl::expected { + EXPECT_EQ(keys, std::vector{"old_key"}); + replacement_result.emplace(storage_backend.BatchOffload( + replacement_batch, [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + })); + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + }); + + ASSERT_FALSE(failed_eviction_result.has_value()); + EXPECT_EQ(failed_eviction_result.error(), ErrorCode::INTERNAL_ERROR); + + ASSERT_TRUE(replacement_result.has_value()); + ASSERT_FALSE(replacement_result->has_value()); + EXPECT_EQ(replacement_result->error(), ErrorCode::OBJECT_ALREADY_EXISTS); + + std::vector load_buffer(old_value.size()); + std::unordered_map load_batch; + load_batch.emplace( + "old_key", + Slice{load_buffer.data(), static_cast(load_buffer.size())}); + ASSERT_TRUE(storage_backend.BatchLoad(load_batch).has_value()); + EXPECT_EQ(std::string(load_buffer.begin(), load_buffer.end()), old_value); +} + +TEST_F(StorageBackendTest, + BucketRollbackPreservesCapacityAgainstConcurrentWrite) { + FileStorageConfig config; + config.storage_filepath = data_path; + + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 10; + bucket_config.bucket_size_limit = 8 * 1024; + bucket_config.max_total_size = 10 * 1024; + bucket_config.eviction_policy = BucketEvictionPolicy::FIFO; + + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::string old_value(6 * 1024, 'A'); + std::unordered_map> old_batch; + old_batch.emplace("old_key", std::vector{Slice{old_value.data(), + old_value.size()}}); + ASSERT_TRUE(storage_backend + .BatchOffload(old_batch, + [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }) + .has_value()); + + std::string incoming_value(6 * 1024, 'B'); + std::unordered_map> incoming_batch; + incoming_batch.emplace("incoming_key", + std::vector{Slice{incoming_value.data(), + incoming_value.size()}}); + + std::string concurrent_value(3 * 1024, 'C'); + std::unordered_map> concurrent_batch; + concurrent_batch.emplace( + "concurrent_key", std::vector{Slice{concurrent_value.data(), + concurrent_value.size()}}); + + std::optional> concurrent_result; + auto failed_eviction_result = storage_backend.BatchOffload( + incoming_batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }, + [&](const std::vector& keys) + -> tl::expected { + EXPECT_EQ(keys, std::vector{"old_key"}); + concurrent_result.emplace(storage_backend.BatchOffload( + concurrent_batch, [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + })); + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + }); + + ASSERT_TRUE(concurrent_result.has_value()); + ASSERT_FALSE(concurrent_result->has_value()); + EXPECT_EQ(concurrent_result->error(), ErrorCode::FILE_WRITE_FAIL); + ASSERT_FALSE(failed_eviction_result.has_value()); + EXPECT_EQ(failed_eviction_result.error(), ErrorCode::INTERNAL_ERROR); + + auto old_exists = storage_backend.IsExist("old_key"); + ASSERT_TRUE(old_exists.has_value()); + EXPECT_TRUE(old_exists.value()); + + ASSERT_TRUE(storage_backend + .BatchOffload(concurrent_batch, + [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }) + .has_value()); +} + +TEST_F(StorageBackendTest, BucketPendingWriteRejectsReentrantDuplicate) { + FileStorageConfig config; + config.storage_filepath = data_path; + + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 10; + bucket_config.bucket_size_limit = 8 * 1024; + bucket_config.max_total_size = 20 * 1024; + bucket_config.eviction_policy = BucketEvictionPolicy::FIFO; + + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::string outer_value(3 * 1024, 'A'); + std::unordered_map> outer_batch; + outer_batch.emplace( + "shared_key", + std::vector{Slice{outer_value.data(), outer_value.size()}}); + + std::string nested_value(1024, 'B'); + std::unordered_map> nested_batch; + nested_batch.emplace( + "shared_key", + std::vector{Slice{nested_value.data(), nested_value.size()}}); + + std::optional> nested_result; + auto outer_result = storage_backend.BatchOffload( + outer_batch, [&](const std::vector&, + std::vector&) { + nested_result.emplace(storage_backend.BatchOffload( + nested_batch, [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + })); + return ErrorCode::OK; + }); + + ASSERT_TRUE(outer_result.has_value()); + ASSERT_TRUE(nested_result.has_value()); + ASSERT_FALSE(nested_result->has_value()); + EXPECT_EQ(nested_result->error(), ErrorCode::OBJECT_ALREADY_EXISTS); + + std::vector load_buffer(outer_value.size()); + std::unordered_map load_batch; + load_batch.emplace( + "shared_key", + Slice{load_buffer.data(), static_cast(load_buffer.size())}); + ASSERT_TRUE(storage_backend.BatchLoad(load_batch).has_value()); + EXPECT_EQ(std::string(load_buffer.begin(), load_buffer.end()), outer_value); +} + +TEST_F(StorageBackendTest, BucketBatchOffloadContinuesAfterFinalizeFailure) { + FileStorageConfig config; + config.storage_filepath = data_path; + + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 10; + bucket_config.bucket_size_limit = 8 * 1024; + bucket_config.max_total_size = 10 * 1024; + bucket_config.eviction_policy = BucketEvictionPolicy::FIFO; + + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::string old_value(6 * 1024, 'A'); + std::unordered_map> old_batch; + old_batch.emplace("old_key", std::vector{Slice{old_value.data(), + old_value.size()}}); + auto old_result = storage_backend.BatchOffload( + old_batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(old_result.has_value()); + + const auto old_bucket_path = + fs::path(data_path) / (std::to_string(old_result.value()) + ".bucket"); + const auto old_metadata_path = + fs::path(data_path) / (std::to_string(old_result.value()) + ".meta"); + ASSERT_TRUE(fs::remove(old_bucket_path)); + ASSERT_TRUE(fs::create_directory(old_bucket_path)); + { + std::ofstream blocker(old_bucket_path / "blocker"); + ASSERT_TRUE(blocker.is_open()); + blocker << "prevent directory removal"; + } + + std::vector notified_keys; + std::string new_value(6 * 1024, 'B'); + std::unordered_map> new_batch; + new_batch.emplace("new_key", std::vector{Slice{new_value.data(), + new_value.size()}}); + auto new_result = storage_backend.BatchOffload( + new_batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }, + [&](const std::vector& keys) { + notified_keys = keys; + return tl::expected{}; + }); + + ASSERT_TRUE(new_result.has_value()); + EXPECT_EQ(notified_keys, std::vector{"old_key"}); + EXPECT_FALSE(fs::exists(old_metadata_path)); + EXPECT_TRUE(fs::exists(old_bucket_path)); + + auto old_exists = storage_backend.IsExist("old_key"); + ASSERT_TRUE(old_exists.has_value()); + EXPECT_FALSE(old_exists.value()); + + auto new_exists = storage_backend.IsExist("new_key"); + ASSERT_TRUE(new_exists.has_value()); + EXPECT_TRUE(new_exists.value()); + + BucketStorageBackend restarted_backend(config, bucket_config); + ASSERT_TRUE(restarted_backend.Init()); + + auto restarted_old_exists = restarted_backend.IsExist("old_key"); + ASSERT_TRUE(restarted_old_exists.has_value()); + EXPECT_FALSE(restarted_old_exists.value()); + + auto restarted_new_exists = restarted_backend.IsExist("new_key"); + ASSERT_TRUE(restarted_new_exists.has_value()); + EXPECT_TRUE(restarted_new_exists.value()); +} + +TEST_F(StorageBackendTest, + BucketWatermarkEvictionReturnsKeysAfterFinalizeFailure) { + FileStorageConfig config; + config.storage_filepath = data_path; + + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 10; + bucket_config.bucket_size_limit = 8 * 1024; + bucket_config.max_total_size = 10 * 1024; + bucket_config.eviction_policy = BucketEvictionPolicy::FIFO; + + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::string value(6 * 1024, 'W'); + std::unordered_map> batch; + batch.emplace("watermark_key", + std::vector{Slice{value.data(), value.size()}}); + auto result = storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(result.has_value()); + + const auto bucket_path = + fs::path(data_path) / (std::to_string(result.value()) + ".bucket"); + const auto metadata_path = + fs::path(data_path) / (std::to_string(result.value()) + ".meta"); + ASSERT_TRUE(fs::remove(bucket_path)); + ASSERT_TRUE(fs::create_directory(bucket_path)); + { + std::ofstream blocker(bucket_path / "blocker"); + ASSERT_TRUE(blocker.is_open()); + blocker << "prevent directory removal"; + } + + std::vector notified_keys; + auto eviction_result = storage_backend.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/0.50, /*low_watermark_ratio=*/0.25, + [&](const std::vector& keys) { + notified_keys = keys; + return tl::expected{}; + }); + + ASSERT_TRUE(eviction_result.has_value()); + EXPECT_EQ(eviction_result.value(), + std::vector{"watermark_key"}); + EXPECT_EQ(notified_keys, eviction_result.value()); + EXPECT_FALSE(fs::exists(metadata_path)); + EXPECT_TRUE(fs::exists(bucket_path)); + + auto exists = storage_backend.IsExist("watermark_key"); + ASSERT_TRUE(exists.has_value()); + EXPECT_FALSE(exists.value()); + + BucketStorageBackend restarted_backend(config, bucket_config); + ASSERT_TRUE(restarted_backend.Init()); + + auto restarted_exists = restarted_backend.IsExist("watermark_key"); + ASSERT_TRUE(restarted_exists.has_value()); + EXPECT_FALSE(restarted_exists.value()); + + std::vector recovered_keys; + auto scan_result = + restarted_backend.ScanMeta([&](const std::vector& keys, + std::vector&) { + recovered_keys.insert(recovered_keys.end(), keys.begin(), + keys.end()); + return ErrorCode::OK; + }); + ASSERT_TRUE(scan_result.has_value()); + EXPECT_TRUE(recovered_keys.empty()); +} + +TEST_F(StorageBackendTest, + BucketWatermarkEvictionRestoresMetadataWhenNotificationFails) { + FileStorageConfig config; + config.storage_filepath = data_path; + + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 10; + bucket_config.bucket_size_limit = 8 * 1024; + bucket_config.max_total_size = 30 * 1024; + bucket_config.eviction_policy = BucketEvictionPolicy::FIFO; + + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + std::vector> buffers; + for (int i = 0; i < 3; ++i) { + const std::string key = "rollback_bucket_key_" + std::to_string(i); + auto buffer = std::make_unique(6 * 1024); + std::memset(buffer.get(), static_cast('A' + i), 6 * 1024); + std::unordered_map> batch; + batch.emplace(key, std::vector{Slice{buffer.get(), 6 * 1024}}); + buffers.push_back(std::move(buffer)); + + auto result = storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(result.has_value()); + } + + auto failed_evict = storage_backend.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/0.50, /*low_watermark_ratio=*/0.25, + [](const std::vector&) -> tl::expected { + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + }); + ASSERT_FALSE(failed_evict.has_value()); + EXPECT_EQ(failed_evict.error(), ErrorCode::INTERNAL_ERROR); + + auto oldest_exists = storage_backend.IsExist("rollback_bucket_key_0"); + ASSERT_TRUE(oldest_exists.has_value()); + EXPECT_TRUE(oldest_exists.value()); + + auto successful_evict = storage_backend.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/0.50, /*low_watermark_ratio=*/0.25, + [](const std::vector&) -> tl::expected { + return {}; + }); + ASSERT_TRUE(successful_evict.has_value()); + oldest_exists = storage_backend.IsExist("rollback_bucket_key_0"); + ASSERT_TRUE(oldest_exists.has_value()); + EXPECT_FALSE(oldest_exists.value()); +} + +TEST_F(StorageBackendTest, BucketWatermarkEvictionNoopsWhenPolicyIsNone) { + FileStorageConfig config; + config.storage_filepath = data_path; + + BucketBackendConfig bucket_config; + bucket_config.bucket_keys_limit = 10; + bucket_config.bucket_size_limit = 8 * 1024; + bucket_config.max_total_size = 30 * 1024; + bucket_config.eviction_policy = BucketEvictionPolicy::NONE; + + BucketStorageBackend storage_backend(config, bucket_config); + ASSERT_TRUE(storage_backend.Init()); + + auto buffer = std::make_unique(6 * 1024); + std::memset(buffer.get(), 'N', 6 * 1024); + std::unordered_map> batch; + batch.emplace("no_evict_key", + std::vector{Slice{buffer.get(), 6 * 1024}}); + + auto result = storage_backend.BatchOffload( + batch, + [](const std::vector&, + std::vector&) { return ErrorCode::OK; }); + ASSERT_TRUE(result.has_value()); + + bool handler_called = false; + auto evict_result = storage_backend.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/0.01, /*low_watermark_ratio=*/0.005, + [&](const std::vector&) -> tl::expected { + handler_called = true; + return {}; + }); + ASSERT_TRUE(evict_result.has_value()); + EXPECT_TRUE(evict_result.value().empty()); + EXPECT_FALSE(handler_called); + + auto exists = storage_backend.IsExist("no_evict_key"); + ASSERT_TRUE(exists.has_value()); + EXPECT_TRUE(exists.value()); +} + +TEST_F(StorageBackendTest, OffsetAllocatorWatermarkEvictionNoops) { + FileStorageConfig config; + config.storage_filepath = data_path; + + OffsetAllocatorStorageBackend storage_backend(config); + + bool handler_called = false; + auto evict_result = storage_backend.EvictAboveDiskWatermark( + /*high_watermark_ratio=*/0.01, /*low_watermark_ratio=*/0.005, + [&](const std::vector&) -> tl::expected { + handler_called = true; + return {}; + }); + + ASSERT_TRUE(evict_result.has_value()); + EXPECT_TRUE(evict_result.value().empty()); + EXPECT_FALSE(handler_called); +} + +//----------------------------------------------------------------------------- + +//----------------------------------------------------------------------------- +// OffsetAllocatorStorageBackend Eviction Tests +//----------------------------------------------------------------------------- + +// Helper: build a BatchOffload request for a single key/value pair. +std::unordered_map> MakeSingleKeyBatch( + const std::string& key, const std::string& value, + std::vector>& buffers) { + auto buf = std::make_unique(value.size()); + std::memcpy(buf.get(), value.data(), value.size()); + buffers.push_back(std::move(buf)); + std::unordered_map> batch; + batch.emplace( + key, std::vector{Slice{buffers.back().get(), value.size()}}); + return batch; +} + +TEST_F(StorageBackendTest, OffsetAllocatorStorageBackend_Eviction_FifoOrder) { + // Verify that when watermark-triggered eviction fires, the oldest + // key (by insertion order) is evicted first. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = 20 * 1024; // 20KB — very small arena + config.total_keys_limit = 100; + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + std::vector evicted_keys; + auto eviction_handler = + [&evicted_keys](const std::vector& keys) { + for (const auto& k : keys) evicted_keys.push_back(k); + return tl::expected{}; + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + // Write A, B, C one by one. The arena is 20 KB; each 1 KB record is + // 4K-aligned (header + key + padding + value = 4096 + 1000 = 5096 + // bytes), so three fit and the fourth write should trigger eviction. + std::string data(1000, 'x'); + std::vector> buffers; + + for (const auto& key : {"key_a", "key_b", "key_c"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + [[maybe_unused]] auto res = storage_backend.BatchOffload( + batch, complete_handler, eviction_handler); + ASSERT_TRUE(res.has_value()) << "key=" << key; + } + + // The fourth write should push total_size_ over high_watermark_bytes_ + // and evict key_a (the oldest). + auto batch = MakeSingleKeyBatch("key_d", data, buffers); + [[maybe_unused]] auto res = + storage_backend.BatchOffload(batch, complete_handler, eviction_handler); + ASSERT_TRUE(res.has_value()); + + ASSERT_FALSE(evicted_keys.empty()) + << "Should have evicted at least one key"; + EXPECT_EQ(evicted_keys[0], "key_a") + << "FIFO eviction must evict the oldest key first"; + EXPECT_FALSE(storage_backend.IsExist("key_a").value_or(true)) + << "key_a should no longer exist after eviction"; + EXPECT_TRUE(storage_backend.IsExist("key_d").value_or(false)) + << "key_d should exist"; +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_NoEvictionWhenNONE) { + // Under default NONE policy, the allocate-fail path still breaks. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = 16 * 1024; // fits two 4K-aligned records + config.total_keys_limit = 100; + + OffsetAllocatorStorageBackend storage_backend(config); + ASSERT_TRUE(storage_backend.Init()); + + std::vector evicted_keys; + auto eviction_handler = + [&evicted_keys](const std::vector& keys) { + for (const auto& k : keys) evicted_keys.push_back(k); + return tl::expected{}; + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + std::string data(1500, 'x'); + std::vector> buffers; + + int offloaded = 0; + for (const auto& key : {"key_a", "key_b", "key_c", "key_d"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + [[maybe_unused]] auto res = storage_backend.BatchOffload( + batch, complete_handler, eviction_handler); + if (res.has_value()) + ++offloaded; + else + break; // allocation failure should break + } + + EXPECT_GT(offloaded, 0); + EXPECT_LT(offloaded, 5) << "NONE policy should break on allocation failure"; + EXPECT_TRUE(evicted_keys.empty()) << "NONE policy should never evict"; +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_MasterNotifiedBeforeReuse) { + // Verify that eviction_handler is called BEFORE allocate() for the + // key whose eviction made room, i.e. the notify-before-reuse contract. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = 20 * 1024; + config.total_keys_limit = 100; + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + bool handler_called_before_allocation = false; + bool allocate_happened = false; + + auto eviction_handler = + [&handler_called_before_allocation, + &allocate_happened](const std::vector& keys) { + if (!keys.empty() && !allocate_happened) { + handler_called_before_allocation = true; + } + return tl::expected{}; + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + // The test doesn't have direct instrumentation for "allocate() just + // happened". But the design guarantees that eviction_handler is called + // at (B) before (C) in BatchOffload. We verify indirectly: + // after writing enough to trigger eviction, the handler must have been + // invoked with at least one key AND that key no longer exists. + std::string data(1000, 'x'); + std::vector> buffers; + + for (const auto& key : {"key_a", "key_b", "key_c"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + storage_backend.BatchOffload(batch, complete_handler, eviction_handler); + } + + std::vector captured_evicted; + auto capture_handler = + [&captured_evicted](const std::vector& keys) { + for (const auto& k : keys) captured_evicted.push_back(k); + return tl::expected{}; + }; + + auto batch = MakeSingleKeyBatch("key_d", data, buffers); + storage_backend.BatchOffload(batch, complete_handler, capture_handler); + + EXPECT_FALSE(captured_evicted.empty()) + << "Should have evicted at least one key"; + for (const auto& ek : captured_evicted) { + EXPECT_FALSE(storage_backend.IsExist(ek).value_or(true)) + << "Evicted key " << ek + << " should not exist (erased before reuse)"; + } + EXPECT_TRUE(storage_backend.IsExist("key_d").value_or(false)); +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_NotificationFailureRollback) { + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = 20 * 1024; + config.total_keys_limit = 100; + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + // Three 4K-aligned records (5096 B each) fit below high; the fourth + // crosses it and eviction drives back to low (one victim). + evict_cfg.high_watermark_bytes = 16384; + evict_cfg.low_watermark_bytes = 15288; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + auto successful_eviction_handler = + [](const std::vector&) -> tl::expected { + return {}; + }; + + std::string data(1000, 'x'); + std::vector> buffers; + for (const auto& key : {"key_a", "key_b", "key_c"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + auto result = storage_backend.BatchOffload(batch, complete_handler, + successful_eviction_handler); + ASSERT_TRUE(result.has_value()) << "key=" << key; + } + + std::vector failed_evictions; + auto failing_eviction_handler = + [&failed_evictions](const std::vector& keys) + -> tl::expected { + failed_evictions = keys; + return tl::make_unexpected(ErrorCode::INTERNAL_ERROR); + }; + + auto batch = MakeSingleKeyBatch("key_d", data, buffers); + auto failed_result = storage_backend.BatchOffload(batch, complete_handler, + failing_eviction_handler); + ASSERT_FALSE(failed_result.has_value()); + EXPECT_EQ(failed_result.error(), ErrorCode::INTERNAL_ERROR); + ASSERT_FALSE(failed_evictions.empty()); + EXPECT_EQ(failed_evictions.front(), "key_a"); + + for (const auto& key : failed_evictions) { + auto exists = storage_backend.IsExist(key); + ASSERT_TRUE(exists.has_value()); + EXPECT_TRUE(exists.value()) << "key=" << key; + + std::vector output(data.size()); + std::unordered_map load_batch; + load_batch.emplace(key, Slice{output.data(), output.size()}); + auto load_result = storage_backend.BatchLoad(load_batch); + ASSERT_TRUE(load_result.has_value()) << "key=" << key; + EXPECT_EQ(std::string(output.begin(), output.end()), data); + } + EXPECT_FALSE(storage_backend.IsExist("key_d").value_or(true)); + + std::vector retry_evictions; + auto retry_eviction_handler = + [&retry_evictions](const std::vector& keys) + -> tl::expected { + retry_evictions = keys; + return {}; + }; + auto retry_result = storage_backend.BatchOffload(batch, complete_handler, + retry_eviction_handler); + ASSERT_TRUE(retry_result.has_value()); + EXPECT_EQ(retry_evictions, failed_evictions); + for (const auto& key : retry_evictions) { + EXPECT_FALSE(storage_backend.IsExist(key).value_or(true)) + << "key=" << key; + } + EXPECT_TRUE(storage_backend.IsExist("key_d").value_or(false)); +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_PostLoopFlush) { + // The last key in a batch may trigger eviction but fail allocate. + // The evicted keys must still be flushed to the handler before return. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + // Arena large enough for several keys but small enough to trigger + // eviction near capacity. + config.total_size_limit = 24 * 1024; + config.total_keys_limit = 100; + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + std::vector all_evicted; + auto eviction_handler = + [&all_evicted](const std::vector& keys) { + for (const auto& k : keys) all_evicted.push_back(k); + return tl::expected{}; + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + std::string data(1500, 'x'); // each 4K-aligned record = 5596 bytes + std::vector> buffers; + + // Write enough to fill the arena and trigger eviction. + for (const auto& key : {"key_a", "key_b", "key_c", "key_d", "key_e"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + [[maybe_unused]] auto res = storage_backend.BatchOffload( + batch, complete_handler, eviction_handler); + // Some may succeed, some may fail — we just care that evicted + // keys are eventually reported. + } + + // If any eviction happened, the evicted keys should be reported. + // We don't assert non-empty because capacity calculations can vary; + // we just assert that if keys were evicted, they're no longer present. + if (!all_evicted.empty()) { + for (const auto& ek : all_evicted) { + EXPECT_FALSE(storage_backend.IsExist(ek).value_or(true)) + << "Evicted key " << ek << " should not exist"; + } + } +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_KeyCountTrigger) { + // Verify that eviction fires when total_keys_ exceeds the key-count + // high watermark, even when bytes are well below the byte watermark. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = 1 * 1024 * 1024; // 1 MB — plenty of bytes + config.total_keys_limit = 10; // only 10 keys allowed + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + std::vector evicted_keys; + auto eviction_handler = + [&evicted_keys](const std::vector& keys) { + for (const auto& k : keys) evicted_keys.push_back(k); + return tl::expected{}; + }; + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + // Each key is tiny (10 bytes), so bytes stay low. + std::string data(10, 'x'); + std::vector> buffers; + + for (int i = 0; i < 15; ++i) { + std::string key = "tiny_key_" + std::to_string(i); + auto batch = MakeSingleKeyBatch(key, data, buffers); + storage_backend.BatchOffload(batch, complete_handler, eviction_handler); + } + + // Key-count watermark should have triggered eviction. + EXPECT_FALSE(evicted_keys.empty()) + << "Key-count overflow should trigger eviction even with low bytes"; +} + +//----------------------------------------------------------------------------- + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Eviction_ConcurrentReadSafety) { + // Validate the load-bearing safety property: AllocationPtr refcount + // must keep an evicted key's physical extent alive (pinned) until + // all in-flight BatchLoad calls release their shared_ptr copies. + // + // Mechanism being tested: + // 1. BatchLoad copies entry.allocation (shared_ptr) into its + // ReadPlan, incrementing the refcount. (storage_backend.cpp + // ~line 3375: "entry.allocation" copy in ReadPlan) + // 2. EvictToMakeRoom retains the allocation shared_ptr in pending + // eviction state before erasing the key from shard.map. The pending + // reference is released only after the master accepts the removal. + // 3. While either pending eviction or a reader holds a shared_ptr, + // freeAllocation does NOT fire → the extent is still marked "used" + // in the allocator → allocate() cannot re-issue that offset. The + // reader always sees the original bytes. + // + // This test interleaves reads and eviction-triggering writes; + // any data corruption means the allocator re-issued a still-read + // offset, which would be a violation of the refcount contract. + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = + 160 * 1024; // tight enough that 80 4K-aligned records (~4.2KB + // each) trigger eviction + config.total_keys_limit = 500; + + OffsetAllocatorBackendConfig evict_cfg; + evict_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + + OffsetAllocatorStorageBackend storage_backend(config, evict_cfg); + ASSERT_TRUE(storage_backend.Init()); + + auto complete_handler = [](const std::vector&, + std::vector&) { + return ErrorCode::OK; + }; + + // Pre-populate with distinct-per-key data so a re-issued offset + // is detectable: if key_i's extent is handed to a new key and the + // reader still reads from it, read_buf[0] won't match 'A' + i%26. + std::vector> buffers; + const int kNumKeys = 30; + for (int i = 0; i < kNumKeys; ++i) { + std::string key = "ckey_" + std::to_string(i); + std::string val(100, static_cast('A' + (i % 26))); + auto batch = MakeSingleKeyBatch(key, val, buffers); + storage_backend.BatchOffload(batch, complete_handler); + } + // Confirm pre-populated keys exist before the stress phase. + for (int i = 0; i < std::min(kNumKeys, 5); ++i) { + EXPECT_TRUE(storage_backend.IsExist("ckey_" + std::to_string(i)) + .value_or(false)); + } + + // Eviction handler that tracks victim keys so we can assert post-hoc. + std::vector all_evicted; + std::mutex evict_mtx; + auto eviction_handler = [&all_evicted, + &evict_mtx](const std::vector& keys) { + std::lock_guard lk(evict_mtx); + for (const auto& k : keys) all_evicted.push_back(k); + return tl::expected{}; + }; + + std::atomic stop{false}; + std::atomic read_errors{0}; + std::atomic read_success{0}; + std::atomic read_not_found{0}; + + std::thread reader([&]() { + while (!stop) { + for (int i = 0; i < kNumKeys; ++i) { + std::string key = "ckey_" + std::to_string(i); + auto read_buf = std::make_unique(100); + std::unordered_map load; + load.emplace(key, Slice{read_buf.get(), 100}); + auto res = storage_backend.BatchLoad(load); + if (res.has_value()) { + read_success++; + if (read_buf[0] != static_cast('A' + (i % 26))) { + read_errors++; + } + } else { + read_not_found++; // expected after eviction + } + } + } + }); + + std::thread writer([&]() { + for (int i = kNumKeys; i < kNumKeys + 50 && !stop; ++i) { + std::string key = "newkey_" + std::to_string(i); + std::string val(100, 'Z'); + std::vector> wbufs; + auto batch = MakeSingleKeyBatch(key, val, wbufs); + storage_backend.BatchOffload(batch, complete_handler, + eviction_handler); + } + stop = true; + }); + + writer.join(); + reader.join(); + + // Core safety assertion: zero data corruption across all reads. + EXPECT_EQ(read_errors.load(), 0) + << "Refcount must prevent allocator from re-issuing in-use extents"; + + // Sanity: some reads succeeded and some keys were evicted. + EXPECT_GT(read_success.load(), 0); + { + std::lock_guard lk(evict_mtx); + EXPECT_FALSE(all_evicted.empty()) + << "Eviction must have occurred during concurrent stress"; + } + + // Post-condition: at least one evicted key is no longer in the map. + { + std::lock_guard lk(evict_mtx); + bool any_gone = false; + for (const auto& ek : all_evicted) { + if (!storage_backend.IsExist(ek).value_or(true)) { + any_gone = true; + break; + } + } + EXPECT_TRUE(any_gone) << "Evicted keys must be removed from shard.map"; + } +} + +//----------------------------------------------------------------------------- +// OffsetAllocatorStorageBackend Persistence / Recovery Tests +//----------------------------------------------------------------------------- + +namespace { + +// Loads a single key into a heap buffer; nullopt when the key is missing +// or the read fails. +std::optional LoadOne(OffsetAllocatorStorageBackend& backend, + const std::string& key, size_t value_size) { + auto buf = std::make_unique(value_size); + std::unordered_map load; + load.emplace(key, Slice{buf.get(), value_size}); + auto res = backend.BatchLoad(load); + if (!res) return std::nullopt; + return std::string(buf.get(), value_size); +} + +auto NoopCompleteHandler() { + return [](const std::vector&, + std::vector&) { return ErrorCode::OK; }; +} + +auto AcceptAllEvictionHandler() { + return + [](const std::vector&) -> tl::expected { + return {}; + }; +} + +FileStorageConfig MakeOffsetPersistConfig(const std::string& data_path, + int64_t size_limit) { + FileStorageConfig config; + config.storage_filepath = data_path; + config.storage_backend_type = StorageBackendType::kOffsetAllocator; + config.total_size_limit = size_limit; + config.total_keys_limit = 100; + return config; +} + +} // namespace + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Persist_RecoveryRoundtrip) { + // kStrict: every batch checkpoints. After a graceful restart the + // recovered backend must serve all keys with their exact values. + auto config = MakeOffsetPersistConfig(data_path, 64 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.persist_mode = OffsetPersistMode::kStrict; + + const std::string value_a(1000, 'a'); + const std::string value_b(1500, 'b'); + const std::string value_c(777, 'c'); + { + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + for (const auto& kv : + {std::pair{"key_a", &value_a}, std::pair{"key_b", &value_b}, + std::pair{"key_c", &value_c}}) { + auto batch = MakeSingleKeyBatch(kv.first, *kv.second, buffers); + auto res = + backend.BatchOffload(batch, NoopCompleteHandler(), nullptr); + ASSERT_TRUE(res.has_value()) << "key=" << kv.first; + ASSERT_EQ(res.value(), 1); + } + } // graceful destruction (final checkpoint is a no-op for kStrict) + + OffsetAllocatorStorageBackend recovered(config, backend_cfg); + ASSERT_TRUE(recovered.Init()); + EXPECT_EQ(LoadOne(recovered, "key_a", value_a.size()), + std::make_optional(value_a)); + EXPECT_EQ(LoadOne(recovered, "key_b", value_b.size()), + std::make_optional(value_b)); + EXPECT_EQ(LoadOne(recovered, "key_c", value_c.size()), + std::make_optional(value_c)); + EXPECT_FALSE(recovered.IsExist("key_missing").value_or(true)); +} + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Persist_RelaxedDestructorCheckpoint) { + // kRelaxed with a long interval only checkpoints the first batch; + // the destructor's final checkpoint must cover the rest. + auto config = MakeOffsetPersistConfig(data_path, 64 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.persist_mode = OffsetPersistMode::kRelaxed; + backend_cfg.persist_interval_seconds = 3600; + + const std::string value_a(1000, 'a'); + const std::string value_b(1000, 'b'); + { + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + // First batch: checkpoint fires (last_persist_time_us_ == 0). + auto batch_a = MakeSingleKeyBatch("key_a", value_a, buffers); + ASSERT_TRUE( + backend.BatchOffload(batch_a, NoopCompleteHandler(), nullptr) + .has_value()); + // Second batch: inside the interval, no checkpoint until shutdown. + auto batch_b = MakeSingleKeyBatch("key_b", value_b, buffers); + ASSERT_TRUE( + backend.BatchOffload(batch_b, NoopCompleteHandler(), nullptr) + .has_value()); + } // destructor must checkpoint key_b + + OffsetAllocatorStorageBackend recovered(config, backend_cfg); + ASSERT_TRUE(recovered.Init()); + EXPECT_EQ(LoadOne(recovered, "key_a", value_a.size()), + std::make_optional(value_a)); + EXPECT_EQ(LoadOne(recovered, "key_b", value_b.size()), + std::make_optional(value_b)); +} + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Persist_EvictionTombstone) { + // An evicted key must NOT be resurrected by recovery: the checkpoint + // records a tombstone for it. + auto config = MakeOffsetPersistConfig(data_path, 20 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + backend_cfg.persist_mode = OffsetPersistMode::kStrict; + + std::vector evicted; + auto eviction_handler = [&evicted](const std::vector& keys) { + for (const auto& k : keys) evicted.push_back(k); + return tl::expected{}; + }; + + const std::string data(1000, 'x'); + { + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + for (const auto& key : {"key_a", "key_b", "key_c", "key_d"}) { + auto batch = MakeSingleKeyBatch(key, data, buffers); + auto res = backend.BatchOffload(batch, NoopCompleteHandler(), + eviction_handler); + ASSERT_TRUE(res.has_value()) << "key=" << key; + } + ASSERT_FALSE(evicted.empty()); + EXPECT_EQ(evicted[0], "key_a"); + EXPECT_FALSE(backend.IsExist("key_a").value_or(true)); + } + + OffsetAllocatorStorageBackend recovered(config, backend_cfg); + ASSERT_TRUE(recovered.Init()); + EXPECT_FALSE(recovered.IsExist("key_a").value_or(true)) + << "evicted key_a must not be resurrected by recovery"; + for (const auto& key : {"key_b", "key_c", "key_d"}) { + EXPECT_EQ(LoadOne(recovered, key, data.size()), + std::make_optional(data)) + << "key=" << key; + } +} + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Persist_DetectsCorruptRecord) { + // A torn/partial write after the last checkpoint must be detected via + // CRC and skipped, while intact records survive recovery. + auto config = MakeOffsetPersistConfig(data_path, 64 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.persist_mode = OffsetPersistMode::kStrict; + + using RecordHeader = OffsetAllocatorStorageBackend::RecordHeader; + + std::unordered_map metas; + auto capturing_handler = + [&metas](const std::vector& keys, + std::vector& metadatas) { + for (size_t i = 0; i < keys.size(); ++i) { + metas.emplace(keys[i], metadatas[i]); + } + return ErrorCode::OK; + }; + + const std::string value_a(1000, 'a'); + const std::string value_b(1000, 'b'); + { + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + auto batch_a = MakeSingleKeyBatch("key_a", value_a, buffers); + ASSERT_TRUE(backend.BatchOffload(batch_a, capturing_handler, nullptr) + .has_value()); + auto batch_b = MakeSingleKeyBatch("key_b", value_b, buffers); + ASSERT_TRUE(backend.BatchOffload(batch_b, capturing_handler, nullptr) + .has_value()); + } + ASSERT_EQ(metas.count("key_a"), 1u); + + // Flip one byte inside key_a's value region (header stays intact, so + // only the CRC can catch this). The value starts at its 4K-aligned + // offset within the record. + { + const auto& meta = metas.at("key_a"); + const off_t corrupt_pos = + static_cast(meta.offset) + + static_cast(RecordHeader::ValueOffsetInRecord( + static_cast(meta.key_size))) + + 10; + const std::string data_file = data_path + "/kv_cache.data"; + int fd = open(data_file.c_str(), O_RDWR); + ASSERT_GE(fd, 0); + char byte; + ASSERT_EQ(pread(fd, &byte, 1, corrupt_pos), 1); + byte ^= 0xFF; + ASSERT_EQ(pwrite(fd, &byte, 1, corrupt_pos), 1); + fsync(fd); + close(fd); + } + + OffsetAllocatorStorageBackend recovered(config, backend_cfg); + ASSERT_TRUE(recovered.Init()); + EXPECT_FALSE(recovered.IsExist("key_a").value_or(true)) + << "CRC-corrupted record must be skipped on recovery"; + EXPECT_EQ(LoadOne(recovered, "key_b", value_b.size()), + std::make_optional(value_b)); +} + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Persist_RejectsPostCheckpointWrite) { + // Crash window: a write (and the eviction that made room for it) + // happens AFTER the last checkpoint. On recovery from that + // checkpoint the new record carries a seq >= the checkpoint's + // insert_seq and must be dropped rather than resurrected as a + // phantom key. + auto config = MakeOffsetPersistConfig(data_path, 6 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.eviction_policy = OffsetEvictionPolicy::FIFO; + backend_cfg.persist_mode = OffsetPersistMode::kRelaxed; + backend_cfg.persist_interval_seconds = 3600; // no periodic checkpoint + + const std::string data(1000, 'x'); + { + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + + // Batch 1: key_k is written and checkpointed (M1). The eviction + // handler must be non-null for the key to enter the FIFO index + // (eviction is inert when the handler is null). + auto batch_k = MakeSingleKeyBatch("key_k", data, buffers); + auto res_k = backend.BatchOffload(batch_k, NoopCompleteHandler(), + AcceptAllEvictionHandler()); + ASSERT_TRUE(res_k.has_value()); + ASSERT_EQ(res_k.value(), 1); + + // Simulate a crash from here on: no more checkpoints. + backend.SetSkipFinalCheckpointForTest(); + + // Batch 2: the arena (6KB) only fits one 4K-aligned record + // (5096 B), so key_j evicts key_k and reuses its extent -- but + // M1 never learns about it. + auto batch_j = MakeSingleKeyBatch("key_j", data, buffers); + auto res = backend.BatchOffload(batch_j, NoopCompleteHandler(), + AcceptAllEvictionHandler()); + ASSERT_TRUE(res.has_value()); + ASSERT_EQ(res.value(), 1) << "key_j must reuse key_k's extent"; + ASSERT_TRUE(backend.IsExist("key_j").value_or(false)); + ASSERT_FALSE(backend.IsExist("key_k").value_or(true)); + } + + OffsetAllocatorStorageBackend recovered(config, backend_cfg); + ASSERT_TRUE(recovered.Init()); + EXPECT_FALSE(recovered.IsExist("key_j").value_or(true)) + << "post-checkpoint write must be dropped (seq >= insert_seq)"; + EXPECT_FALSE(recovered.IsExist("key_k").value_or(true)) + << "key_k's extent was reused; nothing valid remains at M1"; +} + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Persist_OversizedKeyRejected) { + // Keys above kMaxKeyLen (1MB) are rejected at write time so they + // cannot silently vanish on recovery (which enforces the same cap). + auto config = MakeOffsetPersistConfig(data_path, 4 * 1024 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.persist_mode = OffsetPersistMode::kStrict; + + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + + const std::string huge_key(1024 * 1024 + 1, 'k'); + const std::string small_value(32, 'v'); + auto batch = MakeSingleKeyBatch(huge_key, small_value, buffers); + auto res = backend.BatchOffload(batch, NoopCompleteHandler(), nullptr); + ASSERT_TRUE(res.has_value()); + EXPECT_EQ(res.value(), 0) << "oversized key must be skipped"; + EXPECT_FALSE(backend.IsExist(huge_key).value_or(true)); + + // Boundary: exactly 1MB is accepted. + const std::string max_key(1024 * 1024, 'm'); + auto batch_max = MakeSingleKeyBatch(max_key, small_value, buffers); + auto res_max = + backend.BatchOffload(batch_max, NoopCompleteHandler(), nullptr); + ASSERT_TRUE(res_max.has_value()); + EXPECT_EQ(res_max.value(), 1); + EXPECT_TRUE(backend.IsExist(max_key).value_or(false)); +} + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Persist_RecordCrcDisabledRoundtrip) { + // With per-record CRC disabled, recovery must still restore every + // record: validation then rests on the checkpoint's seq guard alone. + auto config = MakeOffsetPersistConfig(data_path, 64 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.persist_mode = OffsetPersistMode::kStrict; + backend_cfg.enable_record_crc = false; + + const std::string value_a(1000, 'a'); + const std::string value_b(1500, 'b'); + { + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + for (const auto& kv : + {std::pair{"key_a", &value_a}, std::pair{"key_b", &value_b}}) { + auto batch = MakeSingleKeyBatch(kv.first, *kv.second, buffers); + auto res = + backend.BatchOffload(batch, NoopCompleteHandler(), nullptr); + ASSERT_TRUE(res.has_value()) << "key=" << kv.first; + ASSERT_EQ(res.value(), 1); + } + } + + OffsetAllocatorStorageBackend recovered(config, backend_cfg); + ASSERT_TRUE(recovered.Init()); + EXPECT_EQ(LoadOne(recovered, "key_a", value_a.size()), + std::make_optional(value_a)); + EXPECT_EQ(LoadOne(recovered, "key_b", value_b.size()), + std::make_optional(value_b)); +} + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Persist_RecordCrcDisabledTornValue) { + // Documents the weaker guarantee of CRC-disabled mode: in-place value + // corruption below the checkpoint's seq watermark is NOT detected on + // recovery and the record is served as-is. Only disable CRC when torn + // writes are otherwise impossible (strict ordering, PLP storage) or + // the value never touches the CPU (DMA writers). + auto config = MakeOffsetPersistConfig(data_path, 64 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.persist_mode = OffsetPersistMode::kStrict; + backend_cfg.enable_record_crc = false; + + using RecordHeader = OffsetAllocatorStorageBackend::RecordHeader; + + std::unordered_map metas; + auto capturing_handler = + [&metas](const std::vector& keys, + std::vector& metadatas) { + for (size_t i = 0; i < keys.size(); ++i) { + metas.emplace(keys[i], metadatas[i]); + } + return ErrorCode::OK; + }; + + const std::string value_a(1000, 'a'); + { + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + auto batch_a = MakeSingleKeyBatch("key_a", value_a, buffers); + ASSERT_TRUE(backend.BatchOffload(batch_a, capturing_handler, nullptr) + .has_value()); + } + ASSERT_EQ(metas.count("key_a"), 1u); + + // Flip one byte inside key_a's value region. + { + const auto& meta = metas.at("key_a"); + const off_t corrupt_pos = + static_cast(meta.offset) + + static_cast(RecordHeader::ValueOffsetInRecord( + static_cast(meta.key_size))) + + 10; + const std::string data_file = data_path + "/kv_cache.data"; + int fd = open(data_file.c_str(), O_RDWR); + ASSERT_GE(fd, 0); + char byte; + ASSERT_EQ(pread(fd, &byte, 1, corrupt_pos), 1); + byte ^= 0xFF; + ASSERT_EQ(pwrite(fd, &byte, 1, corrupt_pos), 1); + fsync(fd); + close(fd); + } + + OffsetAllocatorStorageBackend recovered(config, backend_cfg); + ASSERT_TRUE(recovered.Init()); + // No CRC to fail: the record survives recovery ... + EXPECT_TRUE(recovered.IsExist("key_a").value_or(false)); + // ... and the corrupted byte is served verbatim. + std::string expected = value_a; + expected[10] ^= 0xFF; + EXPECT_EQ(LoadOne(recovered, "key_a", value_a.size()), + std::make_optional(expected)); +} + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_RecordLayout_ValueAlignedTo4K) { + // The value region must start at a kValueAlignment boundary within + // every record (cuFile/DMA requirement), with zero-filled padding + // that is a pure function of key_len. + using RecordHeader = OffsetAllocatorStorageBackend::RecordHeader; + + // Static layout properties, including the padding boundary cases. + static_assert(RecordHeader::SIZE == 24, "v3 header is 24 bytes"); + static_assert(RecordHeader::ValuePadding(4072) == 0, + "24 + 4072 == 4096 needs no padding"); + static_assert(RecordHeader::ValuePadding(4073) == 4095, ""); + static_assert(RecordHeader::ValueOffsetInRecord(0) == 4096, ""); + static_assert(RecordHeader::RecordSize(5, 1000) == 5096, ""); + EXPECT_EQ( + RecordHeader::ValueOffsetInRecord(5) % RecordHeader::kValueAlignment, + 0u); + EXPECT_EQ( + RecordHeader::ValueOffsetInRecord(4073) % RecordHeader::kValueAlignment, + 0u); + + auto config = MakeOffsetPersistConfig(data_path, 256 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.persist_mode = OffsetPersistMode::kStrict; + + std::unordered_map metas; + auto capturing_handler = + [&metas](const std::vector& keys, + std::vector& metadatas) { + for (size_t i = 0; i < keys.size(); ++i) { + metas.emplace(keys[i], metadatas[i]); + } + return ErrorCode::OK; + }; + + const std::string value(1000, 'v'); + const std::string short_key = "key_a"; + const std::string long_key(4073, 'k'); // padding boundary: 4095 B pad + { + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + for (const auto* key : {&short_key, &long_key}) { + auto batch = MakeSingleKeyBatch(*key, value, buffers); + ASSERT_TRUE(backend.BatchOffload(batch, capturing_handler, nullptr) + .has_value()); + } + } + ASSERT_EQ(metas.size(), 2u); + + // On-disk contract: padding is zero-filled and value bytes start + // exactly at the aligned offset. + const std::string data_file = data_path + "/kv_cache.data"; + int fd = open(data_file.c_str(), O_RDONLY); + ASSERT_GE(fd, 0); + for (const auto& [key, meta] : metas) { + const uint32_t key_len = static_cast(meta.key_size); + const uint32_t pad = RecordHeader::ValuePadding(key_len); + std::vector pad_buf(pad); + ASSERT_EQ(pread(fd, pad_buf.data(), pad, + static_cast(meta.offset) + + static_cast(RecordHeader::SIZE) + key_len), + static_cast(pad)) + << "key=" << key; + for (char c : pad_buf) { + ASSERT_EQ(c, '\0') << "padding must be zero-filled, key=" << key; + } + char first_value_byte; + ASSERT_EQ(pread(fd, &first_value_byte, 1, + static_cast(meta.offset) + + static_cast( + RecordHeader::ValueOffsetInRecord(key_len))), + 1) + << "key=" << key; + EXPECT_EQ(first_value_byte, 'v') << "key=" << key; + } + close(fd); + + // Recovery roundtrip still works with aligned records. + OffsetAllocatorStorageBackend recovered(config, backend_cfg); + ASSERT_TRUE(recovered.Init()); + EXPECT_EQ(LoadOne(recovered, "key_a", value.size()), + std::make_optional(value)); + EXPECT_EQ(LoadOne(recovered, long_key, value.size()), + std::make_optional(value)); +} + +TEST_F(StorageBackendTest, + OffsetAllocatorStorageBackend_Persist_RejectsUnknownFlags) { + // A record whose flags carry bits this build does not know must be + // dropped on recovery (it belongs to a newer on-disk format). + auto config = MakeOffsetPersistConfig(data_path, 64 * 1024); + OffsetAllocatorBackendConfig backend_cfg; + backend_cfg.persist_mode = OffsetPersistMode::kStrict; + + std::unordered_map metas; + auto capturing_handler = + [&metas](const std::vector& keys, + std::vector& metadatas) { + for (size_t i = 0; i < keys.size(); ++i) { + metas.emplace(keys[i], metadatas[i]); + } + return ErrorCode::OK; + }; + + const std::string value_a(1000, 'a'); + const std::string value_b(1000, 'b'); + { + OffsetAllocatorStorageBackend backend(config, backend_cfg); + ASSERT_TRUE(backend.Init()); + std::vector> buffers; + auto batch_a = MakeSingleKeyBatch("key_a", value_a, buffers); + ASSERT_TRUE(backend.BatchOffload(batch_a, capturing_handler, nullptr) + .has_value()); + auto batch_b = MakeSingleKeyBatch("key_b", value_b, buffers); + ASSERT_TRUE(backend.BatchOffload(batch_b, capturing_handler, nullptr) + .has_value()); + } + ASSERT_EQ(metas.count("key_a"), 1u); + + // Set flags = 0x02 (unknown bit, kFlagHasCrc clear) in key_a's header: + // the flags field sits at byte offset 16 (u32 key_len, u32 value_len, + // u64 seq), little-endian. + { + const auto& meta = metas.at("key_a"); + const std::string data_file = data_path + "/kv_cache.data"; + int fd = open(data_file.c_str(), O_RDWR); + ASSERT_GE(fd, 0); + const char flag_byte = 0x02; + ASSERT_EQ( + pwrite(fd, &flag_byte, 1, static_cast(meta.offset) + 16), 1); + fsync(fd); + close(fd); + } + + OffsetAllocatorStorageBackend recovered(config, backend_cfg); + ASSERT_TRUE(recovered.Init()); + EXPECT_FALSE(recovered.IsExist("key_a").value_or(true)) + << "record with unknown flags must be dropped on recovery"; + EXPECT_EQ(LoadOne(recovered, "key_b", value_b.size()), + std::make_optional(value_b)); +} } // namespace mooncake::test diff --git a/mooncake-store/tests/stress_cluster_benchmark.py b/mooncake-store/tests/stress_cluster_benchmark.py index 427b8f7ed9..738fec0920 100644 --- a/mooncake-store/tests/stress_cluster_benchmark.py +++ b/mooncake-store/tests/stress_cluster_benchmark.py @@ -2,7 +2,6 @@ import time import statistics import logging -import ctypes import numpy as np from dataclasses import dataclass from typing import List, Dict, Any diff --git a/mooncake-store/tests/tenant_id_test.cpp b/mooncake-store/tests/tenant_id_test.cpp new file mode 100644 index 0000000000..540a8da605 --- /dev/null +++ b/mooncake-store/tests/tenant_id_test.cpp @@ -0,0 +1,65 @@ +#include "tenant_id.h" + +#include +#include + +#include + +namespace mooncake { +namespace { + +TEST(TenantIdTest, NormalizesEmptyTenantToDefault) { + const TenantId empty(""); + EXPECT_EQ(empty, TenantId::Default()); + EXPECT_EQ(empty.value(), "default"); + EXPECT_TRUE(empty.IsDefault()); + + const TenantId named("tenant-a"); + EXPECT_EQ(named.value(), "tenant-a"); + EXPECT_FALSE(named.IsDefault()); +} + +TEST(TenantIdTest, ValidatesCanonicalTenantName) { + EXPECT_TRUE(TenantId().IsValid()); + EXPECT_TRUE(TenantId("tenant-a").IsValid()); + EXPECT_FALSE(TenantId("_reserved").IsValid()); + EXPECT_FALSE(TenantId(std::string("bad\nname")).IsValid()); + EXPECT_FALSE(TenantId(std::string("bad\x7f", 4)).IsValid()); +} + +TEST(TenantIdTest, SupportsOrderedAndHashedKeys) { + EXPECT_LT(TenantId("tenant-a"), TenantId("tenant-b")); + + std::unordered_map tenants; + tenants.emplace(TenantId("tenant-a"), 1); + EXPECT_EQ(tenants.at(TenantId("tenant-a")), 1); +} + +TEST(TenantIdTest, ScopedKeyPreservesExistingEncoding) { + const TenantId tenant("tenant:with:colon"); + const std::string scoped = tenant.MakeScopedKey("path/key:with:colon"); + const std::string expected = + std::string("tenant:with:colon") + '\0' + "path/key:with:colon"; + EXPECT_EQ(scoped, expected); + + auto [parsed_tenant, parsed_key] = TenantId::ParseScopedKey(scoped); + EXPECT_EQ(parsed_tenant, tenant); + EXPECT_EQ(parsed_key, "path/key:with:colon"); +} + +TEST(TenantIdTest, ScopedKeyRoundTripPreservesEmbeddedNullInLocalKey) { + const std::string local_key("part1\0part2", 11); + const std::string scoped = TenantId("tenant-a").MakeScopedKey(local_key); + auto [tenant, parsed_key] = TenantId::ParseScopedKey(scoped); + EXPECT_EQ(tenant, TenantId("tenant-a")); + EXPECT_EQ(parsed_key, local_key); +} + +TEST(TenantIdTest, ParsesLegacyUnscopedKeyAsDefaultTenant) { + auto [tenant, key] = TenantId::ParseScopedKey("legacy-key"); + EXPECT_EQ(tenant, TenantId::Default()); + EXPECT_EQ(key, "legacy-key"); +} + +} // namespace +} // namespace mooncake diff --git a/mooncake-store/tests/tenant_quota_test.cpp b/mooncake-store/tests/tenant_quota_test.cpp index c61e67ccb2..e88b9d90e6 100644 --- a/mooncake-store/tests/tenant_quota_test.cpp +++ b/mooncake-store/tests/tenant_quota_test.cpp @@ -1,23 +1,52 @@ -#include "tenant_quota.h" +#include "tenant_quota_sharded.h" +#include "tenant_quota_policy_store.h" #include "types.h" +#ifdef STORE_USE_ETCD +#include "etcd_helper.h" +#endif + +#include +#include +#include +#include +#include #include +#include +#include #include +#include +#include #include #include +#include namespace mooncake { namespace { -TenantQuotaSnapshot Snapshot(const TenantQuotaTable& table, - const std::string& tenant_id) { - auto snapshot = table.GetTenantSnapshot(tenant_id); +#ifdef STORE_USE_ETCD +constexpr const char* kTenantQuotaEtcdEndpoints = "127.0.0.1:2379"; +constexpr std::string_view kTenantQuotaEtcdProbeKey = "tenant_quota_probe"; + +std::string GetTenantQuotaEtcdEndpoints() { + const char* endpoints = std::getenv("MOONCAKE_TENANT_QUOTA_ETCD_ENDPOINTS"); + if (endpoints != nullptr && endpoints[0] != '\0') { + return endpoints; + } + return kTenantQuotaEtcdEndpoints; +} +#endif + +template +TenantQuotaSnapshot Snapshot(const Table& table, const std::string& tenant_id) { + auto snapshot = table.GetTenantSnapshot(TenantId(tenant_id)); EXPECT_TRUE(snapshot.has_value()); return *snapshot; } -uint64_t SumEffectiveQuotas(const TenantQuotaTable& table) { +template +uint64_t SumEffectiveQuotas(const Table& table) { uint64_t sum = 0; for (const auto& snapshot : table.ListTenantSnapshots()) { sum += snapshot.effective_quota_bytes; @@ -25,33 +54,81 @@ uint64_t SumEffectiveQuotas(const TenantQuotaTable& table) { return sum; } -void MakeInheritedTenantActive(TenantQuotaTable* table, - const std::string& tenant_id, - uint64_t capacity) { - ASSERT_TRUE(table->UpsertTenantPolicy(tenant_id, capacity).has_value()); - table->RecomputeEffectiveQuotas(capacity); - ASSERT_TRUE(table->Reserve(tenant_id, 1).has_value()); - ASSERT_TRUE(table->Commit(tenant_id, 1).has_value()); - table->EraseTenantPolicy(tenant_id); +std::filesystem::path MakeTempPolicyPath(const std::string& suffix) { + return std::filesystem::temp_directory_path() / + ("mooncake_tenant_quota_policy_store_test_" + + std::to_string(::getpid()) + "_" + suffix + ".yaml"); +} + +#ifdef STORE_USE_ETCD +std::string PrefixEnd(std::string prefix) { + for (int i = static_cast(prefix.size()) - 1; i >= 0; --i) { + unsigned char c = static_cast(prefix[i]); + if (c < 0xFF) { + prefix[i] = static_cast(c + 1); + prefix.resize(i + 1); + return prefix; + } + } + return std::string(1, '\0'); +} + +std::optional GetTenantQuotaEtcdSkipReason() { + const std::string endpoints = GetTenantQuotaEtcdEndpoints(); + ErrorCode error = EtcdHelper::ConnectToEtcdStoreClient(endpoints); + if (error != ErrorCode::OK) { + return "Etcd server not reachable at " + endpoints + ": " + + toString(error); + } + std::string value; + EtcdRevisionId revision_id = 0; + error = + EtcdHelper::Get(kTenantQuotaEtcdProbeKey.data(), + kTenantQuotaEtcdProbeKey.size(), value, revision_id); + if (error == ErrorCode::ETCD_OPERATION_ERROR) { + return "Etcd server not reachable at " + endpoints + ": " + + toString(error); + } + return std::nullopt; +} + +void CleanupTenantQuotaEtcdCluster(const std::string& cluster_id) { + std::string prefix = "mooncake-store/" + cluster_id + "/"; + std::string end = PrefixEnd(prefix); + (void)EtcdHelper::DeleteRange(prefix.c_str(), prefix.size(), end.c_str(), + end.size()); +} +#endif + +void MakeOrphanTenant(TenantQuotaTable* table, const std::string& tenant_id, + uint64_t bytes) { + const TenantId canonical_tenant(tenant_id); + ASSERT_TRUE(table->UpsertTenantPolicy(canonical_tenant, bytes).has_value()); + table->RecomputeEffectiveQuotas(bytes); + ASSERT_TRUE(table->Reserve(canonical_tenant, bytes).has_value()); + ASSERT_TRUE(table->Commit(canonical_tenant, bytes).has_value()); + table->ApplyTenantPolicies({}); } -TEST(TenantQuotaTableTest, NormalizesEmptyTenantIdToDefault) { +TEST(TenantQuotaTableTest, NormalizesEmptyExplicitTenantIdToDefault) { TenantQuotaTable table; - EXPECT_EQ(NormalizeTenantId(""), "default"); - MakeInheritedTenantActive(&table, "", 1024); - table.RecomputeEffectiveQuotas(1024); + ASSERT_TRUE(table.UpsertTenantPolicy(TenantId(""), 1024).has_value()); + table.RecomputeEffectiveQuotas(4096); auto snapshot = Snapshot(table, ""); - EXPECT_EQ(snapshot.tenant_id, "default"); + EXPECT_EQ(snapshot.tenant_id, TenantId::Default()); + EXPECT_TRUE(snapshot.has_explicit_policy); + EXPECT_EQ(snapshot.requested_quota_bytes, 1024); EXPECT_EQ(snapshot.effective_quota_bytes, 1024); } TEST(TenantQuotaTableTest, RejectsZeroExplicitQuotaWithoutChangingState) { TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - auto result = table.UpsertTenantPolicy("tenant-a", 0); + const TenantId tenant_id("tenant-a"); + ASSERT_TRUE(table.UpsertTenantPolicy(tenant_id, 100).has_value()); + auto result = table.UpsertTenantPolicy(tenant_id, 0); ASSERT_FALSE(result.has_value()); EXPECT_EQ(result.error(), TenantQuotaError::kInvalidArgument); @@ -60,383 +137,519 @@ TEST(TenantQuotaTableTest, RejectsZeroExplicitQuotaWithoutChangingState) { EXPECT_EQ(snapshot.requested_quota_bytes, 100); } -TEST(TenantQuotaTableTest, ExplicitPolicyOverridesDefaultAndEraseFallsBack) { +TEST(TenantQuotaTableTest, ApplyPoliciesCreatesOrphanState) { TenantQuotaTable table; - table.SetDefaultRequestedQuota(50); - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(1000); - - EXPECT_TRUE(Snapshot(table, "tenant-a").has_explicit_policy); - EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 100); - - ASSERT_TRUE(table.Reserve("tenant-a", 40).has_value()); - ASSERT_TRUE(table.Commit("tenant-a", 40).has_value()); - table.EraseTenantPolicy("tenant-a"); + MakeOrphanTenant(&table, "tenant-a", 40); table.RecomputeEffectiveQuotas(1000); auto snapshot = Snapshot(table, "tenant-a"); EXPECT_FALSE(snapshot.has_explicit_policy); - EXPECT_EQ(snapshot.requested_quota_bytes, 50); - EXPECT_EQ(snapshot.effective_quota_bytes, 1000); + EXPECT_EQ(snapshot.requested_quota_bytes, 0); + EXPECT_EQ(snapshot.effective_quota_bytes, 0); EXPECT_EQ(snapshot.used_bytes, 40); EXPECT_EQ(snapshot.committed_count, 1); + EXPECT_TRUE(snapshot.over_quota); +} + +TEST(TenantQuotaTableTest, ApplyPoliciesReplacesCanonicalPolicySet) { + TenantQuotaTable table; + const TenantId tenant_a("tenant-a"); + const TenantId tenant_b("tenant-b"); + const TenantId tenant_c("tenant-c"); + table.ApplyTenantPolicies({{tenant_a, 100}, {tenant_b, 200}}); + table.IncrementMetadataObjectCount(tenant_a); + + table.ApplyTenantPolicies({{tenant_b, 300}, {tenant_c, 400}}); + + EXPECT_FALSE(table.IsTenantRegistered(tenant_a)); + EXPECT_TRUE(table.IsTenantRegistered(tenant_b)); + EXPECT_TRUE(table.IsTenantRegistered(tenant_c)); + EXPECT_TRUE(Snapshot(table, tenant_a.value()).over_quota); + EXPECT_EQ(table.GetTenantPolicies(), + (TenantQuotaPolicyMap{{tenant_b, 300}, {tenant_c, 400}})); } -TEST(TenantQuotaTableTest, EraseMissingPolicyDoesNotCreateLazyState) { +TEST(TenantQuotaTableTest, DisableMissingPolicyDoesNotCreateLazyState) { TenantQuotaTable table; - table.EraseTenantPolicy("missing"); + auto result = table.DisableTenantPolicyIfEmpty(TenantId("missing")); - EXPECT_FALSE(table.GetTenantSnapshot("missing").has_value()); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), TenantQuotaError::kTenantNotFound); + EXPECT_FALSE(table.GetTenantSnapshot(TenantId("missing")).has_value()); EXPECT_TRUE(table.ListTenantSnapshots().empty()); } TEST(TenantQuotaTableTest, PolicyMutationDoesNotRecomputeEffectiveQuota) { TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); + const TenantId tenant_id("tenant-a"); + ASSERT_TRUE(table.UpsertTenantPolicy(tenant_id, 100).has_value()); table.RecomputeEffectiveQuotas(1000); EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 100); - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 200).has_value()); + ASSERT_TRUE(table.UpsertTenantPolicy(tenant_id, 200).has_value()); EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 100); table.RecomputeEffectiveQuotas(1000); EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 200); } -TEST(TenantQuotaTableTest, - DefaultPolicyMutationDoesNotRecomputeEffectiveQuota) { - TenantQuotaTable table; - MakeInheritedTenantActive(&table, "default", 1000); - EXPECT_EQ(Snapshot(table, "default").effective_quota_bytes, 1000); - - table.SetDefaultRequestedQuota(100); - EXPECT_EQ(Snapshot(table, "default").requested_quota_bytes, 100); - EXPECT_EQ(Snapshot(table, "default").effective_quota_bytes, 1000); - - table.RecomputeEffectiveQuotas(500); - EXPECT_EQ(Snapshot(table, "default").effective_quota_bytes, 500); -} - -TEST(TenantQuotaTableTest, ListSnapshotsSortedAndSkipsLazyEmptyTenants) { +TEST(TenantQuotaTableTest, ListSnapshotsSortedAndCleansLazyEmptyTenants) { TenantQuotaTable table; - ASSERT_TRUE(table.Reserve("z-empty", 0).has_value()); - ASSERT_TRUE(table.UpsertTenantPolicy("b", 10).has_value()); - ASSERT_TRUE(table.UpsertTenantPolicy("a", 10).has_value()); + ASSERT_TRUE(table.UpsertTenantPolicy(TenantId("z-empty"), 10).has_value()); + ASSERT_TRUE( + table.DisableTenantPolicyIfEmpty(TenantId("z-empty")).has_value()); + ASSERT_TRUE(table.UpsertTenantPolicy(TenantId("b"), 10).has_value()); + ASSERT_TRUE(table.UpsertTenantPolicy(TenantId("a"), 10).has_value()); table.RecomputeEffectiveQuotas(100); auto snapshots = table.ListTenantSnapshots(); ASSERT_EQ(snapshots.size(), 2); - EXPECT_EQ(snapshots[0].tenant_id, "a"); - EXPECT_EQ(snapshots[1].tenant_id, "b"); -} - -TEST(TenantQuotaTableTest, SingleDefaultTenantReceivesFullCapacity) { - TenantQuotaTable table; - MakeInheritedTenantActive(&table, "default", 1234); - - table.RecomputeEffectiveQuotas(1234); - - EXPECT_EQ(Snapshot(table, "default").effective_quota_bytes, 1234); + EXPECT_EQ(snapshots[0].tenant_id, TenantId("a")); + EXPECT_EQ(snapshots[1].tenant_id, TenantId("b")); } -TEST(TenantQuotaTableTest, - ExplicitTenantsGetRequestedAndDefaultSharesRemainder) { +TEST(TenantQuotaTableTest, ExplicitTenantsReceiveRequestedWhenCapacityFits) { TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - MakeInheritedTenantActive(&table, "default", 250); + ASSERT_TRUE( + table.UpsertTenantPolicy(TenantId("tenant-a"), 100).has_value()); + ASSERT_TRUE( + table.UpsertTenantPolicy(TenantId("tenant-b"), 200).has_value()); - table.RecomputeEffectiveQuotas(250); + table.RecomputeEffectiveQuotas(1000); EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 100); - EXPECT_EQ(Snapshot(table, "default").effective_quota_bytes, 150); -} - -TEST(TenantQuotaTableTest, DefaultTenantsSplitRemainderWithTenantIdTieBreak) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("explicit", 100).has_value()); - MakeInheritedTenantActive(&table, "b", 203); - MakeInheritedTenantActive(&table, "a", 203); - - table.RecomputeEffectiveQuotas(203); - - EXPECT_EQ(Snapshot(table, "explicit").effective_quota_bytes, 100); - EXPECT_EQ(Snapshot(table, "a").effective_quota_bytes, 52); - EXPECT_EQ(Snapshot(table, "b").effective_quota_bytes, 51); - EXPECT_LE(SumEffectiveQuotas(table), 203); + EXPECT_EQ(Snapshot(table, "tenant-b").effective_quota_bytes, 200); + EXPECT_EQ(SumEffectiveQuotas(table), 300); } TEST(TenantQuotaTableTest, OverCapacityScalesOnlyExplicitTenants) { TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("b", 200).has_value()); - ASSERT_TRUE(table.UpsertTenantPolicy("a", 100).has_value()); - MakeInheritedTenantActive(&table, "default", 150); + MakeOrphanTenant(&table, "orphan", 20); + ASSERT_TRUE(table.UpsertTenantPolicy(TenantId("b"), 200).has_value()); + ASSERT_TRUE(table.UpsertTenantPolicy(TenantId("a"), 100).has_value()); table.RecomputeEffectiveQuotas(150); EXPECT_EQ(Snapshot(table, "a").effective_quota_bytes, 50); EXPECT_EQ(Snapshot(table, "b").effective_quota_bytes, 100); - EXPECT_EQ(Snapshot(table, "default").effective_quota_bytes, 0); + EXPECT_EQ(Snapshot(table, "orphan").effective_quota_bytes, 0); + EXPECT_TRUE(Snapshot(table, "orphan").over_quota); } -TEST(TenantQuotaTableTest, LeavesRemainderUnallocatedWithoutDefaultTenants) { +TEST(TenantQuotaTableTest, LazyEmptyOrphansDoNotAppearInList) { TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); + ASSERT_TRUE(table.UpsertTenantPolicy(TenantId("team-a"), 30).has_value()); + ASSERT_TRUE(table.UpsertTenantPolicy(TenantId("ghost"), 10).has_value()); + ASSERT_TRUE( + table.DisableTenantPolicyIfEmpty(TenantId("ghost")).has_value()); - table.RecomputeEffectiveQuotas(1000); + table.RecomputeEffectiveQuotas(100); - EXPECT_EQ(Snapshot(table, "tenant-a").effective_quota_bytes, 100); - EXPECT_EQ(SumEffectiveQuotas(table), 100); + EXPECT_EQ(Snapshot(table, "team-a").effective_quota_bytes, 30); + EXPECT_FALSE(table.GetTenantSnapshot(TenantId("ghost")).has_value()); + + auto snapshots = table.ListTenantSnapshots(); + ASSERT_EQ(snapshots.size(), 1); + EXPECT_EQ(snapshots[0].tenant_id, TenantId("team-a")); } -TEST(TenantQuotaTableTest, DefaultUnlimitedTenantDoesNotSqueezeExplicitQuota) { +TEST(TenantQuotaTableTest, ReserveRequiresRegisteredTenantIncludingZeroBytes) { TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("small", 1).has_value()); - MakeInheritedTenantActive(&table, "default", 10); - table.RecomputeEffectiveQuotas(10); + auto regular = table.Reserve(TenantId("missing"), 1); + auto zero = table.Reserve(TenantId("missing"), 0); - EXPECT_EQ(Snapshot(table, "small").effective_quota_bytes, 1); - EXPECT_EQ(Snapshot(table, "default").effective_quota_bytes, 9); + ASSERT_FALSE(regular.has_value()); + ASSERT_FALSE(zero.has_value()); + EXPECT_EQ(regular.error(), TenantQuotaError::kTenantNotRegistered); + EXPECT_EQ(zero.error(), TenantQuotaError::kTenantNotRegistered); } -TEST(TenantQuotaTableTest, LazyEmptyTenantsDoNotDiluteActiveDefaultTenant) { +TEST(TenantQuotaTableTest, TracksAdditionalCommitAndMetadataCount) { TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("team-a", 30).has_value()); - MakeInheritedTenantActive(&table, "default", 100); - ASSERT_TRUE(table.UpsertTenantPolicy("ghost", 10).has_value()); - table.EraseTenantPolicy("ghost"); + const TenantId tenant_id("tenant-a"); + ASSERT_TRUE(table.UpsertTenantPolicy(tenant_id, 300).has_value()); + table.RecomputeEffectiveQuotas(300); - table.RecomputeEffectiveQuotas(100); + ASSERT_TRUE(table.Reserve(tenant_id, 200).has_value()); + ASSERT_TRUE(table.Commit(tenant_id, 100).has_value()); + ASSERT_TRUE(table.CommitAdditional(tenant_id, 100).has_value()); + table.IncrementMetadataObjectCount(tenant_id); - EXPECT_EQ(Snapshot(table, "team-a").effective_quota_bytes, 30); - EXPECT_EQ(Snapshot(table, "default").effective_quota_bytes, 70); - EXPECT_EQ(Snapshot(table, "ghost").effective_quota_bytes, 0); + auto snapshot = Snapshot(table, "tenant-a"); + EXPECT_EQ(snapshot.used_bytes, 200); + EXPECT_EQ(snapshot.reserved_bytes, 0); + EXPECT_EQ(snapshot.committed_count, 1); + EXPECT_EQ(snapshot.metadata_object_count, 1); +} - auto snapshots = table.ListTenantSnapshots(); - ASSERT_EQ(snapshots.size(), 2); - EXPECT_EQ(snapshots[0].tenant_id, "default"); - EXPECT_EQ(snapshots[1].tenant_id, "team-a"); +TEST(TenantQuotaTableTest, AccountingMismatchDoesNotMutateState) { + TenantQuotaTable table; + const TenantId tenant_id("tenant-a"); + ASSERT_TRUE(table.UpsertTenantPolicy(tenant_id, 100).has_value()); + table.RecomputeEffectiveQuotas(100); + ASSERT_TRUE(table.Reserve(tenant_id, 10).has_value()); + + auto commit = table.Commit(tenant_id, 11); + auto abort = table.Abort(tenant_id, 11); + auto before_commit = Snapshot(table, "tenant-a"); + ASSERT_FALSE(commit.has_value()); + ASSERT_FALSE(abort.has_value()); + EXPECT_EQ(commit.error(), TenantQuotaError::kAccountingMismatch); + EXPECT_EQ(abort.error(), TenantQuotaError::kAccountingMismatch); + EXPECT_EQ(before_commit.used_bytes, 0); + EXPECT_EQ(before_commit.reserved_bytes, 10); + + ASSERT_TRUE(table.Commit(tenant_id, 10).has_value()); + auto release = table.Release(tenant_id, 11); + auto partial = table.ReleasePartial(tenant_id, 11); + ASSERT_FALSE(release.has_value()); + ASSERT_FALSE(partial.has_value()); + EXPECT_EQ(release.error(), TenantQuotaError::kAccountingMismatch); + EXPECT_EQ(partial.error(), TenantQuotaError::kAccountingMismatch); + auto after_commit = Snapshot(table, "tenant-a"); + EXPECT_EQ(after_commit.used_bytes, 10); + EXPECT_EQ(after_commit.committed_count, 1); + + table.RebuildUsage({{tenant_id, + {.used_bytes = 10, + .committed_count = 0, + .metadata_object_count = 1}}}); + auto inconsistent_release = table.Release(tenant_id, 5); + ASSERT_FALSE(inconsistent_release.has_value()); + EXPECT_EQ(inconsistent_release.error(), + TenantQuotaError::kAccountingMismatch); + auto after_inconsistent_release = Snapshot(table, "tenant-a"); + EXPECT_EQ(after_inconsistent_release.used_bytes, 10); + EXPECT_EQ(after_inconsistent_release.committed_count, 0); } -TEST(TenantQuotaTableTest, LargestRemainderTieBreakUsesTenantId) { +TEST(TenantQuotaTableTest, DisablePolicyRejectsNonEmptyTenant) { TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("b", 1).has_value()); - ASSERT_TRUE(table.UpsertTenantPolicy("a", 1).has_value()); + const TenantId tenant_id("tenant-a"); + ASSERT_TRUE(table.UpsertTenantPolicy(tenant_id, 100).has_value()); + table.RecomputeEffectiveQuotas(100); + ASSERT_TRUE(table.Reserve(tenant_id, 1).has_value()); - table.RecomputeEffectiveQuotas(1); + auto result = table.DisableTenantPolicyIfEmpty(tenant_id); - EXPECT_EQ(Snapshot(table, "a").effective_quota_bytes, 1); - EXPECT_EQ(Snapshot(table, "b").effective_quota_bytes, 0); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), TenantQuotaError::kTenantNotEmpty); + EXPECT_TRUE(table.IsTenantRegistered(tenant_id)); } -TEST(TenantQuotaTableTest, CapacityShrinkAndGrowthRefreshOverQuota) { +TEST(TenantQuotaTableTest, RebuildUsageCreatesAndRemovesOrphans) { TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); + const TenantId explicit_tenant("tenant-a"); + const TenantId orphan("orphan"); + ASSERT_TRUE(table.UpsertTenantPolicy(explicit_tenant, 100).has_value()); + + TenantQuotaUsageMap usage{ + {explicit_tenant, + {.used_bytes = 40, .committed_count = 1, .metadata_object_count = 1}}, + {orphan, + {.used_bytes = 20, .committed_count = 1, .metadata_object_count = 1}}, + }; + table.RebuildUsage(usage); table.RecomputeEffectiveQuotas(100); - ASSERT_TRUE(table.Reserve("tenant-a", 80).has_value()); - ASSERT_TRUE(table.Commit("tenant-a", 80).has_value()); - table.RecomputeEffectiveQuotas(50); - EXPECT_TRUE(Snapshot(table, "tenant-a").over_quota); + EXPECT_TRUE(Snapshot(table, "tenant-a").has_explicit_policy); + EXPECT_FALSE(Snapshot(table, "orphan").has_explicit_policy); + EXPECT_TRUE(Snapshot(table, "orphan").over_quota); - table.RecomputeEffectiveQuotas(100); - EXPECT_FALSE(Snapshot(table, "tenant-a").over_quota); + table.RebuildUsage({}); + EXPECT_FALSE(table.GetTenantSnapshot(orphan).has_value()); + EXPECT_TRUE(table.GetTenantSnapshot(explicit_tenant).has_value()); } -TEST(TenantQuotaTableTest, LargeValuesDoNotOverflowDuringRecompute) { +TEST(TenantQuotaTableTest, OverflowChecksDoNotWrapAccounting) { TenantQuotaTable table; + const TenantId tenant_id("tenant-a"); const uint64_t max = std::numeric_limits::max(); - ASSERT_TRUE(table.UpsertTenantPolicy("a", max).has_value()); - ASSERT_TRUE(table.UpsertTenantPolicy("b", max).has_value()); - + ASSERT_TRUE(table.UpsertTenantPolicy(tenant_id, max).has_value()); + table.RebuildUsage({{tenant_id, + {.used_bytes = max - 5, + .committed_count = max, + .metadata_object_count = max}}}); table.RecomputeEffectiveQuotas(max); - EXPECT_EQ(Snapshot(table, "a").effective_quota_bytes, max / 2 + max % 2); - EXPECT_EQ(Snapshot(table, "b").effective_quota_bytes, max / 2); - EXPECT_LE(SumEffectiveQuotas(table), max); -} + EXPECT_EQ(table.ComputeDeficit(tenant_id, 10), 5); + auto overflow_reserve = table.Reserve(tenant_id, 10); + ASSERT_FALSE(overflow_reserve.has_value()); + EXPECT_EQ(overflow_reserve.error(), TenantQuotaError::kQuotaExceeded); -TEST(TenantQuotaTableTest, ReserveCommitUpdatesAccounting) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(100); + ASSERT_TRUE(table.Reserve(tenant_id, 5).has_value()); + ASSERT_TRUE(table.Commit(tenant_id, 5).has_value()); + table.IncrementMetadataObjectCount(tenant_id); - ASSERT_TRUE(table.Reserve("tenant-a", 40).has_value()); auto snapshot = Snapshot(table, "tenant-a"); - EXPECT_EQ(snapshot.reserved_bytes, 40); - EXPECT_EQ(snapshot.used_bytes, 0); - - ASSERT_TRUE(table.Commit("tenant-a", 40).has_value()); - snapshot = Snapshot(table, "tenant-a"); + EXPECT_EQ(snapshot.used_bytes, max); EXPECT_EQ(snapshot.reserved_bytes, 0); - EXPECT_EQ(snapshot.used_bytes, 40); - EXPECT_EQ(snapshot.committed_count, 1); + EXPECT_EQ(snapshot.committed_count, max); + EXPECT_EQ(snapshot.metadata_object_count, max); } -TEST(TenantQuotaTableTest, ReserveOverQuotaDoesNotModifyState) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(100); - ASSERT_TRUE(table.Reserve("tenant-a", 80).has_value()); +TEST(ShardedTenantQuotaTableTest, ConcurrentReserveNeverExceedsQuota) { + ShardedTenantQuotaTable<8> table; + const TenantId tenant_id("tenant-a"); + table.ApplyTenantPolicies({{tenant_id, 1000}}, 1000); + + std::atomic successes = 0; + std::vector workers; + for (int i = 0; i < 20; ++i) { + workers.emplace_back([&] { + if (table.Reserve(tenant_id, 100).has_value()) { + ++successes; + } + }); + } + for (auto& worker : workers) { + worker.join(); + } - auto before = Snapshot(table, "tenant-a"); - auto result = table.Reserve("tenant-a", 21); + EXPECT_EQ(successes.load(), 10); + EXPECT_EQ(Snapshot(table, "tenant-a").reserved_bytes, 1000); +} - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error(), TenantQuotaError::kQuotaExceeded); - auto after = Snapshot(table, "tenant-a"); - EXPECT_EQ(after.reserved_bytes, before.reserved_bytes); - EXPECT_EQ(after.used_bytes, before.used_bytes); - EXPECT_EQ(after.committed_count, before.committed_count); +TEST(ShardedTenantQuotaTableTest, DifferentShardsUpdateIndependently) { + using TestTable = ShardedTenantQuotaTable<2>; + const TenantId tenant_a("tenant-a"); + TenantId tenant_b("tenant-b"); + for (int suffix = 0; TenantIdHash{}(tenant_a) % TestTable::kNumShards == + TenantIdHash{}(tenant_b) % TestTable::kNumShards; + ++suffix) { + tenant_b = TenantId("tenant-b-" + std::to_string(suffix)); + } + + TestTable table; + table.ApplyTenantPolicies({{tenant_a, 1000}, {tenant_b, 1000}}, 2000); + + std::atomic failures = 0; + auto update = [&](const TenantId& tenant_id) { + for (int i = 0; i < 1000; ++i) { + if (!table.Reserve(tenant_id, 1) || !table.Abort(tenant_id, 1)) { + ++failures; + } + } + }; + std::thread first(update, std::cref(tenant_a)); + std::thread second(update, std::cref(tenant_b)); + first.join(); + second.join(); + + EXPECT_EQ(failures.load(), 0); + EXPECT_EQ(Snapshot(table, tenant_a.value()).reserved_bytes, 0); + EXPECT_EQ(Snapshot(table, tenant_b.value()).reserved_bytes, 0); } -TEST(TenantQuotaTableTest, ReserveUsesOverflowSafeHeadroomCheck) { - TenantQuotaTable table; - const uint64_t max = std::numeric_limits::max(); - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", max).has_value()); - table.RecomputeEffectiveQuotas(max); - ASSERT_TRUE(table.Reserve("tenant-a", max).has_value()); - ASSERT_TRUE(table.Commit("tenant-a", max).has_value()); +TEST(ShardedTenantQuotaTableTest, + DisabledPolicyRejectsRegularAndZeroByteReservations) { + ShardedTenantQuotaTable<8> table; + const TenantId tenant_id("tenant-a"); + table.ApplyTenantPolicies({{tenant_id, 100}}, 100); + ASSERT_TRUE(table.DisableTenantPolicyIfEmpty(tenant_id).has_value()); - auto before = Snapshot(table, "tenant-a"); - auto result = table.Reserve("tenant-a", 1); + auto regular = table.Reserve(tenant_id, 1); + auto zero = table.Reserve(tenant_id, 0); - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error(), TenantQuotaError::kQuotaExceeded); - auto after = Snapshot(table, "tenant-a"); - EXPECT_EQ(after.used_bytes, before.used_bytes); - EXPECT_EQ(after.reserved_bytes, before.reserved_bytes); - EXPECT_EQ(after.committed_count, before.committed_count); + ASSERT_FALSE(regular.has_value()); + ASSERT_FALSE(zero.has_value()); + EXPECT_EQ(regular.error(), TenantQuotaError::kTenantNotRegistered); + EXPECT_EQ(zero.error(), TenantQuotaError::kTenantNotRegistered); } -TEST(TenantQuotaTableTest, ReserveMissingTenantDoesNotCreateStateOnFailure) { - TenantQuotaTable table; +TEST(ShardedTenantQuotaTableTest, RecomputeCanRunWithAccounting) { + ShardedTenantQuotaTable<8> table; + const TenantId tenant_id("tenant-a"); + table.ApplyTenantPolicies({{tenant_id, 1000}}, 1000); + + std::atomic failures = 0; + std::thread accounting([&] { + for (int i = 0; i < 1000; ++i) { + if (!table.Reserve(tenant_id, 1) || !table.Abort(tenant_id, 1)) { + ++failures; + } + } + }); + std::thread recompute([&] { + for (int i = 0; i < 1000; ++i) { + table.RecomputeEffectiveQuotas(1000); + } + }); + + accounting.join(); + recompute.join(); + + EXPECT_EQ(failures.load(), 0); + auto snapshot = Snapshot(table, "tenant-a"); + EXPECT_EQ(snapshot.reserved_bytes, 0); + EXPECT_EQ(snapshot.effective_quota_bytes, 1000); +} - auto result = table.Reserve("missing", 1); +TEST(TenantQuotaPolicyStoreTest, ParsesValidYamlUnits) { + const char* yaml = R"yaml( +version: 1 + +tenants: + - name: tenant-a + quota: 200GB + - name: tenant-b + quota: 500MB + - name: experiment + quota: 12345 +)yaml"; + + auto snapshot = ParseTenantQuotaPolicyYaml(yaml); + + ASSERT_TRUE(snapshot.has_value()) << snapshot.error(); + EXPECT_EQ(snapshot->tenant_quotas.at("tenant-a"), + 200ULL * 1024 * 1024 * 1024); + EXPECT_EQ(snapshot->tenant_quotas.at("tenant-b"), 500ULL * 1024 * 1024); + EXPECT_EQ(snapshot->tenant_quotas.at("experiment"), 12345); +} - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error(), TenantQuotaError::kQuotaExceeded); - EXPECT_FALSE(table.GetTenantSnapshot("missing").has_value()); - EXPECT_TRUE(table.ListTenantSnapshots().empty()); +TEST(TenantQuotaPolicyStoreTest, RejectsInvalidYamlPolicies) { + std::vector invalid_policies = { + "version: 2\n\ntenants: []\n", + "version: 1\n\ntenants:\n - name: tenant-a\n quota: 1XB\n", + "version: 1\n\ntenants:\n - name: tenant-a\n quota: 0\n", + "version: 1\n\ntenants:\n - name: \"\"\n quota: 1KB\n", + "version: 1\n\ntenants:\n - name: _system\n quota: 1KB\n", + "version: 1\n\ntenants:\n - name: \"tenant\\0bad\"\n quota: " + "1KB\n", + "version: 1\n\ntenants:\n - name: \"tenant\\nline\"\n quota: " + "1KB\n", + "version: 1\n\ntenants:\n - name: \"tenant\\x7f\"\n quota: 1KB\n", + "version: 1\n\ntenants:\n - name: tenant-a\n quota: 1KB\n - name: " + "tenant-a\n quota: 2KB\n", + "version: 1\n\ntenants:\n - name: tenant-a\n quota: " + "18446744073709551616\n", + "version: 1\n\ntenants:\n - name: tenant-a\n quota: " + "18446744073709551615TB\n", + }; + + for (const auto& policy : invalid_policies) { + auto snapshot = ParseTenantQuotaPolicyYaml(policy); + EXPECT_FALSE(snapshot.has_value()) << policy; + } } -TEST(TenantQuotaTableTest, ReserveAbortReleasesReservation) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(100); +TEST(TenantQuotaPolicyStoreTest, RoundTripsYamlFile) { + const auto path = MakeTempPolicyPath("roundtrip"); + std::filesystem::remove(path); - ASSERT_TRUE(table.Reserve("tenant-a", 40).has_value()); - ASSERT_TRUE(table.Abort("tenant-a", 40).has_value()); + YamlTenantQuotaPolicyStore store(path.string()); + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = {{"tenant-a", 1024}, {"tenant-b", 2048}}; - EXPECT_EQ(Snapshot(table, "tenant-a").reserved_bytes, 0); + auto save = store.Save(snapshot); + ASSERT_TRUE(save.has_value()) << save.error(); + + auto loaded = store.Load(); + ASSERT_TRUE(loaded.has_value()) << loaded.error(); + EXPECT_EQ(loaded->tenant_quotas, snapshot.tenant_quotas); + + std::filesystem::remove(path); } -TEST(TenantQuotaTableTest, - CommitWithoutEnoughReservationDoesNotModifyStateAndReportsMismatch) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(100); - ASSERT_TRUE(table.Reserve("tenant-a", 5).has_value()); +TEST(TenantQuotaPolicyStoreTest, FileFactoryCreatesYamlStore) { + const auto path = MakeTempPolicyPath("factory-file"); + auto store = + CreateTenantQuotaPolicyStore("file", path.string(), "test_cluster"); + ASSERT_TRUE(store.has_value()) << store.error(); +} - auto before = Snapshot(table, "tenant-a"); - auto result = table.Commit("tenant-a", 10); +TEST(TenantQuotaPolicyStoreTest, FileFactoryRequiresUri) { + auto store = CreateTenantQuotaPolicyStore("file", "", "test_cluster"); + ASSERT_FALSE(store.has_value()); + EXPECT_NE(store.error().find("non-empty uri"), std::string::npos); +} - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error(), TenantQuotaError::kAccountingMismatch); - auto after = Snapshot(table, "tenant-a"); - EXPECT_EQ(after.reserved_bytes, before.reserved_bytes); - EXPECT_EQ(after.used_bytes, before.used_bytes); - EXPECT_EQ(after.committed_count, before.committed_count); +#ifndef STORE_USE_ETCD +TEST(TenantQuotaPolicyStoreTest, EtcdFactoryRequiresStoreUseEtcd) { + auto store = + CreateTenantQuotaPolicyStore("etcd", "127.0.0.1:2379", "test_cluster"); + ASSERT_FALSE(store.has_value()); + EXPECT_NE(store.error().find("STORE_USE_ETCD"), std::string::npos); } +#endif -TEST(TenantQuotaTableTest, - AbortWithoutEnoughReservationDoesNotModifyStateAndReportsMismatch) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(100); - ASSERT_TRUE(table.Reserve("tenant-a", 5).has_value()); +#ifdef STORE_USE_ETCD +TEST(TenantQuotaPolicyStoreTest, EtcdMissingKeyLoadsEmptySnapshot) { + if (auto skip_reason = GetTenantQuotaEtcdSkipReason(); + skip_reason.has_value()) { + GTEST_SKIP() << skip_reason.value(); + } - auto before = Snapshot(table, "tenant-a"); - auto result = table.Abort("tenant-a", 10); + const std::string cluster_id = + "tenant_quota_missing_" + std::to_string(::getpid()); + CleanupTenantQuotaEtcdCluster(cluster_id); - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error(), TenantQuotaError::kAccountingMismatch); - auto after = Snapshot(table, "tenant-a"); - EXPECT_EQ(after.reserved_bytes, before.reserved_bytes); - EXPECT_EQ(after.used_bytes, before.used_bytes); - EXPECT_EQ(after.committed_count, before.committed_count); + auto store = CreateTenantQuotaPolicyStore( + "etcd", GetTenantQuotaEtcdEndpoints(), cluster_id); + ASSERT_TRUE(store.has_value()) << store.error(); + + auto loaded = store.value()->Load(); + ASSERT_TRUE(loaded.has_value()) << loaded.error(); + EXPECT_TRUE(loaded->tenant_quotas.empty()); + + CleanupTenantQuotaEtcdCluster(cluster_id); } -TEST(TenantQuotaTableTest, - ReleaseWithoutEnoughUsedDoesNotModifyStateAndReportsMismatch) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(100); - ASSERT_TRUE(table.Reserve("tenant-a", 5).has_value()); - ASSERT_TRUE(table.Commit("tenant-a", 5).has_value()); +TEST(TenantQuotaPolicyStoreTest, EtcdRoundTripsSnapshot) { + if (auto skip_reason = GetTenantQuotaEtcdSkipReason(); + skip_reason.has_value()) { + GTEST_SKIP() << skip_reason.value(); + } - auto before = Snapshot(table, "tenant-a"); - auto result = table.Release("tenant-a", 10); + const std::string cluster_id = + "tenant_quota_roundtrip_" + std::to_string(::getpid()); + CleanupTenantQuotaEtcdCluster(cluster_id); - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error(), TenantQuotaError::kAccountingMismatch); - auto after = Snapshot(table, "tenant-a"); - EXPECT_EQ(after.reserved_bytes, before.reserved_bytes); - EXPECT_EQ(after.used_bytes, before.used_bytes); - EXPECT_EQ(after.committed_count, before.committed_count); -} + auto store = CreateTenantQuotaPolicyStore( + "etcd", GetTenantQuotaEtcdEndpoints(), cluster_id); + ASSERT_TRUE(store.has_value()) << store.error(); -TEST(TenantQuotaTableTest, ReleasePartialDoesNotChangeCommittedCount) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(100); - ASSERT_TRUE(table.Reserve("tenant-a", 50).has_value()); - ASSERT_TRUE(table.Commit("tenant-a", 50).has_value()); + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = {{"tenant-a", 1024}, {"tenant-b", 2048}}; + auto save = store.value()->Save(snapshot); + ASSERT_TRUE(save.has_value()) << save.error(); - ASSERT_TRUE(table.ReleasePartial("tenant-a", 20).has_value()); + auto loaded = store.value()->Load(); + ASSERT_TRUE(loaded.has_value()) << loaded.error(); + EXPECT_EQ(loaded->tenant_quotas, snapshot.tenant_quotas); - auto snapshot = Snapshot(table, "tenant-a"); - EXPECT_EQ(snapshot.used_bytes, 30); - EXPECT_EQ(snapshot.committed_count, 1); + CleanupTenantQuotaEtcdCluster(cluster_id); } +#endif -TEST(TenantQuotaTableTest, - ReleasePartialUnderflowReportsMismatchInReleaseBuild) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(100); +TEST(TenantQuotaPolicyStoreTest, RoundTripsYamlSpecialScalarNames) { + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = {{"foo#bar", 1}, + {"true", 2}, + {"[a, b]", 3}, + {"key: val", 4}, + {"quote\"slash\\", 5}}; -#ifdef NDEBUG - auto before = Snapshot(table, "tenant-a"); - auto result = table.ReleasePartial("tenant-a", 10); + auto parsed = + ParseTenantQuotaPolicyYaml(FormatTenantQuotaPolicyYaml(snapshot)); - ASSERT_FALSE(result.has_value()); - EXPECT_EQ(result.error(), TenantQuotaError::kAccountingMismatch); - auto after = Snapshot(table, "tenant-a"); - EXPECT_EQ(after.reserved_bytes, before.reserved_bytes); - EXPECT_EQ(after.used_bytes, before.used_bytes); - EXPECT_EQ(after.committed_count, before.committed_count); -#else - EXPECT_DEATH({ (void)table.ReleasePartial("tenant-a", 10); }, ""); -#endif + ASSERT_TRUE(parsed.has_value()) << parsed.error(); + EXPECT_EQ(parsed->tenant_quotas, snapshot.tenant_quotas); } -TEST(TenantQuotaTableTest, ZeroByteAccountingOperationsAreNoOpSuccess) { - TenantQuotaTable table; - ASSERT_TRUE(table.UpsertTenantPolicy("tenant-a", 100).has_value()); - table.RecomputeEffectiveQuotas(100); +TEST(TenantQuotaPolicyStoreTest, SaveFailureReturnsError) { + const auto path = MakeTempPolicyPath("missing-dir").parent_path() / + ("missing_dir_" + std::to_string(::getpid())) / + "policy.yaml"; + YamlTenantQuotaPolicyStore store(path.string()); - EXPECT_TRUE(table.Reserve("tenant-a", 0).has_value()); - EXPECT_TRUE(table.Commit("tenant-a", 0).has_value()); - EXPECT_TRUE(table.Abort("tenant-a", 0).has_value()); - EXPECT_TRUE(table.Release("tenant-a", 0).has_value()); - EXPECT_TRUE(table.ReleasePartial("tenant-a", 0).has_value()); + TenantQuotaPolicySnapshot snapshot; + snapshot.tenant_quotas = {{"tenant-a", 1024}}; - auto snapshot = Snapshot(table, "tenant-a"); - EXPECT_EQ(snapshot.used_bytes, 0); - EXPECT_EQ(snapshot.reserved_bytes, 0); - EXPECT_EQ(snapshot.committed_count, 0); + auto save = store.Save(snapshot); + EXPECT_FALSE(save.has_value()); } } // namespace diff --git a/mooncake-store/tests/transfer_task_test.cpp b/mooncake-store/tests/transfer_task_test.cpp index f841831c25..8753757ca0 100644 --- a/mooncake-store/tests/transfer_task_test.cpp +++ b/mooncake-store/tests/transfer_task_test.cpp @@ -7,10 +7,14 @@ #include #include #include +#include #include #include #include "types.h" +#ifdef USE_CUDA +#include +#endif namespace mooncake { @@ -144,6 +148,139 @@ TEST_F(TransferTaskTest, MemcpyWorkerPoolMultipleOperations) { } } +TEST_F(TransferTaskTest, TransferScatterHandlesFragmentedCpuBuffers) { + constexpr size_t kBufferSize = 512; + constexpr size_t kFragmentCount = 128; + std::vector source(kBufferSize), destination(kBufferSize, 0); + std::iota(source.begin(), source.end(), 0); + std::vector destination_offsets, source_offsets, + lengths(kFragmentCount, 1); + for (size_t i = 0; i < kFragmentCount; ++i) { + destination_offsets.push_back(i * 2); + source_offsets.push_back(i * 3); + } + + TransferEngine engine(false); + ASSERT_EQ(engine.init("P2PHANDSHAKE", "localhost:17931"), 0); + if (!engine.isUsingTent()) { + ASSERT_NE(engine.installTransport("tcp", nullptr), nullptr); + } + ASSERT_EQ(engine.registerLocalMemory(source.data(), source.size(), "cpu:0"), + 0); + ASSERT_EQ(engine.registerLocalMemory(destination.data(), destination.size(), + "cpu:0"), + 0); + + ASSERT_TRUE(engine + .transferScatter({{ + .opcode = TransferRequest::READ, + .remote_segment = engine.getLocalIpAndPort(), + .remote_base_offset = + reinterpret_cast(source.data()), + .remote_size = source.size(), + .local_buffer = destination.data(), + .local_capacity = destination.size(), + .local_offsets = destination_offsets, + .remote_offsets = source_offsets, + .lengths = lengths, + .on_fragment_complete = {}, + }}) + .ok()); + for (size_t i = 0; i < kFragmentCount; ++i) { + EXPECT_EQ(destination[destination_offsets[i]], + source[source_offsets[i]]); + } +} + +#ifdef USE_CUDA +TEST_F(TransferTaskTest, TransferScatterHandlesFragmentedGpuBuffers) { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA device is unavailable"; + } + + constexpr size_t kBufferSize = 512; + constexpr size_t kFragmentCount = 128; + std::vector source(kBufferSize), cpu_destination(kBufferSize, 0); + std::iota(source.begin(), source.end(), 0); + void *gpu_source = nullptr, *gpu_destination = nullptr; + ASSERT_EQ(cudaMalloc(&gpu_source, kBufferSize), cudaSuccess); + ASSERT_EQ(cudaMalloc(&gpu_destination, kBufferSize), cudaSuccess); + ASSERT_EQ(cudaMemcpy(gpu_source, source.data(), kBufferSize, + cudaMemcpyHostToDevice), + cudaSuccess); + + TransferEngine engine(false); + ASSERT_EQ(engine.init("P2PHANDSHAKE", "localhost:17932"), 0); + if (!engine.isUsingTent()) { + ASSERT_NE(engine.installTransport("tcp", nullptr), nullptr); + } + ASSERT_EQ(engine.registerLocalMemory(source.data(), source.size(), "cpu:0"), + 0); + ASSERT_EQ(engine.registerLocalMemory(cpu_destination.data(), + cpu_destination.size(), "cpu:0"), + 0); + ASSERT_EQ(engine.registerLocalMemory(gpu_source, kBufferSize, "cuda:0"), 0); + ASSERT_EQ( + engine.registerLocalMemory(gpu_destination, kBufferSize, "cuda:0"), 0); + + std::vector destination_offsets, source_offsets, + lengths(kFragmentCount, 1); + std::vector expected(kBufferSize, 0), actual(kBufferSize); + for (size_t i = 0; i < kFragmentCount; ++i) { + destination_offsets.push_back(i * 2); + source_offsets.push_back(i * 3); + expected[i * 2] = source[i * 3]; + } + auto make_transfer = [&](void* remote_source, void* local_destination) { + return TransferEngine::ScatterTransferRange{ + .opcode = TransferRequest::READ, + .remote_segment = engine.getLocalIpAndPort(), + .remote_base_offset = reinterpret_cast(remote_source), + .remote_size = kBufferSize, + .local_buffer = local_destination, + .local_capacity = kBufferSize, + .local_offsets = destination_offsets, + .remote_offsets = source_offsets, + .lengths = lengths, + .on_fragment_complete = {}, + }; + }; + + auto invalid_transfer = make_transfer(source.data(), gpu_destination); + invalid_transfer.remote_offsets = {}; + EXPECT_TRUE(engine.transferScatter({invalid_transfer}).IsInvalidArgument()); + + ASSERT_EQ(cudaMemset(gpu_destination, 0, kBufferSize), cudaSuccess); + ASSERT_TRUE( + engine.transferScatter({make_transfer(source.data(), gpu_destination)}) + .ok()); + ASSERT_EQ(cudaMemcpy(actual.data(), gpu_destination, kBufferSize, + cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(actual, expected); + + ASSERT_TRUE(engine + .transferScatter( + {make_transfer(gpu_source, cpu_destination.data())}) + .ok()); + EXPECT_EQ(cpu_destination, expected); + + ASSERT_EQ(cudaMemset(gpu_destination, 0, kBufferSize), cudaSuccess); + auto operation = + engine.submitScatter({make_transfer(gpu_source, gpu_destination)}); + ASSERT_EQ(engine.freeEngine(), 0); + ASSERT_TRUE(operation.wait().ok()); + ASSERT_EQ(cudaMemcpy(actual.data(), gpu_destination, kBufferSize, + cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(actual, expected); + + EXPECT_EQ(cudaFree(gpu_source), cudaSuccess); + EXPECT_EQ(cudaFree(gpu_destination), cudaSuccess); +} +#endif + // Test the locality decision used by TransferSubmitter::isLocalTransfer. // Same-host different-process pairs share an IP but have distinct ports; // they must NOT be treated as locally addressable, otherwise memcpy in the diff --git a/mooncake-store/tests/uds_transport_test.cpp b/mooncake-store/tests/uds_transport_test.cpp new file mode 100644 index 0000000000..073a02a7ce --- /dev/null +++ b/mooncake-store/tests/uds_transport_test.cpp @@ -0,0 +1,262 @@ +#include "uds_transport.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace mooncake { +namespace { + +std::string testSocketPath(const std::string &suffix) { + return "uds_transport_test_" + std::to_string(getpid()) + "_" + suffix; +} + +int createTempFd(const char *content) { + char path[] = "/tmp/uds_transport_test_XXXXXX"; + int fd = mkstemp(path); + if (fd < 0) return -1; + unlink(path); + if (write(fd, content, strlen(content)) < 0) { + close(fd); + return -1; + } + lseek(fd, 0, SEEK_SET); + return fd; +} + +std::string readFdContent(int fd) { + char buffer[64] = {}; + lseek(fd, 0, SEEK_SET); + ssize_t n = read(fd, buffer, sizeof(buffer) - 1); + if (n < 0) return ""; + return std::string(buffer, static_cast(n)); +} + +bool waitForFlag(const std::atomic &flag) { + for (int i = 0; i < 100; ++i) { + if (flag.load()) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + return flag.load(); +} + +} // namespace + +TEST(UdsTransportTest, SendsRawPayload) { + UdsAcceptor acceptor(testSocketPath("raw")); + std::atomic handled{false}; + + acceptor.registerHandler([&](UdsConnection &connection) { + uint32_t value = 0; + EXPECT_EQ(connection.recvRaw(&value, sizeof(value)), 0); + value += 1; + EXPECT_EQ(connection.sendRaw(&value, sizeof(value)), 0); + handled = true; + }); + auto start_result = acceptor.start(); + ASSERT_TRUE(start_result) << start_result.error(); + + UdsConnector connector(testSocketPath("raw")); + auto connection_result = connector.connect(); + ASSERT_TRUE(connection_result) << connection_result.error(); + auto connection = std::move(connection_result.value()); + + uint32_t value = 41; + ASSERT_EQ(connection->sendRaw(&value, sizeof(value)), 0); + ASSERT_EQ(connection->recvRaw(&value, sizeof(value)), 0); + EXPECT_EQ(value, 42); + + acceptor.stop(); + EXPECT_TRUE(handled.load()); +} + +TEST(UdsTransportTest, SendsFdFromClientToServer) { + UdsAcceptor acceptor(testSocketPath("client_fd")); + std::atomic received{false}; + + acceptor.registerHandler([&](UdsConnection &connection) { + uint32_t marker = 0; + int fd = connection.recvFd(&marker, sizeof(marker)); + ASSERT_GE(fd, 0); + EXPECT_EQ(marker, 7u); + EXPECT_EQ(readFdContent(fd), "client-to-server"); + close(fd); + received = true; + }); + auto start_result = acceptor.start(); + ASSERT_TRUE(start_result) << start_result.error(); + + UdsConnector connector(testSocketPath("client_fd")); + auto connection_result = connector.connect(); + ASSERT_TRUE(connection_result) << connection_result.error(); + auto connection = std::move(connection_result.value()); + + int fd = createTempFd("client-to-server"); + ASSERT_GE(fd, 0); + uint32_t marker = 7; + ASSERT_EQ(connection->sendFd(fd, &marker, sizeof(marker)), 0); + close(fd); + + ASSERT_TRUE(waitForFlag(received)); + acceptor.stop(); + EXPECT_TRUE(received.load()); +} + +TEST(UdsTransportTest, RejectsFdWithPartialPayload) { + int sockets[2] = {-1, -1}; + ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0, sockets), 0); + + UdsConnection receiver(sockets[1]); + int fd = createTempFd("partial-payload"); + ASSERT_GE(fd, 0); + + uint32_t marker = 7; + iovec iov; + iov.iov_base = ▮ + iov.iov_len = sizeof(marker) - 1; + + char control[CMSG_SPACE(sizeof(int))]; + memset(control, 0, sizeof(control)); + + msghdr msg; + memset(&msg, 0, sizeof(msg)); + msg.msg_iov = &iov; + msg.msg_iovlen = 1; + msg.msg_control = control; + msg.msg_controllen = sizeof(control); + + cmsghdr *cmsg = CMSG_FIRSTHDR(&msg); + cmsg->cmsg_level = SOL_SOCKET; + cmsg->cmsg_type = SCM_RIGHTS; + cmsg->cmsg_len = CMSG_LEN(sizeof(int)); + memcpy(CMSG_DATA(cmsg), &fd, sizeof(int)); + + ASSERT_EQ(sendmsg(sockets[0], &msg, 0), + static_cast(sizeof(marker) - 1)); + close(fd); + close(sockets[0]); + + uint32_t received_marker = 0; + EXPECT_LT(receiver.recvFd(&received_marker, sizeof(received_marker)), 0); +} + +TEST(UdsTransportTest, SendsFdFromServerToClient) { + UdsAcceptor acceptor(testSocketPath("server_fd")); + + acceptor.registerHandler([&](UdsConnection &connection) { + int fd = createTempFd("server-to-client"); + ASSERT_GE(fd, 0); + uint32_t marker = 9; + EXPECT_EQ(connection.sendFd(fd, &marker, sizeof(marker)), 0); + close(fd); + }); + auto start_result = acceptor.start(); + ASSERT_TRUE(start_result) << start_result.error(); + + UdsConnector connector(testSocketPath("server_fd")); + auto connection_result = connector.connect(); + ASSERT_TRUE(connection_result) << connection_result.error(); + auto connection = std::move(connection_result.value()); + + uint32_t marker = 0; + int fd = connection->recvFd(&marker, sizeof(marker)); + ASSERT_GE(fd, 0); + EXPECT_EQ(marker, 9u); + EXPECT_EQ(readFdContent(fd), "server-to-client"); + close(fd); + + acceptor.stop(); +} + +TEST(UdsTransportTest, StopWakesAcceptLoop) { + UdsAcceptor acceptor(testSocketPath("stop")); + acceptor.registerHandler([](UdsConnection &) {}); + auto start_result = acceptor.start(); + ASSERT_TRUE(start_result) << start_result.error(); + acceptor.stop(); + + UdsConnector connector(testSocketPath("stop")); + EXPECT_FALSE(connector.connect()); +} + +TEST(UdsTransportTest, StopWakesActiveClientHandler) { + UdsAcceptor acceptor(testSocketPath("active_stop")); + std::atomic entered{false}; + std::atomic exited{false}; + + acceptor.registerHandler([&](UdsConnection &connection) { + entered = true; + uint32_t value = 0; + EXPECT_EQ(connection.recvRaw(&value, sizeof(value)), -1); + exited = true; + }); + auto start_result = acceptor.start(); + ASSERT_TRUE(start_result) << start_result.error(); + + UdsConnector connector(testSocketPath("active_stop")); + auto connection_result = connector.connect(); + ASSERT_TRUE(connection_result) << connection_result.error(); + + ASSERT_TRUE(waitForFlag(entered)); + acceptor.stop(); + EXPECT_TRUE(exited.load()); +} + +TEST(UdsTransportTest, ConnectRejectsNonPositiveTimeout) { + UdsConnector connector(testSocketPath("invalid_timeout"), + std::chrono::milliseconds(0)); + auto connection_result = connector.connect(); + ASSERT_FALSE(connection_result); + EXPECT_NE(connection_result.error().find("timeout"), std::string::npos); +} + +TEST(UdsTransportTest, ConnectFailureReturnsPromptlyWithShortTimeout) { + std::string socket_name = testSocketPath("missing"); + UdsConnector connector(socket_name, std::chrono::milliseconds(10)); + + auto start = std::chrono::steady_clock::now(); + auto connection_result = connector.connect(); + auto elapsed = std::chrono::steady_clock::now() - start; + + EXPECT_FALSE(connection_result); + EXPECT_LT(elapsed, std::chrono::seconds(1)); + EXPECT_NE(connection_result.error().find(socket_name), std::string::npos); +} + +TEST(UdsTransportTest, ConnectClearsSendTimeout) { + UdsAcceptor acceptor(testSocketPath("send_timeout")); + acceptor.registerHandler([](UdsConnection &) {}); + auto start_result = acceptor.start(); + ASSERT_TRUE(start_result) << start_result.error(); + + UdsConnector connector(testSocketPath("send_timeout"), + std::chrono::milliseconds(10)); + auto connection_result = connector.connect(); + ASSERT_TRUE(connection_result) << connection_result.error(); + auto connection = std::move(connection_result.value()); + + timeval tv; + socklen_t tv_len = sizeof(tv); + ASSERT_EQ( + getsockopt(connection->fd(), SOL_SOCKET, SO_SNDTIMEO, &tv, &tv_len), 0); + EXPECT_EQ(tv.tv_sec, 0); + EXPECT_EQ(tv.tv_usec, 0); + + acceptor.stop(); +} + +TEST(UdsTransportTest, StartRequiresRegisteredHandler) { + UdsAcceptor acceptor(testSocketPath("no_handler")); + auto start_result = acceptor.start(); + ASSERT_FALSE(start_result); + EXPECT_NE(start_result.error().find("handler"), std::string::npos); +} + +} // namespace mooncake diff --git a/mooncake-store/tests/utils_test.cpp b/mooncake-store/tests/utils_test.cpp index b3102fd9fa..ef224eb203 100644 --- a/mooncake-store/tests/utils_test.cpp +++ b/mooncake-store/tests/utils_test.cpp @@ -1,13 +1,84 @@ #include "utils.h" +#include "random.h" +#include +#include #include #include +#include #include +#include #include #include using namespace mooncake; +TEST(RandomTest, ReusesEngineWithinThread) { + EXPECT_EQ(&threadLocalRandomEngine(), &threadLocalRandomEngine()); +} + +TEST(RandomTest, UsesDifferentEngineAcrossThreads) { + auto* main_engine = &threadLocalRandomEngine(); + std::promise worker_address; + std::promise release_worker; + auto release_future = release_worker.get_future(); + + std::thread worker([&] { + worker_address.set_value(&threadLocalRandomEngine()); + release_future.wait(); + }); + + EXPECT_NE(main_engine, worker_address.get_future().get()); + release_worker.set_value(); + worker.join(); +} + +TEST(RandomTest, RandomIndexStaysWithinBounds) { + std::mt19937_64 engine(42); + for (size_t i = 0; i < 10000; ++i) { + EXPECT_LT(randomIndex(17, engine), 17); + } +} + +TEST(RandomTest, ExplicitEngineIsDeterministic) { + std::mt19937_64 first(42); + std::mt19937_64 second(42); + for (size_t i = 0; i < 100; ++i) { + EXPECT_EQ(randomIndex(1000, first), randomIndex(1000, second)); + } +} + +TEST(RandomTest, RandomUniformIncludesRequestedBounds) { + std::mt19937_64 engine(42); + for (size_t i = 0; i < 10000; ++i) { + int value = randomUniform(-5, 9, engine); + EXPECT_GE(value, -5); + EXPECT_LE(value, 9); + } + EXPECT_EQ(randomUniform(7, 7, engine), 7); +} + +TEST(RandomTest, RandomUniformSupportsNarrowIntegers) { + std::mt19937_64 engine(42); + for (size_t i = 0; i < 1000; ++i) { + auto signed_value = randomUniform(-5, 9, engine); + EXPECT_GE(signed_value, -5); + EXPECT_LE(signed_value, 9); + + auto unsigned_value = randomUniform(2, 7, engine); + EXPECT_GE(unsigned_value, 2); + EXPECT_LE(unsigned_value, 7); + } + EXPECT_FALSE(randomUniform(false, false, engine)); + EXPECT_TRUE(randomUniform(true, true, engine)); +} + +TEST(RandomTest, RejectsInvalidBounds) { + std::mt19937_64 engine(42); + EXPECT_THROW(randomIndex(0, engine), std::invalid_argument); + EXPECT_THROW(randomUniform(2, 1, engine), std::invalid_argument); +} + TEST(UtilsTest, ByteSizeToString) { EXPECT_EQ(byte_size_to_string(999), "999 B"); EXPECT_EQ(byte_size_to_string(2048), "2.00 KB"); @@ -37,19 +108,6 @@ TEST(UtilsTest, StringToByteSize) { EXPECT_EQ(string_to_byte_size("-5"), 0); } -TEST(UtilsTest, StringToBool) { - EXPECT_EQ(string_to_bool("1"), true); - EXPECT_EQ(string_to_bool("true"), true); - EXPECT_EQ(string_to_bool("YES"), true); - EXPECT_EQ(string_to_bool(" on "), true); - EXPECT_EQ(string_to_bool("0"), false); - EXPECT_EQ(string_to_bool("false"), false); - EXPECT_EQ(string_to_bool("No"), false); - EXPECT_EQ(string_to_bool(" off "), false); - EXPECT_EQ(string_to_bool("maybe"), std::nullopt); - EXPECT_EQ(string_to_bool(""), std::nullopt); -} - TEST(UtilsTest, IsPortAvailable) { // Find an available port int test_port = -1; @@ -148,17 +206,6 @@ TEST(UtilsTest, AutoPortBinderMultipleInstances) { EXPECT_NE(port1, port2); } -TEST(UtilsTest, SplitStringBasic) { - std::string input = "a, b ,c, d"; - auto tokens = splitString(input, ',', true, false); - - ASSERT_EQ(tokens.size(), 4); - EXPECT_EQ(tokens[0], "a"); - EXPECT_EQ(tokens[1], "b"); - EXPECT_EQ(tokens[2], "c"); - EXPECT_EQ(tokens[3], "d"); -} - TEST(UtilsTest, GetInterfaceIPv4AddressLoopback) { auto address = GetInterfaceIPv4Address("lo"); ASSERT_TRUE(address.has_value()) << address.error(); diff --git a/mooncake-store/tools/CMakeLists.txt b/mooncake-store/tools/CMakeLists.txt new file mode 100644 index 0000000000..df983d66b7 --- /dev/null +++ b/mooncake-store/tools/CMakeLists.txt @@ -0,0 +1,6 @@ +add_executable(oplog_batch_inspector oplog_batch_inspector.cpp + oplog_batch_auditor.cpp) +target_include_directories(oplog_batch_inspector + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/..) +target_link_libraries(oplog_batch_inspector PRIVATE mooncake_store + JsonCpp::JsonCpp) diff --git a/mooncake-store/tools/oplog_batch_auditor.cpp b/mooncake-store/tools/oplog_batch_auditor.cpp new file mode 100644 index 0000000000..fb3fea6a4b --- /dev/null +++ b/mooncake-store/tools/oplog_batch_auditor.cpp @@ -0,0 +1,257 @@ +#include "tools/oplog_batch_auditor.h" + +#include +#include +#include +#include +#include + +#if __has_include() +#include // Ubuntu +#else +#include // CentOS +#endif + +#include "ha/oplog/oplog_batch_codec.h" +#include "ha/oplog/oplog_types.h" + +namespace mooncake { +namespace { + +constexpr size_t kMaxAuditErrors = 100; + +void AddError(OpLogAuditReport* report, std::string error) { + if (report->errors.size() < kMaxAuditErrors) { + report->errors.push_back(std::move(error)); + } else { + report->truncated_errors = true; + } +} + +bool ParsePaddedId(std::string_view value, uint64_t* id) { + if (value.size() != kOpLogBatchIdWidth || + !std::all_of(value.begin(), value.end(), + [](char c) { return c >= '0' && c <= '9'; })) { + return false; + } + auto result = + std::from_chars(value.data(), value.data() + value.size(), *id); + return result.ec == std::errc() && + result.ptr == value.data() + value.size(); +} + +std::string PrefixEnd(std::string prefix) { + for (size_t i = prefix.size(); i > 0; --i) { + auto value = static_cast(prefix[i - 1]); + if (value != 0xFF) { + prefix[i - 1] = static_cast(value + 1); + prefix.resize(i); + return prefix; + } + } + return std::string(1, '\0'); +} + +} // namespace + +ErrorCode ReadOpLogNamespace(const std::string& cluster_id, + HaKvBackend& backend, size_t max_keys, + OpLogNamespaceRead* result) { + if (result == nullptr || max_keys == 0 || + max_keys == std::numeric_limits::max()) { + return ErrorCode::INVALID_PARAMS; + } + result->kvs.clear(); + result->truncated = false; + + std::string normalized = cluster_id; + if (!NormalizeAndValidateClusterId(normalized) || normalized.empty()) { + return ErrorCode::INVALID_PARAMS; + } + const std::string prefix = "/oplog/" + normalized + "/"; + const ErrorCode err = + backend.Range(prefix, PrefixEnd(prefix), max_keys + 1, result->kvs); + if (err != ErrorCode::OK) { + return err; + } + if (result->kvs.size() > max_keys) { + result->kvs.resize(max_keys); + result->truncated = true; + } + return ErrorCode::OK; +} + +std::optional ComputeOpLogDumpLimit(uint64_t from_batch, + uint64_t to_batch, + size_t max_batches) { + if (from_batch == 0 || max_batches == 0 || + (to_batch != 0 && to_batch < from_batch)) { + return std::nullopt; + } + if (to_batch == 0) { + return max_batches; + } + const uint64_t difference = to_batch - from_batch; + if (difference >= max_batches) { + return std::nullopt; + } + return static_cast(difference + 1); +} + +OpLogAuditReport AuditOpLogNamespace(const std::string& cluster_id, + const std::vector& kvs, + size_t max_batches) { + OpLogAuditReport report; + report.cluster_id = cluster_id; + + std::string normalized = cluster_id; + if (!NormalizeAndValidateClusterId(normalized) || normalized.empty()) { + AddError(&report, "invalid cluster_id"); + return report; + } + report.cluster_id = normalized; + + const std::string prefix = "/oplog/" + normalized + "/"; + const std::string batch_prefix = prefix + "batches/"; + const std::string durable_key = prefix + "durable_prefix"; + std::map batches; + size_t raw_batch_count = 0; + + for (const auto& kv : kvs) { + if (kv.key == durable_key) { + DurablePrefix durable; + std::string reason; + if (!DecodeDurablePrefix(kv.value, &durable, &reason)) { + AddError(&report, "invalid durable_prefix: " + reason); + } else { + report.durable_prefix = durable; + } + continue; + } + + if (kv.key.starts_with(batch_prefix)) { + ++raw_batch_count; + uint64_t key_id = 0; + const std::string_view suffix(kv.key.data() + batch_prefix.size(), + kv.key.size() - batch_prefix.size()); + if (!ParsePaddedId(suffix, &key_id)) { + AddError(&report, "malformed batch key: " + kv.key); + continue; + } + OpLogBatchRecord batch; + std::string reason; + if (!DecodeOpLogBatchRecord(kv.value, &batch, &reason)) { + AddError(&report, "invalid batch " + std::to_string(key_id) + + ": " + reason); + continue; + } + if (batch.batch_id != key_id) { + AddError(&report, "batch id does not match key: " + kv.key); + continue; + } + batches.emplace(key_id, std::move(batch)); + continue; + } + + if (kv.key.starts_with(prefix)) { + uint64_t legacy_seq = 0; + const std::string_view suffix(kv.key.data() + prefix.size(), + kv.key.size() - prefix.size()); + if (ParsePaddedId(suffix, &legacy_seq)) { + report.legacy_max_seq = + std::max(report.legacy_max_seq, legacy_seq); + } else if (suffix.size() == kOpLogBatchIdWidth) { + AddError(&report, "malformed legacy key: " + kv.key); + } else if (suffix != "latest" && !suffix.starts_with("snapshot/") && + !suffix.starts_with("cleanup/")) { + report.warnings.push_back("unknown OpLog key: " + kv.key); + } + } + } + + report.batch_count = raw_batch_count; + if (raw_batch_count > max_batches) { + AddError(&report, "batch count exceeds audit limit"); + return report; + } + if (!batches.empty() && !report.durable_prefix.has_value()) { + for (const auto& [batch_id, batch] : batches) { + report.orphan_batches.push_back(batch_id); + report.entry_count += batch.entries.size(); + } + AddError(&report, "batch records exist without durable_prefix"); + return report; + } + + uint64_t expected_batch_id = 1; + uint64_t expected_first_seq = report.legacy_max_seq + 1; + for (const auto& [batch_id, batch] : batches) { + if (report.durable_prefix.has_value() && + batch_id > report.durable_prefix->batch_id) { + report.orphan_batches.push_back(batch_id); + continue; + } + if (batch_id != expected_batch_id) { + AddError(&report, + "batch id gap before " + std::to_string(batch_id)); + break; + } + if (batch.first_seq != expected_first_seq) { + AddError(&report, + "sequence gap before batch " + std::to_string(batch_id)); + break; + } + report.entry_count += batch.entries.size(); + expected_batch_id = batch_id + 1; + expected_first_seq = batch.last_seq + 1; + } + + if (!report.orphan_batches.empty()) { + AddError(&report, "batch records exist beyond durable_prefix"); + } + if (report.durable_prefix.has_value() && + (expected_batch_id != report.durable_prefix->batch_id + 1 || + expected_first_seq != report.durable_prefix->last_seq + 1)) { + AddError(&report, "durable_prefix does not match batch history"); + } + + report.ok = report.errors.empty(); + return report; +} + +std::string OpLogAuditReportToJson(const OpLogAuditReport& report) { + Json::Value root; + root["schema_version"] = 1; + root["ok"] = report.ok; + root["cluster_id"] = report.cluster_id; + root["legacy_max_seq"] = Json::UInt64(report.legacy_max_seq); + if (report.durable_prefix.has_value()) { + root["durable_prefix"]["batch_id"] = + Json::UInt64(report.durable_prefix->batch_id); + root["durable_prefix"]["last_seq"] = + Json::UInt64(report.durable_prefix->last_seq); + } else { + root["durable_prefix"] = Json::nullValue; + } + root["batch_count"] = Json::UInt64(report.batch_count); + root["entry_count"] = Json::UInt64(report.entry_count); + root["orphan_batches"] = Json::arrayValue; + for (uint64_t batch_id : report.orphan_batches) { + root["orphan_batches"].append(Json::UInt64(batch_id)); + } + root["warnings"] = Json::arrayValue; + for (const auto& warning : report.warnings) { + root["warnings"].append(warning); + } + root["errors"] = Json::arrayValue; + for (const auto& error : report.errors) { + root["errors"].append(error); + } + root["truncated_errors"] = report.truncated_errors; + + Json::StreamWriterBuilder builder; + builder["indentation"] = ""; + return Json::writeString(builder, root); +} + +} // namespace mooncake diff --git a/mooncake-store/tools/oplog_batch_auditor.h b/mooncake-store/tools/oplog_batch_auditor.h new file mode 100644 index 0000000000..e9a1b82c72 --- /dev/null +++ b/mooncake-store/tools/oplog_batch_auditor.h @@ -0,0 +1,43 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "ha/kv/ha_kv_backend.h" +#include "ha/oplog/oplog_batch_types.h" + +namespace mooncake { + +struct OpLogAuditReport { + bool ok{false}; + std::string cluster_id; + uint64_t legacy_max_seq{0}; + std::optional durable_prefix; + size_t batch_count{0}; + size_t entry_count{0}; + std::vector orphan_batches; + std::vector warnings; + std::vector errors; + bool truncated_errors{false}; +}; + +struct OpLogNamespaceRead { + std::vector kvs; + bool truncated{false}; +}; + +ErrorCode ReadOpLogNamespace(const std::string& cluster_id, + HaKvBackend& backend, size_t max_keys, + OpLogNamespaceRead* result); +std::optional ComputeOpLogDumpLimit(uint64_t from_batch, + uint64_t to_batch, + size_t max_batches); +OpLogAuditReport AuditOpLogNamespace(const std::string& cluster_id, + const std::vector& kvs, + size_t max_batches); +std::string OpLogAuditReportToJson(const OpLogAuditReport& report); + +} // namespace mooncake diff --git a/mooncake-store/tools/oplog_batch_inspector.cpp b/mooncake-store/tools/oplog_batch_inspector.cpp new file mode 100644 index 0000000000..4ca888a3b0 --- /dev/null +++ b/mooncake-store/tools/oplog_batch_inspector.cpp @@ -0,0 +1,217 @@ +#include +#include +#if __has_include() +#include // Ubuntu +#else +#include // CentOS +#endif + +#include +#include +#include +#include +#include +#include +#include + +#include "etcd_helper.h" +#include "ha/kv/etcd_ha_kv_backend.h" +#include "ha/oplog/oplog_batch_storage.h" +#include "ha/oplog/oplog_types.h" +#include "tools/oplog_batch_auditor.h" + +DEFINE_string(endpoints, "127.0.0.1:2379", "Etcd endpoints"); +DEFINE_string(cluster_id, "", "OpLog cluster ID"); +DEFINE_bool(json, false, "Emit JSON to stdout"); +DEFINE_uint64(max_batches, 1000000, + "Maximum namespace keys to audit or batches to dump"); +DEFINE_uint64(timeout_sec, 30, "Wait timeout in seconds"); +DEFINE_uint64(last_seq, 0, "Wait for at least this durable sequence"); +DEFINE_uint64(batch_id, 0, "Wait for at least this durable batch ID"); +DEFINE_uint64(from_batch, 1, "First batch ID to dump"); +DEFINE_uint64(to_batch, 0, "Last batch ID to dump; zero means no limit"); +DEFINE_bool(entries, false, "Include entry details in dump output"); + +namespace mooncake { +namespace { + +void PrintUsage(const char* program) { + std::cerr << "Usage: " << program + << " --cluster_id=ID " + "[--endpoints=HOST:PORT]\n"; +} + +void PrintTextReport(const OpLogAuditReport& report) { + std::cout << "cluster=" << report.cluster_id + << " ok=" << (report.ok ? "true" : "false") + << " legacy_max_seq=" << report.legacy_max_seq; + if (report.durable_prefix.has_value()) { + std::cout << " durable_batch=" << report.durable_prefix->batch_id + << " durable_seq=" << report.durable_prefix->last_seq; + } else { + std::cout << " durable_prefix=missing"; + } + std::cout << " batches=" << report.batch_count + << " entries=" << report.entry_count << '\n'; + for (const auto& warning : report.warnings) { + std::cout << "warning: " << warning << '\n'; + } + for (const auto& error : report.errors) { + std::cout << "error: " << error << '\n'; + } + if (report.truncated_errors) { + std::cout << "error: additional errors truncated\n"; + } +} + +int RunAudit(const std::string& command, const std::string& cluster_id, + HaKvBackend& backend) { + OpLogNamespaceRead input; + const ErrorCode err = ReadOpLogNamespace( + cluster_id, backend, static_cast(FLAGS_max_batches), &input); + if (err != ErrorCode::OK) { + LOG(ERROR) << "Failed to read OpLog namespace: " << toString(err); + return 1; + } + auto report = AuditOpLogNamespace(cluster_id, input.kvs, + static_cast(FLAGS_max_batches)); + if (input.truncated) { + report.ok = false; + if (report.errors.size() < 100) { + report.errors.push_back("namespace key count exceeds audit limit"); + } else { + report.truncated_errors = true; + } + } + if (FLAGS_json) { + std::cout << OpLogAuditReportToJson(report) << '\n'; + } else { + PrintTextReport(report); + } + return command == "verify" && !report.ok ? 3 : 0; +} + +int RunWait(const std::string& cluster_id, HaKvBackend& backend) { + OpLogBatchStorage storage(cluster_id, backend); + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::seconds(FLAGS_timeout_sec); + DurablePrefix prefix; + while (std::chrono::steady_clock::now() < deadline) { + const ErrorCode err = storage.ReadDurablePrefix(prefix); + if (err == ErrorCode::OK) { + if (prefix.last_seq >= FLAGS_last_seq && + prefix.batch_id >= FLAGS_batch_id) { + return RunAudit("verify", cluster_id, backend); + } + } else if (err != ErrorCode::ETCD_KEY_NOT_EXIST) { + LOG(ERROR) << "Failed to read durable prefix: " << toString(err); + return 1; + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + LOG(ERROR) << "Timed out waiting for durable prefix: last_seq>=" + << FLAGS_last_seq << ", batch_id>=" << FLAGS_batch_id; + return 4; +} + +int RunDump(const std::string& cluster_id, HaKvBackend& backend) { + const auto requested = + ComputeOpLogDumpLimit(FLAGS_from_batch, FLAGS_to_batch, + static_cast(FLAGS_max_batches)); + if (!requested.has_value()) { + LOG(ERROR) << "Invalid dump batch range or range exceeds " + "--max_batches"; + return 1; + } + + OpLogBatchStorage storage(cluster_id, backend); + std::vector batches; + const ErrorCode err = + storage.ReadBatchesAfter(FLAGS_from_batch - 1, *requested, batches); + if (err != ErrorCode::OK) { + LOG(ERROR) << "Failed to read batches: " << toString(err); + return 1; + } + + Json::Value root(Json::arrayValue); + for (const auto& batch : batches) { + if (FLAGS_to_batch != 0 && batch.batch_id > FLAGS_to_batch) { + break; + } + Json::Value item; + item["batch_id"] = Json::UInt64(batch.batch_id); + item["first_seq"] = Json::UInt64(batch.first_seq); + item["last_seq"] = Json::UInt64(batch.last_seq); + item["entry_count"] = Json::UInt64(batch.entries.size()); + item["checksum"] = batch.checksum; + if (FLAGS_entries) { + item["entries"] = Json::arrayValue; + for (const auto& entry : batch.entries) { + Json::Value encoded_entry; + encoded_entry["sequence_id"] = Json::UInt64(entry.sequence_id); + encoded_entry["op_type"] = + static_cast(entry.op_type); + encoded_entry["tenant_id"] = entry.tenant_id; + encoded_entry["object_key"] = entry.object_key; + encoded_entry["payload_bytes"] = + Json::UInt64(entry.payload.size()); + item["entries"].append(std::move(encoded_entry)); + } + } + root.append(std::move(item)); + } + + Json::StreamWriterBuilder builder; + builder["indentation"] = FLAGS_json ? "" : " "; + std::cout << Json::writeString(builder, root) << '\n'; + return 0; +} + +} // namespace +} // namespace mooncake + +int main(int argc, char** argv) { + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = true; + gflags::SetUsageMessage( + "Inspect and verify Mooncake batch-record OpLog history"); + gflags::ParseCommandLineFlags(&argc, &argv, true); + + if (argc != 2 || FLAGS_cluster_id.empty()) { + mooncake::PrintUsage(argv[0]); + return 1; + } + const std::string command = argv[1]; + if (command != "verify" && command != "summary" && command != "wait" && + command != "dump") { + mooncake::PrintUsage(argv[0]); + return 1; + } + if (FLAGS_max_batches == 0 || + FLAGS_max_batches >= std::numeric_limits::max()) { + LOG(ERROR) << "max_batches must fit in size_t and be greater than zero"; + return 1; + } + + std::string cluster_id = FLAGS_cluster_id; + if (!mooncake::NormalizeAndValidateClusterId(cluster_id) || + cluster_id.empty()) { + LOG(ERROR) << "Invalid cluster_id"; + return 1; + } + const auto err = + mooncake::EtcdHelper::ConnectToEtcdStoreClient(FLAGS_endpoints); + if (err != mooncake::ErrorCode::OK) { + LOG(ERROR) << "Failed to connect to etcd: " << mooncake::toString(err); + return 1; + } + + mooncake::EtcdHaKvBackend backend; + if (command == "wait") { + return mooncake::RunWait(cluster_id, backend); + } + if (command == "dump") { + return mooncake::RunDump(cluster_id, backend); + } + return mooncake::RunAudit(command, cluster_id, backend); +} diff --git a/mooncake-transfer-engine/CMakeLists.txt b/mooncake-transfer-engine/CMakeLists.txt index 121c3972d4..f43e814622 100644 --- a/mooncake-transfer-engine/CMakeLists.txt +++ b/mooncake-transfer-engine/CMakeLists.txt @@ -1,6 +1,8 @@ cmake_minimum_required(VERSION 3.16) project(mooncake-transfer-engine) +include(CheckCXXSourceCompiles) + if (NOT GLOBAL_CONFIG) if (USE_ETCD) message(FATAL_ERROR "Cannot enable USE_ETCD while building transfer engine independently") @@ -16,6 +18,31 @@ if (NOT GLOBAL_CONFIG) ${CMAKE_CURRENT_BINARY_DIR}/mooncake-common-src) endif() # GLOBAL_CONFIG +if (USE_CUDA) + set(CMAKE_REQUIRED_LIBRARIES mlx5) + check_cxx_source_compiles( + "#include + #include + int main() { + mlx5dv_devx_umem_in input{}; + input.comp_mask = MLX5DV_UMEM_MASK_DMABUF; + input.dmabuf_fd = -1; + auto reg_ex = &mlx5dv_devx_umem_reg_ex; + (void)reg_ex; + return 0; + }" + MOONCAKE_HAVE_MLX5_DMABUF_UMEM) + unset(CMAKE_REQUIRED_LIBRARIES) + + if (MOONCAKE_HAVE_MLX5_DMABUF_UMEM) + add_compile_definitions(MOONCAKE_HAVE_MLX5_DMABUF_UMEM) + message(STATUS "CUDA IBGDA DMA-BUF DevX UMEM support enabled") + else() + message(WARNING + "CUDA IBGDA DMA-BUF DevX UMEM support disabled: the installed mlx5 headers or libmlx5 do not provide MLX5DV_UMEM_MASK_DMABUF, mlx5dv_devx_umem_in::dmabuf_fd, and mlx5dv_devx_umem_reg_ex. Update rdma-core/libmlx5 to enable DMA-BUF control memory; Device API will use GPU-VA and host-mapped fallback.") + endif() +endif() + if (USE_ASCEND) include(../mooncake-common/FindMpi.cmake) file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux") @@ -63,7 +90,7 @@ endif() # Non-peermem requires cuMem allocated memory pointers; since WITH_NVIDIA_PEERMEM # is now a runtime environment variable, always build when GPU support is present. -if (USE_MNNVL OR USE_CUDA OR USE_HIP OR USE_MUSA OR USE_MACA) +if (USE_MNNVL OR USE_CUDA OR USE_HIP OR USE_MUSA OR USE_MACA OR USE_SUPA) add_subdirectory(nvlink-allocator) endif() diff --git a/mooncake-transfer-engine/benchmark/CMakeLists.txt b/mooncake-transfer-engine/benchmark/CMakeLists.txt index b17a8a803d..e13773ec5b 100644 --- a/mooncake-transfer-engine/benchmark/CMakeLists.txt +++ b/mooncake-transfer-engine/benchmark/CMakeLists.txt @@ -30,6 +30,8 @@ file(GLOB TEBENCH_SOURCES "*.cpp") # translation unit (which pulls in tent/ headers) from non-TENT builds. if(NOT USE_TENT) list(REMOVE_ITEM TEBENCH_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/tent_backend.cpp") + list(APPEND TEBENCH_SOURCES + "${CMAKE_CURRENT_SOURCE_DIR}/../tent/src/common/qos_metrics.cpp") endif() add_executable(tebench ${TEBENCH_SOURCES}) target_link_libraries(tebench PUBLIC transfer_engine) @@ -63,3 +65,23 @@ endif() set_target_properties( tebench PROPERTIES BUILD_WITH_INSTALL_RPATH TRUE INSTALL_RPATH "$ORIGIN/../lib:$ORIGIN/../../mooncake-common${TANGRT_RPATH}") + +if(BUILD_UNIT_TESTS) + add_executable(tebench_qos_metrics_test tests/qos_metrics_test.cpp + qos_metrics_adapter.cpp + workload_config.cpp utils.cpp) + if(NOT USE_TENT) + target_sources(tebench_qos_metrics_test PRIVATE + ../tent/src/common/qos_metrics.cpp) + endif() + target_link_libraries(tebench_qos_metrics_test + PRIVATE transfer_engine gtest gtest_main) + if(USE_TENT) + target_link_libraries(tebench_qos_metrics_test PRIVATE tent_common) + endif() + target_include_directories( + tebench_qos_metrics_test + PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}" + "${CMAKE_CURRENT_SOURCE_DIR}/../tent/include") + add_test(NAME tebench_qos_metrics_test COMMAND tebench_qos_metrics_test) +endif() diff --git a/mooncake-transfer-engine/benchmark/bench_runner.h b/mooncake-transfer-engine/benchmark/bench_runner.h index f0bd69e4e2..0023ed0b72 100644 --- a/mooncake-transfer-engine/benchmark/bench_runner.h +++ b/mooncake-transfer-engine/benchmark/bench_runner.h @@ -28,6 +28,8 @@ namespace mooncake { namespace tent { +enum class IntentType : int; + class BenchRunner { public: BenchRunner() {} @@ -57,10 +59,11 @@ class BenchRunner { virtual double runSingleTransfer(uint64_t local_addr, uint64_t target_addr, uint64_t block_size, uint64_t batch_size, - OpCode opcode) = 0; + OpCode opcode, uint64_t deadline_ns, + IntentType intent_type) = 0; }; } // namespace tent } // namespace mooncake -#endif // BENCH_RUNNER_H \ No newline at end of file +#endif // BENCH_RUNNER_H diff --git a/mooncake-transfer-engine/benchmark/main.cpp b/mooncake-transfer-engine/benchmark/main.cpp index 6e9d76cb0e..ad756e980b 100644 --- a/mooncake-transfer-engine/benchmark/main.cpp +++ b/mooncake-transfer-engine/benchmark/main.cpp @@ -15,15 +15,19 @@ #include "utils.h" #include "bench_runner.h" +#include "qos_metrics_adapter.h" #include "te_backend.h" +#include "workload_config.h" #ifdef USE_TENT #include "tent_backend.h" #endif using namespace mooncake::tent; -int processBatchSizes(BenchRunner& runner, size_t block_size, size_t batch_size, - int num_threads) { +int processBatchSizes( + BenchRunner& runner, size_t block_size, size_t batch_size, int num_threads, + const std::vector& qos_classes, + const std::vector& workload_classes = {}) { bool mixed_opcode = false; OpCode opcode = READ; if (XferBenchConfig::check_consistency || XferBenchConfig::op_type == "mix") @@ -38,22 +42,57 @@ int processBatchSizes(BenchRunner& runner, size_t block_size, size_t batch_size, } XferBenchStats stats; + std::vector qos_stats(qos_classes.size()); + XferBenchStats tight_stats; + XferBenchStats loose_stats; std::mutex mutex; + size_t address_stride_bytes = + XferBenchConfig::max_block_size * XferBenchConfig::max_batch_size; + if (!workload_classes.empty()) { + address_stride_bytes = 0; + for (const auto& config : workload_classes) { + address_stride_bytes = std::max( + address_stride_bytes, config.block_size * config.batch_size); + } + } int rc = runner.runInitiatorTasks([&](int thread_id) -> int { runner.pinThread(thread_id); - auto max_block_size = XferBenchConfig::max_block_size; - auto max_batch_size = XferBenchConfig::max_batch_size; auto local_gpu_offset = std::max(0, XferBenchConfig::local_gpu_id); auto target_gpu_offset = std::max(0, XferBenchConfig::target_gpu_id); uint64_t local_addr = runner.getLocalBufferBase( - local_gpu_offset + thread_id, max_block_size, max_batch_size); + local_gpu_offset + thread_id, address_stride_bytes, 1); uint64_t target_addr = runner.getTargetBufferBase( - target_gpu_offset + thread_id, max_block_size, max_batch_size); + target_gpu_offset + thread_id, address_stride_bytes, 1); + const bool qos_enabled = !qos_classes.empty(); + const size_t qos_class = + qos_enabled ? qosClassForThread(qos_classes, thread_id) : 0; + const auto* workload = + workload_classes.empty() ? nullptr : &workload_classes[qos_class]; + const size_t thread_block_size = + workload ? workload->block_size : block_size; + const size_t thread_batch_size = + workload ? workload->batch_size : batch_size; + const IntentType intent_type = + workload ? workload->intent_type : IntentType::INTENT_UNSPEC; + const bool tight = XferBenchConfig::deadline_us > 0 && + thread_id < XferBenchConfig::deadline_tight_threads; + auto deadlineNs = [&]() -> uint64_t { + const uint64_t deadline_us = + workload ? workload->deadline_us + : (tight ? XferBenchConfig::deadline_us : 0); + if (deadline_us == 0) return 0; + const auto now = + std::chrono::steady_clock::now().time_since_epoch(); + return std::chrono::duration_cast(now) + .count() + + deadline_us * 1000ull; + }; XferBenchTimer timer; while (timer.lap_us(false) < 1000000ull) { - runner.runSingleTransfer(local_addr, target_addr, block_size, - batch_size, opcode); + runner.runSingleTransfer(local_addr, target_addr, thread_block_size, + thread_batch_size, opcode, deadlineNs(), + intent_type); } timer.reset(); std::vector transfer_duration; @@ -62,37 +101,75 @@ int processBatchSizes(BenchRunner& runner, size_t block_size, size_t batch_size, XferBenchConfig::duration * 1000000ull) { uint8_t pattern = 0; if (XferBenchConfig::check_consistency) - pattern = - fillData((void*)local_addr, block_size * batch_size); + pattern = fillData((void*)local_addr, + thread_block_size * thread_batch_size); auto val = runner.runSingleTransfer( - local_addr, target_addr, block_size, batch_size, WRITE); + local_addr, target_addr, thread_block_size, + thread_batch_size, WRITE, deadlineNs(), intent_type); transfer_duration.push_back(val); - fillData((void*)local_addr, block_size * batch_size); - val = runner.runSingleTransfer(local_addr, target_addr, - block_size, batch_size, READ); + fillData((void*)local_addr, + thread_block_size * thread_batch_size); + val = runner.runSingleTransfer( + local_addr, target_addr, thread_block_size, + thread_batch_size, READ, deadlineNs(), intent_type); if (XferBenchConfig::check_consistency) - verifyData((void*)local_addr, block_size * batch_size, - pattern); + verifyData((void*)local_addr, + thread_block_size * thread_batch_size, pattern); transfer_duration.push_back(val); } } else { while (timer.lap_us(false) < XferBenchConfig::duration * 1000000ull) { auto val = runner.runSingleTransfer( - local_addr, target_addr, block_size, batch_size, opcode); + local_addr, target_addr, thread_block_size, + thread_batch_size, opcode, deadlineNs(), intent_type); transfer_duration.push_back(val); } } auto total_duration = timer.lap_us(); - mutex.lock(); + std::lock_guard lock(mutex); stats.total_duration.add(total_duration); - for (auto val : transfer_duration) stats.transfer_duration.add(val); - mutex.unlock(); + stats.transfer_duration.add(transfer_duration); + if (qos_enabled) { + qos_stats[qos_class].total_duration.add(total_duration); + qos_stats[qos_class].transfer_duration.add(transfer_duration); + } + auto& group_stats = tight ? tight_stats : loose_stats; + group_stats.total_duration.add(total_duration); + group_stats.transfer_duration.add(transfer_duration); return 0; }); if (rc != 0) return -1; - printStats(block_size, batch_size, stats, num_threads); + if (workload_classes.empty()) + printStats(block_size, batch_size, stats, num_threads); + if (!qos_classes.empty()) { + std::vector bytes_per_operation; + for (const auto& config : workload_classes) { + bytes_per_operation.push_back(config.block_size * + config.batch_size); + } + auto report = calculateQosMetricsFromBenchStats( + block_size, batch_size, num_threads, qos_classes, &qos_stats, + XferBenchConfig::qos_link_capacity_gbps, bytes_per_operation); + printQosMetrics(report); + if (!XferBenchConfig::qos_output_jsonl.empty()) { + std::string error; + if (!appendQosMetricsJsonl(XferBenchConfig::qos_output_jsonl, + report, &error)) { + LOG(ERROR) << error; + return -1; + } + } + } + if (XferBenchConfig::deadline_us > 0 && workload_classes.empty()) { + const int tight_threads = + std::min(num_threads, XferBenchConfig::deadline_tight_threads); + printDeadlineGroupStats("tight", block_size, batch_size, tight_stats, + tight_threads, XferBenchConfig::deadline_us); + printDeadlineGroupStats("loose", block_size, batch_size, loose_stats, + num_threads - tight_threads, 0); + } return 0; } @@ -102,6 +179,103 @@ int main(int argc, char* argv[]) { "Usage: ./tebench [options]"); gflags::ParseCommandLineFlags(&argc, &argv, true); XferBenchConfig::loadFromFlags(); + std::vector workload_classes; + std::vector qos_classes; + if (!XferBenchConfig::workload_classes_json.empty() && + (!XferBenchConfig::qos_classes.empty() || + !XferBenchConfig::qos_classes_json.empty())) { + LOG(ERROR) + << "Use --workload_classes_json or QoS class flags, not both"; + return EXIT_FAILURE; + } + if (!XferBenchConfig::qos_classes.empty() && + !XferBenchConfig::qos_classes_json.empty()) { + LOG(ERROR) << "Use only one of --qos_classes or --qos_classes_json"; + return EXIT_FAILURE; + } + if (!XferBenchConfig::workload_classes_json.empty()) { + std::string error; + if (!parseWorkloadClassesJson(XferBenchConfig::workload_classes_json, + &workload_classes, &error)) { + LOG(ERROR) << "Invalid --workload_classes_json: " << error; + return EXIT_FAILURE; + } + qos_classes = qosClassesFromWorkload(workload_classes); + } else if (!XferBenchConfig::qos_classes_json.empty()) { + std::string error; + if (!parseQosClassesJson(XferBenchConfig::qos_classes_json, + &qos_classes, &error)) { + LOG(ERROR) << "Invalid --qos_classes_json: " << error; + return EXIT_FAILURE; + } + } else if (!XferBenchConfig::qos_classes.empty()) { + std::string error; + if (!parseQosClasses(XferBenchConfig::qos_classes, &qos_classes, + &error)) { + LOG(ERROR) << "Invalid --qos_classes: " << error; + return EXIT_FAILURE; + } + } + if (!qos_classes.empty()) { + std::string error; + if (XferBenchConfig::start_num_threads != + XferBenchConfig::max_num_threads) { + LOG(ERROR) + << "QoS metrics require start_num_threads == max_num_threads"; + return EXIT_FAILURE; + } + if (!validateQosClasses(qos_classes, XferBenchConfig::start_num_threads, + &error)) { + LOG(ERROR) << "Invalid QoS classes: " << error; + return EXIT_FAILURE; + } + if (!workload_classes.empty() && + !validateWorkloadClasses( + workload_classes, XferBenchConfig::start_num_threads, + XferBenchConfig::total_buffer_size, &error)) { + LOG(ERROR) << "Invalid workload classes: " << error; + return EXIT_FAILURE; + } + } + if (XferBenchConfig::qos_link_capacity_gbps < 0.0 || + !std::isfinite(XferBenchConfig::qos_link_capacity_gbps)) { + LOG(ERROR) << "qos_link_capacity_gbps must be finite and non-negative"; + return EXIT_FAILURE; + } + if (XferBenchConfig::deadline_tight_threads < 0 || + XferBenchConfig::deadline_tight_threads > + XferBenchConfig::max_num_threads) { + LOG(ERROR) << "deadline_tight_threads must be in [0, max_num_threads]"; + return EXIT_FAILURE; + } + if (XferBenchConfig::deadline_us > 0 && + XferBenchConfig::backend != "tent") { + LOG(ERROR) << "deadline tagging is supported only by the tent backend"; + return EXIT_FAILURE; + } + if (!workload_classes.empty() && + XferBenchConfig::tent_intent_type != "unspec") { + LOG(ERROR) << "--tent_intent_type cannot be combined with per-class " + "intent_type"; + return EXIT_FAILURE; + } + if (!workload_classes.empty() && + (XferBenchConfig::deadline_us != 0 || + XferBenchConfig::deadline_tight_threads != 0)) { + LOG(ERROR) << "--deadline_us and --deadline_tight_threads cannot be " + "combined with per-class workload deadlines"; + return EXIT_FAILURE; + } + if (!workload_classes.empty() && XferBenchConfig::backend != "tent") { + for (const auto& config : workload_classes) { + if (config.deadline_us != 0 || + config.intent_type != IntentType::INTENT_UNSPEC) { + LOG(ERROR) << "per-class deadline_us and intent_type require " + "the tent backend"; + return EXIT_FAILURE; + } + } + } std::unique_ptr runner; if (XferBenchConfig::backend == "classic") { runner = std::make_unique(); @@ -124,6 +298,14 @@ int main(int argc, char* argv[]) { << "Press Ctrl-C to terminate\033[0m" << std::endl; return runner->runTarget(); } + if (!workload_classes.empty()) { + const int num_threads = XferBenchConfig::start_num_threads; + if (runner->startInitiator(num_threads) != 0) return EXIT_FAILURE; + const int rc = processBatchSizes(*runner, 0, 0, num_threads, + qos_classes, workload_classes); + runner->stopInitiator(); + return rc == 0 ? EXIT_SUCCESS : EXIT_FAILURE; + } printStatsHeader(); bool interrupted = false; for (int num_threads = XferBenchConfig::start_num_threads; @@ -145,7 +327,7 @@ int main(int argc, char* argv[]) { << " batch_size " << batch_size; } else { if (processBatchSizes(*runner, block_size, batch_size, - num_threads) != 0) + num_threads, qos_classes) != 0) interrupted = true; } } diff --git a/mooncake-transfer-engine/benchmark/qos_metrics_adapter.cpp b/mooncake-transfer-engine/benchmark/qos_metrics_adapter.cpp new file mode 100644 index 0000000000..94882f09bb --- /dev/null +++ b/mooncake-transfer-engine/benchmark/qos_metrics_adapter.cpp @@ -0,0 +1,56 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "qos_metrics_adapter.h" + +namespace mooncake { +namespace tent { + +std::vector makeQosClassSamples( + const std::vector& classes, + std::vector* stats) { + std::vector samples(classes.size()); + for (size_t i = 0; i < classes.size(); ++i) { + auto& class_stats = (*stats)[i]; + auto& sample = samples[i]; + sample.operations = class_stats.transfer_duration.count(); + sample.total_duration_us = class_stats.total_duration.avg(); + sample.p99_us = class_stats.transfer_duration.p99(); + if (classes[i].slo_us != 0) { + sample.slo_attainment = + class_stats.transfer_duration.fractionAtOrBelow( + static_cast(classes[i].slo_us)); + } + } + return samples; +} + +QosMetricsReport calculateQosMetricsFromBenchStats( + size_t block_size, size_t batch_size, int num_threads, + const std::vector& classes, + std::vector* stats, double link_capacity_gbps, + const std::vector& bytes_per_operation) { + auto samples = makeQosClassSamples(classes, stats); + if (bytes_per_operation.size() == samples.size()) { + for (size_t i = 0; i < samples.size(); ++i) { + samples[i].transferred_bytes = + bytes_per_operation[i] * samples[i].operations; + } + } + return calculateQosMetrics(block_size, batch_size, num_threads, classes, + samples, link_capacity_gbps); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/benchmark/qos_metrics_adapter.h b/mooncake-transfer-engine/benchmark/qos_metrics_adapter.h new file mode 100644 index 0000000000..ac1f4b15e2 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/qos_metrics_adapter.h @@ -0,0 +1,39 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TEBENCH_QOS_METRICS_ADAPTER_H +#define TEBENCH_QOS_METRICS_ADAPTER_H + +#include + +#include "tent/common/qos_metrics.h" +#include "utils.h" + +namespace mooncake { +namespace tent { + +std::vector makeQosClassSamples( + const std::vector& classes, + std::vector* stats); + +QosMetricsReport calculateQosMetricsFromBenchStats( + size_t block_size, size_t batch_size, int num_threads, + const std::vector& classes, + std::vector* stats, double link_capacity_gbps, + const std::vector& bytes_per_operation = {}); + +} // namespace tent +} // namespace mooncake + +#endif // TEBENCH_QOS_METRICS_ADAPTER_H diff --git a/mooncake-transfer-engine/benchmark/te_backend.cpp b/mooncake-transfer-engine/benchmark/te_backend.cpp index 7b742e94ac..5e3632aa46 100644 --- a/mooncake-transfer-engine/benchmark/te_backend.cpp +++ b/mooncake-transfer-engine/benchmark/te_backend.cpp @@ -349,7 +349,11 @@ int TEBenchRunner::runInitiatorTasks( double TEBenchRunner::runSingleTransfer(uint64_t local_addr, uint64_t target_addr, uint64_t block_size, - uint64_t batch_size, OpCode opcode) { + uint64_t batch_size, OpCode opcode, + uint64_t deadline_ns, + IntentType intent_type) { + (void)deadline_ns; + (void)intent_type; auto batch_id = engine_->allocateBatchID(batch_size); std::vector requests; for (uint64_t i = 0; i < batch_size; ++i) { diff --git a/mooncake-transfer-engine/benchmark/te_backend.h b/mooncake-transfer-engine/benchmark/te_backend.h index 7d48a1d96d..9685f37421 100644 --- a/mooncake-transfer-engine/benchmark/te_backend.h +++ b/mooncake-transfer-engine/benchmark/te_backend.h @@ -69,7 +69,8 @@ class TEBenchRunner : public BenchRunner { double runSingleTransfer(uint64_t local_addr, uint64_t target_addr, uint64_t block_size, uint64_t batch_size, - OpCode opcode); + OpCode opcode, uint64_t deadline_ns, + IntentType intent_type); private: int allocateBuffers(); diff --git a/mooncake-transfer-engine/benchmark/tent_backend.cpp b/mooncake-transfer-engine/benchmark/tent_backend.cpp index 3aa24db0d4..2c0b731f17 100644 --- a/mooncake-transfer-engine/benchmark/tent_backend.cpp +++ b/mooncake-transfer-engine/benchmark/tent_backend.cpp @@ -51,6 +51,8 @@ std::shared_ptr loadConfig() { config->set("metadata_type", XferBenchConfig::metadata_type); config->set("metadata_servers", XferBenchConfig::metadata_url_list); config->set("rpc_server_port", XferBenchConfig::rpc_server_port); + config->set("transports/rdma/deadline_bw_arbitration", + XferBenchConfig::deadline_bw_arbitration); // Configure transport types based on xport_type parameter if (!XferBenchConfig::xport_type.empty()) { @@ -92,20 +94,36 @@ static TransportType getTransportType(const std::string& xport_type) { return UNSPEC; } -int TENTBenchRunner::allocateBuffers() { - const auto total_buffer_size = XferBenchConfig::total_buffer_size; - const auto& seg_type = XferBenchConfig::seg_type; - const auto& xport_type = XferBenchConfig::xport_type; - - // Resolve device prefix, start index, and buffer count per seg_type - std::string device_prefix; - int start_idx = 0, num_buffers = 0; +static IntentType getIntentType(const std::string& intent_type) { + std::string normalized_intent = intent_type; + for (auto& c : normalized_intent) c = to_lower(c); + + static const std::unordered_map kIntentTypes = { + {"unspec", IntentType::INTENT_UNSPEC}, + {"intent_unspec", IntentType::INTENT_UNSPEC}, + {"foreground_get", IntentType::FOREGROUND_GET}, + {"background_prefetch", IntentType::BACKGROUND_PREFETCH}, + {"migration", IntentType::MIGRATION}, + {"checkpoint", IntentType::CHECKPOINT}, + {"weight_loading", IntentType::WEIGHT_LOADING}, + {"staging_internal", IntentType::STAGING_INTERNAL}, + }; + auto it = kIntentTypes.find(normalized_intent); + LOG_ASSERT(it != kIntentTypes.end()) + << "Invalid --tent_intent_type=" << intent_type; + return it->second; +} - if (seg_type == "DRAM") { +// Resolve device prefix, start index, and buffer count for a single seg_type. +// Returns 0 on success, -1 on failure. +static int resolveSegTypeParams(const std::string& seg_type, + std::string& device_prefix, int& start_idx, + int& num_buffers) { + if (seg_type == "DRAM" || seg_type == "dram") { device_prefix = "cpu"; num_buffers = numa_num_configured_nodes(); #if defined(USE_CUDA) || defined(USE_SUNRISE) - } else if (seg_type == "VRAM") { + } else if (seg_type == "VRAM" || seg_type == "vram") { device_prefix = "cuda"; int gpu_count = 0; auto err = cudaGetDeviceCount(&gpu_count); @@ -121,7 +139,7 @@ int TENTBenchRunner::allocateBuffers() { << gpu_count << ")"; } #elif defined(USE_HIP) - } else if (seg_type == "VRAM") { + } else if (seg_type == "VRAM" || seg_type == "vram") { device_prefix = "rocm"; int gpu_count = 0; hipGetDeviceCount(&gpu_count); @@ -139,48 +157,123 @@ int TENTBenchRunner::allocateBuffers() { LOG(ERROR) << "Unknown seg_type: " << seg_type; return -1; } + return 0; +} + +// Parse a comma-separated list (e.g. "dram,vram") into a vector of lowercased +// seg_type names. Empty input returns an empty vector. +static std::vector parseSegTypeMix(const std::string& mix) { + std::vector result; + if (mix.empty()) return result; + std::stringstream ss(mix); + std::string token; + while (std::getline(ss, token, ',')) { + size_t b = token.find_first_not_of(" \t"); + size_t e = token.find_last_not_of(" \t"); + if (b == std::string::npos) continue; + std::string name = token.substr(b, e - b + 1); + for (auto& c : name) c = to_lower(c); + result.push_back(name); + } + return result; +} - pinned_buffer_list_.resize(num_buffers, nullptr); +int TENTBenchRunner::allocateBuffers() { + const auto total_buffer_size = XferBenchConfig::total_buffer_size; + const auto& xport_type = XferBenchConfig::xport_type; + + // Determine the list of seg_types to allocate. --seg_type_mix wins over + // --seg_type when set; otherwise fall back to --seg_type (single type). + seg_type_mix_ = parseSegTypeMix(XferBenchConfig::seg_type_mix); + std::vector seg_types; + if (!seg_type_mix_.empty()) { + seg_types = seg_type_mix_; + } else { + std::string s = XferBenchConfig::seg_type; + for (auto& c : s) c = to_lower(c); + seg_types.push_back(s); + } + + pinned_buffer_list_.clear(); + pinned_buffer_seg_type_.clear(); uint64_t alloc_ns = 0, reg_ns = 0; - for (int i = 0; i < num_buffers; ++i) { - auto location = device_prefix + ":" + std::to_string(start_idx + i); - MemoryOptions options; - if (!xport_type.empty()) { - options.type = getTransportType(xport_type); - options.location = location; + uint64_t total_bytes = 0; + + for (const auto& seg_type : seg_types) { + std::string device_prefix; + int start_idx = 0, num_buffers = 0; + if (resolveSegTypeParams(seg_type, device_prefix, start_idx, + num_buffers) != 0) { + return -1; } - auto t0 = getCurrentTimeInNano(); - if (!xport_type.empty()) { + for (int i = 0; i < num_buffers; ++i) { + auto location = device_prefix + ":" + std::to_string(start_idx + i); + MemoryOptions options; options.location = location; - CHECK_FAIL(engine_->allocateLocalMemory( - &pinned_buffer_list_[i], total_buffer_size, options)); - } else { - CHECK_FAIL(engine_->allocateLocalMemory( - &pinned_buffer_list_[i], total_buffer_size, location)); - } - auto t1 = getCurrentTimeInNano(); + if (!xport_type.empty()) { + // Explicit single transport: honor it for both allocate and + // register (existing behavior). + options.type = getTransportType(xport_type); + } else if (!seg_type_mix_.empty()) { + // Multi-transport mode (--seg_type_mix set, transports + // enabled via MC_TENT_CONF). Pick the allocate transport + // based on seg_type so the right transport-specific allocator + // runs (e.g. SHM sets shm_path, NVLink sets up GPU memory + // handle). registerLocalMemory below then resets options.type + // to UNSPEC so the buffer registers to ALL enabled transports, + // not just this one. + options.type = (seg_type == "vram") ? NVLINK : SHM; + } + // else: single-seg_type fallback (--seg_type without + // --seg_type_mix, no --xport_type). Leave options.type as UNSPEC + // so the engine picks the default transport (original behavior). + + auto t0 = getCurrentTimeInNano(); + void* buf = nullptr; + CHECK_FAIL( + engine_->allocateLocalMemory(&buf, total_buffer_size, options)); + pinned_buffer_list_.push_back(buf); + pinned_buffer_seg_type_.push_back(seg_type); + auto t1 = getCurrentTimeInNano(); #ifdef USE_SUNRISE - if (seg_type == "VRAM") { - auto err = cudaSetDevice(start_idx + i); - CHECK_FAIL(err == cudaSuccess ? Status::OK() - : Status::InternalError( - "Failed to set Sunrise device " - "before registerLocalMemory")); - } + if (seg_type == "vram") { + auto err = cudaSetDevice(start_idx + i); + CHECK_FAIL(err == cudaSuccess + ? Status::OK() + : Status::InternalError("Failed to set Sunrise " + "device before " + "registerLocalMemory")); + } #endif - CHECK_FAIL(engine_->registerLocalMemory(pinned_buffer_list_[i], - total_buffer_size, options)); - auto t2 = getCurrentTimeInNano(); - - alloc_ns += (t1 - t0); - reg_ns += (t2 - t1); + // Reset options.type to UNSPEC so registerLocalMemory registers + // the buffer to ALL enabled transports (via + // getSupportedTransports(UNSPEC)), not just the one that + // allocated it. This is what makes multi-transport work: the + // buffer ends up registered to both SHM (which has shm_path set + // by the allocate call above) and TCP/NVLink (which don't need + // it), so requests targeting this buffer resolve regardless of + // which transport the engine picks. Only applies in multi-transport + // mode (--seg_type_mix set, no --xport_type); single-seg_type + // fallback keeps options.type as set above (UNSPEC or explicit). + if (xport_type.empty() && !seg_type_mix_.empty()) { + options.type = UNSPEC; + } + CHECK_FAIL( + engine_->registerLocalMemory(buf, total_buffer_size, options)); + auto t2 = getCurrentTimeInNano(); + + alloc_ns += (t1 - t0); + reg_ns += (t2 - t1); + total_bytes += total_buffer_size; + } } - LOG(INFO) << "Allocated " << total_buffer_size * num_buffers << " bytes " - << seg_type << " buffers in " << alloc_ns / 1e6 - << " ms, registered in " << reg_ns / 1e6 << " ms"; + LOG(INFO) << "Allocated " << total_bytes << " bytes across " + << pinned_buffer_list_.size() << " buffers (" << seg_types.size() + << " seg_types) in " << alloc_ns / 1e6 << " ms, registered in " + << reg_ns / 1e6 << " ms"; return 0; } @@ -199,8 +292,8 @@ TENTBenchRunner::TENTBenchRunner() { signal(SIGINT, signalHandlerV1); signal(SIGTERM, signalHandlerV1); engine_ = std::make_unique(loadConfig()); - transport_hint_ = TransportSelector::parseTransportType( - XferBenchConfig::tent_transport_hint); + transport_hint_ = parseTransportType(XferBenchConfig::tent_transport_hint); + intent_type_ = getIntentType(XferBenchConfig::tent_intent_type); allocateBuffers(); } @@ -337,7 +430,9 @@ int TENTBenchRunner::runInitiatorTasks( double TENTBenchRunner::runSingleTransfer(uint64_t local_addr, uint64_t target_addr, uint64_t block_size, - uint64_t batch_size, OpCode opcode) { + uint64_t batch_size, OpCode opcode, + uint64_t deadline_ns, + IntentType intent_type) { auto batch_id = engine_->allocateBatch(batch_size); std::vector requests; for (uint64_t i = 0; i < batch_size; ++i) { @@ -348,6 +443,10 @@ double TENTBenchRunner::runSingleTransfer(uint64_t local_addr, entry.target_id = handle_; entry.target_offset = target_addr + block_size * i; entry.transport_hint = transport_hint_; + entry.deadline_ns = deadline_ns; + entry.intent_type = intent_type == IntentType::INTENT_UNSPEC + ? intent_type_ + : intent_type; requests.emplace_back(entry); } XferBenchTimer timer; diff --git a/mooncake-transfer-engine/benchmark/tent_backend.h b/mooncake-transfer-engine/benchmark/tent_backend.h index 1648ad9269..8f021cbb21 100644 --- a/mooncake-transfer-engine/benchmark/tent_backend.h +++ b/mooncake-transfer-engine/benchmark/tent_backend.h @@ -34,6 +34,7 @@ #include "tent/transfer_engine.h" #include "tent/common/utils/random.h" #include "tent/common/utils/os.h" +#include "tent/runtime/topology.h" // LocationParser namespace mooncake { namespace tent { @@ -59,20 +60,70 @@ class TENTBenchRunner : public BenchRunner { uint64_t getLocalBufferBase(int thread_id, uint64_t block_size, uint64_t batch_size) const { - const size_t num_buffers = pinned_buffer_list_.size(); - return (uint64_t)pinned_buffer_list_[thread_id % num_buffers] + - block_size * batch_size * (thread_id / num_buffers); + if (seg_type_mix_.empty()) { + // Single seg_type mode: original behavior. + const size_t num_buffers = pinned_buffer_list_.size(); + return (uint64_t)pinned_buffer_list_[thread_id % num_buffers] + + block_size * batch_size * (thread_id / num_buffers); + } + // Mixed mode: pick this thread's seg_type by round-robin, then index + // into the subset of pinned_buffer_list_ matching that seg_type. + const std::string& seg_type = + seg_type_mix_[thread_id % seg_type_mix_.size()]; + size_t base = 0, count = 0; + for (size_t i = 0; i < pinned_buffer_seg_type_.size(); ++i) { + if (pinned_buffer_seg_type_[i] == seg_type) { + if (count == 0) base = i; + ++count; + } + } + if (count == 0) { + LOG(FATAL) << "No local buffer of seg_type " << seg_type + << " for thread " << thread_id; + } + size_t slot = (thread_id / seg_type_mix_.size()) % count; + return (uint64_t)pinned_buffer_list_[base + slot] + + block_size * batch_size * + ((thread_id / seg_type_mix_.size()) / count); } uint64_t getTargetBufferBase(int thread_id, uint64_t block_size, uint64_t batch_size) const { - return info_.buffers[thread_id % info_.buffers.size()].base + - block_size * batch_size * (thread_id / info_.buffers.size()); + if (seg_type_mix_.empty()) { + // Single seg_type mode: original behavior. + return info_.buffers[thread_id % info_.buffers.size()].base + + block_size * batch_size * (thread_id / info_.buffers.size()); + } + // Mixed mode: pick this thread's seg_type, then index into the + // subset of target segment buffers matching that seg_type. Use + // LocationParser to classify each buffer as host (cpu) vs device + // (cuda/rocm/supa/...) rather than hard-coding a vendor prefix. + const std::string& seg_type = + seg_type_mix_[thread_id % seg_type_mix_.size()]; + bool want_device = (seg_type == "vram"); + std::vector matches; + for (size_t i = 0; i < info_.buffers.size(); ++i) { + const auto& loc = info_.buffers[i].location; + bool is_device = + LocationParser(loc).type() != "cpu" && loc != kWildcardLocation; + if (is_device == want_device) { + matches.push_back(i); + } + } + if (matches.empty()) { + LOG(FATAL) << "No target buffer of seg_type " << seg_type + << " for thread " << thread_id; + } + size_t slot = (thread_id / seg_type_mix_.size()) % matches.size(); + return info_.buffers[matches[slot]].base + + block_size * batch_size * + ((thread_id / seg_type_mix_.size()) / matches.size()); } double runSingleTransfer(uint64_t local_addr, uint64_t target_addr, uint64_t block_size, uint64_t batch_size, - OpCode opcode); + OpCode opcode, uint64_t deadline_ns, + IntentType intent_type); private: int allocateBuffers(); @@ -84,9 +135,17 @@ class TENTBenchRunner : public BenchRunner { private: std::unique_ptr engine_; std::vector pinned_buffer_list_; + // Parallel to pinned_buffer_list_: the seg_type each buffer belongs to + // ("dram" or "vram"). Populated by allocateBuffers; used by + // getLocalBufferBase to pick the right buffer per thread. + std::vector pinned_buffer_seg_type_; + // Parsed --seg_type_mix (e.g. ["dram", "vram"]). Empty when --seg_type_mix + // is not set → single-seg_type mode (existing behavior). + std::vector seg_type_mix_; SegmentID handle_; SegmentInfo info_; TransportType transport_hint_{UNSPEC}; + IntentType intent_type_{IntentType::INTENT_UNSPEC}; std::vector> current_task_; std::vector threads_; @@ -99,4 +158,4 @@ class TENTBenchRunner : public BenchRunner { } // namespace tent } // namespace mooncake -#endif // TEV1_BACKEND_H \ No newline at end of file +#endif // TEV1_BACKEND_H diff --git a/mooncake-transfer-engine/benchmark/tests/qos_metrics_test.cpp b/mooncake-transfer-engine/benchmark/tests/qos_metrics_test.cpp new file mode 100644 index 0000000000..0c890dab92 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/tests/qos_metrics_test.cpp @@ -0,0 +1,222 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "qos_metrics_adapter.h" + +#include +#include +#include +#include + +#include + +#include "tent/thirdparty/nlohmann/json.h" +#include "workload_config.h" + +namespace mooncake { +namespace tent { +namespace { + +TEST(QosMetricsTest, SupportsBulkSamples) { + XferMetricStats stats; + stats.add(std::vector{3.0, 1.0, 2.0}); + + EXPECT_EQ(stats.count(), 3u); + EXPECT_DOUBLE_EQ(stats.min(), 1.0); + EXPECT_DOUBLE_EQ(stats.max(), 3.0); + EXPECT_DOUBLE_EQ(stats.avg(), 2.0); +} + +TEST(QosMetricsTest, ParsesClassContract) { + std::vector classes; + std::string error; + ASSERT_TRUE(parseQosClasses("fg:4:1000:2:12.5,bg:12:0:1", &classes, &error)) + << error; + ASSERT_EQ(classes.size(), 2); + EXPECT_EQ(classes[0].name, "fg"); + EXPECT_EQ(classes[0].threads, 4); + EXPECT_EQ(classes[0].slo_us, 1000); + EXPECT_DOUBLE_EQ(classes[0].weight, 2.0); + ASSERT_TRUE(classes[0].isolated_throughput_gbps); + EXPECT_DOUBLE_EQ(*classes[0].isolated_throughput_gbps, 12.5); + EXPECT_FALSE(classes[1].isolated_throughput_gbps); + EXPECT_TRUE(validateQosClasses(classes, 16, &error)); + EXPECT_EQ(qosClassForThread(classes, 0), 0); + EXPECT_EQ(qosClassForThread(classes, 3), 0); + EXPECT_EQ(qosClassForThread(classes, 4), 1); + EXPECT_EQ(qosClassForThread(classes, 15), 1); +} + +TEST(QosMetricsTest, ParsesJsonClassContract) { + std::vector classes; + std::string error; + ASSERT_TRUE(parseQosClassesJson( + R"json([ + {"name":"fg","threads":4,"slo_us":1000,"weight":2,"isolated_gbps":12.5}, + {"name":"bg","threads":12,"slo_us":0,"weight":1} + ])json", + &classes, &error)) + << error; + ASSERT_EQ(classes.size(), 2); + EXPECT_EQ(classes[0].name, "fg"); + EXPECT_EQ(classes[0].threads, 4); + EXPECT_EQ(classes[0].slo_us, 1000); + EXPECT_DOUBLE_EQ(classes[0].weight, 2.0); + ASSERT_TRUE(classes[0].isolated_throughput_gbps); + EXPECT_DOUBLE_EQ(*classes[0].isolated_throughput_gbps, 12.5); + EXPECT_FALSE(classes[1].isolated_throughput_gbps); +} + +TEST(QosMetricsTest, RejectsAmbiguousOrInvalidContracts) { + std::vector classes; + std::string error; + EXPECT_FALSE(parseQosClasses("fg:1:100:1,fg:1:0:1", &classes, &error)); + EXPECT_FALSE(parseQosClasses("fg:0:100:1", &classes, &error)); + EXPECT_FALSE(parseQosClasses("fg:1::1", &classes, &error)); + EXPECT_FALSE(parseQosClasses("fg:1:-1:1", &classes, &error)); + EXPECT_FALSE(parseQosClasses("fg:1:100:0", &classes, &error)); + EXPECT_FALSE(parseQosClasses("fg:1:100:1:0", &classes, &error)); + ASSERT_TRUE(parseQosClasses("fg:1:100:1", &classes, &error)); + EXPECT_FALSE(validateQosClasses(classes, 2, &error)); +} + +TEST(QosMetricsTest, ParsesMixedWorkloadClasses) { + std::vector classes; + std::string error; + ASSERT_TRUE(parseWorkloadClassesJson( + R"json([ + {"name":"foreground","threads":2,"block_size":4096,"batch_size":1, + "intent_type":"foreground_get","deadline_us":250, + "slo_us":300,"weight":4}, + {"name":"migration","threads":6,"block_size":4194304,"batch_size":2, + "intent_type":"migration","slo_us":0,"weight":1} + ])json", + &classes, &error)) + << error; + ASSERT_EQ(classes.size(), 2); + EXPECT_EQ(classes[0].qos.name, "foreground"); + EXPECT_EQ(classes[0].block_size, 4096u); + EXPECT_EQ(classes[0].deadline_us, 250u); + EXPECT_EQ(classes[0].intent_type, IntentType::FOREGROUND_GET); + EXPECT_EQ(classes[1].intent_type, IntentType::MIGRATION); + EXPECT_TRUE(validateWorkloadClasses(classes, 8, 1UL << 30, &error)) + << error; + const auto qos_classes = qosClassesFromWorkload(classes); + ASSERT_EQ(qos_classes.size(), 2); + EXPECT_EQ(qos_classes[1].threads, 6); +} + +TEST(QosMetricsTest, RejectsInvalidMixedWorkloadClasses) { + std::vector classes; + std::string error; + EXPECT_FALSE(parseWorkloadClassesJson( + R"json([{"name":"fg","threads":1,"block_size":0,"batch_size":1, + "intent_type":"foreground_get","slo_us":100,"weight":1}])json", + &classes, &error)); + EXPECT_FALSE(parseWorkloadClassesJson( + R"json([{"name":"fg","threads":1,"block_size":4096,"batch_size":1, + "intent_type":"unknown","slo_us":100,"weight":1}])json", + &classes, &error)); +} + +TEST(QosMetricsTest, CalculatesSloFairnessIsolationAndUtilization) { + std::vector classes = { + {"foreground", 1, 100, 2.0, 0.006}, + {"background", 1, 0, 1.0, 0.004}, + }; + std::vector stats(2); + stats[0].total_duration.add(1000.0); + stats[0].transfer_duration.add(50.0); + stats[0].transfer_duration.add(100.0); + stats[0].transfer_duration.add(150.0); + stats[1].total_duration.add(1000.0); + stats[1].transfer_duration.add(100.0); + stats[1].transfer_duration.add(100.0); + + const auto report = + calculateQosMetricsFromBenchStats(1000, 1, 2, classes, &stats, 0.01); + ASSERT_EQ(report.classes.size(), 2); + EXPECT_NEAR(report.aggregate_throughput_gbps, 0.005, 1e-12); + EXPECT_NEAR(report.weighted_goodput_gbps, 0.006, 1e-12); + EXPECT_NEAR(report.jain_fairness, 0.98, 1e-12); + ASSERT_TRUE(report.max_isolation_leakage); + EXPECT_NEAR(*report.max_isolation_leakage, 0.5, 1e-12); + ASSERT_TRUE(report.total_utilization); + EXPECT_NEAR(*report.total_utilization, 0.5, 1e-12); + ASSERT_TRUE(report.link_capacity_gbps); + EXPECT_NEAR(*report.link_capacity_gbps, 0.01, 1e-12); + + const auto& foreground = report.classes[0]; + EXPECT_EQ(foreground.operations, 3); + EXPECT_NEAR(foreground.p99_us, 149.0, 1e-12); + ASSERT_TRUE(foreground.slo_attainment); + EXPECT_NEAR(*foreground.slo_attainment, 2.0 / 3.0, 1e-12); + EXPECT_NEAR(foreground.goodput_gbps, 0.002, 1e-12); + EXPECT_NEAR(foreground.weighted_goodput_gbps, 0.004, 1e-12); + ASSERT_TRUE(foreground.isolated_throughput_gbps); + EXPECT_NEAR(*foreground.isolated_throughput_gbps, 0.006, 1e-12); + + EXPECT_FALSE(report.classes[1].slo_attainment); +} + +TEST(QosMetricsTest, UsesPerClassTransferredBytes) { + std::vector classes = { + {"foreground", 1, 0, 1.0, std::nullopt}, + {"migration", 1, 0, 1.0, std::nullopt}, + }; + std::vector stats(2); + for (auto& class_stats : stats) { + class_stats.total_duration.add(1000.0); + class_stats.transfer_duration.add(100.0); + } + + const auto report = calculateQosMetricsFromBenchStats( + 0, 0, 2, classes, &stats, 0.0, {4096, 4UL << 20}); + ASSERT_EQ(report.classes.size(), 2); + EXPECT_EQ(report.classes[0].transferred_bytes, 4096u); + EXPECT_EQ(report.classes[1].transferred_bytes, 4UL << 20); + EXPECT_NEAR(report.classes[0].throughput_gbps, 0.004096, 1e-12); + EXPECT_NEAR(report.classes[1].throughput_gbps, 4.194304, 1e-12); +} + +TEST(QosMetricsTest, UsesNullForUnavailableJsonMetrics) { + std::vector classes = { + {"best_effort", 1, 0, 1.0, std::nullopt}}; + std::vector stats(1); + stats[0].total_duration.add(1000.0); + stats[0].transfer_duration.add(100.0); + const auto report = + calculateQosMetricsFromBenchStats(1000, 1, 1, classes, &stats, 0.0); + + const std::string path = "tebench_qos_metrics_test.jsonl"; + std::remove(path.c_str()); + std::string error; + ASSERT_TRUE(appendQosMetricsJsonl(path, report, &error)) << error; + + std::ifstream input(path); + nlohmann::json record; + ASSERT_NO_THROW(input >> record); + EXPECT_EQ(record["schema_version"], 1); + EXPECT_TRUE(record["total_utilization"].is_null()); + EXPECT_TRUE(record["link_capacity_gbps"].is_null()); + EXPECT_TRUE(record["max_isolation_leakage"].is_null()); + EXPECT_TRUE(record["classes"][0]["slo_attainment"].is_null()); + EXPECT_TRUE(record["classes"][0]["isolation_leakage"].is_null()); + EXPECT_TRUE(record["classes"][0]["isolated_throughput_gbps"].is_null()); + std::remove(path.c_str()); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/benchmark/utils.cpp b/mooncake-transfer-engine/benchmark/utils.cpp index d907a2bd68..cbc839d3ac 100644 --- a/mooncake-transfer-engine/benchmark/utils.cpp +++ b/mooncake-transfer-engine/benchmark/utils.cpp @@ -20,6 +20,15 @@ DEFINE_string(seg_name, "", "Memory segment name for the local side"); DEFINE_string(seg_type, "DRAM", "Memory segment type for the target side: DRAM|VRAM"); +DEFINE_string( + seg_type_mix, "", + "Comma-separated segment types for mixed DRAM+VRAM runs, e.g. " + "\"dram,vram\". When set, target registers buffers of each listed type " + "in one segment, and initiator threads round-robin across them so a " + "single tebench process drives traffic over multiple memory types (and " + "thus multiple transports — SHM for DRAM, NVLink for VRAM) " + "concurrently. Empty falls back to --seg_type (single type, existing " + "behavior). Requires transports to be enabled via MC_TENT_CONF."); DEFINE_string(target_seg_name, "", "Memory segment name for the target side"); DEFINE_string(op_type, "read", "Operation type to benchmark: read|write|mix"); DEFINE_bool(check_consistency, false, @@ -35,6 +44,32 @@ DEFINE_int32(start_num_threads, 1, "Start number of concurrent worker threads."); DEFINE_int32(max_num_threads, 1, "Maximum number of concurrent worker threads."); +DEFINE_string( + qos_classes, "", + "QoS classes as name:threads:slo_us:weight[:isolated_gbps],...; " + "enables per-class QoS metrics and requires a fixed thread count."); +DEFINE_string(qos_classes_json, "", + "QoS classes as a JSON array of objects with name, threads, " + "slo_us, weight, and optional isolated_gbps fields."); +DEFINE_string( + workload_classes_json, "", + "Mixed traffic classes as a JSON array with name, threads, block_size, " + "batch_size, intent_type, slo_us, weight, and optional deadline_us and " + "isolated_gbps fields. Overrides the block/batch sweep."); +DEFINE_double(qos_link_capacity_gbps, 0.0, + "Link capacity in GB/s for total utilization (0 reports N/A)."); +DEFINE_string(qos_output_jsonl, "", + "Append versioned QoS metric records to this JSONL file."); +DEFINE_uint64(deadline_us, 0, + "tent only: relative per-transfer deadline in microseconds for " + "tight worker threads (0 disables deadline tagging); cannot be " + "combined with --workload_classes_json."); +DEFINE_int32(deadline_tight_threads, 0, + "tent only: workers [0, N) that carry --deadline_us; remaining " + "workers have no deadline; cannot be combined with " + "--workload_classes_json."); +DEFINE_bool(deadline_bw_arbitration, false, + "tent only: enable deadline-aware RDMA bandwidth arbitration."); DEFINE_int32(local_gpu_id, 0, "Local GPU ID to be used, -1 for all GPUs"); DEFINE_int32(target_gpu_id, 0, "Target GPU ID to be used, -1 for all GPUs"); DEFINE_string(metadata_type, "p2p", @@ -53,11 +88,16 @@ DEFINE_string( tent_transport_hint, "unspec", "tent only: per-request transport_hint. " "unspec|rdma|tcp|shm|nvlink|gds|io_uring|mnnvl|ascend|sunrise_link"); +DEFINE_string(tent_intent_type, "unspec", + "tent only: intent_type attached to every benchmark request. " + "unspec|foreground_get|background_prefetch|migration|checkpoint|" + "weight_loading|staging_internal"); namespace mooncake { namespace tent { std::string XferBenchConfig::seg_name; std::string XferBenchConfig::seg_type; +std::string XferBenchConfig::seg_type_mix; std::string XferBenchConfig::target_seg_name; std::string XferBenchConfig::op_type; bool XferBenchConfig::check_consistency = false; @@ -70,6 +110,14 @@ size_t XferBenchConfig::max_batch_size = 0; int XferBenchConfig::duration = 0; int XferBenchConfig::max_num_threads = 0; int XferBenchConfig::start_num_threads = 0; +std::string XferBenchConfig::qos_classes; +std::string XferBenchConfig::qos_classes_json; +std::string XferBenchConfig::workload_classes_json; +double XferBenchConfig::qos_link_capacity_gbps = 0.0; +std::string XferBenchConfig::qos_output_jsonl; +uint64_t XferBenchConfig::deadline_us = 0; +int XferBenchConfig::deadline_tight_threads = 0; +bool XferBenchConfig::deadline_bw_arbitration = false; std::string XferBenchConfig::metadata_type; std::string XferBenchConfig::metadata_url_list; @@ -78,12 +126,14 @@ std::string XferBenchConfig::xport_type; std::string XferBenchConfig::backend; bool XferBenchConfig::notifi = false; std::string XferBenchConfig::tent_transport_hint; +std::string XferBenchConfig::tent_intent_type; int XferBenchConfig::local_gpu_id = 0; int XferBenchConfig::target_gpu_id = 0; void XferBenchConfig::loadFromFlags() { seg_type = FLAGS_seg_type; + seg_type_mix = FLAGS_seg_type_mix; seg_name = FLAGS_seg_name; target_seg_name = FLAGS_target_seg_name; op_type = FLAGS_op_type; @@ -96,6 +146,14 @@ void XferBenchConfig::loadFromFlags() { max_batch_size = FLAGS_max_batch_size; start_num_threads = FLAGS_start_num_threads; max_num_threads = FLAGS_max_num_threads; + qos_classes = FLAGS_qos_classes; + qos_classes_json = FLAGS_qos_classes_json; + workload_classes_json = FLAGS_workload_classes_json; + qos_link_capacity_gbps = FLAGS_qos_link_capacity_gbps; + qos_output_jsonl = FLAGS_qos_output_jsonl; + deadline_us = FLAGS_deadline_us; + deadline_tight_threads = FLAGS_deadline_tight_threads; + deadline_bw_arbitration = FLAGS_deadline_bw_arbitration; duration = FLAGS_duration; metadata_type = FLAGS_metadata_type; @@ -106,6 +164,7 @@ void XferBenchConfig::loadFromFlags() { backend = FLAGS_backend; notifi = FLAGS_notifi; tent_transport_hint = FLAGS_tent_transport_hint; + tent_intent_type = FLAGS_tent_intent_type; local_gpu_id = FLAGS_local_gpu_id; target_gpu_id = FLAGS_target_gpu_id; @@ -168,5 +227,20 @@ void printStats(size_t block_size, size_t batch_size, XferBenchStats& stats, // clang-format on } +void printDeadlineGroupStats(const char* group, size_t block_size, + size_t batch_size, XferBenchStats& stats, + int num_threads, uint64_t deadline_us) { + if (num_threads <= 0 || stats.transfer_duration.count() == 0) return; + const double duration_s = stats.total_duration.avg() / 1e6; + const double bytes = static_cast(block_size) * batch_size * + stats.transfer_duration.count(); + const double throughput_gbs = bytes / 1e9 / duration_s; + std::cout << " [deadline-" << group << "] threads=" << num_threads; + if (deadline_us != 0) std::cout << " deadline_us=" << deadline_us; + std::cout << " operations=" << stats.transfer_duration.count() + << " throughput=" << std::fixed << std::setprecision(6) + << throughput_gbs << " GB/s" << std::endl; +} + } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/benchmark/utils.h b/mooncake-transfer-engine/benchmark/utils.h index 6c862cd6c5..9d6cab1e6c 100644 --- a/mooncake-transfer-engine/benchmark/utils.h +++ b/mooncake-transfer-engine/benchmark/utils.h @@ -56,6 +56,9 @@ struct XferBenchConfig { static std::string seg_name; static std::string seg_type; + // Comma-separated segment types for mixed DRAM+VRAM runs, e.g. + // "dram,vram". Empty falls back to --seg_type (single type). + static std::string seg_type_mix; static std::string target_seg_name; static std::string op_type; static bool check_consistency; @@ -68,6 +71,14 @@ struct XferBenchConfig { static int duration; static int max_num_threads; static int start_num_threads; + static std::string qos_classes; + static std::string qos_classes_json; + static std::string workload_classes_json; + static double qos_link_capacity_gbps; + static std::string qos_output_jsonl; + static uint64_t deadline_us; + static int deadline_tight_threads; + static bool deadline_bw_arbitration; static std::string metadata_type; static std::string metadata_url_list; @@ -76,6 +87,7 @@ struct XferBenchConfig { static std::string backend; static bool notifi; static std::string tent_transport_hint; + static std::string tent_intent_type; static int local_gpu_id; static int target_gpu_id; @@ -107,8 +119,20 @@ struct XferMetricStats { double p999() { return percentile(99.9); } + double fractionAtOrBelow(double threshold) const { + if (samples.empty()) return 0.0; + const auto count = std::count_if( + samples.begin(), samples.end(), + [threshold](double value) { return value <= threshold; }); + return static_cast(count) / samples.size(); + } + void add(double value) { samples.push_back(value); } + void add(const std::vector& values) { + samples.insert(samples.end(), values.begin(), values.end()); + } + void clear() { samples.clear(); } size_t count() { return samples.size(); } @@ -153,6 +177,10 @@ void printStatsHeader(); void printStats(size_t block_size, size_t batch_size, XferBenchStats& stats, int num_threads); +void printDeadlineGroupStats(const char* group, size_t block_size, + size_t batch_size, XferBenchStats& stats, + int num_threads, uint64_t deadline_us); + #if defined(USE_CUDA) || defined(USE_SUNRISE) static inline bool isCudaMemory(void* ptr) { cudaPointerAttributes attr; diff --git a/mooncake-transfer-engine/benchmark/workload_config.cpp b/mooncake-transfer-engine/benchmark/workload_config.cpp new file mode 100644 index 0000000000..9f10d162be --- /dev/null +++ b/mooncake-transfer-engine/benchmark/workload_config.cpp @@ -0,0 +1,149 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "workload_config.h" + +#include +#include +#include + +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake { +namespace tent { +namespace { + +bool parseIntentType(const std::string& value, IntentType* intent_type) { + static const std::unordered_map kIntentTypes = { + {"unspec", IntentType::INTENT_UNSPEC}, + {"intent_unspec", IntentType::INTENT_UNSPEC}, + {"foreground_get", IntentType::FOREGROUND_GET}, + {"background_prefetch", IntentType::BACKGROUND_PREFETCH}, + {"migration", IntentType::MIGRATION}, + {"checkpoint", IntentType::CHECKPOINT}, + {"weight_loading", IntentType::WEIGHT_LOADING}, + {"staging_internal", IntentType::STAGING_INTERNAL}, + }; + const auto it = kIntentTypes.find(value); + if (it == kIntentTypes.end()) return false; + *intent_type = it->second; + return true; +} + +bool parsePositiveSize(const nlohmann::json& node, const char* field, + const std::string& path, size_t* result, + std::string* error) { + if (!node.contains(field) || !node[field].is_number_unsigned()) { + *error = path + "." + field + " must be an unsigned integer"; + return false; + } + const uint64_t value = node[field].get(); + if (value == 0 || value > std::numeric_limits::max()) { + *error = path + "." + field + " must be positive and fit in size_t"; + return false; + } + *result = static_cast(value); + return true; +} + +} // namespace + +bool parseWorkloadClassesJson(const std::string& spec, + std::vector* classes, + std::string* error) { + classes->clear(); + try { + const auto root = nlohmann::json::parse(spec); + if (!root.is_array() || root.empty()) { + *error = "workload_classes_json must be a non-empty array"; + return false; + } + std::vector qos_classes; + if (!parseQosClassesJson(spec, &qos_classes, error)) return false; + + for (size_t i = 0; i < root.size(); ++i) { + const auto& node = root[i]; + const std::string path = + "workload_classes_json[" + std::to_string(i) + "]"; + + WorkloadClassConfig config; + config.qos = qos_classes[i]; + if (!parsePositiveSize(node, "block_size", path, &config.block_size, + error) || + !parsePositiveSize(node, "batch_size", path, &config.batch_size, + error)) { + return false; + } + if (config.block_size > + std::numeric_limits::max() / config.batch_size) { + *error = path + " block_size * batch_size overflows size_t"; + return false; + } + + if (node.contains("deadline_us")) { + if (!node["deadline_us"].is_number_unsigned()) { + *error = path + ".deadline_us must be an unsigned integer"; + return false; + } + config.deadline_us = node["deadline_us"].get(); + } + if (!node.contains("intent_type") || + !node["intent_type"].is_string() || + !parseIntentType(node["intent_type"].get(), + &config.intent_type)) { + *error = path + ".intent_type is invalid"; + return false; + } + classes->push_back(std::move(config)); + } + return true; + } catch (const std::exception& e) { + *error = + std::string("failed to parse workload_classes_json: ") + e.what(); + return false; + } +} + +bool validateWorkloadClasses(const std::vector& classes, + int num_threads, size_t total_buffer_size, + std::string* error) { + size_t configured_threads = 0; + size_t max_bytes_per_thread = 0; + for (const auto& config : classes) { + configured_threads += static_cast(config.qos.threads); + const size_t bytes = config.block_size * config.batch_size; + max_bytes_per_thread = std::max(max_bytes_per_thread, bytes); + } + if (configured_threads != static_cast(num_threads)) { + *error = "workload class thread count must equal benchmark threads"; + return false; + } + if (max_bytes_per_thread > total_buffer_size || + configured_threads > total_buffer_size / max_bytes_per_thread) { + *error = "mixed workload address stride exceeds total_buffer_size"; + return false; + } + return true; +} + +std::vector qosClassesFromWorkload( + const std::vector& classes) { + std::vector qos_classes; + qos_classes.reserve(classes.size()); + for (const auto& config : classes) qos_classes.push_back(config.qos); + return qos_classes; +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/benchmark/workload_config.h b/mooncake-transfer-engine/benchmark/workload_config.h new file mode 100644 index 0000000000..6821f5e657 --- /dev/null +++ b/mooncake-transfer-engine/benchmark/workload_config.h @@ -0,0 +1,51 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TEBENCH_WORKLOAD_CONFIG_H +#define TEBENCH_WORKLOAD_CONFIG_H + +#include +#include +#include +#include + +#include "tent/common/qos_metrics.h" +#include "tent/common/types.h" + +namespace mooncake { +namespace tent { + +struct WorkloadClassConfig { + QosClassConfig qos; + size_t block_size = 0; + size_t batch_size = 0; + uint64_t deadline_us = 0; + IntentType intent_type = IntentType::INTENT_UNSPEC; +}; + +bool parseWorkloadClassesJson(const std::string& spec, + std::vector* classes, + std::string* error); + +bool validateWorkloadClasses(const std::vector& classes, + int num_threads, size_t total_buffer_size, + std::string* error); + +std::vector qosClassesFromWorkload( + const std::vector& classes); + +} // namespace tent +} // namespace mooncake + +#endif // TEBENCH_WORKLOAD_CONFIG_H diff --git a/mooncake-transfer-engine/example/CMakeLists.txt b/mooncake-transfer-engine/example/CMakeLists.txt index 4adfcdba93..cdd681ad67 100644 --- a/mooncake-transfer-engine/example/CMakeLists.txt +++ b/mooncake-transfer-engine/example/CMakeLists.txt @@ -1,5 +1,8 @@ set(WORKSPACE "${CMAKE_CURRENT_SOURCE_DIR}") +include(CheckPIESupported) +check_pie_supported(LANGUAGES CXX) + if(USE_HIP) file(GLOB EXAMPLE_SOURCES "*.cpp") hipify_files(EXAMPLE_SOURCES) @@ -10,6 +13,8 @@ if(USE_HIP) endif() add_executable(transfer_engine_bench ${WORKSPACE}/transfer_engine_bench.cpp) +set_target_properties(transfer_engine_bench PROPERTIES POSITION_INDEPENDENT_CODE + ON) target_link_libraries(transfer_engine_bench PUBLIC transfer_engine) if(USE_TENT) target_link_libraries(transfer_engine_bench PUBLIC tent_link_group) @@ -106,3 +111,53 @@ if(USE_CUDA AND BUILD_DEVICE_TRANSPORT_EXAMPLE) CUDA_ARCHITECTURES "80;90") endif() endif() + +add_executable(show_link ${WORKSPACE}/show_link.cpp) +target_link_libraries(show_link PUBLIC transfer_engine gflags::gflags + glog::glog) + +# NCCL DeviceTransport example. This is a manual two-rank validation because +# communicator initialization and symmetric-window registration are collective. +option(BUILD_NCCL_DEVICE_TRANSPORT_EXAMPLE + "Build the two-rank NCCL DeviceTransport LSA/GIN example" + OFF) +if(USE_NCCL_DEVICE AND BUILD_NCCL_DEVICE_TRANSPORT_EXAMPLE) + enable_language(CUDA) + add_executable(nccl_device_transport_example + ${WORKSPACE}/nccl_device_transport_example.cu) + target_include_directories(nccl_device_transport_example PRIVATE + ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include) + target_link_libraries(nccl_device_transport_example PRIVATE + transfer_engine NCCL::nccl gflags::gflags glog::glog) + target_compile_options(nccl_device_transport_example PRIVATE + $<$:--expt-relaxed-constexpr>) + set_target_properties(nccl_device_transport_example PROPERTIES + CUDA_STANDARD 20 CUDA_STANDARD_REQUIRED ON CUDA_EXTENSIONS OFF) + + if(NOT CMAKE_CUDA_ARCHITECTURES) + if(TORCH_CUDA_ARCH_LIST) + set(_nccl_cuda_arch_list "") + foreach(_arch IN LISTS TORCH_CUDA_ARCH_LIST) + string(REPLACE "." "" _arch_clean "${_arch}") + list(APPEND _nccl_cuda_arch_list "${_arch_clean}") + endforeach() + set_target_properties(nccl_device_transport_example PROPERTIES + CUDA_ARCHITECTURES "${_nccl_cuda_arch_list}") + else() + set_target_properties(nccl_device_transport_example PROPERTIES + CUDA_ARCHITECTURES "80;90") + endif() + endif() +endif() + +option(BUILD_NCCL_HOST_TRANSPORT_EXAMPLE + "Build the two-GPU NCCL host RMA Transfer Engine example" + OFF) +if(USE_NCCL_HOST AND BUILD_NCCL_HOST_TRANSPORT_EXAMPLE) + add_executable(nccl_host_transport_example + ${WORKSPACE}/nccl_host_transport_example.cpp) + target_include_directories(nccl_host_transport_example PRIVATE + ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include) + target_link_libraries(nccl_host_transport_example PRIVATE + transfer_engine NCCL::nccl glog::glog gflags::gflags) +endif() diff --git a/mooncake-transfer-engine/example/batch_register_bench.py b/mooncake-transfer-engine/example/batch_register_bench.py index 3b0cfca1e1..384a1c3a82 100644 --- a/mooncake-transfer-engine/example/batch_register_bench.py +++ b/mooncake-transfer-engine/example/batch_register_bench.py @@ -25,11 +25,8 @@ import ctypes import ctypes.util import json -import os -import random import signal import statistics -import sys import time @@ -224,7 +221,6 @@ def run_initiator(args): """Run as initiator: pull data from random blocks.""" from mooncake.engine import TransferEngine - block_bytes = int(args.block_size_gb * 1024 * 1024 * 1024) transfer_bytes = int(args.transfer_size_mb * 1024 * 1024) print(f"=== Initiator Node ===") diff --git a/mooncake-transfer-engine/example/http-metadata-server-python/bootstrap_server.py b/mooncake-transfer-engine/example/http-metadata-server-python/bootstrap_server.py index 4032d40906..d82f714ba2 100644 --- a/mooncake-transfer-engine/example/http-metadata-server-python/bootstrap_server.py +++ b/mooncake-transfer-engine/example/http-metadata-server-python/bootstrap_server.py @@ -1,5 +1,4 @@ from enum import Enum -from time import sleep from aiohttp import web import threading import asyncio diff --git a/mooncake-transfer-engine/example/http-metadata-server/go.mod b/mooncake-transfer-engine/example/http-metadata-server/go.mod index 84a942aa88..3641bc48c0 100644 --- a/mooncake-transfer-engine/example/http-metadata-server/go.mod +++ b/mooncake-transfer-engine/example/http-metadata-server/go.mod @@ -1,8 +1,6 @@ module github.com/kvcache-ai/Mooncake/mooncake-transfer-engine/example/http-metadata-server -go 1.23.0 - -toolchain go1.23.8 +go 1.25.0 require github.com/gin-gonic/gin v1.10.0 @@ -27,10 +25,10 @@ require ( github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.2.12 // indirect golang.org/x/arch v0.8.0 // indirect - golang.org/x/crypto v0.36.0 // indirect - golang.org/x/net v0.38.0 // indirect - golang.org/x/sys v0.31.0 // indirect - golang.org/x/text v0.23.0 // indirect + golang.org/x/crypto v0.52.0 // indirect + golang.org/x/net v0.55.0 // indirect + golang.org/x/sys v0.45.0 // indirect + golang.org/x/text v0.37.0 // indirect google.golang.org/protobuf v1.34.1 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/mooncake-transfer-engine/example/kvcache_prefix_bench.py b/mooncake-transfer-engine/example/kvcache_prefix_bench.py index 71d675060d..dbf52f44a5 100644 --- a/mooncake-transfer-engine/example/kvcache_prefix_bench.py +++ b/mooncake-transfer-engine/example/kvcache_prefix_bench.py @@ -27,10 +27,8 @@ import ctypes import ctypes.util import json -import os import signal import statistics -import sys import time from concurrent.futures import ThreadPoolExecutor, as_completed diff --git a/mooncake-transfer-engine/example/nccl_device_transport_example.cu b/mooncake-transfer-engine/example/nccl_device_transport_example.cu new file mode 100644 index 0000000000..30d3726a6c --- /dev/null +++ b/mooncake-transfer-engine/example/nccl_device_transport_example.cu @@ -0,0 +1,338 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Two-rank NCCL DeviceTransport correctness example. +// +// Launch one process per rank with the same unique --run_id and a bootstrap +// directory visible to both ranks. Use --force_gin to exercise GIN even when +// both ranks are in one LSA team. + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "transfer_engine.h" +#include "transport/device/cuda/cuda_ops.cuh" +#include "transport/device/nccl_device.cuh" + +DEFINE_int32(rank, -1, "Rank of this process (0 or 1)"); +DEFINE_int32(world_size, 2, "Number of ranks; this example requires 2"); +DEFINE_int32(gpu_id, -1, + "CUDA device (default: rank modulo visible device count)"); +DEFINE_int32(data_bytes, 4096, "Number of bytes transferred from rank 0"); +DEFINE_int32(timeout_seconds, 120, "Bootstrap file wait timeout"); +DEFINE_string(bootstrap_dir, "/tmp", + "Directory visible to both ranks for bootstrap"); +DEFINE_string(run_id, "", + "Unique identifier shared by ranks for this invocation"); +DEFINE_bool(force_gin, false, "Use GIN even when the peer is LSA reachable"); +DEFINE_bool(disable_gin, false, + "Create an LSA-only device communicator without GIN resources"); +DEFINE_bool(require_lsa_multimem, false, + "Require LSA multimem resources and use them for LSA barriers"); + +namespace { + +constexpr size_t kSignalAlignment = 256; + +void checkCuda(cudaError_t result, const char* operation) { + CHECK_EQ(result, cudaSuccess) + << operation << ": " << cudaGetErrorString(result); +} + +size_t alignUp(size_t value, size_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +std::string bootstrapPath() { + CHECK(!FLAGS_run_id.empty()) + << "--run_id must be unique for every invocation"; + for (unsigned char c : FLAGS_run_id) { + CHECK(std::isalnum(c) || c == '-' || c == '_' || c == '.') + << "--run_id may contain only letters, digits, '.', '-', and '_'"; + } + return FLAGS_bootstrap_dir + "/mooncake_nccl_device_" + FLAGS_run_id + + ".bin"; +} + +void writeUniqueId(const std::string& path, const std::vector& id) { + const std::string temporary = path + ".tmp"; + std::ofstream output(temporary, std::ios::binary | std::ios::trunc); + CHECK(output) << "Cannot open " << temporary; + const uint32_t words = static_cast(id.size()); + output.write(reinterpret_cast(&words), sizeof(words)); + output.write(reinterpret_cast(id.data()), + id.size() * sizeof(int32_t)); + output.close(); + CHECK_EQ(std::rename(temporary.c_str(), path.c_str()), 0) + << "Cannot publish " << path; +} + +std::vector readUniqueId(const std::string& path) { + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::seconds(FLAGS_timeout_seconds); + while (std::chrono::steady_clock::now() < deadline) { + std::ifstream input(path, std::ios::binary); + if (input) { + uint32_t words = 0; + input.read(reinterpret_cast(&words), sizeof(words)); + if (input && words > 0 && words <= 256) { + std::vector id(words); + input.read(reinterpret_cast(id.data()), + id.size() * sizeof(int32_t)); + if (input) return id; + } + } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + LOG(FATAL) << "Timed out waiting for " << path; + return {}; +} + +void publishMarker(const std::string& path) { + const std::string temporary = path + ".tmp"; + std::ofstream output(temporary, std::ios::trunc); + CHECK(output) << "Cannot open " << temporary; + output << "done\n"; + output.close(); + CHECK_EQ(std::rename(temporary.c_str(), path.c_str()), 0) + << "Cannot publish " << path; +} + +void waitForMarker(const std::string& path) { + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::seconds(FLAGS_timeout_seconds); + while (std::chrono::steady_clock::now() < deadline) { + std::ifstream input(path); + std::string marker; + if (input >> marker && marker == "done") return; + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + LOG(FATAL) << "Timed out waiting for " << path; +} + +__global__ void senderKernel(mooncake::device::NcclDeviceContext ctx, + char* local_buffer, size_t bytes, + size_t data_signal_offset, + size_t ack_signal_offset, bool force_gin, + int* result) { + using namespace mooncake::device; + if (blockIdx.x != 0 || threadIdx.x != 0) return; + + constexpr int peer = 1; + mc_nccl_lsa_barrier(ctx, 0); + + const NcclDeviceRoute route = mc_nccl_route(ctx, peer); + const bool use_gin = force_gin || route == NcclDeviceRoute::kGin; + auto* data_signal = + reinterpret_cast(local_buffer + data_signal_offset); + auto* ack_signal = + reinterpret_cast(local_buffer + ack_signal_offset); + + if (use_gin) { + if (!mc_nccl_gin_available(ctx)) { + *result = -1; + return; + } + mc_nccl_put_with_signal(ctx, 0, peer, 1, local_buffer, local_buffer, + static_cast(bytes), data_signal, 1, + 0); + mc_nccl_flush(ctx, 0, 0); + } else if (route == NcclDeviceRoute::kLsa) { + char* peer_buffer = + static_cast(mc_nccl_peer_ptr(ctx, peer, local_buffer)); + for (size_t i = 0; i < bytes; ++i) peer_buffer[i] = local_buffer[i]; + mc_nccl_signal_add(ctx, peer, 0, 1, data_signal, 1, 0); + } else { + *result = -2; + return; + } + + mc_nccl_wait_signal(ctx, 0, ack_signal, 1, 0); + mc_nccl_lsa_barrier(ctx, 0); + *result = 0; +} + +__global__ void receiverKernel(mooncake::device::NcclDeviceContext ctx, + char* local_buffer, size_t bytes, + size_t data_signal_offset, + size_t ack_signal_offset, int* result) { + using namespace mooncake::device; + if (blockIdx.x != 0 || threadIdx.x != 0) return; + + constexpr int peer = 0; + mc_nccl_lsa_barrier(ctx, 0); + + auto* data_signal = + reinterpret_cast(local_buffer + data_signal_offset); + auto* ack_signal = + reinterpret_cast(local_buffer + ack_signal_offset); + mc_nccl_wait_signal(ctx, 0, data_signal, 1, 0); + + int mismatches = 0; + for (size_t i = 0; i < bytes; ++i) { + if (local_buffer[i] != static_cast(i & 0xff)) ++mismatches; + } + *result = mismatches; + + mc_nccl_signal_add(ctx, peer, 0, 1, ack_signal, 1, 0); + if (mc_nccl_route(ctx, peer) == NcclDeviceRoute::kGin) + mc_nccl_flush(ctx, 0, 0); + + mc_nccl_lsa_barrier(ctx, 0); +} + +} // namespace + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + + CHECK_EQ(FLAGS_world_size, 2); + CHECK(FLAGS_rank == 0 || FLAGS_rank == 1); + CHECK_GT(FLAGS_data_bytes, 0); + CHECK(!(FLAGS_force_gin && FLAGS_disable_gin)); + const std::string bootstrap_file = bootstrapPath(); + const std::string local_done = + bootstrap_file + ".done." + std::to_string(FLAGS_rank); + const std::string peer_done = + bootstrap_file + ".done." + std::to_string(1 - FLAGS_rank); + const std::string completion_ack = bootstrap_file + ".done.ack"; + + int device_count = 0; + checkCuda(cudaGetDeviceCount(&device_count), "cudaGetDeviceCount"); + CHECK_GT(device_count, 0); + const int gpu_id = + FLAGS_gpu_id >= 0 ? FLAGS_gpu_id : FLAGS_rank % device_count; + checkCuda(cudaSetDevice(gpu_id), "cudaSetDevice"); + + auto engine = std::make_unique(false); + auto* transport = engine->getOrCreateNcclTransport(); + CHECK_NOTNULL(transport); + + std::vector unique_id; + if (FLAGS_rank == 0) { + std::remove(bootstrap_file.c_str()); + std::remove(local_done.c_str()); + std::remove(peer_done.c_str()); + std::remove(completion_ack.c_str()); + unique_id = transport->createUniqueId(); + CHECK(!unique_id.empty()); + writeUniqueId(bootstrap_file, unique_id); + } else { + unique_id = readUniqueId(bootstrap_file); + } + + mooncake::device::NcclTransportConfig config; + config.rank = FLAGS_rank; + config.num_ranks = FLAGS_world_size; + config.enable_gin = !FLAGS_disable_gin; + config.gin_context_count = config.enable_gin ? 1 : 0; + config.lsa_barrier_count = 1; + config.require_lsa_multimem = FLAGS_require_lsa_multimem; + CHECK_EQ(transport->initialize(config, unique_id), 0); + + const size_t data_signal_offset = + alignUp(static_cast(FLAGS_data_bytes), kSignalAlignment); + const size_t ack_signal_offset = data_signal_offset + kSignalAlignment; + const size_t allocation_bytes = ack_signal_offset + kSignalAlignment; + + void* allocation = nullptr; + mooncake::device::NcclBufferRegistration registration; + CHECK_EQ(transport->allocateAndRegisterBuffer(allocation_bytes, &allocation, + ®istration), + 0); + CHECK_NOTNULL(allocation); + CHECK(registration.valid()); + char* buffer = static_cast(allocation); + checkCuda(cudaMemset(buffer, 0, allocation_bytes), "cudaMemset"); + + if (FLAGS_rank == 0) { + std::vector pattern(FLAGS_data_bytes); + for (int i = 0; i < FLAGS_data_bytes; ++i) { + pattern[i] = static_cast(i & 0xff); + } + checkCuda(cudaMemcpy(buffer, pattern.data(), pattern.size(), + cudaMemcpyHostToDevice), + "cudaMemcpy(pattern)"); + } + + int* device_result = nullptr; + checkCuda(cudaMalloc(&device_result, sizeof(int)), "cudaMalloc(result)"); + checkCuda(cudaMemset(device_result, 0xff, sizeof(int)), + "cudaMemset(result)"); + + const auto context = transport->deviceContext(registration); + CHECK(context.valid()); + if (FLAGS_rank == 0) { + senderKernel<<<1, 1>>>(context, buffer, FLAGS_data_bytes, + data_signal_offset, ack_signal_offset, + FLAGS_force_gin, device_result); + } else { + receiverKernel<<<1, 1>>>(context, buffer, FLAGS_data_bytes, + data_signal_offset, ack_signal_offset, + device_result); + } + checkCuda(cudaGetLastError(), "kernel launch"); + checkCuda(cudaDeviceSynchronize(), "cudaDeviceSynchronize"); + + int result = -1; + checkCuda(cudaMemcpy(&result, device_result, sizeof(result), + cudaMemcpyDeviceToHost), + "cudaMemcpy(result)"); + CHECK_EQ(result, 0) << "Device transfer validation failed"; + + // The Mooncake API intentionally leaves world synchronization to the + // caller. Rendezvous before deregistration so neither rank invalidates a + // window while the peer can still access it. The acknowledgement lets + // rank 0 remove the run-scoped marker files without racing rank 1. + publishMarker(local_done); + waitForMarker(peer_done); + if (FLAGS_rank == 1) { + publishMarker(completion_ack); + } else { + waitForMarker(completion_ack); + } + + const auto properties = transport->properties(); + LOG(INFO) << "PASS rank=" << FLAGS_rank + << " gin_backend=" << static_cast(properties.gin_backend) + << " gin_connections=" << properties.gin_connection_count + << " multimem_supported=" << properties.multimem_supported + << " multimem_enabled=" << properties.lsa_multimem_enabled + << " force_gin=" << FLAGS_force_gin + << " disable_gin=" << FLAGS_disable_gin; + + checkCuda(cudaFree(device_result), "cudaFree(result)"); + CHECK_EQ(transport->deregisterBuffer(®istration), 0); + CHECK_EQ(transport->freeBuffer(buffer), 0); + CHECK_EQ(transport->shutdown(), 0); + + if (FLAGS_rank == 0) { + std::remove(bootstrap_file.c_str()); + std::remove(local_done.c_str()); + std::remove(peer_done.c_str()); + std::remove(completion_ack.c_str()); + } + return 0; +} diff --git a/mooncake-transfer-engine/example/nccl_host_transport_example.cpp b/mooncake-transfer-engine/example/nccl_host_transport_example.cpp new file mode 100644 index 0000000000..f396548887 --- /dev/null +++ b/mooncake-transfer-engine/example/nccl_host_transport_example.cpp @@ -0,0 +1,248 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "transfer_engine.h" + +namespace { + +using mooncake::BatchID; +using mooncake::SegmentHandle; +using mooncake::TransferEngine; +using mooncake::TransferRequest; +using mooncake::TransferStatus; +using mooncake::TransferStatusEnum; + +constexpr size_t kBufferBytes = 1 << 20; +constexpr size_t kTransferBytes = 256 << 10; + +bool checkCuda(cudaError_t result, const char* operation) { + if (result == cudaSuccess) return true; + LOG(ERROR) << operation << " failed: " << cudaGetErrorString(result); + return false; +} + +bool checkNccl(ncclResult_t result, const char* operation) { + if (result == ncclSuccess) return true; + LOG(ERROR) << operation << " failed: " << ncclGetErrorString(result); + return false; +} + +bool waitForTransfer(TransferEngine* engine, BatchID batch) { + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(60); + while (std::chrono::steady_clock::now() < deadline) { + TransferStatus status{}; + auto result = engine->getTransferStatus(batch, 0, status); + if (!result.ok()) { + LOG(ERROR) << result.ToString(); + return false; + } + if (status.s == TransferStatusEnum::COMPLETED) return true; + if (status.s == TransferStatusEnum::FAILED || + status.s == TransferStatusEnum::TIMEOUT) { + LOG(ERROR) << "Transfer failed with status " << status.s; + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + LOG(ERROR) << "Timed out waiting for NCCL host transfer"; + return false; +} + +bool submitOne(TransferEngine* engine, SegmentHandle target, void* local, + uint64_t remote, TransferRequest::OpCode opcode) { + BatchID batch = engine->allocateBatchID(1); + if (batch == mooncake::INVALID_BATCH_ID) return false; + + TransferRequest request{}; + request.opcode = opcode; + request.source = local; + request.target_id = target; + request.target_offset = remote; + request.length = kTransferBytes; + auto result = engine->submitTransfer(batch, {request}); + if (!result.ok()) { + LOG(ERROR) << result.ToString(); + return false; + } + if (!waitForTransfer(engine, batch)) return false; + result = engine->freeBatchID(batch); + if (!result.ok()) { + LOG(ERROR) << result.ToString(); + return false; + } + return true; +} + +bool expectReadRejected(TransferEngine* engine, SegmentHandle target, + void* local, uint64_t remote) { + BatchID batch = engine->allocateBatchID(1); + if (batch == mooncake::INVALID_BATCH_ID) return false; + + TransferRequest request{}; + request.opcode = TransferRequest::READ; + request.source = local; + request.target_id = target; + request.target_offset = remote; + request.length = kTransferBytes; + auto result = engine->submitTransfer(batch, {request}); + const bool rejected = result.IsNotSupportedTransport(); + if (!rejected) { + LOG(ERROR) << "NCCL host READ returned " << result.ToString() + << ", expected NotSupportedTransport"; + } + + TransferStatus status{}; + auto status_result = engine->getTransferStatus(batch, 0, status); + const bool failed = + status_result.ok() && status.s == TransferStatusEnum::FAILED; + if (!failed) { + LOG(ERROR) << "Rejected NCCL host READ did not reach FAILED status"; + } + + auto free_result = engine->freeBatchID(batch); + if (!free_result.ok()) LOG(ERROR) << free_result.ToString(); + return rejected && failed && free_result.ok(); +} + +bool verifyPattern(void* device_ptr, int device, uint8_t expected) { + std::vector host(kTransferBytes); + if (!checkCuda(cudaSetDevice(device), "cudaSetDevice")) return false; + if (!checkCuda(cudaMemcpy(host.data(), device_ptr, host.size(), + cudaMemcpyDeviceToHost), + "cudaMemcpy device-to-host")) { + return false; + } + for (uint8_t value : host) { + if (value != expected) { + LOG(ERROR) << "Pattern mismatch: expected " + << static_cast(expected) << ", got " + << static_cast(value); + return false; + } + } + return true; +} + +} // namespace + +int main(int argc, char** argv) { + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = true; + + int device_count = 0; + if (!checkCuda(cudaGetDeviceCount(&device_count), "cudaGetDeviceCount") || + device_count < 2) { + LOG(ERROR) << "The NCCL host example requires two GPUs"; + return 1; + } + + void* buffers[2] = {nullptr, nullptr}; + std::unique_ptr engines[2]; + std::string names[2]; + + for (int rank = 0; rank < 2; ++rank) { + if (!checkCuda(cudaSetDevice(rank), "cudaSetDevice") || + !checkNccl(ncclMemAlloc(&buffers[rank], kBufferBytes), + "ncclMemAlloc") || + !checkCuda(cudaMemset(buffers[rank], 0, kBufferBytes), + "cudaMemset")) { + return 1; + } + + engines[rank] = std::make_unique(false); + if (engines[rank]->init(P2PHANDSHAKE, "127.0.0.1:0", "127.0.0.1", 0) != + 0 || + !engines[rank]->installTransport("nccl", nullptr) || + engines[rank]->registerLocalMemory( + buffers[rank], kBufferBytes, "cuda:" + std::to_string(rank)) != + 0) { + LOG(ERROR) << "Failed to initialize NCCL TE rank " << rank; + return 1; + } + names[rank] = engines[rank]->getLocalIpAndPort(); + } + + SegmentHandle peer_on_rank0 = engines[0]->openSegment(names[1]); + SegmentHandle peer_on_rank1 = engines[1]->openSegment(names[0]); + if (peer_on_rank0 == static_cast(-1) || + peer_on_rank1 == static_cast(-1)) { + LOG(ERROR) << "Failed to exchange P2P segment metadata"; + return 1; + } + + if (!expectReadRejected(engines[0].get(), peer_on_rank0, buffers[0], + reinterpret_cast(buffers[1]))) { + return 1; + } + LOG(INFO) << "NCCL host READ rejection validation passed"; + + std::atomic writes_ok{true}; + std::vector writers; + if (!checkCuda(cudaSetDevice(0), "cudaSetDevice") || + !checkCuda(cudaMemset(buffers[0], 0x5a, kTransferBytes), + "cudaMemset write source")) { + return 1; + } + for (int i = 0; i < 4; ++i) { + writers.emplace_back([&] { + if (!submitOne(engines[0].get(), peer_on_rank0, buffers[0], + reinterpret_cast(buffers[1]), + TransferRequest::WRITE)) { + writes_ok.store(false, std::memory_order_relaxed); + } + }); + } + for (auto& writer : writers) writer.join(); + if (!writes_ok.load(std::memory_order_relaxed) || + !verifyPattern(buffers[1], 1, 0x5a)) { + return 1; + } + LOG(INFO) << "NCCL host WRITE validation passed"; + + if (!checkCuda(cudaSetDevice(1), "cudaSetDevice") || + !checkCuda(cudaMemset(buffers[1], 0xa5, kTransferBytes), + "cudaMemset reverse write source") || + !checkCuda(cudaSetDevice(0), "cudaSetDevice") || + !checkCuda(cudaMemset(buffers[0], 0, kTransferBytes), + "cudaMemset reverse write destination") || + !submitOne(engines[1].get(), peer_on_rank1, buffers[1], + reinterpret_cast(buffers[0]), + TransferRequest::WRITE) || + !verifyPattern(buffers[0], 0, 0xa5)) { + return 1; + } + LOG(INFO) << "NCCL host bidirectional WRITE validation passed"; + + engines[0].reset(); + engines[1].reset(); + for (int rank = 0; rank < 2; ++rank) { + cudaSetDevice(rank); + if (!checkNccl(ncclMemFree(buffers[rank]), "ncclMemFree")) return 1; + } + return 0; +} diff --git a/mooncake-transfer-engine/example/show_link.cpp b/mooncake-transfer-engine/example/show_link.cpp new file mode 100644 index 0000000000..ee0e3f0867 --- /dev/null +++ b/mooncake-transfer-engine/example/show_link.cpp @@ -0,0 +1,69 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + +#include +#include + +#include "show_links.h" +#include "transfer_engine.h" + +DEFINE_string(metadata_server, "etcd://127.0.0.1:2379", + "Metadata server connection string"); +DEFINE_string(local_server_name, "", "Local server name (ip:port)"); +DEFINE_string(ip_or_host_name, "", "IP or hostname for RPC"); +DEFINE_int32(rpc_port, 12345, "RPC port"); +DEFINE_bool(json, false, "Output in JSON format"); +DEFINE_bool(discover_only, false, + "Only discover local topology without connecting"); + +using namespace mooncake; + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = 1; + + if (FLAGS_discover_only) { + auto engine = std::make_unique(/*auto_discover=*/false); + auto topology = engine->getLocalTopology(); + if (topology) { + topology->discover({}); + } + std::cout << engine->showLinks(FLAGS_json) << std::endl; + return 0; + } + + if (FLAGS_local_server_name.empty()) { + LOG(ERROR) << "Must specify --local_server_name or --discover_only"; + return 1; + } + + auto engine = std::make_unique(/*auto_discover=*/true); + int ret = engine->init(FLAGS_metadata_server, FLAGS_local_server_name, + FLAGS_ip_or_host_name, FLAGS_rpc_port); + if (ret) { + LOG(ERROR) << "Failed to initialize TransferEngine"; + return 1; + } + + engine->installTransport("rdma", nullptr); + + std::cout << engine->showLinks(FLAGS_json) << std::endl; + + engine->freeEngine(); + return 0; +} diff --git a/mooncake-transfer-engine/example/transfer_engine_bench.cpp b/mooncake-transfer-engine/example/transfer_engine_bench.cpp index 5c653227b3..13e4b0f7bf 100644 --- a/mooncake-transfer-engine/example/transfer_engine_bench.cpp +++ b/mooncake-transfer-engine/example/transfer_engine_bench.cpp @@ -44,7 +44,7 @@ #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ - defined(USE_UBSHMEM) || defined(USE_SUNRISE) + defined(USE_UBSHMEM) || defined(USE_SUPA) || defined(USE_SUNRISE) #include #if defined(USE_MNNVL) || defined(USE_UBSHMEM) @@ -110,7 +110,7 @@ DEFINE_string(backend, "classic", "Backend to use: classic|tent"); #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ - defined(USE_UBSHMEM) || defined(USE_SUNRISE) + defined(USE_UBSHMEM) || defined(USE_SUPA) || defined(USE_SUNRISE) DEFINE_bool(use_vram, true, "Allocate memory from GPU/NPU VRAM"); DEFINE_bool(init_mem, true, "Initialize allocated memory"); DEFINE_int32(gpu_id, 0, @@ -123,7 +123,7 @@ static void* allocateMemoryPool(size_t size, int buffer_id, bool from_vram = false) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ - defined(USE_UBSHMEM) || defined(USE_SUNRISE) + defined(USE_UBSHMEM) || defined(USE_SUPA) || defined(USE_SUNRISE) if (from_vram) { int gpu_id; if (FLAGS_gpu_id == -1) { @@ -195,7 +195,7 @@ static void* allocateMemoryPool(size_t size, int buffer_id, static void freeMemoryPool(void* addr, size_t size) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ - defined(USE_UBSHMEM) || defined(USE_SUNRISE) + defined(USE_UBSHMEM) || defined(USE_SUPA) || defined(USE_SUNRISE) if (FLAGS_protocol == "nvlink" || FLAGS_protocol == "hip") { #ifdef USE_MNNVL if (FLAGS_use_vram) { @@ -296,7 +296,7 @@ std::atomic total_batch_count(0); static inline void setWorkerDeviceIfNeeded() { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ - defined(USE_SUNRISE) + defined(USE_SUPA) || defined(USE_SUNRISE) if (FLAGS_use_vram && FLAGS_gpu_id >= 0) { checkCudaError(cudaSetDevice(FLAGS_gpu_id), "Failed to set device in worker"); @@ -308,7 +308,7 @@ static inline void setWorkerDeviceIfNeeded() { static int determineBufferCount() { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ - defined(USE_SUNRISE) + defined(USE_SUPA) || defined(USE_SUNRISE) if (FLAGS_use_vram) { int gpu_num; LOG(INFO) << "VRAM is used"; @@ -340,7 +340,7 @@ static std::vector allocateBuffers() { std::vector addr(buffer_num); #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ - defined(USE_UBSHMEM) || defined(USE_SUNRISE) + defined(USE_UBSHMEM) || defined(USE_SUPA) || defined(USE_SUNRISE) for (int i = 0; i < buffer_num; ++i) { addr[i] = allocateMemoryPool(FLAGS_buffer_size, i, FLAGS_use_vram); } @@ -364,7 +364,7 @@ static void freeBuffers(std::vector& addr) { static std::string getLocationName(int buffer_id) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ - defined(USE_UBSHMEM) || defined(USE_SUNRISE) + defined(USE_UBSHMEM) || defined(USE_SUPA) || defined(USE_SUNRISE) if (FLAGS_use_vram) { int name_suffix = (FLAGS_gpu_id == -1) ? buffer_id : FLAGS_gpu_id; return std::string(GPU_PREFIX) + std::to_string(name_suffix); diff --git a/mooncake-transfer-engine/fabric_allocator.cmake b/mooncake-transfer-engine/fabric_allocator.cmake index 62ee1095bf..a04208a42d 100644 --- a/mooncake-transfer-engine/fabric_allocator.cmake +++ b/mooncake-transfer-engine/fabric_allocator.cmake @@ -1,7 +1,7 @@ function(add_fabric_allocator_build_target) set(options) - set(oneValueArgs TARGET_NAME BUILD_SCRIPT COMMENT ENABLE_BUILD) - set(multiValueArgs BUILD_ARGS) + set(oneValueArgs TARGET_NAME BUILD_SCRIPT OUTPUT_NAME COMMENT ENABLE_BUILD) + set(multiValueArgs BUILD_ARGS BUILD_DEPENDS) cmake_parse_arguments(FAB "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) @@ -16,25 +16,29 @@ function(add_fabric_allocator_build_target) "BUILD_SCRIPT is required for add_fabric_allocator_build_target") endif() - add_custom_target(${FAB_TARGET_NAME} DEPENDS transfer_engine) - - get_target_property(_include_dirs ${FAB_TARGET_NAME} INCLUDE_DIRECTORIES) - if(NOT _include_dirs) - set(_include_dirs "") - endif() - string(REPLACE ";" " " _include_dirs_str "${_include_dirs}") - if(FAB_ENABLE_BUILD) + if(NOT FAB_OUTPUT_NAME) + message( + FATAL_ERROR + "OUTPUT_NAME is required for enabled fabric allocator build targets") + endif() + + set(_output_path "${CMAKE_CURRENT_BINARY_DIR}/${FAB_OUTPUT_NAME}") + # Track the allocator artifact so install targets do not rebuild it after a + # successful normal build. add_custom_command( - TARGET ${FAB_TARGET_NAME} - POST_BUILD - COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR} - COMMAND bash ${FAB_BUILD_SCRIPT} ${FAB_BUILD_ARGS} - ${CMAKE_CURRENT_BINARY_DIR} "${_include_dirs_str}" - WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR} + OUTPUT "${_output_path}" + COMMAND "${CMAKE_COMMAND}" -E make_directory "${CMAKE_CURRENT_BINARY_DIR}" + COMMAND bash "${FAB_BUILD_SCRIPT}" ${FAB_BUILD_ARGS} + "${CMAKE_CURRENT_BINARY_DIR}" "" + DEPENDS "${FAB_BUILD_SCRIPT}" ${FAB_BUILD_DEPENDS} + WORKING_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}" COMMENT "${FAB_COMMENT}" VERBATIM) + add_custom_target(${FAB_TARGET_NAME} ALL DEPENDS "${_output_path}") + else() + add_custom_target(${FAB_TARGET_NAME} ALL) endif() - set_property(TARGET ${FAB_TARGET_NAME} PROPERTY EXCLUDE_FROM_ALL FALSE) + add_dependencies(${FAB_TARGET_NAME} transfer_engine) endfunction() diff --git a/mooncake-transfer-engine/include/CMakeLists.txt b/mooncake-transfer-engine/include/CMakeLists.txt index c9486b9b52..7d5444d452 100644 --- a/mooncake-transfer-engine/include/CMakeLists.txt +++ b/mooncake-transfer-engine/include/CMakeLists.txt @@ -17,8 +17,11 @@ install(FILES transport/device/device_ops.cuh DESTINATION include/transport/devi install(FILES transport/device/comm_device.cuh DESTINATION include/transport/device) install(FILES transport/device/p2p_device.cuh DESTINATION include/transport/device) install(FILES transport/device/ibgda_device.cuh DESTINATION include/transport/device) +install(FILES transport/device/nccl_device_transport.h DESTINATION include/transport/device) +install(FILES transport/device/nccl_device.cuh DESTINATION include/transport/device) install(FILES transport/device/cuda/cuda_ops.cuh DESTINATION include/transport/device/cuda) install(FILES transport/device/musa/musa_ops.cuh DESTINATION include/transport/device/musa) +install(FILES transport/device/maca/maca_ops.cuh DESTINATION include/transport/device/maca) # IBGDA library headers install(DIRECTORY transport/device/ibgda/ DESTINATION include/transport/device/ibgda diff --git a/mooncake-transfer-engine/include/ascend_allocator.h b/mooncake-transfer-engine/include/ascend_allocator.h index c712718758..4f4c1f7392 100644 --- a/mooncake-transfer-engine/include/ascend_allocator.h +++ b/mooncake-transfer-engine/include/ascend_allocator.h @@ -9,6 +9,13 @@ namespace mooncake { void* ascend_allocate_memory(size_t total_size, const std::string& protocol); +// Fabric-mem best-effort: try 100%/90%/.../50% of target (1G-aligned). +// Sets *actual_size on success; returns nullptr if 50% also fails. +// Non-fabric paths fall back to exact-size ascend_allocate_memory. +void* ascend_allocate_memory_best_effort(size_t target_size, + const std::string& protocol, + size_t* actual_size); + void ascend_free_memory(const std::string& protocol, void* ptr); // Check if [addr, addr+length) overlaps with any store memory range. diff --git a/mooncake-transfer-engine/include/common.h b/mooncake-transfer-engine/include/common.h index 5ae9a66471..6ad2fb582e 100644 --- a/mooncake-transfer-engine/include/common.h +++ b/mooncake-transfer-engine/include/common.h @@ -32,10 +32,12 @@ #include #include #include +#include #include #include #include #include +#include #include "error.h" @@ -61,6 +63,7 @@ enum class HandShakeRequestType { Metadata = 1, Notify = 2, Probe = 3, + Invalid = 0xfe, // placeholder for old protocol without RequestType OldProtocol = 0xff, }; @@ -74,6 +77,16 @@ static inline std::string getHostname() { return hostname; } +// libnuma fills the cache numa_node_to_cpus() reads lazily and without locking, +// so concurrent first callers each allocate it and all but one are orphaned -- +// a leak LeakSanitizer fails the build on. Worker pools bind every thread at +// startup, so they hit that window. An inline function, not a static local: +// bindToSocket() has internal linkage, so a static local would be per-TU. +inline std::mutex &numaNodeCpuCacheMutex() { + static std::mutex mutex; + return mutex; +} + static inline int bindToSocket(int socket_id) { if (unlikely(numa_available() < 0)) { LOG(WARNING) << "The platform does not support NUMA"; @@ -84,7 +97,10 @@ static inline int bindToSocket(int socket_id) { if (socket_id < 0 || socket_id >= numa_num_configured_nodes()) socket_id = 0; struct bitmask *cpu_list = numa_allocate_cpumask(); - numa_node_to_cpus(socket_id, cpu_list); + { + std::lock_guard guard(numaNodeCpuCacheMutex()); + numa_node_to_cpus(socket_id, cpu_list); + } int nr_possible_cpus = numa_num_possible_cpus(); int nr_cpus = 0; for (int cpu = 0; cpu < nr_possible_cpus; ++cpu) { @@ -366,9 +382,13 @@ static inline ssize_t readFully(int fd, void *buf, size_t len) { size_t nbytes = len; while (nbytes && std::chrono::steady_clock::now() < deadline) { ssize_t rc = read(fd, pos, nbytes); - if (rc < 0 && (errno == EAGAIN || errno == EINTR)) + if (rc < 0 && errno == EINTR) continue; - else if (rc < 0) { + else if (rc < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) { + LOG(WARNING) << "Socket read timed out: expected " << len + << " bytes, actual " << len - nbytes << " bytes"; + return len - nbytes; + } else if (rc < 0) { PLOG(ERROR) << "Socket read failed"; return rc; } else if (rc == 0) { @@ -443,13 +463,13 @@ static inline size_t getHandshakeMaxLength() { } static inline std::pair readString(int fd) { - HandShakeRequestType type = HandShakeRequestType::Connection; + HandShakeRequestType type = HandShakeRequestType::Invalid; const size_t kMaxLength = getHandshakeMaxLength(); uint64_t length = 0; ssize_t n = readFully(fd, &length, sizeof(length)); if (n != (ssize_t)sizeof(length)) { - LOG(ERROR) << "readString: failed to read length, got: " << n; + LOG(WARNING) << "readString: incomplete handshake length, got: " << n; return {type, ""}; } @@ -528,6 +548,7 @@ static inline bool overlap(const void *a, size_t a_len, const void *b, class RWSpinlock { union RWTicket { constexpr RWTicket() : whole(0) {} + constexpr RWTicket(uint64_t v) : whole(v) {} uint64_t whole; uint32_t readWrite; struct { @@ -535,26 +556,12 @@ class RWSpinlock { uint16_t read; uint16_t users; }; - } ticket; - - private: - static void asm_volatile_memory() { asm volatile("" ::: "memory"); } - - template - static T load_acquire(T *addr) { - T t = *addr; - asm_volatile_memory(); - return t; - } + }; - template - static void store_release(T *addr, T v) { - asm_volatile_memory(); - *addr = v; - } + std::atomic ticket; public: - RWSpinlock() {} + RWSpinlock() : ticket(0) {} RWSpinlock(RWSpinlock const &) = delete; RWSpinlock &operator=(RWSpinlock const &) = delete; @@ -562,17 +569,21 @@ class RWSpinlock { void lock() { writeLockNice(); } bool tryLock() { - RWTicket t; - uint64_t old = t.whole = load_acquire(&ticket.whole); + RWTicket t, expected; + expected.whole = ticket.load(std::memory_order_acquire); + t.whole = expected.whole; if (t.users != t.write) return false; ++t.users; - return __sync_bool_compare_and_swap(&ticket.whole, old, t.whole); + return ticket.compare_exchange_weak(expected.whole, t.whole, + std::memory_order_acquire); } void writeLockAggressive() { uint32_t count = 0; - uint16_t val = __sync_fetch_and_add(&ticket.users, 1); - while (val != load_acquire(&ticket.write)) { + uint16_t val = fetch_add_users(1); + RWTicket t; + while (val != + (t.whole = ticket.load(std::memory_order_acquire), t.write)) { PAUSE(); if (++count > 1000) std::this_thread::yield(); } @@ -587,16 +598,22 @@ class RWSpinlock { } void unlockAndLockShared() { - uint16_t val = __sync_fetch_and_add(&ticket.read, 1); + uint16_t val = fetch_add_read(1); (void)val; } void unlock() { + uint64_t expected = ticket.load(std::memory_order_relaxed); + uint64_t new_val; RWTicket t; - t.whole = load_acquire(&ticket.whole); - ++t.read; - ++t.write; - store_release(&ticket.readWrite, t.readWrite); + do { + t.whole = expected; + ++t.read; + ++t.write; + new_val = t.whole; + } while (!ticket.compare_exchange_weak(expected, new_val, + std::memory_order_release, + std::memory_order_relaxed)); } void lockShared() { @@ -608,15 +625,58 @@ class RWSpinlock { } bool tryLockShared() { - RWTicket t, old; - old.whole = t.whole = load_acquire(&ticket.whole); - old.users = old.read; + RWTicket t, expected; + expected.whole = ticket.load(std::memory_order_acquire); + t.whole = expected.whole; + expected.users = expected.read; ++t.read; ++t.users; - return __sync_bool_compare_and_swap(&ticket.whole, old.whole, t.whole); + return ticket.compare_exchange_weak(expected.whole, t.whole, + std::memory_order_acquire); } - void unlockShared() { __sync_fetch_and_add(&ticket.write, 1); } + void unlockShared() { fetch_add_write(1); } + + private: + uint16_t fetch_add_users(uint16_t delta) { + uint64_t expected = ticket.load(std::memory_order_relaxed); + uint64_t new_val; + RWTicket t; + do { + t.whole = expected; + t.users += delta; + new_val = t.whole; + } while (!ticket.compare_exchange_weak(expected, new_val, + std::memory_order_acquire, + std::memory_order_relaxed)); + return static_cast(t.users - delta); + } + + uint16_t fetch_add_read(uint16_t delta) { + uint64_t expected = ticket.load(std::memory_order_relaxed); + uint64_t new_val; + RWTicket t; + do { + t.whole = expected; + t.read += delta; + new_val = t.whole; + } while (!ticket.compare_exchange_weak(expected, new_val, + std::memory_order_release)); + return static_cast(t.read - delta); + } + + uint16_t fetch_add_write(uint16_t delta) { + uint64_t expected = ticket.load(std::memory_order_relaxed); + uint64_t new_val; + RWTicket t; + do { + t.whole = expected; + t.write += delta; + new_val = t.whole; + } while (!ticket.compare_exchange_weak(expected, new_val, + std::memory_order_release)); + return static_cast(t.write - delta); + } public: struct WriteGuard { diff --git a/mooncake-transfer-engine/include/config.h b/mooncake-transfer-engine/include/config.h index 82eb050460..6686e9de37 100644 --- a/mooncake-transfer-engine/include/config.h +++ b/mooncake-transfer-engine/include/config.h @@ -22,6 +22,7 @@ #include #include #include +#include #include #include @@ -32,6 +33,12 @@ enum class EndpointStoreType { SIEVE = 1, }; +// Which NICs the EFA transport registers a buffer on. +enum class EfaNicSelection { + ALL = 0, // every NIC, the historical behavior + LOCAL = 1, // device memory only on that GPU's topology-local NICs +}; + struct GlobalConfig { size_t num_cq_per_ctx = 1; size_t num_comp_channels_per_ctx = 1; @@ -44,6 +51,12 @@ struct GlobalConfig { size_t num_qp_per_ep = 2; size_t max_sge = 4; size_t max_wr = 256; + // Set when MC_MAX_WR was given explicitly. EFA's transmit depth is a + // per-device attribute (2048 on p6-b300, 4096 on p5), so with no override + // the EFA transport adopts the provider's depth rather than max_wr; with + // one it honors the operator's value, clamped to the hardware. The RDMA + // transport passes max_wr to ibv_create_qp() and is unaffected. + bool max_wr_from_env = false; size_t max_inline = 64; ibv_mtu mtu_length = IBV_MTU_4096; uint16_t handshake_port = 12001; @@ -58,19 +71,53 @@ struct GlobalConfig { // torn-down pod IP) it stalls for the kernel's full SYN-retry cycle, // which is minutes. Override via MC_HANDSHAKE_CONNECT_TIMEOUT. int handshake_connect_timeout = 5; + // Cooldown before retrying a failed RDMA peer rail. Override via + // MC_RDMA_RAIL_PAUSE_SECONDS. + uint64_t rdma_rail_pause_seconds = 30; bool metacache = true; + // Periodically refresh Transfer Engine metadata-derived local caches. 0 + // disables the background poller and preserves the manual + // syncSegmentCache() behavior. Currently refreshes cached remote segment + // descriptors. Override via MC_TE_METADATA_REFRESH_INTERVAL_SECONDS. + uint64_t te_metadata_refresh_interval_seconds = 0; int log_level = google::INFO; bool trace = false; int64_t slice_timeout = -1; + // Active-connect circuit-breaker. After an endpoint to a peer is torn down + // (path failure / QP fatal), pause active reconnection to that peer's + // address for this many milliseconds, so the posting worker is not + // blocked re-handshaking a likely-gone peer (a k8s rolling restart brings + // the pod back at a different podIP:port, so the old address is dead). The + // not-yet-posted slices fail/redispatch instead of hanging. 0 disables. + // Override via MC_CONN_PAUSE_TTL_MS. + int conn_pause_ttl_ms = 0; uint16_t rpc_min_port = 15000; uint16_t rpc_max_port = 17000; bool use_ipv6 = false; size_t fragment_limit = 16384; bool enable_dest_device_affinity = false; + bool enable_hca_peer_affinity = false; + std::unordered_map> nic_peer_affinity; + bool log_rdma_slice_affinity = false; + bool track_rdma_posted_slices = false; int parallel_reg_mr = -1; + // Cap on concurrent buffer registrations in registerLocalMemoryBatch(). + // 0 (default) = unbounded, one thread per buffer. Set via + // MC_MAX_CONCURRENT_REG_MR; the best value is platform-specific, see the + // measured tables in efa_transport.cpp before choosing one. + size_t max_concurrent_reg_mr = 0; + // Which NICs a buffer is registered on in the EFA transport. ALL (default) + // registers every buffer on every NIC; LOCAL restricts device memory to the + // NICs the topology reports as closest to that GPU. Set via + // MC_EFA_NIC_SELECTION=all|local; see efa_transport.cpp for the trade-off. + EfaNicSelection efa_nic_selection = EfaNicSelection::ALL; size_t eic_max_block_size = 64UL * 1024 * 1024; EndpointStoreType endpoint_store_type = EndpointStoreType::SIEVE; int ib_traffic_class = -1; + // InfiniBand Service Level (SL), 0-15. -1 = use default (0). + // Maps to a Virtual Lane on the switch for QoS isolation, e.g. to + // steer KV-cache traffic into a different VL than EP all-to-all. + int ib_service_level = -1; // mlx5 QP UDP source ports for ECMP path diversification. // Empty = no modification. QP at index i uses // mlx5_qp_udp_sports[i % size]. Requires mlx5 device + RoCEv2, @@ -82,9 +129,10 @@ struct GlobalConfig { // mode the setting is a no-op. Requires USE_MLX5DV. bool mlx5_qp_lag_port_balance = false; // ib_pci_relaxed_ordering_mode: 0: off, 1: on if supported, 2: auto - int ib_pci_relaxed_ordering_mode = 0; + int ib_pci_relaxed_ordering_mode = 1; bool ascend_use_fabric_mem = false; bool ascend_agent_mode = false; + bool sunrise_use_device_mem = false; // Transient flag scoped to a single TE init: set true by the Store entry // (Client::InitTransferEngine) before installing the ascend transport, and // reset to false right after. Lets ascend_direct distinguish a Store-init @@ -105,7 +153,10 @@ struct RpcCommunicatorConfig { std::string listen_address; size_t thread_count = 0; size_t timeout_seconds = 30; - size_t pool_size = 10; + // Maximum number of cached RPC client connections per target endpoint. + // RPC client I/O threads are configured by + // MC_TE_RPC_CLIENT_IO_THREADS/MC_RPC_CLIENT_IO_THREADS. + size_t pool_size = 100; }; void loadGlobalConfig(GlobalConfig& config); diff --git a/mooncake-transfer-engine/include/cuda_alike.h b/mooncake-transfer-engine/include/cuda_alike.h index 2ec35f4196..bcec46db45 100644 --- a/mooncake-transfer-engine/include/cuda_alike.h +++ b/mooncake-transfer-engine/include/cuda_alike.h @@ -13,6 +13,8 @@ #include "gpu_vendor/ubshmem.h" #elif defined(USE_MACA) #include "gpu_vendor/maca.h" +#elif defined(USE_SUPA) +#include "gpu_vendor/supa.h" #elif defined(USE_SUNRISE) #include "gpu_vendor/sunrise.h" #elif defined(USE_HYGON) @@ -23,8 +25,9 @@ #include #endif -#if !defined(USE_HIP) && !defined(USE_MUSA) && !defined(USE_MLU) && \ - !defined(USE_UBSHMEM) && !defined(USE_MACA) && !defined(USE_SUNRISE) +#if !defined(USE_HIP) && !defined(USE_MUSA) && !defined(USE_MLU) && \ + !defined(USE_UBSHMEM) && !defined(USE_MACA) && !defined(USE_SUNRISE) && \ + !defined(USE_SUPA) #include const static std::string GPU_PREFIX = "cuda:"; #endif diff --git a/mooncake-transfer-engine/include/gpu_vendor/maca.h b/mooncake-transfer-engine/include/gpu_vendor/maca.h index 03c54f9886..04c96004de 100644 --- a/mooncake-transfer-engine/include/gpu_vendor/maca.h +++ b/mooncake-transfer-engine/include/gpu_vendor/maca.h @@ -120,6 +120,7 @@ static inline CUresult cuGetErrorString(CUresult error, const char **err_str) { #define cudaMemcpyDeviceToHost mcMemcpyDeviceToHost #define cudaMemcpyHostToDevice mcMemcpyHostToDevice #define cudaMemcpyKind mcMemcpyKind +#define cudaMemcpyPeerAsync mcMemcpyPeerAsync #define cudaMemset mcMemset #define cudaMemsetAsync mcMemsetAsync #define cudaMemoryTypeDevice mcMemoryTypeDevice @@ -133,6 +134,7 @@ static inline CUresult cuGetErrorString(CUresult error, const char **err_str) { #define cudaStreamDestroy mcStreamDestroy #define cudaStreamNonBlocking mcStreamNonBlocking #define cudaStreamPerThread mcStreamPerThread +#define cudaStreamQuery mcStreamQuery #define cudaStreamSynchronize mcStreamSynchronize #define cudaStream_t mcStream_t #define cudaStreamWaitEvent mcStreamWaitEvent diff --git a/mooncake-transfer-engine/include/gpu_vendor/musa.h b/mooncake-transfer-engine/include/gpu_vendor/musa.h index fd46043f1b..4fd47def0b 100644 --- a/mooncake-transfer-engine/include/gpu_vendor/musa.h +++ b/mooncake-transfer-engine/include/gpu_vendor/musa.h @@ -69,6 +69,8 @@ const static std::string GPU_PREFIX = "musa:"; #define cudaGetDeviceCount musaGetDeviceCount #define cudaGetErrorString musaGetErrorString #define cudaGetLastError musaGetLastError +#define cudaHostAlloc musaHostAlloc +#define cudaHostAllocMapped musaHostAllocMapped #define cudaHostRegister musaHostRegister #define cudaHostRegisterPortable musaHostRegisterPortable #define cudaHostUnregister musaHostUnregister @@ -84,6 +86,7 @@ const static std::string GPU_PREFIX = "musa:"; #define cudaMemcpyDefault musaMemcpyDefault #define cudaMemcpyDeviceToHost musaMemcpyDeviceToHost #define cudaMemcpyHostToDevice musaMemcpyHostToDevice +#define cudaMemcpyKind musaMemcpyKind #define cudaMemset musaMemset #define cudaMemsetAsync musaMemsetAsync #define cudaMemoryTypeDevice musaMemoryTypeDevice @@ -93,13 +96,30 @@ const static std::string GPU_PREFIX = "musa:"; #define cudaPointerGetAttributes musaPointerGetAttributes #define cudaSetDevice musaSetDevice #define cudaStreamCreate musaStreamCreate +#define cudaStreamCreateWithFlags musaStreamCreateWithFlags +#define cudaStreamNonBlocking musaStreamNonBlocking #define cudaStreamDestroy musaStreamDestroy +#define cudaStreamPerThread musaStreamPerThread +#define cudaStreamQuery musaStreamQuery +#define cudaStreamCaptureStatus musaStreamCaptureStatus +#define cudaStreamCaptureStatusNone musaStreamCaptureStatusNone +#define cudaStreamIsCapturing musaStreamIsCapturing +#define cudaStreamWaitEvent musaStreamWaitEvent #define cudaDeviceSynchronize musaDeviceSynchronize #define cudaStreamSynchronize musaStreamSynchronize #define cudaStream_t musaStream_t #define cudaSuccess musaSuccess +#define cudaErrorNotReady musaErrorNotReady #define cudaDeviceGetAttribute musaDeviceGetAttribute #define cudaEvent_t musaEvent_t +#define cudaEventCreateWithFlags musaEventCreateWithFlags +#define cudaEventDisableTiming MU_EVENT_DISABLE_TIMING +#define cudaEventDestroy musaEventDestroy +#define cudaEventQuery musaEventQuery +#define cudaEventRecord musaEventRecord +#define cudaEventSynchronize musaEventSynchronize +#define cudaDeviceProp musaDeviceProp +#define cudaGetDeviceProperties musaGetDeviceProperties #define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice #define cudaDevAttrClockRate musaDevAttrClockRate #define cudaLaunchConfig_t musaLaunchConfig_t diff --git a/mooncake-transfer-engine/include/gpu_vendor/supa.h b/mooncake-transfer-engine/include/gpu_vendor/supa.h new file mode 100644 index 0000000000..1c8e0040c7 --- /dev/null +++ b/mooncake-transfer-engine/include/gpu_vendor/supa.h @@ -0,0 +1,137 @@ +#pragma once + +#include +#include +#include + +const static std::string GPU_PREFIX = "supa:"; + +// ===================== Runtime API types ===================== +#define cudaError_t supaError_t +#define cudaSuccess supaSuccess +#define cudaErrorNotReady supaErrorNotReady +#define cudaErrorPeerAccessAlreadyEnabled supaErrorPeerAccessAlreadyEnabled +#define cudaGetErrorString supaGetErrorString +#define cudaMemoryTypeHost supaMemoryTypeHost +#define cudaMemoryTypeDevice supaMemoryTypeDevice +#define cudaMemoryTypeUnregistered supaMemoryTypeUnregistered +#define cudaPointerAttributes supaPointerAttributes +#define cudaDeviceProp supaDeviceProp +#define cudaIpcMemHandle_t supaIpcMemHandle_t +#define cudaIpcMemLazyEnablePeerAccess supaIpcMemLazyEnablePeerAccess +#define cudaStream_t supaStream_t +#define cudaEvent_t supaEvent_t +#define cudaMemcpyDefault supaMemcpyDefault +#define cudaMemcpyHostToDevice supaMemcpyHostToDevice +#define cudaMemcpyDeviceToHost supaMemcpyDeviceToHost +#define cudaMemcpyDeviceToDevice supaMemcpyDeviceToDevice +#define cudaHostAllocDefault supaHostAllocDefault +#define cudaHostAllocPortable supaHostAllocPortable +#define cudaHostAllocMapped supaHostAllocMapped +#define cudaHostRegisterDefault supaHostRegisterDefault +#define cudaHostRegisterPortable supaHostRegisterPortable +#define cudaHostRegisterMapped supaHostRegisterMapped +#define cudaHostRegisterIoMemory supaHostRegisterIoMemory +#define cudaStreamNonBlocking supaStreamNonBlocking +#define cudaEventDefault supaEventDefault +#define cudaEventDisableTiming supaEventDisableTiming +#define cudaDevAttrClockRate supaDevAttrClockRate + +// ===================== Runtime API functions ===================== +#define cudaGetDeviceCount supaGetDeviceCount +#define cudaSetDevice supaSetDevice +#define cudaGetDevice supaGetDevice +#define cudaDeviceGetAttribute supaDeviceGetAttribute +#define cudaDeviceGetPCIBusId supaDeviceGetPCIBusId +#define cudaMalloc supaMalloc +#define cudaFree supaFree +#define cudaMallocHost supaMallocHost +#define cudaHostAlloc supaHostAlloc +#define cudaFreeHost supaFreeHost +#define cudaHostRegister supaHostRegister +#define cudaHostUnregister supaHostUnregister +#define cudaHostGetDevicePointer supaHostGetDevicePointer +#define cudaMemcpy supaMemcpy +#define cudaMemcpyAsync supaMemcpyAsync +#define cudaMemset supaMemset +#define cudaMemsetAsync supaMemsetAsync +#define cudaPointerGetAttributes supaPointerGetAttributes +#define cudaDeviceCanAccessPeer supaDeviceCanAccessPeer +#define cudaDeviceEnablePeerAccess supaDeviceEnablePeerAccess +#define cudaIpcGetMemHandle supaIpcGetMemHandle +#define cudaIpcOpenMemHandle supaIpcOpenMemHandle +#define cudaIpcCloseMemHandle supaIpcCloseMemHandle +#define cudaDeviceSynchronize supaDeviceSynchronize +#define cudaStreamCreate supaStreamCreate +#define cudaStreamCreateWithFlags supaStreamCreateWithFlags +#define cudaStreamDestroy supaStreamDestroy +#define cudaStreamQuery supaStreamQuery +#define cudaStreamSynchronize supaStreamSynchronize +#define cudaLaunchHostFunc supaLaunchHostFunc +#define cudaEventCreate supaEventCreate +#define cudaEventCreateWithFlags supaEventCreateWithFlags +#define cudaEventDestroy supaEventDestroy +#define cudaEventRecord supaEventRecord +#define cudaEventQuery supaEventQuery +#define cudaGetLastError supaGetLastError +#define cudaGetDeviceProperties supaGetDeviceProperties +#define cudaStreamPerThread supaStreamPerThread +#define cudaEventSynchronize supaEventSynchronize + +// ===================== Driver API types ===================== +#define CUresult SUresult +#define CUDA_SUCCESS SUPA_SUCCESS +#define CUDA_ERROR_NOT_PERMITTED SUPA_ERROR_NOT_PERMITTED +#define CUDA_ERROR_NOT_SUPPORTED SUPA_ERROR_NOT_SUPPORTED +#define CUdevice SUdevice +#define CUdeviceptr SUdeviceptr +#define CUcontext SUcontext +#define CUmemGenericAllocationHandle SUmemGenericAllocationHandle +#define CUmemAllocationProp SUmemAllocationProp +#define CUmemAccessDesc SUmemAccessDesc +#define CUmemFabricHandle SUmemFabricHandle +#define CUmemAllocationHandleType SUmemAllocationHandleType +#define CUmemorytype SUmemorytype +#define CUmemRangeHandleType SUmemRangeHandleType + +// ===================== Driver API enums ===================== +#define CU_MEM_ALLOCATION_TYPE_PINNED SU_MEM_ALLOCATION_TYPE_PINNED +#define CU_MEM_LOCATION_TYPE_DEVICE SU_MEM_LOCATION_TYPE_DEVICE +#define CU_MEM_HANDLE_TYPE_FABRIC SU_MEM_HANDLE_TYPE_FABRIC +#define CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD SU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD +#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE SU_MEM_ACCESS_FLAGS_PROT_READWRITE +#define CU_MEM_ALLOC_GRANULARITY_MINIMUM SU_MEM_ALLOC_GRANULARITY_MINIMUM +#define CU_MEMORYTYPE_HOST SU_MEMORYTYPE_HOST +#define CU_MEMORYTYPE_DEVICE SU_MEMORYTYPE_DEVICE +#define CU_POINTER_ATTRIBUTE_MEMORY_TYPE SU_POINTER_ATTRIBUTE_MEMORY_TYPE +#define CU_POINTER_ATTRIBUTE_RANGE_SIZE SU_POINTER_ATTRIBUTE_RANGE_SIZE +#define CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED \ + SU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED +#define CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED \ + SU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED +#define CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED \ + SU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_SUPA_VMM_SUPPORTED + +// ===================== Driver API functions ===================== +#define cuInit suInit +#define cuDeviceGet suDeviceGet +#define cuDeviceGetAttribute suDeviceGetAttribute +#define cuPointerGetAttribute suPointerGetAttribute +#define cuGetErrorString suGetErrorString +#define cuMemCreate suMemCreate +#define cuMemRelease suMemRelease +#define cuMemAddressReserve suMemAddressReserve +#define cuMemAddressFree suMemAddressFree +#define cuMemMap suMemMap +#define cuMemUnmap suMemUnmap +#define cuMemSetAccess suMemSetAccess +#define cuMemGetAddressRange suMemGetAddressRange +#define cuMemGetHandleForAddressRange suMemGetHandleForAddressRange +#define cuMemRetainAllocationHandle suMemRetainAllocationHandle +#define cuMemExportToShareableHandle suMemExportToShareableHandle +#define cuMemImportFromShareableHandle suMemImportFromShareableHandle +#define cuMemGetAllocationGranularity suMemGetAllocationGranularity +#define cuDevicePrimaryCtxRetain suDevicePrimaryCtxRetain +#define cuDevicePrimaryCtxRelease suDevicePrimaryCtxRelease +#define cuCtxSetCurrent suCtxSetCurrent +#define CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL SU_POINTER_ATTRIBUTE_DEVICE_ORDINAL diff --git a/mooncake-transfer-engine/include/graceful_shutdown.h b/mooncake-transfer-engine/include/graceful_shutdown.h new file mode 100644 index 0000000000..41a637c382 --- /dev/null +++ b/mooncake-transfer-engine/include/graceful_shutdown.h @@ -0,0 +1,41 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MOONCAKE_GRACEFUL_SHUTDOWN_H_ +#define MOONCAKE_GRACEFUL_SHUTDOWN_H_ + +#include + +namespace mooncake { + +class TransferEngineImpl; + +class ShutdownToken { + public: + virtual ~ShutdownToken() = default; + + virtual void shutdown() = 0; + + virtual void detach() = 0; +}; + +void registerTokenForShutdown(std::shared_ptr token); + +void registerEngineForShutdown(std::shared_ptr impl); + +void installGracefulShutdownHandlers(); + +} // namespace mooncake + +#endif // MOONCAKE_GRACEFUL_SHUTDOWN_H_ diff --git a/mooncake-transfer-engine/include/hip_device_guard.h b/mooncake-transfer-engine/include/hip_device_guard.h new file mode 100644 index 0000000000..49b6b5069e --- /dev/null +++ b/mooncake-transfer-engine/include/hip_device_guard.h @@ -0,0 +1,52 @@ +// Copyright 2025 Mooncake Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#if defined(USE_HIP) || defined(USE_HIP_DMABUF) + +#include + +namespace mooncake { + +// RAII: saves the active HIP device and restores it on scope exit, so a +// function that calls hipSetDevice() doesn't leave the caller on the wrong GPU +// (which would fail its next kernel launch with hipErrorInvalidDevice). +class HipDeviceGuard { + public: + HipDeviceGuard() { saved_ = (hipGetDevice(&prev_device_) == hipSuccess); } + + // Also switch to target_device; set_ok() reports whether that succeeded. + explicit HipDeviceGuard(int target_device) : HipDeviceGuard() { + set_ok_ = (hipSetDevice(target_device) == hipSuccess); + } + + ~HipDeviceGuard() { + if (saved_) (void)hipSetDevice(prev_device_); + } + + bool set_ok() const { return set_ok_; } + + HipDeviceGuard(const HipDeviceGuard&) = delete; + HipDeviceGuard& operator=(const HipDeviceGuard&) = delete; + + private: + int prev_device_ = 0; + bool saved_ = false; + bool set_ok_ = true; +}; + +} // namespace mooncake + +#endif // USE_HIP || USE_HIP_DMABUF diff --git a/mooncake-transfer-engine/include/multi_transport.h b/mooncake-transfer-engine/include/multi_transport.h index c556541bcd..541a7391ea 100644 --- a/mooncake-transfer-engine/include/multi_transport.h +++ b/mooncake-transfer-engine/include/multi_transport.h @@ -20,7 +20,11 @@ #include "transport/transport.h" namespace mooncake { +class TransferEngineImplTestPeer; + class MultiTransport { + friend class TransferEngineImplTestPeer; + public: using BatchID = Transport::BatchID; using TransferRequest = Transport::TransferRequest; @@ -85,4 +89,4 @@ class MultiTransport { }; } // namespace mooncake -#endif // MULTI_TRANSPORT_H_ \ No newline at end of file +#endif // MULTI_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/include/multi_transport_locality.h b/mooncake-transfer-engine/include/multi_transport_locality.h new file mode 100644 index 0000000000..c298b4a702 --- /dev/null +++ b/mooncake-transfer-engine/include/multi_transport_locality.h @@ -0,0 +1,82 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MULTI_TRANSPORT_LOCALITY_H +#define MULTI_TRANSPORT_LOCALITY_H + +#include +#include + +namespace mooncake { + +// Extract the host portion of a segment name. Segment names carry an optional +// ":port" suffix, and the host may be an IPv4 address, a hostname, or an IPv6 +// literal. IPv6 literals contain multiple colons, so a naive rfind(':') would +// corrupt them; the following forms are handled explicitly: +// "10.0.0.1:8000" -> "10.0.0.1" (IPv4 / hostname with port) +// "node-a" -> "node-a" (no port) +// "[2001:db8::1]:8000" -> "2001:db8::1" (bracketed IPv6 with port) +// "[2001:db8::1]" -> "2001:db8::1" (bracketed IPv6, no port) +// "2001:db8::1" -> "2001:db8::1" (bare IPv6, no port) +inline std::string segmentHost(const std::string& segment_name) { + if (!segment_name.empty() && segment_name.front() == '[') { + // Bracketed IPv6 literal: strip the brackets and ignore any ":port". + auto close = segment_name.find(']'); + if (close != std::string::npos) { + return segment_name.substr(1, close - 1); + } + return segment_name; // Malformed; return as-is. + } + auto first = segment_name.find(':'); + if (first == std::string::npos) { + return segment_name; // No port and no colon. + } + if (first != segment_name.rfind(':')) { + // More than one colon and not bracketed: a bare IPv6 literal without a + // port (an IPv6 host with a port must use brackets). Treat it all as + // the host. + return segment_name; + } + return segment_name.substr(0, first); // "host:port". +} + +// Case-insensitive equality. Hostnames are case-insensitive per DNS, and IPv6 +// hex literals may differ only in letter case, so a plain "==" would spuriously +// classify the same host as remote and lose the intra-node hip fast path. +inline bool hostEquals(const std::string& a, const std::string& b) { + if (a.size() != b.size()) return false; + for (size_t i = 0; i < a.size(); ++i) { + if (std::tolower(static_cast(a[i])) != + std::tolower(static_cast(b[i]))) { + return false; + } + } + return true; +} + +// The device KV pool is registered under both rdma and hip in a "rdma,hip" +// multi-protocol segment. hip transport uses GPU IPC, which only works between +// processes on the same physical host, so a hip buffer is a valid transport for +// a target only when that target lives on the same host as the initiator. +// A cross-host target must fall back to rdma. Two engines co-located on one +// host share the host portion of their segment name (they differ only in port). +inline bool isHipReachableTarget(const std::string& target_segment_name, + const std::string& local_server_name) { + return hostEquals(segmentHost(target_segment_name), + segmentHost(local_server_name)); +} + +} // namespace mooncake + +#endif // MULTI_TRANSPORT_LOCALITY_H diff --git a/mooncake-transfer-engine/include/show_links.h b/mooncake-transfer-engine/include/show_links.h new file mode 100644 index 0000000000..41726ddef6 --- /dev/null +++ b/mooncake-transfer-engine/include/show_links.h @@ -0,0 +1,42 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MOONCAKE_SHOW_LINKS_H_ +#define MOONCAKE_SHOW_LINKS_H_ + +#include +#include +#include + +namespace mooncake { + +struct NicDiagInfo { + std::string device_name; + int numa_node; + std::string gid; + int gid_index; + uint8_t port; + int active_speed; + int active_width; +}; + +class TransferEngineImpl; + +std::string buildShowLinksJson(TransferEngineImpl* impl); + +std::string buildShowLinksReadable(TransferEngineImpl* impl); + +} // namespace mooncake + +#endif // MOONCAKE_SHOW_LINKS_H_ diff --git a/mooncake-transfer-engine/include/sunrise_allocator.h b/mooncake-transfer-engine/include/sunrise_allocator.h new file mode 100644 index 0000000000..be528e2630 --- /dev/null +++ b/mooncake-transfer-engine/include/sunrise_allocator.h @@ -0,0 +1,207 @@ +#pragma once + +#include +#include +#include + +#if defined(USE_SUNRISE) +#include +#include +#include +#include +#include +#endif + +#if defined(USE_SUNRISE) +namespace mooncake { +namespace sunrise_alloc_detail { +inline std::unordered_set& tangHostAllocatedSet() { + static std::unordered_set s; + return s; +} + +inline std::unordered_set& tangDeviceAllocatedSet() { + static std::unordered_set s; + return s; +} + +inline std::mutex& tangAllocMutex() { + static std::mutex m; + return m; +} + +struct MemRange { + void* base; + size_t size; +}; + +inline std::vector& storeMemRanges() { + static std::vector v; + return v; +} + +inline std::mutex& storeMemMutex() { + static std::mutex m; + return m; +} + +inline void addStoreMemRange(void* ptr, size_t size) { + std::lock_guard lock(storeMemMutex()); + storeMemRanges().push_back({ptr, size}); +} + +inline void removeStoreMemRange(void* ptr) { + std::lock_guard lock(storeMemMutex()); + auto& ranges = storeMemRanges(); + ranges.erase( + std::remove_if(ranges.begin(), ranges.end(), + [ptr](const MemRange& r) { return r.base == ptr; }), + ranges.end()); +} + +} // namespace sunrise_alloc_detail +} // namespace mooncake +#endif + +inline void* sunrise_allocate_memory(size_t total_size, size_t alignment = 4096, + bool use_device_mem = false) { +#if defined(USE_SUNRISE) + if (use_device_mem) { + int cur_dev = -1; + tangGetDevice(&cur_dev); + if (cur_dev < 0) { + tangError_t sd_ret = tangSetDevice(0); + if (sd_ret != tangSuccess) { + LOG(ERROR) << "sunrise_allocate_memory: tangSetDevice(0) " + << "failed: " << sd_ret << " " + << tangGetErrorString(sd_ret); + return nullptr; + } + } + void* ptr = nullptr; + tangError_t ret = tangMalloc(&ptr, total_size); + if (ret == tangSuccess && ptr) { + { + std::lock_guard guard( + mooncake::sunrise_alloc_detail::tangAllocMutex()); + mooncake::sunrise_alloc_detail::tangDeviceAllocatedSet().insert( + ptr); + } + mooncake::sunrise_alloc_detail::addStoreMemRange(ptr, total_size); + return ptr; + } + LOG(ERROR) << "tangMalloc failed (" << ret + << ") for size=" << total_size; + return nullptr; + } + + void* ptr = nullptr; + tangError_t ret = tangHostAlloc(&ptr, total_size, 0); + if (ret == tangSuccess && ptr) { + { + std::lock_guard guard( + mooncake::sunrise_alloc_detail::tangAllocMutex()); + mooncake::sunrise_alloc_detail::tangHostAllocatedSet().insert(ptr); + } + mooncake::sunrise_alloc_detail::addStoreMemRange(ptr, total_size); + return ptr; + } + LOG(WARNING) << "tangHostAlloc failed (" << ret + << "), falling back to posix_memalign for size=" << total_size + << " alignment=" << alignment; +#endif + (void)use_device_mem; + void* fallback = nullptr; + if (posix_memalign(&fallback, alignment, total_size) != 0) return nullptr; + return fallback; +} + +inline void sunrise_free_memory(void* ptr) { + if (!ptr) return; +#if defined(USE_SUNRISE) + { + std::lock_guard guard( + mooncake::sunrise_alloc_detail::tangAllocMutex()); + if (mooncake::sunrise_alloc_detail::tangDeviceAllocatedSet().count( + ptr)) { + mooncake::sunrise_alloc_detail::tangDeviceAllocatedSet().erase(ptr); + mooncake::sunrise_alloc_detail::removeStoreMemRange(ptr); + tangError_t ret = tangFree(ptr); + if (ret != tangSuccess) { + LOG(WARNING) + << "tangFree failed (" << ret << ") for tracked ptr=" << ptr + << " (leaking to avoid crash)"; + } + return; + } + if (mooncake::sunrise_alloc_detail::tangHostAllocatedSet().count(ptr)) { + mooncake::sunrise_alloc_detail::tangHostAllocatedSet().erase(ptr); + mooncake::sunrise_alloc_detail::removeStoreMemRange(ptr); + tangError_t ret = tangFreeHost(ptr); + if (ret != tangSuccess) { + LOG(WARNING) << "tangFreeHost failed (" << ret + << ") for tracked ptr=" << ptr + << " (leaking to avoid crash)"; + } + return; + } + } +#endif + free(ptr); +} + +inline bool sunrise_is_device_memory_range(void* ptr) { +#if defined(USE_SUNRISE) + if (!ptr) return false; + auto addr = reinterpret_cast(ptr); + void* matched_base = nullptr; + { + std::lock_guard lock( + mooncake::sunrise_alloc_detail::storeMemMutex()); + for (const auto& range : + mooncake::sunrise_alloc_detail::storeMemRanges()) { + auto base_addr = reinterpret_cast(range.base); + if (addr >= base_addr && addr < base_addr + range.size) { + matched_base = range.base; + break; + } + } + } + if (!matched_base) return false; + std::lock_guard guard( + mooncake::sunrise_alloc_detail::tangAllocMutex()); + return mooncake::sunrise_alloc_detail::tangDeviceAllocatedSet().count( + matched_base) > 0; +#else + (void)ptr; + return false; +#endif +} + +inline bool sunrise_is_host_allocated(void* ptr) { +#if defined(USE_SUNRISE) + if (!ptr) return false; + auto addr = reinterpret_cast(ptr); + void* matched_base = nullptr; + { + std::lock_guard lock( + mooncake::sunrise_alloc_detail::storeMemMutex()); + for (const auto& range : + mooncake::sunrise_alloc_detail::storeMemRanges()) { + auto base_addr = reinterpret_cast(range.base); + if (addr >= base_addr && addr < base_addr + range.size) { + matched_base = range.base; + break; + } + } + } + if (!matched_base) return false; + std::lock_guard guard( + mooncake::sunrise_alloc_detail::tangAllocMutex()); + return mooncake::sunrise_alloc_detail::tangHostAllocatedSet().count( + matched_base) > 0; +#else + (void)ptr; + return false; +#endif +} diff --git a/mooncake-transfer-engine/include/topology.h b/mooncake-transfer-engine/include/topology.h index 123c174920..2efbe0f4d5 100644 --- a/mooncake-transfer-engine/include/topology.h +++ b/mooncake-transfer-engine/include/topology.h @@ -30,6 +30,7 @@ #include #include #include +#include #include "common.h" @@ -39,6 +40,8 @@ struct TopologyEntry { std::vector preferred_hca; std::vector avail_hca; + bool operator==(const TopologyEntry &other) const = default; + Json::Value toJson() const { Json::Value matrix(Json::arrayValue); Json::Value hca_list(Json::arrayValue); @@ -83,9 +86,14 @@ class Topology { Json::Value toJson() const; + bool operator==(const Topology &other) const; + bool operator!=(const Topology &other) const { return !(*this == other); } + int selectDevice(const std::string storage_type, int retry_count = 0); int selectDevice(const std::string storage_type, std::string_view hint, int retry_count = 0); + int selectDeviceByLocalHca(const std::string storage_type, + std::string_view local_hca, int retry_count = 0); TopologyMatrix getMatrix() const { return matrix_; } @@ -121,8 +129,12 @@ class Topology { }; std::unordered_map resolved_matrix_; + std::unordered_map /* peer HCA ids */>> + resolved_hca_peer_affinity_by_local_; }; } // namespace mooncake -#endif // TOPOLOGY_H \ No newline at end of file +#endif // TOPOLOGY_H diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index 8d5b8d2e71..cfb06d3d22 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -15,22 +15,38 @@ #ifndef MULTI_TRANSFER_ENGINE_H_ #define MULTI_TRANSFER_ENGINE_H_ +#include +#include +#include +#include +#include +#include +#include +#include + #include "memory_location.h" #include "multi_transport.h" #include "transfer_metadata.h" #include "transport/transport.h" namespace mooncake { +class ShutdownToken; class TransferEngineImpl; namespace tent { class TransferEngine; }; -#if defined(USE_CUDA) || defined(USE_MUSA) +#if (defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)) && \ + !defined(USE_CXI) namespace device { class P2pTransport; class RdmaTransport; } // namespace device #endif +#ifdef USE_NCCL_DEVICE +namespace device { +class NcclTransport; +} // namespace device +#endif using TransferRequest = Transport::TransferRequest; using TransferStatus = Transport::TransferStatus; using TransferStatusEnum = Transport::TransferStatusEnum; @@ -39,12 +55,18 @@ using SegmentID = Transport::SegmentID; using BatchID = Transport::BatchID; const static BatchID INVALID_BATCH_ID = UINT64_MAX; using BufferEntry = Transport::BufferEntry; +using NicLoadStats = Transport::NicLoadStats; enum class PeerLiveness : uint8_t { Alive = 0, Unreachable = 1, }; +struct AutoDiscoverConfig { + bool enabled = false; + std::string protocol; +}; + class TransferEngine { public: #ifdef ENABLE_MULTI_PROTOCOL @@ -71,6 +93,10 @@ class TransferEngine { TransferEngine(bool auto_discover, const std::vector& filter); + TransferEngine(TransferEngine&& other) noexcept; + + TransferEngine& operator=(TransferEngine&& other) noexcept; + ~TransferEngine(); int init(const std::string& metadata_conn_string, @@ -88,6 +114,8 @@ class TransferEngine { int getRpcPort(); + bool isUsingTent() const { return use_tent_; } + SegmentHandle openSegment(const std::string& segment_name); Status CheckSegmentStatus(SegmentID sid); @@ -106,6 +134,49 @@ class TransferEngine { Status submitTransfer(BatchID batch_id, const std::vector& entries); + struct ScatterTransferRange { + TransferRequest::OpCode opcode; + std::string remote_segment; + uint64_t remote_base_offset; + size_t remote_size; + void* local_buffer; + size_t local_capacity; + std::span local_offsets; + std::span remote_offsets; + std::span lengths; + std::function on_fragment_complete; + }; + + class ScatterTransferOperation { + public: + ScatterTransferOperation(ScatterTransferOperation&&) noexcept; + ScatterTransferOperation& operator=( + ScatterTransferOperation&&) noexcept; + + // Destruction waits for physical completion before releasing state. + ~ScatterTransferOperation(); + + ScatterTransferOperation(const ScatterTransferOperation&) = delete; + ScatterTransferOperation& operator=(const ScatterTransferOperation&) = + delete; + + Status wait(); + + // A wait timeout does not cancel the transfer. Keep this operation and + // its CPU/GPU buffers alive until a later wait reaches completion. + Status waitFor(std::chrono::nanoseconds timeout); + + private: + class Impl; + explicit ScatterTransferOperation(std::unique_ptr impl); + std::unique_ptr impl_; + friend class TransferEngine; + }; + + ScatterTransferOperation submitScatter( + const std::vector& ranges); + Status transferScatter(const std::vector& ranges); + Status submitTransferWithNotify(BatchID batch_id, const std::vector& entries, TransferMetadata::NotifyDesc notify_msg); @@ -154,9 +225,12 @@ class TransferEngine { Status getBatchTransferStatus(BatchID batch_id, TransferStatus& status); + Status getNicLoadStats(std::vector& stats) const; + Transport* getTransport(const std::string& proto); -#if defined(USE_CUDA) || defined(USE_MUSA) +#if (defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)) && \ + !defined(USE_CXI) // Device transport accessors (P2P + IBGDA). Lazily created on first // call and owned by the TransferEngine. These allow EP (and future // CPU-proxy paths) to obtain device transports from an engine instance @@ -165,6 +239,10 @@ class TransferEngine { device::RdmaTransport* getOrCreateRdmaTransport( const std::vector& device_filter = {}); #endif +#ifdef USE_NCCL_DEVICE + // NCCL is CUDA-only and independent of the host network transport. + device::NcclTransport* getOrCreateNcclTransport(); +#endif /** * @brief Check if TCP is the only installed transport. @@ -181,6 +259,7 @@ class TransferEngine { bool checkOverlap(void* addr, uint64_t length); void setAutoDiscover(bool auto_discover); + void setAutoDiscover(const AutoDiscoverConfig& config); void* getBaseAddr(); @@ -190,9 +269,13 @@ class TransferEngine { std::shared_ptr getLocalTopology(); + void enableGracefulShutdown(); + std::string showLinks(bool json = false) const; + private: std::shared_ptr impl_; std::shared_ptr impl_tent_; + std::shared_ptr shutdown_token_; bool use_tent_{false}; }; } // namespace mooncake diff --git a/mooncake-transfer-engine/include/transfer_engine_c.h b/mooncake-transfer-engine/include/transfer_engine_c.h index 82baf8ee43..236a180a91 100644 --- a/mooncake-transfer-engine/include/transfer_engine_c.h +++ b/mooncake-transfer-engine/include/transfer_engine_c.h @@ -165,6 +165,21 @@ int freeBatchID(transfer_engine_t engine, batch_id_t batch_id); int syncSegmentCache(transfer_engine_t engine); +struct nic_load_stat { + char device_name[64]; + uint64_t inflight_bytes; + double ewma_bandwidth_bps; +}; + +typedef struct nic_load_stat nic_load_stat_t; + +int getNicLoadStats(transfer_engine_t engine, nic_load_stat_t *stats, + size_t *count); + +void enableGracefulShutdown(transfer_engine_t engine); +int showLinks(transfer_engine_t engine, char *buf_out, size_t buf_len, + int json); + #ifdef __cplusplus } #endif // __cplusplus diff --git a/mooncake-transfer-engine/include/transfer_engine_impl.h b/mooncake-transfer-engine/include/transfer_engine_impl.h index b1e4fff7e8..12c0bbf8de 100644 --- a/mooncake-transfer-engine/include/transfer_engine_impl.h +++ b/mooncake-transfer-engine/include/transfer_engine_impl.h @@ -33,15 +33,21 @@ #include "transfer_metadata.h" #include "transfer_engine.h" #include "transport/transport.h" -#if defined(USE_CUDA) || defined(USE_MUSA) +#if (defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)) && \ + !defined(USE_CXI) #include "transport/device/device_transport.h" #endif +#ifdef USE_NCCL_DEVICE +#include "transport/device/nccl_device_transport.h" +#endif #ifdef WITH_METRICS #include "ylt/metric/counter.hpp" #include "ylt/metric/histogram.hpp" #endif namespace mooncake { +class TransferEngineImplTestPeer; + using TransferRequest = Transport::TransferRequest; using TransferStatus = Transport::TransferStatus; using TransferStatusEnum = Transport::TransferStatusEnum; @@ -55,11 +61,13 @@ using RegisteredBuffer = TransferEngine::RegisteredBuffer; #endif class TransferEngineImpl { + friend class TransferEngineImplTestPeer; + public: TransferEngineImpl(bool auto_discover = false) : metadata_(nullptr), local_topology_(std::make_shared()), - auto_discover_(auto_discover) { + auto_discover_config_{.enabled = auto_discover, .protocol = ""} { #ifdef WITH_METRICS InitializeMetricsConfig(); StartMetricsReportingThread(); @@ -70,7 +78,7 @@ class TransferEngineImpl { const std::vector& filter) : metadata_(nullptr), local_topology_(std::make_shared()), - auto_discover_(auto_discover), + auto_discover_config_{.enabled = auto_discover, .protocol = ""}, filter_(filter) { #ifdef WITH_METRICS InitializeMetricsConfig(); @@ -341,15 +349,20 @@ class TransferEngineImpl { } Transport* getTransport(const std::string& proto) { + if (!multi_transports_) return nullptr; return multi_transports_->getTransport(proto); } -#if defined(USE_CUDA) || defined(USE_MUSA) +#if (defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)) && \ + !defined(USE_CXI) // Device transport accessors — lazily created, owned by this impl. device::P2pTransport* getOrCreateP2pTransport(int num_ranks); device::RdmaTransport* getOrCreateRdmaTransport( const std::vector& device_filter = {}); #endif +#ifdef USE_NCCL_DEVICE + device::NcclTransport* getOrCreateNcclTransport(); +#endif bool isTcpOnly() const { return multi_transports_->isTcpOnly(); } @@ -372,7 +385,13 @@ class TransferEngineImpl { void rollbackAllRegistrations(const std::vector& records); #endif - void setAutoDiscover(bool auto_discover) { auto_discover_ = auto_discover; } + void setAutoDiscover(bool auto_discover) { + auto_discover_config_ = {.enabled = auto_discover, .protocol = ""}; + } + + void setAutoDiscover(const AutoDiscoverConfig& config) { + auto_discover_config_ = config; + } void* getBaseAddr() { return multi_transports_->getBaseAddr(); } @@ -398,12 +417,16 @@ class TransferEngineImpl { using MemoryRegionMap = std::map; - MemoryRegionMap::iterator findMemoryRegionContaining(uintptr_t addr); + bool hasOverlapLocked(uintptr_t addr, uint64_t length) const; + + bool hasOverlapInMapLocked(const MemoryRegionMap& regions, uintptr_t addr, + uint64_t length) const; - MemoryRegionMap::const_iterator findMemoryRegionContaining( - uintptr_t addr) const; + bool tryReserveMemoryRegions(const std::vector& regions); - bool hasOverlapLocked(uintptr_t addr, uint64_t length) const; + void commitMemoryRegions(const std::vector& regions); + + void releaseMemoryRegions(const std::vector& regions); void insertMemoryRegionLocked(const MemoryRegion& region); @@ -414,6 +437,7 @@ class TransferEngineImpl { std::shared_ptr multi_transports_; std::shared_mutex mutex_; MemoryRegionMap local_memory_regions_; + MemoryRegionMap registering_memory_regions_; std::shared_ptr local_topology_; RWSpinlock send_notifies_lock_; @@ -421,18 +445,31 @@ class TransferEngineImpl { std::pair> notifies_to_send_; - // Discover topology and install transports automatically when it's true. - // Set it to false only for testing. - bool auto_discover_; + std::string autoDiscoverTransport() const { + if (use_barex_) { + return "barex"; + } + if (auto_discover_config_.protocol == "efa") { + return "efa"; + } + return "rdma"; + } + + // Discover topology and install transports automatically when enabled. + AutoDiscoverConfig auto_discover_config_; std::vector filter_; bool use_barex_ = false; -#if defined(USE_CUDA) || defined(USE_MUSA) +#if (defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)) && \ + !defined(USE_CXI) // Device transports (P2P + IBGDA) — lazily created, owned by this impl. // Referenced by EP and future CPU-proxy paths. std::unique_ptr p2p_transport_; std::unique_ptr rdma_transport_; #endif +#ifdef USE_NCCL_DEVICE + std::unique_ptr nccl_transport_; +#endif #ifdef WITH_METRICS // Latency bucket in microseconds diff --git a/mooncake-transfer-engine/include/transfer_engine_rpc_client_io_context.h b/mooncake-transfer-engine/include/transfer_engine_rpc_client_io_context.h new file mode 100644 index 0000000000..e1fa1319a5 --- /dev/null +++ b/mooncake-transfer-engine/include/transfer_engine_rpc_client_io_context.h @@ -0,0 +1,23 @@ +#pragma once + +#include "environ.h" +#include "rpc_client_io_context.h" + +namespace mooncake { + +inline uint32_t GetTransferEngineRpcClientIoThreads() { + return Environ::Get().GetTransferEngineRpcClientIoThreads(); +} + +namespace detail { +struct TransferEngineRpcClientIoContextPoolTag {}; +} // namespace detail + +inline coro_io::io_context_pool& GetTransferEngineRpcClientIoContextPool() { + static auto& io_pool = GetRpcClientIoContextPool< + detail::TransferEngineRpcClientIoContextPoolTag>( + GetTransferEngineRpcClientIoThreads()); + return io_pool; +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/include/transfer_metadata.h b/mooncake-transfer-engine/include/transfer_metadata.h index 0233a35ce4..69ce49da5f 100644 --- a/mooncake-transfer-engine/include/transfer_metadata.h +++ b/mooncake-transfer-engine/include/transfer_metadata.h @@ -24,9 +24,11 @@ #include #include +#include #include #include #include +#include #include #include #include @@ -47,27 +49,42 @@ class TransferMetadata { uint16_t lid; std::string gid; std::string eid; // for ub + + bool operator==(const DeviceDesc &other) const = default; }; struct BufferDesc { std::string name; uint64_t addr; uint64_t length; + int32_t device_id = -1; // CUDA device for NCCL buffers #ifdef ENABLE_MULTI_PROTOCOL std::string protocol; // for multi-protocol mode (cxl/tcp/rdma) #endif - std::vector lkey; // for rdma - std::vector rkey; // for rdma + // EFA/CXI's libfabric provider returns 64-bit MR keys (fi_mr_key()), so + // these must be 64-bit wide to avoid truncation. RDMA verbs keys are + // 32-bit and the non-EFA/CXI path keeps them as such. +#if defined(USE_EFA) || defined(USE_CXI) + using mr_key_t = uint64_t; +#else + using mr_key_t = uint32_t; +#endif + std::vector lkey; // for rdma/efa + std::vector rkey; // for rdma/efa std::string shm_name; // for nvlink and hip uint64_t offset; // for cxl std::vector tseg; // for ub/urma std::vector l_seg_index; // for ub/urma + + bool operator==(const BufferDesc &other) const = default; }; struct NVMeoFBufferDesc { std::string file_path; uint64_t length; std::unordered_map local_path_map; + + bool operator==(const NVMeoFBufferDesc &other) const = default; }; struct RankInfoDesc { @@ -81,6 +98,8 @@ class TransferMetadata { uint64_t devicePort; uint64_t pid; std::vector endpoints; + + bool operator==(const RankInfoDesc &other) const = default; }; using SegmentID = uint64_t; @@ -103,6 +122,10 @@ class TransferMetadata { RankInfoDesc rank_info; int tcp_data_port; + // TCP data-plane protocol version advertised by this segment's + // server. v2 adds acknowledged WRITE framing and status-prefixed + // READ responses (#2086); absent/1 = legacy unacknowledged framing. + int tcp_proto_version{1}; // In dual-NIC setups (MC_RDMA_BIND_ADDRESS), the RDMA-reachable // address may differ from the TCP-routable segment name. When @@ -117,6 +140,10 @@ class TransferMetadata { return rdma_server_name.empty() ? name : rdma_server_name; } + bool operator==(const SegmentDesc &other) const; + bool operator!=(const SegmentDesc &other) const { + return !(*this == other); + } void dump() const; }; @@ -130,6 +157,7 @@ class TransferMetadata { }; struct HandShakeDesc { + std::string payload; // opaque transport-specific handshake data std::string local_nic_path; uint16_t local_lid = 0; std::string local_gid; @@ -141,9 +169,16 @@ class TransferMetadata { uint16_t barex_port; #endif std::vector qp_num; + bool ready_ack = false; + // Capability marker. Encoded only by transports that opt into + // ready_ack; decoded from field presence to detect peer support. + bool ready_ack_supported = false; std::string reply_msg; // on error #ifdef USE_EFA std::string efa_addr; // EFA endpoint address (hex encoded) +#endif +#ifdef USE_CXI + std::string cxi_addr; #endif }; @@ -227,6 +262,9 @@ class TransferMetadata { Json::Value &local_json); int receivePeerProbe(const Json::Value &peer_json, Json::Value &local_json); std::string getFullMetadataKey(const std::string &segment_name) const; + void startMetadataRefreshPollingIfNeeded(); + void stopMetadataRefreshPollingThread(); + void metadataRefreshPollingLoop(uint64_t refresh_interval_seconds); bool p2p_handshake_mode_{false}; std::string common_key_prefix_; @@ -247,6 +285,10 @@ class TransferMetadata { std::shared_ptr handshake_plugin_; std::shared_ptr storage_plugin_; + std::mutex metadata_refresh_mutex_; + std::condition_variable metadata_refresh_cv_; + std::atomic should_stop_metadata_refresh_thread_{false}; + std::thread metadata_refresh_thread_; }; } // namespace mooncake diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.h b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.h index 01d30af97e..4b96d9c27e 100644 --- a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.h +++ b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.h @@ -58,6 +58,9 @@ class AsyncTransferExecutor : public TransferExecutorBase { // pending_batches for this route. Batches enqueued after disconnect // are not affected. bool fail_entire_route = false; + // When true with auto_connect, ADXL DisconnectOnError already ran; + // only forgetConnectedSegment is needed (if the route was recorded). + bool adxl_already_disconnected = false; std::string fail_reason; }; @@ -66,7 +69,8 @@ class AsyncTransferExecutor : public TransferExecutorBase { void processOneBatch(QueryBatch& batch, BatchPollResult& result); void failAllPendingOnRoute(size_t engine_idx, const std::string& target, std::vector& pending, - const std::string& reason); + const std::string& reason, + bool adxl_already_disconnected); void markBatchFailed(QueryBatch& batch, const std::string& reason, bool log_error = true); void handleTaskFinished(); diff --git a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.h b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.h index 3c84f9757e..e07e97162e 100644 --- a/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.h +++ b/mooncake-transfer-engine/include/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.h @@ -95,6 +95,7 @@ class TransferExecutorBase { void finalizeEngines(); void disconnectAllForEngine(size_t engine_idx); void recordConnectedSegment(size_t engine_idx, const std::string& remote); + void forgetConnectedSegment(size_t engine_idx, const std::string& remote); int checkAndConnect(size_t engine_idx, const std::string& target_adxl_engine_name); diff --git a/mooncake-transfer-engine/include/transport/cxi_transport/cxi_context.h b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_context.h new file mode 100644 index 0000000000..e606e45b1a --- /dev/null +++ b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_context.h @@ -0,0 +1,216 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef CXI_CONTEXT_H +#define CXI_CONTEXT_H + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "cxi_transport.h" +#include "transport/transport.h" + +namespace mooncake { + +class CxiEndpoint; +class CxiTransport; + +struct CxiCq { + CxiCq() : cq(nullptr), outstanding(0) {} + struct fid_cq* cq; + std::atomic outstanding; +}; + +struct CxiMemoryRegionMeta { + void* addr; + size_t length; + struct fid_mr* mr; + uint64_t key; + bool associated; +}; + +// CxiContext represents the set of resources controlled by each local Cxi +// device, including Memory Region, CQ, EndPoint, etc. using libfabric +class CxiContext { + public: + CxiContext(CxiTransport& engine, const std::string& device_name); + + ~CxiContext(); + + int construct(size_t num_cq_list = 1, size_t max_cqe = 4096, + int max_endpoints = 65536); + + private: + int deconstruct(); + int buildSharedEndpoint(size_t max_wr, size_t max_inline); + + public: + // Memory Region Management + int registerMemoryRegion(void* addr, size_t length, int access); + int unregisterMemoryRegion(void* addr); + int preTouchMemory(void* addr, size_t length); + uint64_t rkey(void* addr); + uint64_t lkey(void* addr); + void* mrDesc(void* addr); // Get MR descriptor for fi_write local_desc + + private: + int registerMemoryRegionInternal(void* addr, size_t length, int access, + CxiMemoryRegionMeta& mrMeta); + + public: + bool active() const { return active_; } + void set_active(bool flag) { active_ = flag; } + + public: + // Get or create a per-peer handle. Does NOT open an fid_ep or call + // fi_enable — the shared endpoint was created once at construct() time. + // The returned CxiEndpoint only carries {peer_fi_addr_t, status, mutex} + // and, when connected, routes sends through this context's shared_ep_. + std::shared_ptr endpoint(const std::string& peer_nic_path); + + // Non-creating lookup under the normalized key. Returns nullptr if the + // peer handle does not yet exist. Safe for idempotency checks. + std::shared_ptr peekEndpoint(const std::string& peer_nic_path); + + int deleteEndpoint(const std::string& peer_nic_path); + int disconnectAllEndpoints(); + + // Number of live peer handles. Historically named "QP number"; with the + // shared endpoint model the actual QP count is always 1 per context. + size_t getTotalQPNumber() const; + + public: + // Access to engine for endpoint handshake + CxiTransport& engine() { return engine_; } + const CxiTransport& engine() const { return engine_; } + + // Submit slices for transfer + int submitPostSend(const std::vector& slice_list); + + // Hot-path submit: post a batch of slices to `peer_fi_addr` via the + // shared endpoint. Handles WR / CQ reservation, MR descriptor prep, + // and the fi_write / fi_read burst under post_lock_. Called by + // CxiEndpoint::submitPostSend once the peer is connected. + int submitSlicesOnPeer(fi_addr_t peer_fi_addr, + std::vector& slice_list, + std::vector& failed_slice_list); + + // Poll completion queue for completed operations + int pollCq(int max_entries, int cq_index = 0); + + // Get CQ count + size_t cqCount() const { return cq_list_.size(); } + + public: + // Device name, such as `rdmap0s2` + std::string deviceName() const { return device_name_; } + + // NIC Path, such as `192.168.3.76@rdmap0s2` + std::string nicPath() const; + + public: + // Libfabric accessors + struct fid_fabric* fabric() const { return fabric_; } + struct fid_domain* domain() const { return domain_; } + struct fid_av* av() const { return av_; } + struct fi_info* info() const { return fi_info_; } + std::string localAddr() const; + + // Local (shared-endpoint) address in hex, for inclusion in handshake. + // Populated by construct() -> buildSharedEndpoint(). + std::string localEpAddr() const; + + // Raw bytes of the local endpoint address. Use this for loopback + // (skip the hex encode/decode round-trip) or for any caller that + // already has the bytes. + const std::vector& localEpAddrBytes() const { + return local_ep_addr_; + } + + // Insert a peer's hex-encoded CXI address into this context's AV and + // return the resulting fi_addr_t. Thread-safe (fi_av_insert is safe + // under libfabric's domain-level threading). + int insertPeerAddr(const std::string& peer_hex_addr, fi_addr_t& out); + + // Binary variant — avoids the hex-decode when the caller already + // has the raw address bytes (e.g. loopback). + int insertPeerAddrBytes(const uint8_t* addr, size_t len, fi_addr_t& out); + + // Remove a peer from the AV. No-op if fi_addr is FI_ADDR_UNSPEC. + void removePeerAddr(fi_addr_t fi_addr); + + // Compatibility methods (libfabric doesn't use lid/gid like ibverbs) + uint16_t lid() const { return 0; } + std::string gid() const { return localAddr(); } + + private: + CxiTransport& engine_; + std::string device_name_; + + // Libfabric objects + struct fi_info* fi_info_; + struct fi_info* hints_; + struct fid_fabric* fabric_; + struct fid_domain* domain_; + struct fid_av* av_; // Address vector for peer addressing + + bool active_; + + // ---- Shared endpoint (one per local NIC, serves ALL peers) ---- + struct fid_ep* shared_ep_; + std::vector local_ep_addr_; // bytes returned by fi_getname() + // Pacing for outstanding work requests on the shared endpoint. Shared + // across all peers routed through this context. std::atomic so the + // submit-path fetch_add and the CQ-poller fetch_sub obey the C++ memory + // model; plain `volatile int` + __sync_* was UB under the current + // standard. + std::atomic wr_depth_; + int max_wr_depth_; + // CQ that shared_ep_ is bound to (FI_TRANSMIT|FI_RECV). Points into + // cq_list_[0]; kept here to avoid re-indexing on the hot path. + std::shared_ptr shared_cq_; + // Serializes fi_write / fi_read calls on shared_ep_. libfabric's RDM + // endpoints are not thread-safe for concurrent post, even with + // FI_THREAD_SAFE at the domain level. + std::atomic_flag post_lock_; + + std::vector> cq_list_; + + // ---- Peer handles (one entry per peer, each ~constant size) ---- + mutable RWSpinlock peer_map_lock_; + std::unordered_map> peer_map_; + + RWSpinlock mr_lock_; + std::map mr_map_; +}; + +} // namespace mooncake + +#endif // CXI_CONTEXT_H diff --git a/mooncake-transfer-engine/include/transport/cxi_transport/cxi_endpoint.h b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_endpoint.h new file mode 100644 index 0000000000..9a9febaf3a --- /dev/null +++ b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_endpoint.h @@ -0,0 +1,132 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef CXI_ENDPOINT_H +#define CXI_ENDPOINT_H + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "common.h" +#include "cxi_context.h" +#include "transfer_metadata.h" +#include "transport/transport.h" + +namespace mooncake { + +class CxiContext; + +// Custom context for libfabric operations - stores slice pointer for +// completion handling. This struct MUST have fi_context as its first member. +struct CxiOpContext { + struct fi_context fi_ctx; // Must be first member + Transport::Slice* slice; // Slice pointer for completion handling + // Pointer to the context's wr_depth_ for CQ completion decrement. + // std::atomic (not volatile int) so the CQ-poller decrement + // and the submit-path fetch_add play by the C++ memory model. + std::atomic* wr_depth; +}; + +// Per-peer handle in the shared-endpoint SRD model. +// +// This class does NOT own an fid_ep. A single `shared_ep_` on the owning +// CxiTransport services every peer; each CxiEndpoint just carries one +// fi_addr_t (AV index) plus the handshake state needed to populate it. +// +// Lifecycle: +// 1. construct() : trivial; no libfabric resources allocated. +// 2. setupConnectionsByActive(peer): handshake -> fi_av_insert -> CONNECTED. +// 3. submitPostSend() : delegates to CxiTransport::submitSlicesOnPeer. +// 4. disconnect() / dtor : fi_av_remove(peer_fi_addr_). +class CxiEndpoint { + public: + using HandShakeDesc = TransferMetadata::HandShakeDesc; + + enum Status { INITIALIZING, UNCONNECTED, CONNECTED }; + + explicit CxiEndpoint(CxiContext& context); + ~CxiEndpoint(); + + public: + void setPeerNicPath(const std::string& peer_nic_path); + + int setupConnectionsByActive(); + + int setupConnectionsByActive(const std::string& peer_nic_path) { + setPeerNicPath(peer_nic_path); + return setupConnectionsByActive(); + } + + int setupConnectionsByPassive(const HandShakeDesc& peer_desc, + HandShakeDesc& local_desc); + + // Always false under the shared-endpoint model: outstanding work is + // tracked at the context (shared) level, not per-peer. Retained for + // API compatibility with callers that still ask. + bool hasOutstandingSlice() const { return false; } + + bool connected() const { + return status_.load(std::memory_order_relaxed) == CONNECTED; + } + + void disconnect(); + + // Called during CxiTransport teardown: forget the AV slot WITHOUT calling + // fi_av_remove(). The AV itself is about to be closed, which invalidates + // every slot in one shot; calling fi_av_remove after the shared endpoint + // has been closed trips an assertion inside the CXI provider. + void markDetachedForTeardown(); + + private: + void disconnectUnlocked(); + + public: + const std::string toString() const; + + // Submit a batch of slices bound for this peer. Internally establishes + // the connection if needed, then delegates to + // CxiTransport::submitSlicesOnPeer using this peer's fi_addr_t. + int submitPostSend(std::vector& slice_list, + std::vector& failed_slice_list); + + // Always 1: the shared endpoint backs every peer. Kept so callers that + // sum "QP count" across endpoints don't break. + size_t getQPNumber() const { return 1; } + + fi_addr_t getPeerFiAddr() const { return peer_fi_addr_; } + + CxiContext& context() { return context_; } + + private: + CxiContext& context_; + std::atomic status_; + + RWSpinlock lock_; // protects peer_nic_path_ and status_ + std::string peer_nic_path_; + fi_addr_t peer_fi_addr_; // slot in context_.av() +}; + +} // namespace mooncake + +#endif // CXI_ENDPOINT_H diff --git a/mooncake-transfer-engine/include/transport/cxi_transport/cxi_transport.h b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_transport.h new file mode 100644 index 0000000000..0243b0384b --- /dev/null +++ b/mooncake-transfer-engine/include/transport/cxi_transport/cxi_transport.h @@ -0,0 +1,174 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef CXI_TRANSPORT_H_ +#define CXI_TRANSPORT_H_ + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "topology.h" +#include "transfer_metadata.h" +#include "transport/transport.h" + +namespace mooncake { + +class CxiContext; +class CxiEndpoint; +class TransferMetadata; + +enum NicReplicaPolicy { REPLICATE_ALL, NUMA_AWARE, DEBUG }; + +class CxiTransport : public Transport { + friend class CxiContext; + friend class CxiEndpoint; + + public: + using BufferDesc = TransferMetadata::BufferDesc; + using SegmentDesc = TransferMetadata::SegmentDesc; + using HandShakeDesc = TransferMetadata::HandShakeDesc; + + public: + CxiTransport(); + + ~CxiTransport(); + + int install(std::string& local_server_name, + std::shared_ptr meta, + std::shared_ptr topo) override; + + const char* getName() const override { return "cxi"; } + + int registerLocalMemory(void* addr, size_t length, + const std::string& location, bool remote_accessible, + bool update_metadata) override; + + int unregisterLocalMemory(void* addr, bool update_metadata = true) override; + + int registerLocalMemoryBatch(const std::vector& buffer_list, + const std::string& location) override; + + int unregisterLocalMemoryBatch( + const std::vector& addr_list) override; + + // Eagerly populate the address vector with every (local_ctx, peer_nic) + // handshake for `segment_name`. + // + // Under the shared-endpoint model each warmup is handshake RPC + + // fi_av_insert — no fi_endpoint/fi_enable. This removes the first- + // submit stall (handshakes on the data path) without consuming QPs. + // + // Safe to call multiple times (idempotent: endpoint() + setup are both + // idempotent). Re-run after any openSegment() on a new peer. + int warmupSegment(const std::string& segment_name); + + private: + // Internal version with force_sequential option to avoid nested parallelism + int registerLocalMemoryInternal(void* addr, size_t length, + const std::string& location, + bool remote_accessible, + bool update_metadata, + bool force_sequential); + + int unregisterLocalMemoryInternal(void* addr, bool update_metadata, + bool force_sequential); + + // TRANSFER + + Status submitTransfer(BatchID batch_id, + const std::vector& entries) override; + + Status submitTransferTask( + const std::vector& task_list) override; + + Status getTransferStatus(BatchID batch_id, + std::vector& status); + + Status getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus& status) override; + + SegmentID getSegmentID(const std::string& segment_name); + + NicReplicaPolicy getReplicaPolicy(); + + private: + int allocateLocalSegmentID(); + + int preTouchMemory(void* addr, size_t length); + + public: + int onSetupCxiConnections(const HandShakeDesc& peer_desc, + HandShakeDesc& local_desc); + + int sendHandshake(const std::string& peer_server_name, + const HandShakeDesc& local_desc, + HandShakeDesc& peer_desc) { + return metadata_->sendHandshake(peer_server_name, local_desc, + peer_desc); + } + + const std::string& local_server_name() const { return local_server_name_; } + + std::shared_ptr meta() { return metadata_; } + + private: + int initializeCxiResources(); + + int startHandshakeDaemon(std::string& local_server_name); + + public: + static int selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, + int& buffer_id, int& device_id, int retry_cnt = 0); + static int selectDevice(SegmentDesc* desc, uint64_t offset, size_t length, + std::string_view hint, int& buffer_id, + int& device_id, int retry_cnt = 0); + + private: + // Start/stop CQ polling worker threads + void startWorkerThreads(); + void stopWorkerThreads(); + void workerThreadFunc(int thread_id); + + private: + std::vector> context_list_; + std::shared_ptr local_topology_; + + // Track chunked MR registrations for per-NIC partitioned buffers. + // When a buffer exceeds max_mr_size, it is split into chunks, each + // registered on a disjoint subset of NICs (per-NIC partition). + struct ChunkRegistration { + uint64_t addr; + std::vector nic_indices; + }; + std::mutex chunk_map_mutex_; + std::unordered_map> chunk_map_; + + // CQ polling worker threads + std::atomic worker_running_{false}; + std::vector worker_threads_; +}; + +} // namespace mooncake + +#endif // CXI_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/include/transport/device/comm_device.cuh b/mooncake-transfer-engine/include/transport/device/comm_device.cuh index 29e6c6c6f4..a3be4bd619 100644 --- a/mooncake-transfer-engine/include/transport/device/comm_device.cuh +++ b/mooncake-transfer-engine/include/transport/device/comm_device.cuh @@ -36,7 +36,11 @@ __device__ __forceinline__ CommCtx make_comm_ctx( ctx.p2p.peer_ptrs = ipc_peer_ptrs; ctx.p2p.local_base = gdr_buffer; +#ifdef MOONCAKE_EP_USE_MACA + ctx.ibgda.qp_devctxs = qp_devctxs; +#else ctx.ibgda.qp_devctxs = reinterpret_cast(qp_devctxs); +#endif ctx.ibgda.raddrs = reinterpret_cast(raddrs); ctx.ibgda.rkeys = reinterpret_cast(rkeys); ctx.ibgda.local_atomic_base = rdma_send_signal_buffer; diff --git a/mooncake-transfer-engine/include/transport/device/device_ops.cuh b/mooncake-transfer-engine/include/transport/device/device_ops.cuh index 2835f81533..29bf4b7d01 100644 --- a/mooncake-transfer-engine/include/transport/device/device_ops.cuh +++ b/mooncake-transfer-engine/include/transport/device/device_ops.cuh @@ -7,6 +7,8 @@ #ifdef MOONCAKE_EP_USE_MUSA #include "transport/device/musa/musa_ops.cuh" +#elif defined(MOONCAKE_EP_USE_MACA) +#include "transport/device/maca/maca_ops.cuh" #else #include "transport/device/cuda/cuda_ops.cuh" #endif diff --git a/mooncake-transfer-engine/include/transport/device/device_transport.h b/mooncake-transfer-engine/include/transport/device/device_transport.h index 82556e7dbf..4b71860c2f 100644 --- a/mooncake-transfer-engine/include/transport/device/device_transport.h +++ b/mooncake-transfer-engine/include/transport/device/device_transport.h @@ -55,16 +55,18 @@ class P2pTransport { // Free a buffer previously returned by allocateBuffer. virtual void freeBuffer(void* ptr) = 0; - // Export an IPC handle for the buffer allocated by this rank. - // Returns a byte blob (serialised as int32_t array for Python compat). - // Returns empty vector if IPC is not needed (e.g. fabric memory). + // Export the metadata peers need to access this rank's buffer. + // Returns an opaque byte blob (serialised as int32_t array for Python + // compat). For classic IPC this is a cudaIpcMemHandle-like payload; for + // fabric-backed paths this carries a CUmemFabricHandle plus mapping size. virtual std::vector exportIpcHandle(void* ptr) = 0; - // Import peer IPC handles and populate the device-visible tables. - // remote_handles[i] is the handle exported by rank i (may be empty for - // ranks that use fabric memory or are on a different node). + // Import peer access metadata and populate the device-visible tables. + // remote_handles[i] is the blob exported by rank i. // active_ranks_mask[i] == 1 means rank i is participating. // After this call, availableTablePtr() and peerPtrsTablePtr() are valid. + // The tables are transport-scoped: a later call replaces the previously + // imported peer mapping context. virtual void importPeerHandles( void* local_ptr, int rank, int num_ranks, const std::vector>& remote_handles, diff --git a/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h index 82030954d2..24154e6bf1 100644 --- a/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h +++ b/mooncake-transfer-engine/include/transport/device/ibgda/mlx5gda.h @@ -27,6 +27,19 @@ struct mlx5gda_wqebb { uint64_t qwords[8]; // 64 bytes }; +struct mlx5gda_control_region { + void *addr; + size_t size; + struct mlx5dv_devx_umem *umem; +}; + +struct mlx5gda_control_region_allocator { + void *context; + int (*allocate)(void *context, size_t size, + struct mlx5gda_control_region *region); + void (*release)(void *context, struct mlx5gda_control_region *region); +}; + struct mlx5gda_rdma_write_wqe { struct mlx5_wqe_ctrl_seg ctrl; struct mlx5_wqe_raddr_seg raddr; @@ -47,13 +60,18 @@ struct mlx5gda_cq { uint32_t cqe; size_t cq_offset; size_t dbr_offset; + void *cq_buf; + void *dbr; + struct mlx5gda_control_region cq_region; + struct mlx5gda_control_region dbr_region; + struct mlx5gda_control_region_allocator region_allocator; }; -struct mlx5gda_cq *mlx5gda_create_cq(void *ctrl_buf, - struct mlx5dv_devx_umem *ctrl_buf_umem, - struct memheap *ctrl_buf_heap, - struct ibv_pd *pd, int num_cqe, - cudaStream_t stream); +struct mlx5gda_cq *mlx5gda_create_cq( + void *ctrl_buf, struct mlx5dv_devx_umem *ctrl_buf_umem, + struct memheap *ctrl_buf_heap, struct ibv_pd *pd, int num_cqe, + cudaStream_t stream, + const struct mlx5gda_control_region_allocator *region_allocator); void mlx5gda_destroy_cq(struct memheap *ctrl_buf_heap, struct mlx5gda_cq *cq); static const size_t MLX5GDA_BF_SIZE = 256; @@ -72,6 +90,11 @@ struct mlx5gda_qp { uint32_t num_wqebb; size_t wq_offset; size_t dbr_offset; + void *wq; + void *dbr; + struct mlx5gda_control_region wq_region; + struct mlx5gda_control_region dbr_region; + struct mlx5gda_control_region_allocator region_allocator; }; struct mlx5gda_qp_devctx { @@ -87,11 +110,21 @@ struct mlx5gda_qp_devctx { uint16_t wq_tail; // last non-completed wqeid }; -struct mlx5gda_qp *mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void *ctrl_buf, - struct mlx5dv_devx_umem *ctrl_buf_umem, - struct memheap *ctrl_buf_heap, - struct ibv_pd *pd, int wqe, - uint8_t port_num, cudaStream_t stream); +struct mlx5gda_create_qp_failure { + bool valid; + uint32_t status; + uint32_t syndrome; + int sys_errno; +}; + +void mlx5gda_reset_create_qp_failure(); +mlx5gda_create_qp_failure mlx5gda_last_create_qp_failure(); + +struct mlx5gda_qp *mlx5gda_create_rc_qp( + struct mlx5dv_pd mpd, void *ctrl_buf, + struct mlx5dv_devx_umem *ctrl_buf_umem, struct memheap *ctrl_buf_heap, + struct ibv_pd *pd, int wqe, uint8_t port_num, cudaStream_t stream, + const struct mlx5gda_control_region_allocator *region_allocator); void mlx5gda_destroy_qp(struct memheap *ctrl_buf_heap, struct mlx5gda_qp *qp); int mlx5gda_modify_rc_qp_rst2init(struct mlx5gda_qp *qp, uint16_t pkey_index); diff --git a/mooncake-transfer-engine/include/transport/device/ibgda_device.cuh b/mooncake-transfer-engine/include/transport/device/ibgda_device.cuh index 15fcb990c8..2376f3b42e 100644 --- a/mooncake-transfer-engine/include/transport/device/ibgda_device.cuh +++ b/mooncake-transfer-engine/include/transport/device/ibgda_device.cuh @@ -11,6 +11,32 @@ #include #include "transport/device/device_ops.cuh" +#ifdef MOONCAKE_EP_USE_MACA + +namespace mooncake { +namespace device { + +struct IbgdaContext { + void* qp_devctxs; + const uint64_t* raddrs; + const uint32_t* rkeys; + const void* local_atomic_base; + const void* remote_atomic_base; +}; + +__device__ __forceinline__ void mc_ibgda_put(const IbgdaContext&, int, int, int, + int, const void*, uint64_t, + uint32_t) {} + +__device__ __forceinline__ void mc_ibgda_red_add(const IbgdaContext&, int, int, + int, int, uint64_t, uint64_t, + int32_t) {} + +} // namespace device +} // namespace mooncake + +#else // !MOONCAKE_EP_USE_MACA + #ifndef MOONCAKE_EP_USE_MUSA #include #endif @@ -204,3 +230,5 @@ __device__ __forceinline__ void mc_ibgda_red_add( } // namespace device } // namespace mooncake + +#endif // MOONCAKE_EP_USE_MACA diff --git a/mooncake-transfer-engine/include/transport/device/maca/maca_ops.cuh b/mooncake-transfer-engine/include/transport/device/maca/maca_ops.cuh new file mode 100644 index 0000000000..50fdff72f7 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/device/maca/maca_ops.cuh @@ -0,0 +1,96 @@ +// MACA implementations of device-side memory ordering primitives. +// +// MACA's cu-bridge compiler accepts CUDA-like intrinsics, but does not reliably +// compile the PTX acquire/release/barrier instructions used by the CUDA path. +#pragma once + +#include + +namespace mooncake { +namespace device { + +__device__ __forceinline__ int mc_ld_acquire(const int* ptr) { + __threadfence_system(); + return *const_cast(ptr); +} + +__device__ __forceinline__ uint64_t mc_ld_acquire_u64(const uint64_t* ptr) { + __threadfence_system(); + return *const_cast(ptr); +} + +__device__ __forceinline__ void mc_st_release(const int* ptr, int val) { + *const_cast(ptr) = val; + __threadfence_system(); +} + +__device__ __forceinline__ void mc_st_release_u32(const uint32_t* ptr, + uint32_t val) { + *const_cast(ptr) = val; + __threadfence_system(); +} + +__device__ __forceinline__ void mc_st_release_u64(const uint64_t* ptr, + uint64_t val) { + *const_cast(ptr) = val; + __threadfence_system(); +} + +__device__ __forceinline__ int mc_atomic_add_release(const int* ptr, int val) { + int ret = atomicAdd(const_cast(ptr), val); + __threadfence_system(); + return ret; +} + +__device__ __forceinline__ int4 mc_ld_nc(const int4* ptr) { return __ldg(ptr); } + +__device__ __forceinline__ int mc_ld_nc_s32(const int* ptr) { + return __ldg(ptr); +} + +__device__ __forceinline__ float mc_ld_nc_f32(const float* ptr) { + return __ldg(ptr); +} + +__device__ __forceinline__ int64_t mc_ld_nc_s64(const int64_t* ptr) { + return __ldg(ptr); +} + +__device__ __forceinline__ void mc_st_na(const int4* ptr, const int4& val) { + *const_cast(ptr) = val; +} + +__device__ __forceinline__ void mc_bar_init() {} + +__device__ __forceinline__ void mc_bar_sync(int /*bar_id*/, + int /*num_threads*/) { + __syncthreads(); +} + +__device__ __forceinline__ void mc_grid_sync() {} + +__device__ __forceinline__ void mc_fence() { __threadfence_system(); } + +__device__ __forceinline__ void mc_fence_barrier_fence() { + mc_fence(); + mc_bar_sync(0, 0); + mc_fence(); +} + +__device__ __forceinline__ uint16_t mc_bswap16(uint16_t x) { + return (uint16_t)(((x & 0x00FFu) << 8) | ((x & 0xFF00u) >> 8)); +} + +__device__ __forceinline__ uint32_t mc_bswap32(uint32_t x) { + return ((x & 0x000000FFu) << 24) | ((x & 0x0000FF00u) << 8) | + ((x & 0x00FF0000u) >> 8) | ((x & 0xFF000000u) >> 24); +} + +__device__ __forceinline__ uint64_t mc_bswap64(uint64_t x) { + uint32_t hi = mc_bswap32((uint32_t)(x >> 32)); + uint32_t lo = mc_bswap32((uint32_t)(x)); + return ((uint64_t)lo << 32) | hi; +} + +} // namespace device +} // namespace mooncake diff --git a/mooncake-transfer-engine/include/transport/device/nccl_device.cuh b/mooncake-transfer-engine/include/transport/device/nccl_device.cuh new file mode 100644 index 0000000000..ce84adba15 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/device/nccl_device.cuh @@ -0,0 +1,302 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#ifndef USE_NCCL_DEVICE +#error "nccl_device.cuh requires USE_NCCL_DEVICE" +#endif + +#include +#include + +#include +#include + +#include "transport/device/nccl_device_transport.h" + +#if NCCL_VERSION_CODE < 23004 +#error "Mooncake NCCL DeviceTransport requires NCCL 2.30.4 or newer" +#endif + +// NCCL device code uses version-specific Device API definitions and layouts +// from nccl_device.h. Mooncake therefore requires these headers to exactly +// match the loaded runtime libnccl. After an NCCL upgrade, rebuild AOT kernels +// that include this header and invalidate and re-JIT cached NCCL Device API +// kernels before running. + +namespace mooncake { +namespace device { +namespace detail { + +struct NcclDeviceContextAccess { + __device__ __forceinline__ static const ncclDevComm_t& comm( + const NcclDeviceContext& ctx) { + return *static_cast(ctx.native_comm_); + } + + __device__ __forceinline__ static ncclWindow_t window( + const NcclDeviceContext& ctx) { + return reinterpret_cast( + const_cast(ctx.native_window_)); + } + + __device__ __forceinline__ static const char* localBase( + const NcclDeviceContext& ctx) { + return static_cast(ctx.local_base_); + } + + __device__ __forceinline__ static int rank(const NcclDeviceContext& ctx) { + return ctx.rank_; + } + + __device__ __forceinline__ static int ginContextCount( + const NcclDeviceContext& ctx) { + return ctx.gin_context_count_; + } + + __device__ __forceinline__ static bool ginEnabled( + const NcclDeviceContext& ctx) { + return ctx.gin_enabled_; + } + + __device__ __forceinline__ static bool lsaMultimemEnabled( + const NcclDeviceContext& ctx) { + return ctx.lsa_multimem_enabled_; + } +}; + +__device__ __forceinline__ size_t +mc_nccl_pointer_offset(const NcclDeviceContext& ctx, const void* ptr) { + return static_cast(static_cast(ptr) - + NcclDeviceContextAccess::localBase(ctx)); +} + +__device__ __forceinline__ int mc_nccl_gin_context(const NcclDeviceContext& ctx, + unsigned int channel) { + const int count = NcclDeviceContextAccess::ginContextCount(ctx); + return static_cast(channel % static_cast(count)); +} + +} // namespace detail + +// Device operations in this header are intentionally unchecked, matching the +// existing Mooncake and NCCL device APIs. The caller must provide a live +// context, a valid and reachable rank, the required NCCL resources, and pointer +// ranges wholly contained in the registration bound to ctx. Violating a +// precondition is undefined behavior. Hoist capability and route queries out of +// hot loops when the route is already known. + +// Pointer and route queries are per-thread operations. peer must be a valid +// world rank. +__device__ __forceinline__ bool mc_nccl_lsa_available( + const NcclDeviceContext& ctx, int peer) { + const auto& comm = detail::NcclDeviceContextAccess::comm(ctx); + return ncclTeamRankIsMember(ncclTeamLsa(comm), ncclTeamWorld(comm), peer); +} + +__device__ __forceinline__ bool mc_nccl_gin_available( + const NcclDeviceContext& ctx) { + return detail::NcclDeviceContextAccess::ginEnabled(ctx); +} + +__device__ __forceinline__ NcclDeviceRoute +mc_nccl_route(const NcclDeviceContext& ctx, int peer) { + if (peer == detail::NcclDeviceContextAccess::rank(ctx)) + return NcclDeviceRoute::kLocal; + if (mc_nccl_lsa_available(ctx, peer)) return NcclDeviceRoute::kLsa; + return mc_nccl_gin_available(ctx) ? NcclDeviceRoute::kGin + : NcclDeviceRoute::kUnavailable; +} + +// Resolve the peer address corresponding to local_ptr. local_ptr must lie in +// the registration used to create ctx, and peer must be local or LSA reachable. +__device__ __forceinline__ void* mc_nccl_peer_ptr(const NcclDeviceContext& ctx, + int peer, + const void* local_ptr) { + if (peer == detail::NcclDeviceContextAccess::rank(ctx)) + return const_cast(local_ptr); + + const size_t offset = detail::mc_nccl_pointer_offset(ctx, local_ptr); + return ncclGetPeerPointer(detail::NcclDeviceContextAccess::window(ctx), + offset, peer); +} + +__device__ __forceinline__ bool mc_nccl_multimem_available( + const NcclDeviceContext& ctx) { + return detail::NcclDeviceContextAccess::lsaMultimemEnabled(ctx); +} + +// Return the LSA multicast pointer corresponding to local_ptr. The context must +// have been initialized with require_lsa_multimem=true. +__device__ __forceinline__ void* mc_nccl_multimem_ptr( + const NcclDeviceContext& ctx, const void* local_ptr) { + const size_t offset = detail::mc_nccl_pointer_offset(ctx, local_ptr); + return ncclGetLsaMultimemPointer( + detail::NcclDeviceContextAccess::window(ctx), offset, + detail::NcclDeviceContextAccess::comm(ctx)); +} + +// Synchronize the local LSA team. Every thread in the CTA must call this +// convergently, and barrier_index must be below config.lsa_barrier_count. +// World/cross-LSA synchronization remains the caller's responsibility. +__device__ __forceinline__ void mc_nccl_lsa_barrier( + const NcclDeviceContext& ctx, unsigned int barrier_index) { + const auto& comm = detail::NcclDeviceContextAccess::comm(ctx); + ncclLsaBarrierSession barrier{ + ncclCoopCta{}, comm, ncclTeamTagLsa{}, barrier_index, + detail::NcclDeviceContextAccess::lsaMultimemEnabled(ctx)}; + barrier.sync(ncclCoopCta{}, cuda::memory_order_acq_rel); +} + +// The GIN helpers below mirror Mooncake's lane-selected IBGDA call shape. +// Exactly lane 0 issues an NCCL ncclCoopThread operation; other lanes return. +// If a warp staged send_ptr, the caller must synchronize the warp before this +// call. mc_nccl_flush() makes sources issued on the context safe to reuse, but +// remote settlement requires mc_nccl_put_with_signal() and a receiver wait. +// GIN calls require GIN to be enabled, a valid nonlocal destination, and a +// nonnegative channel. Signal pointers must be naturally aligned uint64_t +// objects inside the registration. + +__device__ __forceinline__ void mc_nccl_put( + const NcclDeviceContext& ctx, int channel, int dst_rank, int qps_per_rank, + const void* send_ptr, void* recv_ptr, uint32_t nbytes, int lane_id) { + (void)qps_per_rank; + if (lane_id != 0) return; + + const size_t src_offset = detail::mc_nccl_pointer_offset(ctx, send_ptr); + const size_t dst_offset = detail::mc_nccl_pointer_offset(ctx, recv_ptr); + const auto& comm = detail::NcclDeviceContextAccess::comm(ctx); + const auto window = detail::NcclDeviceContextAccess::window(ctx); + ncclGin gin{comm, detail::mc_nccl_gin_context( + ctx, static_cast(channel))}; + gin.put(ncclTeamWorld(comm), dst_rank, window, dst_offset, window, + src_offset, nbytes, ncclGin_None{}, ncclGin_None{}, + ncclCoopThread{}); +} + +// Issue a GIN put and add completion_delta to a 64-bit destination signal only +// after this payload and preceding puts to the same peer/context have settled. +// recv_ptr and completion_ptr identify offsets in the peer's matching buffer. +__device__ __forceinline__ void mc_nccl_put_with_signal( + const NcclDeviceContext& ctx, int channel, int dst_rank, int qps_per_rank, + const void* send_ptr, void* recv_ptr, uint32_t nbytes, + uint64_t* completion_ptr, uint64_t completion_delta, int lane_id) { + (void)qps_per_rank; + if (lane_id != 0) return; + + const size_t src_offset = detail::mc_nccl_pointer_offset(ctx, send_ptr); + const size_t dst_offset = detail::mc_nccl_pointer_offset(ctx, recv_ptr); + const size_t completion_offset = + detail::mc_nccl_pointer_offset(ctx, completion_ptr); + const auto& comm = detail::NcclDeviceContextAccess::comm(ctx); + const auto window = detail::NcclDeviceContextAccess::window(ctx); + ncclGin gin{comm, detail::mc_nccl_gin_context( + ctx, static_cast(channel))}; + gin.put(ncclTeamWorld(comm), dst_rank, window, dst_offset, window, + src_offset, nbytes, + ncclGin_VASignalAdd{window, completion_offset, completion_delta}, + ncclGin_None{}, ncclCoopThread{}); +} + +// Add value to a 64-bit signal in the destination's matching buffer. This is +// explicitly increment semantics, not assignment. Local/LSA routes use a +// system-scope atomic; GIN uses a VA signal on the selected context. +__device__ __forceinline__ void mc_nccl_signal_add( + const NcclDeviceContext& ctx, int dst_rank, int channel, int qps_per_rank, + uint64_t* signal_ptr, uint64_t value, int lane_id) { + (void)qps_per_rank; + if (lane_id != 0) return; + + if (dst_rank == detail::NcclDeviceContextAccess::rank(ctx)) { + cuda::atomic_ref signal( + *signal_ptr); + signal.fetch_add(value, cuda::memory_order_release); + return; + } + const auto& comm = detail::NcclDeviceContextAccess::comm(ctx); + if (ncclTeamRankIsMember(ncclTeamLsa(comm), ncclTeamWorld(comm), + dst_rank)) { + const size_t signal_offset = + detail::mc_nccl_pointer_offset(ctx, signal_ptr); + auto* target = static_cast( + ncclGetPeerPointer(detail::NcclDeviceContextAccess::window(ctx), + signal_offset, dst_rank)); + cuda::atomic_ref signal(*target); + signal.fetch_add(value, cuda::memory_order_release); + return; + } + + const size_t signal_offset = + detail::mc_nccl_pointer_offset(ctx, signal_ptr); + const auto window = detail::NcclDeviceContextAccess::window(ctx); + ncclGin gin{comm, detail::mc_nccl_gin_context( + ctx, static_cast(channel))}; + gin.signal(ncclTeamWorld(comm), dst_rank, + ncclGin_VASignalAdd{window, signal_offset, value}, + ncclCoopThread{}); +} + +__device__ __forceinline__ void mc_nccl_flush(const NcclDeviceContext& ctx, + int channel, int lane_id) { + if (lane_id != 0) return; + const auto& comm = detail::NcclDeviceContextAccess::comm(ctx); + ncclGin gin{comm, detail::mc_nccl_gin_context( + ctx, static_cast(channel))}; + gin.flush(ncclCoopThread{}); +} + +__device__ __forceinline__ uint64_t mc_nccl_read_signal( + const NcclDeviceContext& ctx, int channel, const uint64_t* signal_ptr) { + if (!mc_nccl_gin_available(ctx)) { + auto& value = *const_cast(signal_ptr); + cuda::atomic_ref signal(value); + return signal.load(cuda::memory_order_acquire); + } + + const size_t signal_offset = + detail::mc_nccl_pointer_offset(ctx, signal_ptr); + const auto& comm = detail::NcclDeviceContextAccess::comm(ctx); + ncclGin gin{comm, detail::mc_nccl_gin_context( + ctx, static_cast(channel))}; + return gin.readSignal(detail::NcclDeviceContextAccess::window(ctx), + signal_offset); +} + +// Only lane 0 waits. Callers that need a warp-wide dependency must synchronize +// after this function returns. +__device__ __forceinline__ void mc_nccl_wait_signal( + const NcclDeviceContext& ctx, int channel, const uint64_t* signal_ptr, + uint64_t least, int lane_id) { + if (lane_id != 0) return; + if (!mc_nccl_gin_available(ctx)) { + auto& value = *const_cast(signal_ptr); + cuda::atomic_ref signal(value); + while (signal.load(cuda::memory_order_acquire) < least) { + } + return; + } + + const size_t signal_offset = + detail::mc_nccl_pointer_offset(ctx, signal_ptr); + const auto& comm = detail::NcclDeviceContextAccess::comm(ctx); + ncclGin gin{comm, detail::mc_nccl_gin_context( + ctx, static_cast(channel))}; + gin.waitSignal(ncclCoopThread{}, + detail::NcclDeviceContextAccess::window(ctx), signal_offset, + least); +} + +} // namespace device +} // namespace mooncake diff --git a/mooncake-transfer-engine/include/transport/device/nccl_device_transport.h b/mooncake-transfer-engine/include/transport/device/nccl_device_transport.h new file mode 100644 index 0000000000..a32fd5a543 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/device/nccl_device_transport.h @@ -0,0 +1,166 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#pragma once + +#include +#include +#include +#include + +namespace mooncake { +namespace device { + +enum class NcclGinBackend : uint8_t { + kNone = 0, + kProxy, + kGdaki, + kGpi, +}; + +enum class NcclDeviceRoute : uint8_t { + kUnavailable = 0, + kLocal, + kLsa, + kGin, +}; + +struct NcclTransportConfig { + int rank = -1; + int num_ranks = 0; + + // Mooncake currently supports either full world-team GIN connectivity or + // no GIN. Rail connectivity is deferred until Mooncake has a rail-team + // rank contract. + bool enable_gin = true; + int gin_context_count = 4; + bool gin_exclusive_contexts = false; + + // LSA barriers synchronize only the local LSA team. Cross-LSA/world + // synchronization remains the caller's responsibility. + int lsa_barrier_count = 0; + bool require_lsa_multimem = false; +}; + +struct NcclTransportProperties { + int runtime_version = 0; + int rank = -1; + int num_ranks = 0; + int cuda_device = -1; + bool device_api_supported = false; + bool multimem_supported = false; + bool lsa_multimem_enabled = false; + int lsa_team_count = 0; + int lsa_barrier_count = 0; + bool gin_enabled = false; + NcclGinBackend gin_backend = NcclGinBackend::kNone; + int gin_connection_count = 0; + int gin_context_count = 0; +}; + +namespace detail { +struct NcclDeviceContextAccess; +} // namespace detail + +class NcclDeviceTransportImpl; + +// Opaque token for one collectively registered symmetric buffer. The token +// does not own the allocation; deregister it before freeing the buffer. +class NcclBufferRegistration { + public: + bool valid() const { return id_ != 0; } + + private: + uint64_t id_ = 0; + + friend class NcclDeviceTransportImpl; +}; + +// Pass this small Mooncake context by value to kernels. Native NCCL +// communicators and windows remain behind opaque pointers. +class NcclDeviceContext { + public: + bool valid() const { return native_comm_ != nullptr; } + + private: + const void* native_comm_ = nullptr; + const void* native_window_ = nullptr; + const void* local_base_ = nullptr; + int rank_ = -1; + int gin_context_count_ = 0; + bool gin_enabled_ = false; + bool lsa_multimem_enabled_ = false; + + friend class NcclDeviceTransportImpl; + friend struct detail::NcclDeviceContextAccess; +}; + +// Host calls on one transport instance are not thread-safe and must be +// externally serialized. initialize() binds the transport to the current CUDA +// device; keep that device current for every later host call, including +// shutdown and destruction. +// +class NcclTransport { + public: + virtual ~NcclTransport() = default; + + // Generate the NCCL bootstrap ID on one rank. Exchange this int32_t blob + // through the caller's existing control plane before initialize(). + virtual std::vector createUniqueId() = 0; + + // Create the host and device communicators. Every rank must call this with + // the same unique ID and compatible config. The NCCL headers used to build + // Mooncake and device kernels must exactly match the runtime libnccl; + // initialization rejects a mismatch. Rebuild AOT kernels and regenerate + // cached JIT kernels after every NCCL upgrade. + virtual int initialize(const NcclTransportConfig& config, + const std::vector& unique_id) = 0; + + // NCCL-compatible VMM allocation. These low-level methods are local; every + // rank must verify allocation success before entering registerBuffer(). + virtual void* allocateBuffer(size_t bytes) = 0; + virtual int freeBuffer(void* ptr) = 0; + + // Collectively register a symmetric buffer. Calls must occur in the same + // order on every rank. Deregistration is local after all device work and + // remote access have completed. The registration is invalidated on + // successful deregistration. + virtual int registerBuffer(void* ptr, size_t bytes, + NcclBufferRegistration* registration) = 0; + virtual int deregisterBuffer(NcclBufferRegistration* registration) = 0; + + // Safe common path: every rank allocates, collectively checks that all + // allocations succeeded, and only then enters registration. All ranks + // must call this method in the same order. + virtual int allocateAndRegisterBuffer( + size_t bytes, void** ptr, NcclBufferRegistration* registration) = 0; + + // Snapshot passed by value to a CUDA kernel and bound to one registered + // buffer. It remains valid until registration removal or shutdown. + virtual NcclDeviceContext deviceContext( + const NcclBufferRegistration& registration) const = 0; + + virtual NcclTransportProperties properties() const = 0; + virtual bool initialized() const = 0; + + // The caller must first ensure that no kernel can access the context or + // any registered buffer. + virtual int shutdown() = 0; +}; + +// Create the CUDA-only NCCL LSA/GIN device transport. +std::unique_ptr createNcclDeviceTransport(); + +} // namespace device +} // namespace mooncake diff --git a/mooncake-transfer-engine/include/transport/device/p2p_device.cuh b/mooncake-transfer-engine/include/transport/device/p2p_device.cuh index d47ca0a2cc..824a33c2ac 100644 --- a/mooncake-transfer-engine/include/transport/device/p2p_device.cuh +++ b/mooncake-transfer-engine/include/transport/device/p2p_device.cuh @@ -19,7 +19,11 @@ struct P2PContext { __device__ __forceinline__ bool mc_p2p_available(const P2PContext& ctx, int dst_rank) { +#ifdef MOONCAKE_EP_USE_MUSA + return ctx.peer_ptrs[dst_rank] != nullptr; +#else return ctx.available[dst_rank] != 0 && ctx.peer_ptrs[dst_rank] != nullptr; +#endif } // Translate a local pointer (within the GDR buffer) to the peer's mapped VA. diff --git a/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h b/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h index 067c46d187..1a606c86e5 100644 --- a/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h +++ b/mooncake-transfer-engine/include/transport/efa_transport/efa_context.h @@ -216,7 +216,13 @@ class EfaContext { // model; plain `volatile int` + __sync_* was UB under the current // standard. std::atomic wr_depth_; + // Ceilings for the two pacing counters (wr_depth_ above, EfaCq::outstanding + // on shared_cq_). Both are the depths the provider gave us for this + // device, not GlobalConfig values: a counter that disagrees with the queue + // it paces either stalls submission early or hands out credit fi_write has + // to refuse. Set in buildSharedEndpoint() and construct() respectively. int max_wr_depth_; + size_t max_cqe_ = 0; // CQ that shared_ep_ is bound to (FI_TRANSMIT|FI_RECV). Points into // cq_list_[0]; kept here to avoid re-indexing on the hot path. std::shared_ptr shared_cq_; diff --git a/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h b/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h index a171198618..c081f4b93d 100644 --- a/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h +++ b/mooncake-transfer-engine/include/transport/efa_transport/efa_transport.h @@ -70,6 +70,17 @@ class EfaTransport : public Transport { int unregisterLocalMemoryBatch( const std::vector& addr_list) override; + // For each memory location the topology knows about (e.g. "cuda:3"), the + // indices into `device_names` of the NICs it reports as closest. Locations + // whose preferred set is empty, or whose preferred NICs are all absent from + // `device_names`, are omitted rather than mapped to an empty vector. + // + // Exposed for testing; used to build the lookup MC_EFA_NIC_SELECTION=local + // consults. See registerLocalMemoryInternal() for the trade-off. + static std::unordered_map> + buildLocalNicMap(const TopologyMatrix& matrix, + const std::vector& device_names); + // Eagerly populate the address vector with every (local_ctx, peer_nic) // handshake for `segment_name`. // @@ -150,6 +161,12 @@ class EfaTransport : public Transport { std::vector> context_list_; std::shared_ptr local_topology_; + // Memory location -> context_list_ indices of the topology-local NICs, + // built once in initializeEfaResources(). Only read when + // MC_EFA_NIC_SELECTION=local; empty otherwise. Const after install(), so no + // lock is needed on the registration path. + std::unordered_map> local_nic_map_; + // Track chunked MR registrations for per-NIC partitioned buffers. // When a buffer exceeds max_mr_size, it is split into chunks, each // registered on a disjoint subset of NICs (per-NIC partition). diff --git a/mooncake-transfer-engine/include/transport/nccl_transport/nccl_transport.h b/mooncake-transfer-engine/include/transport/nccl_transport/nccl_transport.h new file mode 100644 index 0000000000..9ff27818e2 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/nccl_transport/nccl_transport.h @@ -0,0 +1,83 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef NCCL_TRANSPORT_H_ +#define NCCL_TRANSPORT_H_ + +#include + +#include "transport/transport.h" + +namespace mooncake { + +// Host-submitted NCCL RMA backend for the classic Transfer Engine API. +// Native NCCL communicators and windows are intentionally private. +// +// Install the NCCL transport before registering its buffers so registration +// order is retained. +// +// Each peer pair must register the same number of CUDA VMM buffers in the same +// order, with matching lengths, before its first transfer. Corresponding +// buffers become one NCCL symmetric window; ncclMemAlloc is the supported +// allocation path. Both endpoints must remain available while the first +// transfer initializes the peer session and its collective symmetric windows. +// The first valid WRITE freezes the registered-buffer catalog before bootstrap, +// so registration cannot change afterward even if bootstrap fails. Session +// initialization is attempted once per endpoint/device pair. A terminal +// session failure remains cached, and later transfers return that error +// without retrying bootstrap. Recovery requires recreating the NCCL transport +// (normally its TransferEngine instance) on both peers and registering the +// buffers again; one-sided restart is unsupported. +// +// Only WRITE is supported; it uses ncclPutSignal. NCCL 2.30 has no public host +// Get operation, so READ is rejected with Status::NotSupportedTransport +// without submitting an NCCL operation. +class NcclHostTransport final : public Transport { + public: + NcclHostTransport(); + ~NcclHostTransport() override; + + Status submitTransfer(BatchID batch_id, + const std::vector& entries) override; + + Status submitTransferTask( + const std::vector& task_list) override; + + Status getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus& status) override; + + protected: + int install(std::string& local_server_name, + std::shared_ptr metadata, + std::shared_ptr topology) override; + + private: + int registerLocalMemory(void* addr, size_t length, + const std::string& location, bool remote_accessible, + bool update_metadata = true) override; + int unregisterLocalMemory(void* addr, bool update_metadata = true) override; + int registerLocalMemoryBatch(const std::vector& buffer_list, + const std::string& location) override; + int unregisterLocalMemoryBatch( + const std::vector& addr_list) override; + + const char* getName() const override { return "nccl"; } + + class Impl; + std::unique_ptr impl_; +}; + +} // namespace mooncake + +#endif // NCCL_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/include/transport/nvmeof_transport/cufile_desc_pool.h b/mooncake-transfer-engine/include/transport/nvmeof_transport/cufile_desc_pool.h index 498941f701..0f44801758 100644 --- a/mooncake-transfer-engine/include/transport/nvmeof_transport/cufile_desc_pool.h +++ b/mooncake-transfer-engine/include/transport/nvmeof_transport/cufile_desc_pool.h @@ -27,6 +27,8 @@ namespace mooncake { +class CUFileDescPoolTestPeer; + // Wrapper for reusable CUfileBatchHandle_t // cuFileBatchIOSetUp is expensive, so we reuse handles (similar to GDS // transport) @@ -40,10 +42,17 @@ struct BatchHandle { struct CUFileBatchDesc { BatchHandle* batch_handle; // Pointer to reusable handle from pool std::vector io_params; + // Completion events returned by cuFile are correlated by cookie and cached + // by submission index. cuFileBatchIOGetStatus only returns completed I/Os, + // so its output cannot be treated as a positional status snapshot. std::vector io_events; + std::vector polled_events; + bool reusable = true; }; class CUFileDescPool { + friend class CUFileDescPoolTestPeer; + public: explicit CUFileDescPool(size_t max_batch_size = 128); ~CUFileDescPool(); @@ -65,6 +74,18 @@ class CUFileDescPool { // Get transfer status for a specific slice CUfileIOEvents_t getTransferStatus(int idx, int slice_id); + // Poll cuFile once for the batch and update cached completion events. + bool updateBatchStatus(int idx); + + // Get cached transfer status for a specific slice. + CUfileIOEvents_t getCachedTransferStatus(int idx, int slice_id); + + // Best-effort cancellation for a submitted batch. + bool cancelBatch(int idx); + + // Prevent an unsafe batch handle from being returned to the reusable pool. + void markUnreusable(int idx); + // Get current number of slices in the descriptor int getSliceNum(int idx); @@ -75,13 +96,24 @@ class CUFileDescPool { CUFileBatchDesc* getDesc(int idx); private: + static bool cachePolledEvent(std::vector& io_events, + const CUfileIOEvents_t& event); + static bool isTerminalStatus(CUfileStatus_t status); + static CUfileIOEvents_t failedEvent(); + static void destroyDesc(CUFileBatchDesc* desc); + static bool updateBatchStatus(CUFileBatchDesc* desc, int idx); + static const size_t MAX_NR_DESC = 256; // Max number of descriptors + void cleanupQuarantinedDescs(); + size_t max_batch_size_; // Object pool for BatchHandle to avoid frequent cuFileBatchIOSetUp/Destroy std::vector handle_pool_; std::mutex handle_pool_lock_; + std::vector quarantined_descs_; + // Array of descriptors (nullptr = free slot) CUFileBatchDesc* descs_[MAX_NR_DESC]; RWSpinlock mutex_; @@ -89,4 +121,4 @@ class CUFileDescPool { } // namespace mooncake -#endif \ No newline at end of file +#endif diff --git a/mooncake-transfer-engine/include/transport/nvmeof_transport/nvmeof_transport.h b/mooncake-transfer-engine/include/transport/nvmeof_transport/nvmeof_transport.h index 120e25db61..123e1dcce2 100644 --- a/mooncake-transfer-engine/include/transport/nvmeof_transport/nvmeof_transport.h +++ b/mooncake-transfer-engine/include/transport/nvmeof_transport/nvmeof_transport.h @@ -30,6 +30,8 @@ namespace mooncake { +class NVMeoFTransportTestPeer; + struct NVMeoFBatchDesc { int desc_idx_; std::vector transfer_status; @@ -37,6 +39,8 @@ struct NVMeoFBatchDesc { }; class NVMeoFTransport : public Transport { + friend class NVMeoFTransportTestPeer; + public: NVMeoFTransport(); @@ -60,6 +64,16 @@ class NVMeoFTransport : public Transport { TransferTask &task, const char *file_path); private: + explicit NVMeoFTransport(std::shared_ptr desc_pool); + + static TransferStatus aggregateTransferStatus( + const std::vector &slice_statuses, bool &is_finished); + + void collectSliceStatuses(int desc_idx, size_t slice_id, size_t slice_num, + std::vector &slice_statuses); + + static bool isTerminalFailure(TransferStatusEnum status); + void startTransfer(Slice *slice); private: diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/connect_pause_tracker.h b/mooncake-transfer-engine/include/transport/rdma_transport/connect_pause_tracker.h new file mode 100644 index 0000000000..0574c3df14 --- /dev/null +++ b/mooncake-transfer-engine/include/transport/rdma_transport/connect_pause_tracker.h @@ -0,0 +1,86 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef CONNECT_PAUSE_TRACKER_H +#define CONNECT_PAUSE_TRACKER_H + +#include +#include +#include +#include +#include + +namespace mooncake { + +// Active-connect circuit-breaker state: per peer server name (peer IP), a +// "pause active reconnection until" timestamp. After an endpoint to a peer is +// torn down (path failure / QP fatal), the peer is paused so the CQ poller is +// not blocked re-handshaking a likely-gone peer. The clock is injected so the +// TTL/expiry/prune logic is unit-testable without RDMA hardware or real sleeps. +// All methods are thread-safe. +class ConnectPauseTracker { + public: + using Clock = std::function; // nanosecond timestamp source + + explicit ConnectPauseTracker(Clock clock) : clock_(std::move(clock)) {} + + // Arm/refresh the pause for `server` for `duration_ns`. The deadline is + // derived from the injected clock so production and tests use one time + // source. + void pauseFor(const std::string &server, uint64_t duration_ns) { + std::lock_guard lock(mu_); + until_ns_[server] = clock_() + duration_ns; + } + + // Whether `server` is currently paused. Expired entries are deleted on + // access (lazy expiry). + bool isPaused(const std::string &server) { + std::lock_guard lock(mu_); + auto it = until_ns_.find(server); + if (it == until_ns_.end()) return false; + if (clock_() >= it->second) { + until_ns_.erase(it); // expired -> delete + return false; + } + return true; + } + + // Drop all expired entries (so the map doesn't grow for peers that are + // never re-checked after their pause lapses). Intended for a periodic tick. + void prune() { + uint64_t now = clock_(); + std::lock_guard lock(mu_); + for (auto it = until_ns_.begin(); it != until_ns_.end();) { + if (now >= it->second) + it = until_ns_.erase(it); + else + ++it; + } + } + + // For tests / diagnostics. + size_t size() { + std::lock_guard lock(mu_); + return until_ns_.size(); + } + + private: + const Clock clock_; + std::mutex mu_; + std::unordered_map until_ns_; +}; + +} // namespace mooncake + +#endif // CONNECT_PAUSE_TRACKER_H diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h b/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h index 3d99bb2a88..f3bd20dd01 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/endpoint_store.h @@ -43,7 +43,15 @@ class EndpointStore { virtual std::shared_ptr insertEndpoint( const std::string &peer_nic_path, RdmaContext *context) = 0; virtual int deleteEndpoint(const std::string &peer_nic_path) = 0; - virtual int deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) = 0; + // Deletes the endpoint matching endpoint_ptr (by pointer identity, under + // the store lock -- the pointer is never dereferenced, so a stale/freed + // pointer is safe and simply does not match). If found and + // deleted_peer_nic_path is non-null, it is set to that endpoint's peer NIC + // path (read from the live map key) so callers can act on the path without + // touching the raw pointer. + virtual int deleteEndpointByPtr( + const RdmaEndPoint *endpoint_ptr, + std::string *deleted_peer_nic_path = nullptr) = 0; virtual void evictEndpoint() = 0; // Takes endpoint_map_lock_; caller must not hold it (RWSpinlock is // non-reentrant, so recursive acquisition deadlocks). @@ -77,7 +85,9 @@ class FIFOEndpointStore : public EndpointStore { std::shared_ptr insertEndpoint( const std::string &peer_nic_path, RdmaContext *context) override; int deleteEndpoint(const std::string &peer_nic_path) override; - int deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) override; + int deleteEndpointByPtr( + const RdmaEndPoint *endpoint_ptr, + std::string *deleted_peer_nic_path = nullptr) override; void evictEndpoint() override; void reclaimEndpoint() override; size_t getSize() override; @@ -117,7 +127,9 @@ class SIEVEEndpointStore : public EndpointStore { std::shared_ptr insertEndpoint( const std::string &peer_nic_path, RdmaContext *context) override; int deleteEndpoint(const std::string &peer_nic_path) override; - int deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) override; + int deleteEndpointByPtr( + const RdmaEndPoint *endpoint_ptr, + std::string *deleted_peer_nic_path = nullptr) override; void evictEndpoint() override; void reclaimEndpoint() override; size_t getSize() override; diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h index 162b7fc99a..370d6c2660 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_context.h @@ -37,6 +37,7 @@ #include "common.h" #include "rdma_gid_probe.h" #include "rdma_transport.h" +#include "transport/rdma_transport/connect_pause_tracker.h" #include "transport/transport.h" namespace mooncake { @@ -59,10 +60,25 @@ struct GidSelectionSnapshot { int gid_index = -1; }; +enum class GidRefreshResult { + UNCHANGED = 0, + CHANGED = 1, + FAILED = 2, +}; + struct RdmaCq { RdmaCq() : native(nullptr), outstanding(0) {} + RdmaCq(const RdmaCq &) = delete; + RdmaCq &operator=(const RdmaCq &) = delete; + RdmaCq(RdmaCq &&other) noexcept + : native(other.native), + outstanding(other.outstanding.load(std::memory_order_relaxed)) { + other.native = nullptr; + } + RdmaCq &operator=(RdmaCq &&) = delete; + ibv_cq *native; - volatile int outstanding; + std::atomic outstanding; }; struct MemoryRegionMeta { @@ -72,11 +88,25 @@ struct MemoryRegionMeta { struct ibv_mr *mr; }; +// A dma_buf handle exported once for a buffer and shared across every NIC's +// registration of that buffer. Exporting a single fd (instead of one per NIC) +// collapses the per-NIC dma_buf objects into one kernel object, so the GPU +// driver reserves a single BAR1 window for the buffer rather than one window +// per NIC. Host memory (and the nvidia-peermem path) yields kHostReg with no +// fd, taking the plain ibv_reg_mr path. +struct DmabufExport { + enum class Method { kHostReg, kDmabufReg }; + Method method = Method::kHostReg; + int fd = -1; // live dma_buf fd; -1 when not applicable + uint64_t offset = 0; // offset of addr within the exported allocation +}; + // RdmaContext represents the set of resources controlled by each local NIC, // including Memory Region, CQ, EndPoint (QPs), etc. class RdmaContext { public: friend class RdmaContextTestPeer; + friend class WorkerPool; RdmaContext(RdmaTransport &engine, const std::string &device_name); @@ -93,6 +123,26 @@ class RdmaContext { // Memory Region Management int registerMemoryRegion(void *addr, size_t length, int access); + // Shared-fd variant: the caller exports a single dma_buf fd for the buffer + // via exportDmabuf(), passes the same handle to every NIC's registration, + // then closes the fd once via closeDmabufExport() AFTER all registrations + // have completed. This keeps one dma_buf object alive across all NICs so + // the GPU driver reserves a single BAR1 window for the buffer. + int registerMemoryRegion(void *addr, size_t length, int access, + const DmabufExport &exp); + + // Exports a single dma_buf fd for the allocation backing addr. GPU device + // memory yields kDmabufReg with a live fd; host memory and the + // nvidia-peermem path yield kHostReg with no fd. Any fd placed in out.fd + // MUST be closed by the caller (via closeDmabufExport) AFTER every + // registerMemoryRegion() call consuming it has returned — each successful + // registration takes its own reference, so closing earlier would invalidate + // the fd for the remaining NICs. + static int exportDmabuf(void *addr, DmabufExport &out); + + // Closes the fd held by a DmabufExport, if any. Idempotent. + static void closeDmabufExport(DmabufExport &exp); + int unregisterMemoryRegion(void *addr); int preTouchMemory(void *addr, size_t length); @@ -103,6 +153,7 @@ class RdmaContext { private: int registerMemoryRegionInternal(void *addr, size_t length, int access, + const DmabufExport &exp, MemoryRegionMeta &mrMeta); using MemoryRegionMap = std::map; @@ -113,9 +164,11 @@ class RdmaContext { uintptr_t addr) const; public: - bool active() const { return active_; } + bool active() const { return active_.load(std::memory_order_acquire); } - void set_active(bool flag) { active_ = flag; } + void set_active(bool flag) { + active_.store(flag, std::memory_order_release); + } public: // EndPoint Management @@ -127,6 +180,16 @@ class RdmaContext { int deleteEndpoint(const std::string &peer_nic_path); int deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr); + // Active-connect circuit-breaker. After deleteEndpointByPtr tears an + // endpoint down, active reconnection to that peer's address is paused for + // globalConfig().conn_pause_ttl_ms so the CQ poller isn't blocked + // re-handshaking a likely-gone peer. Entries expire (lazily on + // isConnectPaused, and via pruneConnectPause from the monitor tick). All + // no-ops when the TTL is 0. Keyed by peer server name (the peer IP). + void pauseConnect(const std::string &peer_nic_path); + bool isConnectPaused(const std::string &peer_nic_path); + void pruneConnectPause(); + // Drain the endpoint store's waiting list. Safe to call on any thread; // intended to be invoked periodically from monitorWorker so reclaim is // not gated on new endpoint insertions (which can stall under failure @@ -170,6 +233,12 @@ class RdmaContext { const std::vector &tried_selections = {}, std::string *previous_gid = nullptr, std::string *next_gid = nullptr); + // Refresh the runtime GID after IBV_EVENT_GID_CHANGE. Auto-GID mode uses + // the same candidate filtering/ranking as initial device open; explicit + // MC_GID_INDEX keeps the configured index and refreshes only its value. + GidRefreshResult refreshCurrentGid(std::string *previous_gid = nullptr, + std::string *next_gid = nullptr); + ibv_context *context() const { return context_; } RdmaTransport &engine() const { return engine_; } @@ -181,6 +250,7 @@ class RdmaContext { uint8_t numLagPorts() const { return num_lag_ports_; } int activeSpeed() const { return active_speed_; } + int activeWidth() const { return active_width_; } ibv_mtu activeMTU() const { return active_mtu_; } @@ -192,7 +262,7 @@ class RdmaContext { ibv_cq *cq(); - volatile int *cqOutstandingCount(int cq_index) { + std::atomic *cqOutstandingCount(int cq_index) { return &cq_list_[cq_index].outstanding; } @@ -216,6 +286,11 @@ class RdmaContext { public: int submitPostSend(const std::vector &slice_list); + void trackPostedSlices(const std::vector &slice_list, + size_t first, size_t count); + void untrackPostedSlices(const std::vector &slice_list, + size_t first, size_t count); + private: const std::string device_name_; RdmaTransport &engine_; @@ -232,6 +307,7 @@ class RdmaContext { uint16_t lid_ = 0; int gid_index_ = -1; int active_speed_ = -1; + int active_width_ = 1; ibv_mtu active_mtu_; uint8_t num_lag_ports_ = 0; // 0/1 = not in LAG; ≥2 = LAG active ibv_gid gid_; @@ -245,6 +321,9 @@ class RdmaContext { std::shared_ptr endpoint_store_; + // Active-connect circuit-breaker (keyed by peer server name). + ConnectPauseTracker connect_pause_; + std::vector background_thread_; std::atomic threads_running_; @@ -254,7 +333,7 @@ class RdmaContext { std::shared_ptr worker_pool_; - volatile bool active_; + std::atomic active_; }; } // namespace mooncake diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h index d1ae09df06..255ae8efda 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_endpoint.h @@ -15,12 +15,15 @@ #ifndef RDMA_ENDPOINT_H #define RDMA_ENDPOINT_H +#include #include #include "rdma_context.h" namespace mooncake { +class RdmaEndPointTestPeer; + // RdmaEndPoint represents all QP connections between the local NIC1 (identified // by its RdmaContext) and the remote NIC2 (identified by peer_nic_path). // 1. After construct, resources are allocated without specifying the peers. @@ -31,7 +34,9 @@ namespace mooncake { // which can be obtained from RdmaContext::nicPath() on the remote side // - Remote side calls the setupConnectionsByPassive() function in its RPC // service. -// After above steps, the RdmaEndPoint state is set to CONNECTED +// After above steps, the RdmaEndPoint state is set to CONNECTED. With RDMA +// ready ACK enabled, QPs first enter CONNECTED_WAIT_READY_ACK after reaching +// RTS and become CONNECTED only after the ready-ACK phase completes. // // If the user initiates a disconnect() call or an error is detected internally, // the connection is closed and the RdmaEndPoint state is set to UNCONNECTED. @@ -42,11 +47,14 @@ class RdmaEndPoint { INITIALIZING, UNCONNECTED, CONNECTING, + CONNECTED_WAIT_READY_ACK, CONNECTED, DESTROYING, DESTROYED, }; + friend class RdmaEndPointTestPeer; + public: RdmaEndPoint(RdmaContext &context); @@ -59,6 +67,7 @@ class RdmaEndPoint { int reconstruct(); int deconstruct(); int deconstructLocked(); + void beginDestroyLocked(); public: void setPeerNicPath(const std::string &peer_nic_path); @@ -74,22 +83,35 @@ class RdmaEndPoint { int setupConnectionsByPassive(const HandShakeDesc &peer_desc, HandShakeDesc &local_desc); - bool active() const { return active_; } + bool active() const { return active_.load(std::memory_order_acquire); } void set_active(bool flag) { RWSpinlock::WriteGuard guard(lock_); - active_ = flag; - if (!flag) inactive_time_ = getCurrentTimeInNano(); + if (!flag) + inactive_time_.store(getCurrentTimeInNano(), + std::memory_order_relaxed); + active_.store(flag, std::memory_order_release); } double inactiveTime() { - if (active_) return 0.0; - return (getCurrentTimeInNano() - inactive_time_) / 1000000000.0; + if (active_.load(std::memory_order_acquire)) return 0.0; + return (getCurrentTimeInNano() - + inactive_time_.load(std::memory_order_relaxed)) / + 1000000000.0; } public: - bool connected() const { - return status_.load(std::memory_order_relaxed) == CONNECTED; + bool connected() const { return isConnectedStatus(status()); } + + // CONNECTED_WAIT_READY_ACK means local QPs have reached RTS but the RDMA + // ready-ACK phase has not completed yet. Only CONNECTED can post WRs. + bool readyToSend() const { return status() == CONNECTED; } + + bool readyAckTimedOut() const; + + bool retired() const { + auto status = status_.load(std::memory_order_relaxed); + return status == DESTROYING || status == DESTROYED; } // Interrupts the connection, which can be triggered by user or by internal @@ -113,26 +135,11 @@ class RdmaEndPoint { private: int disconnectUnlocked(); - // Resets the connection. - // - // The main difference between this function and `disconnectUnlocked` - // is that it will reconstruct QPs when `CONFIG_ERDMA` is defined. - // Without `CONFIG_ERDMA`, it is essentially the same as - // `disconnectUnlocked` but with additional logging. - // - // This serves as a workaround for Aliyun eRDMA devices (i.e., once a QP is - // transitioned to the RTS state, it cannot be reset to RTS again directly). - // For more details: - // https://github.com/kvcache-ai/Mooncake/pull/1733#discussion_r2992088663 - // - // In practice: - // - Call `resetConnection` if the QPs' state may have transitioned to RTS. - // - Call `disconnectUnlocked` otherwise. - // - // This is mainly used in `setupConnectionsByActive` or - // `setupConnectionsByPassive`. It is NOT invoked in the normal execution - // flow, so a `reason` argument is passed for internal logging purposes. + // Resets only pre-connected handshake attempts. Once an endpoint has ever + // reached CONNECTED, it is retired instead of being reused. int resetConnection(const std::string &reason); + int sendReadyAck(const std::string &peer_server_name, + const HandShakeDesc &local_desc); public: const std::string toString() const; @@ -162,10 +169,17 @@ class RdmaEndPoint { int sys_errno = 0; }; + Status status() const { return status_.load(std::memory_order_relaxed); } + + static bool isConnectedStatus(Status status) { + return status == CONNECTED_WAIT_READY_ACK || status == CONNECTED; + } + std::vector qpNum() const; int doSetupConnection(const std::string &peer_gid, uint16_t peer_lid, std::vector peer_qp_num_list, + Status connected_status = CONNECTED, std::string *reply_msg = nullptr, SetupConnectionFailureInfo *failure_info = nullptr); @@ -177,6 +191,8 @@ class RdmaEndPoint { private: static constexpr uint64_t kWaitExistingHandshakeTimeoutNano = 10 * 1000000000ull; // 10 seconds + static constexpr uint64_t kReadyAckTimeoutNano = + 10 * 1000000000ull; // 10 seconds static constexpr uint32_t kWaitExistingHandshakeSpinCount = 500; static constexpr uint32_t kWaitExistingHandshakeInitialSleepUs = 50; static constexpr uint32_t kWaitExistingHandshakeMaxSleepUs = 2000; @@ -199,15 +215,17 @@ class RdmaEndPoint { std::string peer_nic_path_; std::vector peer_qp_num_list_; + bool has_connected_; + std::atomic ready_wait_start_ts_; - volatile int *wr_depth_list_; + std::atomic *wr_depth_list_; int max_wr_depth_; size_t max_sge_per_wr_; size_t max_inline_bytes_; - volatile bool active_; - volatile int *cq_outstanding_; - volatile uint64_t inactive_time_; + std::atomic active_; + std::atomic *cq_outstanding_; + std::atomic inactive_time_; int finish_destroy_retries_ = 0; }; diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h index a4ee60dbd3..1927998ade 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_gid_probe.h @@ -29,9 +29,18 @@ namespace mooncake { enum class AutoGidCandidateClass { kNetworkRoutable = 0, kNoNetworkRoutable = 1, - kNetworkDegraded = 2, - kNoNetworkDegraded = 3, - kFallbackNonzero = 4, + // Private-range (RFC 1918 / CGNAT) IPv4-mapped GIDs. Split out of the + // degraded tier (#2729): datacenter RoCE fabrics routinely address NICs + // from 10/8, and such a GID is at worst same-subnet-usable — unlike a + // link-local fe80:: GID, which can never work across subnets. The + // relative order of all pre-existing tiers is unchanged; only the former + // tie between private-range IPv4 and link-local/overlay entries within + // "degraded" is split. + kNetworkPrivateV4 = 2, + kNetworkDegraded = 3, + kNoNetworkPrivateV4 = 4, + kNoNetworkDegraded = 5, + kFallbackNonzero = 6, }; enum class AutoGidRetryAction { @@ -72,8 +81,12 @@ inline const char* autoGidCandidateClassToString( return "network-routable"; case AutoGidCandidateClass::kNoNetworkRoutable: return "no-network-routable"; + case AutoGidCandidateClass::kNetworkPrivateV4: + return "network-private-v4"; case AutoGidCandidateClass::kNetworkDegraded: return "network-degraded"; + case AutoGidCandidateClass::kNoNetworkPrivateV4: + return "no-network-private-v4"; case AutoGidCandidateClass::kNoNetworkDegraded: return "no-network-degraded"; case AutoGidCandidateClass::kFallbackNonzero: @@ -95,20 +108,26 @@ inline std::optional classifyAutoGidCandidate( } const bool is_roce_v2 = candidate.gid_type == IBV_GID_TYPE_ROCE_V2; - const bool is_overlay = - is_roce_v2 && (candidate.is_overlay_network || - (candidate.is_ipv4_mapped && candidate.is_overlay_ipv4)); + // Interface-name-based overlay detection (docker*/cni*/...) is the + // reliable demotion signal; a private address range alone is not — DC + // RoCE fabrics commonly use 10/8 (#2729). + const bool is_overlay_iface = is_roce_v2 && candidate.is_overlay_network; + const bool is_private_v4 = is_roce_v2 && !is_overlay_iface && + candidate.is_ipv4_mapped && + candidate.is_overlay_ipv4; const bool is_link_local = is_roce_v2 && !candidate.is_ipv4_mapped && candidate.is_link_local_ipv6; - const bool is_degraded = is_overlay || is_link_local; + const bool is_degraded = is_overlay_iface || is_link_local; if (candidate.has_network_device) { - return is_degraded ? AutoGidCandidateClass::kNetworkDegraded - : AutoGidCandidateClass::kNetworkRoutable; + if (is_degraded) return AutoGidCandidateClass::kNetworkDegraded; + if (is_private_v4) return AutoGidCandidateClass::kNetworkPrivateV4; + return AutoGidCandidateClass::kNetworkRoutable; } - return is_degraded ? AutoGidCandidateClass::kNoNetworkDegraded - : AutoGidCandidateClass::kNoNetworkRoutable; + if (is_degraded) return AutoGidCandidateClass::kNoNetworkDegraded; + if (is_private_v4) return AutoGidCandidateClass::kNoNetworkPrivateV4; + return AutoGidCandidateClass::kNoNetworkRoutable; } inline int autoGidCandidateClassPriority( @@ -118,14 +137,18 @@ inline int autoGidCandidateClassPriority( return 0; case AutoGidCandidateClass::kNoNetworkRoutable: return 1; - case AutoGidCandidateClass::kNetworkDegraded: + case AutoGidCandidateClass::kNetworkPrivateV4: return 2; - case AutoGidCandidateClass::kNoNetworkDegraded: + case AutoGidCandidateClass::kNetworkDegraded: return 3; - case AutoGidCandidateClass::kFallbackNonzero: + case AutoGidCandidateClass::kNoNetworkPrivateV4: return 4; + case AutoGidCandidateClass::kNoNetworkDegraded: + return 5; + case AutoGidCandidateClass::kFallbackNonzero: + return 6; } - return 5; + return 7; } inline std::vector rankAutoGidCandidates( diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h index c999b89cda..d5d8b3b5e5 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h @@ -130,6 +130,14 @@ class RdmaTransport : public Transport { static int selectDevice(SegmentDesc *desc, uint64_t offset, size_t length, std::string_view hint, int &buffer_id, int &device_id, int retry_cnt = 0); + static int selectDeviceByLocalHca(SegmentDesc *desc, uint64_t offset, + size_t length, std::string_view local_hca, + int &buffer_id, int &device_id, + int retry_cnt = 0); + + const std::vector> &getContextList() const { + return context_list_; + } private: std::vector> context_list_; @@ -140,6 +148,13 @@ class RdmaTransport : public Transport { // local_server_name_ keeps the TCP-reachable address for P2P routing. std::string rdma_server_name_; std::mutex local_desc_lock_; + // Mooncake#2017: buffers larger than the device max_mr_size are split into + // multiple sub-max_mr_size MRs (one BufferDesc per chunk) so that + // ibv_reg_mr is never silently truncated. unregisterLocalMemory() only + // receives the base addr, so remember each base buffer's chunk + // start-addresses for cleanup. + std::mutex chunk_map_mutex_; + std::unordered_map> chunk_map_; }; using TransferRequest = Transport::TransferRequest; diff --git a/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h b/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h index 08ee083ba9..3ce036217b 100644 --- a/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h +++ b/mooncake-transfer-engine/include/transport/rdma_transport/worker_pool.h @@ -22,7 +22,10 @@ #include "rdma_context.h" namespace mooncake { +class WorkerPoolTestPeer; class WorkerPool { + friend class WorkerPoolTestPeer; + public: WorkerPool(RdmaContext &context, int numa_socket_id = 0); @@ -31,18 +34,39 @@ class WorkerPool { // Add slices to queue, called by Transport int submitPostSend(const std::vector &slice_list); + void trackPostedSlices(const std::vector &slice_list, + size_t first, size_t count); + void untrackPostedSlices(const std::vector &slice_list, + size_t first, size_t count); + private: + using SliceList = std::vector; + const static int kShardCount = 8; + + // Enqueue slices that were prepared by another WorkerPool. Used for + // local-NIC failure handoff: the original worker keeps the remote path + // fixed, updates the local lkey, and pushes the slice to this context's + // worker queue. + int submitPreparedPostSend( + const std::vector &slice_list); + void enqueuePreparedSlices(SliceList (&slice_list_map)[kShardCount], + uint64_t submitted_slice_count); + void performPostSend(int thread_id); void performPollCq(int thread_id); - void redispatch(std::vector &slice_list, int thread_id); + void redispatch(std::vector &slice_list, int thread_id, + bool handoff_to_local_worker = false); void transferWorker(int thread_id); + bool hasOutstandingCq(int thread_id); + void monitorWorker(); int doProcessContextEvents(); + void processContextEventForTest(ibv_event_type event_type); // Simplified rail monitor: pause problematic paths for a cooldown period struct RailState { @@ -50,34 +74,56 @@ class WorkerPool { uint64_t pause_until_ns = 0; // Timestamp (ns) when pause expires }; - void markRailFailed(const std::string &peer_nic_path); + void markRailFailed(const std::string &peer_nic_path, + bool immediate_pause = false); bool isRailAvailable(const std::string &peer_nic_path); // Retry helper: increment retry count and return whether retry is allowed static bool shouldRetrySlice(Transport::Slice *slice); - // Unified path failure handler: marks rail failed, notifies other workers, - // and optionally deletes the endpoint - void handlePathFailure(const std::string &peer_nic_path, - RdmaEndPoint *endpoint = nullptr); + void refreshPublishedLocalTopology(); + GidRefreshResult refreshPublishedLocalGid(); + bool handleContextEvent(ibv_event_type event_type, bool injected_for_test, + struct ibv_async_event *event = nullptr); + void scheduleContextRecovery(uint64_t delay_ns = kContextRecoveryDelayNs); + void maybeActivateRecoveredContext(); + bool hasAvailablePeerRailAlternative(Transport::Slice *slice, + const std::string &failed_peer_path); + + static bool isLocalWcFailure(const ibv_wc &wc); + + // Local-side failure handler: degrade current context so retries are + // handed off to another local context's worker pool. + void handleLocalFailure(const std::string &peer_nic_path, + RdmaEndPoint *endpoint = nullptr); + + bool tryHandoffToAnotherLocalWorker(Transport::Slice *slice); // Context-level health tracking for catastrophic hardware failure. // When all rails through a local RNIC are unavailable, increment the // failure counter. Reset on any success. Mark context inactive after // consecutive failures exceed threshold. bool contextHealthy() const { - return context_failure_count_ < kContextFailureThreshold; + return context_failure_count_.load(std::memory_order_relaxed) < + kContextFailureThreshold; + } + void markContextSuccess() { + context_failure_count_.store(0, std::memory_order_relaxed); } - void markContextSuccess() { context_failure_count_ = 0; } - void markContextFailure() { - context_failure_count_++; - if (context_failure_count_ >= kContextFailureThreshold) { - LOG(WARNING) << "All rails failed for context " - << context_.deviceName() << " for " - << context_failure_count_ - << " consecutive attempts, marking inactive"; - context_.set_active(false); + bool markContextFailure() { + auto failure_count = + context_failure_count_.fetch_add(1, std::memory_order_relaxed) + 1; + if (failure_count >= kContextFailureThreshold) { + if (context_.active()) { + LOG(WARNING) + << "All rails failed for context " << context_.deviceName() + << " for " << failure_count + << " consecutive attempts, marking inactive"; + context_.set_active(false); + return true; + } } + return false; } private: @@ -86,16 +132,24 @@ class WorkerPool { std::vector worker_thread_; std::atomic workers_running_; - std::atomic suspended_flag_; + + std::atomic parked_worker_count_; + + // The poll worker updates these on every poll pass. The monitor worker + // reads them when CQ entries stay outstanding, so a transfer timeout can + // be distinguished from a stalled poller. + std::atomic last_poll_ts_ns_{0}; + std::atomic last_poll_interval_ns_{0}; + std::atomic max_poll_interval_ns_{0}; + + std::mutex posted_slices_mutex_; + std::unordered_set posted_slices_; std::atomic redispatch_counter_; std::mutex cond_mutex_; std::condition_variable cond_var_; - using SliceList = std::vector; - - const static int kShardCount = 8; std::unordered_map slice_queue_[kShardCount]; std::atomic slice_queue_count_[kShardCount]; TicketLock slice_queue_lock_[kShardCount]; @@ -104,17 +158,19 @@ class WorkerPool { collective_slice_queue_; std::atomic submitted_slice_count_, processed_slice_count_; + std::atomic recovery_activate_after_ns_{0}; // Rail state management: peer_nic_path -> RailState std::unordered_map rail_states_; std::mutex rail_state_lock_; // Rail monitor configuration - const static int kRailErrorThreshold = 5; // Errors before pause - const static uint64_t kRailPauseNs = 1000000000ull; // 1 second pause + const static int kRailErrorThreshold = 5; // Errors before pause + const static uint64_t kContextRecoveryDelayNs = + 30000000000ull; // 30 seconds before a recovered local RNIC is reused // Context-level health tracking - int context_failure_count_ = 0; + std::atomic context_failure_count_{0}; const static int kContextFailureThreshold = 32; // consecutive all-rails-failed }; diff --git a/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h index 88733b53e4..bafd3a09db 100644 --- a/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h +++ b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_communicator.h @@ -74,9 +74,10 @@ class RpcCommunicator { std::unique_ptr server_; std::function data_receive_callback_; - pybind11::handle py_callback_; + struct PyCallbackHolder; + std::unique_ptr py_callback_; std::shared_ptr> client_pools_; }; -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_interface.h b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_interface.h index eb6e19e559..752f829336 100644 --- a/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_interface.h +++ b/mooncake-transfer-engine/include/transport/rpc_communicator/rpc_interface.h @@ -43,12 +43,13 @@ class RpcInterface { RpcInterface(); ~RpcInterface(); + // pool_size limits cached RPC client connections per target endpoint. bool initialize(const std::string& listen_address = "", size_t thread_count = 0, size_t timeout_seconds = 30, - size_t pool_size = 10); + size_t pool_size = 100); // Convenience methods for common use cases - bool initializeClient(size_t pool_size = 10, size_t timeout_seconds = 30); + bool initializeClient(size_t pool_size = 100, size_t timeout_seconds = 30); bool initializeServer(const std::string& listen_address, size_t thread_count = 8, size_t timeout_seconds = 30); @@ -88,4 +89,4 @@ std::unique_ptr createRpcServer(uint64_t local_rank = 0, void bind_rpc_interface(pybind11::module_& m); -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h index 3e537cfec2..88b51c9169 100644 --- a/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h +++ b/mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -75,6 +76,9 @@ class TcpTransport : public Transport { Status submitTransferTask( const std::vector &task_list) override; + Status submitTransferTaskGroup( + const std::vector &task_list) override; + Status getTransferStatus(BatchID batch_id, size_t task_id, TransferStatus &status) override; @@ -102,7 +106,13 @@ class TcpTransport : public Transport { void worker(); - void startTransfer(Slice *slice); + Slice *prepareTransfer(TransferTask *task, const TransferRequest &request); + + void startTransfer(Slice *slice, + std::function continuation = nullptr, + bool reuse_connection = false); + + void startTransferSequence(std::vector slices); bool validateAddress(uint64_t addr, uint64_t size) const; @@ -139,9 +149,13 @@ class TcpTransport : public Transport { std::mutex pool_mutex_; std::shared_ptr getConnection( - const std::string &host, uint16_t port); + const std::string &host, uint16_t port, bool use_pool); void returnConnection(const std::string &host, uint16_t port, std::shared_ptr socket); + // Close `socket` and drop it from the pool: used for requests that did + // not terminate in a well-defined protocol state (#2086). + void discardConnection(const std::string &host, uint16_t port, + std::shared_ptr socket); void cleanupIdleConnections(); static constexpr std::chrono::seconds kConnectionIdleTimeout{60}; diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index 8e787f9c4b..4233705e26 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -60,6 +60,8 @@ class Transport { struct TransferRequest { enum OpCode { READ, WRITE }; + static constexpr uint64_t kNoTaskGroup = 0; + OpCode opcode; void *source; SegmentID target_id; @@ -68,6 +70,8 @@ class Transport { int advise_retry_cnt = 0; // Per-request transport pin, TENT only. int transport_hint = 0; + // Adjacent requests in the same group are one transport submission. + uint64_t task_group_id = kNoTaskGroup; }; enum TransferStatusEnum { @@ -85,6 +89,12 @@ class Transport { size_t transferred_bytes; }; + struct NicLoadStats { + std::string device_name; + uint64_t inflight_bytes{0}; + double ewma_bandwidth_bps{0.0}; + }; + struct BatchDesc; struct TransferTask; @@ -113,19 +123,33 @@ class Transport { TransferRequest::OpCode opcode; SegmentID target_id; std::string peer_nic_path; + std::string source_location; SliceStatus status; TransferTask *task; - std::vector dest_rkeys; + // EFA/CXI's libfabric MR keys are 64-bit (fi_mr_key()); RDMA verbs keys + // are 32-bit. Use a scoped alias so the width is defined in one place. +#if defined(USE_EFA) || defined(USE_CXI) + using mr_key_t = uint64_t; +#else + using mr_key_t = uint32_t; +#endif + std::vector dest_rkeys; bool from_cache; + // Optional resource cleanup invoked exactly once before the slice is + // deleted or returned to the thread-local cache. The callback must not + // delete the slice. + using CleanupCallback = void (*)(Slice *); + CleanupCallback cleanup_callback = nullptr; + union { struct { uint64_t dest_addr; - uint32_t source_lkey; - uint32_t dest_rkey; + mr_key_t source_lkey; + mr_key_t dest_rkey; int lkey_index; int rkey_index; - volatile int *qp_depth; + std::atomic *qp_depth; uint32_t retry_cnt; uint32_t max_retry_cnt; RdmaEndPoint *endpoint; // Endpoint used for this transfer @@ -144,6 +168,10 @@ class Transport { void *cuda_stream; // cudaStream_t, used by async NVLink // transport } local; + struct { + void *event; // cudaEvent_t + int device_id; + } nccl; struct { uint64_t dest_addr; } tcp; @@ -273,6 +301,12 @@ class Transport { } void deallocate(Slice *slice) { + // Clear before invoking so a cached slice cannot carry a + // transport-specific cleanup callback into its next use. + auto cleanup = slice->cleanup_callback; + slice->cleanup_callback = nullptr; + if (cleanup) cleanup(slice); + if (head_ - tail_ == kLazyDeleteSliceCapacity) { delete slice; return; @@ -369,6 +403,11 @@ class Transport { "Transport::submitTransferTask is not implemented"); } + virtual Status submitTransferTaskGroup( + const std::vector &task_list) { + return submitTransferTask(task_list); + } + /// @brief Get the status of a submitted transfer. This function shall not /// be called again after completion. /// @return Return 1 on completed (either success or failure); 0 if still in diff --git a/mooncake-transfer-engine/nvlink-allocator/CMakeLists.txt b/mooncake-transfer-engine/nvlink-allocator/CMakeLists.txt index 69e0d7149f..0ee627d7fd 100644 --- a/mooncake-transfer-engine/nvlink-allocator/CMakeLists.txt +++ b/mooncake-transfer-engine/nvlink-allocator/CMakeLists.txt @@ -1,12 +1,19 @@ include(${CMAKE_CURRENT_SOURCE_DIR}/../fabric_allocator.cmake) set(_extra_build_opts "") +set(_extra_build_depends "") if(USE_HIP) list(APPEND _extra_build_opts --use-hipcc) + list(APPEND _extra_build_depends + ${CMAKE_CURRENT_SOURCE_DIR}/../include/gpu_vendor/hip.h) elseif(USE_MUSA) list(APPEND _extra_build_opts --use-mcc) + list(APPEND _extra_build_depends + ${CMAKE_CURRENT_SOURCE_DIR}/../include/gpu_vendor/musa.h) elseif(USE_MACA) list(APPEND _extra_build_opts --use-maca) + list(APPEND _extra_build_depends + ${CMAKE_CURRENT_SOURCE_DIR}/../include/gpu_vendor/maca.h) endif() set(_enable_nvlink_allocator_build FALSE) @@ -22,8 +29,15 @@ add_fabric_allocator_build_target( build_nvlink_allocator BUILD_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/build.sh + OUTPUT_NAME + nvlink_allocator.so BUILD_ARGS ${_extra_build_opts} + BUILD_DEPENDS + ${CMAKE_CURRENT_SOURCE_DIR}/nvlink_allocator.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/../include/cuda_alike.h + ${CMAKE_CURRENT_SOURCE_DIR}/../scripts/allocator_build_common.sh + ${_extra_build_depends} COMMENT "Building nvlink allocator to ${CMAKE_CURRENT_BINARY_DIR}" ENABLE_BUILD diff --git a/mooncake-transfer-engine/scripts/real_rdma_link_failover.sh b/mooncake-transfer-engine/scripts/real_rdma_link_failover.sh new file mode 100755 index 0000000000..53eebb474f --- /dev/null +++ b/mooncake-transfer-engine/scripts/real_rdma_link_failover.sh @@ -0,0 +1,389 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +BUILD_DIR="${BUILD_DIR:-$ROOT_DIR/build}" +BENCH="${BENCH:-$BUILD_DIR/mooncake-transfer-engine/example/transfer_engine_bench}" +VALIDATOR="${VALIDATOR:-$BUILD_DIR/mooncake-transfer-engine/example/transfer_engine_validator}" +LOG_DIR="${LOG_DIR:-/tmp/mooncake-real-rdma-link-failover-$(date +%s)}" + +FAULT_SIDE="receiver" +DOWN_AFTER_SEC=5 +DOWN_FOR_SEC=7 +DURATION_SEC=200 +TARGET_PRIMARY="mlx5_3" +TARGET_AVAILABLE="mlx5_4" +SENDER_PRIMARY="mlx5_2" +SENDER_AVAILABLE="mlx5_1" +NETDEV_OVERRIDE="" +VALIDATE_DATA="${VALIDATE_DATA:-1}" + +usage() { + cat <&2; usage; exit 2 ;; + esac +done + +if [[ "$FAULT_SIDE" != "sender" && "$FAULT_SIDE" != "receiver" ]]; then + echo "--fault-side must be sender or receiver" >&2 + exit 2 +fi + +validate_single_rnic() { + local label="$1" + local value="$2" + if [[ -z "$value" ]]; then + echo "$label RNIC must not be empty" >&2 + exit 2 + fi + if [[ "$value" == *,* ]]; then + echo "$label RNIC must be exactly one device for this failover validation: $value" >&2 + exit 2 + fi +} + +validate_pair() { + local side="$1" + local primary="$2" + local available="$3" + validate_single_rnic "$side primary" "$primary" + validate_single_rnic "$side available" "$available" + if [[ "$primary" == "$available" ]]; then + echo "$side primary and available RNIC must be different: $primary" >&2 + exit 2 + fi +} + +validate_pair "target" "$TARGET_PRIMARY" "$TARGET_AVAILABLE" +validate_pair "sender" "$SENDER_PRIMARY" "$SENDER_AVAILABLE" + +run_privileged() { + if [[ "${EUID:-$(id -u)}" -eq 0 ]]; then + "$@" + else + sudo "$@" + fi +} + +require_privilege() { + if [[ "${EUID:-$(id -u)}" -eq 0 ]]; then + return + fi + if ! sudo -n true 2>/dev/null; then + echo "This test must change link state. Run it with sudo, for example:" >&2 + echo " sudo $0 --fault-side $FAULT_SIDE" >&2 + exit 1 + fi +} + +rdma_to_netdev() { + local rdma_dev="$1" + local ndev_dir="/sys/class/infiniband/$rdma_dev/ports/1/gid_attrs/ndevs" + if [[ ! -d "$ndev_dir" ]]; then + echo "Cannot find $ndev_dir for RDMA device $rdma_dev" >&2 + return 1 + fi + + local netdev="" + local entry value + while IFS= read -r entry; do + value="$(cat "$entry" 2>/dev/null || true)" + value="${value%%[[:space:]]*}" + if [[ -n "$value" && "$value" != "(null)" && + -e "/sys/class/net/$value" ]]; then + netdev="$value" + break + fi + done < <(find "$ndev_dir" -maxdepth 1 -type f -printf '%f %p\n' | + sort -n | + awk '{ print $2 }') + if [[ -z "$netdev" ]]; then + echo "No Linux netdev found under $ndev_dir for RDMA device $rdma_dev" >&2 + return 1 + fi + printf '%s\n' "$netdev" +} + +make_matrix() { + local preferred="$1" + local fallback="$2" + printf '{"cpu:0":[[%s],[%s]],"cpu:1":[[%s],[%s]]}\n' \ + "$(quote_csv "$preferred")" "$(quote_csv "$fallback")" \ + "$(quote_csv "$preferred")" "$(quote_csv "$fallback")" +} + +quote_csv() { + local csv="$1" + local out="" + local item + IFS=',' read -ra parts <<< "$csv" + for item in "${parts[@]}"; do + [[ -z "$item" ]] && continue + if [[ -n "$out" ]]; then out+=","; fi + out+="\"$item\"" + done + printf '%s' "$out" +} + +wait_for_log() { + local file="$1" + local pattern="$2" + local timeout="$3" + local deadline=$((SECONDS + timeout)) + while (( SECONDS < deadline )); do + if [[ -f "$file" ]] && grep -q "$pattern" "$file"; then + return 0 + fi + sleep 0.2 + done + return 1 +} + +RUNNER="$BENCH" +BUILD_TARGET="transfer_engine_bench" +if [[ "$VALIDATE_DATA" == "1" ]]; then + RUNNER="$VALIDATOR" + BUILD_TARGET="transfer_engine_validator" +fi + +if [[ "${SKIP_BUILD:-0}" != "1" ]]; then + echo "Building $BUILD_TARGET at $RUNNER ..." + cmake --build "$BUILD_DIR" --target "$BUILD_TARGET" -j"$(nproc)" +elif [[ ! -x "$RUNNER" ]]; then + echo "$BUILD_TARGET not found at $RUNNER and SKIP_BUILD=1" >&2 + exit 1 +fi + +require_privilege + +FAULT_RDMA="$TARGET_PRIMARY" +if [[ "$FAULT_SIDE" == "sender" ]]; then + FAULT_RDMA="$SENDER_PRIMARY" +fi + +FAULT_NETDEV="$NETDEV_OVERRIDE" +if [[ -z "$FAULT_NETDEV" ]]; then + FAULT_NETDEV="$(rdma_to_netdev "$FAULT_RDMA")" +fi + +mkdir -p "$LOG_DIR" +TARGET_MATRIX_FILE="$LOG_DIR/target_nic_priority.json" +SENDER_MATRIX_FILE="$LOG_DIR/sender_nic_priority.json" +TARGET_LOG="$LOG_DIR/target.log" +INITIATOR_LOG="$LOG_DIR/initiator.log" + +make_matrix "$TARGET_PRIMARY" "$TARGET_AVAILABLE" > "$TARGET_MATRIX_FILE" +make_matrix "$SENDER_PRIMARY" "$SENDER_AVAILABLE" > "$SENDER_MATRIX_FILE" + +TARGET_PORT="${TARGET_PORT:-$((30000 + RANDOM % 10000))}" +INITIATOR_PORT="${INITIATOR_PORT:-$((40000 + RANDOM % 10000))}" +TARGET_SERVER="127.0.0.1:$TARGET_PORT" +INITIATOR_SERVER="127.0.0.1:$INITIATOR_PORT" + +COMMON_FLAGS=( + --metadata_server=P2PHANDSHAKE + --protocol=rdma + --buffer_size=$((256 * 1024 * 1024)) + --block_size=$((4 * 1024 * 1024)) + --batch_size=1 + --threads=1 + --duration="$DURATION_SEC" +) + +if [[ "$VALIDATE_DATA" != "1" ]]; then + COMMON_FLAGS+=(--operation=write) +fi + +if "$RUNNER" --help 2>&1 | grep -q -- '--use_vram'; then + COMMON_FLAGS+=(--use_vram=false) +fi + +TARGET_PID="" +LINK_WAS_DOWN=0 + +cleanup() { + local rc=$? + if [[ "$LINK_WAS_DOWN" -eq 1 ]]; then + echo "[cleanup] Restoring $FAULT_NETDEV up" + run_privileged ip link set "$FAULT_NETDEV" up || true + fi + if [[ -n "$TARGET_PID" ]] && kill -0 "$TARGET_PID" 2>/dev/null; then + kill -TERM "$TARGET_PID" 2>/dev/null || true + wait "$TARGET_PID" 2>/dev/null || true + fi + exit "$rc" +} +trap cleanup EXIT INT TERM + +echo "=== Real RDMA link failover test ===" +echo "fault_side : $FAULT_SIDE" +echo "fault_rdma : $FAULT_RDMA" +echo "fault_netdev : $FAULT_NETDEV" +echo "target primary : $TARGET_PRIMARY" +echo "target available : $TARGET_AVAILABLE" +echo "sender primary : $SENDER_PRIMARY" +echo "sender available : $SENDER_AVAILABLE" +echo "target matrix : $(cat "$TARGET_MATRIX_FILE")" +echo "sender matrix : $(cat "$SENDER_MATRIX_FILE")" +echo "target server : $TARGET_SERVER" +echo "initiator server : $INITIATOR_SERVER" +echo "runner : $RUNNER" +echo "validate_data : $VALIDATE_DATA" +echo "logs : $LOG_DIR" +echo +echo "WARNING: this will run: ip link set $FAULT_NETDEV down; then up." +echo + +"$RUNNER" "${COMMON_FLAGS[@]}" \ + --mode=target \ + --local_server_name="$TARGET_SERVER" \ + --device_name="$TARGET_PRIMARY,$TARGET_AVAILABLE" \ + --nic_priority_matrix="$TARGET_MATRIX_FILE" \ + >"$TARGET_LOG" 2>&1 & +TARGET_PID=$! + +if ! wait_for_log "$TARGET_LOG" "RDMA device" 15; then + echo "Target did not start successfully. See $TARGET_LOG" >&2 + exit 1 +fi + +TARGET_SEGMENT="$( + awk '/Transfer Engine RPC using P2P handshake, listening on / { + print $NF; + exit; + }' "$TARGET_LOG" +)" +if [[ -z "$TARGET_SEGMENT" ]]; then + TARGET_SEGMENT="$TARGET_SERVER" +fi +echo "target segment : $TARGET_SEGMENT" + +"$RUNNER" "${COMMON_FLAGS[@]}" \ + --mode=initiator \ + --local_server_name="$INITIATOR_SERVER" \ + --segment_id="$TARGET_SEGMENT" \ + --device_name="$SENDER_PRIMARY,$SENDER_AVAILABLE" \ + --nic_priority_matrix="$SENDER_MATRIX_FILE" \ + >"$INITIATOR_LOG" 2>&1 & +INITIATOR_PID=$! + +sleep "$DOWN_AFTER_SEC" + +echo "[$(date '+%F %T')] Bringing $FAULT_NETDEV down ($FAULT_RDMA)" +run_privileged ip link set "$FAULT_NETDEV" down +LINK_WAS_DOWN=1 + +sleep "$DOWN_FOR_SEC" + +echo "[$(date '+%F %T')] Bringing $FAULT_NETDEV up ($FAULT_RDMA)" +run_privileged ip link set "$FAULT_NETDEV" up +LINK_WAS_DOWN=0 + +set +e +wait "$INITIATOR_PID" +INITIATOR_RC=$? +set -e + +kill -TERM "$TARGET_PID" 2>/dev/null || true +wait "$TARGET_PID" 2>/dev/null || true +TARGET_PID="" + +echo +echo "=== Result ===" +echo "initiator exit code: $INITIATOR_RC" +echo "target log : $TARGET_LOG" +echo "initiator log : $INITIATOR_LOG" + +grep_output="$(grep -E "FAILED|failed|Transport retry counter exceeded|Cannot make connection|KVTransferError|AddressNotRegistered|Detect data integrity problem" "$INITIATOR_LOG" "$TARGET_LOG" 2>/dev/null | tail -n 80 || true)" +if [[ -n "$grep_output" ]]; then + echo + echo "Important failure-like log lines:" + echo "$grep_output" +fi + +if [[ "$FAULT_SIDE" == "sender" ]]; then + CHECK_LOG="$INITIATOR_LOG" + CHECK_PRIMARY="$SENDER_PRIMARY" + CHECK_AVAILABLE="$SENDER_AVAILABLE" +else + CHECK_LOG="$TARGET_LOG" + CHECK_PRIMARY="$TARGET_PRIMARY" + CHECK_AVAILABLE="$TARGET_AVAILABLE" +fi + +if ! grep -q "RDMA device: $CHECK_PRIMARY" "$CHECK_LOG"; then + echo "FAIL: $FAULT_SIDE primary RNIC $CHECK_PRIMARY was not initialized" >&2 + exit 1 +fi +if ! grep -q "RDMA device: $CHECK_AVAILABLE" "$CHECK_LOG"; then + echo "FAIL: $FAULT_SIDE available RNIC $CHECK_AVAILABLE was not initialized" >&2 + exit 1 +fi +if ! grep -q "Context $CHECK_PRIMARY is now inactive" "$CHECK_LOG"; then + echo "FAIL: did not observe $FAULT_SIDE primary RNIC $CHECK_PRIMARY going inactive" >&2 + exit 1 +fi +if grep -q "Context $CHECK_AVAILABLE is now inactive" "$CHECK_LOG"; then + echo "FAIL: $FAULT_SIDE available RNIC $CHECK_AVAILABLE also went inactive" >&2 + exit 1 +fi + +if [[ "$INITIATOR_RC" -ne 0 ]]; then + echo "FAIL: initiator exited non-zero" + exit "$INITIATOR_RC" +fi + +if ! grep -q "Test completed" "$INITIATOR_LOG"; then + echo "FAIL: initiator did not report completion" + exit 1 +fi + +echo "PASS: transfer completed across real $FAULT_RDMA/$FAULT_NETDEV down/up" diff --git a/mooncake-transfer-engine/scripts/register.py b/mooncake-transfer-engine/scripts/register.py index 5db9853846..f178c20a7d 100644 --- a/mooncake-transfer-engine/scripts/register.py +++ b/mooncake-transfer-engine/scripts/register.py @@ -31,7 +31,6 @@ etcd_host = sys.argv[1] segment_name = "mooncake/nvmeof/" + sys.argv[2] files = sys.argv[3:] - local_server_name = socket.gethostname() server_name = socket.gethostname() etcd = etcd3.client(host=etcd_host, port=2379) diff --git a/mooncake-transfer-engine/src/CMakeLists.txt b/mooncake-transfer-engine/src/CMakeLists.txt index 1bd77e1d80..c91b02147f 100644 --- a/mooncake-transfer-engine/src/CMakeLists.txt +++ b/mooncake-transfer-engine/src/CMakeLists.txt @@ -63,13 +63,25 @@ if(USE_BAREX) endif() if(USE_CUDA) - target_include_directories(transfer_engine PRIVATE /usr/local/cuda/include) - target_link_libraries(transfer_engine PUBLIC cuda cudart rt mlx5) + target_include_directories(transfer_engine PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + if(USE_CXI) + target_link_libraries(transfer_engine PUBLIC CUDA::cuda_driver CUDA::cudart rt) + else() + target_link_libraries(transfer_engine PUBLIC CUDA::cuda_driver CUDA::cudart rt mlx5) + endif() if(USE_NVMEOF) target_link_libraries(transfer_engine PUBLIC nvmeof_transport cufile) endif() endif() +if(USE_NCCL_DEVICE OR USE_NCCL_HOST) + target_link_libraries(transfer_engine PUBLIC NCCL::nccl) +endif() + +if(USE_SUPA) + target_link_libraries(transfer_engine PUBLIC supa supart) +endif() + if(USE_MACA) if(NOT DEFINED MACA_RUNTIME_LIBS) set(MACA_RUNTIME_LIBS mcruntime mxc-runtime64 rt) @@ -97,17 +109,27 @@ if(USE_HIP) target_include_directories(transfer_engine PRIVATE ${HIP_INCLUDE_DIRS}) target_link_libraries(transfer_engine PUBLIC hip::host rt) - # Optional dmabuf MR registration path (requires hsa-runtime64 and a - # kernel with CONFIG_PCI_P2PDMA + CONFIG_DMABUF_MOVE_NOTIFY). + # Optional dmabuf MR registration path (requires hsa-runtime64 and a kernel + # with CONFIG_PCI_P2PDMA + CONFIG_DMABUF_MOVE_NOTIFY). option(USE_HIP_DMABUF "Enable HIP dmabuf RDMA MR registration" ON) if(USE_HIP_DMABUF) find_package(hsa-runtime64 CONFIG) if(hsa-runtime64_FOUND) - target_compile_definitions(transfer_engine PRIVATE USE_HIP_DMABUF) + # The dmabuf MR-registration code lives in rdma_context.cpp, which is + # compiled in the rdma_transport OBJECT library and pulled into + # transfer_engine via $. A PRIVATE define on + # transfer_engine never reaches that object compilation, so the dmabuf + # path is silently compiled out and GPU MRs fall back to plain + # ibv_reg_mr (EINVAL on device memory). Put the define (and the hsa + # include/link usage requirements) on the target that actually compiles + # rdma_context.cpp. + target_compile_definitions(rdma_transport PRIVATE USE_HIP_DMABUF) + target_link_libraries(rdma_transport PUBLIC hsa-runtime64::hsa-runtime64) target_link_libraries(transfer_engine PUBLIC hsa-runtime64::hsa-runtime64) message(STATUS "HIP dmabuf MR registration enabled (hsa-runtime64 found)") else() - message(STATUS "HIP dmabuf MR registration disabled (hsa-runtime64 not found)") + message( + STATUS "HIP dmabuf MR registration disabled (hsa-runtime64 not found)") endif() else() message(STATUS "HIP dmabuf MR registration disabled (USE_HIP_DMABUF=OFF)") @@ -158,6 +180,11 @@ if(USE_EFA) message(STATUS "Enabled USE_EFA (AWS Elastic Fabric Adapter) support") target_link_libraries(transfer_engine PUBLIC fabric efa_transport) endif() +if(USE_CXI) + message(STATUS "Enabled USE_CXI (HPE Cray Slingshot) support") + target_link_libraries(transfer_engine PUBLIC fabric cxi_transport) +endif() + if(USE_UB) message(STATUS "Enabled USE_UB protocol support") target_link_libraries(transfer_engine PUBLIC ub_transport) diff --git a/mooncake-transfer-engine/src/config.cpp b/mooncake-transfer-engine/src/config.cpp index aef962c230..e3b0154659 100644 --- a/mooncake-transfer-engine/src/config.cpp +++ b/mooncake-transfer-engine/src/config.cpp @@ -14,14 +14,92 @@ #include "config.h" +#include #include #include #include #include #include +#include #include namespace mooncake { +namespace { + +std::string trimConfigToken(const std::string& value) { + const auto begin = value.find_first_not_of(" \t\n\r"); + if (begin == std::string::npos) return ""; + const auto end = value.find_last_not_of(" \t\n\r"); + return value.substr(begin, end - begin + 1); +} + +std::vector splitConfigString(const std::string& value, + char delim) { + std::vector result; + std::stringstream stream(value); + std::string item; + while (std::getline(stream, item, delim)) { + result.push_back(trimConfigToken(item)); + } + return result; +} + +bool parseBoolConfigEnv(const char* value, const char* env_name, bool& output) { + if (strcmp(value, "1") == 0 || strcasecmp(value, "true") == 0) { + output = true; + return true; + } + if (strcmp(value, "0") == 0 || strcasecmp(value, "false") == 0) { + output = false; + return true; + } + LOG(WARNING) << "Ignore value from environment variable " << env_name + << ", it should be 0|1|true|false"; + return false; +} + +void parseNicPeerAffinity( + const char* env, + std::unordered_map>& affinity) { + affinity.clear(); + if (!env || env[0] == '\0') return; + + for (const auto& raw_rule : splitConfigString(env, ';')) { + const auto rule = trimConfigToken(raw_rule); + if (rule.empty()) continue; + + auto delim = rule.find('='); + if (delim == std::string::npos) { + LOG(WARNING) << "Invalid MC_NIC_PEER_AFFINITY rule '" << rule + << "'. Expected local_hca=peer_hca[,peer_hca]."; + continue; + } + + auto local_hca = trimConfigToken(rule.substr(0, delim)); + if (local_hca.empty()) { + LOG(WARNING) << "Invalid MC_NIC_PEER_AFFINITY rule '" << rule + << "': local HCA is empty."; + continue; + } + + std::vector peer_hcas; + for (const auto& peer_hca : + splitConfigString(rule.substr(delim + 1), ',')) { + if (!peer_hca.empty()) peer_hcas.push_back(peer_hca); + } + + if (peer_hcas.empty()) { + LOG(WARNING) << "Invalid MC_NIC_PEER_AFFINITY rule '" << rule + << "': peer HCA list is empty."; + continue; + } + + affinity[local_hca] = std::move(peer_hcas); + } +} + +} // namespace + void loadGlobalConfig(GlobalConfig& config) { const char* num_cq_per_ctx_env = std::getenv("MC_NUM_CQ_PER_CTX"); if (num_cq_per_ctx_env) { @@ -128,10 +206,12 @@ void loadGlobalConfig(GlobalConfig& config) { const char* max_wr_env = std::getenv("MC_MAX_WR"); if (max_wr_env) { size_t val = atoi(max_wr_env); - if (val > 0 && val <= UINT16_MAX) + if (val > 0 && val <= UINT16_MAX) { config.max_wr = val; - else + config.max_wr_from_env = true; + } else { LOG(WARNING) << "Ignore value from environment variable MC_MAX_WR"; + } } const char* max_inline_env = std::getenv("MC_MAX_INLINE"); @@ -240,6 +320,27 @@ void loadGlobalConfig(GlobalConfig& config) { config.metacache = false; } + const char* te_metadata_refresh_interval_seconds = + std::getenv("MC_TE_METADATA_REFRESH_INTERVAL_SECONDS"); + if (te_metadata_refresh_interval_seconds) { + try { + int val = std::stoi(te_metadata_refresh_interval_seconds); + if (val >= 0) { + config.te_metadata_refresh_interval_seconds = + static_cast(val); + } else { + LOG(WARNING) << "Ignore value from environment variable " + "MC_TE_METADATA_REFRESH_INTERVAL_SECONDS"; + } + } catch (const std::exception& e) { + LOG(WARNING) << "Invalid MC_TE_METADATA_REFRESH_INTERVAL_SECONDS " + "environment " + "value: " + << te_metadata_refresh_interval_seconds + << ". Error: " << e.what(); + } + } + const char* handshake_listen_backlog = std::getenv("MC_HANDSHAKE_LISTEN_BACKLOG"); if (handshake_listen_backlog) { @@ -269,6 +370,24 @@ void loadGlobalConfig(GlobalConfig& config) { "MC_HANDSHAKE_CONNECT_TIMEOUT"; } + const char* rdma_rail_pause_seconds = + std::getenv("MC_RDMA_RAIL_PAUSE_SECONDS"); + if (rdma_rail_pause_seconds) { + try { + int val = std::stoi(rdma_rail_pause_seconds); + if (val > 0 && val < 3600) { + config.rdma_rail_pause_seconds = static_cast(val); + } else { + LOG(WARNING) << "Ignore value from environment variable " + "MC_RDMA_RAIL_PAUSE_SECONDS"; + } + } catch (const std::exception& e) { + LOG(WARNING) << "Invalid MC_RDMA_RAIL_PAUSE_SECONDS environment " + "value: " + << rdma_rail_pause_seconds << ". Error: " << e.what(); + } + } + const char* log_level = std::getenv("MC_LOG_LEVEL"); config.trace = false; if (log_level) { @@ -295,6 +414,31 @@ void loadGlobalConfig(GlobalConfig& config) { << "Ignore value from environment variable MC_SLICE_TIMEOUT"; } + const char* conn_pause_ttl_env = std::getenv("MC_CONN_PAUSE_TTL_MS"); + if (conn_pause_ttl_env) { + // Robust parse (not atoi): a non-numeric typo must keep the default + // rather than silently resolve to 0. 0 is a valid explicit "disable"; + // negative / out-of-range / garbage are rejected, preserving the + // default. + int val = 0; + const char* end = conn_pause_ttl_env + strlen(conn_pause_ttl_env); + auto [ptr, ec] = std::from_chars(conn_pause_ttl_env, end, val); + if (ec == std::errc() && ptr == end) { + if (val >= 0 && val <= 600000) { + config.conn_pause_ttl_ms = val; + } else { + LOG(WARNING) << "Ignore value from environment variable " + "MC_CONN_PAUSE_TTL_MS, value " + << conn_pause_ttl_env + << " out of range (should be 0-600000)"; + } + } else { + LOG(WARNING) << "Invalid MC_CONN_PAUSE_TTL_MS environment value: " + << conn_pause_ttl_env + << ". Expected an integer in range 0-600000"; + } + } + const char* log_dir_path = std::getenv("MC_LOG_DIR"); if (log_dir_path) { google::InitGoogleLogging("mooncake-transfer-engine"); @@ -348,6 +492,42 @@ void loadGlobalConfig(GlobalConfig& config) { config.enable_dest_device_affinity = true; } + const char* enable_hca_peer_affinity_env = + std::getenv("MC_ENABLE_HCA_PEER_AFFINITY"); + if (enable_hca_peer_affinity_env) { + parseBoolConfigEnv(enable_hca_peer_affinity_env, + "MC_ENABLE_HCA_PEER_AFFINITY", + config.enable_hca_peer_affinity); + } + + parseNicPeerAffinity(std::getenv("MC_NIC_PEER_AFFINITY"), + config.nic_peer_affinity); + + if (config.enable_hca_peer_affinity && config.enable_dest_device_affinity) { + LOG(ERROR) << "MC_ENABLE_HCA_PEER_AFFINITY and " + "MC_ENABLE_DEST_DEVICE_AFFINITY cannot be enabled at " + "the same time; falling back to default peer device " + "selection."; + config.enable_hca_peer_affinity = false; + config.enable_dest_device_affinity = false; + } + + const char* log_rdma_slice_affinity_env = + std::getenv("MC_LOG_RDMA_SLICE_AFFINITY"); + if (log_rdma_slice_affinity_env) { + parseBoolConfigEnv(log_rdma_slice_affinity_env, + "MC_LOG_RDMA_SLICE_AFFINITY", + config.log_rdma_slice_affinity); + } + + const char* track_rdma_posted_slices_env = + std::getenv("MC_TRACK_RDMA_POSTED_SLICES"); + if (track_rdma_posted_slices_env) { + parseBoolConfigEnv(track_rdma_posted_slices_env, + "MC_TRACK_RDMA_POSTED_SLICES", + config.track_rdma_posted_slices); + } + const char* enable_parallel_reg_mr = std::getenv("MC_ENABLE_PARALLEL_REG_MR"); if (enable_parallel_reg_mr) { @@ -360,6 +540,38 @@ void loadGlobalConfig(GlobalConfig& config) { } } + const char* max_concurrent_reg_mr = std::getenv("MC_MAX_CONCURRENT_REG_MR"); + if (max_concurrent_reg_mr) { + // Robust parse (not atol): a non-numeric typo must keep the default + // rather than silently resolve to 0, which here means "no cap" and so + // would read as a deliberate request for the old unbounded behavior. + // 0 is a valid explicit way to ask for no cap; negative and garbage are + // rejected. + size_t val = 0; + const char* end = max_concurrent_reg_mr + strlen(max_concurrent_reg_mr); + auto [ptr, ec] = std::from_chars(max_concurrent_reg_mr, end, val); + if (ec == std::errc() && ptr == end) { + config.max_concurrent_reg_mr = val; + } else { + LOG(WARNING) << "Invalid MC_MAX_CONCURRENT_REG_MR environment " + "value: " + << max_concurrent_reg_mr << ", keeping default"; + } + } + + const char* efa_nic_selection = std::getenv("MC_EFA_NIC_SELECTION"); + if (efa_nic_selection) { + if (strcasecmp(efa_nic_selection, "all") == 0) { + config.efa_nic_selection = EfaNicSelection::ALL; + } else if (strcasecmp(efa_nic_selection, "local") == 0) { + config.efa_nic_selection = EfaNicSelection::LOCAL; + } else { + LOG(WARNING) << "Invalid MC_EFA_NIC_SELECTION environment value: " + << efa_nic_selection + << ", expected all|local, keeping default"; + } + } + const char* endpoint_store_type_env = std::getenv("MC_ENDPOINT_STORE_TYPE"); if (endpoint_store_type_env) { if (strcmp(endpoint_store_type_env, "FIFO") == 0) { @@ -390,6 +602,24 @@ void loadGlobalConfig(GlobalConfig& config) { } } + const char* service_level_env = std::getenv("MC_IB_SL"); + if (service_level_env) { + try { + int val = std::stoi(service_level_env); + if (val >= 0 && val <= 15) { + config.ib_service_level = val; + } else { + LOG(WARNING) + << "Ignore value from environment variable MC_IB_SL, " + << "value " << service_level_env + << " out of range (should be 0-15)"; + } + } catch (const std::exception& e) { + LOG(WARNING) << "Invalid MC_IB_SL environment value: " + << service_level_env << ". Error: " << e.what(); + } + } + const char* ib_relaxed_ordering_env = std::getenv("MC_IB_PCI_RELAXED_ORDERING"); if (ib_relaxed_ordering_env) { @@ -506,7 +736,14 @@ void dumpGlobalConfig() { LOG(INFO) << "max_inline = " << config.max_inline; LOG(INFO) << "mtu_length = " << mtuLengthToString(config.mtu_length); LOG(INFO) << "parallel_reg_mr = " << config.parallel_reg_mr; + LOG(INFO) << "efa_nic_selection = " + << (config.efa_nic_selection == EfaNicSelection::LOCAL ? "local" + : "all"); LOG(INFO) << "ib_traffic_class = " << config.ib_traffic_class; + LOG(INFO) << "ib_service_level = " << config.ib_service_level; + LOG(INFO) << "te_metadata_refresh_interval_seconds = " + << config.te_metadata_refresh_interval_seconds; + LOG(INFO) << "rdma_rail_pause_seconds = " << config.rdma_rail_pause_seconds; { std::ostringstream oss; for (size_t i = 0; i < config.mlx5_qp_udp_sports.size(); ++i) { @@ -519,6 +756,10 @@ void dumpGlobalConfig() { } LOG(INFO) << "mlx5_qp_lag_port_balance = " << (config.mlx5_qp_lag_port_balance ? "true" : "false"); + LOG(INFO) << "log_rdma_slice_affinity = " + << (config.log_rdma_slice_affinity ? "true" : "false"); + LOG(INFO) << "track_rdma_posted_slices = " + << (config.track_rdma_posted_slices ? "true" : "false"); } GlobalConfig& globalConfig() { diff --git a/mooncake-transfer-engine/src/graceful_shutdown.cpp b/mooncake-transfer-engine/src/graceful_shutdown.cpp new file mode 100644 index 0000000000..39187add49 --- /dev/null +++ b/mooncake-transfer-engine/src/graceful_shutdown.cpp @@ -0,0 +1,206 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "graceful_shutdown.h" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "transfer_engine_impl.h" + +namespace mooncake { + +namespace { + +std::mutex g_registry_mutex; +std::vector> g_tokens; +std::atomic g_cleanup_started{false}; + +std::mutex g_install_mutex; +bool g_handlers_installed = false; +volatile sig_atomic_t g_handlers_pid = 0; +bool g_atexit_registered = false; +int g_signal_pipe[2] = {-1, -1}; +volatile sig_atomic_t g_signal_seen = 0; + +class TransferEngineImplShutdownToken : public ShutdownToken { + public: + explicit TransferEngineImplShutdownToken( + std::shared_ptr impl) + : impl_(std::move(impl)) {} + + void shutdown() override { + if (auto impl = impl_.lock()) { + impl->freeEngine(); + } + } + + void detach() override { impl_.reset(); } + + private: + std::weak_ptr impl_; +}; + +std::vector> collectTokensForCleanup() { + std::lock_guard lock(g_registry_mutex); + auto tokens = std::move(g_tokens); + g_tokens.clear(); + return tokens; +} + +void cleanupEngines() { + if (g_cleanup_started.exchange(true)) return; + + auto tokens = collectTokensForCleanup(); + for (auto& token : tokens) { + token->shutdown(); + } +} + +void atexitCleanup() { cleanupEngines(); } + +void signalWatcher() { + unsigned char signal_byte = 0; + ssize_t bytes_read = 0; + do { + bytes_read = read(g_signal_pipe[0], &signal_byte, sizeof(signal_byte)); + } while (bytes_read < 0 && errno == EINTR); + + if (bytes_read == static_cast(sizeof(signal_byte))) { + cleanupEngines(); + _Exit(128 + static_cast(signal_byte)); + } + + _Exit(1); +} + +bool startSignalWatcherLocked(pid_t current_pid) { + if (g_signal_pipe[0] >= 0) close(g_signal_pipe[0]); + if (g_signal_pipe[1] >= 0) close(g_signal_pipe[1]); + g_signal_pipe[0] = -1; + g_signal_pipe[1] = -1; + g_signal_seen = 0; + + if (pipe(g_signal_pipe) != 0) return false; + + // The watcher must never take SIGTERM/SIGINT itself: the handler writes + // to the pipe and then pauses forever, so if it ran on the watcher + // thread no reader would be left and cleanup would never run. Block both + // signals around thread creation so the watcher inherits the mask, and + // restore the caller's original mask on every path (it may have had its + // own reasons to block these signals). + sigset_t watcher_block_set; + sigset_t saved_mask; + sigemptyset(&watcher_block_set); + sigaddset(&watcher_block_set, SIGTERM); + sigaddset(&watcher_block_set, SIGINT); + const bool mask_blocked = + pthread_sigmask(SIG_BLOCK, &watcher_block_set, &saved_mask) == 0; + + bool watcher_started = false; + try { + std::thread(signalWatcher).detach(); + watcher_started = true; + } catch (const std::exception& e) { + LOG(WARNING) << "action=start_signal_watcher_failed, error=" + << e.what(); + } catch (...) { + } + + if (mask_blocked) pthread_sigmask(SIG_SETMASK, &saved_mask, nullptr); + + if (!watcher_started) { + close(g_signal_pipe[0]); + close(g_signal_pipe[1]); + g_signal_pipe[0] = -1; + g_signal_pipe[1] = -1; + return false; + } + + g_handlers_pid = static_cast(current_pid); + return true; +} + +void shutdownSignalHandler(int signo) { + // After fork(), only the calling thread survives. If the parent had already + // installed handlers, the child inherits the handler and pipe fds but not + // the watcher thread. Exit directly instead of blocking forever in pause(). + if (g_handlers_pid != static_cast(getpid())) { + _Exit(128 + signo); + } + + if (g_signal_seen == 0) { + g_signal_seen = signo; + unsigned char signal_byte = static_cast(signo); + if (g_signal_pipe[1] < 0 || + write(g_signal_pipe[1], &signal_byte, sizeof(signal_byte)) != + static_cast(sizeof(signal_byte))) { + _Exit(128 + signo); + } + } + + for (;;) { + pause(); + } +} + +} // namespace + +void registerTokenForShutdown(std::shared_ptr token) { + std::lock_guard lock(g_registry_mutex); + g_tokens.push_back(std::move(token)); +} + +void registerEngineForShutdown(std::shared_ptr impl) { + registerTokenForShutdown( + std::make_shared(std::move(impl))); +} + +void installGracefulShutdownHandlers() { + std::lock_guard lock(g_install_mutex); + pid_t current_pid = getpid(); + if (g_handlers_installed && + g_handlers_pid == static_cast(current_pid)) + return; + + if (!g_atexit_registered) { + atexit(atexitCleanup); + g_atexit_registered = true; + } + + if (!startSignalWatcherLocked(current_pid)) return; + + struct sigaction sa{}; + sa.sa_handler = shutdownSignalHandler; + sigemptyset(&sa.sa_mask); + sigaddset(&sa.sa_mask, SIGTERM); + sigaddset(&sa.sa_mask, SIGINT); + sa.sa_flags = 0; + sigaction(SIGTERM, &sa, nullptr); + sigaction(SIGINT, &sa, nullptr); + + g_handlers_installed = true; +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/memory_location.cpp b/mooncake-transfer-engine/src/memory_location.cpp index a18b246052..32b488b05e 100644 --- a/mooncake-transfer-engine/src/memory_location.cpp +++ b/mooncake-transfer-engine/src/memory_location.cpp @@ -30,6 +30,24 @@ std::string genGpuNodeName(int node) { return kWildcardLocation; } +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ + defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ + defined(USE_COREX) || defined(USE_SUNRISE) +// A GPU-less host (e.g. an RDMA-only sidecar) has no device for +// cudaPointerGetAttributes to classify, yet a CUDA-enabled build probes +// every buffer, each call failing and logging a per-buffer ERROR that +// buries real logs. Detect the absence once and skip the probe. +static bool detectCudaDevicePresent() { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + LOG(WARNING) << "No CUDA device detected; treating buffers as " + "host memory"; + return false; + } + return true; +} +#endif + const std::vector getMemoryLocation(void *start, size_t len, bool only_first_page) { @@ -38,19 +56,24 @@ const std::vector getMemoryLocation(void *start, #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ defined(USE_COREX) || defined(USE_SUNRISE) - cudaPointerAttributes attributes; - cudaError_t result = cudaPointerGetAttributes(&attributes, start); - if (result != cudaSuccess) { - LOG(ERROR) << "cudaPointerGetAttributes failed (Error code: " << result - << " - " << cudaGetErrorString(result) << ")" << std::endl; - entries.push_back({(uint64_t)start, len, kWildcardLocation}); - return entries; - } + static const bool cuda_device_present = detectCudaDevicePresent(); + + if (cuda_device_present) { + cudaPointerAttributes attributes; + cudaError_t result = cudaPointerGetAttributes(&attributes, start); + if (result != cudaSuccess) { + LOG(ERROR) << "cudaPointerGetAttributes failed (Error code: " + << result << " - " << cudaGetErrorString(result) << ")" + << std::endl; + entries.push_back({(uint64_t)start, len, kWildcardLocation}); + return entries; + } - if (attributes.type == cudaMemoryTypeDevice) { - entries.push_back( - {(uint64_t)start, len, genGpuNodeName(attributes.device)}); - return entries; + if (attributes.type == cudaMemoryTypeDevice) { + entries.push_back( + {(uint64_t)start, len, genGpuNodeName(attributes.device)}); + return entries; + } } #endif diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 577712fd31..cfb13ca637 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -14,10 +14,12 @@ #include "multi_transport.h" #include +#include #include #include #include "config.h" +#include "multi_transport_locality.h" #include "transport/rdma_transport/rdma_transport.h" #ifdef USE_BAREX #include "transport/barex_transport/barex_transport.h" @@ -29,6 +31,9 @@ #ifdef USE_NVMEOF #include "transport/nvmeof_transport/nvmeof_transport.h" #endif +#ifdef USE_NCCL_HOST +#include "transport/nccl_transport/nccl_transport.h" +#endif #ifdef USE_ASCEND_DIRECT #include "transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.h" #endif @@ -59,6 +64,9 @@ #ifdef USE_EFA #include "transport/efa_transport/efa_transport.h" #endif +#ifdef USE_CXI +#include "transport/cxi_transport/cxi_transport.h" +#endif #ifdef USE_SUNRISE #include "transport/sunrise_link_transport/sunrise_link_transport.h" #endif @@ -115,16 +123,30 @@ Status MultiTransport::submitTransfer( "Exceed the limitation of batch capacity"); } + std::vector transports; + transports.reserve(entries.size()); + for (const auto& request : entries) { + Transport* transport = nullptr; + auto status = selectTransport(request, transport); + if (!status.ok()) return status; + assert(transport); + transports.push_back(transport); + } + size_t task_id = batch_desc.task_list.size(); batch_desc.task_list.resize(task_id + entries.size()); + struct TaskGroup { + uint64_t id; + Transport* transport; + std::vector tasks; + }; std::unordered_map > submit_tasks; - for (auto& request : entries) { - Transport* transport = nullptr; - auto status = selectTransport(request, transport); - if (!status.ok()) return status; - assert(transport); + std::vector task_groups; + for (size_t i = 0; i < entries.size(); ++i) { + const auto& request = entries[i]; + auto* transport = transports[i]; auto& task = batch_desc.task_list[task_id]; task.batch_id = batch_id; task.transport_ = transport; @@ -134,7 +156,15 @@ Status MultiTransport::submitTransfer( task.request = &request; #endif ++task_id; - submit_tasks[transport].push_back(&task); + if (request.task_group_id == TransferRequest::kNoTaskGroup) { + submit_tasks[transport].push_back(&task); + } else if (!task_groups.empty() && + task_groups.back().id == request.task_group_id && + task_groups.back().transport == transport) { + task_groups.back().tasks.push_back(&task); + } else { + task_groups.push_back({request.task_group_id, transport, {&task}}); + } } Status overall_status = Status::OK(); for (auto& entry : submit_tasks) { @@ -145,6 +175,10 @@ Status MultiTransport::submitTransfer( overall_status = status; } } + for (auto& group : task_groups) { + auto status = group.transport->submitTransferTaskGroup(group.tasks); + if (!status.ok()) overall_status = status; + } return overall_status; } @@ -230,7 +264,7 @@ Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id, task.transport_->getTransferStatus(batch_id, task_id, status); if (!ret.ok()) return ret; - // Apply timeout check on top of the transport's result. + // Apply timeout check on top of the transport result. if (status.s == Transport::TransferStatusEnum::WAITING && checkSliceTimeout(task)) { status.s = Transport::TransferStatusEnum::TIMEOUT; @@ -287,7 +321,8 @@ Status MultiTransport::getBatchTransferStatus(BatchID batch_id, if (task_status.s == Transport::TransferStatusEnum::COMPLETED) { status.transferred_bytes += task_status.transferred_bytes; success_count++; - } else if (task_status.s == Transport::TransferStatusEnum::FAILED) { + } else if (task_status.s == Transport::TransferStatusEnum::FAILED || + task_status.s == Transport::TransferStatusEnum::TIMEOUT) { status.s = Transport::TransferStatusEnum::FAILED; return Status::OK(); } @@ -308,6 +343,14 @@ Status MultiTransport::getBatchTransferStatus(BatchID batch_id, Transport* MultiTransport::installTransport(const std::string& proto, std::shared_ptr topo) { +#ifdef USE_NCCL_HOST + if ((proto == "nccl" && !transport_map_.empty()) || + (proto != "nccl" && transport_map_.count("nccl") != 0)) { + LOG(ERROR) << "NCCL host transport must be the only transport " + "installed in a Transfer Engine instance"; + return nullptr; + } +#endif Transport* transport = nullptr; if (std::string(proto) == "rdma") { transport = new RdmaTransport(); @@ -332,6 +375,11 @@ Transport* MultiTransport::installTransport(const std::string& proto, transport = new NVMeoFTransport(); } #endif +#ifdef USE_NCCL_HOST + else if (std::string(proto) == "nccl") { + transport = new NcclHostTransport(); + } +#endif #ifdef USE_ASCEND_DIRECT else if (std::string(proto) == "ascend") { transport = new AscendDirectTransport(); @@ -384,6 +432,11 @@ Transport* MultiTransport::installTransport(const std::string& proto, transport = new EfaTransport(); } #endif +#ifdef USE_CXI + else if (std::string(proto) == "cxi") { + transport = new CxiTransport(); + } +#endif #ifdef USE_SUNRISE else if (std::string(proto) == "sunrise_link") { transport = new SunriseLinkTransport(); @@ -447,6 +500,73 @@ Status MultiTransport::selectTransport(const TransferRequest& entry, std::to_string(entry.target_id)); } auto proto = target_segment_desc->protocol; +#ifdef ENABLE_MULTI_PROTOCOL + // Multi-protocol segment (e.g. "rdma,hip"): a single batch may target + // buffers owned by different transports (the device KV pool via hip, the + // host aux/metadata buffers via rdma). Route each request to the transport + // that owns the buffer covering the target address. When a buffer is + // registered under more than one protocol (the device KV pool is registered + // by both rdma and hip), pick the highest-performance transport by a fixed + // priority instead of relying on buffer registration order. + if (proto.find(',') != std::string::npos) { + auto protocol_priority = [](const std::string& p) { + // hip is intra-node GPU-IPC only. On a cross-node request a + // hip+rdma segment must fall through to rdma; allow deployments + // that know they need the cross-node path to de-prioritize hip. + if (p == "hip") return std::getenv("MC_DISABLE_HIP") ? 0 : 4; + if (p == "maca") return std::getenv("MC_DISABLE_MACA") ? 0 : 4; + if (p == "cxl") return 3; + if (p == "rdma") return 2; + if (p == "tcp") return 1; + return 0; + }; + // hip transport uses GPU IPC, which cannot reach a GPU on another host. + // The device KV pool is registered under both rdma and hip, so a + // cross-host target must skip its hip buffers and fall back to rdma. + // This makes the intra-node fast path (hip) and the cross-node path + // (rdma) work automatically from a single multi-protocol segment, + // without requiring the operator to set MC_DISABLE_HIP. + const bool hip_reachable = + isHipReachableTarget(target_segment_desc->name, local_server_name_); + std::string chosen; + int chosen_priority = -1; + for (const auto& buffer : target_segment_desc->buffers) { + // CXL buffers locate via offset + cxl_base_addr; all other + // protocols use the absolute virtual address in buffer.addr. + uint64_t start = + (buffer.protocol == "cxl") + ? buffer.offset + target_segment_desc->cxl_base_addr + : buffer.addr; + if (entry.target_offset >= start && + entry.target_offset < start + buffer.length) { + if (buffer.protocol == "hip" && !hip_reachable) continue; + int priority = protocol_priority(buffer.protocol); + if (priority > chosen_priority) { + chosen = buffer.protocol; + chosen_priority = priority; + } + } + } + if (chosen.empty()) { + return Status::InvalidArgument( + "No matching buffer for target offset in multi-protocol " + "segment " + + std::to_string(entry.target_id)); + } + if (!transport_map_.count(chosen)) { + return Status::NotSupportedTransport("Transport " + chosen + + " not installed"); + } + if (globalConfig().trace) { + LOG(INFO) << "MultiTransport::selectTransport route: target_id=" + << entry.target_id << " segment_protocol=\"" << proto + << "\" hip_reachable=" << hip_reachable + << " chosen=" << chosen; + } + transport = transport_map_[chosen].get(); + return Status::OK(); + } +#endif #ifdef USE_ASCEND_HETEROGENEOUS // When USE_ASCEND_HETEROGENEOUS is enabled: // - Target side directly reuses RDMA Transport @@ -480,6 +600,28 @@ Status MultiTransport::mp_selectTransport(const TransferRequest& entry, if (!item.empty()) protos.push_back(item); } + // hip GPU IPC cannot reach a remote host; downgrade an explicit hip + // preference to a cross-host-capable transport for a cross-host target + // (mirrors the locality gate in selectTransport). Prefer rdma, then tcp. + if (preferred_proto == "hip" && + !isHipReachableTarget(target_segment_desc->name, local_server_name_)) { + std::string fallback; + for (const char* candidate : {"rdma", "tcp"}) { + if (std::find(protos.begin(), protos.end(), candidate) != + protos.end()) { + fallback = candidate; + break; + } + } + if (fallback.empty()) { + return Status::NotSupportedTransport( + "hip target is cross-host but segment " + + std::to_string(entry.target_id) + + " offers no cross-host transport (rdma/tcp)"); + } + preferred_proto = fallback; + } + #ifdef USE_ASCEND_HETEROGENEOUS // When USE_ASCEND_HETEROGENEOUS is enabled: // - Target side directly reuses RDMA Transport diff --git a/mooncake-transfer-engine/src/show_links.cpp b/mooncake-transfer-engine/src/show_links.cpp new file mode 100644 index 0000000000..b50b8bfccd --- /dev/null +++ b/mooncake-transfer-engine/src/show_links.cpp @@ -0,0 +1,197 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "show_links.h" + +#include + +#include +#include + +#include "transfer_engine_impl.h" +#include "transport/rdma_transport/rdma_context.h" +#include "transport/rdma_transport/rdma_transport.h" + +namespace mooncake { + +static std::vector collectLocalNics(TransferEngineImpl* impl) { + std::vector nics; + auto* transport = impl->getTransport("rdma"); + if (!transport) { + auto topo = impl->getLocalTopology(); + if (!topo) return nics; + + for (const auto& hca : topo->getHcaList()) { + NicDiagInfo info; + info.device_name = hca; + info.numa_node = -1; + info.gid_index = -1; + info.port = 0; + info.active_speed = 0; + info.active_width = 0; + nics.push_back(info); + } + return nics; + } + + auto* rdma = static_cast(transport); + for (auto& ctx : rdma->getContextList()) { + NicDiagInfo info; + info.device_name = ctx->deviceName(); + info.numa_node = ctx->socketId(); + info.gid = ctx->gid(); + info.gid_index = ctx->gidIndex(); + info.port = ctx->portNum(); + info.active_speed = ctx->activeSpeed(); + info.active_width = ctx->activeWidth(); + nics.push_back(info); + } + return nics; +} + +static bool hasTransportDetails(const NicDiagInfo& nic) { + return nic.gid_index >= 0; +} + +static int computeSpeedGbps(int active_speed, int active_width) { + if (active_width <= 0) return 0; + + int lane_gbps = 0; + switch (active_speed) { + case 1: + lane_gbps = 2; + break; + case 2: + lane_gbps = 5; + break; + case 4: + lane_gbps = 10; + break; + case 8: + lane_gbps = 10; + break; + case 16: + lane_gbps = 14; + break; + case 32: + lane_gbps = 25; + break; + case 64: + lane_gbps = 50; + break; + case 128: + lane_gbps = 100; + break; + case 256: + lane_gbps = 200; + break; + default: + lane_gbps = std::max(0, active_speed); + break; + } + int lanes = 1; + switch (active_width) { + case 1: + lanes = 1; + break; + case 2: + lanes = 4; + break; + case 4: + lanes = 8; + break; + case 8: + lanes = 12; + break; + case 16: + lanes = 2; + break; + default: + lanes = 1; + break; + } + return lane_gbps * lanes; +} + +std::string buildShowLinksJson(TransferEngineImpl* impl) { + Json::Value root; + + auto nics = collectLocalNics(impl); + Json::Value nic_arr(Json::arrayValue); + for (auto& nic : nics) { + Json::Value n; + n["device"] = nic.device_name; + n["numa"] = nic.numa_node; + n["gid"] = nic.gid; + n["gid_index"] = nic.gid_index; + n["port"] = nic.port; + n["speed_gbps"] = computeSpeedGbps(nic.active_speed, nic.active_width); + n["source"] = hasTransportDetails(nic) ? "rdma_transport" : "topology"; + nic_arr.append(n); + } + root["local_nics"] = nic_arr; + + auto topo = impl->getLocalTopology(); + if (topo) { + root["topology"] = topo->toJson(); + } + + Json::StreamWriterBuilder builder; + builder["indentation"] = " "; + return Json::writeString(builder, root); +} + +std::string buildShowLinksReadable(TransferEngineImpl* impl) { + std::ostringstream os; + + auto nics = collectLocalNics(impl); + os << "=== Local NICs ===\n"; + if (nics.empty()) { + os << " (no RDMA devices found)\n"; + } + for (auto& nic : nics) { + os << " " << nic.device_name; + if (!hasTransportDetails(nic)) { + os << " (topology only; RDMA transport not initialized)\n"; + continue; + } + os << " NUMA=" << nic.numa_node << " GID=" << nic.gid + << " (idx=" << nic.gid_index << ")" + << " Port=" << (int)nic.port << " " + << computeSpeedGbps(nic.active_speed, nic.active_width) << "Gbps\n"; + } + + auto topo = impl->getLocalTopology(); + if (topo && !topo->empty()) { + os << "\n=== Topology (NIC Selection) ===\n"; + auto matrix = topo->getMatrix(); + for (auto& [location, entry] : matrix) { + os << " " << location << " -> preferred: ["; + for (size_t i = 0; i < entry.preferred_hca.size(); i++) { + if (i > 0) os << ", "; + os << entry.preferred_hca[i]; + } + os << "] available: ["; + for (size_t i = 0; i < entry.avail_hca.size(); i++) { + if (i > 0) os << ", "; + os << entry.avail_hca[i]; + } + os << "]\n"; + } + } + + return os.str(); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/topology.cpp b/mooncake-transfer-engine/src/topology.cpp index bc61a722f3..5383f1e3d7 100644 --- a/mooncake-transfer-engine/src/topology.cpp +++ b/mooncake-transfer-engine/src/topology.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -31,6 +32,7 @@ #include #include "cuda_alike.h" +#include "config.h" #include "memory_location.h" #include "topology.h" #include "char_util.h" @@ -189,6 +191,7 @@ static std::set getIbvDeviceWhitelist() { static std::vector listInfiniBandDevices( const std::vector &filter) { +#ifndef USE_CXI int num_devices = 0; std::vector devices; const auto whitelist = getIbvDeviceWhitelist(); @@ -246,6 +249,101 @@ static std::vector listInfiniBandDevices( } ibv_free_device_list(device_list); return devices; +#else + // assume devices available and working, later use libcxi to check if they + // actually work + std::vector device_names; + std::vector devices; + const char *nic_path = "/sys/class/cxi"; + + DIR *dir = opendir(nic_path); + if (!dir) { + perror("opendir"); + LOG(WARNING) << "No CXI devices were found, check you system"; + return {}; + } + + struct dirent *entry; + while ((entry = readdir(dir)) != nullptr) { + if (entry->d_type == DT_LNK && std::string(entry->d_name) != "." && + std::string(entry->d_name) != "..") { + device_names.emplace_back(entry->d_name); + } + } + closedir(dir); + + for (auto &dv_name : device_names) { + if (!filter.empty() && + std::find(filter.begin(), filter.end(), dv_name) == filter.end()) + continue; + + char path[PATH_MAX + 32]; + char resolved_path[PATH_MAX]; + + snprintf(path, sizeof(path), "/sys/class/cxi/%s/device", + dv_name.c_str()); + if (realpath(path, resolved_path) == NULL) { + PLOG(ERROR) << "Can't resolve CXI device path for " << path; + continue; + } + std::string pci_bus_id = basename(resolved_path); + + int numa_node = -1; + + snprintf(path, sizeof(path), "%s/numa_node", resolved_path); + std::ifstream(path) >> numa_node; + + devices.push_back(InfinibandDevice{.name = std::move(dv_name), + .pci_bus_id = std::move(pci_bus_id), + .numa_node = numa_node}); + } + return devices; +#endif +} + +static void precomputeResolvedHcaPeerAffinity( + const TopologyMatrix &matrix, const std::map &hca_id_map, + std::unordered_map>> + &affinity_by_local) { + affinity_by_local.clear(); + const auto &config = globalConfig(); + if (!config.enable_hca_peer_affinity) return; + if (config.nic_peer_affinity.empty()) { + LOG(WARNING) << "MC_ENABLE_HCA_PEER_AFFINITY is set, but " + "MC_NIC_PEER_AFFINITY has no valid mappings."; + return; + } + + for (const auto &entry : matrix) { + // Match the vendor-specific GPU location prefix (cuda:/hip:/...). + // Hardcoding "cuda:" silently disables HCA peer affinity under USE_HIP. + if (entry.first.rfind(GPU_PREFIX, 0) != 0) continue; + + std::unordered_set preferred_set( + entry.second.preferred_hca.begin(), + entry.second.preferred_hca.end()); + + for (const auto &mapping : config.nic_peer_affinity) { + const auto &local_hca = mapping.first; + const auto &peer_hcas = mapping.second; + + std::vector peer_preferred_hca; + std::unordered_set peer_preferred_set; + for (const auto &peer_hca : peer_hcas) { + if (!preferred_set.count(peer_hca)) continue; + auto hca_id_it = hca_id_map.find(peer_hca); + if (hca_id_it == hca_id_map.end()) continue; + if (peer_preferred_set.insert(peer_hca).second) { + peer_preferred_hca.push_back(hca_id_it->second); + } + } + + if (peer_preferred_hca.empty()) continue; + affinity_by_local[local_hca][entry.first] = + std::move(peer_preferred_hca); + } + } } #ifdef USE_UB @@ -521,10 +619,12 @@ void Topology::clear() { matrix_.clear(); hca_list_.clear(); resolved_matrix_.clear(); + resolved_hca_peer_affinity_by_local_.clear(); } int Topology::discover(const std::vector &filter) { matrix_.clear(); + resolved_hca_peer_affinity_by_local_.clear(); auto all_hca = listInfiniBandDevices(filter); for (auto &ent : discoverCpuTopology(all_hca)) { matrix_[ent.name] = ent; @@ -566,6 +666,7 @@ int Topology::parse(const std::string &topology_json) { } matrix_.clear(); + resolved_hca_peer_affinity_by_local_.clear(); for (const auto &key : root.getMemberNames()) { const Json::Value &value = root[key]; if (value.isArray() && value.size() == 2) { @@ -596,6 +697,10 @@ std::string Topology::toString() const { return value.toStyledString(); } +bool Topology::operator==(const Topology &other) const { + return matrix_ == other.matrix_ && hca_list_ == other.hca_list_; +} + Json::Value Topology::toJson() const { Json::Value root; for (const auto &pair : matrix_) { @@ -623,10 +728,44 @@ int Topology::selectDevice(const std::string storage_type, return selectDevice(storage_type, retry_count); } +int Topology::selectDeviceByLocalHca(const std::string storage_type, + std::string_view local_hca, + int retry_count) { + if (!local_hca.empty()) { + auto local_it = + resolved_hca_peer_affinity_by_local_.find(std::string(local_hca)); + if (local_it != resolved_hca_peer_affinity_by_local_.end()) { + auto hints_it = local_it->second.find(std::string(storage_type)); + if (hints_it != local_it->second.end() && + !hints_it->second.empty()) { + const auto &candidates = hints_it->second; + if (retry_count == 0) { + int rand_value; + if (use_round_robin_) { + thread_local int tl_counter = 0; + rand_value = tl_counter; + tl_counter = (tl_counter + 1) % 10000; + } else { + rand_value = SimpleRandom::Get().next(); + } + return candidates[rand_value % candidates.size()]; + } + return candidates[(retry_count - 1) % candidates.size()]; + } + } + } + + return selectDevice(storage_type, retry_count); +} + int Topology::selectDevice(const std::string storage_type, int retry_count) { if (resolved_matrix_.count(storage_type) == 0) return ERR_DEVICE_NOT_FOUND; auto &entry = resolved_matrix_[storage_type]; + if (entry.preferred_hca.empty() && entry.avail_hca.empty()) { + return ERR_DEVICE_NOT_FOUND; + } + if (retry_count == 0) { int rand_value; if (use_round_robin_) { @@ -649,11 +788,12 @@ int Topology::selectDevice(const std::string storage_type, int retry_count) { return entry.avail_hca[index]; } } - return 0; + return ERR_DEVICE_NOT_FOUND; } int Topology::resolve() { resolved_matrix_.clear(); + resolved_hca_peer_affinity_by_local_.clear(); hca_list_.clear(); std::map hca_id_map; int next_hca_map_index = 0; @@ -688,10 +828,19 @@ int Topology::resolve() { hca_id_map[hca]; } } + precomputeResolvedHcaPeerAffinity(matrix_, hca_id_map, + resolved_hca_peer_affinity_by_local_); return 0; } int Topology::disableDevice(const std::string &device_name) { + int disabled_hca_index = -1; + auto hca_iter = std::find(hca_list_.begin(), hca_list_.end(), device_name); + if (hca_iter != hca_list_.end()) { + disabled_hca_index = + static_cast(std::distance(hca_list_.begin(), hca_iter)); + } + for (auto &record : matrix_) { auto &preferred_hca = record.second.preferred_hca; auto preferred_hca_iter = @@ -703,6 +852,35 @@ int Topology::disableDevice(const std::string &device_name) { std::find(avail_hca.begin(), avail_hca.end(), device_name); if (avail_hca_iter != avail_hca.end()) avail_hca.erase(avail_hca_iter); } - return resolve(); + + if (disabled_hca_index < 0) return 0; + + // Keep existing HCA indexes stable. RDMA transport stores lkey/rkey and + // context arrays by the resolved HCA index, so re-running resolve() here + // could compact indexes after a disabled device and make those arrays point + // at the wrong RNIC. + for (auto &record : resolved_matrix_) { + auto &preferred_hca = record.second.preferred_hca; + preferred_hca.erase( + std::remove(preferred_hca.begin(), preferred_hca.end(), + disabled_hca_index), + preferred_hca.end()); + record.second.preferred_hca_name_to_index_map_.erase(device_name); + + auto &avail_hca = record.second.avail_hca; + avail_hca.erase( + std::remove(avail_hca.begin(), avail_hca.end(), disabled_hca_index), + avail_hca.end()); + record.second.avail_hca_name_to_index_map_.erase(device_name); + } + for (auto &local_entry : resolved_hca_peer_affinity_by_local_) { + for (auto &storage_entry : local_entry.second) { + auto &candidates = storage_entry.second; + candidates.erase(std::remove(candidates.begin(), candidates.end(), + disabled_hca_index), + candidates.end()); + } + } + return 0; } } // namespace mooncake diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index 3747244e76..1f6ef80007 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -12,12 +12,61 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include +#include +#include +#include + #ifndef USE_TENT #include "transfer_engine.h" +#include "show_links.h" #include "transfer_engine_impl.h" +#include "graceful_shutdown.h" +#include #include namespace mooncake { +namespace { + +class TransferEngineShutdownToken : public ShutdownToken { + public: + explicit TransferEngineShutdownToken(TransferEngine* engine) + : engine_(engine) {} + + void shutdown() override { + TransferEngine* engine = nullptr; + { + std::lock_guard lock(mutex_); + engine = engine_; + engine_ = nullptr; + } + if (engine) engine->freeEngine(); + } + + void detach() override { + std::lock_guard lock(mutex_); + engine_ = nullptr; + } + + private: + std::mutex mutex_; + TransferEngine* engine_; +}; + +std::shared_ptr registerTransferEngineShutdownToken( + TransferEngine* engine) { + auto token = std::make_shared(engine); + registerTokenForShutdown(token); + return token; +} + +void detachShutdownToken(std::shared_ptr& token) { + if (!token) return; + token->detach(); + token.reset(); +} + +} // namespace TransferEngine::TransferEngine(bool auto_discover) : impl_(std::make_shared(auto_discover)) {} @@ -26,6 +75,33 @@ TransferEngine::TransferEngine(bool auto_discover, const std::vector& filter) : impl_(std::make_shared(auto_discover, filter)) {} +TransferEngine::TransferEngine(TransferEngine&& other) noexcept + : impl_(std::move(other.impl_)), + impl_tent_(std::move(other.impl_tent_)), + use_tent_(other.use_tent_) { + const bool shutdown_enabled = static_cast(other.shutdown_token_); + detachShutdownToken(other.shutdown_token_); + if (shutdown_enabled) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + installGracefulShutdownHandlers(); + } +} + +TransferEngine& TransferEngine::operator=(TransferEngine&& other) noexcept { + if (this == &other) return *this; + freeEngine(); + impl_ = std::move(other.impl_); + impl_tent_ = std::move(other.impl_tent_); + use_tent_ = other.use_tent_; + const bool shutdown_enabled = static_cast(other.shutdown_token_); + detachShutdownToken(other.shutdown_token_); + if (shutdown_enabled) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + installGracefulShutdownHandlers(); + } + return *this; +} + TransferEngine::~TransferEngine() { freeEngine(); } int TransferEngine::init(const std::string& metadata_conn_string, @@ -37,8 +113,9 @@ int TransferEngine::init(const std::string& metadata_conn_string, } int TransferEngine::freeEngine() { + detachShutdownToken(shutdown_token_); if (impl_) { - impl_->freeEngine(); + if (impl_.use_count() == 1) impl_->freeEngine(); impl_.reset(); } return 0; @@ -175,11 +252,17 @@ Status TransferEngine::getBatchTransferStatus(BatchID batch_id, return impl_->getBatchTransferStatus(batch_id, status); } +Status TransferEngine::getNicLoadStats(std::vector& stats) const { + stats.clear(); + return Status::OK(); +} + Transport* TransferEngine::getTransport(const std::string& proto) { return impl_->getTransport(proto); } -#if defined(USE_CUDA) || defined(USE_MUSA) +#if (defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)) && \ + !defined(USE_CXI) device::P2pTransport* TransferEngine::getOrCreateP2pTransport(int num_ranks) { return impl_->getOrCreateP2pTransport(num_ranks); } @@ -190,6 +273,12 @@ device::RdmaTransport* TransferEngine::getOrCreateRdmaTransport( } #endif +#ifdef USE_NCCL_DEVICE +device::NcclTransport* TransferEngine::getOrCreateNcclTransport() { + return impl_->getOrCreateNcclTransport(); +} +#endif + bool TransferEngine::isTcpOnly() const { return impl_->isTcpOnly(); } int TransferEngine::syncSegmentCache(const std::string& segment_name) { @@ -208,6 +297,10 @@ void TransferEngine::setAutoDiscover(bool auto_discover) { impl_->setAutoDiscover(auto_discover); } +void TransferEngine::setAutoDiscover(const AutoDiscoverConfig& config) { + impl_->setAutoDiscover(config); +} + void* TransferEngine::getBaseAddr() { return impl_->getBaseAddr(); } void TransferEngine::setWhitelistFilters(std::vector&& filters) { @@ -220,6 +313,19 @@ std::shared_ptr TransferEngine::getLocalTopology() { return impl_->getLocalTopology(); } +void TransferEngine::enableGracefulShutdown() { + if (!shutdown_token_) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + } + installGracefulShutdownHandlers(); +} + +std::string TransferEngine::showLinks(bool json) const { + if (!impl_) return "{}"; + return json ? buildShowLinksJson(impl_.get()) + : buildShowLinksReadable(impl_.get()); +} + } // namespace mooncake #else #include "transfer_engine.h" @@ -227,9 +333,53 @@ std::shared_ptr TransferEngine::getLocalTopology() { #include "tent/transfer_engine.h" #include "tent/common/config.h" +#include #include +#include "graceful_shutdown.h" +#include "show_links.h" namespace mooncake { +namespace { + +class TransferEngineShutdownToken : public ShutdownToken { + public: + explicit TransferEngineShutdownToken(TransferEngine* engine) + : engine_(engine) {} + + void shutdown() override { + TransferEngine* engine = nullptr; + { + std::lock_guard lock(mutex_); + engine = engine_; + engine_ = nullptr; + } + if (engine) engine->freeEngine(); + } + + void detach() override { + std::lock_guard lock(mutex_); + engine_ = nullptr; + } + + private: + std::mutex mutex_; + TransferEngine* engine_; +}; + +std::shared_ptr registerTransferEngineShutdownToken( + TransferEngine* engine) { + auto token = std::make_shared(engine); + registerTokenForShutdown(token); + return token; +} + +void detachShutdownToken(std::shared_ptr& token) { + if (!token) return; + token->detach(); + token.reset(); +} + +} // namespace TransferEngine::TransferEngine(bool auto_discover) { if (getenv("MC_USE_TENT") || getenv("MC_USE_TEV1")) { @@ -250,6 +400,34 @@ TransferEngine::TransferEngine(bool auto_discover, } } +TransferEngine::TransferEngine(TransferEngine&& other) noexcept + : impl_(std::move(other.impl_)), + impl_tent_(std::move(other.impl_tent_)), + shutdown_token_(nullptr), + use_tent_(other.use_tent_) { + const bool shutdown_enabled = static_cast(other.shutdown_token_); + detachShutdownToken(other.shutdown_token_); + if (shutdown_enabled) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + installGracefulShutdownHandlers(); + } +} + +TransferEngine& TransferEngine::operator=(TransferEngine&& other) noexcept { + if (this == &other) return *this; + freeEngine(); + impl_ = std::move(other.impl_); + impl_tent_ = std::move(other.impl_tent_); + use_tent_ = other.use_tent_; + const bool shutdown_enabled = static_cast(other.shutdown_token_); + detachShutdownToken(other.shutdown_token_); + if (shutdown_enabled) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + installGracefulShutdownHandlers(); + } + return *this; +} + TransferEngine::~TransferEngine() { freeEngine(); } static std::pair parseConnectionStringInternal( @@ -299,8 +477,9 @@ int TransferEngine::init(const std::string& metadata_conn_string, } int TransferEngine::freeEngine() { + detachShutdownToken(shutdown_token_); if (!use_tent_ && impl_) { - impl_->freeEngine(); + if (impl_.use_count() == 1) impl_->freeEngine(); impl_.reset(); } else { impl_tent_.reset(); @@ -348,7 +527,8 @@ SegmentHandle TransferEngine::openSegment(const std::string& segment_name) { if (use_tent_) { SegmentHandle handle; auto status = impl_tent_->openSegment(handle, segment_name); - if (!status.ok()) return (SegmentHandle)(-1); + if (!status.ok()) + return static_cast(ERR_INVALID_ARGUMENT); return handle; } else return impl_->openSegment(segment_name); @@ -431,7 +611,8 @@ int TransferEngine::unregisterLocalMemoryBatch( BatchID TransferEngine::allocateBatchID(size_t batch_size) { if (use_tent_) { - return impl_tent_->allocateBatch(batch_size); + const auto batch_id = impl_tent_->allocateBatch(batch_size); + return batch_id == 0 ? INVALID_BATCH_ID : batch_id; } else { return impl_->allocateBatchID(batch_size); } @@ -588,6 +769,24 @@ Status TransferEngine::getBatchTransferStatus(BatchID batch_id, return impl_->getBatchTransferStatus(batch_id, status); } +Status TransferEngine::getNicLoadStats(std::vector& stats) const { + stats.clear(); + if (use_tent_) { + std::vector tent_stats; + auto status = impl_tent_->getNicLoadStats(tent_stats); + if (!status.ok()) return Status::Context(status.ToString()); + stats.reserve(tent_stats.size()); + for (const auto& stat : tent_stats) { + NicLoadStats load_stats; + load_stats.device_name = stat.device_name; + load_stats.inflight_bytes = stat.inflight_bytes; + load_stats.ewma_bandwidth_bps = stat.ewma_bandwidth_bps; + stats.push_back(load_stats); + } + } + return Status::OK(); +} + Transport* TransferEngine::getTransport(const std::string& proto) { if (use_tent_) return nullptr; @@ -595,7 +794,8 @@ Transport* TransferEngine::getTransport(const std::string& proto) { return impl_->getTransport(proto); } -#if defined(USE_CUDA) || defined(USE_MUSA) +#if (defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)) && \ + !defined(USE_CXI) device::P2pTransport* TransferEngine::getOrCreateP2pTransport(int num_ranks) { if (use_tent_) return nullptr; return impl_->getOrCreateP2pTransport(num_ranks); @@ -608,6 +808,13 @@ device::RdmaTransport* TransferEngine::getOrCreateRdmaTransport( } #endif +#ifdef USE_NCCL_DEVICE +device::NcclTransport* TransferEngine::getOrCreateNcclTransport() { + if (use_tent_) return nullptr; + return impl_->getOrCreateNcclTransport(); +} +#endif + bool TransferEngine::isTcpOnly() const { if (use_tent_) // TENT already rejects TCP loopback transfers when MC_STORE_MEMCPY @@ -641,6 +848,10 @@ void TransferEngine::setAutoDiscover(bool auto_discover) { if (!use_tent_) impl_->setAutoDiscover(auto_discover); } +void TransferEngine::setAutoDiscover(const AutoDiscoverConfig& config) { + if (!use_tent_) impl_->setAutoDiscover(config); +} + void TransferEngine::setWhitelistFilters(std::vector&& filters) { if (!use_tent_) impl_->setWhitelistFilters(std::move(filters)); } @@ -668,5 +879,400 @@ void* TransferEngine::getBaseAddr() { return impl_->getBaseAddr(); } +void TransferEngine::enableGracefulShutdown() { + if (!shutdown_token_) { + shutdown_token_ = registerTransferEngineShutdownToken(this); + } + installGracefulShutdownHandlers(); +} + +std::string TransferEngine::showLinks(bool json) const { + if (use_tent_ || !impl_) { + return json ? "{}" : "(TENT mode or not initialized)"; + } + return json ? buildShowLinksJson(impl_.get()) + : buildShowLinksReadable(impl_.get()); +} + } // namespace mooncake #endif + +namespace mooncake { + +class TransferEngine::ScatterTransferOperation::Impl { + public: + struct Backend { + std::shared_ptr legacy; +#ifdef USE_TENT + std::shared_ptr tent; +#endif + }; + + Impl(TransferEngine& engine, Backend backend, + const std::vector& ranges) + : backend_(std::move(backend)) { + callbacks_.reserve(ranges.size()); + for (const auto& range : ranges) + callbacks_.push_back(range.on_fragment_complete); + build(engine, ranges); + } + + ~Impl() { wait(); } + + Status wait() { + while (!completed_) { + poll(); + if (!completed_) std::this_thread::sleep_for(kPollInterval); + } + return aggregate_status_; + } + + Status waitFor(std::chrono::nanoseconds timeout) { + const auto now = std::chrono::steady_clock::now(); + const auto until_max = + std::chrono::steady_clock::time_point::max() - now; + const auto max_timeout = + std::chrono::duration_cast(until_max); + auto deadline = now; + if (timeout > std::chrono::nanoseconds::zero()) { + deadline = timeout >= max_timeout + ? std::chrono::steady_clock::time_point::max() + : now + timeout; + } + while (!completed_) { + poll(); + if (completed_) break; + if (std::chrono::steady_clock::now() >= deadline) + return Status::Clock("scatter transfer wait timed out"); + std::this_thread::sleep_for(kPollInterval); + } + return aggregate_status_; + } + + private: + static constexpr auto kPollInterval = std::chrono::microseconds(10); + + bool useTent() const { +#ifdef USE_TENT + return static_cast(backend_.tent); +#else + return false; +#endif + } + + int closeSegment(SegmentHandle handle) { +#ifdef USE_TENT + if (backend_.tent) + return static_cast(backend_.tent->closeSegment(handle).code()); +#endif + return backend_.legacy->closeSegment(handle); + } + + Status getStatus(BatchID batch_id, size_t task_id, TransferStatus& status) { +#ifdef USE_TENT + if (backend_.tent) { + mooncake::tent::TransferStatus tent_status; + auto result = backend_.tent->getTransferStatus(batch_id, task_id, + tent_status); + if (!result.ok()) return Status::Context(result.ToString()); + status.s = static_cast(tent_status.s); + status.transferred_bytes = tent_status.transferred_bytes; + return Status::OK(); + } +#endif + return backend_.legacy->getTransferStatus(batch_id, task_id, status); + } + + Status freeBatch(BatchID batch_id) { +#ifdef USE_TENT + if (backend_.tent) { + auto result = backend_.tent->freeBatch(batch_id); + return result.ok() ? Status::OK() + : Status::Context(result.ToString()); + } +#endif + return backend_.legacy->freeBatchID(batch_id); + } + + void remember(const Status& status) { + if (aggregate_status_.ok() && !status.ok()) aggregate_status_ = status; + } + + Status closeSegments(Status status) { + for (const auto& entry : segment_handles_) { + if (entry.second == + static_cast(ERR_INVALID_ARGUMENT)) + continue; + if (closeSegment(entry.second) != 0 && status.ok()) + status = + Status::Context("failed to close scatter transfer segment"); + } + segment_handles_.clear(); + return status; + } + + void finish() { + aggregate_status_ = closeSegments(aggregate_status_); + completed_ = true; + callbacks_.clear(); + requests_.clear(); + request_fragments_.clear(); + done_.clear(); + backend_ = {}; + } + + void complete(size_t range_index, size_t fragment_index, + const Status& status) { + remember(status); + const auto& callback = callbacks_[range_index]; + if (!callback) return; + try { + callback(fragment_index, status); + } catch (...) { + LOG(ERROR) << "scatter transfer callback failed"; + remember(Status::Context("scatter transfer callback failed")); + } + } + + void failPending(const Status& status) { + for (size_t i = 0; i < request_fragments_.size(); ++i) { + if (done_[i]) continue; + const auto [range_index, fragment_index] = request_fragments_[i]; + complete(range_index, fragment_index, status); + done_[i] = true; + --remaining_; + } + } + + void requestAbort(const Status& status) { + remember(status); + if (abort_requested_) return; + abort_requested_ = true; +#ifdef USE_TENT + if (!backend_.tent) return; + for (size_t i = 0; i < requests_.size(); ++i) { + if (done_[i]) continue; + auto cancel_status = backend_.tent->cancelTransfer(batch_id_, i); + if (!cancel_status.ok() && !cancel_status.IsNotImplemented()) { + LOG(WARNING) << "failed to cancel scatter transfer task " << i + << ": " << cancel_status.ToString(); + } + } +#endif + } + + void build(TransferEngine& engine, + const std::vector& ranges) { + for (size_t range_index = 0; range_index < ranges.size(); + ++range_index) { + const auto& range = ranges[range_index]; + const size_t fragment_count = range.local_offsets.size(); + if (range.remote_offsets.size() != fragment_count || + range.lengths.size() != fragment_count || + range.local_buffer == nullptr || range.remote_segment.empty()) { + const auto status = + Status::InvalidArgument("invalid scatter transfer range"); + remember(status); + for (size_t i = 0; i < fragment_count; ++i) + complete(range_index, i, status); + continue; + } + + for (size_t fragment_index = 0; fragment_index < fragment_count; + ++fragment_index) { + const size_t length = range.lengths[fragment_index]; + const size_t local_offset = range.local_offsets[fragment_index]; + const size_t remote_offset = + range.remote_offsets[fragment_index]; + if (local_offset > range.local_capacity || + length > range.local_capacity - local_offset || + remote_offset > range.remote_size || + length > range.remote_size - remote_offset || + range.remote_base_offset > + std::numeric_limits::max() - remote_offset || + length > std::numeric_limits::max() - + (range.remote_base_offset + remote_offset)) { + complete(range_index, fragment_index, + Status::InvalidArgument( + "invalid scatter transfer fragment")); + continue; + } + + if (length == 0) { + complete(range_index, fragment_index, Status::OK()); + continue; + } + + auto [segment, inserted] = segment_handles_.emplace( + range.remote_segment, + static_cast(ERR_INVALID_ARGUMENT)); + if (inserted) + segment->second = engine.openSegment(range.remote_segment); + if (segment->second == + static_cast(ERR_INVALID_ARGUMENT)) { + complete(range_index, fragment_index, + Status::InvalidArgument( + "failed to open scatter transfer segment")); + continue; + } + + requests_.push_back(TransferRequest{ + .opcode = range.opcode, + .source = + static_cast(range.local_buffer) + local_offset, + .target_id = segment->second, + .target_offset = range.remote_base_offset + remote_offset, + .length = length, + .task_group_id = range_index + 1, + }); + request_fragments_.emplace_back(range_index, fragment_index); + } + } + + if (requests_.empty()) { + finish(); + return; + } + + done_.assign(requests_.size(), false); + remaining_ = requests_.size(); + batch_id_ = engine.allocateBatchID(requests_.size()); + if (batch_id_ == INVALID_BATCH_ID) { + failPending(Status::InvalidArgument( + "failed to allocate scatter transfer batch")); + finish(); + return; + } + + auto submit_status = engine.submitTransfer(batch_id_, requests_); + if (submit_status.ok()) return; + +#ifdef USE_TENT + if (backend_.tent) { + // TENT publishes all task slots together. Drain them if a rare + // post-publication error escapes submitTransfer(). + mooncake::tent::TransferStatus status; + auto probe = backend_.tent->getTransferStatus(batch_id_, 0, status); + if (probe.ok() || !probe.IsInvalidArgument()) { + requestAbort(submit_status); + return; + } + remember(submit_status); + remember(freeBatch(batch_id_)); + batch_id_ = INVALID_BATCH_ID; + failPending(submit_status); + finish(); + return; + } +#endif + + requestAbort(submit_status); + auto free_status = freeBatch(batch_id_); + if (free_status.ok()) { + batch_id_ = INVALID_BATCH_ID; + failPending(submit_status); + finish(); + } else if (!free_status.IsBatchBusy()) { + remember(free_status); + } + } + + void poll() { + for (size_t i = 0; i < requests_.size(); ++i) { + if (done_[i]) continue; + TransferStatus status; + auto result = getStatus(batch_id_, i, status); + if (!result.ok()) { + requestAbort(result); + continue; + } + + const auto [range_index, fragment_index] = request_fragments_[i]; + Status fragment_status; + if (status.s == TransferStatusEnum::COMPLETED) { + fragment_status = Status::OK(); + } else if (status.s == TransferStatusEnum::WAITING || + status.s == TransferStatusEnum::PENDING) { + continue; + } else if (status.s == TransferStatusEnum::TIMEOUT) { + fragment_status = + Status::Socket("scatter transfer fragment timed out"); + requestAbort(fragment_status); + if (!useTent()) continue; + } else { + fragment_status = + Status::Socket("scatter transfer fragment failed"); + } + if (!fragment_status.ok()) requestAbort(fragment_status); + complete(range_index, fragment_index, fragment_status); + done_[i] = true; + --remaining_; + } + + if (remaining_ != 0) return; + auto free_status = freeBatch(batch_id_); + if (free_status.IsBatchBusy()) return; + remember(free_status); + batch_id_ = INVALID_BATCH_ID; + finish(); + } + + Backend backend_; + std::vector requests_; + std::vector> request_fragments_; + std::vector> callbacks_; + std::unordered_map segment_handles_; + std::vector done_; + BatchID batch_id_ = INVALID_BATCH_ID; + size_t remaining_ = 0; + Status aggregate_status_; + bool abort_requested_ = false; + bool completed_ = false; +}; + +TransferEngine::ScatterTransferOperation::ScatterTransferOperation( + std::unique_ptr impl) + : impl_(std::move(impl)) {} + +TransferEngine::ScatterTransferOperation::ScatterTransferOperation( + ScatterTransferOperation&&) noexcept = default; + +TransferEngine::ScatterTransferOperation& +TransferEngine::ScatterTransferOperation::operator=( + ScatterTransferOperation&&) noexcept = default; + +TransferEngine::ScatterTransferOperation::~ScatterTransferOperation() = default; + +Status TransferEngine::ScatterTransferOperation::wait() { + return impl_ + ? impl_->wait() + : Status::InvalidArgument("invalid scatter transfer operation"); +} + +Status TransferEngine::ScatterTransferOperation::waitFor( + std::chrono::nanoseconds timeout) { + return impl_ + ? impl_->waitFor(timeout) + : Status::InvalidArgument("invalid scatter transfer operation"); +} + +TransferEngine::ScatterTransferOperation TransferEngine::submitScatter( + const std::vector& ranges) { + ScatterTransferOperation::Impl::Backend backend; + backend.legacy = impl_; +#ifdef USE_TENT + backend.tent = impl_tent_; +#endif + + return ScatterTransferOperation( + std::make_unique( + *this, std::move(backend), ranges)); +} + +Status TransferEngine::transferScatter( + const std::vector& ranges) { + auto operation = submitScatter(ranges); + return operation.wait(); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transfer_engine_c.cpp b/mooncake-transfer-engine/src/transfer_engine_c.cpp index b643a15386..568f75d7b1 100644 --- a/mooncake-transfer-engine/src/transfer_engine_c.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_c.cpp @@ -14,6 +14,7 @@ #include "transfer_engine_c.h" +#include #include #include @@ -250,3 +251,36 @@ int syncSegmentCache(transfer_engine_t engine) { TransferEngine *native = (TransferEngine *)engine; return native->syncSegmentCache(); } + +void enableGracefulShutdown(transfer_engine_t engine) { + TransferEngine *native = (TransferEngine *)engine; + native->enableGracefulShutdown(); +} + +int getNicLoadStats(transfer_engine_t engine, nic_load_stat_t *stats, + size_t *count) { + if (!engine || !stats || !count) return -1; + TransferEngine *native = (TransferEngine *)engine; + std::vector native_stats; + Status s = native->getNicLoadStats(native_stats); + if (!s.ok()) return (int)s.code(); + size_t to_copy = std::min(native_stats.size(), *count); + for (size_t i = 0; i < to_copy; ++i) { + snprintf(stats[i].device_name, sizeof(stats[i].device_name), "%s", + native_stats[i].device_name.c_str()); + stats[i].inflight_bytes = native_stats[i].inflight_bytes; + stats[i].ewma_bandwidth_bps = native_stats[i].ewma_bandwidth_bps; + } + *count = native_stats.size(); + return 0; +} + +int showLinks(transfer_engine_t engine, char *buf_out, size_t buf_len, + int json) { + if (!engine || !buf_out || buf_len == 0) return -1; + + TransferEngine *native = (TransferEngine *)engine; + auto result = native->showLinks(json != 0); + snprintf(buf_out, buf_len, "%s", result.c_str()); + return 0; +} diff --git a/mooncake-transfer-engine/src/transfer_engine_impl.cpp b/mooncake-transfer-engine/src/transfer_engine_impl.cpp index 83960ad575..a7930021f5 100644 --- a/mooncake-transfer-engine/src/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/src/transfer_engine_impl.cpp @@ -202,6 +202,27 @@ int TransferEngineImpl::init(const std::string& metadata_conn_string, int ret = metadata_->addRpcMetaEntry(local_server_name_, desc); if (ret) return ret; + // Universal TCP force mechanism: if MC_FORCE_TCP is set, skip all other + // transport installation logic and use TCP transport only. This allows + // running metadata-only instances without requiring specialized hardware + // (e.g., NPU for Ascend Direct, RDMA HCAs, etc.). + if (getenv("MC_FORCE_TCP")) { +#ifdef USE_TCP + Transport* tcp_transport = + multi_transports_->installTransport("tcp", nullptr); + if (!tcp_transport) { + LOG(ERROR) + << "MC_FORCE_TCP is set but failed to install TCP transport"; + return -1; + } + LOG(INFO) << "MC_FORCE_TCP is set, using TCP transport only"; + return 0; +#else + LOG(ERROR) << "MC_FORCE_TCP is set but USE_TCP is not compiled in"; + return -1; +#endif + } + #if defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) Transport* ascend_transport = multi_transports_->installTransport("ascend", local_topology_); @@ -218,7 +239,7 @@ int TransferEngineImpl::init(const std::string& metadata_conn_string, LOG(ERROR) << "Failed to install UBShmem transport"; return -1; } - auto_discover_ = false; + auto_discover_config_.enabled = false; #endif #if defined(USE_CXL) && !defined(USE_ASCEND) && \ @@ -233,7 +254,7 @@ int TransferEngineImpl::init(const std::string& metadata_conn_string, } #endif - if (auto_discover_) { + if (auto_discover_config_.enabled) { LOG(INFO) << "Auto-discovering topology..."; if (getenv("MC_CUSTOM_TOPO_JSON")) { auto path = getenv("MC_CUSTOM_TOPO_JSON"); @@ -345,27 +366,26 @@ int TransferEngineImpl::init(const std::string& metadata_conn_string, if ((local_topology_->getHcaList().size() > 0 && !getenv("MC_FORCE_TCP")) || getenv("MC_FORCE_HCA")) { - // only install RDMA transport when there is at least one HCA - Transport* rdma_transport = nullptr; - if (use_barex_) { + const std::string transport_type = autoDiscoverTransport(); + Transport* transport = nullptr; + if (transport_type == "barex") { #ifdef USE_BAREX - rdma_transport = multi_transports_->installTransport( + transport = multi_transports_->installTransport( "barex", local_topology_); #else LOG(ERROR) << "Set USE BAREX while barex not compiled"; return -1; #endif } else { - rdma_transport = multi_transports_->installTransport( - "rdma", local_topology_); + transport = multi_transports_->installTransport( + transport_type, local_topology_); } - if (rdma_transport == nullptr) { - LOG(ERROR) << "Failed to install RDMA transport, type=" - << (use_barex_ ? "barex" : "rdma"); + if (transport == nullptr) { + LOG(ERROR) << "Failed to install transport, type=" + << transport_type; return -1; } else { - LOG(INFO) << "installTransport, type=" - << (use_barex_ ? "barex" : "rdma"); + LOG(INFO) << "installTransport, type=" << transport_type; } } else { Transport* tcp_transport = @@ -414,6 +434,13 @@ Transport* TransferEngineImpl::installTransport(const std::string& proto, LOG(WARNING) << "Transport " << proto << " already installed"; return transport; } +#ifdef USE_NCCL_HOST + if (proto == "nccl" && !local_memory_regions_.empty()) { + LOG(ERROR) << "Install NCCL before registering local memory so peer " + "buffer order remains deterministic"; + return nullptr; + } +#endif if (args != nullptr && args[0] != nullptr) { const std::string nic_priority_matrix = static_cast(args[0]); @@ -444,7 +471,8 @@ int TransferEngineImpl::uninstallTransport(const std::string& proto) { return 0; } -#if defined(USE_CUDA) || defined(USE_MUSA) +#if (defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_MACA)) && \ + !defined(USE_CXI) device::P2pTransport* TransferEngineImpl::getOrCreateP2pTransport( int num_ranks) { if (!p2p_transport_) { @@ -462,6 +490,15 @@ device::RdmaTransport* TransferEngineImpl::getOrCreateRdmaTransport( } #endif +#ifdef USE_NCCL_DEVICE +device::NcclTransport* TransferEngineImpl::getOrCreateNcclTransport() { + if (!nccl_transport_) { + nccl_transport_ = device::createNcclDeviceTransport(); + } + return nccl_transport_.get(); +} +#endif + int TransferEngineImpl::getRpcPort() { return metadata_->localRpcMeta().rpc_port; } @@ -565,33 +602,59 @@ int TransferEngineImpl::registerLocalMemory(void* addr, size_t length, const std::string& location, bool remote_accessible, bool update_metadata) { - if (checkOverlap(addr, length)) { - LOG(ERROR) - << "Transfer Engine does not support overlapped memory region"; - return ERR_ADDRESS_OVERLAPPED; - } if (length == 0) { LOG(ERROR) << "Transfer Engine does not support zero length memory region"; return ERR_INVALID_ARGUMENT; } + + std::vector regions = { + {addr, length, location, remote_accessible}}; + if (!tryReserveMemoryRegions(regions)) { + LOG(ERROR) + << "Transfer Engine does not support overlapped memory region"; + return ERR_ADDRESS_OVERLAPPED; + } + + std::vector attempted_transports; for (auto transport : multi_transports_->listTransports()) { + attempted_transports.push_back(transport); int ret = transport->registerLocalMemory( addr, length, location, remote_accessible, update_metadata); - if (ret < 0) return ret; + if (ret < 0) { + // Roll back the transports that already registered so a partial + // failure doesn't leave the region registered on some of them. + // Mirrors registerLocalMemoryBatch (#2869). + for (auto it = attempted_transports.rbegin(); + it != attempted_transports.rend(); ++it) { + int rollback_ret = + (*it)->unregisterLocalMemory(addr, update_metadata); + if (rollback_ret != 0 && + rollback_ret != ERR_ADDRESS_NOT_REGISTERED) { + LOG(WARNING) + << "Failed to roll back registration for " + << (*it)->getName() << ", ret=" << rollback_ret; + } + } + releaseMemoryRegions(regions); + return ret; + } } - std::unique_lock lock(mutex_); - insertMemoryRegionLocked({addr, length, location, remote_accessible}); + commitMemoryRegions(regions); return 0; } int TransferEngineImpl::unregisterLocalMemory(void* addr, bool update_metadata) { + // Best-effort: try every transport so one failure can't leave the region + // registered on the others; mirrors unregisterLocalMemoryBatch (#2869). + int first_error = 0; for (auto& transport : multi_transports_->listTransports()) { int ret = transport->unregisterLocalMemory(addr, update_metadata); - if (ret) return ret; + if (ret && !first_error) first_error = ret; } + if (first_error) return first_error; std::unique_lock lock(mutex_); eraseMemoryRegionLocked(addr); @@ -720,31 +783,79 @@ int TransferEngineImpl::mp_unregisterLocalMemory( int TransferEngineImpl::registerLocalMemoryBatch( const std::vector& buffer_list, const std::string& location) { - for (auto& buffer : buffer_list) { - if (checkOverlap(buffer.addr, buffer.length)) { + std::vector sorted_buffers = buffer_list; + std::sort(sorted_buffers.begin(), sorted_buffers.end(), + [](const BufferEntry& lhs, const BufferEntry& rhs) { + return reinterpret_cast(lhs.addr) < + reinterpret_cast(rhs.addr); + }); + + for (size_t i = 0; i < sorted_buffers.size(); ++i) { + const auto& buffer = sorted_buffers[i]; + if (buffer.length == 0) { LOG(ERROR) - << "Transfer Engine does not support overlapped memory region"; - return ERR_ADDRESS_OVERLAPPED; + << "Transfer Engine does not support zero length memory region"; + return ERR_INVALID_ARGUMENT; } + + if (i > 0) { + const auto& previous = sorted_buffers[i - 1]; + auto address = reinterpret_cast(buffer.addr); + auto previous_address = reinterpret_cast(previous.addr); + if (address - previous_address < previous.length) { + LOG(ERROR) << "Transfer Engine does not support overlapped " + "memory region"; + return ERR_ADDRESS_OVERLAPPED; + } + } + } + + std::vector regions; + std::vector addr_list; + regions.reserve(buffer_list.size()); + addr_list.reserve(buffer_list.size()); + for (const auto& buffer : buffer_list) { + regions.push_back({buffer.addr, buffer.length, location, true}); + addr_list.push_back(buffer.addr); } + if (!tryReserveMemoryRegions(regions)) { + LOG(ERROR) + << "Transfer Engine does not support overlapped memory region"; + return ERR_ADDRESS_OVERLAPPED; + } + + std::vector attempted_transports; for (auto transport : multi_transports_->listTransports()) { + attempted_transports.push_back(transport); int ret = transport->registerLocalMemoryBatch(buffer_list, location); - if (ret < 0) return ret; + if (ret) { + for (auto it = attempted_transports.rbegin(); + it != attempted_transports.rend(); ++it) { + int rollback_ret = (*it)->unregisterLocalMemoryBatch(addr_list); + if (rollback_ret != 0 && + rollback_ret != ERR_ADDRESS_NOT_REGISTERED) { + LOG(WARNING) + << "Failed to roll back batch registration for " + << (*it)->getName() << ", ret=" << rollback_ret; + } + } + releaseMemoryRegions(regions); + return ret; + } } - std::unique_lock lock(mutex_); - for (auto& buffer : buffer_list) { - insertMemoryRegionLocked({buffer.addr, buffer.length, location, true}); - } + commitMemoryRegions(regions); return 0; } int TransferEngineImpl::unregisterLocalMemoryBatch( const std::vector& addr_list) { + int first_error = 0; for (auto transport : multi_transports_->listTransports()) { int ret = transport->unregisterLocalMemoryBatch(addr_list); - if (ret < 0) return ret; + if (ret && !first_error) first_error = ret; } + if (first_error) return first_error; std::unique_lock lock(mutex_); for (auto& addr : addr_list) { @@ -753,51 +864,27 @@ int TransferEngineImpl::unregisterLocalMemoryBatch( return 0; } -TransferEngineImpl::MemoryRegionMap::iterator -TransferEngineImpl::findMemoryRegionContaining(uintptr_t addr) { - auto upper = local_memory_regions_.upper_bound(addr); - if (upper == local_memory_regions_.begin()) { - return local_memory_regions_.end(); - } - auto candidate = std::prev(upper); - return overlapWithRegion(addr, 1, candidate->second.addr, - candidate->second.length) - ? candidate - : local_memory_regions_.end(); -} - -TransferEngineImpl::MemoryRegionMap::const_iterator -TransferEngineImpl::findMemoryRegionContaining(uintptr_t addr) const { - auto upper = local_memory_regions_.upper_bound(addr); - if (upper == local_memory_regions_.begin()) { - return local_memory_regions_.end(); - } - auto candidate = std::prev(upper); - return overlapWithRegion(addr, 1, candidate->second.addr, - candidate->second.length) - ? candidate - : local_memory_regions_.end(); -} - bool TransferEngineImpl::hasOverlapLocked(uintptr_t addr, uint64_t length) const { + return hasOverlapInMapLocked(local_memory_regions_, addr, length) || + hasOverlapInMapLocked(registering_memory_regions_, addr, length); +} + +bool TransferEngineImpl::hasOverlapInMapLocked(const MemoryRegionMap& regions, + uintptr_t addr, + uint64_t length) const { if (length == 0) { return false; } - auto containing = findMemoryRegionContaining(addr); - if (containing != local_memory_regions_.end()) { - return true; - } - - auto next = local_memory_regions_.lower_bound(addr); - if (next != local_memory_regions_.end() && + auto next = regions.lower_bound(addr); + if (next != regions.end() && overlapWithRegion(addr, length, next->second.addr, next->second.length)) { return true; } - if (next != local_memory_regions_.begin()) { + if (next != regions.begin()) { auto prev = std::prev(next); if (overlapWithRegion(addr, length, prev->second.addr, prev->second.length)) { @@ -808,6 +895,45 @@ bool TransferEngineImpl::hasOverlapLocked(uintptr_t addr, return false; } +bool TransferEngineImpl::tryReserveMemoryRegions( + const std::vector& regions) { + std::unique_lock lock(mutex_); + std::vector reserved; + reserved.reserve(regions.size()); + + for (const auto& region : regions) { + auto addr = reinterpret_cast(region.addr); + if (hasOverlapLocked(addr, region.length)) { + for (auto reserved_addr : reserved) { + registering_memory_regions_.erase(reserved_addr); + } + return false; + } + registering_memory_regions_[addr] = region; + reserved.push_back(addr); + } + return true; +} + +void TransferEngineImpl::commitMemoryRegions( + const std::vector& regions) { + std::unique_lock lock(mutex_); + for (const auto& region : regions) { + registering_memory_regions_.erase( + reinterpret_cast(region.addr)); + insertMemoryRegionLocked(region); + } +} + +void TransferEngineImpl::releaseMemoryRegions( + const std::vector& regions) { + std::unique_lock lock(mutex_); + for (const auto& region : regions) { + registering_memory_regions_.erase( + reinterpret_cast(region.addr)); + } +} + void TransferEngineImpl::insertMemoryRegionLocked(const MemoryRegion& region) { local_memory_regions_[reinterpret_cast(region.addr)] = region; } diff --git a/mooncake-transfer-engine/src/transfer_metadata.cpp b/mooncake-transfer-engine/src/transfer_metadata.cpp index ebf0fe0bd8..48bb44fef7 100644 --- a/mooncake-transfer-engine/src/transfer_metadata.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata.cpp @@ -19,6 +19,8 @@ #include #include #include +#include +#include #include "common.h" #include "config.h" @@ -66,6 +68,7 @@ struct TransferNotifyUtil { struct TransferHandshakeUtil { static Json::Value encode(const TransferMetadata::HandShakeDesc &desc) { Json::Value root; + root["payload"] = desc.payload; root["local_nic_path"] = desc.local_nic_path; root["local_lid"] = desc.local_lid; root["local_gid"] = desc.local_gid; @@ -76,10 +79,15 @@ struct TransferHandshakeUtil { Json::Value qpNums(Json::arrayValue); for (const auto &qp : desc.qp_num) qpNums.append(qp); root["qp_num"] = qpNums; + if (desc.ready_ack_supported || desc.ready_ack) + root["ready_ack"] = desc.ready_ack; root["reply_msg"] = desc.reply_msg; #ifdef USE_EFA root["efa_addr"] = desc.efa_addr; // EFA endpoint address #endif +#ifdef USE_CXI + root["cxi_addr"] = desc.cxi_addr; +#endif #ifdef USE_UB Json::Value jettyNums(Json::arrayValue); @@ -93,6 +101,7 @@ struct TransferHandshakeUtil { } static int decode(Json::Value root, TransferMetadata::HandShakeDesc &desc) { + desc.payload = root["payload"].asString(); desc.local_nic_path = root["local_nic_path"].asString(); if (root.isMember("local_lid") && root["local_lid"].isUInt()) { desc.local_lid = root["local_lid"].asUInt(); @@ -110,11 +119,21 @@ struct TransferHandshakeUtil { #endif for (const auto &qp : root["qp_num"]) desc.qp_num.push_back(qp.asUInt()); + desc.ready_ack_supported = root.isMember("ready_ack"); + if (desc.ready_ack_supported && root["ready_ack"].isBool()) { + desc.ready_ack = root["ready_ack"].asBool(); + } else { + desc.ready_ack = false; + } desc.reply_msg = root["reply_msg"].asString(); #ifdef USE_EFA desc.efa_addr = root["efa_addr"].asString(); // EFA endpoint address #endif +#ifdef USE_CXI + desc.cxi_addr = root["cxi_addr"].asString(); +#endif + #ifdef USE_UB for (const auto &jetty : root["jetty_num"]) { desc.jetty_num.push_back(jetty.asUInt()); @@ -154,6 +173,7 @@ TransferMetadata::TransferMetadata(const std::string &conn_string) { } if (conn_string == P2PHANDSHAKE) { p2p_handshake_mode_ = true; + startMetadataRefreshPollingIfNeeded(); return; } storage_plugin_ = MetadataStoragePlugin::Create(conn_string); @@ -162,9 +182,68 @@ TransferMetadata::TransferMetadata(const std::string &conn_string) { << "Unable to create metadata storage plugin with conn string " << conn_string; } + startMetadataRefreshPollingIfNeeded(); +} + +TransferMetadata::~TransferMetadata() { + stopMetadataRefreshPollingThread(); + handshake_plugin_.reset(); + storage_plugin_.reset(); +} + +void TransferMetadata::startMetadataRefreshPollingIfNeeded() { + const auto &config = globalConfig(); + if (!config.metacache || config.te_metadata_refresh_interval_seconds == 0) { + return; + } + if (!p2p_handshake_mode_ && !storage_plugin_) { + return; + } + + const auto refresh_interval_seconds = + config.te_metadata_refresh_interval_seconds; + should_stop_metadata_refresh_thread_ = false; + metadata_refresh_thread_ = std::thread([this, refresh_interval_seconds]() { + metadataRefreshPollingLoop(refresh_interval_seconds); + }); + LOG(INFO) << "TE metadata refresh polling enabled, interval_seconds=" + << refresh_interval_seconds; +} + +void TransferMetadata::stopMetadataRefreshPollingThread() { + should_stop_metadata_refresh_thread_ = true; + metadata_refresh_cv_.notify_all(); + if (metadata_refresh_thread_.joinable()) { + metadata_refresh_thread_.join(); + } } -TransferMetadata::~TransferMetadata() { handshake_plugin_.reset(); } +void TransferMetadata::metadataRefreshPollingLoop( + uint64_t refresh_interval_seconds) { + std::unique_lock lock(metadata_refresh_mutex_); + while (!should_stop_metadata_refresh_thread_) { + if (metadata_refresh_cv_.wait_for( + lock, std::chrono::seconds(refresh_interval_seconds), [this]() { + return should_stop_metadata_refresh_thread_.load(); + })) { + break; + } + lock.unlock(); + try { + int ret = syncSegmentCache(""); + if (ret) { + LOG(WARNING) + << "TE metadata refresh polling failed, ret=" << ret; + } + } catch (const std::exception &e) { + LOG(ERROR) << "Exception in TE metadata refresh polling: " + << e.what(); + } catch (...) { + LOG(ERROR) << "Unknown exception in TE metadata refresh polling"; + } + lock.lock(); + } +} std::string TransferMetadata::getFullMetadataKey( const std::string &segment_name) const { @@ -262,12 +341,16 @@ static int encodeMultiProtocolSegmentDesc( bufferJSON["lkey"] = lkeyJSON; } else if (buffer.protocol == "tcp") { bufferJSON["addr"] = static_cast(buffer.addr); + } else if (buffer.protocol == "hip" || buffer.protocol == "maca") { + bufferJSON["addr"] = static_cast(buffer.addr); + bufferJSON["shm_name"] = buffer.shm_name; } buffersJSON.append(bufferJSON); } segmentJSON["buffers"] = buffersJSON; segmentJSON["protocol"] = protocolJSON; segmentJSON["tcp_data_port"] = desc.tcp_data_port; + segmentJSON["tcp_proto_version"] = desc.tcp_proto_version; segmentJSON["timestamp"] = getCurrentDateTime(); return 0; @@ -277,35 +360,28 @@ static int encodeMultiProtocolSegmentDesc( int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, Json::Value &segmentJSON) { #ifdef ENABLE_MULTI_PROTOCOL - // Check if this is a multi-protocol scenario (CXL+TCP or CXL+RDMA) + // A segment is multi-protocol when more than one transport registered + // buffers on it (e.g. tcp+hip or rdma+hip for intra-node disagg). Every + // protocol must have a per-buffer emitter below; otherwise its buffers + // would be silently dropped. std::vector protocols = splitProtocols(desc.protocol); bool is_multi_protocol = false; - if (protocols.size() == 2) { - // Only support CXL+TCP or CXL+RDMA combinations - bool has_cxl = false, has_tcp = false, has_rdma = false; + if (protocols.size() >= 2) { + is_multi_protocol = true; for (const auto &proto : protocols) { - if (proto == "cxl") - has_cxl = true; - else if (proto == "tcp") - has_tcp = true; - else if (proto == "rdma") - has_rdma = true; - } - // Multi-protocol only supported for CXL+TCP or CXL+RDMA - if (has_cxl && (has_tcp || has_rdma)) { - is_multi_protocol = true; + if (proto != "cxl" && proto != "tcp" && proto != "rdma" && + proto != "hip" && proto != "maca") { + is_multi_protocol = false; + break; + } } - // If not valid multi-protocol combination, return error if (!is_multi_protocol) { LOG(ERROR) << "Unsupported multi-protocol combination: " << desc.protocol - << ". Only CXL+TCP or CXL+RDMA are supported."; + << ". Only cxl, tcp, rdma, hip and maca may be " + "combined."; return ERR_INVALID_ARGUMENT; } - } else if (protocols.size() > 2) { - LOG(ERROR) << "Unsupported multi-protocol combination: " - << desc.protocol << ". Maximum 2 protocols allowed."; - return ERR_INVALID_ARGUMENT; } // If multi-protocol scenario, use multi-protocol encoding @@ -317,6 +393,7 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, segmentJSON["name"] = desc.name; segmentJSON["protocol"] = desc.protocol; segmentJSON["tcp_data_port"] = desc.tcp_data_port; + segmentJSON["tcp_proto_version"] = desc.tcp_proto_version; segmentJSON["timestamp"] = getCurrentDateTime(); if (!desc.rdma_server_name.empty()) { segmentJSON["rdma_server_name"] = desc.rdma_server_name; @@ -324,7 +401,7 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, if (segmentJSON["protocol"] == "rdma" || segmentJSON["protocol"] == "barex" || - segmentJSON["protocol"] == "efa") { + segmentJSON["protocol"] == "efa" || segmentJSON["protocol"] == "cxi") { Json::Value devicesJSON(Json::arrayValue); for (const auto &device : desc.devices) { Json::Value deviceJSON; @@ -384,6 +461,17 @@ int TransferMetadata::encodeSegmentDesc(const SegmentDesc &desc, buffersJSON.append(bufferJSON); } segmentJSON["buffers"] = buffersJSON; + } else if (segmentJSON["protocol"] == "nccl") { + Json::Value buffersJSON(Json::arrayValue); + for (const auto &buffer : desc.buffers) { + Json::Value bufferJSON; + bufferJSON["name"] = buffer.name; + bufferJSON["addr"] = static_cast(buffer.addr); + bufferJSON["length"] = static_cast(buffer.length); + bufferJSON["device_id"] = buffer.device_id; + buffersJSON.append(bufferJSON); + } + segmentJSON["buffers"] = buffersJSON; } else if (segmentJSON["protocol"] == "ascend") { Json::Value devicesJSON(Json::arrayValue); for (const auto &device : desc.devices) { @@ -513,6 +601,9 @@ decodeMultiProtocolSegmentDesc(Json::Value &segmentJSON, auto desc = std::make_shared(); desc->name = segmentJSON["name"].asString(); desc->tcp_data_port = segmentJSON["tcp_data_port"].asInt(); + desc->tcp_proto_version = segmentJSON.isMember("tcp_proto_version") + ? segmentJSON["tcp_proto_version"].asInt() + : 1; if (segmentJSON.isMember("timestamp")) desc->timestamp = segmentJSON["timestamp"].asString(); if (segmentJSON.isMember("rdma_server_name")) @@ -574,18 +665,27 @@ decodeMultiProtocolSegmentDesc(Json::Value &segmentJSON, buffer.length = bufferJSON["length"].asUInt64(); buffer.protocol = buffer_protocol; for (const auto &rkeyJSON : bufferJSON["rkey"]) - buffer.rkey.push_back(rkeyJSON.asUInt()); + buffer.rkey.push_back( + static_cast( + rkeyJSON.asUInt64())); for (const auto &lkeyJSON : bufferJSON["lkey"]) - buffer.lkey.push_back(lkeyJSON.asUInt()); + buffer.lkey.push_back( + static_cast( + lkeyJSON.asUInt64())); + // The same topology-derived device_id indexes both rkey and + // devices, and a publisher emits exactly one key per device, so a + // key vector longer than the device list lets that index run past + // the end of devices[]. See the comment in decodeSegmentDesc. if (buffer.name.empty() || !buffer.addr || !buffer.length || - buffer.rkey.empty() || - buffer.rkey.size() != buffer.lkey.size()) { + (!buffer.rkey.empty() && + (buffer.rkey.size() != buffer.lkey.size() || + buffer.rkey.size() != desc->devices.size()))) { LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " buffer_protocol " << buffer_protocol << ", " << buffer.name << ", " << buffer.addr << ", " << buffer.length << ", " << buffer.rkey.size() << ", " - << buffer.lkey.size(); + << buffer.lkey.size() << ", " << desc->devices.size(); return nullptr; } desc->buffers.push_back(buffer); @@ -602,6 +702,21 @@ decodeMultiProtocolSegmentDesc(Json::Value &segmentJSON, return nullptr; } desc->buffers.push_back(buffer); + } else if (buffer_protocol == "hip" || buffer_protocol == "maca") { + TransferMetadata::BufferDesc buffer; + buffer.name = bufferJSON["name"].asString(); + buffer.addr = bufferJSON["addr"].asUInt64(); + buffer.length = bufferJSON["length"].asUInt64(); + buffer.shm_name = bufferJSON["shm_name"].asString(); + buffer.protocol = buffer_protocol; + if (buffer.name.empty() || !buffer.addr || !buffer.length || + buffer.shm_name.empty()) { + LOG(WARNING) + << "Corrupted segment descriptor, name " << segment_name + << " buffer_protocol " << buffer_protocol; + return nullptr; + } + desc->buffers.push_back(buffer); } } @@ -617,34 +732,26 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, bool is_multi_protocol = false; if (segmentJSON["protocol"].isArray()) { size_t proto_count = segmentJSON["protocol"].size(); - if (proto_count == 2) { - // Only support CXL+TCP or CXL+RDMA combinations - bool has_cxl = false, has_tcp = false, has_rdma = false; + if (proto_count >= 2) { + // Every protocol must have a per-buffer parser in + // decodeMultiProtocolSegmentDesc below. + is_multi_protocol = true; for (const auto &protocolStr : segmentJSON["protocol"]) { std::string proto = protocolStr.asString(); - if (proto == "cxl") - has_cxl = true; - else if (proto == "tcp") - has_tcp = true; - else if (proto == "rdma") - has_rdma = true; - } - // Multi-protocol only supported for CXL+TCP or CXL+RDMA - if (has_cxl && (has_tcp || has_rdma)) { - is_multi_protocol = true; + if (proto != "cxl" && proto != "tcp" && proto != "rdma" && + proto != "hip" && proto != "maca") { + is_multi_protocol = false; + break; + } } - // If not valid multi-protocol combination, return error if (!is_multi_protocol) { - LOG(ERROR) - << "Unsupported multi-protocol combination in segment: " - << segment_name - << ". Only CXL+TCP or CXL+RDMA are supported."; + LOG(ERROR) << "Unsupported multi-protocol combination in " + "segment: " + << segment_name + << ". Only cxl, tcp, rdma, hip and maca may be " + "combined."; return nullptr; } - } else if (proto_count > 2) { - LOG(ERROR) << "Unsupported multi-protocol combination in segment: " - << segment_name << ". Maximum 2 protocols allowed."; - return nullptr; } } @@ -658,13 +765,16 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, desc->name = segmentJSON["name"].asString(); desc->protocol = segmentJSON["protocol"].asString(); desc->tcp_data_port = segmentJSON["tcp_data_port"].asInt(); + desc->tcp_proto_version = segmentJSON.isMember("tcp_proto_version") + ? segmentJSON["tcp_proto_version"].asInt() + : 1; if (segmentJSON.isMember("timestamp")) desc->timestamp = segmentJSON["timestamp"].asString(); if (segmentJSON.isMember("rdma_server_name")) desc->rdma_server_name = segmentJSON["rdma_server_name"].asString(); if (desc->protocol == "rdma" || desc->protocol == "barex" || - desc->protocol == "efa") { + desc->protocol == "efa" || desc->protocol == "cxi") { for (const auto &deviceJSON : segmentJSON["devices"]) { DeviceDesc device; device.name = deviceJSON["name"].asString(); @@ -678,23 +788,43 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, desc->devices.push_back(device); } + // rdma, efa and cxi pick a device with topology.selectDevice() and + // then use that one device_id to index both buffers[].rkey and + // devices[] (RdmaTransport::selectDevice / WorkerPool::selectPeerDevice + // and their efa/cxi counterparts). A publisher emits exactly one key + // per device, so a descriptor whose key vector is longer than its + // device list can drive device_id past the end of devices[]. barex is + // excluded: it hands the whole rkey vector to the peer rather than + // indexing it by device_id, and its key count comes from the mempool + // MR set rather than the device list. + const bool keys_indexed_by_device = desc->protocol == "rdma" || + desc->protocol == "efa" || + desc->protocol == "cxi"; + for (const auto &bufferJSON : segmentJSON["buffers"]) { BufferDesc buffer; buffer.name = bufferJSON["name"].asString(); buffer.addr = bufferJSON["addr"].asUInt64(); buffer.length = bufferJSON["length"].asUInt64(); for (const auto &rkeyJSON : bufferJSON["rkey"]) - buffer.rkey.push_back(rkeyJSON.asUInt()); + buffer.rkey.push_back( + static_cast( + rkeyJSON.asUInt64())); for (const auto &lkeyJSON : bufferJSON["lkey"]) - buffer.lkey.push_back(lkeyJSON.asUInt()); + buffer.lkey.push_back( + static_cast( + lkeyJSON.asUInt64())); if (buffer.name.empty() || !buffer.addr || !buffer.length || - buffer.rkey.empty() || - buffer.rkey.size() != buffer.lkey.size()) { + (!buffer.rkey.empty() && + (buffer.rkey.size() != buffer.lkey.size() || + (keys_indexed_by_device && + buffer.rkey.size() != desc->devices.size())))) { LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol << ", " << buffer.name << ", " << buffer.addr << ", " << buffer.length << ", " - << buffer.rkey.size() << ", " << buffer.lkey.size(); + << buffer.rkey.size() << ", " << buffer.lkey.size() << ", " + << desc->devices.size(); return nullptr; } desc->buffers.push_back(buffer); @@ -755,6 +885,23 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, } desc->buffers.push_back(buffer); } + } else if (desc->protocol == "nccl") { + for (const auto &bufferJSON : segmentJSON["buffers"]) { + BufferDesc buffer; + buffer.name = bufferJSON["name"].asString(); + buffer.addr = bufferJSON["addr"].asUInt64(); + buffer.length = bufferJSON["length"].asUInt64(); + buffer.device_id = bufferJSON.isMember("device_id") + ? bufferJSON["device_id"].asInt() + : -1; + if (buffer.name.empty() || !buffer.addr || !buffer.length || + buffer.device_id < 0) { + LOG(WARNING) << "Corrupted segment descriptor, name " + << segment_name << " protocol " << desc->protocol; + return nullptr; + } + desc->buffers.push_back(buffer); + } } else if (desc->protocol == "nvlink" || desc->protocol == "nvlink_intra" || desc->protocol == "hip" || desc->protocol == "maca" || desc->protocol == "ubshmem" || @@ -765,8 +912,15 @@ TransferMetadata::decodeSegmentDesc(Json::Value &segmentJSON, buffer.addr = bufferJSON["addr"].asUInt64(); buffer.length = bufferJSON["length"].asUInt64(); buffer.shm_name = bufferJSON["shm_name"].asString(); - if (buffer.name.empty() || !buffer.addr || !buffer.length || - buffer.shm_name.empty()) { + if (buffer.shm_name.empty()) { + // In a multi-protocol build every transport registers each + // buffer into this node's single shared segment. Buffers owned + // by another transport (e.g. RDMA) carry no HIP IPC handle and + // are not reachable via this transport. Skip them instead of + // rejecting the whole segment, which would tear down sessions. + continue; + } + if (buffer.name.empty() || !buffer.addr || !buffer.length) { LOG(WARNING) << "Corrupted segment descriptor, name " << segment_name << " protocol " << desc->protocol << "buffer name " << buffer.name << "buffer addr " @@ -935,7 +1089,21 @@ std::shared_ptr TransferMetadata::getSegmentDesc( return result; } +bool TransferMetadata::SegmentDesc::operator==(const SegmentDesc &other) const { + // timestamp is intentionally excluded: metadata encoding may refresh it + // even when the operational descriptor is unchanged. + return name == other.name && protocol == other.protocol && + devices == other.devices && topology == other.topology && + buffers == other.buffers && nvmeof_buffers == other.nvmeof_buffers && + cxl_name == other.cxl_name && cxl_base_addr == other.cxl_base_addr && + rank_info == other.rank_info && + tcp_data_port == other.tcp_data_port && + rdma_server_name == other.rdma_server_name; +} + int TransferMetadata::syncSegmentCache(const std::string &segment_name) { + const auto sync_start = std::chrono::steady_clock::now(); + // Collect segment names to sync first, then release lock before network I/O std::vector names_to_sync; { @@ -948,25 +1116,75 @@ int TransferMetadata::syncSegmentCache(const std::string &segment_name) { } } + size_t fetched_count = 0; + size_t failed_count = 0; + size_t updated_count = 0; + size_t unchanged_count = 0; + size_t skipped_count = 0; + // Fetch updates without holding lock (may involve network I/O) std::vector>> updates; for (const auto &name : names_to_sync) { auto segment_desc = getSegmentDesc(name); if (segment_desc) { updates.emplace_back(name, segment_desc); + ++fetched_count; } else { + ++failed_count; LOG(WARNING) << "segment " << name << " is now invalid"; } } - // Apply updates with write lock - RWSpinlock::WriteGuard guard(segment_lock_); - for (const auto &[name, desc] : updates) { - auto it = segment_name_to_id_map_.find(name); - if (it != segment_name_to_id_map_.end()) { - segment_id_to_desc_map_[it->second] = desc; + { + // Apply updates with write lock + RWSpinlock::WriteGuard guard(segment_lock_); + for (const auto &[name, desc] : updates) { + auto it = segment_name_to_id_map_.find(name); + if (it == segment_name_to_id_map_.end()) { + ++skipped_count; + continue; + } + + const auto segment_id = it->second; + auto current_it = segment_id_to_desc_map_.find(segment_id); + const auto old_desc = current_it == segment_id_to_desc_map_.end() + ? nullptr + : current_it->second; + bool changed = true; + if (old_desc) { + changed = *old_desc != *desc; + } + + if (!changed) { + ++unchanged_count; + continue; + } + + segment_id_to_desc_map_[segment_id] = desc; + ++updated_count; + LOG(WARNING) << "Segment cache descriptor changed, name=" << name + << ", segment_id=" << segment_id; + if (old_desc) { + LOG(INFO) << "Old segment descriptor:"; + old_desc->dump(); + } else { + LOG(INFO) << "Old segment descriptor: "; + } + LOG(INFO) << "New segment descriptor:"; + desc->dump(); } } + const auto sync_duration_ms = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - sync_start) + .count(); + LOG(INFO) << "Segment cache sync finished, requested_segment=" + << (segment_name.empty() ? "" : segment_name) + << ", scanned=" << names_to_sync.size() + << ", fetched=" << fetched_count << ", updated=" << updated_count + << ", unchanged=" << unchanged_count + << ", failed=" << failed_count << ", skipped=" << skipped_count + << ", sync_duration_ms=" << sync_duration_ms; return 0; } @@ -1190,10 +1408,9 @@ int TransferMetadata::rePublishRpcMetaEntry(const std::string &server_name) { Json::Value existing; if (storage_plugin_->get(full_key, existing)) { - Json::Value desired; - desired["ip_or_host_name"] = local_rpc_meta_.ip_or_host_name; - desired["rpc_port"] = static_cast(local_rpc_meta_.rpc_port); - if (existing == desired) { + if (existing["ip_or_host_name"].asString() == + local_rpc_meta_.ip_or_host_name && + existing["rpc_port"].asUInt() == local_rpc_meta_.rpc_port) { return 0; } } @@ -1246,7 +1463,18 @@ int TransferMetadata::startHandshakeDaemon( TransferHandshakeUtil::decode(peer, peer_desc); if (on_receive_handshake) { int ret = on_receive_handshake(peer_desc, local_desc); - if (ret) return ret; + if (ret) { + if (local_desc.reply_msg.empty()) { + local_desc.reply_msg = + "Handshake callback failed: " + std::to_string(ret); + } + // The callback failure is a handshake-level rejection, not + // an RPC handler failure. Return a structured reply so the + // peer can report the rejection reason instead of seeing an + // empty/undecodable handshake response. + local = TransferHandshakeUtil::encode(local_desc); + return 0; + } } local = TransferHandshakeUtil::encode(local_desc); return 0; diff --git a/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp b/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp index 4df028d239..10e99e7e64 100644 --- a/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp +++ b/mooncake-transfer-engine/src/transfer_metadata_plugin.cpp @@ -746,7 +746,7 @@ struct SocketHandShakePlugin : public HandShakePlugin { } struct timeval timeout; - timeout.tv_sec = 60; + timeout.tv_sec = 5; timeout.tv_usec = 0; if (setsockopt(conn_fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout))) { @@ -762,6 +762,11 @@ struct SocketHandShakePlugin : public HandShakePlugin { Json::Value local, peer; auto [type, json_str] = readString(conn_fd); + if (type == HandShakeRequestType::Invalid) { + close(conn_fd); + continue; + } + std::string errs; if (!parseJsonString(json_str, peer, &errs)) { LOG(ERROR) diff --git a/mooncake-transfer-engine/src/transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/CMakeLists.txt index 47a1dc99cf..43baa24cc8 100644 --- a/mooncake-transfer-engine/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/CMakeLists.txt @@ -2,86 +2,106 @@ file(GLOB XPORT_SOURCES "*.cpp") add_subdirectory(rdma_transport) add_subdirectory(rpc_communicator) -add_library(transport OBJECT ${XPORT_SOURCES} $ $) -target_link_libraries(transport PRIVATE JsonCpp::JsonCpp yalantinglibs::yalantinglibs glog::glog pthread) +add_library(transport OBJECT ${XPORT_SOURCES} $ + $) +target_link_libraries( + transport PRIVATE JsonCpp::JsonCpp yalantinglibs::yalantinglibs glog::glog + pthread) if(USE_UB) add_subdirectory(kunpeng_transport) target_sources(transport PUBLIC $) endif() -if (USE_TCP) +if(USE_TCP) add_subdirectory(tcp_transport) target_sources(transport PUBLIC $) endif() -if (USE_BAREX) +if(USE_BAREX) add_subdirectory(barex_transport) target_sources(transport PUBLIC $) endif() -if (USE_NVMEOF) +if(USE_NVMEOF) add_subdirectory(nvmeof_transport) target_sources(transport PUBLIC $) endif() -if (USE_CXL) +if(USE_NCCL_HOST) + add_subdirectory(nccl_transport) + target_sources(transport PUBLIC $) +endif() + +if(USE_CXL) add_subdirectory(cxl_transport) target_sources(transport PUBLIC $) endif() -if (USE_ASCEND_DIRECT) +if(USE_ASCEND_DIRECT) add_subdirectory(ascend_transport) elseif(USE_ASCEND) add_subdirectory(ascend_transport) target_sources(transport PUBLIC $) endif() -if (USE_ASCEND_HETEROGENEOUS) +if(USE_ASCEND_HETEROGENEOUS) add_subdirectory(ascend_transport) target_sources(transport PUBLIC $) endif() -if (USE_HIP) +if(USE_HIP) add_subdirectory(hip_transport) target_sources(transport PUBLIC $) endif() -if (USE_MACA) +if(USE_MACA) add_subdirectory(maca_transport) target_sources(transport PUBLIC $) endif() -if (USE_SUNRISE) +if(USE_SUNRISE) add_subdirectory(sunrise_link) target_sources(transport PUBLIC $) endif() -if (USE_MNNVL AND NOT USE_HIP) +if(USE_MNNVL AND NOT USE_HIP) add_subdirectory(nvlink_transport) target_sources(transport PUBLIC $) endif() -if (USE_INTRA_NVLINK) - add_subdirectory(intranode_nvlink_transport) - target_sources(transport PUBLIC $) +if(USE_INTRA_NVLINK) + add_subdirectory(intranode_nvlink_transport) + target_sources(transport PUBLIC $) endif() -if (USE_UBSHMEM) +if(USE_UBSHMEM) add_subdirectory(ascend_transport) target_sources(transport PUBLIC $) endif() -if (USE_EFA) +if(USE_EFA) add_subdirectory(efa_transport) target_sources(transport PUBLIC $) target_link_libraries(transport PRIVATE fabric) endif() -if(USE_CUDA OR USE_MUSA) +if(USE_CXI) + add_subdirectory(cxi_transport) + target_sources(transport PUBLIC $) + target_link_libraries(transport PRIVATE fabric) +endif() + +if(USE_CUDA OR USE_MUSA OR USE_MACA) add_subdirectory(device) target_sources(transport PUBLIC $) # device_transport (ibgda_device_transport.cpp / mlx5gda.cpp) calls libmlx5 # DevX symbols (mlx5dv_devx_*, mlx5dv_init_obj) directly. - target_link_libraries(transport PUBLIC mlx5) + if((USE_CUDA OR USE_MUSA) AND NOT USE_CXI) + target_link_libraries(transport PUBLIC mlx5) + endif() + if(USE_CUDA AND MOONCAKE_HAVE_MLX5_DMABUF_UMEM) + target_compile_definitions(device_transport PRIVATE _GNU_SOURCE) + target_link_libraries(transport PUBLIC ${CMAKE_DL_LIBS}) + endif() endif() diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_allocator.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_allocator.cpp index 1071999a2b..76adc636e9 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_allocator.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_allocator.cpp @@ -4,13 +4,37 @@ #include #include +#include "ascend_allocator.h" #include "config.h" #include "acl/acl.h" #include "transport/ascend_transport/ascend_direct_transport/adxl_compat.h" namespace mooncake { namespace { -constexpr size_t kFabricMemPageSize = 1024 * 1024 * 1024; // 1G +constexpr size_t kFabricMemPageSize = 1024ULL * 1024 * 1024; // 1G +constexpr int kBestEffortStartPercent = 100; +constexpr int kBestEffortMinPercent = 50; +constexpr int kBestEffortPercentStep = 10; + +size_t align_down_1g(size_t n) { + return (n / kFabricMemPageSize) * kFabricMemPageSize; +} + +struct StoreMemRange { + void *base; + size_t size; +}; +std::mutex g_store_mem_mutex; +std::vector g_store_mem_ranges; + +void remove_store_memory_range(void *ptr) { + std::lock_guard store_lock(g_store_mem_mutex); + auto it = std::remove_if( + g_store_mem_ranges.begin(), g_store_mem_ranges.end(), + [ptr](const StoreMemRange &range) { return range.base == ptr; }); + g_store_mem_ranges.erase(it, g_store_mem_ranges.end()); +} + #ifdef ASCEND_SUPPORT_FABRIC_MEM struct AllocRecord { aclrtDrvMemHandle handle; @@ -19,17 +43,22 @@ struct AllocRecord { std::mutex g_vmm_alloc_mutex; std::unordered_map g_vmm_alloc_records; -int allocate_physical_memory(size_t total_size, aclrtDrvMemHandle &handle) { +int allocate_physical_memory(size_t total_size, aclrtDrvMemHandle &handle, + bool quiet) { int32_t user_dev_id; auto ret = aclrtGetDevice(&user_dev_id); if (ret != ACL_ERROR_NONE) { - LOG(ERROR) << "Failed to get device: " << ret; + if (!quiet) { + LOG(ERROR) << "Failed to get device: " << ret; + } return -1; } int32_t physical_dev_id; ret = aclrtGetPhyDevIdByLogicDevId(user_dev_id, &physical_dev_id); if (ret != ACL_ERROR_NONE) { - LOG(ERROR) << "Failed to get physical dev id: " << ret; + if (!quiet) { + LOG(ERROR) << "Failed to get physical dev id: " << ret; + } return -1; } aclrtPhysicalMemProp prop = {}; @@ -42,20 +71,28 @@ int allocate_physical_memory(size_t total_size, aclrtDrvMemHandle &handle) { const int32_t kNumaNodeStep = 2; prop.location.id = (physical_dev_id / kDevicesPerChip) * kNumaNodeStep; prop.reserve = 0; - LOG(INFO) << "Malloc host memory for numa:" << prop.location.id; + if (!quiet) { + LOG(INFO) << "Malloc host memory for numa:" << prop.location.id; + } ret = aclrtMallocPhysical(&handle, total_size, &prop, 0); if (ret != ACL_ERROR_NONE) { - LOG(INFO) << "Malloc host memory for numa:" << prop.location.id - << " failed, try common allocate instead."; + if (!quiet) { + LOG(INFO) << "Malloc host memory for numa:" << prop.location.id + << " failed, try common allocate instead."; + } prop.location.type = ACL_MEM_LOCATION_TYPE_HOST; prop.location.id = 0; ret = aclrtMallocPhysical(&handle, total_size, &prop, 0); if (ret != ACL_ERROR_NONE) { - LOG(INFO) << "Malloc failed, try smaller page instead."; + if (!quiet) { + LOG(INFO) << "Malloc failed, try smaller page instead."; + } prop.memAttr = ACL_MEM_P2P_HUGE; ret = aclrtMallocPhysical(&handle, total_size, &prop, 0); if (ret != ACL_ERROR_NONE) { - LOG(ERROR) << "Failed to allocate memory: " << ret; + if (!quiet) { + LOG(ERROR) << "Failed to allocate memory: " << ret; + } return -1; } } @@ -65,21 +102,25 @@ int allocate_physical_memory(size_t total_size, aclrtDrvMemHandle &handle) { // Direct ACL VMM allocation (always bypasses adxl MallocMem). // Used by shm_helper when ascend_agent_mode && ascend_use_fabric_mem. -void *allocate_vmm_memory_direct_impl(size_t total_size) { +void *allocate_vmm_memory_direct_impl(size_t total_size, bool quiet = false) { aclrtDrvMemHandle handle = nullptr; - if (allocate_physical_memory(total_size, handle) != 0) { + if (allocate_physical_memory(total_size, handle, quiet) != 0) { return nullptr; } void *va = nullptr; auto ret = aclrtReserveMemAddress(&va, total_size, 0, nullptr, 1); if (ret != ACL_ERROR_NONE) { - LOG(ERROR) << "Failed to reserve memory: " << ret; + if (!quiet) { + LOG(ERROR) << "Failed to reserve memory: " << ret; + } (void)aclrtFreePhysical(handle); return nullptr; } ret = aclrtMapMem(va, total_size, 0, handle, 0); if (ret != ACL_ERROR_NONE) { - LOG(ERROR) << "Failed to map memory: " << ret; + if (!quiet) { + LOG(ERROR) << "Failed to map memory: " << ret; + } (void)aclrtReleaseMemAddress(va); (void)aclrtFreePhysical(handle); return nullptr; @@ -88,22 +129,37 @@ void *allocate_vmm_memory_direct_impl(size_t total_size) { g_vmm_alloc_records.emplace(va, AllocRecord{handle, true}); return va; } -#endif -struct StoreMemRange { - void *base; - size_t size; -}; -std::mutex g_store_mem_mutex; -std::vector g_store_mem_ranges; - -void remove_store_memory_range(void *ptr) { - std::lock_guard store_lock(g_store_mem_mutex); - auto it = std::remove_if( - g_store_mem_ranges.begin(), g_store_mem_ranges.end(), - [ptr](const StoreMemRange &range) { return range.base == ptr; }); - g_store_mem_ranges.erase(it, g_store_mem_ranges.end()); +// Exact-size fabric alloc: use adxl when the weak symbol is present; otherwise +// direct VMM. Matches original ascend_allocate_memory behavior — adxl failure +// does not fall back to direct VMM. When quiet=true, skip failure ERROR logs +// (used by best-effort percentile probes). +void *allocate_fabric_exact(size_t total_size, bool quiet = false) { + void *va = nullptr; + if (&adxl::AdxlEngine::MallocMem != nullptr) { + auto status = adxl::AdxlEngine::MallocMem(adxl::MemType::MEM_HOST, + total_size, &va); + if (status != adxl::SUCCESS) { + if (!quiet) { + LOG(ERROR) << "Failed to allocate fabric memory, errmsg: " + << aclGetRecentErrMsg(); + } + return nullptr; + } + LOG(INFO) << "Call adxl MallocMem suc, va:" << va + << ", size:" << total_size; + std::lock_guard lock(g_vmm_alloc_mutex); + g_vmm_alloc_records.emplace(va, AllocRecord{nullptr, false}); + return va; + } + va = allocate_vmm_memory_direct_impl(total_size, quiet); + if (va) { + std::lock_guard store_lock(g_store_mem_mutex); + g_store_mem_ranges.push_back({va, total_size}); + } + return va; } +#endif } // namespace void *ascend_allocate_vmm_memory_direct(size_t total_size) { @@ -131,31 +187,13 @@ aclrtDrvMemHandle ascend_get_physical_handle_from_va(void *va) { void *ascend_allocate_memory(size_t total_size, const std::string &protocol) { if (globalConfig().ascend_use_fabric_mem) { - void *va = nullptr; #ifdef ASCEND_SUPPORT_FABRIC_MEM - if (&adxl::AdxlEngine::MallocMem != nullptr) { - // Try to use adxl_engine's MallocMem - auto status = adxl::AdxlEngine::MallocMem(adxl::MemType::MEM_HOST, - total_size, &va); - if (status != adxl::SUCCESS) { - LOG(ERROR) << "Failed to allocate fabric memory, errmsg: " - << aclGetRecentErrMsg(); - return nullptr; - } - LOG(INFO) << "Call adxl MallocMem suc, va:" << va; - std::lock_guard lock(g_vmm_alloc_mutex); - g_vmm_alloc_records.emplace(va, AllocRecord{nullptr, false}); - return va; - } - va = allocate_vmm_memory_direct_impl(total_size); - if (va) { - std::lock_guard store_lock(g_store_mem_mutex); - g_store_mem_ranges.push_back({va, total_size}); - } - return va; + return allocate_fabric_exact(total_size); #else + (void)total_size; LOG(ERROR) << "Fabric mem mode is not supported, please upgrade Ascend " "HDK and CANN."; + return nullptr; #endif } if (protocol == "ubshmem") { @@ -178,6 +216,64 @@ void *ascend_allocate_memory(size_t total_size, const std::string &protocol) { return buffer; } +void *ascend_allocate_memory_best_effort(size_t target_size, + const std::string &protocol, + size_t *actual_size) { + if (actual_size == nullptr) { + LOG(ERROR) << "ascend_allocate_memory_best_effort: actual_size is null"; + return nullptr; + } + *actual_size = 0; + + if (!globalConfig().ascend_use_fabric_mem) { + void *ptr = ascend_allocate_memory(target_size, protocol); + if (ptr) { + *actual_size = target_size; + } + return ptr; + } + +#ifdef ASCEND_SUPPORT_FABRIC_MEM + (void)protocol; + // Try 100%, 90%, ... down to 50% of configured size (1G-aligned). + // Skip candidates that fall below the unaligned 50% floor after round-down + // (e.g. 2.1 GiB target → 50% is 1.05 GiB → 1 GiB must not be accepted). + // Probe attempts are quiet; only the final failure logs ERROR. + const size_t min_size = + target_size * static_cast(kBestEffortMinPercent) / 100; + size_t last_tried = 0; + for (int pct = kBestEffortStartPercent; pct >= kBestEffortMinPercent; + pct -= kBestEffortPercentStep) { + size_t sz = + (pct == kBestEffortStartPercent) + ? align_down_1g(target_size) + : align_down_1g(target_size * static_cast(pct) / 100); + if (sz == 0 || sz < min_size || sz == last_tried) { + continue; + } + last_tried = sz; + void *ptr = allocate_fabric_exact(sz, /*quiet=*/true); + if (ptr != nullptr) { + *actual_size = sz; + if (sz < target_size) { + LOG(WARNING) << "Fabric mem best-effort: target=" << target_size + << ", actual=" << sz << " (" << pct << "%)"; + } + return ptr; + } + } + + LOG(ERROR) << "Fabric mem best-effort failed: cannot allocate at least " + << kBestEffortMinPercent << "% of target=" << target_size; + return nullptr; +#else + (void)protocol; + LOG(ERROR) << "Fabric mem mode is not supported, please upgrade Ascend " + "HDK and CANN."; + return nullptr; +#endif +} + bool ascend_is_store_memory(void *addr, size_t length) { if (!addr || length == 0) return false; auto addr_start = reinterpret_cast(addr); diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp index f5869b34e1..365505cbb0 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp @@ -343,8 +343,8 @@ int AscendDirectTransport::registerLocalMemory(void *addr, size_t length, } else if (attributes.location.type == ACL_MEM_LOCATION_TYPE_DEVICE) { mem_type = adxl::MEM_DEVICE; } else { - LOG(ERROR) << "mem addr:" << addr - << " can not be recognized, try set to host mem."; + LOG(INFO) << "mem addr:" << addr + << " can not be recognized, try set to host mem."; mem_type = adxl::MEM_HOST; } } else { @@ -420,16 +420,18 @@ int AscendDirectTransport::unregisterLocalMemoryBatch( "with addr count: " << addr_list.size(); + int first_error = 0; for (void *addr : addr_list) { int ret = unregisterLocalMemory(addr, false); if (ret != 0) { LOG(ERROR) << "Failed to unregister memory in batch, addr: " << addr; - return ret; + if (!first_error) first_error = ret; } } // Update metadata once for the entire batch - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } } // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp index 535023ebfb..3455d12170 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/async_transfer_executor.cpp @@ -139,8 +139,10 @@ TransferExecutorBase::ExecuteResult AsyncTransferExecutor::execute( } async_task_cv_.notify_one(); } - disconnect(local_engine_idx, target_adxl_engine_name, - kDefaultDisconnectTime); + if (!params_.auto_connect) { + disconnect(local_engine_idx, target_adxl_engine_name, + kDefaultDisconnectTime); + } return {.ret = -1, .status = status, .retryable = true}; } @@ -172,7 +174,8 @@ void AsyncTransferExecutor::queryThreadLoop() { : poll_result.fail_reason; failAllPendingOnRoute(batch.engine_idx, batch.target_adxl_engine_name, - pending_batches, reason); + pending_batches, reason, + poll_result.adxl_already_disconnected); // failAllPendingOnRoute swap-pops may move unvisited batches to // indices < i; restart the scan so none are skipped this cycle. i = 0; @@ -220,9 +223,10 @@ void AsyncTransferExecutor::markBatchFailed(QueryBatch& batch, void AsyncTransferExecutor::failAllPendingOnRoute( size_t engine_idx, const std::string& target, - std::vector& pending, const std::string& reason) { + std::vector& pending, const std::string& reason, + bool adxl_already_disconnected) { LOG(ERROR) << reason; - bool disconnected = false; + bool cleaned_up = false; for (size_t j = 0; j < pending.size();) { auto& batch = pending[j]; if (batch.engine_idx != engine_idx || @@ -232,10 +236,14 @@ void AsyncTransferExecutor::failAllPendingOnRoute( } markBatchFailed(batch, reason, false); - if (!disconnected && engine_idx < adxl_engines_.size() && + if (!cleaned_up && engine_idx < adxl_engines_.size() && !target.empty()) { - disconnect(engine_idx, target, params_.connect_timeout); - disconnected = true; + if (params_.auto_connect && adxl_already_disconnected) { + forgetConnectedSegment(engine_idx, target); + } else { + disconnect(engine_idx, target, params_.connect_timeout); + } + cleaned_up = true; } std::swap(pending[j], pending.back()); @@ -284,6 +292,7 @@ void AsyncTransferExecutor::processOneBatch(QueryBatch& batch, if (ret != adxl::SUCCESS || task_status == adxl::TransferStatus::FAILED) { result.done = true; result.fail_entire_route = true; + result.adxl_already_disconnected = true; result.fail_reason = "Get transfer status failed, ret: " + std::to_string(static_cast(ret)) + ", errmsg: " + aclGetRecentErrMsg(); diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/sync_transfer_executor.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/sync_transfer_executor.cpp index 8e19431ccd..434235a270 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/sync_transfer_executor.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/sync_transfer_executor.cpp @@ -76,8 +76,10 @@ TransferExecutorBase::ExecuteResult SyncTransferExecutor::execute( op_descs, params_.transfer_timeout); if (status != adxl::SUCCESS) { - disconnect(local_engine_idx, target_adxl_engine_name, - kDefaultDisconnectTime); + if (!params_.auto_connect) { + disconnect(local_engine_idx, target_adxl_engine_name, + kDefaultDisconnectTime); + } return {.ret = -1, .status = status, .retryable = true}; } diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp index 9cef2fa85d..6f08a3257c 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/transfer_executor_base.cpp @@ -290,6 +290,19 @@ void TransferExecutorBase::recordConnectedSegment(size_t engine_idx, connected_segments_[engine_idx].insert(remote); } +void TransferExecutorBase::forgetConnectedSegment(size_t engine_idx, + const std::string& remote) { + std::lock_guard lock(connection_mutex_); + auto it = connected_segments_.find(engine_idx); + if (it == connected_segments_.end()) { + return; + } + it->second.erase(remote); + if (it->second.empty()) { + connected_segments_.erase(it); + } +} + int TransferExecutorBase::disconnect(size_t engine_idx, const std::string& target_adxl_engine_name, int32_t timeout_in_millis) { diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/utils.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/utils.cpp index 014ad3e734..e1e623722e 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/utils.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/utils.cpp @@ -17,7 +17,9 @@ #include #include +#include #include +#include #include #if __has_include() @@ -163,7 +165,9 @@ bool IsRoceModeEnabled() { // AscendThreadPool implementation AscendThreadPool::AscendThreadPool(size_t num_threads) : running_(true) { for (size_t i = 0; i < num_threads; ++i) { - workers_.emplace_back([this] { + workers_.emplace_back([this, i] { + auto thread_name = "ascend-wk-" + std::to_string(i); + pthread_setname_np(pthread_self(), thread_name.c_str()); while (true) { std::function task; { diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp index a1c6958502..ef81cb6f30 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/hccl_transport/hccl_transport.cpp @@ -607,15 +607,23 @@ int HcclTransport::allocateLocalSegmentID() { int HcclTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, -1); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int HcclTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, -1); - return metadata_->updateLocalSegmentDesc(); + int first_error = 0; + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret && !first_error) first_error = ret; + } + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } } // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp index 943332a919..1f0962d245 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp @@ -208,16 +208,18 @@ int HeterogeneousRdmaTransport::registerLocalMemoryBatch( int HeterogeneousRdmaTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { + int first_error = 0; for (auto &addr : addr_list) { int ret = unregisterLocalMemory(addr, false); if (ret) { LOG(ERROR) << "HeterogeneousRdmaTransport " "unregisterLocalMemoryBatch error, ret: " << ret; - return ret; + if (!first_error) first_error = ret; } } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } int HeterogeneousRdmaTransport::checkAndCreateStreamCopy() { diff --git a/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp b/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp index 484c698335..1125e5d466 100644 --- a/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp +++ b/mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp @@ -722,7 +722,7 @@ int UBShmemTransport::registerLocalMemoryBatch( for (auto &buffer : buffer_list) { int rc = registerLocalMemory(buffer.addr, buffer.length, location, true, false); - if (rc < 0) { + if (rc) { return rc; } } @@ -731,13 +731,13 @@ int UBShmemTransport::registerLocalMemoryBatch( int UBShmemTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { + int first_error = 0; for (auto &addr : addr_list) { int rc = unregisterLocalMemory(addr, false); - if (rc < 0) { - return rc; - } + if (rc && !first_error) first_error = rc; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void *UBShmemTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp b/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp index e575dfc39d..567bef44af 100644 --- a/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp +++ b/mooncake-transfer-engine/src/transport/barex_transport/barex_transport.cpp @@ -306,7 +306,7 @@ int BarexTransport::registerLocalMemoryBatch( if (ret) { LOG(ERROR) << "BarexTransport: Failed to register memory: addr " << buffer.addr << " length " << buffer.length; - return ERR_ADDRESS_NOT_REGISTERED; + return ret; } } @@ -323,13 +323,17 @@ int BarexTransport::unregisterLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < addr_list.size(); ++i) { - if (results[i].get()) + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "BarexTransport: Failed to unregister memory: addr " << addr_list[i]; + if (!first_error) first_error = ret; + } } - - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status BarexTransport::submitTransfer( @@ -1169,8 +1173,13 @@ Status BarexTransport::submitTransferTask( } if (!found_device) { auto source_addr = slice->source_addr; - for (auto &entry : slices_to_post) - for (auto s : entry.second) getSliceCache().deallocate(s); + // Do not deallocate slices already queued in slices_to_post + // here: every slice is also recorded in its owning + // TransferTask::slice_list right after allocation, and + // ~TransferTask() returns everything in slice_list to the + // cache exactly once. Deallocating here double-frees them + // into ThreadLocalSliceCache, letting a later allocate() + // hand the same Slice* to two unrelated transfers. LOG(ERROR) << "Memory region not registered by any active device(s): " << source_addr; diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/cxi_transport/CMakeLists.txt new file mode 100644 index 0000000000..51d4197776 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/cxi_transport/CMakeLists.txt @@ -0,0 +1,14 @@ +file(GLOB CXI_SOURCES "*.cpp") +add_library(cxi_transport OBJECT ${CXI_SOURCES}) + +# Link against libfabric (fabric) instead of ibverbs for AWS EFA +target_link_libraries(cxi_transport PRIVATE fabric glog::glog) + +target_include_directories( + cxi_transport + PRIVATE + ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include + ${CMAKE_SOURCE_DIR}/mooncake-transfer-engine/include/transport/cxi_transport + ${CMAKE_SOURCE_DIR}/mooncake-common/include ${LIBFABRIC_INCLUDE_DIR}) + +# libfabric library path is set globally via common.cmake (LIBFABRIC_LIB_DIR) diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_context.cpp b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_context.cpp new file mode 100644 index 0000000000..ea1aeee3c4 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_context.cpp @@ -0,0 +1,981 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "transport/cxi_transport/cxi_transport.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "config.h" +#include "cuda_alike.h" +#include "transport/cxi_transport/cxi_endpoint.h" +#include "transport/cxi_transport/cxi_transport.h" +#include "transport/transport.h" + +namespace mooncake { + +// CxiContext implementation +CxiContext::CxiContext(CxiTransport& engine, const std::string& device_name) + : engine_(engine), + device_name_(device_name), + fi_info_(nullptr), + hints_(nullptr), + fabric_(nullptr), + domain_(nullptr), + av_(nullptr), + active_(true), + shared_ep_(nullptr), + wr_depth_(0), + max_wr_depth_(0), + post_lock_(ATOMIC_FLAG_INIT) {} + +CxiContext::~CxiContext() { + if (fabric_) { + deconstruct(); + } +} + +int CxiContext::construct(size_t num_cq_list, size_t max_cqe, + int max_endpoints) { + // Setup hints for CXI provider + hints_ = fi_allocinfo(); + if (!hints_) { + LOG(ERROR) << "Failed to allocate fi_info hints"; + return ERR_CONTEXT; + } + + hints_->caps = + FI_MSG | FI_RMA | FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE +#if defined(USE_CUDA) || defined(USE_HIP) + | FI_HMEM // must be added to the fabric caps to allow cuda/hip on cxi +#endif + ; + hints_->mode = FI_CONTEXT; + hints_->ep_attr->type = FI_EP_RDM; // CXI uses RDM endpoints + hints_->fabric_attr->prov_name = strdup("cxi"); + + // Specify the domain (device) name - append "-rdm" for RDM endpoint + std::string domain_name = device_name_; + // IMPORTANT: FI_MR_ENDPOINT must be specified when using CXI, otherwise + // fi_getinfo returns ENODATA CXI does not support FI_MR_VIRT_ADDR, so + // addresses must be offsets compared to MR base address + hints_->domain_attr->mr_mode = FI_MR_ENDPOINT | FI_MR_LOCAL | + FI_MR_ALLOCATED | FI_MR_PROV_KEY +#if defined(USE_CUDA) || defined(USE_HIP) + | FI_MR_HMEM +#endif + ; + hints_->domain_attr->name = strdup(domain_name.c_str()); + hints_->domain_attr->threading = FI_THREAD_SAFE; + + // Get fabric info + int ret = fi_getinfo(fi_version(), nullptr, nullptr, 0, hints_, &fi_info_); + if (ret) { + LOG(ERROR) << "fi_getinfo failed for device " << device_name_ << ": " + << fi_strerror(-ret); + fi_freeinfo(hints_); + hints_ = nullptr; + return ERR_CONTEXT; + } + + // Open fabric + ret = fi_fabric(fi_info_->fabric_attr, &fabric_, nullptr); + if (ret) { + LOG(ERROR) << "fi_fabric failed: " << fi_strerror(-ret); + fi_freeinfo(fi_info_); + fi_freeinfo(hints_); + fi_info_ = nullptr; + hints_ = nullptr; + return ERR_CONTEXT; + } + + // Open domain + ret = fi_domain(fabric_, fi_info_, &domain_, nullptr); + if (ret) { + LOG(ERROR) << "fi_domain failed: " << fi_strerror(-ret); + fi_close(&fabric_->fid); + fi_freeinfo(fi_info_); + fi_freeinfo(hints_); + fabric_ = nullptr; + fi_info_ = nullptr; + hints_ = nullptr; + return ERR_CONTEXT; + } + + // Create address vector + struct fi_av_attr av_attr = {}; + av_attr.type = FI_AV_TABLE; + av_attr.count = max_endpoints; + + ret = fi_av_open(domain_, &av_attr, &av_, nullptr); + if (ret) { + LOG(ERROR) << "fi_av_open failed: " << fi_strerror(-ret); + fi_close(&domain_->fid); + fi_close(&fabric_->fid); + fi_freeinfo(fi_info_); + fi_freeinfo(hints_); + domain_ = nullptr; + fabric_ = nullptr; + fi_info_ = nullptr; + hints_ = nullptr; + return ERR_CONTEXT; + } + + // Create completion queues + cq_list_.resize(num_cq_list); + for (size_t i = 0; i < num_cq_list; ++i) { + auto cq = std::make_shared(); + + struct fi_cq_attr cq_attr = {}; + cq_attr.size = max_cqe; + cq_attr.format = FI_CQ_FORMAT_DATA; + cq_attr.wait_obj = FI_WAIT_NONE; + + ret = fi_cq_open(domain_, &cq_attr, &cq->cq, nullptr); + if (ret) { + LOG(ERROR) << "fi_cq_open failed: " << fi_strerror(-ret); + return ERR_CONTEXT; + } + cq_list_[i] = cq; + } + + ret = buildSharedEndpoint(globalConfig().max_wr, 16); + if (ret) { + LOG(ERROR) << "CxiContext::construct: buildSharedEndpoint failed for " + << device_name_; + return ret; + } + + LOG(INFO) << "CXI device (libfabric): " << device_name_ + << ", domain: " << fi_info_->domain_attr->name + << ", provider: " << fi_info_->fabric_attr->prov_name; + + return 0; +} + +int CxiContext::buildSharedEndpoint(size_t max_wr, size_t max_inline) { + (void)max_inline; + if (shared_ep_) return 0; + + shared_cq_ = cq_list_.empty() ? nullptr : cq_list_[0]; + if (!shared_cq_) { + LOG(ERROR) << "CxiContext::buildSharedEndpoint: no CQ available"; + return ERR_CONTEXT; + } + max_wr_depth_ = static_cast(max_wr); + + int ret = fi_endpoint(domain_, fi_info_, &shared_ep_, nullptr); + if (ret) { + LOG(ERROR) << "fi_endpoint failed: " << fi_strerror(-ret); + shared_ep_ = nullptr; + return ERR_ENDPOINT; + } + + ret = fi_ep_bind(shared_ep_, &av_->fid, 0); + if (ret) { + LOG(ERROR) << "fi_ep_bind(av) failed: " << fi_strerror(-ret); + fi_close(&shared_ep_->fid); + shared_ep_ = nullptr; + return ERR_ENDPOINT; + } + + ret = fi_ep_bind(shared_ep_, &shared_cq_->cq->fid, FI_TRANSMIT); + if (ret) { + LOG(ERROR) << "fi_ep_bind(tx_cq) failed: " << fi_strerror(-ret); + fi_close(&shared_ep_->fid); + shared_ep_ = nullptr; + return ERR_ENDPOINT; + } + + ret = fi_ep_bind(shared_ep_, &shared_cq_->cq->fid, FI_RECV); + if (ret) { + LOG(ERROR) << "fi_ep_bind(rx_cq) failed: " << fi_strerror(-ret); + fi_close(&shared_ep_->fid); + shared_ep_ = nullptr; + return ERR_ENDPOINT; + } + + ret = fi_enable(shared_ep_); + if (ret) { + LOG(ERROR) << "fi_enable failed: " << fi_strerror(-ret); + fi_close(&shared_ep_->fid); + shared_ep_ = nullptr; + return ERR_ENDPOINT; + } + + // Cache our own libfabric address for handshake advertisement. + size_t addr_len = 8; + local_ep_addr_.assign(addr_len, 0); + ret = fi_getname(&shared_ep_->fid, local_ep_addr_.data(), &addr_len); + if (ret) { + LOG(ERROR) << "fi_getname failed: " << fi_strerror(-ret); + fi_close(&shared_ep_->fid); + shared_ep_ = nullptr; + local_ep_addr_.clear(); + return ERR_ENDPOINT; + } + LOG(INFO) << "device " << deviceName() + << " ptr address: " << &shared_ep_->fid; + local_ep_addr_.resize(addr_len); + return 0; +} + +int CxiContext::deconstruct() { + // Destroy all endpoints before closing domain/fabric/AV. + // Endpoints hold fi_ep handles that reference the domain, so they must + // be closed first. + if (shared_ep_) { + fi_close(&shared_ep_->fid); + shared_ep_ = nullptr; + } + + { + RWSpinlock::WriteGuard guard(peer_map_lock_); + for (auto& entry : peer_map_) { + if (entry.second) entry.second->markDetachedForTeardown(); + } + peer_map_.clear(); + } + + { + RWSpinlock::WriteGuard guard(mr_lock_); + for (auto& entry : mr_map_) { + if (entry.second.mr) { + fi_close(&entry.second.mr->fid); + } + } + mr_map_.clear(); + } + + for (auto& cq : cq_list_) { + if (cq && cq->cq) { + fi_close(&cq->cq->fid); + cq->cq = nullptr; + } + } + cq_list_.clear(); + shared_cq_.reset(); + + if (av_) { + fi_close(&av_->fid); + av_ = nullptr; + } + + if (domain_) { + fi_close(&domain_->fid); + domain_ = nullptr; + } + + if (fabric_) { + fi_close(&fabric_->fid); + fabric_ = nullptr; + } + + if (fi_info_) { + fi_freeinfo(fi_info_); + fi_info_ = nullptr; + } + + if (hints_) { + fi_freeinfo(hints_); + hints_ = nullptr; + } + + return 0; +} + +int CxiContext::registerMemoryRegionInternal(void* addr, size_t length, + int access, + CxiMemoryRegionMeta& mrMeta) { + (void)access; + if (length > (size_t)globalConfig().max_mr_size) { + LOG(ERROR) << "Buffer length " << length + << " exceeds device max_mr_size " + << globalConfig().max_mr_size + << ". Use CxiTransport::registerLocalMemory() which " + "auto-splits large buffers."; + return ERR_CONTEXT; + } else if (!shared_ep_) { + LOG(ERROR) + << "shared ep must be built before registering memory regions!"; + } + + mrMeta.addr = addr; + mrMeta.length = length; + + uint64_t fi_access = FI_READ | FI_WRITE | FI_REMOTE_READ | FI_REMOTE_WRITE; + + enum fi_hmem_iface iface = FI_HMEM_SYSTEM; + int device_ordinal = 0; + int current_device = 0; + +#if defined(USE_CUDA) + cudaPointerAttributes attributes; + cudaError_t cuda_ret = cudaPointerGetAttributes(&attributes, addr); + if (cuda_ret == cudaSuccess && attributes.type == cudaMemoryTypeDevice) { + iface = FI_HMEM_CUDA; + device_ordinal = attributes.device; + cudaGetDevice(¤t_device); + cudaSetDevice(device_ordinal); // if by any chance we are setting a + // device pointer from a different + // device, fi_mr_regattr will fail + } +#elif defined(USE_HIP) + hipPointerAttribute_t attributes; + hipError_t hip_ret = hipPointerGetAttributes(&attributes, addr); + if (hip_ret == hipSuccess && attributes.type == hipMemoryTypeDevice) { + iface = FI_HMEM_ROCR; + device_ordinal = attributes.device; + hipGetDevice(¤t_device); + hipSetDevice(device_ordinal); + } +#endif + + int ret; + if (iface != FI_HMEM_SYSTEM) { + // GPU memory: use fi_mr_regattr with explicit iface and device + struct iovec iov = {.iov_base = addr, .iov_len = length}; + struct fi_mr_attr attr = {}; + attr.mr_iov = &iov; + attr.iov_count = 1; + attr.access = fi_access; + attr.iface = iface; + attr.device.cuda = device_ordinal; + + ret = fi_mr_regattr(domain_, &attr, 0, &mrMeta.mr); + if (ret) { + LOG(ERROR) << "fi_mr_regattr failed for GPU memory " << addr + << " (device " << device_ordinal + << "): " << fi_strerror(-ret); +#if defined(USE_CUDA) + cudaSetDevice(current_device); +#elif defined(USE_HIP) + hipSetDevice(current_device); +#endif + return ERR_CONTEXT; + } +#if defined(USE_CUDA) + cudaSetDevice(current_device); +#elif defined(USE_HIP) + hipSetDevice(current_device); +#endif + + } else { + // CPU memory: fi_mr_reg is sufficient + ret = fi_mr_reg(domain_, addr, length, fi_access, 0, 0, 0, &mrMeta.mr, + nullptr); + if (ret) { + LOG(ERROR) << "fi_mr_reg failed for " << addr << ": " + << fi_strerror(-ret); + return ERR_CONTEXT; + } + } + + // cxi requires FI_MR_ENDPOINT + ret = fi_mr_bind(mrMeta.mr, &shared_ep_->fid, 0); + if (ret) { + LOG(ERROR) << "fi_mr_bind failed for " << addr << ": " + << fi_strerror(-ret); + return ERR_CONTEXT; + } + ret = fi_mr_enable(mrMeta.mr); + if (ret) { + LOG(ERROR) << "fi_mr_enable failed for " << addr << ": " + << fi_strerror(-ret); + return ERR_CONTEXT; + } + mrMeta.key = fi_mr_key(mrMeta.mr); + return 0; +} + +int CxiContext::registerMemoryRegion(void* addr, size_t length, int access) { + CxiMemoryRegionMeta mrMeta; + int ret = registerMemoryRegionInternal(addr, length, access, mrMeta); + if (ret != 0) { + return ret; + } + { + RWSpinlock::WriteGuard guard(mr_lock_); + mr_map_[(uint64_t)addr] = mrMeta; + } + return 0; +} + +int CxiContext::unregisterMemoryRegion(void* addr) { + RWSpinlock::WriteGuard guard(mr_lock_); + auto it = mr_map_.find((uint64_t)addr); + if (it == mr_map_.end()) { + return 0; + } + if (it->second.mr) { + fi_close(&it->second.mr->fid); + } + mr_map_.erase(it); + return 0; +} + +int CxiContext::preTouchMemory(void* addr, size_t length) { + volatile char* ptr = (volatile char*)addr; + for (size_t i = 0; i < length; i += 4096) { + ptr[i] = ptr[i]; + } + return 0; +} + +uint64_t CxiContext::rkey(void* addr) { + RWSpinlock::ReadGuard guard(mr_lock_); + auto it = mr_map_.upper_bound((uint64_t)addr); + if (it != mr_map_.begin()) { + --it; + if ((uint64_t)addr < it->first + it->second.length && it->second.mr) { + return it->second.key; + } + } + return 0; +} + +uint64_t CxiContext::lkey(void* addr) { + RWSpinlock::ReadGuard guard(mr_lock_); + auto it = mr_map_.upper_bound((uint64_t)addr); + if (it != mr_map_.begin()) { + --it; + if ((uint64_t)addr < it->first + it->second.length && it->second.mr) { + return it->second.key; + } + } + return 0; +} + +void* CxiContext::mrDesc(void* addr) { + RWSpinlock::ReadGuard guard(mr_lock_); + // Find the MR that contains this address + auto it = mr_map_.upper_bound((uint64_t)addr); + if (it != mr_map_.begin()) { + --it; + if ((uint64_t)addr < it->first + it->second.length && it->second.mr) { + return fi_mr_desc(it->second.mr); + } + } + return nullptr; +} + +std::shared_ptr CxiContext::endpoint( + const std::string& peer_nic_path) { + // Use normalized key (strip port) so the same physical peer reuses its + // handle across reconnections. Each P2PHANDSHAKE run picks a random + // port, producing a different peer_nic_path for the same peer host+NIC. + std::string key = peer_nic_path; + + { + RWSpinlock::ReadGuard guard(peer_map_lock_); + auto it = peer_map_.find(key); + if (it != peer_map_.end()) { + it->second->setPeerNicPath(peer_nic_path); + return it->second; + } + } + + auto new_ep = std::make_shared(*this); + new_ep->setPeerNicPath(peer_nic_path); + + RWSpinlock::WriteGuard guard(peer_map_lock_); + auto it = peer_map_.find(key); + if (it != peer_map_.end()) { + it->second->setPeerNicPath(peer_nic_path); + return it->second; + } + peer_map_[key] = new_ep; + return new_ep; +} + +std::shared_ptr CxiContext::peekEndpoint( + const std::string& peer_nic_path) { + RWSpinlock::ReadGuard guard(peer_map_lock_); + auto it = peer_map_.find(peer_nic_path); + if (it == peer_map_.end()) return nullptr; + return it->second; +} + +int CxiContext::deleteEndpoint(const std::string& peer_nic_path) { + std::shared_ptr ep; + { + RWSpinlock::WriteGuard guard(peer_map_lock_); + auto it = peer_map_.find(peer_nic_path); + if (it == peer_map_.end()) return 0; + ep = it->second; + peer_map_.erase(it); + } + if (ep) ep->disconnect(); // runs fi_av_remove + return 0; +} + +int CxiContext::disconnectAllEndpoints() { + RWSpinlock::WriteGuard guard(peer_map_lock_); + for (auto& entry : peer_map_) { + if (entry.second) entry.second->disconnect(); + } + return 0; +} + +size_t CxiContext::getTotalQPNumber() const { + RWSpinlock::ReadGuard guard(peer_map_lock_); + return peer_map_.size(); +} + +std::string CxiContext::nicPath() const { + return engine_.local_server_name() + "@" + device_name_; +} + +std::string CxiContext::localAddr() const { + // Return a hex string representation of the local address info + if (!fi_info_ || !fi_info_->src_addr) { + return ""; + } + + std::ostringstream oss; + const uint8_t* addr = static_cast(fi_info_->src_addr); + for (size_t i = 0; i < fi_info_->src_addrlen; ++i) { + oss << std::hex << std::setw(2) << std::setfill('0') << (int)addr[i]; + } + return oss.str(); +} + +std::string CxiContext::localEpAddr() const { + static constexpr char kHex[] = "0123456789abcdef"; + std::string out; + out.resize(local_ep_addr_.size() * 2); + for (size_t i = 0; i < local_ep_addr_.size(); ++i) { + out[2 * i] = kHex[(local_ep_addr_[i] >> 4) & 0xF]; + out[2 * i + 1] = kHex[local_ep_addr_[i] & 0xF]; + } + return out; +} + +// Decode one hex nibble, -1 on invalid input. +static inline int hexNibble(char c) { + if (c >= '0' && c <= '9') return c - '0'; + if (c >= 'a' && c <= 'f') return c - 'a' + 10; + if (c >= 'A' && c <= 'F') return c - 'A' + 10; + return -1; +} + +int CxiContext::insertPeerAddr(const std::string& peer_hex_addr, + fi_addr_t& out) { + if (peer_hex_addr.empty() || (peer_hex_addr.size() % 2) != 0) { + LOG(ERROR) << "insertPeerAddr: invalid hex length " + << peer_hex_addr.size(); + return ERR_INVALID_ARGUMENT; + } + const size_t n = peer_hex_addr.size() / 2; + std::vector bin(n); + for (size_t i = 0; i < n; ++i) { + int hi = hexNibble(peer_hex_addr[2 * i]); + int lo = hexNibble(peer_hex_addr[2 * i + 1]); + if (hi < 0 || lo < 0) { + LOG(ERROR) << "insertPeerAddr: non-hex char at offset " << (2 * i); + return ERR_INVALID_ARGUMENT; + } + bin[i] = static_cast((hi << 4) | lo); + } + return insertPeerAddrBytes(bin.data(), bin.size(), out); +} + +int CxiContext::insertPeerAddrBytes(const uint8_t* addr, size_t len, + fi_addr_t& out) { + if (!addr || len == 0) return ERR_INVALID_ARGUMENT; + int ret = fi_av_insert(av_, addr, 1, &out, 0, nullptr); + if (ret != 1) { + LOG(ERROR) << "fi_av_insert failed: " << fi_strerror(-ret); + return ERR_ENDPOINT; + } + return 0; +} + +void CxiContext::removePeerAddr(fi_addr_t fi_addr) { + if (fi_addr == FI_ADDR_UNSPEC) return; + int ret = fi_av_remove(av_, &fi_addr, 1, 0); + if (ret) { + LOG(WARNING) << "fi_av_remove failed: " << fi_strerror(-ret); + } +} + +int CxiContext::submitPostSend( + const std::vector& slice_list) { + // Route slices to appropriate peer handles. Group by peer NIC path. + std::unordered_map> + slices_by_peer; + + for (auto* slice : slice_list) { + if (!slice) continue; + + // Fast path: peer info already filled in by the caller + // (dest_rkey and peer_nic_path set on the slice before dispatch). + if (!slice->peer_nic_path.empty()) { + slices_by_peer[slice->peer_nic_path].push_back(slice); + continue; + } + + // Slow path: resolve peer info per-slice. + auto peer_segment_desc = + engine_.meta()->getSegmentDescByID(slice->target_id); + if (!peer_segment_desc) { + LOG(ERROR) << "Cannot get segment descriptor for target " + << slice->target_id; + slice->markFailed(); + continue; + } + + int buffer_id = -1, device_id = -1; + if (CxiTransport::selectDevice(peer_segment_desc.get(), + slice->rdma.dest_addr, slice->length, + buffer_id, device_id)) { + // The cached peer descriptor may be stale: the remote can free + // and re-register buffers at new virtual addresses between + // iterations (e.g. a fresh per-iter H2D buffer). Force-refresh + // the segment descriptor once and retry before giving up. + // Mirrors the RDMA transport's retry-on-miss in worker_pool.cpp. + peer_segment_desc = + engine_.meta()->getSegmentDescByID(slice->target_id, true); + if (!peer_segment_desc || + CxiTransport::selectDevice(peer_segment_desc.get(), + slice->rdma.dest_addr, slice->length, + buffer_id, device_id)) { + LOG(ERROR) << "Cannot select device for dest_addr " + << (void*)slice->rdma.dest_addr; + slice->markFailed(); + continue; + } + } + + // device_id comes from the peer-supplied topology, whose HCA list is + // independent of the peer 'devices' array, and selectDevice() bounds it + // against rkey only. decodeSegmentDesc() now rejects a descriptor whose + // key count and device count disagree; bound the value used to index + // devices[] locally as well. + if (static_cast(device_id) >= + peer_segment_desc->devices.size()) { + LOG(ERROR) << "Peer device index out of range for target " + << slice->target_id << ": device_id=" << device_id + << " devices=" << peer_segment_desc->devices.size(); + slice->markFailed(); + continue; + } + + // no FI_VIRT_ADDR support on slingshot, must be offset of memory region + slice->rdma.dest_addr -= peer_segment_desc->buffers[buffer_id].addr; + slice->rdma.dest_rkey = + peer_segment_desc->buffers[buffer_id].rkey[device_id]; + + std::string peer_nic_path = peer_segment_desc->name + "@" + + peer_segment_desc->devices[device_id].name; + slice->peer_nic_path = peer_nic_path; + slices_by_peer[peer_nic_path].push_back(slice); + } + + for (auto& entry : slices_by_peer) { + const std::string& peer_nic_path = entry.first; + auto& peer_slices = entry.second; + + auto ep = endpoint(peer_nic_path); + if (!ep) { + LOG(ERROR) << "Cannot create peer handle for " << peer_nic_path; + for (auto* slice : peer_slices) slice->markFailed(); + continue; + } + + std::vector failed_slice_list; + int rc = ep->submitPostSend(peer_slices, failed_slice_list); + for (auto* slice : failed_slice_list) slice->markFailed(); + + // Drop peer handle if it is no longer connected after submit, + // freeing its AV entry for reuse. Under the shared-endpoint model + // this is cheap (no fid_ep to destroy). + if (rc != 0 && !ep->connected()) { + deleteEndpoint(peer_nic_path); + } + } + + return 0; +} + +int CxiContext::submitSlicesOnPeer( + fi_addr_t peer_fi_addr, std::vector& slice_list, + std::vector& failed_slice_list) { + // Batched submission against the shared endpoint. Mirrors the previous + // per-endpoint submit path but uses context-level wr_depth / post_lock. + // + // 1. Reserve N WR+CQ slots in bulk (single CAS each) + // 2. Prepare MR descriptors and op contexts outside the lock + // 3. Hold post_lock_ once for the entire batch of fi_write calls + const int kMaxBackoffYields = 100000; + const int cq_limit = static_cast(globalConfig().max_cqe); + std::atomic* cq_outstanding = + shared_cq_ ? &shared_cq_->outstanding : nullptr; + + struct BatchEntry { + Transport::Slice* slice; + void* local_desc; + CxiOpContext* op_ctx; + }; + + // Consume slice_list via a moving index instead of erase-from-front, + // which was O(N^2) on large batches. retry_slices accumulate across + // passes and are applied by rewinding the cursor. + size_t cursor = 0; + std::vector retry_slices; + while (cursor < slice_list.size() || !retry_slices.empty()) { + if (!retry_slices.empty()) { + // Splice retry slices back in at the current cursor so the next + // pass picks them up. O(retry_slices.size()) per retry wave, + // which is bounded by valid_count of the last batch. + slice_list.insert(slice_list.begin() + cursor, retry_slices.begin(), + retry_slices.end()); + retry_slices.clear(); + std::this_thread::yield(); + } + + const size_t remaining = slice_list.size() - cursor; + int batch_count = 0; + int backoff = 0; + bool timed_out = false; + while (batch_count == 0) { + int cur_wr = wr_depth_.load(std::memory_order_relaxed); + int wr_avail = max_wr_depth_ - cur_wr; + if (wr_avail <= 0) { + if (++backoff > kMaxBackoffYields) { + timed_out = true; + break; + } + std::this_thread::yield(); + continue; + } + int want = std::min(wr_avail, (int)remaining); + if (cq_outstanding) { + int cur_cq = cq_outstanding->load(std::memory_order_relaxed); + int cq_avail = cq_limit - cur_cq; + if (cq_avail <= 0) { + if (++backoff > kMaxBackoffYields) { + timed_out = true; + break; + } + std::this_thread::yield(); + continue; + } + want = std::min(want, cq_avail); + if (!wr_depth_.compare_exchange_weak( + cur_wr, cur_wr + want, std::memory_order_acq_rel, + std::memory_order_relaxed)) { + continue; + } + cur_cq = cq_outstanding->load(std::memory_order_relaxed); + cq_avail = cq_limit - cur_cq; + if (cq_avail < want) { + wr_depth_.fetch_sub(want, std::memory_order_acq_rel); + continue; + } + if (!cq_outstanding->compare_exchange_weak( + cur_cq, cur_cq + want, std::memory_order_acq_rel, + std::memory_order_relaxed)) { + wr_depth_.fetch_sub(want, std::memory_order_acq_rel); + continue; + } + } else { + if (!wr_depth_.compare_exchange_weak( + cur_wr, cur_wr + want, std::memory_order_acq_rel, + std::memory_order_relaxed)) { + continue; + } + } + batch_count = want; + } + + if (timed_out) { + LOG(WARNING) << "CXI submitSlicesOnPeer: timed out waiting for CQ" + << " drain (wr_depth=" + << wr_depth_.load(std::memory_order_relaxed) + << ", max=" << max_wr_depth_ << ", cq_outstanding=" + << (cq_outstanding ? cq_outstanding->load( + std::memory_order_relaxed) + : -1) + << ", max_cqe=" << cq_limit << ")"; + for (size_t i = cursor; i < slice_list.size(); ++i) { + failed_slice_list.push_back(slice_list[i]); + } + slice_list.clear(); + return 0; + } + + std::vector batch(batch_count); + int valid_count = 0; + + for (int i = 0; i < batch_count; i++) { + Transport::Slice* slice = slice_list[cursor + i]; + void* local_desc = mrDesc(slice->source_addr); + if (!local_desc) { + LOG(ERROR) << "No MR descriptor found for address " + << slice->source_addr; + failed_slice_list.push_back(slice); + continue; + } + CxiOpContext* op_ctx = new CxiOpContext(); + memset(op_ctx, 0, sizeof(CxiOpContext)); + op_ctx->slice = slice; + op_ctx->wr_depth = &wr_depth_; + batch[valid_count++] = {slice, local_desc, op_ctx}; + } + + int mr_failures = batch_count - valid_count; + if (mr_failures > 0) { + wr_depth_.fetch_sub(mr_failures, std::memory_order_acq_rel); + if (cq_outstanding) + cq_outstanding->fetch_sub(mr_failures, + std::memory_order_acq_rel); + } + + if (valid_count > 0) { + while (post_lock_.test_and_set(std::memory_order_acquire)) { + } + for (int i = 0; i < valid_count; i++) { + auto& entry = batch[i]; + ssize_t ret; + if (entry.slice->opcode == Transport::TransferRequest::READ) { + ret = fi_read(shared_ep_, (void*)entry.slice->source_addr, + entry.slice->length, entry.local_desc, + peer_fi_addr, entry.slice->rdma.dest_addr, + entry.slice->rdma.dest_rkey, + &entry.op_ctx->fi_ctx); + } else { + ret = fi_write(shared_ep_, (void*)entry.slice->source_addr, + entry.slice->length, entry.local_desc, + peer_fi_addr, entry.slice->rdma.dest_addr, + entry.slice->rdma.dest_rkey, + &entry.op_ctx->fi_ctx); + } + if (ret == 0) { + entry.slice->status = Transport::Slice::PENDING; + } else if (ret == -FI_EAGAIN) { + delete entry.op_ctx; + int not_posted = valid_count - i; + wr_depth_.fetch_sub(not_posted, std::memory_order_acq_rel); + if (cq_outstanding) + cq_outstanding->fetch_sub(not_posted, + std::memory_order_acq_rel); + for (int j = i; j < valid_count; j++) { + if (j > i) delete batch[j].op_ctx; + retry_slices.push_back(batch[j].slice); + } + break; + } else { + LOG(ERROR) + << "fi_read/fi_write failed: " << fi_strerror(-ret) + << " (source=" << entry.slice->source_addr + << ", len=" << entry.slice->length + << ", dest=" << (void*)entry.slice->rdma.dest_addr + << ", rkey=" << entry.slice->rdma.dest_rkey << ")"; + delete entry.op_ctx; + wr_depth_.fetch_sub(1, std::memory_order_acq_rel); + if (cq_outstanding) + cq_outstanding->fetch_sub(1, std::memory_order_acq_rel); + failed_slice_list.push_back(entry.slice); + } + } + post_lock_.clear(std::memory_order_release); + } + + cursor += batch_count; + } + + slice_list.clear(); + return 0; +} + +int CxiContext::pollCq(int max_entries, int cq_index) { + if (cq_index < 0 || (size_t)cq_index >= cq_list_.size()) { + return 0; + } + + struct fid_cq* cq = cq_list_[cq_index]->cq; + if (!cq) return 0; + + struct fi_cq_data_entry entries[64]; + int to_poll = std::min(max_entries, 64); + + ssize_t ret = fi_cq_read(cq, entries, to_poll); + + if (ret > 0) { + std::unordered_map*, int> wr_depth_set; + for (ssize_t i = 0; i < ret; i++) { + CxiOpContext* op_ctx = + reinterpret_cast(entries[i].op_context); + if (op_ctx && op_ctx->slice) { + op_ctx->slice->markSuccess(); + if (op_ctx->wr_depth) { + wr_depth_set[op_ctx->wr_depth]++; + } + delete op_ctx; + } + } + for (auto& entry : wr_depth_set) { + entry.first->fetch_sub(entry.second, std::memory_order_acq_rel); + } + cq_list_[cq_index]->outstanding.fetch_sub(static_cast(ret), + std::memory_order_acq_rel); + return static_cast(ret); + } else if (ret == -FI_EAGAIN) { + return 0; + } else if (ret < 0) { + int err_count = 0; + struct fi_cq_err_entry err_entry; + std::unordered_map*, int> wr_depth_set; + + while ((ret = fi_cq_readerr(cq, &err_entry, 0)) > 0) { + CxiOpContext* op_ctx = + reinterpret_cast(err_entry.op_context); + if (op_ctx && op_ctx->slice) { + LOG(ERROR) << "CXI CQ error: " + << fi_cq_strerror(cq, err_entry.prov_errno, + err_entry.err_data, nullptr, 0) + << " for slice at " << op_ctx->slice->source_addr; + op_ctx->slice->markFailed(); + if (op_ctx->wr_depth) { + wr_depth_set[op_ctx->wr_depth]++; + } + delete op_ctx; + } + err_count++; + } + + for (auto& entry : wr_depth_set) { + entry.first->fetch_sub(entry.second, std::memory_order_acq_rel); + } + if (err_count > 0) { + cq_list_[cq_index]->outstanding.fetch_sub( + err_count, std::memory_order_acq_rel); + } + return err_count; + } + + return 0; +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_endpoint.cpp b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_endpoint.cpp new file mode 100644 index 0000000000..f17326e6a3 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_endpoint.cpp @@ -0,0 +1,186 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "transport/cxi_transport/cxi_endpoint.h" + +#include + +#include +#include +#include +#include +#include +#include + +#include "config.h" + +namespace mooncake { + +CxiEndpoint::CxiEndpoint(CxiContext& context) + : context_(context), status_(INITIALIZING), peer_fi_addr_(FI_ADDR_UNSPEC) {} + +CxiEndpoint::~CxiEndpoint() { disconnect(); } + +void CxiEndpoint::setPeerNicPath(const std::string& peer_nic_path) { + RWSpinlock::WriteGuard guard(lock_); + if (peer_nic_path_ == peer_nic_path) return; // No change + if (status_.load(std::memory_order_relaxed) == CONNECTED) { + LOG(INFO) << "Peer reconnected with new address, re-establishing: " + << peer_nic_path_ << " -> " << peer_nic_path; + disconnectUnlocked(); + } + peer_nic_path_ = peer_nic_path; +} + +int CxiEndpoint::setupConnectionsByActive() { + RWSpinlock::WriteGuard guard(lock_); + if (status_.load(std::memory_order_relaxed) == CONNECTED) return 0; + + // Loopback: handshake against ourselves. Use the binary overload + // so we avoid hex-encoding the local address just to have + // insertPeerAddr decode it right back. + if (context_.nicPath() == peer_nic_path_) { + const auto& bytes = context_.localEpAddrBytes(); + int ret = context_.insertPeerAddrBytes(bytes.data(), bytes.size(), + peer_fi_addr_); + if (ret != 0) return ret; + status_.store(CONNECTED, std::memory_order_release); + LOG(INFO) << "CXI loopback connection established: " << toString(); + return 0; + } + + // Exchange addresses via the transfer-metadata handshake RPC. + TransferMetadata::HandShakeDesc local_desc, peer_desc; + local_desc.local_nic_path = context_.nicPath(); + local_desc.peer_nic_path = peer_nic_path_; + local_desc.cxi_addr = context_.localEpAddr(); + + auto peer_server_name = getServerNameFromNicPath(peer_nic_path_); + auto peer_nic_name = getNicNameFromNicPath(peer_nic_path_); + if (peer_server_name.empty() || peer_nic_name.empty()) { + LOG(ERROR) << "Parse peer CXI nic path failed: " << peer_nic_path_; + return ERR_INVALID_ARGUMENT; + } + + int rc = context_.engine().sendHandshake(peer_server_name, local_desc, + peer_desc); + if (rc) return rc; + + if (peer_desc.cxi_addr.empty()) { + LOG(ERROR) << "Peer did not provide CXI address in handshake"; + return ERR_REJECT_HANDSHAKE; + } + + rc = context_.insertPeerAddr(peer_desc.cxi_addr, peer_fi_addr_); + if (rc != 0) return rc; + + status_.store(CONNECTED, std::memory_order_release); + VLOG(1) << "CXI connection established: " << toString() + << " peer_fi_addr=" << peer_fi_addr_; + return 0; +} + +int CxiEndpoint::setupConnectionsByPassive(const HandShakeDesc& peer_desc, + HandShakeDesc& local_desc) { + RWSpinlock::WriteGuard guard(lock_); + if (status_.load(std::memory_order_relaxed) == CONNECTED) { + LOG(WARNING) << "Re-establish CXI connection: " << toString(); + disconnectUnlocked(); + } + + if (peer_desc.peer_nic_path != context_.nicPath() || + peer_desc.local_nic_path != peer_nic_path_) { + local_desc.reply_msg = "CXI nic path inconsistency"; + LOG(ERROR) << "Invalid argument: peer CXI nic path inconsistency" + << " peer_nic_path=" << peer_desc.peer_nic_path + << " context_.nicPath()=" << context_.nicPath() + << " local_nic_path=" << peer_desc.local_nic_path + << " peer_nic_path_=" << peer_nic_path_; + return ERR_REJECT_HANDSHAKE; + } + + if (peer_desc.cxi_addr.empty()) { + local_desc.reply_msg = "No CXI address provided"; + LOG(ERROR) << "Peer did not provide CXI address"; + return ERR_REJECT_HANDSHAKE; + } + + int ret = context_.insertPeerAddr(peer_desc.cxi_addr, peer_fi_addr_); + if (ret != 0) { + local_desc.reply_msg = "Failed to insert peer address"; + return ret; + } + + local_desc.local_nic_path = context_.nicPath(); + local_desc.peer_nic_path = peer_nic_path_; + local_desc.cxi_addr = context_.localEpAddr(); + // reply_msg empty on success + + status_.store(CONNECTED, std::memory_order_release); + VLOG(1) << "CXI connection established (passive): " << toString(); + return 0; +} + +const std::string CxiEndpoint::toString() const { + return "CxiEndpoint[" + context_.nicPath() + " <-> " + peer_nic_path_ + "]"; +} + +void CxiEndpoint::disconnect() { + RWSpinlock::WriteGuard guard(lock_); + disconnectUnlocked(); +} + +void CxiEndpoint::disconnectUnlocked() { + if (peer_fi_addr_ != FI_ADDR_UNSPEC) { + context_.removePeerAddr(peer_fi_addr_); + peer_fi_addr_ = FI_ADDR_UNSPEC; + } + status_.store(UNCONNECTED, std::memory_order_release); +} + +void CxiEndpoint::markDetachedForTeardown() { + RWSpinlock::WriteGuard guard(lock_); + peer_fi_addr_ = FI_ADDR_UNSPEC; + status_.store(UNCONNECTED, std::memory_order_release); +} + +int CxiEndpoint::submitPostSend( + std::vector& slice_list, + std::vector& failed_slice_list) { + if (status_.load(std::memory_order_relaxed) != CONNECTED) { + int ret = setupConnectionsByActive(); + if (ret != 0) { + for (auto* slice : slice_list) { + failed_slice_list.push_back(slice); + } + slice_list.clear(); + return ret; + } + } + + fi_addr_t peer; + { + RWSpinlock::ReadGuard guard(lock_); + peer = peer_fi_addr_; + } + if (peer == FI_ADDR_UNSPEC) { + for (auto* slice : slice_list) failed_slice_list.push_back(slice); + slice_list.clear(); + return ERR_ENDPOINT; + } + + return context_.submitSlicesOnPeer(peer, slice_list, failed_slice_list); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp new file mode 100644 index 0000000000..3c42fe66da --- /dev/null +++ b/mooncake-transfer-engine/src/transport/cxi_transport/cxi_transport.cpp @@ -0,0 +1,1015 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "transport/cxi_transport/cxi_transport.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "common.h" +#include "config.h" +#include "memory_location.h" +#include "topology.h" +#include "transport/cxi_transport/cxi_context.h" +#include "transport/cxi_transport/cxi_endpoint.h" + +namespace mooncake { + +NicReplicaPolicy CxiTransport::getReplicaPolicy() { + return NicReplicaPolicy::NUMA_AWARE; +} + +CxiTransport::CxiTransport() { + LOG(INFO) << "[CXI] CXI Slingshot transport initialized"; +} + +CxiTransport::~CxiTransport() { + stopWorkerThreads(); + metadata_->removeSegmentDesc(local_server_name_); + batch_desc_set_.clear(); + context_list_.clear(); +} + +void CxiTransport::startWorkerThreads() { + if (worker_running_) return; + + worker_running_ = true; + // One poller thread per context for responsive CQ draining under load + size_t num_threads = context_list_.size(); + for (size_t i = 0; i < num_threads; i++) { + worker_threads_.emplace_back(&CxiTransport::workerThreadFunc, this, i); + } + LOG(INFO) << "CxiTransport: Started " << num_threads + << " CQ polling worker threads"; +} + +void CxiTransport::stopWorkerThreads() { + if (!worker_running_) return; + + worker_running_ = false; + for (auto& thread : worker_threads_) { + if (thread.joinable()) { + thread.join(); + } + } + worker_threads_.clear(); + LOG(INFO) << "CxiTransport: Stopped CQ polling worker threads"; +} + +void CxiTransport::workerThreadFunc(int thread_id) { + const int kPollBatchSize = 64; + + while (worker_running_) { + bool did_work = false; + + // Poll CQs from all contexts + for (size_t ctx_idx = thread_id; ctx_idx < context_list_.size(); + ctx_idx += worker_threads_.size()) { + auto& context = context_list_[ctx_idx]; + if (!context || !context->active()) continue; + + for (size_t cq_idx = 0; cq_idx < context->cqCount(); cq_idx++) { + int completed = context->pollCq(kPollBatchSize, cq_idx); + if (completed > 0) { + did_work = true; + } + } + } + + // Under the shared-endpoint model there is no per-peer QP to evict: + // a new peer costs one AV entry (~bytes), not an fi_endpoint slot. + // Stale peers are reclaimed when submitPostSend() drops a peer whose + // handshake failed. + + // If no work was done, yield CPU briefly + if (!did_work) { + std::this_thread::yield(); + } + } +} + +int CxiTransport::install(std::string& local_server_name, + std::shared_ptr meta, + std::shared_ptr topo) { + if (topo == nullptr) { + LOG(ERROR) << "CxiTransport: missing topology"; + return ERR_INVALID_ARGUMENT; + } + + metadata_ = meta; + local_server_name_ = local_server_name; + local_topology_ = topo; + + auto ret = initializeCxiResources(); + if (ret) { + LOG(ERROR) << "CxiTransport: cannot initialize CXI resources"; + return ret; + } + + ret = allocateLocalSegmentID(); + if (ret) { + LOG(ERROR) << "Transfer engine cannot be initialized: cannot " + "allocate local segment"; + return ret; + } + + ret = startHandshakeDaemon(local_server_name); + if (ret) { + LOG(ERROR) << "CxiTransport: cannot start handshake daemon"; + return ret; + } + + ret = metadata_->updateLocalSegmentDesc(); + if (ret) { + LOG(ERROR) << "CxiTransport: cannot publish segments"; + return ret; + } + + // Start CQ polling worker threads + startWorkerThreads(); + + return 0; +} + +int CxiTransport::preTouchMemory(void* addr, size_t length) { + if (context_list_.size() == 0) { + return 0; + } + + unsigned int hwc = std::thread::hardware_concurrency(); + unsigned int num_threads = hwc > 64 ? 16 : std::min(hwc, 8u); + num_threads = std::max( + num_threads, 1u); // guard in case hardware_concurrency() returns 0 + size_t block_size = length / num_threads; + if (block_size == 0) { + return 0; + } + + std::vector threads; + threads.reserve(num_threads); + std::vector thread_results(num_threads, 0); + + for (size_t thread_i = 0; thread_i < num_threads; ++thread_i) { + void* block_addr = static_cast(addr) + thread_i * block_size; + threads.emplace_back([this, thread_i, block_addr, block_size, + &thread_results]() { + int ret = context_list_[0]->preTouchMemory(block_addr, block_size); + thread_results[thread_i] = ret; + }); + } + + for (auto& thread : threads) { + thread.join(); + } + + for (size_t i = 0; i < num_threads; ++i) { + if (thread_results[i] != 0) { + return thread_results[i]; + } + } + + return 0; +} + +int CxiTransport::registerLocalMemory(void* addr, size_t length, + const std::string& name, + bool remote_accessible, + bool update_metadata) { + return registerLocalMemoryInternal(addr, length, name, remote_accessible, + update_metadata, false); +} + +int CxiTransport::registerLocalMemoryInternal(void* addr, size_t length, + const std::string& name, + bool remote_accessible, + bool update_metadata, + bool force_sequential) { + (void)remote_accessible; + const int kBaseAccessRights = + IBV_ACCESS_LOCAL_WRITE | IBV_ACCESS_REMOTE_WRITE | + IBV_ACCESS_REMOTE_READ; // this is not used, ignore + + int access_rights = kBaseAccessRights; + size_t max_mr = (size_t)globalConfig().max_mr_size; + + // max_mr must be set, slingshot handles mempages differently compared to + // efa + size_t chunk_limit = max_mr; + LOG(INFO) << "Auto-split params: max_mr_size=" << max_mr + << ", chunk_limit=" << chunk_limit; + + // Determine chunk boundaries + std::vector> chunks; + if (length > chunk_limit) { + size_t offset = 0; + while (offset < length) { + size_t chunk_len = std::min(chunk_limit, length - offset); + chunks.emplace_back(static_cast(addr) + offset, chunk_len); + offset += chunk_len; + } + LOG(WARNING) << "Auto-splitting buffer " << addr << " (" << length + << " bytes) into " << chunks.size() + << " chunks of <= " << chunk_limit << " bytes each"; + } else { + chunks.emplace_back(addr, length); + } + + // Resolve location name once (based on original buffer) + std::string resolved_name; + if (name == kWildcardLocation) { + bool only_first_page = true; + const std::vector entries = getMemoryLocation( + addr, length, only_first_page); // check only first page + if (entries.empty()) return ERR_DEVICE_NOT_FOUND; + resolved_name = entries[0].location; + } else { + resolved_name = name; + } + + size_t num_nics = context_list_.size(); + size_t num_chunks = chunks.size(); + + auto policy = getReplicaPolicy(); + + std::vector> nic_assignments(num_chunks); + + int selected_device = local_topology_->selectDevice(resolved_name); + if (selected_device == ERR_DEVICE_NOT_FOUND) { + LOG(ERROR) << "could not select a NIC for data, resolved_name=" + << resolved_name; + return ERR_DEVICE_NOT_FOUND; + } + std::string nic = local_topology_->getHcaList().at( + local_topology_->selectDevice(resolved_name)); + LOG(INFO) << "for this allocation, selected NIC " << nic; + int id = -1; + for (size_t i = 0; i < context_list_.size(); i++) { + auto& entry = context_list_[i]; + if (entry->deviceName() == nic) { + id = i; + break; + } + } + if (id == -1) { + LOG(ERROR) << "possible mismatch between context list devices and " + "topology HCA names!"; + return -1; + } + int numa_node = id; + + // implicit assumption here, the address and relative chunks all lie in the + // same numa node + for (size_t ci = 0; ci < num_chunks; ci++) { + // based on replication policy, assign chunk to nic + switch (policy) { + case NUMA_AWARE: + nic_assignments[ci].push_back(numa_node); + break; + case REPLICATE_ALL: + for (size_t nic_idx = 0; nic_idx < num_nics; nic_idx++) + nic_assignments[ci].push_back(nic_idx); + break; + default: + nic_assignments[ci].push_back( + 0); // register everything to nic 0 + break; + } + } + + auto rollbackChunks = [&](size_t up_to_ci) { + for (size_t ri = 0; ri <= up_to_ci; ++ri) { + for (size_t nic_idx : nic_assignments[ri]) { + context_list_[nic_idx]->unregisterMemoryRegion( + chunks[ri].first); + } + } + }; + + // Register each chunk on its assigned NICs + for (size_t ci = 0; ci < chunks.size(); ++ci) { + void* chunk_addr = chunks[ci].first; + size_t chunk_len = chunks[ci].second; + + const auto& assigned_nics = nic_assignments[ci]; + + // preTouchMemory does a CPU-side store to each page, which segfaults + // on GPU VRAM (cudaMalloc'd pointers). Restrict it to host memory. + bool is_host_mem = resolved_name.rfind("cpu", 0) == 0; + bool do_pre_touch = is_host_mem && context_list_.size() > 0 && + std::thread::hardware_concurrency() >= 4 && + chunk_len >= (size_t)4 * 1024 * 1024 * 1024; + if (do_pre_touch) { + int ret = preTouchMemory(chunk_addr, chunk_len); + if (ret != 0) { + if (ci > 0) rollbackChunks(ci - 1); + return ret; + } + } + + int use_parallel_reg = 0; + if (!force_sequential) { + use_parallel_reg = globalConfig().parallel_reg_mr; + if (use_parallel_reg == -1) { + use_parallel_reg = assigned_nics.size() > 1 && do_pre_touch; + } + } + + auto reg_start = std::chrono::steady_clock::now(); + + if (use_parallel_reg) { + std::vector reg_threads; + reg_threads.reserve(assigned_nics.size()); + std::vector ret_codes(assigned_nics.size(), 0); + const int ar = access_rights; + + for (size_t j = 0; j < assigned_nics.size(); ++j) { + size_t nic_idx = assigned_nics[j]; + reg_threads.emplace_back([this, &ret_codes, j, nic_idx, + chunk_addr, chunk_len, ar]() { + ret_codes[j] = context_list_[nic_idx]->registerMemoryRegion( + chunk_addr, chunk_len, ar); + }); + } + + for (auto& thread : reg_threads) { + thread.join(); + } + + for (size_t j = 0; j < ret_codes.size(); ++j) { + if (ret_codes[j] != 0) { + LOG(ERROR) + << "Failed to register memory region chunk " << ci + << " with CXI context " << assigned_nics[j]; + rollbackChunks(ci); + return ret_codes[j]; + } + } + } else { + for (size_t nic_idx : assigned_nics) { + int ret = context_list_[nic_idx]->registerMemoryRegion( + chunk_addr, chunk_len, access_rights); + if (ret) { + LOG(ERROR) << "Failed to register memory region chunk " + << ci << " with CXI context " << nic_idx; + rollbackChunks(ci); + return ret; + } + } + } + + auto reg_end = std::chrono::steady_clock::now(); + auto reg_duration_ms = + std::chrono::duration_cast(reg_end - + reg_start) + .count(); + + if (globalConfig().trace) { + LOG(INFO) << "CXI registerMemoryRegion: chunk " << ci + << ", addr=" << chunk_addr << ", length=" << chunk_len + << ", nics=" << assigned_nics.size() << "/" + << context_list_.size() + << ", parallel=" << (use_parallel_reg ? "true" : "false") + << ", duration=" << reg_duration_ms << "ms"; + } + + LOG(WARNING) << "Chunk " << ci << "/" << chunks.size() + << " registered on " << assigned_nics.size() << " NICs" + << ", addr=" << chunk_addr << ", length=" << chunk_len + << ", duration=" << reg_duration_ms << "ms"; + + // Collect keys: assigned NICs have valid keys, others get 0 + BufferDesc buffer_desc; + for (auto& context : context_list_) { + buffer_desc.lkey.push_back(context->lkey(chunk_addr)); + buffer_desc.rkey.push_back(context->rkey(chunk_addr)); + } + + buffer_desc.name = resolved_name; + buffer_desc.addr = (uint64_t)chunk_addr; + buffer_desc.length = chunk_len; + int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata); + if (rc) { + rollbackChunks(ci); + return rc; + } + } + + // Track chunks and NIC assignments for unregistration + if (chunks.size() > 1) { + std::lock_guard lock(chunk_map_mutex_); + std::vector regs; + regs.reserve(chunks.size()); + for (size_t ci = 0; ci < chunks.size(); ++ci) { + regs.push_back({(uint64_t)chunks[ci].first, nic_assignments[ci]}); + } + chunk_map_[(uint64_t)addr] = std::move(regs); + } + + return 0; +} + +int CxiTransport::unregisterLocalMemory(void* addr, bool update_metadata) { + return unregisterLocalMemoryInternal(addr, update_metadata, false); +} + +int CxiTransport::unregisterLocalMemoryInternal(void* addr, + bool update_metadata, + bool force_sequential) { + // Check if this buffer was split into chunks (per-NIC partition) + std::vector chunk_regs; + { + std::lock_guard lock(chunk_map_mutex_); + auto it = chunk_map_.find((uint64_t)addr); + if (it != chunk_map_.end()) { + chunk_regs = std::move(it->second); + chunk_map_.erase(it); + } + } + + if (!chunk_regs.empty()) { + // Unregister each chunk from its assigned NICs only + for (auto& reg : chunk_regs) { + void* ca = (void*)reg.addr; + int rc = metadata_->removeLocalMemoryBuffer(ca, update_metadata); + if (rc) { + LOG(ERROR) << "Failed to remove chunk metadata at " << ca; + return rc; + } + + for (size_t nic_idx : reg.nic_indices) { + int ret = context_list_[nic_idx]->unregisterMemoryRegion(ca); + if (ret) { + LOG(ERROR) << "Failed to unregister chunk " << ca + << " with CXI context " << nic_idx; + return ret; + } + } + } + return 0; + } + + // Non-chunked buffer: original path + int rc = metadata_->removeLocalMemoryBuffer(addr, update_metadata); + if (rc) return rc; + + int use_parallel_unreg = 0; + if (!force_sequential) { + use_parallel_unreg = globalConfig().parallel_reg_mr; + if (use_parallel_unreg == -1) { + use_parallel_unreg = context_list_.size() > 1; + } + } + + if (use_parallel_unreg) { + std::vector unreg_threads; + unreg_threads.reserve(context_list_.size()); + std::vector ret_codes(context_list_.size(), 0); + + for (size_t i = 0; i < context_list_.size(); ++i) { + unreg_threads.emplace_back([this, &ret_codes, i, addr]() { + ret_codes[i] = context_list_[i]->unregisterMemoryRegion(addr); + }); + } + + for (auto& thread : unreg_threads) { + thread.join(); + } + + for (size_t i = 0; i < ret_codes.size(); ++i) { + if (ret_codes[i] != 0) { + LOG(ERROR) + << "Failed to unregister memory region with CXI context " + << i; + return ret_codes[i]; + } + } + } else { + for (size_t i = 0; i < context_list_.size(); ++i) { + int ret = context_list_[i]->unregisterMemoryRegion(addr); + if (ret) { + LOG(ERROR) + << "Failed to unregister memory region with CXI context " + << i; + return ret; + } + } + } + + return 0; +} + +int CxiTransport::allocateLocalSegmentID() { + auto desc = std::make_shared(); + if (!desc) return ERR_MEMORY; + desc->name = local_server_name_; + desc->protocol = "cxi"; + for (auto& entry : context_list_) { + TransferMetadata::DeviceDesc device_desc; + device_desc.name = entry->deviceName(); + device_desc.lid = entry->lid(); + device_desc.gid = entry->gid(); + desc->devices.push_back(device_desc); + } + desc->topology = *(local_topology_.get()); + metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_, + std::move(desc)); + return 0; +} + +int CxiTransport::registerLocalMemoryBatch( + const std::vector& buffer_list, + const std::string& location) { + std::vector> results; + for (auto& buffer : buffer_list) { + results.emplace_back( + std::async(std::launch::async, [this, buffer, location]() -> int { + return registerLocalMemoryInternal(buffer.addr, buffer.length, + location, true, false, true); + })); + } + + int first_error = 0; + for (size_t i = 0; i < buffer_list.size(); ++i) { + int ret = results[i].get(); + if (ret) { + LOG(WARNING) << "CxiTransport: Failed to register memory: addr " + << buffer_list[i].addr << " length " + << buffer_list[i].length; + if (!first_error) first_error = ret; + } + } + if (first_error) return first_error; + + return metadata_->updateLocalSegmentDesc(); +} + +int CxiTransport::unregisterLocalMemoryBatch( + const std::vector& addr_list) { + std::vector> results; + for (auto& addr : addr_list) { + results.emplace_back( + std::async(std::launch::async, [this, addr]() -> int { + return unregisterLocalMemoryInternal(addr, false, true); + })); + } + + int first_error = 0; + for (size_t i = 0; i < addr_list.size(); ++i) { + int ret = results[i].get(); + if (ret) { + LOG(WARNING) << "CxiTransport: Failed to unregister memory: addr " + << addr_list[i]; + if (!first_error) first_error = ret; + } + } + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; +} + +int CxiTransport::warmupSegment(const std::string& segment_name) { + if (!metadata_) { + LOG(ERROR) << "CxiTransport::warmupSegment: metadata_ is null"; + return ERR_INVALID_ARGUMENT; + } + if (segment_name.empty() || segment_name == local_server_name_) { + // Loopback / empty name — nothing to pre-connect. + return 0; + } + + auto desc = metadata_->getSegmentDescByName(segment_name); + if (!desc) { + LOG(ERROR) << "CxiTransport::warmupSegment: segment '" << segment_name + << "' not found in metadata (did you openSegment() first?)"; + return ERR_INVALID_ARGUMENT; + } + if (desc->devices.empty()) { + LOG(WARNING) << "CxiTransport::warmupSegment: segment '" << segment_name + << "' has no devices"; + return 0; + } + + // Build peer_nic_path list: "@" for each NIC. + std::vector peer_paths; + peer_paths.reserve(desc->devices.size()); + for (const auto& dev : desc->devices) { + peer_paths.emplace_back(segment_name + "@" + dev.name); + } + + auto t0 = std::chrono::steady_clock::now(); + size_t n_pairs = context_list_.size() * peer_paths.size(); + + // Idempotent short-circuit: if every (local_ctx, peer_nic) pair already + // has a connected endpoint, skip the whole async dispatch. Matters for + // callers that invoke warmupSegment per request loop — without this the + // 256-thread fan-out runs every time even though there is no work to do. + size_t already_ready = 0; + for (auto& ctx : context_list_) { + for (const auto& path : peer_paths) { + auto ep = ctx->peekEndpoint(path); + if (ep && ep->connected()) ++already_ready; + } + } + if (already_ready == n_pairs) { + VLOG(1) << "CxiTransport::warmupSegment('" << segment_name << "'): all " + << n_pairs << " endpoints already connected, " + << "skipping"; + return 0; + } + + // Warm up every (local_ctx, peer_nic) pair concurrently. Under the + // shared-endpoint model each warmup is just a handshake RPC + + // fi_av_insert (no fi_endpoint, no fi_enable), so the critical path is + // max(handshake RTT), not sum. We still dispatch with std::async for + // concurrency, but total wall time is typically ms-level. + std::vector> futs; + futs.reserve(n_pairs); + for (auto& ctx : context_list_) { + for (const auto& path : peer_paths) { + futs.emplace_back( + std::async(std::launch::async, [ctx, path]() -> int { + auto ep = ctx->endpoint(path); + if (!ep) { + LOG(WARNING) << "warmupSegment: endpoint() returned " + "null for " + << path; + return -1; + } + if (ep->connected()) return 0; + int rc = ep->setupConnectionsByActive(); + if (rc != 0) { + // Handshake failed: drop the peer handle so the AV + // slot is freed and the next warmup retry starts + // clean. Cheap under the shared-endpoint model — + // no fi_endpoint teardown required. + ctx->deleteEndpoint(path); + } + return rc; + })); + } + } + int ok = 0, fail = 0; + for (auto& f : futs) { + int rc = f.get(); + if (rc == 0) + ++ok; + else + ++fail; + } + auto elapsed = + std::chrono::duration(std::chrono::steady_clock::now() - t0) + .count(); + LOG(INFO) << "CxiTransport::warmupSegment('" << segment_name << "'): " << ok + << "/" << n_pairs << " endpoints connected (" << fail + << " failed) in " << elapsed << "s (" << context_list_.size() + << " local NICs x " << peer_paths.size() << " peer NICs)"; + return fail == 0 ? 0 : ERR_ENDPOINT; +} + +Status CxiTransport::submitTransfer( + BatchID batch_id, const std::vector& entries) { + auto& batch_desc = *((BatchDesc*)(batch_id)); + if (batch_desc.task_list.size() + entries.size() > batch_desc.batch_size) { + LOG(ERROR) << "CxiTransport: Exceed the limitation of current batch's " + "capacity"; + return Status::InvalidArgument( + "CxiTransport: Exceed the limitation of capacity, batch id: " + + std::to_string(batch_id)); + } + + size_t task_id = batch_desc.task_list.size(); + batch_desc.task_list.resize(task_id + entries.size()); + std::vector task_list; + for (auto& task : batch_desc.task_list) task_list.push_back(&task); + return submitTransferTask(task_list); +} + +Status CxiTransport::submitTransferTask( + const std::vector& task_list) { + std::unordered_map, std::vector> + slices_to_post; + auto local_segment_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); + assert(local_segment_desc.get()); + const int kMaxRetryCount = globalConfig().retry_cnt; + + for (size_t index = 0; index < task_list.size(); ++index) { + assert(task_list[index]); + auto& task = *task_list[index]; + assert(task.request); + auto& request = *task.request; + + if (request.length == 0) continue; + + // Find which buffer and preferred device covers this request + int request_buffer_id = -1, request_device_id = -1; + if (selectDevice(local_segment_desc.get(), (uint64_t)request.source, + request.length, request_buffer_id, + request_device_id)) { + request_buffer_id = -1; + request_device_id = -1; + } + + if (request_buffer_id >= 0 && request_device_id >= 0) { + // One slice per request. Round-robin NIC selection is + // handled by selectDevice above. + auto& context = context_list_[request_device_id]; + if (!context || !context->active()) { + LOG(ERROR) << "CXI Device " << request_device_id + << " is not active"; + return Status::InvalidArgument( + "CXI Device " + std::to_string(request_device_id) + + " is not active"); + } + + Slice* slice = getSliceCache().allocate(); + assert(slice); + slice->peer_nic_path.clear(); + slice->rdma.dest_rkey = 0; // will be set later + + slice->source_addr = (char*)request.source; + slice->length = request.length; + slice->opcode = request.opcode; + slice->rdma.dest_addr = request.target_offset; + slice->rdma.retry_cnt = request.advise_retry_cnt; + slice->rdma.max_retry_cnt = kMaxRetryCount; + slice->task = &task; + slice->target_id = request.target_id; + slice->status = Slice::PENDING; + slice->ts = 0; + task.slice_list.push_back(slice); + + slice->rdma.source_lkey = + local_segment_desc->buffers[request_buffer_id] + .lkey[request_device_id]; + slices_to_post[context].push_back(slice); + __sync_fetch_and_add(&task.total_bytes, slice->length); + __sync_fetch_and_add(&task.slice_count, 1); + } else { + // FALLBACK: device not found via initial selectDevice. + // Try per-slice retry with increasing retry_cnt to find any + // available device (handles edge cases like multi-buffer spans). + int buffer_id = -1, device_id = -1; + int retry_cnt = request.advise_retry_cnt; + bool found_device = false; + while (retry_cnt < kMaxRetryCount && !found_device) { + if (selectDevice(local_segment_desc.get(), + (uint64_t)request.source, request.length, + buffer_id, device_id, retry_cnt++)) + continue; + if (device_id >= 0 && + static_cast(device_id) < context_list_.size() && + context_list_[device_id] && + context_list_[device_id]->active()) { + found_device = true; + break; + } + } + if (!found_device) { + LOG(ERROR) << "Memory region not registered by any active CXI " + "device(s): " + << request.source; + for (auto& entry : slices_to_post) + for (auto s : entry.second) s->markFailed(); + return Status::AddressNotRegistered( + "Memory region not registered by any active CXI " + "device(s): " + + std::to_string( + reinterpret_cast(request.source))); + } + + // Found a device via retry — create single slice + Slice* slice = getSliceCache().allocate(); + assert(slice); + slice->peer_nic_path.clear(); + slice->rdma.dest_rkey = 0; + + slice->source_addr = (char*)request.source; + slice->length = request.length; + slice->opcode = request.opcode; + slice->rdma.dest_addr = request.target_offset; + slice->rdma.retry_cnt = request.advise_retry_cnt; + slice->rdma.max_retry_cnt = kMaxRetryCount; + slice->task = &task; + slice->target_id = request.target_id; + slice->status = Slice::PENDING; + slice->ts = 0; + task.slice_list.push_back(slice); + + auto& context = context_list_[device_id]; + slice->rdma.source_lkey = + local_segment_desc->buffers[buffer_id].lkey[device_id]; + slices_to_post[context].push_back(slice); + __sync_fetch_and_add(&task.total_bytes, slice->length); + __sync_fetch_and_add(&task.slice_count, 1); + } + } + + for (auto& entry : slices_to_post) + if (!entry.second.empty()) entry.first->submitPostSend(entry.second); + return Status::OK(); +} + +Status CxiTransport::getTransferStatus(BatchID batch_id, + std::vector& status) { + auto& batch_desc = *((BatchDesc*)(batch_id)); + const size_t task_count = batch_desc.task_list.size(); + status.resize(task_count); + for (size_t task_id = 0; task_id < task_count; task_id++) { + auto& task = batch_desc.task_list[task_id]; + status[task_id].transferred_bytes = task.transferred_bytes; + uint64_t success_slice_count = task.success_slice_count; + uint64_t failed_slice_count = task.failed_slice_count; + if (success_slice_count + failed_slice_count == task.slice_count) { + if (failed_slice_count) + status[task_id].s = TransferStatusEnum::FAILED; + else + status[task_id].s = TransferStatusEnum::COMPLETED; + task.is_finished = true; + } else { + status[task_id].s = TransferStatusEnum::WAITING; + } + } + return Status::OK(); +} + +Status CxiTransport::getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus& status) { + auto& batch_desc = *((BatchDesc*)(batch_id)); + const size_t task_count = batch_desc.task_list.size(); + if (task_id >= task_count) { + return Status::InvalidArgument( + "CxiTransport::getTransportStatus invalid argument, batch id: " + + std::to_string(batch_id)); + } + auto& task = batch_desc.task_list[task_id]; + status.transferred_bytes = task.transferred_bytes; + uint64_t success_slice_count = task.success_slice_count; + uint64_t failed_slice_count = task.failed_slice_count; + if (success_slice_count + failed_slice_count == task.slice_count) { + if (failed_slice_count) + status.s = TransferStatusEnum::FAILED; + else + status.s = TransferStatusEnum::COMPLETED; + task.is_finished = true; + } else { + status.s = TransferStatusEnum::WAITING; + } + return Status::OK(); +} + +CxiTransport::SegmentID CxiTransport::getSegmentID( + const std::string& segment_name) { + return metadata_->getSegmentID(segment_name); +} + +int CxiTransport::onSetupCxiConnections(const HandShakeDesc& peer_desc, + HandShakeDesc& local_desc) { + auto local_nic_name = getNicNameFromNicPath(peer_desc.peer_nic_path); + if (local_nic_name.empty()) return ERR_INVALID_ARGUMENT; + + // Find context by device name instead of using hca_list index, since + // context_list_ only contains CXI devices and may have different + // indexing than the full hca_list. + std::shared_ptr context; + for (auto& entry : context_list_) { + if (entry->deviceName() == local_nic_name) { + context = entry; + break; + } + } + if (!context) return ERR_INVALID_ARGUMENT; + + auto endpoint = context->endpoint(peer_desc.local_nic_path); + if (!endpoint) return ERR_ENDPOINT; + return endpoint->setupConnectionsByPassive(peer_desc, local_desc); +} + +int CxiTransport::initializeCxiResources() { + auto hca_list = local_topology_->getHcaList(); + + std::vector cxi_devices; + std::vector non_cxi_devices; + for (auto& device_name : hca_list) { + if (device_name.find("cxi") != std::string::npos) { + cxi_devices.push_back(device_name); + } else { + non_cxi_devices.push_back(device_name); + } + } + + if (cxi_devices.empty()) { + LOG(WARNING) << "CxiTransport: No CXI devices found, falling back to " + "all devices"; + cxi_devices = hca_list; + non_cxi_devices.clear(); + } + + // Disable non-CXI devices + for (auto& device_name : non_cxi_devices) { + local_topology_->disableDevice(device_name); + LOG(INFO) << "CxiTransport: Disabled non-CXI device " << device_name + << " in topology"; + } + + for (auto& device_name : cxi_devices) { + auto context = std::make_shared(*this, device_name); + auto& config = globalConfig(); + int ret = context->construct(config.num_cq_per_ctx, config.max_cqe, + config.max_ep_per_ctx); + if (ret) { + local_topology_->disableDevice(device_name); + LOG(WARNING) << "CxiTransport: Disable device " << device_name; + } else { + context_list_.push_back(context); + LOG(INFO) << "CxiTransport: Initialized CXI device " << device_name; + } + } + if (context_list_.empty()) { + LOG(ERROR) << "CxiTransport: No available CXI devices"; + return ERR_DEVICE_NOT_FOUND; + } + + return 0; +} + +int CxiTransport::startHandshakeDaemon(std::string& local_server_name) { + return metadata_->startHandshakeDaemon( + std::bind(&CxiTransport::onSetupCxiConnections, this, + std::placeholders::_1, std::placeholders::_2), + metadata_->localRpcMeta().rpc_port, metadata_->localRpcMeta().sockfd); +} + +int CxiTransport::selectDevice(SegmentDesc* desc, uint64_t offset, + size_t length, std::string_view hint, + int& buffer_id, int& device_id, + int retry_count) { + if (desc == nullptr) return ERR_ADDRESS_NOT_REGISTERED; + const auto& buffers = desc->buffers; + for (buffer_id = 0; buffer_id < static_cast(buffers.size()); + ++buffer_id) { + const auto& buffer = buffers[buffer_id]; + + if (offset < buffer.addr || length > buffer.length || + offset - buffer.addr > buffer.length - length) { + continue; + } + + int num_devices = static_cast(desc->devices.size()); + // this must be <=, because the first attempt will be random, so to + // guarantee that the device is found we need num_device+1 attempts + for (int attempt = 0; attempt <= num_devices; ++attempt) { + int try_count = retry_count + attempt; + device_id = + hint.empty() + ? desc->topology.selectDevice(buffer.name, try_count) + : desc->topology.selectDevice(buffer.name, hint, try_count); + if (device_id >= 0 && + static_cast(device_id) < buffer.rkey.size() && + buffer.rkey[device_id] != 0) { + return 0; + } + device_id = hint.empty() ? desc->topology.selectDevice( + kWildcardLocation, try_count) + : desc->topology.selectDevice( + kWildcardLocation, hint, try_count); + if (device_id >= 0 && + static_cast(device_id) < buffer.rkey.size() && + buffer.rkey[device_id] != 0) { + return 0; + } + } + } + return ERR_ADDRESS_NOT_REGISTERED; +} + +int CxiTransport::selectDevice(SegmentDesc* desc, uint64_t offset, + size_t length, int& buffer_id, int& device_id, + int retry_count) { + return selectDevice(desc, offset, length, "", buffer_id, device_id, + retry_count); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp b/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp index a12c649abb..23984f7d47 100644 --- a/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp +++ b/mooncake-transfer-engine/src/transport/cxl_transport/cxl_transport.cpp @@ -272,15 +272,23 @@ int CxlTransport::unregisterLocalMemory(void *addr, bool update_metadata) { int CxlTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int CxlTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, false); - return metadata_->updateLocalSegmentDesc(); + int first_error = 0; + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret && !first_error) first_error = ret; + } + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status CxlTransport::getTransferStatus(BatchID batch_id, size_t task_id, diff --git a/mooncake-transfer-engine/src/transport/device/CMakeLists.txt b/mooncake-transfer-engine/src/transport/device/CMakeLists.txt index 20ec48c850..8e13979486 100644 --- a/mooncake-transfer-engine/src/transport/device/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/device/CMakeLists.txt @@ -1,22 +1,32 @@ -# Device transport sources — conditional on GPU vendor. -# mlx5gda.cpp (IBGDA QP lifecycle) is compiled directly into the -# device_transport OBJECT library, like every other transport module, so its -# mlx5gda_* symbols flow into transfer_engine without a separate static lib. -# This keeps them visible to consumers that link transfer_engine via hand- -# written ldflags (Go p2p-store / mooncake-store), which bypass CMake's -# target_link_libraries propagation. +# Device transport sources — conditional on GPU vendor. mlx5gda.cpp (IBGDA QP +# lifecycle) is compiled directly into the device_transport OBJECT library, like +# every other transport module, so its mlx5gda_* symbols flow into +# transfer_engine without a separate static lib. This keeps them visible to +# consumers that link transfer_engine via hand- written ldflags (Go p2p-store / +# mooncake-store), which bypass CMake's target_link_libraries propagation. set(DEVICE_TRANSPORT_SOURCES p2p_device_transport.cpp) -if(USE_CUDA OR USE_MUSA) +if((USE_CUDA OR USE_MUSA) AND NOT USE_CXI) list(APPEND DEVICE_TRANSPORT_SOURCES ibgda_device_transport.cpp mlx5gda.cpp) +elseif(USE_MACA) + list(APPEND DEVICE_TRANSPORT_SOURCES ibgda_device_transport_maca_stub.cpp) +endif() +if(USE_NCCL_DEVICE) + list(APPEND DEVICE_TRANSPORT_SOURCES nccl_device_transport.cpp) endif() add_library(device_transport OBJECT ${DEVICE_TRANSPORT_SOURCES}) target_include_directories(device_transport PRIVATE ${CMAKE_SOURCE_DIR}/include) if(USE_CUDA) - target_include_directories(device_transport PRIVATE /usr/local/cuda/include) + target_include_directories(device_transport PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) +endif() +if(USE_NCCL_DEVICE) + target_link_libraries(device_transport PRIVATE NCCL::nccl) endif() if(USE_MUSA) - target_include_directories(device_transport PRIVATE /usr/local/musa/include) - target_compile_definitions(device_transport PRIVATE USE_MUSA) + target_include_directories(device_transport PRIVATE /usr/local/musa/include) + target_compile_definitions(device_transport PRIVATE USE_MUSA) +endif() +if(USE_MACA) + target_compile_definitions(device_transport PRIVATE USE_MACA) endif() diff --git a/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp b/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp index 2dcdac17dc..18533b0fee 100644 --- a/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp +++ b/mooncake-transfer-engine/src/transport/device/ibgda_device_transport.cpp @@ -21,15 +21,20 @@ #include "transport/device/device_transport.h" #include +#include #include #include #include +#include +#include #include #include +#include #include "cuda_alike.h" #include "transport/device/ibgda/memheap.h" +#include "transport/device/ibgda/mlx5_ifc.h" #include "transport/device/ibgda/mlx5gda.h" #include "topology.h" @@ -38,6 +43,46 @@ namespace device { static constexpr size_t kCtrlBufSize = 1024ULL * 1024 * 1024; // 1 GiB +enum class ControlMemoryMode { + kGpuDmabuf, + kGpuVa, + kHostMapped, +}; + +static const char* controlMemoryModeName(ControlMemoryMode mode) { + switch (mode) { + case ControlMemoryMode::kGpuDmabuf: + return "gpu-dmabuf"; + case ControlMemoryMode::kGpuVa: + return "gpu-va"; + case ControlMemoryMode::kHostMapped: + return "host-mapped"; + } + return "unknown"; +} + +#if defined(USE_CUDA) && defined(MOONCAKE_HAVE_MLX5_DMABUF_UMEM) +using Mlx5DevxUmemRegEx = mlx5dv_devx_umem* (*)(ibv_context*, + mlx5dv_devx_umem_in*); + +static Mlx5DevxUmemRegEx dmabufUmemRegEx() { + static const Mlx5DevxUmemRegEx reg_ex = [] { + dlerror(); + void* symbol = + dlvsym(RTLD_DEFAULT, "mlx5dv_devx_umem_reg_ex", "MLX5_1.19"); + if (!symbol) { + const char* error = dlerror(); + LOG(WARNING) << "[EP IBGDA] Runtime libmlx5 does not provide " + "mlx5dv_devx_umem_reg_ex@MLX5_1.19; disabling " + "DMA-BUF control memory" + << (error ? std::string(": ") + error : ""); + } + return reinterpret_cast(symbol); + }(); + return reg_ex; +} +#endif + // Check if IPv6 address is IPv4-mapped (::ffff:x.x.x.x) static bool isIpv4Mapped(const struct in6_addr* a) { return ((a->s6_addr32[0] | a->s6_addr32[1]) == 0 && @@ -199,38 +244,40 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { } int allocateControlBuffer() override { - cudaError_t err = cudaMalloc(&ctrl_buf_, kCtrlBufSize); - if (err != cudaSuccess) { - LOG(ERROR) << "[EP IBGDA] cudaMalloc ctrl_buf failed: " - << cudaGetErrorString(err); - return -1; - } - - ctrl_buf_umem_ = mlx5dv_devx_umem_reg(ctx_, ctrl_buf_, kCtrlBufSize, - IBV_ACCESS_LOCAL_WRITE); - if (!ctrl_buf_umem_) { - LOG(ERROR) << "[EP IBGDA] mlx5dv_devx_umem_reg failed (errno=" - << errno << ")"; - return -1; - } - LOG(INFO) << "[EP IBGDA] ctrl_buf UMEM registered via VA path"; - - ctrl_buf_heap_ = memheap_create(kCtrlBufSize); - if (!ctrl_buf_heap_) { - LOG(ERROR) << "[EP IBGDA] memheap_create failed"; - return -1; - } - return 0; +#if defined(USE_CUDA) + if (cudaSupportsDmabuf()) { + ctrl_buf_mode_ = ControlMemoryMode::kGpuDmabuf; + LOG(INFO) << "[EP IBGDA] Selected per-QP GPU DMA-BUF control " + "regions"; + return 0; + } + LOG(WARNING) + << "[EP IBGDA] CUDA DMA-BUF control memory is unavailable; " + "trying GPU-VA control buffer"; + return allocateGpuVaOrHostControlBuffer(); +#endif + return allocateControlBuffer(ControlMemoryMode::kGpuVa); } int createQueuePairs(void* stream_ptr) override { auto stream = static_cast(stream_ptr); + mlx5gda_control_region_allocator region_allocator{ + .context = this, + .allocate = allocateDmabufControlRegionThunk, + .release = releaseDmabufControlRegionThunk, + }; + const mlx5gda_control_region_allocator* allocator = + ctrl_buf_mode_ == ControlMemoryMode::kGpuDmabuf ? ®ion_allocator + : nullptr; for (int i = 0; i < num_qps_; ++i) { - mlx5gda_qp* qp = - mlx5gda_create_rc_qp(mpd_, ctrl_buf_, ctrl_buf_umem_, - ctrl_buf_heap_, pd_, 16384, 1, stream); + mlx5gda_qp* qp = mlx5gda_create_rc_qp( + mpd_, ctrl_buf_, ctrl_buf_umem_, ctrl_buf_heap_, pd_, 16384, 1, + stream, allocator); if (!qp) { + const int qp_errno = errno; LOG(ERROR) << "[EP IBGDA] mlx5gda_create_rc_qp failed at " << i; + if (retryControlBuffer(qp_errno)) + return createQueuePairs(stream_ptr); return -1; } if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) { @@ -239,31 +286,50 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { return -1; } cudaStreamSynchronize(stream); + const bool split_regions = + ctrl_buf_mode_ == ControlMemoryMode::kGpuDmabuf; mlx5gda_qp_devctx devctx{ .qpn = qp->qpn, .wqeid_mask = qp->num_wqebb - 1, - .wq = reinterpret_cast( - static_cast(ctrl_buf_) + qp->wq_offset), - .cq = reinterpret_cast( - static_cast(ctrl_buf_) + qp->send_cq->cq_offset), - .dbr = reinterpret_cast( - static_cast(ctrl_buf_) + qp->dbr_offset), + .wq = split_regions ? reinterpret_cast(qp->wq) + : reinterpret_cast( + static_cast(ctrl_buf_dev_) + + qp->wq_offset), + .cq = split_regions + ? reinterpret_cast(qp->send_cq->cq_buf) + : reinterpret_cast( + static_cast(ctrl_buf_dev_) + + qp->send_cq->cq_offset), + .dbr = split_regions + ? reinterpret_cast(qp->dbr) + : reinterpret_cast( + static_cast(ctrl_buf_dev_) + + qp->dbr_offset), .bf = static_cast(qp->uar->reg_addr), }; cudaMemcpy( static_cast(qp_devctxs_) + i * sizeof(mlx5gda_qp_devctx), &devctx, sizeof(mlx5gda_qp_devctx), cudaMemcpyHostToDevice); + if (i == 0 && ctrl_buf_mode_ == ControlMemoryMode::kGpuDmabuf) { + LOG(INFO) << "[EP IBGDA] QP 0 DMA-BUF regions: cq_umem=" + << qp->send_cq->cq_region.umem->umem_id + << " cq_dbr_umem=" + << qp->send_cq->dbr_region.umem->umem_id + << " wq_umem=" << qp->wq_region.umem->umem_id + << " qp_dbr_umem=" << qp->dbr_region.umem->umem_id + << " offsets=0"; + } qps_.push_back(qp); } + if (ctrl_buf_mode_ == ControlMemoryMode::kGpuDmabuf) { + LOG(INFO) << "[EP IBGDA] Control buffer mode: gpu-dmabuf" + << " layout=per-qp-regions qps=" << qps_.size(); + } return 0; } int recreateQueuePairs(void* stream_ptr) override { - auto stream = static_cast(stream_ptr); - for (auto* qp : qps_) { - if (qp) mlx5gda_destroy_qp(ctrl_buf_heap_, qp); - } - qps_.clear(); + destroyQueuePairs(); return createQueuePairs(stream_ptr); } @@ -349,11 +415,278 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { int gidIndex() const override { return gid_index_; } private: - void teardown() { + static bool isCreateQpBadParam(const mlx5gda_create_qp_failure& failure) { + return failure.valid && failure.status == MLX5_CMD_STAT_BAD_PARAM_ERR; + } + + bool cudaSupportsDmabuf() const { +#if defined(USE_CUDA) && defined(MOONCAKE_HAVE_MLX5_DMABUF_UMEM) + if (!dmabufUmemRegEx()) return false; + CUdevice device; + CUresult result = cuCtxGetDevice(&device); + if (result != CUDA_SUCCESS) { + LOG(WARNING) << "[EP IBGDA] cuCtxGetDevice failed while probing " + "DMA-BUF control support: " + << cudaDriverError(result); + return false; + } + + int supported = 0; + result = cuDeviceGetAttribute( + &supported, CU_DEVICE_ATTRIBUTE_DMA_BUF_SUPPORTED, device); + if (result != CUDA_SUCCESS) { + LOG(WARNING) << "[EP IBGDA] CUDA DMA-BUF capability query failed: " + << cudaDriverError(result); + return false; + } + if (!supported) { + LOG(INFO) << "[EP IBGDA] Active CUDA device does not support " + "DMA-BUF control memory"; + } + return supported != 0; +#else + LOG(WARNING) << "[EP IBGDA] DMA-BUF control memory was not compiled: " + "the build-time mlx5 headers or libmlx5 lack the " + "required DevX UMEM API"; + return false; +#endif + } + + static std::string cudaDriverError(CUresult result) { +#if defined(USE_CUDA) + const char* name = nullptr; + const char* message = nullptr; + cuGetErrorName(result, &name); + cuGetErrorString(result, &message); + return std::string(name ? name : "CUDA_ERROR_UNKNOWN") + ": " + + (message ? message : "unknown CUDA driver error"); +#else + return std::to_string(static_cast(result)); +#endif + } + + int allocateDmabufControlRegion(size_t requested_size, + mlx5gda_control_region* region) { +#if defined(USE_CUDA) && defined(MOONCAKE_HAVE_MLX5_DMABUF_UMEM) + if (!region || requested_size == 0) { + errno = EINVAL; + return -1; + } + + static const long page_size = sysconf(_SC_PAGESIZE); + if (page_size <= 0) { + LOG(ERROR) << "[EP IBGDA] Failed to query host page size"; + errno = EIO; + return -1; + } + const size_t page_mask = static_cast(page_size) - 1; + const size_t size = (requested_size + page_mask) & ~page_mask; + + void* addr = nullptr; + cudaError_t cuda_error = cudaMalloc(&addr, size); + if (cuda_error != cudaSuccess) { + LOG(ERROR) << "[EP IBGDA] cudaMalloc failed for DMA-BUF control " + "region size=" + << size << ": " << cudaGetErrorString(cuda_error); + errno = cuda_error == cudaErrorMemoryAllocation ? ENOMEM : EIO; + return -1; + } + + int sync_memops = 1; + CUresult result = cuPointerSetAttribute( + &sync_memops, CU_POINTER_ATTRIBUTE_SYNC_MEMOPS, + reinterpret_cast(addr)); + if (result != CUDA_SUCCESS) { + LOG(ERROR) << "[EP IBGDA] Failed to enable synchronous memory " + "operations for DMA-BUF control memory: " + << cudaDriverError(result); + cudaFree(addr); + errno = EIO; + return -1; + } + + int dmabuf_fd = -1; + result = cuMemGetHandleForAddressRange( + &dmabuf_fd, reinterpret_cast(addr), size, + CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, 0); + if (result != CUDA_SUCCESS) { + LOG(ERROR) << "[EP IBGDA] Failed to export DMA-BUF control memory " + "at " + << addr << " size=" << size << ": " + << cudaDriverError(result); + cudaFree(addr); + errno = EIO; + return -1; + } + + struct mlx5dv_devx_umem_in input{}; + input.addr = 0; + input.size = size; + input.access = IBV_ACCESS_LOCAL_WRITE; + input.pgsz_bitmap = + UINT64_MAX & ~(static_cast(page_size) - 1); + input.comp_mask = MLX5DV_UMEM_MASK_DMABUF; + input.dmabuf_fd = dmabuf_fd; + + mlx5dv_devx_umem* umem = dmabufUmemRegEx()(ctx_, &input); + const int registration_errno = errno; + if (close(dmabuf_fd) != 0) { + PLOG(WARNING) << "[EP IBGDA] Failed to close control DMA-BUF fd"; + } + if (!umem) { + errno = registration_errno; + PLOG(ERROR) << "[EP IBGDA] mlx5dv_devx_umem_reg_ex failed for " + "GPU DMA-BUF control memory"; + cudaFree(addr); + return -1; + } + *region = mlx5gda_control_region{ + .addr = addr, + .size = size, + .umem = umem, + }; + return 0; +#else + (void)requested_size; + (void)region; + errno = ENOTSUP; + return -1; +#endif + } + + void releaseDmabufControlRegion(mlx5gda_control_region* region) { + if (!region) return; + if (region->umem) mlx5dv_devx_umem_dereg(region->umem); + if (region->addr) cudaFree(region->addr); + *region = {}; + } + + static int allocateDmabufControlRegionThunk( + void* context, size_t size, mlx5gda_control_region* region) { + return static_cast(context) + ->allocateDmabufControlRegion(size, region); + } + + static void releaseDmabufControlRegionThunk( + void* context, mlx5gda_control_region* region) { + static_cast(context) + ->releaseDmabufControlRegion(region); + } + + int allocateControlBuffer(ControlMemoryMode mode) { + ctrl_buf_mode_ = mode; + if (mode == ControlMemoryMode::kHostMapped) { + void* ptr = nullptr; + int ret = posix_memalign(&ptr, 4096, kCtrlBufSize); + if (ret != 0) { + LOG(ERROR) << "[EP IBGDA] posix_memalign ctrl_buf failed: " + << ret; + return -1; + } + ctrl_buf_ = ptr; + std::memset(ctrl_buf_, 0, kCtrlBufSize); + + cudaError_t err = cudaHostRegister( + ctrl_buf_, kCtrlBufSize, + cudaHostRegisterPortable | cudaHostRegisterMapped); + if (err != cudaSuccess) { + LOG(ERROR) << "[EP IBGDA] cudaHostRegister ctrl_buf failed: " + << cudaGetErrorString(err); + free(ctrl_buf_); + ctrl_buf_ = nullptr; + return -1; + } + + err = cudaHostGetDevicePointer(&ctrl_buf_dev_, ctrl_buf_, 0); + if (err != cudaSuccess) { + LOG(ERROR) + << "[EP IBGDA] cudaHostGetDevicePointer ctrl_buf failed: " + << cudaGetErrorString(err); + cudaHostUnregister(ctrl_buf_); + free(ctrl_buf_); + ctrl_buf_ = nullptr; + return -1; + } + } else { + cudaError_t err = cudaMalloc(&ctrl_buf_, kCtrlBufSize); + if (err != cudaSuccess) { + LOG(ERROR) << "[EP IBGDA] cudaMalloc ctrl_buf failed: " + << cudaGetErrorString(err); + return -1; + } + ctrl_buf_dev_ = ctrl_buf_; + } + + ctrl_buf_umem_ = mlx5dv_devx_umem_reg(ctx_, ctrl_buf_, kCtrlBufSize, + IBV_ACCESS_LOCAL_WRITE); + if (!ctrl_buf_umem_) { + LOG(ERROR) << "[EP IBGDA] Control UMEM registration failed for " + << controlMemoryModeName(mode) << " (errno=" << errno + << ")"; + freeControlBuffer(); + return -1; + } + LOG(INFO) << "[EP IBGDA] Control buffer mode: " + << controlMemoryModeName(mode) << " addr=" << ctrl_buf_dev_ + << " size=" << kCtrlBufSize + << " umem_id=" << ctrl_buf_umem_->umem_id; + + ctrl_buf_heap_ = memheap_create(kCtrlBufSize); + if (!ctrl_buf_heap_) { + LOG(ERROR) << "[EP IBGDA] memheap_create failed"; + freeControlBuffer(); + return -1; + } + return 0; + } + + int allocateGpuVaOrHostControlBuffer() { + if (allocateControlBuffer(ControlMemoryMode::kGpuVa) == 0) return 0; + + LOG(WARNING) << "[EP IBGDA] GPU-VA control buffer is unavailable; " + "trying host-backed mapped control buffer"; + return allocateControlBuffer(ControlMemoryMode::kHostMapped); + } + + bool retryControlBuffer(int failure_errno) { + auto failure = mlx5gda_last_create_qp_failure(); + const bool using_dmabuf_control_regions = + ctrl_buf_mode_ == ControlMemoryMode::kGpuDmabuf; + if (ctrl_buf_mode_ == ControlMemoryMode::kHostMapped || + (!isCreateQpBadParam(failure) && !using_dmabuf_control_regions)) + return false; + + if (using_dmabuf_control_regions) { + LOG(WARNING) << "[EP IBGDA] QP setup failed while using DMA-BUF " + "control regions " + "(errno=" + << failure_errno + << "); retrying with GPU-VA control buffer"; + destroyQueuePairs(); + freeControlBuffer(); + return allocateGpuVaOrHostControlBuffer() == 0; + } else { + LOG(WARNING) + << "[EP IBGDA] GPU-backed control buffer was rejected " + "by DevX CREATE_QP" + << " (status=0x" << std::hex << failure.status << " syndrome=0x" + << failure.syndrome << std::dec + << "); retrying with host-backed mapped control buffer"; + } + + destroyQueuePairs(); + freeControlBuffer(); + return allocateControlBuffer(ControlMemoryMode::kHostMapped) == 0; + } + + void destroyQueuePairs() { for (auto* qp : qps_) { if (qp) mlx5gda_destroy_qp(ctrl_buf_heap_, qp); } qps_.clear(); + } + + void freeControlBuffer() { if (ctrl_buf_heap_) { memheap_destroy(ctrl_buf_heap_); ctrl_buf_heap_ = nullptr; @@ -363,9 +696,21 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { ctrl_buf_umem_ = nullptr; } if (ctrl_buf_) { - cudaFree(ctrl_buf_); + if (ctrl_buf_mode_ == ControlMemoryMode::kHostMapped) { + cudaHostUnregister(ctrl_buf_); + free(ctrl_buf_); + } else { + cudaFree(ctrl_buf_); + } ctrl_buf_ = nullptr; + ctrl_buf_dev_ = nullptr; } + ctrl_buf_mode_ = ControlMemoryMode::kGpuVa; + } + + void teardown() { + destroyQueuePairs(); + freeControlBuffer(); if (mr_) { ibv_dereg_mr(mr_); mr_ = nullptr; @@ -406,7 +751,9 @@ class IbgdaDeviceTransportImpl : public RdmaTransport { std::vector device_filter_; // Control buffer - void* ctrl_buf_ = nullptr; // GPU VA + void* ctrl_buf_ = nullptr; // Allocation address used for UMEM registration + void* ctrl_buf_dev_ = nullptr; + ControlMemoryMode ctrl_buf_mode_ = ControlMemoryMode::kGpuVa; mlx5dv_devx_umem* ctrl_buf_umem_ = nullptr; memheap* ctrl_buf_heap_ = nullptr; diff --git a/mooncake-transfer-engine/src/transport/device/ibgda_device_transport_maca_stub.cpp b/mooncake-transfer-engine/src/transport/device/ibgda_device_transport_maca_stub.cpp new file mode 100644 index 0000000000..d59f8409e6 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/device/ibgda_device_transport_maca_stub.cpp @@ -0,0 +1,34 @@ +#include "transport/device/device_transport.h" + +namespace mooncake { +namespace device { + +class NullRdmaTransport : public RdmaTransport { + public: + int initialize(const std::string&, int, int) override { return -1; } + int registerMemory(void*, size_t) override { return -1; } + int allocateControlBuffer() override { return -1; } + int createQueuePairs(void*) override { return -1; } + int recreateQueuePairs(void*) override { return -1; } + int connectPeers(int, bool, const std::vector&, + const std::vector&, const std::vector&, + const std::vector&, const std::vector&, + const std::vector&, + const std::vector&) override { + return -1; + } + RdmaLocalMetadata localMetadata() const override { return {}; } + void* raddrsPtr() override { return nullptr; } + void* rkeysPtr() override { return nullptr; } + void* qpDevCtxsPtr() override { return nullptr; } + bool isRoce() const override { return false; } + int gidIndex() const override { return -1; } +}; + +std::unique_ptr createIbgdaDeviceTransport( + const std::vector&) { + return std::make_unique(); +} + +} // namespace device +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp b/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp index ca601b2536..37b22c0853 100644 --- a/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp +++ b/mooncake-transfer-engine/src/transport/device/mlx5gda.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -38,6 +39,48 @@ static void print_cuda_error(const char* msg) { fprintf(stderr, "%s: %s\n", msg, err_str); } +static thread_local mlx5gda_create_qp_failure g_last_create_qp_failure{}; + +void mlx5gda_reset_create_qp_failure() { g_last_create_qp_failure = {}; } + +mlx5gda_create_qp_failure mlx5gda_last_create_qp_failure() { + return g_last_create_qp_failure; +} + +static bool is_host_control_buffer(void* ptr) { + cudaPointerAttributes attr{}; + cudaError_t err = cudaPointerGetAttributes(&attr, ptr); + if (err != cudaSuccess) { + cudaGetLastError(); + return false; + } + return attr.type == cudaMemoryTypeHost; +} + +static void print_devx_create_qp_failure(const uint8_t* cmd_out, + uint32_t num_wqebb, uint8_t port_num, + uint32_t pdn, uint32_t uar_page, + uint32_t cqn, uint32_t wq_umem_id, + uint32_t dbr_umem_id, size_t wq_offset, + size_t dbr_offset, int sys_errno) { + uint32_t status = DEVX_GET(create_qp_out, cmd_out, status); + uint32_t syndrome = DEVX_GET(create_qp_out, cmd_out, syndrome); + g_last_create_qp_failure = mlx5gda_create_qp_failure{ + .valid = true, + .status = status, + .syndrome = syndrome, + .sys_errno = sys_errno, + }; + fprintf(stderr, + "mlx5dv_devx_obj_create(create_qp) failed: errno=%d (%s) " + "status=0x%x syndrome=0x%x port=%u pdn=0x%x uar_page=0x%x " + "cqn=0x%x wq_umem_id=0x%x dbr_umem_id=0x%x num_wqebb=%u " + "wq_offset=0x%zx dbr_offset=0x%zx\n", + sys_errno, strerror(sys_errno), status, syndrome, port_num, pdn, + uar_page, cqn, wq_umem_id, dbr_umem_id, num_wqebb, wq_offset, + dbr_offset); +} + // Create UAR for BF (Blue Flame) doorbell ringing. // On CUDA: registers the BF MMIO region into GPU address space so the // GPU kernel can directly write the doorbell (lowest latency). @@ -74,6 +117,20 @@ static struct mlx5dv_devx_uar* create_uar(struct ibv_context* ctx) { return uar; } +static bool uses_control_regions( + const struct mlx5gda_control_region_allocator* allocator) { + return allocator && allocator->allocate && allocator->release; +} + +static void release_control_region( + const struct mlx5gda_control_region_allocator* allocator, + struct mlx5gda_control_region* region) { + if (uses_control_regions(allocator) && + (region->addr || region->umem || region->size != 0)) { + allocator->release(allocator->context, region); + } +} + static void destroy_uar(struct mlx5dv_devx_uar* uar) { if (!uar) return; if (uar->reg_addr) { @@ -84,11 +141,11 @@ static void destroy_uar(struct mlx5dv_devx_uar* uar) { mlx5dv_devx_free_uar(uar); } -struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, - struct mlx5dv_devx_umem* ctrl_buf_umem, - struct memheap* ctrl_buf_heap, - struct ibv_pd* pd, int cqe, - cudaStream_t stream) { +struct mlx5gda_cq* mlx5gda_create_cq( + void* ctrl_buf, struct mlx5dv_devx_umem* ctrl_buf_umem, + struct memheap* ctrl_buf_heap, struct ibv_pd* pd, int cqe, + cudaStream_t stream, + const struct mlx5gda_control_region_allocator* region_allocator) { struct mlx5gda_cq* cq = NULL; struct mlx5dv_devx_uar* uar = NULL; uint32_t eqn = 0; @@ -96,9 +153,15 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, size_t dbr_offset = -1; struct mlx5dv_devx_obj* mlx5_cq = NULL; uint32_t cqn = 0; + void* cq_buf = ctrl_buf; + void* dbr = ctrl_buf; + struct mlx5dv_devx_umem* cq_umem = ctrl_buf_umem; + struct mlx5dv_devx_umem* dbr_umem = ctrl_buf_umem; + const bool split_regions = uses_control_regions(region_allocator); struct ibv_context* ctx = pd->context; void* cq_context = NULL; + bool ctrl_host = !split_regions && is_host_control_buffer(ctrl_buf); if (cqe <= 0) { errno = EINVAL; @@ -109,30 +172,59 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, uint8_t cmd_in[DEVX_ST_SZ_BYTES(create_cq_in)] = {0}; uint8_t cmd_out[DEVX_ST_SZ_BYTES(create_cq_out)] = {0}; - cq_offset = memheap_aligned_alloc(ctrl_buf_heap, - num_cqe * sizeof(struct mlx5_cqe64), - (size_t)1 << MLX5_ADAPTER_PAGE_SHIFT); - if (cq_offset == -1) { - perror("Failed to allocate CQ memory"); - goto fail; + cq = (struct mlx5gda_cq*)calloc(1, sizeof(struct mlx5gda_cq)); + if (!cq) goto fail; + + if (split_regions) { + cq->region_allocator = *region_allocator; + if (region_allocator->allocate(region_allocator->context, + num_cqe * sizeof(struct mlx5_cqe64), + &cq->cq_region) != 0) { + perror("Failed to allocate CQ control region"); + goto fail; + } + if (region_allocator->allocate(region_allocator->context, + sizeof(struct mlx5gda_cq_dbr), + &cq->dbr_region) != 0) { + perror("Failed to allocate CQ DBR control region"); + goto fail; + } + cq_buf = cq->cq_region.addr; + dbr = cq->dbr_region.addr; + cq_umem = cq->cq_region.umem; + dbr_umem = cq->dbr_region.umem; + cq_offset = 0; + dbr_offset = 0; + } else { + cq_offset = memheap_aligned_alloc(ctrl_buf_heap, + num_cqe * sizeof(struct mlx5_cqe64), + (size_t)1 << MLX5_ADAPTER_PAGE_SHIFT); + if (cq_offset == -1) { + perror("Failed to allocate CQ memory"); + goto fail; + } + dbr_offset = + memheap_alloc(ctrl_buf_heap, sizeof(struct mlx5gda_cq_dbr)); + if (dbr_offset == -1) { + perror("Failed to allocate CQ DBR memory"); + goto fail; + } } // MLX5 hardware requirement: CQE (Completion Queue Entry) must be // initialized to 0xFF (-1) to mark them as invalid. The hardware checks the // owner bit in CQE to determine if it's valid. This is mandatory for proper // CQ operation. Use async version to avoid blocking. - if (cudaMemsetAsync(ctrl_buf + cq_offset, -1, - num_cqe * sizeof(struct mlx5_cqe64), - stream) != cudaSuccess) { - print_cuda_error("Failed to memset CQ memory"); - goto fail; - } - dbr_offset = memheap_alloc(ctrl_buf_heap, sizeof(struct mlx5gda_cq_dbr)); - if (dbr_offset == -1) { - perror("Failed to allocate DBR memory"); - goto fail; + if (ctrl_host) { + memset(static_cast(cq_buf) + cq_offset, -1, + num_cqe * sizeof(struct mlx5_cqe64)); + } else { + if (cudaMemsetAsync(static_cast(cq_buf) + cq_offset, -1, + num_cqe * sizeof(struct mlx5_cqe64), + stream) != cudaSuccess) { + print_cuda_error("Failed to memset CQ memory"); + goto fail; + } } - cq = (struct mlx5gda_cq*)malloc(sizeof(struct mlx5gda_cq)); - if (!cq) goto fail; if (mlx5dv_devx_query_eqn(ctx, 0, &eqn)) { perror("Failed to query EQN"); goto fail; @@ -146,7 +238,7 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, } DEVX_SET(create_cq_in, cmd_in, opcode, MLX5_CMD_OP_CREATE_CQ); - DEVX_SET(create_cq_in, cmd_in, cq_umem_id, ctrl_buf_umem->umem_id); + DEVX_SET(create_cq_in, cmd_in, cq_umem_id, cq_umem->umem_id); DEVX_SET(create_cq_in, cmd_in, cq_umem_valid, 1); DEVX_SET64(create_cq_in, cmd_in, cq_umem_offset, cq_offset); cq_context = DEVX_ADDR_OF(create_cq_in, cmd_in, cq_context); @@ -154,7 +246,7 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, DEVX_SET(cqc, cq_context, cqe_sz, MLX5_CQE_SIZE_64B); DEVX_SET(cqc, cq_context, cc, 0x1); // collapsed cq DEVX_SET(cqc, cq_context, oi, 0x1); // cq overrun - DEVX_SET(cqc, cq_context, dbr_umem_id, ctrl_buf_umem->umem_id); + DEVX_SET(cqc, cq_context, dbr_umem_id, dbr_umem->umem_id); DEVX_SET(cqc, cq_context, log_cq_size, __builtin_ctz(num_cqe)); DEVX_SET(cqc, cq_context, uar_page, uar->page_id); DEVX_SET(cqc, cq_context, c_eqn, eqn); @@ -162,9 +254,11 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, // Synchronize stream before creating CQ object, as hardware will read the // CQE memory - if (cudaStreamSynchronize(stream) != cudaSuccess) { - print_cuda_error("Failed to synchronize stream before CQ creation"); - goto fail; + if (!ctrl_host) { + if (cudaStreamSynchronize(stream) != cudaSuccess) { + print_cuda_error("Failed to synchronize stream before CQ creation"); + goto fail; + } } mlx5_cq = mlx5dv_devx_obj_create(ctx, cmd_in, sizeof(cmd_in), cmd_out, @@ -177,6 +271,8 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, cq->cq_offset = cq_offset; cq->dbr_offset = dbr_offset; + cq->cq_buf = static_cast(cq_buf) + cq_offset; + cq->dbr = static_cast(dbr) + dbr_offset; cq->cqe = num_cqe; cq->cqn = cqn; cq->uar = uar; @@ -185,7 +281,15 @@ struct mlx5gda_cq* mlx5gda_create_cq(void* ctrl_buf, fail: int saved_errno = errno; if (uar) mlx5dv_devx_free_uar(uar); - if (cq) free(cq); + if (cq) { + release_control_region(&cq->region_allocator, &cq->dbr_region); + release_control_region(&cq->region_allocator, &cq->cq_region); + free(cq); + } + if (!split_regions) { + if (cq_offset != -1) memheap_free(ctrl_buf_heap, cq_offset); + if (dbr_offset != -1) memheap_free(ctrl_buf_heap, dbr_offset); + } errno = saved_errno; return NULL; } @@ -198,27 +302,40 @@ void mlx5gda_destroy_cq(struct memheap* ctrl_buf_heap, struct mlx5gda_cq* cq) { if (cq->uar) { mlx5dv_devx_free_uar(cq->uar); } - memheap_free(ctrl_buf_heap, cq->cq_offset); - memheap_free(ctrl_buf_heap, cq->dbr_offset); + if (uses_control_regions(&cq->region_allocator)) { + release_control_region(&cq->region_allocator, &cq->dbr_region); + release_control_region(&cq->region_allocator, &cq->cq_region); + } else { + memheap_free(ctrl_buf_heap, cq->cq_offset); + memheap_free(ctrl_buf_heap, cq->dbr_offset); + } free(cq); } -struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, - struct mlx5dv_devx_umem* ctrl_buf_umem, - struct memheap* ctrl_buf_heap, - struct ibv_pd* pd, int wqe, - uint8_t port_num, cudaStream_t stream) { +struct mlx5gda_qp* mlx5gda_create_rc_qp( + struct mlx5dv_pd mpd, void* ctrl_buf, + struct mlx5dv_devx_umem* ctrl_buf_umem, struct memheap* ctrl_buf_heap, + struct ibv_pd* pd, int wqe, uint8_t port_num, cudaStream_t stream, + const struct mlx5gda_control_region_allocator* region_allocator) { + mlx5gda_reset_create_qp_failure(); + struct mlx5gda_qp* qp = NULL; struct mlx5gda_cq* send_cq = NULL; struct mlx5dv_devx_uar* uar = NULL; struct mlx5dv_devx_obj* mlx5_qp = NULL; size_t wq_offset = -1; size_t dbr_offset = -1; + void* wq = ctrl_buf; + void* dbr = ctrl_buf; + struct mlx5dv_devx_umem* wq_umem = ctrl_buf_umem; + struct mlx5dv_devx_umem* dbr_umem = ctrl_buf_umem; + const bool split_regions = uses_control_regions(region_allocator); struct ibv_context* ctx = pd->context; void* qp_context = NULL; void* cap = NULL; uint32_t cqe_version = 0; + bool ctrl_host = !split_regions && is_host_control_buffer(ctrl_buf); if (wqe <= 0) { errno = EINVAL; @@ -237,6 +354,7 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, perror("Failed to allocate QP memory"); goto fail; } + if (split_regions) qp->region_allocator = *region_allocator; qp->port_num = port_num; if (ibv_query_port(ctx, port_num, &qp->port_attr) != 0) { @@ -265,7 +383,7 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, // Create send_cq on GPU memory. send_cq = mlx5gda_create_cq(ctrl_buf, ctrl_buf_umem, ctrl_buf_heap, pd, wqe, - stream); + stream, region_allocator); if (send_cq == NULL) { perror("mlx5gda_create_cq failed"); goto fail; @@ -277,29 +395,57 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, goto fail; } - wq_offset = memheap_aligned_alloc(ctrl_buf_heap, - num_wqebb * sizeof(struct mlx5gda_wqebb), - (size_t)1 << MLX5_ADAPTER_PAGE_SHIFT); - if (wq_offset == -1) { - perror("Failed to allocate WQ memory"); - goto fail; - } + if (split_regions) { + if (region_allocator->allocate(region_allocator->context, + num_wqebb * sizeof(struct mlx5gda_wqebb), + &qp->wq_region) != 0) { + perror("Failed to allocate WQ control region"); + goto fail; + } + if (region_allocator->allocate(region_allocator->context, + sizeof(struct mlx5gda_wq_dbr), + &qp->dbr_region) != 0) { + perror("Failed to allocate QP DBR control region"); + goto fail; + } + wq = qp->wq_region.addr; + dbr = qp->dbr_region.addr; + wq_umem = qp->wq_region.umem; + dbr_umem = qp->dbr_region.umem; + wq_offset = 0; + dbr_offset = 0; + } else { + wq_offset = memheap_aligned_alloc( + ctrl_buf_heap, num_wqebb * sizeof(struct mlx5gda_wqebb), + (size_t)1 << MLX5_ADAPTER_PAGE_SHIFT); + if (wq_offset == -1) { + perror("Failed to allocate WQ memory"); + goto fail; + } - dbr_offset = memheap_alloc(ctrl_buf_heap, sizeof(struct mlx5gda_wq_dbr)); - if (dbr_offset == -1) { - perror("Failed to allocate DBR memory"); - goto fail; + dbr_offset = + memheap_alloc(ctrl_buf_heap, sizeof(struct mlx5gda_wq_dbr)); + if (dbr_offset == -1) { + perror("Failed to allocate DBR memory"); + goto fail; + } } // DBR must be zero-initialized. Use async version to avoid blocking. - if (cudaMemsetAsync(ctrl_buf + dbr_offset, 0, sizeof(struct mlx5gda_wq_dbr), - stream) != cudaSuccess) { - print_cuda_error("Failed to zero DBR memory"); - goto fail; + if (ctrl_host) { + memset(static_cast(dbr) + dbr_offset, 0, + sizeof(struct mlx5gda_wq_dbr)); + } else { + if (cudaMemsetAsync(static_cast(dbr) + dbr_offset, 0, + sizeof(struct mlx5gda_wq_dbr), + stream) != cudaSuccess) { + print_cuda_error("Failed to zero DBR memory"); + goto fail; + } } DEVX_SET(create_qp_in, cmd_in, opcode, MLX5_CMD_OP_CREATE_QP); DEVX_SET(create_qp_in, cmd_in, wq_umem_id, - ctrl_buf_umem->umem_id); // WQ buffer + wq_umem->umem_id); // WQ buffer DEVX_SET64(create_qp_in, cmd_in, wq_umem_offset, wq_offset); DEVX_SET(create_qp_in, cmd_in, wq_umem_valid, 1); // Enable wq_umem_id @@ -320,20 +466,25 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, dbr_offset); // Offset of dbr_umem_id (behavior changed because // of dbr_umem_valid) DEVX_SET(qpc, qp_context, dbr_umem_id, - ctrl_buf_umem->umem_id); // DBR buffer + dbr_umem->umem_id); // DBR buffer DEVX_SET(qpc, qp_context, user_index, 0); DEVX_SET(qpc, qp_context, page_offset, 0); // Synchronize stream before creating QP object, as hardware will read the // DBR memory - if (cudaStreamSynchronize(stream) != cudaSuccess) { - print_cuda_error("Failed to synchronize stream before QP creation"); - goto fail; + if (!ctrl_host) { + if (cudaStreamSynchronize(stream) != cudaSuccess) { + print_cuda_error("Failed to synchronize stream before QP creation"); + goto fail; + } } mlx5_qp = mlx5dv_devx_obj_create(ctx, cmd_in, sizeof(cmd_in), cmd_out, sizeof(cmd_out)); if (mlx5_qp == NULL) { + print_devx_create_qp_failure( + cmd_out, num_wqebb, port_num, mpd.pdn, uar->page_id, send_cq->cqn, + wq_umem->umem_id, dbr_umem->umem_id, wq_offset, dbr_offset, errno); goto fail; } @@ -345,6 +496,8 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, qp->num_wqebb = num_wqebb; qp->wq_offset = wq_offset; qp->dbr_offset = dbr_offset; + qp->wq = static_cast(wq) + wq_offset; + qp->dbr = static_cast(dbr) + dbr_offset; return qp; fail: @@ -358,13 +511,17 @@ struct mlx5gda_qp* mlx5gda_create_rc_qp(struct mlx5dv_pd mpd, void* ctrl_buf, if (send_cq) { mlx5gda_destroy_cq(ctrl_buf_heap, send_cq); } + if (qp) { + release_control_region(&qp->region_allocator, &qp->dbr_region); + release_control_region(&qp->region_allocator, &qp->wq_region); + } if (qp) { free(qp); } - if (wq_offset != -1) { + if (!split_regions && wq_offset != -1) { memheap_free(ctrl_buf_heap, wq_offset); } - if (dbr_offset != -1) { + if (!split_regions && dbr_offset != -1) { memheap_free(ctrl_buf_heap, dbr_offset); } errno = saved_errno; @@ -381,11 +538,16 @@ void mlx5gda_destroy_qp(struct memheap* ctrl_buf_heap, struct mlx5gda_qp* qp) { if (qp->send_cq) { mlx5gda_destroy_cq(ctrl_buf_heap, qp->send_cq); } - if (qp->wq_offset != -1) { - memheap_free(ctrl_buf_heap, qp->wq_offset); - } - if (qp->dbr_offset != -1) { - memheap_free(ctrl_buf_heap, qp->dbr_offset); + if (uses_control_regions(&qp->region_allocator)) { + release_control_region(&qp->region_allocator, &qp->dbr_region); + release_control_region(&qp->region_allocator, &qp->wq_region); + } else { + if (qp->wq_offset != -1) { + memheap_free(ctrl_buf_heap, qp->wq_offset); + } + if (qp->dbr_offset != -1) { + memheap_free(ctrl_buf_heap, qp->dbr_offset); + } } if (qp) { free(qp); @@ -517,4 +679,4 @@ int mlx5gda_modify_rc_qp_rtr2rts(struct mlx5gda_qp* qp) { perror("Failed to modify RC QP (rtr2rts)"); } return ret; -} \ No newline at end of file +} diff --git a/mooncake-transfer-engine/src/transport/device/nccl_device_transport.cpp b/mooncake-transfer-engine/src/transport/device/nccl_device_transport.cpp new file mode 100644 index 0000000000..62638780d6 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/device/nccl_device_transport.cpp @@ -0,0 +1,545 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "transport/device/nccl_device_transport.h" + +#include +#include +#include +#include + +#include +#include +#include +#include + +#if NCCL_VERSION_CODE < 23004 +#error "Mooncake NCCL DeviceTransport requires NCCL 2.30.4 or newer" +#endif + +namespace mooncake { +namespace device { +namespace { + +int reportNcclError(ncclResult_t result, const char* operation) { + if (result == ncclSuccess) return 0; + LOG(ERROR) << "[Device NCCL] " << operation + << " failed: " << ncclGetErrorString(result); + return -1; +} + +int reportCudaError(cudaError_t result, const char* operation) { + if (result == cudaSuccess) return 0; + LOG(ERROR) << "[Device NCCL] " << operation + << " failed: " << cudaGetErrorString(result); + return -1; +} + +NcclGinBackend toGinBackend(ncclGinType_t type) { + switch (type) { + case NCCL_GIN_TYPE_PROXY: + return NcclGinBackend::kProxy; + case NCCL_GIN_TYPE_GDAKI: + return NcclGinBackend::kGdaki; +#if NCCL_VERSION_CODE >= NCCL_VERSION(2, 30, 6) + case NCCL_GIN_TYPE_GPI: + return NcclGinBackend::kGpi; +#endif + case NCCL_GIN_TYPE_NONE: + default: + return NcclGinBackend::kNone; + } +} + +const char* ginBackendName(NcclGinBackend backend) { + switch (backend) { + case NcclGinBackend::kProxy: + return "proxy"; + case NcclGinBackend::kGdaki: + return "gdaki"; + case NcclGinBackend::kGpi: + return "gpi"; + case NcclGinBackend::kNone: + default: + return "none"; + } +} + +bool decodeUniqueId(const std::vector& encoded, ncclUniqueId* id) { + const size_t words = + (sizeof(ncclUniqueId) + sizeof(int32_t) - 1) / sizeof(int32_t); + if (!id || encoded.size() != words) return false; + std::memcpy(id, encoded.data(), sizeof(*id)); + return true; +} + +} // namespace + +class NcclDeviceTransportImpl final : public NcclTransport { + public: + ~NcclDeviceTransportImpl() override { shutdown(); } + + std::vector createUniqueId() override { + ncclUniqueId id{}; + if (reportNcclError(ncclGetUniqueId(&id), "ncclGetUniqueId") != 0) { + return {}; + } + + const size_t words = + (sizeof(id) + sizeof(int32_t) - 1) / sizeof(int32_t); + std::vector encoded(words, 0); + std::memcpy(encoded.data(), &id, sizeof(id)); + return encoded; + } + + int initialize(const NcclTransportConfig& config, + const std::vector& unique_id) override { + if (initialized_) { + LOG(ERROR) << "[Device NCCL] transport is already initialized"; + return -1; + } + if (config.num_ranks <= 0 || config.rank < 0 || + config.rank >= config.num_ranks || config.gin_context_count < 0 || + (config.enable_gin && config.gin_context_count == 0) || + config.gin_context_count > std::numeric_limits::max() / 2 || + config.lsa_barrier_count < 0) { + LOG(ERROR) << "[Device NCCL] invalid communicator configuration"; + return -1; + } + + ncclUniqueId id{}; + if (!decodeUniqueId(unique_id, &id)) { + LOG(ERROR) << "[Device NCCL] invalid NCCL unique ID size"; + return -1; + } + + int runtime_version = 0; + if (reportNcclError(ncclGetVersion(&runtime_version), + "ncclGetVersion") != 0) { + return -1; + } + if (runtime_version != NCCL_VERSION_CODE) { + LOG(ERROR) + << "[Device NCCL] Device API requires matching compile-time " + "and runtime NCCL versions: compiled=" + << NCCL_VERSION_CODE << " runtime=" << runtime_version + << ". Rebuild Mooncake against the runtime NCCL installation, " + "and rebuild AOT NCCL device kernels or invalidate and " + "re-JIT cached NCCL Device API kernels"; + return -1; + } + + int cuda_device = -1; + if (reportCudaError(cudaGetDevice(&cuda_device), "cudaGetDevice") != + 0) { + return -1; + } + + if (reportNcclError( + ncclCommInitRank(&comm_, config.num_ranks, id, config.rank), + "ncclCommInitRank") != 0) { + comm_ = nullptr; + return -1; + } + + ncclCommProperties_t comm_properties = NCCL_COMM_PROPERTIES_INITIALIZER; + if (reportNcclError(ncclCommQueryProperties(comm_, &comm_properties), + "ncclCommQueryProperties") != 0) { + abortPartialInitialization(); + return -1; + } + if (!comm_properties.deviceApiSupport) { + LOG(ERROR) << "[Device NCCL] communicator does not support the " + "device API"; + abortPartialInitialization(); + return -1; + } + if (config.require_lsa_multimem && !comm_properties.multimemSupport) { + LOG(ERROR) << "[Device NCCL] LSA multimem was required but is " + "not supported"; + abortPartialInitialization(); + return -1; + } + if (config.enable_gin && + comm_properties.ginType == NCCL_GIN_TYPE_NONE) { + LOG(ERROR) << "[Device NCCL] full GIN connectivity was requested " + "but GIN is unavailable"; + abortPartialInitialization(); + return -1; + } + + ncclDevCommRequirements_t requirements = + NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER; + requirements.lsaMultimem = config.require_lsa_multimem; + requirements.lsaBarrierCount = config.lsa_barrier_count; + requirements.ginContextCount = + config.enable_gin ? config.gin_context_count : 0; + requirements.ginConnectionType = config.enable_gin + ? NCCL_GIN_CONNECTION_FULL + : NCCL_GIN_CONNECTION_NONE; + requirements.ginExclusiveContexts = + config.enable_gin && config.gin_exclusive_contexts; + + ncclResult_t start_result = ncclGroupStart(); + ncclResult_t create_result = ncclSuccess; + ncclResult_t end_result = ncclSuccess; + if (start_result == ncclSuccess) { + create_result = ncclDevCommCreate(comm_, &requirements, &dev_comm_); + end_result = ncclGroupEnd(); + } + if (reportNcclError(start_result, "ncclGroupStart") != 0 || + reportNcclError(create_result, "ncclDevCommCreate") != 0 || + reportNcclError(end_result, "ncclGroupEnd") != 0) { + abortPartialInitialization(); + return -1; + } + dev_comm_created_ = true; + + if (config.enable_gin && (dev_comm_.ginConnectionCount == 0 || + dev_comm_.ginContextCount == 0)) { + LOG(ERROR) << "[Device NCCL] communicator did not provide the " + "requested full GIN resources"; + abortPartialInitialization(); + return -1; + } + + if (reportCudaError(cudaMalloc(reinterpret_cast(&device_comm_), + sizeof(dev_comm_)), + "cudaMalloc(device communicator)") != 0 || + reportCudaError( + cudaMemcpy(device_comm_, &dev_comm_, sizeof(dev_comm_), + cudaMemcpyHostToDevice), + "cudaMemcpy(device communicator)") != 0 || + reportCudaError(cudaStreamCreateWithFlags(&control_stream_, + cudaStreamNonBlocking), + "cudaStreamCreateWithFlags(control)") != 0 || + reportCudaError( + cudaMalloc(reinterpret_cast(&collective_status_), + sizeof(int)), + "cudaMalloc(collective status)") != 0) { + abortPartialInitialization(); + return -1; + } + + properties_.runtime_version = runtime_version; + properties_.rank = comm_properties.rank; + properties_.num_ranks = comm_properties.nRanks; + properties_.cuda_device = comm_properties.cudaDev; + properties_.device_api_supported = comm_properties.deviceApiSupport; + properties_.multimem_supported = comm_properties.multimemSupport; + properties_.lsa_multimem_enabled = config.require_lsa_multimem; + properties_.lsa_team_count = comm_properties.nLsaTeams; + properties_.lsa_barrier_count = config.lsa_barrier_count; + properties_.gin_enabled = config.enable_gin; + properties_.gin_backend = toGinBackend(comm_properties.ginType); + properties_.gin_connection_count = dev_comm_.ginConnectionCount; + properties_.gin_context_count = + static_cast(dev_comm_.ginContextCount); + initialized_ = true; + + LOG(INFO) << "[Device NCCL] initialized rank=" << properties_.rank + << "/" << properties_.num_ranks + << " cuda_device=" << properties_.cuda_device + << " lsa_teams=" << properties_.lsa_team_count + << " multimem=" << properties_.lsa_multimem_enabled + << " lsa_barriers=" << properties_.lsa_barrier_count + << " gin_backend=" << ginBackendName(properties_.gin_backend) + << " gin_connections=" << properties_.gin_connection_count + << " gin_contexts=" << properties_.gin_context_count; + return 0; + } + + void* allocateBuffer(size_t bytes) override { + if (bytes == 0) { + LOG(ERROR) << "[Device NCCL] cannot allocate a zero-sized buffer"; + return nullptr; + } + + void* ptr = nullptr; + if (reportNcclError(ncclMemAlloc(&ptr, bytes), "ncclMemAlloc") != 0) { + return nullptr; + } + + allocations_.insert(ptr); + return ptr; + } + + int freeBuffer(void* ptr) override { + if (!ptr) return 0; + + for (const auto& entry : registrations_) { + if (entry.second.ptr == ptr) { + LOG(ERROR) << "[Device NCCL] deregister the buffer before " + "freeing its allocation"; + return -1; + } + } + + auto it = allocations_.find(ptr); + if (it == allocations_.end()) { + LOG(ERROR) << "[Device NCCL] buffer was not allocated by this " + "transport"; + return -1; + } + if (reportNcclError(ncclMemFree(ptr), "ncclMemFree") != 0) { + return -1; + } + allocations_.erase(it); + return 0; + } + + int registerBuffer(void* ptr, size_t bytes, + NcclBufferRegistration* registration) override { + if (!ptr || bytes == 0 || !registration || registration->valid()) { + LOG(ERROR) << "[Device NCCL] invalid buffer registration"; + return -1; + } + if (!initialized_) { + LOG(ERROR) << "[Device NCCL] initialize before registering a " + "buffer"; + return -1; + } + + ncclWindow_t window = nullptr; + if (reportNcclError(ncclCommWindowRegister(comm_, ptr, bytes, &window, + NCCL_WIN_COLL_SYMMETRIC), + "ncclCommWindowRegister") != 0) { + return -1; + } + + if (next_registration_id_ == 0) { + LOG(ERROR) << "[Device NCCL] registration ID space exhausted"; + ncclCommWindowDeregister(comm_, window); + return -1; + } + const uint64_t id = next_registration_id_++; + registrations_.emplace(id, WindowRecord{window, ptr, bytes}); + registration->id_ = id; + return 0; + } + + int deregisterBuffer(NcclBufferRegistration* registration) override { + if (!registration || !registration->valid()) return 0; + + auto it = registrations_.find(registration->id_); + if (it == registrations_.end()) { + LOG(ERROR) << "[Device NCCL] unknown buffer registration"; + return -1; + } + if (reportNcclError(ncclCommWindowDeregister(comm_, it->second.window), + "ncclCommWindowDeregister") != 0) { + return -1; + } + registrations_.erase(it); + registration->id_ = 0; + return 0; + } + + int allocateAndRegisterBuffer( + size_t bytes, void** ptr, + NcclBufferRegistration* registration) override { + if (!initialized_ || bytes == 0 || !ptr || !registration || + registration->valid()) { + LOG(ERROR) << "[Device NCCL] invalid collective buffer request"; + return -1; + } + + *ptr = nullptr; + void* local_ptr = allocateBuffer(bytes); + const bool all_allocated = collectiveAllSucceeded(local_ptr != nullptr); + if (!all_allocated) { + if (local_ptr) freeBuffer(local_ptr); + return -1; + } + + NcclBufferRegistration local_registration; + const int register_status = + registerBuffer(local_ptr, bytes, &local_registration); + const bool all_registered = + collectiveAllSucceeded(register_status == 0); + if (!all_registered) { + if (local_registration.valid()) + deregisterBuffer(&local_registration); + freeBuffer(local_ptr); + return -1; + } + + *ptr = local_ptr; + *registration = local_registration; + return 0; + } + + NcclDeviceContext deviceContext( + const NcclBufferRegistration& registration) const override { + NcclDeviceContext context; + if (!initialized_ || !registration.valid()) return context; + + const auto it = registrations_.find(registration.id_); + if (it == registrations_.end()) { + LOG(ERROR) << "[Device NCCL] unknown buffer registration"; + return context; + } + + context.native_comm_ = device_comm_; + context.native_window_ = it->second.window; + context.local_base_ = it->second.ptr; + context.rank_ = properties_.rank; + context.gin_context_count_ = properties_.gin_context_count; + context.gin_enabled_ = properties_.gin_enabled; + context.lsa_multimem_enabled_ = properties_.lsa_multimem_enabled; + return context; + } + + NcclTransportProperties properties() const override { return properties_; } + + bool initialized() const override { return initialized_; } + + int shutdown() override { + int status = 0; + + if (comm_) { + for (const auto& entry : registrations_) { + if (reportNcclError( + ncclCommWindowDeregister(comm_, entry.second.window), + "ncclCommWindowDeregister") != 0) { + status = -1; + } + } + } + registrations_.clear(); + + for (void* ptr : allocations_) { + if (reportNcclError(ncclMemFree(ptr), "ncclMemFree") != 0) { + status = -1; + } + } + allocations_.clear(); + + if (collective_status_) { + if (reportCudaError(cudaFree(collective_status_), + "cudaFree(collective status)") != 0) { + status = -1; + } + collective_status_ = nullptr; + } + if (control_stream_) { + if (reportCudaError(cudaStreamDestroy(control_stream_), + "cudaStreamDestroy(control)") != 0) { + status = -1; + } + control_stream_ = nullptr; + } + if (device_comm_) { + if (reportCudaError(cudaFree(device_comm_), + "cudaFree(device communicator)") != 0) { + status = -1; + } + device_comm_ = nullptr; + } + + if (dev_comm_created_) { + if (reportNcclError(ncclDevCommDestroy(comm_, &dev_comm_), + "ncclDevCommDestroy") != 0) { + status = -1; + } + dev_comm_created_ = false; + dev_comm_ = {}; + } + if (comm_) { + if (reportNcclError(ncclCommDestroy(comm_), "ncclCommDestroy") != + 0) { + status = -1; + } + comm_ = nullptr; + } + + initialized_ = false; + next_registration_id_ = 1; + properties_ = {}; + return status; + } + + private: + struct WindowRecord { + ncclWindow_t window; + void* ptr; + size_t bytes; + }; + + bool collectiveAllSucceeded(bool local_success) { + int value = local_success ? 1 : 0; + if (reportCudaError( + cudaMemcpyAsync(collective_status_, &value, sizeof(value), + cudaMemcpyHostToDevice, control_stream_), + "cudaMemcpyAsync(collective status H2D)") != 0 || + reportNcclError( + ncclAllReduce(collective_status_, collective_status_, 1, + ncclInt, ncclMin, comm_, control_stream_), + "ncclAllReduce(collective status)") != 0 || + reportCudaError( + cudaMemcpyAsync(&value, collective_status_, sizeof(value), + cudaMemcpyDeviceToHost, control_stream_), + "cudaMemcpyAsync(collective status D2H)") != 0 || + reportCudaError(cudaStreamSynchronize(control_stream_), + "cudaStreamSynchronize(collective status)") != 0) { + return false; + } + return value != 0; + } + + void abortPartialInitialization() { + if (collective_status_) { + cudaFree(collective_status_); + collective_status_ = nullptr; + } + if (control_stream_) { + cudaStreamDestroy(control_stream_); + control_stream_ = nullptr; + } + if (device_comm_) { + cudaFree(device_comm_); + device_comm_ = nullptr; + } + if (dev_comm_created_) { + ncclDevCommDestroy(comm_, &dev_comm_); + dev_comm_created_ = false; + dev_comm_ = {}; + } + if (comm_) { + ncclCommAbort(comm_); + comm_ = nullptr; + } + } + + ncclComm_t comm_ = nullptr; + ncclDevComm_t dev_comm_{}; + ncclDevComm_t* device_comm_ = nullptr; + bool dev_comm_created_ = false; + bool initialized_ = false; + + cudaStream_t control_stream_ = nullptr; + int* collective_status_ = nullptr; + + uint64_t next_registration_id_ = 1; + std::unordered_map registrations_; + std::unordered_set allocations_; + NcclTransportProperties properties_; +}; + +std::unique_ptr createNcclDeviceTransport() { + return std::make_unique(); +} + +} // namespace device +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/device/p2p_device_transport.cpp b/mooncake-transfer-engine/src/transport/device/p2p_device_transport.cpp index 23e0225531..31661e1371 100644 --- a/mooncake-transfer-engine/src/transport/device/p2p_device_transport.cpp +++ b/mooncake-transfer-engine/src/transport/device/p2p_device_transport.cpp @@ -21,29 +21,245 @@ #include #include +#include +#include +#include #include +#include +#include +#include #include "cuda_alike.h" namespace mooncake { namespace device { +namespace { + +#if defined(USE_CUDA) +bool supportFabricMem() { + const char* nvlink_ipc = std::getenv("MC_USE_NVLINK_IPC"); + if (nvlink_ipc == nullptr || std::strcmp(nvlink_ipc, "0") != 0) + return false; + + int num_devices = 0; + cudaError_t err = cudaGetDeviceCount(&num_devices); + if (err != cudaSuccess || num_devices == 0) return false; + + for (int device_id = 0; device_id < num_devices; ++device_id) { + CUdevice cu_device; + if (cuDeviceGet(&cu_device, device_id) != CUDA_SUCCESS) return false; + int supported = 0; + if (cuDeviceGetAttribute( + &supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, + cu_device) != CUDA_SUCCESS || + !supported) + return false; + } + return true; +} + +std::vector serializeFabricHandle(const CUmemFabricHandle& handle, + size_t bytes) { + constexpr size_t kPayloadBytes = + sizeof(CUmemFabricHandle) + sizeof(uint64_t); + constexpr size_t kNumInt32s = + (kPayloadBytes + sizeof(int32_t) - 1) / sizeof(int32_t); + std::vector result(kNumInt32s, 0); + auto* payload = reinterpret_cast(result.data()); + uint64_t encoded_size = static_cast(bytes); + memcpy(payload, &handle, sizeof(CUmemFabricHandle)); + memcpy(payload + sizeof(CUmemFabricHandle), &encoded_size, + sizeof(encoded_size)); + return result; +} + +bool deserializeFabricHandle(const std::vector& encoded, + CUmemFabricHandle* handle, size_t* bytes) { + constexpr size_t kPayloadBytes = + sizeof(CUmemFabricHandle) + sizeof(uint64_t); + constexpr size_t kNumInt32s = + (kPayloadBytes + sizeof(int32_t) - 1) / sizeof(int32_t); + if (encoded.size() < kNumInt32s) return false; + const auto* payload = reinterpret_cast(encoded.data()); + uint64_t encoded_size = 0; + memcpy(handle, payload, sizeof(CUmemFabricHandle)); + memcpy(&encoded_size, payload + sizeof(CUmemFabricHandle), + sizeof(encoded_size)); + *bytes = static_cast(encoded_size); + return *bytes != 0; +} + +struct FabricAllocation { + CUdeviceptr ptr = 0; + size_t size = 0; + CUmemGenericAllocationHandle handle{}; +}; +#else +bool supportFabricMem() { return false; } +#endif + +} // namespace + +#ifdef USE_MACA +namespace { + +bool parseBoolEnv(const char* name) { + const char* value = std::getenv(name); + if (value == nullptr) return false; + std::string s(value); + std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return s == "1" || s == "on" || s == "true" || s == "yes"; +} + +std::string getLowerEnv(const char* name) { + const char* value = std::getenv(name); + if (value == nullptr) return ""; + std::string s(value); + std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return s; +} + +int macaAllocFlagFromMode(const std::string& mode, const char* env_name) { + if (mode.empty() || mode == "default" || mode == "cuda") + return mcDeviceMallocDefault; + if (mode == "fine" || mode == "finegrained" || mode == "fine-grained") + return mcDeviceMallocFinegrained; + if (mode == "signal") return mcMallocSignalMemory; + if (mode == "wc" || mode == "writecoherence" || mode == "write-coherence") + return mcDeviceMallocWriteCoherence; + if (mode == "pcie" || mode == "pcie-uncache" || mode == "map-pcie") + return mcDeviceMallocMapPcieDefault; + if (mode == "pcie-wc" || mode == "map-pcie-wc") + return mcDeviceMallocMapPcieCoherence; + if (mode == "fixed" || mode == "fixed-uncache") + return mcDeviceMallocFixedMemDefault; + if (mode == "fixed-wc") return mcDeviceMallocFixedMemCoherence; + LOG(WARNING) << "[EP P2P] unknown " << env_name << "=" << mode + << ", using default cudaMalloc"; + return mcDeviceMallocDefault; +} + +int macaAllocFlagFromEnv() { + return macaAllocFlagFromMode(getLowerEnv("MOONCAKE_EP_MACA_ALLOC"), + "MOONCAKE_EP_MACA_ALLOC"); +} + +std::string macaIpcMode() { + std::string mode = getLowerEnv("MOONCAKE_EP_MACA_IPC"); + return mode.empty() ? "normal" : mode; +} + +bool parseNonNegativeInt(const std::string& token, int* value) { + if (token.empty()) return false; + int result = 0; + for (char c : token) { + if (!std::isdigit(static_cast(c))) return false; + result = result * 10 + (c - '0'); + } + *value = result; + return true; +} + +int physicalDeviceFromVisibleList(int logical_device) { + const char* visible = std::getenv("CUDA_VISIBLE_DEVICES"); + if (visible == nullptr || visible[0] == '\0') return logical_device; + + std::string list(visible); + size_t begin = 0; + int logical = 0; + while (begin <= list.size()) { + size_t end = list.find(',', begin); + if (end == std::string::npos) end = list.size(); + std::string token = list.substr(begin, end - begin); + token.erase(std::remove_if( + token.begin(), token.end(), + [](unsigned char c) { return std::isspace(c) != 0; }), + token.end()); + if (logical == logical_device) { + int physical = logical_device; + return parseNonNegativeInt(token, &physical) ? physical + : logical_device; + } + if (end == list.size()) break; + begin = end + 1; + ++logical; + } + return logical_device; +} + +bool pairListed(const std::string& pairs, int src, int dst) { + size_t begin = 0; + while (begin <= pairs.size()) { + size_t end = pairs.find(',', begin); + if (end == std::string::npos) end = pairs.size(); + std::string item = pairs.substr(begin, end - begin); + item.erase(std::remove_if( + item.begin(), item.end(), + [](unsigned char c) { return std::isspace(c) != 0; }), + item.end()); + + size_t dash = item.find('-'); + if (dash != std::string::npos) { + int a = -1, b = -1; + if (parseNonNegativeInt(item.substr(0, dash), &a) && + parseNonNegativeInt(item.substr(dash + 1), &b)) { + if ((a == src && b == dst) || (a == dst && b == src)) + return true; + } + } + + if (end == pairs.size()) break; + begin = end + 1; + } + return false; +} + +bool macaP2pPairAllowed(int src_physical, int dst_physical) { + if (parseBoolEnv("MOONCAKE_EP_MACA_ALLOW_NODE_P2P")) return true; + + const char* explicit_pairs = std::getenv("MOONCAKE_EP_MACA_P2P_PAIRS"); + if (explicit_pairs != nullptr && explicit_pairs[0] != '\0') + return pairListed(explicit_pairs, src_physical, dst_physical); + + // C500 exposes two direct MetaXLink islands by default: 0<->1 and 2<->3. + // NODE pairs may report canAccessPeer=1, but EP kernel peer stores can hang + // waiting for device-side signals on those paths. + return src_physical / 2 == dst_physical / 2 && + std::abs(src_physical - dst_physical) == 1; +} + +} // namespace +#endif + class P2pDeviceTransportImpl : public P2pTransport { public: - explicit P2pDeviceTransportImpl(int num_ranks) : num_ranks_(num_ranks) { + explicit P2pDeviceTransportImpl(int num_ranks) + : num_ranks_(num_ranks), use_fabric_mem_(supportFabricMem()) { cudaMalloc(&available_table_, num_ranks_ * sizeof(int32_t)); cudaMemset(available_table_, 0, num_ranks_ * sizeof(int32_t)); cudaMallocHost(&peer_ptrs_host_, num_ranks_ * sizeof(void*)); cudaMalloc(&peer_ptrs_dev_, num_ranks_ * sizeof(void*)); for (int i = 0; i < num_ranks_; ++i) peer_ptrs_host_[i] = nullptr; cudaMemset(peer_ptrs_dev_, 0, num_ranks_ * sizeof(void*)); +#if defined(USE_CUDA) + fabric_peer_mappings_.resize(num_ranks_); +#endif } ~P2pDeviceTransportImpl() override { +#if defined(USE_CUDA) + cleanupFabricPeerMappings(); + cleanupFabricAllocations(); +#endif if (available_table_) cudaFree(available_table_); if (peer_ptrs_dev_) cudaFree(peer_ptrs_dev_); if (peer_ptrs_host_) { - for (int i = 0; i < num_ranks_; ++i) { + for (int i = 0; i < num_ranks_ && !use_fabric_mem_; ++i) { if (peer_ptrs_host_[i] && peer_ptrs_host_[i] != local_ptr_) { cudaIpcCloseMemHandle(peer_ptrs_host_[i]); } @@ -54,22 +270,213 @@ class P2pDeviceTransportImpl : public P2pTransport { void* allocateBuffer(size_t bytes) override { void* ptr = nullptr; +#if defined(USE_CUDA) + if (use_fabric_mem_) { + int device_id = 0; + if (cudaGetDevice(&device_id) != cudaSuccess) { + LOG(ERROR) + << "[EP P2P] cudaGetDevice failed before fabric alloc"; + return nullptr; + } + + CUdevice cu_dev; + CUresult res = cuDeviceGet(&cu_dev, device_id); + if (res != CUDA_SUCCESS) { + LOG(ERROR) << "[EP P2P] cuDeviceGet failed: " << res; + return nullptr; + } + + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = cu_dev; + prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC; + + int rdma_capable = 0; + cuDeviceGetAttribute( + &rdma_capable, + CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, + cu_dev); + if (rdma_capable) prop.allocFlags.gpuDirectRDMACapable = 1; + + size_t granularity = 0; + res = cuMemGetAllocationGranularity( + &granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM); + if (res != CUDA_SUCCESS) { + LOG(ERROR) << "[EP P2P] cuMemGetAllocationGranularity failed: " + << res; + return nullptr; + } + + size_t fabric_alloc_size = std::max( + granularity, (bytes + granularity - 1) & ~(granularity - 1)); + CUmemGenericAllocationHandle fabric_mem_handle{}; + res = cuMemCreate(&fabric_mem_handle, fabric_alloc_size, &prop, 0); + if (res != CUDA_SUCCESS) { + LOG(ERROR) << "[EP P2P] cuMemCreate(FABRIC) failed: " << res; + return nullptr; + } + + CUdeviceptr reserved = 0; + res = cuMemAddressReserve(&reserved, fabric_alloc_size, granularity, + 0, 0); + if (res != CUDA_SUCCESS) { + cuMemRelease(fabric_mem_handle); + LOG(ERROR) << "[EP P2P] cuMemAddressReserve failed: " << res; + return nullptr; + } + + res = + cuMemMap(reserved, fabric_alloc_size, 0, fabric_mem_handle, 0); + if (res != CUDA_SUCCESS) { + cuMemAddressFree(reserved, fabric_alloc_size); + cuMemRelease(fabric_mem_handle); + LOG(ERROR) << "[EP P2P] cuMemMap failed: " << res; + return nullptr; + } + + int device_count = 0; + cudaGetDeviceCount(&device_count); + std::vector access(device_count); + for (int i = 0; i < device_count; ++i) { + access[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access[i].location.id = i; + access[i].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + } + res = cuMemSetAccess(reserved, fabric_alloc_size, access.data(), + device_count); + if (res != CUDA_SUCCESS) { + cuMemUnmap(reserved, fabric_alloc_size); + cuMemAddressFree(reserved, fabric_alloc_size); + cuMemRelease(fabric_mem_handle); + LOG(ERROR) << "[EP P2P] cuMemSetAccess failed: " << res; + return nullptr; + } + + ptr = reinterpret_cast(reserved); + fabric_allocations_.emplace( + ptr, FabricAllocation{reserved, fabric_alloc_size, + fabric_mem_handle}); + LOG(INFO) << "[EP P2P] allocated fabric buffer bytes=" << bytes + << " mapped_bytes=" << fabric_alloc_size; + return ptr; + } +#endif +#ifdef USE_MACA + int alloc_flag = macaAllocFlagFromEnv(); + cudaError_t err = alloc_flag == mcDeviceMallocDefault + ? cudaMalloc(&ptr, bytes) + : mcExtMallocWithFlags(&ptr, bytes, alloc_flag); +#else cudaError_t err = cudaMalloc(&ptr, bytes); +#endif if (err != cudaSuccess) { - LOG(ERROR) << "[EP P2P] cudaMalloc(" << bytes + LOG(ERROR) << "[EP P2P] device allocation(" << bytes << ") failed: " << cudaGetErrorString(err); return nullptr; } +#ifdef USE_MACA + if (alloc_flag != mcDeviceMallocDefault) { + LOG(INFO) << "[EP P2P] allocated MACA buffer with " + "mcExtMallocWithFlags flag=" + << alloc_flag; + } +#endif return ptr; } - void freeBuffer(void* ptr) override { cudaFree(ptr); } + void freeBuffer(void* ptr) override { +#if defined(USE_CUDA) + if (ptr == nullptr) return; + if (use_fabric_mem_) { + auto it = fabric_allocations_.find(ptr); + if (it == fabric_allocations_.end()) { + LOG(ERROR) << "[EP P2P] unknown fabric buffer=" << ptr; + return; + } + if (ptr == local_ptr_) { + cleanupFabricPeerMappings(); + local_ptr_ = nullptr; + } + FabricAllocation allocation = it->second; + fabric_allocations_.erase(it); + cuMemUnmap(allocation.ptr, allocation.size); + cuMemAddressFree(allocation.ptr, allocation.size); + cuMemRelease(allocation.handle); + return; + } +#endif + cudaFree(ptr); + } std::vector exportIpcHandle(void* ptr) override { +#if defined(USE_CUDA) + if (use_fabric_mem_) { + auto it = fabric_allocations_.find(ptr); + if (it == fabric_allocations_.end()) { + LOG(ERROR) << "[EP P2P] unknown fabric buffer=" << ptr; + return {}; + } + CUmemFabricHandle export_handle; + CUresult res = + cuMemExportToShareableHandle(&export_handle, it->second.handle, + CU_MEM_HANDLE_TYPE_FABRIC, 0); + if (res != CUDA_SUCCESS) { + LOG(ERROR) + << "[EP P2P] cuMemExportToShareableHandle(FABRIC) failed: " + << res; + return {}; + } + return serializeFabricHandle(export_handle, it->second.size); + } +#endif +#ifdef USE_MACA + if (parseBoolEnv("MOONCAKE_EP_MACA_DISABLE_IPC")) { + LOG(INFO) << "[EP P2P] MACA IPC handle export disabled by " + "MOONCAKE_EP_MACA_DISABLE_IPC"; + return {}; + } + + cudaPointerAttributes attr{}; + cudaError_t attr_err = cudaPointerGetAttributes(&attr, ptr); + if (attr_err != cudaSuccess || attr.type != cudaMemoryTypeDevice || + attr.devicePointer == nullptr) { + LOG(WARNING) << "[EP P2P] skip MACA IPC handle export for " + << "non-device pointer=" << ptr + << ", attr_err=" << cudaGetErrorString(attr_err) + << ", type=" << attr.type + << ", devicePointer=" << attr.devicePointer + << ", allocationFlags=" << attr.allocationFlags; + return {}; + } + + std::string ipc_mode = macaIpcMode(); + if (ipc_mode == "cross-v2" || ipc_mode == "cross_v2") { + mcIpcCrossMemHandle_t handle; + cudaError_t err = mcIpcGetMemHandleCross_v2(&handle, ptr); + if (err != cudaSuccess) { + LOG(ERROR) << "[EP P2P] mcIpcGetMemHandleCross_v2 failed: " + << cudaGetErrorString(err); + return {}; + } + constexpr size_t kHandleBytes = sizeof(mcIpcCrossMemHandle_t); + constexpr size_t kNumInt32s = + (kHandleBytes + sizeof(int32_t) - 1) / sizeof(int32_t); + std::vector result(kNumInt32s); + memcpy(result.data(), &handle, kHandleBytes); + return result; + } +#endif cudaIpcMemHandle_t handle; +#ifdef USE_MACA + cudaError_t err = macaIpcMode() == "cross" + ? mcIpcGetMemHandleCross(&handle, ptr) + : cudaIpcGetMemHandle(&handle, ptr); +#else cudaError_t err = cudaIpcGetMemHandle(&handle, ptr); +#endif if (err != cudaSuccess) { - LOG(ERROR) << "[EP P2P] cudaIpcGetMemHandle failed: " + LOG(ERROR) << "[EP P2P] IPC handle export failed: " << cudaGetErrorString(err); return {}; } @@ -96,6 +503,91 @@ class P2pDeviceTransportImpl : public P2pTransport { available[rank] = 1; peer_ptrs_host_[rank] = local_ptr; +#if defined(USE_CUDA) + if (use_fabric_mem_) { + cleanupFabricPeerMappings(); + all_peers_accessible_ = true; + std::vector access(device_count); + for (int i = 0; i < device_count; ++i) { + access[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access[i].location.id = i; + access[i].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + } + for (int dst = 0; dst < num_ranks_; ++dst) { + if (active_ranks_mask[dst] == 0) continue; + if (dst == rank) continue; + if (dst >= static_cast(remote_handles.size())) { + all_peers_accessible_ = false; + break; + } + CUmemFabricHandle export_handle; + size_t mapped_size = 0; + if (!deserializeFabricHandle(remote_handles[dst], + &export_handle, &mapped_size)) { + all_peers_accessible_ = false; + break; + } + + CUmemGenericAllocationHandle handle; + CUresult res = cuMemImportFromShareableHandle( + &handle, &export_handle, CU_MEM_HANDLE_TYPE_FABRIC); + if (res != CUDA_SUCCESS) { + LOG(ERROR) + << "[EP P2P] cuMemImportFromShareableHandle(FABRIC) " + "failed for rank " + << dst << ": " << res; + all_peers_accessible_ = false; + break; + } + + CUdeviceptr peer_reserved = 0; + res = cuMemAddressReserve(&peer_reserved, mapped_size, 0, 0, 0); + if (res != CUDA_SUCCESS) { + cuMemRelease(handle); + LOG(ERROR) + << "[EP P2P] cuMemAddressReserve peer fabric mapping " + "failed for rank " + << dst << ": " << res; + all_peers_accessible_ = false; + break; + } + res = cuMemMap(peer_reserved, mapped_size, 0, handle, 0); + if (res != CUDA_SUCCESS) { + cuMemAddressFree(peer_reserved, mapped_size); + cuMemRelease(handle); + LOG(ERROR) << "[EP P2P] cuMemMap peer fabric mapping " + "failed for rank " + << dst << ": " << res; + all_peers_accessible_ = false; + break; + } + + res = cuMemSetAccess(peer_reserved, mapped_size, access.data(), + device_count); + if (res != CUDA_SUCCESS) { + cuMemUnmap(peer_reserved, mapped_size); + cuMemAddressFree(peer_reserved, mapped_size); + cuMemRelease(handle); + LOG(ERROR) << "[EP P2P] cuMemSetAccess peer fabric mapping " + "failed for rank " + << dst << ": " << res; + all_peers_accessible_ = false; + break; + } + + void* peer_ptr = reinterpret_cast(peer_reserved); + fabric_peer_mappings_[dst] = {peer_ptr, mapped_size, handle}; + available[dst] = 1; + peer_ptrs_host_[dst] = peer_ptr; + } + cudaMemcpy(available_table_, available.data(), + num_ranks_ * sizeof(int32_t), cudaMemcpyHostToDevice); + cudaMemcpy(peer_ptrs_dev_, peer_ptrs_host_, + num_ranks_ * sizeof(void*), cudaMemcpyHostToDevice); + return; + } +#endif + int node_id = rank / device_count; int group_start = node_id * device_count; int group_end = std::min(group_start + device_count, num_ranks_); @@ -112,6 +604,20 @@ class P2pDeviceTransportImpl : public P2pTransport { << "): canAccessPeer=" << can_access; if (!can_access) continue; +#ifdef USE_MACA + int src_physical = physicalDeviceFromVisibleList(device_id); + int dst_physical = physicalDeviceFromVisibleList(dst_device); + if (!macaP2pPairAllowed(src_physical, dst_physical)) { + LOG(INFO) << "[EP P2P] rank " << rank << " physical GPU" + << src_physical << " -> rank " << dst + << " physical GPU" << dst_physical + << " disabled for MACA EP fast path; set " + "MOONCAKE_EP_MACA_ALLOW_NODE_P2P=1 or " + "MOONCAKE_EP_MACA_P2P_PAIRS to override"; + continue; + } +#endif + cudaError_t err = cudaDeviceEnablePeerAccess(dst_device, 0); if (err != cudaSuccess && err != cudaErrorPeerAccessAlreadyEnabled) { @@ -129,14 +635,50 @@ class P2pDeviceTransportImpl : public P2pTransport { constexpr size_t kHandleBytes = sizeof(cudaIpcMemHandle_t); constexpr size_t kNumInt32s = (kHandleBytes + sizeof(int32_t) - 1) / sizeof(int32_t); +#ifdef USE_MACA + std::string ipc_mode = macaIpcMode(); + if (ipc_mode == "cross-v2" || ipc_mode == "cross_v2") { + constexpr size_t kCrossHandleBytes = + sizeof(mcIpcCrossMemHandle_t); + constexpr size_t kCrossNumInt32s = + (kCrossHandleBytes + sizeof(int32_t) - 1) / sizeof(int32_t); + if (h.size() < kCrossNumInt32s) continue; + mcIpcCrossMemHandle_t handle; + memcpy(&handle, h.data(), kCrossHandleBytes); + void* peer_ptr = nullptr; + err = mcIpcOpenMemHandleCross_v2( + &peer_ptr, &handle, cudaIpcMemLazyEnablePeerAccess); + if (err != cudaSuccess) { + LOG(WARNING) + << "[EP P2P] rank " << rank + << " failed to open cross_v2 IPC handle for rank " + << dst << ": " << cudaGetErrorString(err); + continue; + } + LOG(INFO) << "[EP P2P] rank " << rank + << " opened cross_v2 IPC handle for rank " << dst + << ": peer_ptr=" << peer_ptr; + available[dst] = 1; + peer_ptrs_host_[dst] = peer_ptr; + continue; + } +#endif if (h.size() < kNumInt32s) continue; cudaIpcMemHandle_t handle; memcpy(&handle, h.data(), kHandleBytes); void* peer_ptr = nullptr; +#ifdef USE_MACA + err = ipc_mode == "cross" + ? mcIpcOpenMemHandleCross(&peer_ptr, handle, + cudaIpcMemLazyEnablePeerAccess) + : cudaIpcOpenMemHandle(&peer_ptr, handle, + cudaIpcMemLazyEnablePeerAccess); +#else err = cudaIpcOpenMemHandle(&peer_ptr, handle, cudaIpcMemLazyEnablePeerAccess); +#endif if (err != cudaSuccess) { LOG(WARNING) << "[EP P2P] rank " << rank << " failed to open IPC handle for rank " << dst @@ -252,12 +794,53 @@ class P2pDeviceTransportImpl : public P2pTransport { } private: +#if defined(USE_CUDA) + struct FabricPeerMapping { + void* ptr = nullptr; + size_t size = 0; + CUmemGenericAllocationHandle handle{}; + }; + + void cleanupFabricPeerMappings() { + if (!use_fabric_mem_) return; + for (int i = 0; i < static_cast(fabric_peer_mappings_.size()); + ++i) { + auto& mapping = fabric_peer_mappings_[i]; + if (mapping.ptr == nullptr) continue; + cuMemUnmap(reinterpret_cast(mapping.ptr), + mapping.size); + cuMemAddressFree(reinterpret_cast(mapping.ptr), + mapping.size); + cuMemRelease(mapping.handle); + mapping = {}; + if (peer_ptrs_host_ && peer_ptrs_host_[i] != local_ptr_) { + peer_ptrs_host_[i] = nullptr; + } + } + } + + void cleanupFabricAllocations() { + for (const auto& entry : fabric_allocations_) { + const auto& allocation = entry.second; + cuMemUnmap(allocation.ptr, allocation.size); + cuMemAddressFree(allocation.ptr, allocation.size); + cuMemRelease(allocation.handle); + } + fabric_allocations_.clear(); + } +#endif + int num_ranks_; + bool use_fabric_mem_ = false; void* local_ptr_ = nullptr; int32_t* available_table_ = nullptr; void** peer_ptrs_host_ = nullptr; void** peer_ptrs_dev_ = nullptr; bool all_peers_accessible_ = false; +#if defined(USE_CUDA) + std::vector fabric_peer_mappings_; + std::unordered_map fabric_allocations_; +#endif }; std::unique_ptr createP2pDeviceTransport(int num_ranks) { diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp index 6f8abe495c..c06a202240 100644 --- a/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp +++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_context.cpp @@ -178,7 +178,19 @@ int EfaContext::construct(size_t num_cq_list, size_t max_cqe, return ERR_CONTEXT; } - // Create completion queues + // Create completion queues. + // + // Same design as max_wr_depth_ in buildSharedEndpoint(): the submit path + // also paces against a CQ-occupancy counter (EfaCq::outstanding vs + // max_cqe_), so that ceiling must be the CQ the provider created, not the + // one we asked for. EFA raises any request below + // MAX(rx_attr->size + tx_attr->size, FI_EFA_CQ_SIZE) to that value + // (efa_domain.c), e.g. 12288 on p5 -- 3x the 4096 default, so the counter + // would stop submission long before the CQ is full. + // + // fi_cq_open() writes the size it chose back into cq_attr.size; read that + // instead of recomputing the formula, which folds in the operator-settable + // FI_EFA_CQ_SIZE and may change in a future provider. cq_list_.resize(num_cq_list); for (size_t i = 0; i < num_cq_list; ++i) { auto cq = std::make_shared(); @@ -193,8 +205,17 @@ int EfaContext::construct(size_t num_cq_list, size_t max_cqe, LOG(ERROR) << "fi_cq_open failed: " << fi_strerror(-ret); return ERR_CONTEXT; } + // Every CQ on this domain gets the same treatment, so one value covers + // the list; guard anyway so a future provider that sizes them + // independently cannot leave the counter above a smaller CQ. + if (i == 0 || cq_attr.size < max_cqe_) max_cqe_ = cq_attr.size; cq_list_[i] = cq; } + if (max_cqe_ > max_cqe) { + VLOG(1) << "EFA " << device_name_ << ": provider opened a CQ of " + << max_cqe_ << " entries for the requested " << max_cqe + << "; pacing against the larger real depth."; + } // Build the shared endpoint that services every peer through AV lookup. ret = buildSharedEndpoint(globalConfig().max_wr, 64); @@ -208,7 +229,9 @@ int EfaContext::construct(size_t num_cq_list, size_t max_cqe, << ", domain: " << fi_info_->domain_attr->name << ", fabric: " << fi_info_->fabric_attr->name << ", provider: " << fi_info_->fabric_attr->prov_name - << " (shared endpoint, max_wr=" << max_wr_depth_ << ")"; + << " (shared endpoint, max_wr=" << max_wr_depth_ + << ", provider tx queue=" << fi_info_->tx_attr->size + << ", max_cqe=" << max_cqe_ << ")"; return 0; } @@ -222,6 +245,46 @@ int EfaContext::buildSharedEndpoint(size_t max_wr, size_t max_inline) { LOG(ERROR) << "EfaContext::buildSharedEndpoint: no CQ available"; return ERR_CONTEXT; } + + // Design: adopt the provider's transmit depth instead of pacing against an + // independent number. We never tell libfabric how deep we want the queue + // (fi_endpoint() below takes its depth from fi_info_), and the depth is a + // per-device attribute -- 4096 on p5, 2048 on p6-b300 -- so no compiled-in + // default can match it. Either direction of disagreement hurts: too + // shallow and submitters exhaust credit while the queue is mostly empty; + // too deep and we hand out credit fi_write must refuse with FI_EAGAIN. + // Both were observed in production. + // + // Read-back rather than a hint: asking for more than the device supports + // makes the EFA provider fail fi_getinfo() with -FI_ENODATA, turning a + // misconfigured MC_MAX_WR into a failure to initialize. MC_MAX_WR still + // throttles a NIC when it asks for less; it can no longer ask for more. + const size_t provider_tx_depth = + fi_info_ && fi_info_->tx_attr ? fi_info_->tx_attr->size : 0; + if (provider_tx_depth == 0) { + LOG(ERROR) << "EfaContext::buildSharedEndpoint: provider reported no " + "transmit queue depth for " + << device_name_; + return ERR_CONTEXT; + } + if (!globalConfig().max_wr_from_env) { + // No override: take the provider's depth verbatim, so the default is + // correct on an instance type nobody has measured yet. + max_wr = provider_tx_depth; + } else if (max_wr > provider_tx_depth) { + // FIRST_N(1): MC_MAX_WR is process-wide, so an over-large value is one + // mistake, not one per NIC -- a p5 has 32. The value that took effect + // is in the per-device "EFA device" INFO line below. + LOG_FIRST_N(WARNING, 1) + << "EFA " << device_name_ << ": MC_MAX_WR=" << max_wr + << " exceeds the provider's transmit queue depth (" + << provider_tx_depth + << "); clamping. A larger value cannot increase the number of " + << "operations the NIC accepts -- it only hands out credit that " + << "fi_write must refuse with FI_EAGAIN. Unset MC_MAX_WR to track " + << "the provider's depth automatically."; + max_wr = provider_tx_depth; + } max_wr_depth_ = static_cast(max_wr); int ret = fi_endpoint(domain_, fi_info_, &shared_ep_, nullptr); @@ -344,6 +407,67 @@ int EfaContext::deconstruct() { return 0; } +#if defined(USE_CUDA) +// A silent failure here resurfaces as the provider's opaque "Operation not +// supported" from fi_mr_regattr(), which is the attribution problem this whole +// helper exists to remove -- so name the driver call that actually failed. +static void logCudaFailure(const char* call, CUresult ret, int device_ordinal) { + const char* err = nullptr; + cuGetErrorString(ret, &err); + LOG(WARNING) << "EFA: " << call << " failed for CUDA device " + << device_ordinal << ": " << (err ? err : "unknown") << " (" + << ret + << "); GPU memory registration may fail with a bare" + " \"Operation not supported\""; +} + +// Make `device_ordinal`'s primary context current on the calling thread, unless +// a context for that same device already is. +// +// Why a context is needed at all: fi_mr_regattr() on FI_HMEM_CUDA memory +// reaches libfabric's cuda_get_dmabuf_fd(), which calls the driver API +// cuMemGetHandleForAddressRange() and does no context management of its own. +// The export also runs against the CURRENT context, so a context belonging to +// another device is no better than none -- both end up as a bare "Operation not +// supported" from the provider. +// +// Who arrives here without the right context: registerLocalMemoryBatch() runs +// one std::async(std::launch::async) per buffer and registerLocalMemory() can +// fan out one std::thread per NIC, so a registering thread may have touched no +// CUDA API at all; and since std::async may reuse threads, one thread can +// register buffers on several devices in turn. +// +// cuDevicePrimaryCtxRetain() returns the same primary context the CUDA runtime +// uses, so this attaches to the process's existing context rather than creating +// another. The retain is intentionally not released: the primary context +// outlives every registration, and dropping the last reference here would tear +// down the context the rest of the process is using. +static void bindCudaContextIfNeeded(int device_ordinal) { + CUdevice want; + CUresult ret = cuDeviceGet(&want, device_ordinal); + if (ret != CUDA_SUCCESS) { + logCudaFailure("cuDeviceGet", ret, device_ordinal); + return; + } + + CUcontext cur = nullptr; + CUdevice cur_dev; + if (cuCtxGetCurrent(&cur) == CUDA_SUCCESS && cur != nullptr && + cuCtxGetDevice(&cur_dev) == CUDA_SUCCESS && cur_dev == want) + return; + + CUcontext primary = nullptr; + ret = cuDevicePrimaryCtxRetain(&primary, want); + if (ret != CUDA_SUCCESS) { + logCudaFailure("cuDevicePrimaryCtxRetain", ret, device_ordinal); + return; + } + ret = cuCtxSetCurrent(primary); + if (ret != CUDA_SUCCESS) + logCudaFailure("cuCtxSetCurrent", ret, device_ordinal); +} +#endif + int EfaContext::registerMemoryRegionInternal(void* addr, size_t length, int access, EfaMemoryRegionMeta& mrMeta) { @@ -388,6 +512,9 @@ int EfaContext::registerMemoryRegionInternal(void* addr, size_t length, int ret; if (iface != FI_HMEM_SYSTEM) { +#if defined(USE_CUDA) + bindCudaContextIfNeeded(device_ordinal); +#endif // GPU memory: use fi_mr_regattr with explicit iface and device struct iovec iov = {.iov_base = addr, .iov_len = length}; struct fi_mr_attr attr = {}; @@ -744,6 +871,20 @@ int EfaContext::submitPostSend( continue; } + // device_id comes from the peer-supplied topology, whose HCA list is + // independent of the peer 'devices' array, and selectDevice() bounds it + // against rkey only. decodeSegmentDesc() now rejects a descriptor whose + // key count and device count disagree; bound the value used to index + // devices[] locally as well. + if (static_cast(device_id) >= + peer_segment_desc->devices.size()) { + LOG(ERROR) << "Peer device index out of range for target " + << slice->target_id << ": device_id=" << device_id + << " devices=" << peer_segment_desc->devices.size(); + slice->markFailed(); + continue; + } + slice->rdma.dest_rkey = peer_segment_desc->buffers[buffer_id].rkey[device_id]; @@ -792,7 +933,9 @@ int EfaContext::submitSlicesOnPeer( // 2. Prepare MR descriptors and op contexts outside the lock // 3. Hold post_lock_ once for the entire batch of fi_write calls const int kMaxBackoffYields = 100000; - const int cq_limit = static_cast(globalConfig().max_cqe); + // Not globalConfig().max_cqe: that is the requested value, while max_cqe_ + // is what this device's CQ was actually opened with (see construct()). + const int cq_limit = static_cast(max_cqe_); std::atomic* cq_outstanding = shared_cq_ ? &shared_cq_->outstanding : nullptr; diff --git a/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp index a2247d7a41..9178915af3 100644 --- a/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp +++ b/mooncake-transfer-engine/src/transport/efa_transport/efa_transport.cpp @@ -20,12 +20,15 @@ #include #include +#include #include #include #include #include #include +#include #include +#include #include #include @@ -59,8 +62,12 @@ static size_t detectBufferPageSize(void* addr) { bool in_range = false; while (std::getline(smaps, line)) { - // VMA header: "start-end perms offset dev inode [pathname]" - if (!line.empty() && std::isxdigit(line[0])) { + // VMA header: "start-end perms offset dev inode [pathname]". + // Cast to unsigned char before std::isxdigit: passing a (possibly + // signed) char whose value is > 0x7F is UB, since the argument must be + // representable as unsigned char or equal EOF. + if (!line.empty() && + std::isxdigit(static_cast(line[0]))) { unsigned long start = 0, end = 0; if (sscanf(line.c_str(), "%lx-%lx", &start, &end) == 2) { in_range = (target >= start && target < end); @@ -307,8 +314,37 @@ int EfaTransport::registerLocalMemoryInternal(void* addr, size_t length, size_t total_pages_per_nic = length / page_size; bool use_full_coverage = (total_pages_per_nic <= getMaxPteEntries()); + // Optionally narrow a single-chunk DEVICE buffer to the NICs the topology + // reports as closest to its GPU. See buildLocalNicMap() for where the set + // comes from, and the MC_EFA_NIC_SELECTION docs for the trade-off: fewer + // NICs can serve a transfer touching this buffer, so this is registration + // time bought with potential per-buffer bandwidth, not a free win. + // + // Registering one buffer on N NICs costs ~N times one registration, and for + // device memory the per-domain accumulation described above is paid once + // per domain, so the two multiply. Measured on p5.48xlarge, 48 x 391 MB GPU + // buffers registered serially: 123.7 s on 32 NICs, 17.6 s on 4, 4.4 s on 1. + // + // Host memory is deliberately excluded. It has no accumulation to amplify, + // it is ~4x cheaper to register in the first place, and a host buffer has + // no single owning device -- the topology's "cpu:N" preferred set is a NUMA + // node's NICs, which is half the machine rather than a rail group. + const std::vector* local_nics = nullptr; + if (!local_nic_map_.empty() && resolved_name.rfind("cpu", 0) != 0) { + auto it = local_nic_map_.find(resolved_name); + if (it != local_nic_map_.end()) local_nics = &it->second; + } + std::vector> nic_assignments(num_chunks); - if (chunks.size() <= 1) { + if (chunks.size() <= 1 && local_nics) { + // Single chunk, topology-local NIC selection requested for this device. + nic_assignments[0] = *local_nics; + if (globalConfig().trace) { + LOG(INFO) << "EFA local NIC selection: " << addr << " (" + << resolved_name << ") on " << local_nics->size() << "/" + << num_nics << " NICs"; + } + } else if (chunks.size() <= 1) { // Single chunk: all NICs for (size_t n = 0; n < num_nics; ++n) { nic_assignments[0].push_back(n); @@ -467,20 +503,22 @@ int EfaTransport::registerLocalMemoryInternal(void* addr, size_t length, reg_start) .count(); + // Per-chunk detail is trace-only: a batch registration emits one of + // these per chunk, which is ~1450 lines at Kimi-K3 startup. Note that + // reg_duration_ms is the wall time this thread spent in the call, so + // when several buffers register concurrently it includes time waiting + // on the provider's locks, not just this chunk's own work -- compare it + // against the batch total logged by registerLocalMemoryBatch(). if (globalConfig().trace) { - LOG(INFO) << "EFA registerMemoryRegion: chunk " << ci - << ", addr=" << chunk_addr << ", length=" << chunk_len + LOG(INFO) << "EFA registerMemoryRegion: chunk " << ci << "/" + << chunks.size() << ", addr=" << chunk_addr + << ", length=" << chunk_len << ", nics=" << assigned_nics.size() << "/" << context_list_.size() << ", parallel=" << (use_parallel_reg ? "true" : "false") << ", duration=" << reg_duration_ms << "ms"; } - LOG(WARNING) << "Chunk " << ci << "/" << chunks.size() - << " registered on " << assigned_nics.size() << " NICs" - << ", addr=" << chunk_addr << ", length=" << chunk_len - << ", duration=" << reg_duration_ms << "ms"; - // Collect keys: assigned NICs have valid keys, others get 0 BufferDesc buffer_desc; for (auto& context : context_list_) { @@ -620,46 +658,158 @@ int EfaTransport::allocateLocalSegmentID() { return 0; } +// Optional concurrency cap for the batch register/unregister fan-out below, set +// via MC_MAX_CONCURRENT_REG_MR. Unset (0) means unbounded -- one thread per +// buffer, which is what these entry points have always done. +// +// The cap is PER PROCESS, and registration is CPU-bound page-pinning, so what +// matters is cap x processes against the core count. Measured on p5.48xlarge +// (192 cores) replaying Kimi-K3's KV registration as SGLang issues it -- one +// TransferEngine per TP rank, 182 GPU buffers of 2.5 KB to 391 MB each, slowest +// rank -- the optimum tracks the cores and not the cap: +// +// 8 ranks, 192 cores -> cap 16 (99.8 s); cap 64 is 1.8x slower +// 4 ranks, 192 cores -> cap 32 (89.2 s); same 128 threads as above +// 8 ranks, 64 cores -> cap 8 (121.3 s); 64 threads, tracks the budget +// +// So a good value is roughly cores/processes. Oversubscribing costs more than +// undersubscribing. This cannot be a built-in default because a single engine +// does not know how many peer processes share the node; picking one from the +// core count alone would oversubscribe by exactly the rank count. +// +// Input order matters as much as the cap: at cap 16 the same batch takes 43 s +// in SGLang's pool order but 95-98 s largest-first. Unbounded ignores order +// (99-128 s) since nothing queues. Largest-first being worst is backwards from +// longest-processing-time scheduling and is unexplained, so no sort is applied +// here yet -- another reason the cap stays opt-in rather than a default. +static size_t maxConcurrentRegMr() { + size_t configured = globalConfig().max_concurrent_reg_mr; + // 0 (unset) means unbounded, i.e. one thread per buffer as before. + return configured > 0 ? configured : std::numeric_limits::max(); +} + +// Run `fn(i)` for i in [0, count) on at most maxConcurrentRegMr() threads, +// returning the first non-zero result (all items are still attempted). +// +// With no cap set this spawns count-1 threads and runs the caller as a worker, +// which reproduces what both batch entry points did before: one +// std::async(std::launch::async) per buffer, which libstdc++ honours literally +// as one fresh thread each. Kimi-K3 registers ~1450 KV buffers at once, so a +// 192-core node peaks at ~1400 runnable threads inside fi_mr_regattr, and the +// per-buffer duration this logs inflates with the queueing delay of the ones +// ahead of it -- on 2x p6-b300 the median reached 106 s while the whole batch +// took 138 s. That inflation is not by itself a reason to cap: a badly chosen +// cap is slower still, see maxConcurrentRegMr(). +// +// A plain thread pool rather than a semaphore over std::async, because if a cap +// is set the point is to avoid the thread *creation*, not just to gate entry +// into the provider -- admission control after the thread already exists would +// leave that cost in place. +static int runBoundedParallel(size_t count, + const std::function& fn) { + if (count == 0) return 0; + + size_t workers = std::min(count, maxConcurrentRegMr()); + std::atomic next{0}; + std::atomic first_error{0}; + + auto worker = [&]() { + for (size_t i = next.fetch_add(1); i < count; i = next.fetch_add(1)) { + int ret = fn(i); + if (ret) { + int expected = 0; + first_error.compare_exchange_strong(expected, ret); + } + } + }; + + std::vector threads; + threads.reserve(workers - 1); + for (size_t w = 1; w < workers; ++w) threads.emplace_back(worker); + worker(); // the caller is a worker too + for (auto& t : threads) t.join(); + + return first_error.load(); +} + +// Invert the topology's per-location preferred_hca lists into indices into the +// NIC set the transport actually opened. `Topology::discover()` already +// computes what matters here: for a "cuda:N" entry, preferred_hca holds the +// HCAs at the minimum PCI distance from that GPU, restricted to its NUMA node. +// On p5.48xlarge each GPU and 4 of the 32 EFA devices sit under the same PCIe +// root complex, so preferred_hca is exactly those 4 -- the same set NIXL's +// libfabric plugin derives from hwloc in getEfaDevicesForPci(). No new PCI walk +// needed. +// +// Preferred rather than avail: avail_hca is "every other NIC", which is the +// all-NICs behavior this exists to avoid. +std::unordered_map> +EfaTransport::buildLocalNicMap(const TopologyMatrix& matrix, + const std::vector& device_names) { + std::unordered_map name_to_index; + for (size_t i = 0; i < device_names.size(); ++i) + name_to_index[device_names[i]] = i; + + std::unordered_map> local_nic_map; + for (const auto& entry : matrix) { + std::vector indices; + for (const auto& hca : entry.second.preferred_hca) { + auto it = name_to_index.find(hca); + // A preferred HCA can be missing here: non-EFA devices are filtered + // out of context_list_, and an EFA device whose construct() failed + // was dropped. Skipping it leaves the rest of the set usable. + if (it != name_to_index.end()) indices.push_back(it->second); + } + // No entry at all rather than an empty one, so the caller's lookup miss + // falls back to all NICs instead of registering the buffer nowhere. + if (!indices.empty()) local_nic_map[entry.first] = std::move(indices); + } + return local_nic_map; +} + int EfaTransport::registerLocalMemoryBatch( const std::vector& buffer_list, const std::string& location) { - std::vector> results; - for (auto& buffer : buffer_list) { - results.emplace_back( - std::async(std::launch::async, [this, buffer, location]() -> int { - return registerLocalMemoryInternal(buffer.addr, buffer.length, - location, true, false, true); - })); - } + auto start = std::chrono::steady_clock::now(); - for (size_t i = 0; i < buffer_list.size(); ++i) { - if (results[i].get()) { + int first_error = runBoundedParallel(buffer_list.size(), [&](size_t i) { + int ret = registerLocalMemoryInternal(buffer_list[i].addr, + buffer_list[i].length, location, + true, false, true); + if (ret) { LOG(WARNING) << "EfaTransport: Failed to register memory: addr " << buffer_list[i].addr << " length " << buffer_list[i].length; } - } + return ret; + }); + + auto elapsed = std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count(); + LOG(INFO) << "EfaTransport: registered " << buffer_list.size() + << " buffers on " + << std::min(buffer_list.size(), maxConcurrentRegMr()) + << " threads in " << elapsed << "ms"; + + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } int EfaTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { - std::vector> results; - for (auto& addr : addr_list) { - results.emplace_back( - std::async(std::launch::async, [this, addr]() -> int { - return unregisterLocalMemoryInternal(addr, false, true); - })); - } - - for (size_t i = 0; i < addr_list.size(); ++i) { - if (results[i].get()) + int first_error = runBoundedParallel(addr_list.size(), [&](size_t i) { + int ret = unregisterLocalMemoryInternal(addr_list[i], false, true); + if (ret) { LOG(WARNING) << "EfaTransport: Failed to unregister memory: addr " << addr_list[i]; - } + } + return ret; + }); - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } int EfaTransport::warmupSegment(const std::string& segment_name) { @@ -775,7 +925,17 @@ Status EfaTransport::submitTransfer( size_t task_id = batch_desc.task_list.size(); batch_desc.task_list.resize(task_id + entries.size()); std::vector task_list; - for (auto& task : batch_desc.task_list) task_list.push_back(&task); + for (auto& request : entries) { + auto& task = batch_desc.task_list[task_id]; + ++task_id; + task.batch_id = batch_id; +#ifdef USE_ASCEND_HETEROGENEOUS + task.request = const_cast(&request); +#else + task.request = &request; +#endif + task_list.push_back(&task); + } return submitTransferTask(task_list); } @@ -1030,6 +1190,29 @@ int EfaTransport::initializeEfaResources() { return ERR_DEVICE_NOT_FOUND; } + if (globalConfig().efa_nic_selection == EfaNicSelection::LOCAL) { + std::vector context_names; + context_names.reserve(context_list_.size()); + for (auto& context : context_list_) + context_names.push_back(context->deviceName()); + local_nic_map_ = + buildLocalNicMap(local_topology_->getMatrix(), context_names); + for (auto& entry : local_nic_map_) { + std::string nic_list; + for (size_t i = 0; i < entry.second.size(); ++i) { + if (i > 0) nic_list += ","; + nic_list += context_names[entry.second[i]]; + } + LOG(INFO) << "EfaTransport: local NICs for " << entry.first << ": [" + << nic_list << "]"; + } + if (local_nic_map_.empty()) { + LOG(WARNING) << "EfaTransport: MC_EFA_NIC_SELECTION=local but the " + "topology reports no per-location preferred NICs; " + "every buffer will use all NICs as before"; + } + } + // Query EFA device max_mr_size via ibverbs and clamp globalConfig. // libfabric does not expose max_mr_size, so we go through the ibverbs // layer. diff --git a/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp index 929739f55a..71031c3b9f 100644 --- a/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp +++ b/mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp @@ -27,6 +27,7 @@ #include "common.h" #include "common/serialization.h" #include "config.h" +#include "hip_device_guard.h" #include "transfer_metadata.h" #include "transport/transport.h" @@ -251,16 +252,8 @@ static int setDeviceContext(void* source_ptr, int& device_id) { } static void setupP2PAccess(int num_devices) { - // Save the active device. The loop below calls hipSetDevice once per - // iteration; without restoring it before returning, the calling thread is - // left with its active device pinned to num_devices-1, causing downstream - // HIP calls on the same thread (e.g. PyTorch allocations in TP workers) to - // target the wrong GPU. - int original_device = -1; - if (!checkHip(hipGetDevice(&original_device), - "HipTransport: failed to get current device")) { - return; - } + // The loop switches devices; the guard restores the caller's on return. + HipDeviceGuard device_guard; auto clearStickyPeerAccessError = [](int src_device, int dst_device) { // hipDeviceEnablePeerAccess may leave hipErrorPeerAccessAlreadyEnabled @@ -311,12 +304,6 @@ static void setupP2PAccess(int num_devices) { } } } - - // Restore the active device so this function is transparent to the caller. - if (original_device >= 0) { - (void)checkHip(hipSetDevice(original_device), - "HipTransport: failed to restore device"); - } } static int getNumStreams() { @@ -440,11 +427,26 @@ int HipTransport::install(std::string& local_server_name, metadata_ = metadata; local_server_name_ = local_server_name; + // Compose with any local segment another transport (e.g. RDMA) already + // installed instead of overwriting it, so a single-node segment can + // advertise both protocols (e.g. "rdma,hip"). Work on a copy (the map may + // hand back a descriptor other threads are reading) and publish the new + // one atomically via addLocalSegment. + auto old_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); auto desc = std::make_shared(); if (!desc) return ERR_MEMORY; + if (old_desc) *desc = *old_desc; desc->name = local_server_name_; +#ifdef ENABLE_MULTI_PROTOCOL + if (desc->protocol.empty()) { + desc->protocol = "hip"; + } else if (desc->protocol.find("hip") == std::string::npos) { + desc->protocol += ",hip"; + } +#else desc->protocol = "hip"; +#endif metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_, std::move(desc)); @@ -457,6 +459,8 @@ Status HipTransport::startAsyncTransfer(const TransferRequest& request, hipError_t err; int device_id; + HipDeviceGuard device_guard; + if (setDeviceContext(request.source, device_id) != 0) { return Status::InvalidArgument("Failed to set device context"); } @@ -660,17 +664,23 @@ int HipTransport::registerLocalMemory(void* addr, size_t length, // IPC-based memory registration if (!use_fabric_mem_) { - // Validate memory type + // Only device memory can be exported via HIP IPC. Callers that register + // a batch across all transports (e.g. SGLang's PD metadata/aux buffers, + // which live in host memory) also hand those host buffers to this + // transport. Skip non-device memory gracefully and let RDMA/TCP + // register it; returning an error would roll back the entire + // multi-protocol batch and tear down every session. hipPointerAttribute_t attr; - if (!checkHip(hipPointerGetAttributes(&attr, addr), - "HipTransport: hipPointerGetAttributes failed")) { - return -1; - } - - if (attr.type != hipMemoryTypeDevice) { - LOG(ERROR) << "Unsupported memory type, " << addr << " " - << attr.type; - return -1; + hipError_t attr_err = hipPointerGetAttributes(&attr, addr); + if (attr_err != hipSuccess || attr.type != hipMemoryTypeDevice) { + // Clear any sticky error the failed query latched so it does not + // poison subsequent HIP calls made by the caller. + (void)hipGetLastError(); + if (globalConfig().trace) { + LOG(INFO) << "HipTransport: skipping non-device memory " << addr + << ", leaving it to other transports"; + } + return 0; } // Get IPC handle @@ -687,6 +697,9 @@ int HipTransport::registerLocalMemory(void* addr, size_t length, desc.length = length; desc.name = location; desc.shm_name = serializeBinaryData(&handle, sizeof(hipIpcMemHandle_t)); +#ifdef ENABLE_MULTI_PROTOCOL + desc.protocol = "hip"; +#endif return metadata_->addLocalMemoryBuffer(desc, true); } @@ -810,18 +823,20 @@ int HipTransport::registerLocalMemoryBatch( for (auto& buffer : buffer_list) { int rc = registerLocalMemory(buffer.addr, buffer.length, location, true, false); - if (rc < 0) return rc; + if (rc) return rc; } return metadata_->updateLocalSegmentDesc(); } int HipTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { + int first_error = 0; for (auto& addr : addr_list) { int rc = unregisterLocalMemory(addr, false); - if (rc < 0) return rc; + if (rc && !first_error) first_error = rc; } - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void* HipTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/CMakeLists.txt index 524891a635..5485952429 100644 --- a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/CMakeLists.txt @@ -3,5 +3,5 @@ file(GLOB INTRANODE_NVLINK_SOURCES "*.cpp") add_library(intranode_nvlink_transport OBJECT ${INTRANODE_NVLINK_SOURCES}) if (USE_CUDA) -target_include_directories(intranode_nvlink_transport PUBLIC CUDA::cudart "/usr/local/cuda/include") +target_include_directories(intranode_nvlink_transport PUBLIC ${CUDAToolkit_INCLUDE_DIRS}) endif() diff --git a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp index b4fbd80881..52ca987add 100644 --- a/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp +++ b/mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp @@ -33,58 +33,154 @@ #include "transfer_metadata.h" #include "transport/transport.h" +static bool checkCudaErrorReturn(cudaError_t result, const char *message) { + if (result != cudaSuccess) { + LOG(ERROR) << message << " (Error code: " << result << " - " + << cudaGetErrorString(result) << ")" << std::endl; + return false; + } + return true; +} + +namespace mooncake { + namespace { -struct CudaStreamNVLinkRAII { - cudaStream_t stream_; - CudaStreamNVLinkRAII() : stream_(nullptr) { - auto err = cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking); - if (err != cudaSuccess) { - LOG(FATAL) << "Failed to create NVLink CUDA stream: " << err - << " - " << cudaGetErrorString(err); + +/// Per-device CUDA stream pool (thread-local). +struct CudaStreamEntry { + cudaStream_t stream; + int device_id; +}; + +class PerDeviceStreamPool { + public: + CudaStreamEntry getOrCreate(int device_id) { + auto it = pool_.find(device_id); + if (it != pool_.end()) return it->second; + int saved_device = 0; + cudaGetDevice(&saved_device); + if (cudaSetDevice(device_id) != cudaSuccess) { + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaSetDevice(" + << device_id << ") failed"; + return {nullptr, -1}; + } + CudaStreamEntry entry; + entry.device_id = device_id; + cudaError_t err = + cudaStreamCreateWithFlags(&entry.stream, cudaStreamNonBlocking); + if (err != cudaSuccess) + LOG(FATAL) << "Failed to create NVLink CUDA stream on device " + << device_id << ": " << cudaGetErrorString(err); + cudaSetDevice(saved_device); + cudaDeviceProp prop; + std::string pci = "unknown"; + if (cudaGetDeviceProperties(&prop, device_id) == cudaSuccess) { + pci = std::string(prop.name) + " PCI " + + std::to_string(prop.pciBusID) + ":" + + std::to_string(prop.pciDeviceID); + } + const char *visible = getenv("CUDA_VISIBLE_DEVICES"); + LOG(INFO) << "IntraNodeNvlinkTransport: NVLink CUDA stream created on " + << "device " << device_id << " [physical: " << pci + << "] CUDA_VISIBLE_DEVICES=" + << (visible ? visible : "(not set)") << " pid=" << getpid(); + pool_[device_id] = entry; + return entry; + } + ~PerDeviceStreamPool() { + int saved_device = 0; + cudaGetDevice(&saved_device); + for (auto &kv : pool_) { + cudaSetDevice(kv.first); + if (kv.second.stream) cudaStreamDestroy(kv.second.stream); } + cudaSetDevice(saved_device); } - ~CudaStreamNVLinkRAII() { cudaStreamDestroy(stream_); } + + private: + std::unordered_map pool_; }; -static thread_local CudaStreamNVLinkRAII tl_nvlink_stream; - -/// Thread-local CUDA event for GPU-level stream synchronization. -/// Used to establish a GPU-visible dependency between cudaStreamPerThread -/// and nvlink_stream, which is required by cudaMemcpyBatchAsync's -/// srcAccessOrderStream attribute for cross-stream P2P copies. -struct CudaEventNVLinkRAII { - cudaEvent_t event_; - CudaEventNVLinkRAII() { - auto err = cudaEventCreateWithFlags(&event_, cudaEventDisableTiming); - if (err != cudaSuccess) { - LOG(FATAL) << "Failed to create NVLink CUDA sync event: " << err - << " - " << cudaGetErrorString(err); +static thread_local PerDeviceStreamPool tl_device_stream_pool; + +/// Per-device event pool (thread-local). Caches one event per device to +/// avoid repeated create/destroy on device switches, and ensures proper +/// cleanup when the thread exits. +class PerDeviceEventPool { + public: + cudaEvent_t getOrCreate(int device_id) { + auto it = pool_.find(device_id); + if (it != pool_.end()) return it->second; + int saved_device = 0; + cudaGetDevice(&saved_device); + if (cudaSetDevice(device_id) != cudaSuccess) { + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaSetDevice(" + << device_id << ") failed when creating event"; + return nullptr; + } + cudaEvent_t event = nullptr; + cudaError_t err = + cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + if (err != cudaSuccess) + LOG(FATAL) << "Failed to create NVLink sync event on device " + << device_id << ": " << cudaGetErrorString(err); + cudaSetDevice(saved_device); + pool_[device_id] = event; + return event; + } + ~PerDeviceEventPool() { + int saved_device = 0; + cudaGetDevice(&saved_device); + for (auto &kv : pool_) { + cudaSetDevice(kv.first); + if (kv.second) cudaEventDestroy(kv.second); } + cudaSetDevice(saved_device); } - ~CudaEventNVLinkRAII() { cudaEventDestroy(event_); } + + private: + std::unordered_map pool_; }; -static thread_local CudaEventNVLinkRAII tl_nvlink_sync_event; -} // anonymous namespace +static thread_local PerDeviceEventPool tl_device_event_pool; -static bool checkCudaErrorReturn(cudaError_t result, const char *message) { - if (result != cudaSuccess) { - LOG(ERROR) << message << " (Error code: " << result << " - " - << cudaGetErrorString(result) << ")" << std::endl; - return false; +static cudaEvent_t getCallerSyncEvent() { + int current_device = 0; + cudaGetDevice(¤t_device); + return tl_device_event_pool.getOrCreate(current_device); +} + +static int getDeviceForPointer(const void *ptr) { + cudaPointerAttributes attr; + if (cudaPointerGetAttributes(&attr, ptr) != cudaSuccess) { + cudaGetLastError(); + return -1; } - return true; + return (attr.type == cudaMemoryTypeDevice) ? attr.device : -1; } -namespace mooncake { +static CudaStreamEntry getStreamForRequest(const void *source) { + int device_id = getDeviceForPointer(source); + if (device_id < 0) { + cudaGetDevice(&device_id); + if (device_id < 0) device_id = 0; + } + return tl_device_stream_pool.getOrCreate(device_id); +} -typedef Transport::Slice Slice; +} // anonymous namespace + +using Slice = Transport::Slice; /// Submit batched memcpy operations using cudaMemcpyBatchAsync when available /// (CUDA 12.8+), falling back to per-slice cudaMemcpyAsync otherwise. -/// Uses cudaMemcpySrcAccessOrderStream attribute to respect stream access -/// ordering semantics. Individual slice errors are tracked so that slices -/// whose memcpy failed are marked as FAILED while successful ones are POSTED. +/// Uses cudaMemcpySrcAccessOrderStream to ensure source data visibility for +/// P2P copies. This attribute is REQUIRED — without it, the GPU does not +/// insert the necessary memory barriers for P2P access, causing segfaults. +/// The caller must also establish GPU-level stream synchronization +/// (via cudaEventRecord + cudaStreamWaitEvent) before calling this function. +/// Individual slice errors are tracked so that slices whose memcpy failed +/// are marked as FAILED while successfully submitted ones are POSTED. static void submitBatchMemcpy(const std::vector &slices, const std::vector &srcs, const std::vector &dsts, @@ -113,15 +209,18 @@ static void submitBatchMemcpy(const std::vector &slices, (void)logged_once; #if CUDART_VERSION >= 12080 - // Use srcAccessOrderStream for P2P copies. The GPU-level dependency - // established by cudaEventRecord + cudaStreamWaitEvent in the caller - // ensures the source data is visible through nvlink_stream's access - // order, satisfying srcAccessOrderStream's requirement. + // srcAccessOrderStream is REQUIRED for P2P copies — without it, the GPU + // does not insert necessary memory barriers for cross-device access, + // resulting in segmentation faults. The caller also establishes a + // GPU-level dependency via cudaEventRecord(cudaStreamPerThread) + + // cudaStreamWaitEvent(nvlink_stream) to ensure source data is coherent + // before the memcpy starts; these two mechanisms are complementary. cudaMemcpyAttributes attr{}; attr.srcAccessOrder = cudaMemcpySrcAccessOrderStream; size_t attrs_idx = 0; // cudaMemcpyBatchAsync in CUDA 12.8 takes non-const size_t* for sizes std::vector mutable_sizes(sizes); + size_t fail_idx = count; #endif #if CUDART_VERSION >= 13000 @@ -129,26 +228,56 @@ static void submitBatchMemcpy(const std::vector &slices, const_cast(srcs.data()), mutable_sizes.data(), static_cast(count), &attr, &attrs_idx, 1, stream); + if (err != cudaSuccess) { + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaMemcpyBatchAsync " + << "failed: " << cudaGetErrorString(err); + // CUDA >= 13.0 does not return fail_idx; conservatively mark all + // as FAILED since we cannot determine which copies succeeded. + for (size_t i = 0; i < count; ++i) { + if (slices[i]->status == Slice::PENDING) { + slices[i]->markFailed(); + } + } + } else { + for (size_t i = 0; i < count; ++i) { + slices[i]->status = Slice::POSTED; + slices[i]->local.cuda_stream = (void *)stream; + } + } #elif CUDART_VERSION >= 12080 - { - size_t fail_idx = count; - err = cudaMemcpyBatchAsync( - const_cast(dsts.data()), const_cast(srcs.data()), - mutable_sizes.data(), static_cast(count), &attr, &attrs_idx, - 1, &fail_idx, stream); - if (err != cudaSuccess) { - if (fail_idx < count) { - LOG(ERROR) << "IntraNodeNvlinkTransport: cudaMemcpyBatchAsync " - << "failed at index " << fail_idx - << " (src=" << srcs[fail_idx] - << ", dst=" << dsts[fail_idx] - << ", size=" << sizes[fail_idx] - << "): " << cudaGetErrorString(err); - } else { - LOG(ERROR) << "IntraNodeNvlinkTransport: cudaMemcpyBatchAsync " - << "failed: " << cudaGetErrorString(err); + err = cudaMemcpyBatchAsync(const_cast(dsts.data()), + const_cast(srcs.data()), + mutable_sizes.data(), static_cast(count), + &attr, &attrs_idx, 1, &fail_idx, stream); + if (err != cudaSuccess) { + if (fail_idx < count) { + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaMemcpyBatchAsync " + << "failed at index " << fail_idx + << " (src=" << srcs[fail_idx] + << ", dst=" << dsts[fail_idx] + << ", size=" << sizes[fail_idx] + << "): " << cudaGetErrorString(err); + } else { + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaMemcpyBatchAsync " + << "failed: " << cudaGetErrorString(err); + } + // Copies [0, fail_idx) were submitted successfully → POSTED. + // Copy [fail_idx] failed → FAILED. + // Copies (fail_idx, count) were never submitted → FAILED. + for (size_t i = 0; i < fail_idx; ++i) { + slices[i]->status = Slice::POSTED; + slices[i]->local.cuda_stream = (void *)stream; + } + for (size_t i = fail_idx; i < count; ++i) { + if (slices[i]->status == Slice::PENDING) { + slices[i]->markFailed(); } } + } else { + for (size_t i = 0; i < count; ++i) { + slices[i]->status = Slice::POSTED; + slices[i]->local.cuda_stream = (void *)stream; + } } #else // Fallback for CUDA < 12.8: submit each memcpy individually @@ -167,20 +296,6 @@ static void submitBatchMemcpy(const std::vector &slices, } return; // Slice states already set above #endif - - // For cudaMemcpyBatchAsync paths, update slice states based on result - if (err != cudaSuccess) { - for (size_t i = 0; i < count; ++i) { - if (slices[i]->status == Slice::PENDING) { - slices[i]->markFailed(); - } - } - } else { - for (size_t i = 0; i < count; ++i) { - slices[i]->status = Slice::POSTED; - slices[i]->local.cuda_stream = (void *)stream; - } - } } static int getNumDevices() { @@ -317,33 +432,27 @@ Status IntraNodeNvlinkTransport::submitTransfer( size_t task_id = batch_desc.task_list.size(); batch_desc.task_list.resize(task_id + entries.size()); - // Synchronize with the caller's CUDA stream before issuing any memcpy. - // PyTorch uses cudaStreamPerThread (per-thread default stream), NOT the - // legacy default stream (nullptr). Recording an event on - // cudaStreamPerThread and making nvlink_stream wait for it ensures that all - // PyTorch GPU operations on source/dest buffers complete before the NVLink - // memcpy starts. This GPU-level dependency is also required by - // cudaMemcpyBatchAsync's srcAccessOrderStream attribute, which needs - // the source data to be visible through the nvlink_stream's access order. - // - // Do NOT use the legacy default stream (nullptr) for cudaEventRecord, - // as it would trigger implicit synchronization with blocking streams and - // could cause deadlocks. - cudaStream_t stream = tl_nvlink_stream.stream_; - cudaError_t sync_err = - cudaEventRecord(tl_nvlink_sync_event.event_, cudaStreamPerThread); + // Get per-device transfer stream for the source buffer's device. + CudaStreamEntry stream_entry = + getStreamForRequest(entries.empty() ? nullptr : entries[0].source); + cudaStream_t stream = stream_entry.stream; + if (!stream) return Status::Context("Failed to create NVLink CUDA stream"); + // Synchronize with caller's GPU work via cudaEventSynchronize + // (CPU-blocking) to avoid expensive cross-device cudaStreamWaitEvent on + // non-NVIDIA GPUs. + cudaEvent_t sync_event = getCallerSyncEvent(); + cudaError_t sync_err = cudaEventRecord(sync_event, cudaStreamPerThread); if (sync_err != cudaSuccess) { - LOG(ERROR) << "IntraNodeNvlinkTransport: cudaEventRecord on " - "cudaStreamPerThread failed: " + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaEventRecord failed: " << cudaGetErrorString(sync_err); return Status::Context("cudaEventRecord failed: " + std::string(cudaGetErrorString(sync_err))); } - sync_err = cudaStreamWaitEvent(stream, tl_nvlink_sync_event.event_, 0); + sync_err = cudaEventSynchronize(sync_event); if (sync_err != cudaSuccess) { - LOG(ERROR) << "IntraNodeNvlinkTransport: cudaStreamWaitEvent failed: " + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaEventSynchronize failed: " << cudaGetErrorString(sync_err); - return Status::Context("cudaStreamWaitEvent failed: " + + return Status::Context("cudaEventSynchronize failed: " + std::string(cudaGetErrorString(sync_err))); } @@ -404,17 +513,24 @@ Status IntraNodeNvlinkTransport::getTransferStatus(BatchID batch_id, std::to_string(batch_id)); } auto &task = batch_desc.task_list[task_id]; - // Poll POSTED slices for async completion via cudaStreamQuery + // Poll POSTED slices for async completion via cudaStreamQuery. + std::unordered_map stream_status_cache; for (auto *slice : task.slice_list) { if (slice && slice->status == Slice::POSTED) { cudaStream_t stream = (cudaStream_t)slice->local.cuda_stream; - cudaError_t cuda_err = cudaStreamQuery(stream); + auto it = stream_status_cache.find(stream); + cudaError_t cuda_err; + if (it == stream_status_cache.end()) { + cuda_err = cudaStreamQuery(stream); + stream_status_cache[stream] = cuda_err; + } else { + cuda_err = it->second; + } if (cuda_err == cudaSuccess) { slice->markSuccess(); } else if (cuda_err != cudaErrorNotReady) { slice->markFailed(); } - // cudaErrorNotReady means still in progress, keep POSTED } } status.transferred_bytes = task.transferred_bytes; @@ -435,23 +551,24 @@ Status IntraNodeNvlinkTransport::getTransferStatus(BatchID batch_id, Status IntraNodeNvlinkTransport::submitTransferTask( const std::vector &task_list) { - // Synchronize with the caller's CUDA stream before issuing any memcpy. - // See submitTransfer() for detailed rationale. - cudaStream_t stream = tl_nvlink_stream.stream_; - cudaError_t sync_err = - cudaEventRecord(tl_nvlink_sync_event.event_, cudaStreamPerThread); + // Get per-device transfer stream. See submitTransfer() for rationale. + CudaStreamEntry stream_entry = getStreamForRequest( + task_list.empty() ? nullptr : task_list[0]->request->source); + cudaStream_t stream = stream_entry.stream; + if (!stream) return Status::Context("Failed to create NVLink CUDA stream"); + cudaEvent_t sync_event = getCallerSyncEvent(); + cudaError_t sync_err = cudaEventRecord(sync_event, cudaStreamPerThread); if (sync_err != cudaSuccess) { - LOG(ERROR) << "IntraNodeNvlinkTransport: cudaEventRecord on " - "cudaStreamPerThread failed: " + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaEventRecord failed: " << cudaGetErrorString(sync_err); return Status::Context("cudaEventRecord failed: " + std::string(cudaGetErrorString(sync_err))); } - sync_err = cudaStreamWaitEvent(stream, tl_nvlink_sync_event.event_, 0); + sync_err = cudaEventSynchronize(sync_event); if (sync_err != cudaSuccess) { - LOG(ERROR) << "IntraNodeNvlinkTransport: cudaStreamWaitEvent failed: " + LOG(ERROR) << "IntraNodeNvlinkTransport: cudaEventSynchronize failed: " << cudaGetErrorString(sync_err); - return Status::Context("cudaStreamWaitEvent failed: " + + return Status::Context("cudaEventSynchronize failed: " + std::string(cudaGetErrorString(sync_err))); } @@ -645,15 +762,23 @@ int IntraNodeNvlinkTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, int IntraNodeNvlinkTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int IntraNodeNvlinkTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, false); - return metadata_->updateLocalSegmentDesc(); + int first_error = 0; + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret && !first_error) first_error = ret; + } + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void *IntraNodeNvlinkTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt index 5583b7d813..e8fae49d54 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt @@ -1,32 +1,20 @@ file(GLOB UB_SOURCES "*.cpp" "urma/urma_endpoint.cpp" "ub_allocator.cpp") -# Check if liburma.so exists -find_library(URMA_LIBRARY urma PATHS /usr/lib64) - -# Build ub_transport with real URMA library if available -# Never include mock_urma_api.cpp in ub_transport library -# For test, we will use mock_urma_api.cpp directly -if (URMA_LIBRARY) - add_library(ub_transport OBJECT ${UB_SOURCES}) - target_link_libraries(ub_transport - PUBLIC - ${URMA_LIBRARY} - ) - message(STATUS "Using real URMA library: ${URMA_LIBRARY}") +# Build ub_transport with real URMA library if available Never include +# mock_urma_api.cpp in ub_transport library For test, we will use +# mock_urma_api.cpp directly +if(URMA_LIBRARY) + add_library(ub_transport OBJECT ${UB_SOURCES}) + message(STATUS "Using real URMA library: ${URMA_LIBRARY}") else() - # If liburma.so not found, we'll need to handle this differently - # For now, just build without mock_urma_api.cpp - list(APPEND UB_SOURCES "urma/mock_urma.cpp") - add_library(ub_transport OBJECT ${UB_SOURCES}) - message(WARNING "Not Found liburma.so building ub_transport with Mock URMA library") + # If liburma.so not found, we'll need to handle this differently For now, just + # build without mock_urma_api.cpp + list(APPEND UB_SOURCES "urma/mock_urma.cpp") + add_library(ub_transport OBJECT ${UB_SOURCES}) + message( + WARNING "Not Found liburma.so building ub_transport with Mock URMA library") endif() -target_include_directories(ub_transport - PUBLIC - ${urma_INCLUDE_DIR} -) -target_link_libraries(ub_transport - PRIVATE - JsonCpp::JsonCpp - glog::glog - pthread -) \ No newline at end of file +target_link_libraries( + ub_transport + PUBLIC Urma::urma + PRIVATE JsonCpp::JsonCpp glog::glog pthread) diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp index 77b72c1421..cf443a0ba5 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp @@ -142,13 +142,17 @@ int UbTransport::registerLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < buffer_list.size(); ++i) { - if (results[i].get()) { + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "UbTransport: Failed to register memory: addr " << buffer_list[i].addr << " length " << buffer_list[i].length; + if (!first_error) first_error = ret; } } + if (first_error) return first_error; return metadata_->updateLocalSegmentDesc(); } @@ -164,13 +168,17 @@ int UbTransport::unregisterLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < addr_list.size(); ++i) { - if (results[i].get()) + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "UbTransport: Failed to unregister memory: addr " << addr_list[i]; + if (!first_error) first_error = ret; + } } - - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status UbTransport::submitTransfer( @@ -186,8 +194,18 @@ Status UbTransport::submitTransfer( size_t task_id = batch_desc.task_list.size(); batch_desc.task_list.resize(task_id + entries.size()); std::vector task_list; - task_list.reserve(batch_desc.task_list.size()); - for (auto& task : batch_desc.task_list) task_list.push_back(&task); + task_list.reserve(entries.size()); + for (auto& request : entries) { + auto& task = batch_desc.task_list[task_id]; + ++task_id; + task.batch_id = batch_id; +#ifdef USE_ASCEND_HETEROGENEOUS + task.request = const_cast(&request); +#else + task.request = &request; +#endif + task_list.push_back(&task); + } return submitTransferTask(task_list); } @@ -269,8 +287,13 @@ Status UbTransport::submitTransferTask( } if (device_id < 0) { auto source_addr = slice->source_addr; - for (auto& entry : slices_to_post) - for (auto s : entry.second) getSliceCache().deallocate(s); + // Do not deallocate slices already queued in slices_to_post + // here: every slice is also recorded in its owning + // TransferTask::slice_list right after allocation, and + // ~TransferTask() returns everything in slice_list to the + // cache exactly once. Deallocating here double-frees them + // into ThreadLocalSliceCache, letting a later allocate() + // hand the same Slice* to two unrelated transfers. LOG(ERROR) << "UbTransport: Address not registered by any device(s) " << source_addr; diff --git a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp index c22e5e75a2..06c65d0385 100644 --- a/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/kunpeng_transport/urma/urma_endpoint.cpp @@ -38,7 +38,9 @@ UrmaContext::~UrmaContext() { auto thisString = toString(); worker_pool_.reset(); LOG(INFO) << "destroy worker pool done."; - endpoint_store_->destroy(); + if (endpoint_store_) { + endpoint_store_->destroy(); + } LOG(INFO) << "destroy endpoint store done."; if (urma_context_) deconstruct(); LOG(WARNING) << "finished destroy context : " << thisString; diff --git a/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp b/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp index 517cea889b..c769a9f53b 100644 --- a/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp +++ b/mooncake-transfer-engine/src/transport/maca_transport/maca_transport.cpp @@ -22,8 +22,12 @@ #include #include #include +#include +#include #include #include +#include +#include #include "common.h" #include "common/serialization.h" @@ -125,6 +129,242 @@ static int getDeviceFromPointer(void *ptr) { return -1; } +enum class MacaCallerSyncMode { + Host, + Wait, +}; + +enum class MacaCopyApi { + Auto, + Default, + BatchFlag, +}; + +struct CallerSync { + MacaCallerSyncMode mode; + cudaEvent_t event; +}; + +static MacaCallerSyncMode callerSyncMode() { + static const MacaCallerSyncMode mode = [] { + const char *env = std::getenv("MC_MACA_CALLER_SYNC"); + if (!env) return MacaCallerSyncMode::Host; + + std::string value(env); + if (value == "host") return MacaCallerSyncMode::Host; + if (value == "wait") return MacaCallerSyncMode::Wait; + LOG(WARNING) << "MacaTransport: unknown MC_MACA_CALLER_SYNC=" << value + << ", falling back to host"; + return MacaCallerSyncMode::Host; + }(); + return mode; +} + +static const char *callerSyncModeName(MacaCallerSyncMode mode) { + switch (mode) { + case MacaCallerSyncMode::Host: + return "host"; + case MacaCallerSyncMode::Wait: + return "wait"; + } + return "unknown"; +} + +static MacaCopyApi copyApi() { + static const MacaCopyApi api = [] { + const char *env = std::getenv("MC_MACA_COPY_API"); + if (!env) return MacaCopyApi::Auto; + + std::string value(env); + if (value == "auto") return MacaCopyApi::Auto; + if (value == "default") return MacaCopyApi::Default; + if (value == "batchflag") return MacaCopyApi::BatchFlag; + LOG(WARNING) << "MacaTransport: unknown MC_MACA_COPY_API=" << value + << ", falling back to auto"; + return MacaCopyApi::Auto; + }(); + return api; +} + +static const char *copyApiName(MacaCopyApi api) { + switch (api) { + case MacaCopyApi::Auto: + return "auto"; + case MacaCopyApi::Default: + return "default"; + case MacaCopyApi::BatchFlag: + return "batchflag"; + } + return "unknown"; +} + +static size_t batchFlagMinBytes() { + static const size_t min_bytes = [] { + constexpr size_t kDefaultMinBytes = 1024ULL * 1024ULL; + const char *env = std::getenv("MC_MACA_BATCHFLAG_MIN_BYTES"); + if (!env) return kDefaultMinBytes; + + char *end = nullptr; + unsigned long long value = std::strtoull(env, &end, 0); + if (end == env || *end != '\0') { + LOG(WARNING) << "MacaTransport: unknown " + "MC_MACA_BATCHFLAG_MIN_BYTES=" + << env << ", falling back to " << kDefaultMinBytes; + return kDefaultMinBytes; + } + return static_cast(value); + }(); + return min_bytes; +} + +static bool shouldUseBatchFlag(MacaCopyApi api, size_t length) { + if (api == MacaCopyApi::BatchFlag) return true; + if (api == MacaCopyApi::Auto) return length >= batchFlagMinBytes(); + return false; +} + +class PerDeviceStreamPool { + public: + cudaStream_t getOrCreate(int device_id) { + auto iter = streams_.find(device_id); + if (iter != streams_.end()) return iter->second; + + int original_device = -1; + cudaGetDevice(&original_device); + if (!checkCudaErrorReturn(cudaSetDevice(device_id), + "MacaTransport: failed to set device")) { + if (original_device >= 0) cudaSetDevice(original_device); + return nullptr; + } + + cudaStream_t stream = nullptr; + cudaError_t err = + cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking); + if (original_device >= 0) cudaSetDevice(original_device); + if (!checkCudaErrorReturn( + err, "MacaTransport: cudaStreamCreateWithFlags failed")) { + return nullptr; + } + + streams_[device_id] = stream; + return stream; + } + + ~PerDeviceStreamPool() { + int original_device = -1; + cudaGetDevice(&original_device); + for (auto &entry : streams_) { + cudaSetDevice(entry.first); + cudaStreamDestroy(entry.second); + } + if (original_device >= 0) cudaSetDevice(original_device); + } + + private: + std::unordered_map streams_; +}; + +static thread_local PerDeviceStreamPool thread_local_stream_pool; + +class PerDeviceEventPool { + public: + cudaEvent_t getOrCreate(int device_id) { + auto iter = events_.find(device_id); + if (iter != events_.end()) return iter->second; + + int original_device = -1; + cudaGetDevice(&original_device); + if (!checkCudaErrorReturn(cudaSetDevice(device_id), + "MacaTransport: failed to set device")) { + if (original_device >= 0) cudaSetDevice(original_device); + return nullptr; + } + + cudaEvent_t event = nullptr; + cudaError_t err = + cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + if (original_device >= 0) cudaSetDevice(original_device); + if (!checkCudaErrorReturn( + err, "MacaTransport: cudaEventCreateWithFlags failed")) { + return nullptr; + } + + events_[device_id] = event; + return event; + } + + ~PerDeviceEventPool() { + int original_device = -1; + cudaGetDevice(&original_device); + for (auto &entry : events_) { + cudaSetDevice(entry.first); + cudaEventDestroy(entry.second); + } + if (original_device >= 0) cudaSetDevice(original_device); + } + + private: + std::unordered_map events_; +}; + +static thread_local PerDeviceEventPool thread_local_event_pool; + +static Status prepareCallerSync(CallerSync &sync) { + sync.mode = callerSyncMode(); + sync.event = nullptr; + + int current_device = 0; + cudaGetDevice(¤t_device); + cudaEvent_t event = thread_local_event_pool.getOrCreate(current_device); + if (!event) { + return Status::Context("MacaTransport: failed to create sync event"); + } + + cudaError_t err = cudaEventRecord(event, cudaStreamPerThread); + if (err != cudaSuccess) { + LOG(ERROR) << "MacaTransport: cudaEventRecord failed: " + << cudaGetErrorString(err); + return Status::Context("MacaTransport: cudaEventRecord failed"); + } + + sync.event = event; + if (sync.mode == MacaCallerSyncMode::Host) { + err = cudaEventSynchronize(event); + if (err != cudaSuccess) { + LOG(ERROR) << "MacaTransport: cudaEventSynchronize failed: " + << cudaGetErrorString(err); + return Status::Context( + "MacaTransport: cudaEventSynchronize failed"); + } + } + return Status::OK(); +} + +static void getCopyEndpoints(Transport::Slice *slice, void *&dst, + const void *&src) { + dst = slice->local.dest_addr; + src = slice->source_addr; + if (slice->opcode == Transport::TransferRequest::READ) { + dst = slice->source_addr; + src = slice->local.dest_addr; + } +} + +static cudaError_t submitMemcpyAsync(Transport::Slice *slice, + cudaStream_t stream) { + void *dst; + const void *src; + getCopyEndpoints(slice, dst, src); + return cudaMemcpyAsync(dst, src, slice->length, cudaMemcpyDefault, stream); +} + +static cudaError_t submitBatchFlagAsync(std::vector ©_batch, + cudaStream_t stream) { + if (copy_batch.empty()) return cudaSuccess; + return mcExtBatchCopyFlagAndWaitV2(copy_batch.data(), copy_batch.size(), + nullptr, 0, stream); +} + MacaTransport::MacaTransport() { int num_devices = getNumDevices(); if (globalConfig().trace) { @@ -163,10 +403,21 @@ int MacaTransport::install(std::string &local_server_name, metadata_ = metadata; local_server_name_ = local_server_name; + auto old_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); auto desc = std::make_shared(); if (!desc) return ERR_MEMORY; + if (old_desc) *desc = *old_desc; + desc->name = local_server_name_; +#ifdef ENABLE_MULTI_PROTOCOL + if (desc->protocol.empty()) { + desc->protocol = "maca"; + } else if (desc->protocol.find("maca") == std::string::npos) { + desc->protocol += ",maca"; + } +#else desc->protocol = "maca"; +#endif metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_, std::move(desc)); return 0; @@ -189,6 +440,8 @@ Status MacaTransport::submitTransfer( for (auto &request : entries) { TransferTask &task = batch_desc.task_list[task_id]; ++task_id; + task.batch_id = batch_id; + task.transport_ = this; uint64_t dest_addr = request.target_offset; if (request.target_id != LOCAL_SEGMENT_ID) { int rc = relocateSharedMemoryAddress(dest_addr, request.length, @@ -204,6 +457,7 @@ Status MacaTransport::submitTransfer( slice->task = &task; slice->target_id = request.target_id; slice->status = Slice::PENDING; + task.slice_list.push_back(slice); __sync_fetch_and_add(&task.slice_count, 1); // Set correct device context before memcpy @@ -242,6 +496,28 @@ Status MacaTransport::getTransferStatus(BatchID batch_id, size_t task_id, std::to_string(batch_id)); } auto &task = batch_desc.task_list[task_id]; + std::unordered_map stream_status_cache; + for (auto *slice : task.slice_list) { + if (slice && slice->status == Slice::POSTED) { + cudaStream_t stream = (cudaStream_t)slice->local.cuda_stream; + auto iter = stream_status_cache.find(stream); + cudaError_t err; + if (iter == stream_status_cache.end()) { + err = cudaStreamQuery(stream); + stream_status_cache[stream] = err; + } else { + err = iter->second; + } + + if (err == cudaSuccess) { + slice->markSuccess(); + } else if (err != cudaErrorNotReady) { + LOG(ERROR) << "MacaTransport: cudaStreamQuery failed: " + << cudaGetErrorString(err); + slice->markFailed(); + } + } + } status.transferred_bytes = task.transferred_bytes; uint64_t success_slice_count = task.success_slice_count; uint64_t failed_slice_count = task.failed_slice_count; @@ -260,52 +536,244 @@ Status MacaTransport::getTransferStatus(BatchID batch_id, size_t task_id, Status MacaTransport::submitTransferTask( const std::vector &task_list) { + MacaCallerSyncMode sync_mode = callerSyncMode(); + static const bool logged_sync_mode = [sync_mode] { + LOG(INFO) << "MacaTransport: caller sync mode " + << callerSyncModeName(sync_mode); + return true; + }(); + (void)logged_sync_mode; + MacaCopyApi api = copyApi(); + static const bool logged_copy_api = [api] { + LOG(INFO) << "MacaTransport: copy api " << copyApiName(api); + return true; + }(); + (void)logged_copy_api; + static const bool logged_batchflag_threshold = [api] { + if (api == MacaCopyApi::Auto) { + LOG(INFO) << "MacaTransport: batchflag min bytes " + << batchFlagMinBytes(); + } + return true; + }(); + (void)logged_batchflag_threshold; + + CallerSync caller_sync; + Status sync_status = prepareCallerSync(caller_sync); + if (!sync_status.ok()) return sync_status; + + struct DeviceStream { + int device_id; + cudaStream_t stream; + bool ok; + std::vector copy_batch; + std::vector batch_slices; + }; + + std::vector streams; + std::unordered_map stream_index_by_device; + Status first_error = Status::OK(); + bool has_batchflag_copies = false; + + int original_device = -1; + cudaGetDevice(&original_device); + int active_copy_device = -1; + + auto getStream = [&](int device_id, cudaStream_t &stream, + size_t &stream_index) -> bool { + auto iter = stream_index_by_device.find(device_id); + if (iter != stream_index_by_device.end()) { + stream_index = iter->second; + stream = streams[stream_index].stream; + return true; + } + + if (!checkCudaErrorReturn(cudaSetDevice(device_id), + "MacaTransport: failed to set device")) { + return false; + } + + cudaStream_t new_stream = + thread_local_stream_pool.getOrCreate(device_id); + if (!new_stream) return false; + + if (caller_sync.mode == MacaCallerSyncMode::Wait && caller_sync.event) { + cudaError_t wait_err = + cudaStreamWaitEvent(new_stream, caller_sync.event, 0); + if (wait_err != cudaSuccess) { + LOG(ERROR) << "MacaTransport: cudaStreamWaitEvent failed: " + << cudaGetErrorString(wait_err); + return false; + } + } + + stream_index = streams.size(); + streams.push_back({device_id, new_stream, true, {}, {}}); + stream_index_by_device[device_id] = stream_index; + stream = new_stream; + return true; + }; + for (size_t index = 0; index < task_list.size(); ++index) { assert(task_list[index]); auto &task = *task_list[index]; assert(task.request); auto &request = *task.request; uint64_t dest_addr = request.target_offset; - if (request.target_id != LOCAL_SEGMENT_ID) { - int rc = relocateSharedMemoryAddress(dest_addr, request.length, - request.target_id); - if (rc) return Status::Memory("device memory not registered"); - } + task.total_bytes = request.length; Slice *slice = getSliceCache().allocate(); slice->source_addr = (char *)request.source; - slice->local.dest_addr = (char *)dest_addr; slice->length = request.length; slice->opcode = request.opcode; slice->task = &task; slice->target_id = request.target_id; slice->status = Slice::PENDING; + slice->ts = + globalConfig().slice_timeout > 0 ? getCurrentTimeInNano() : 0; task.slice_list.push_back(slice); __sync_fetch_and_add(&task.slice_count, 1); - // Set correct device context before memcpy - int original_device = -1; - cudaGetDevice(&original_device); + if (request.target_id != LOCAL_SEGMENT_ID) { + int rc = relocateSharedMemoryAddress(dest_addr, request.length, + request.target_id); + if (rc) { + slice->local.dest_addr = nullptr; + slice->markFailed(); + if (first_error.ok()) + first_error = + Status::Memory("device memory not registered"); + continue; + } + } + slice->local.dest_addr = (char *)dest_addr; + int target_device = getDeviceFromPointer(request.source); if (target_device < 0) target_device = getDeviceFromPointer((void *)dest_addr); - if (target_device >= 0) cudaSetDevice(target_device); + if (target_device < 0) { + slice->markFailed(); + if (first_error.ok()) + first_error = + Status::InvalidArgument("Cannot infer MACA device"); + continue; + } - cudaError_t err; - if (slice->opcode == TransferRequest::READ) - err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr, - slice->length, cudaMemcpyDefault); - else - err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr, - slice->length, cudaMemcpyDefault); - if (err != cudaSuccess) + cudaStream_t stream = nullptr; + size_t stream_index = 0; + if (!getStream(target_device, stream, stream_index)) { slice->markFailed(); - else - slice->markSuccess(); + if (first_error.ok()) + first_error = + Status::Memory("MacaTransport: failed to get MACA stream"); + continue; + } - if (original_device >= 0) cudaSetDevice(original_device); + if (active_copy_device != target_device) { + if (!checkCudaErrorReturn(cudaSetDevice(target_device), + "MacaTransport: failed to set device")) { + slice->markFailed(); + streams[stream_index].ok = false; + if (first_error.ok()) + first_error = + Status::Context("MacaTransport: failed to set device"); + continue; + } + active_copy_device = target_device; + } + + if (shouldUseBatchFlag(api, slice->length)) { + void *dst; + const void *src; + getCopyEndpoints(slice, dst, src); + + mcCopyFlag_t copy; + std::memset(©, 0, sizeof(copy)); + copy.dst = dst; + copy.src = src; + copy.engine = ParallelCopyEngineDefault; + copy.count = slice->length; + copy.waitNum = 0; + copy.writeNum = 0; + + streams[stream_index].copy_batch.push_back(copy); + streams[stream_index].batch_slices.push_back(slice); + slice->local.cuda_stream = (void *)stream; + has_batchflag_copies = true; + continue; + } + + cudaError_t err = submitMemcpyAsync(slice, stream); + if (err != cudaSuccess) { + LOG(ERROR) << "MacaTransport: async copy failed: " + << cudaGetErrorString(err); + slice->markFailed(); + streams[stream_index].ok = false; + if (first_error.ok()) + first_error = + Status::Memory("MacaTransport: async copy failed"); + } else { + slice->local.cuda_stream = (void *)stream; + slice->status = Slice::POSTED; + } } - return Status::OK(); + + if (has_batchflag_copies) { + auto failBatchSlices = [](DeviceStream &entry) { + for (auto *slice : entry.batch_slices) { + if (slice->status == Slice::PENDING) { + slice->markFailed(); + } + } + }; + + for (auto &entry : streams) { + if (entry.copy_batch.empty()) continue; + if (!entry.ok) { + failBatchSlices(entry); + if (first_error.ok()) + first_error = + Status::Memory("MacaTransport: batch copy skipped"); + continue; + } + if (!checkCudaErrorReturn(cudaSetDevice(entry.device_id), + "MacaTransport: failed to set device")) { + entry.ok = false; + failBatchSlices(entry); + if (first_error.ok()) + first_error = + Status::Context("MacaTransport: failed to set device"); + continue; + } + + cudaError_t err = + submitBatchFlagAsync(entry.copy_batch, entry.stream); + if (err != cudaSuccess) { + LOG(ERROR) + << "MacaTransport: mcExtBatchCopyFlagAndWaitV2 failed: " + << cudaGetErrorString(err); + entry.ok = false; + failBatchSlices(entry); + if (first_error.ok()) + first_error = + Status::Memory("MacaTransport: batch copy failed"); + } + } + + for (auto &entry : streams) { + if (entry.copy_batch.empty()) continue; + for (auto *slice : entry.batch_slices) { + if (slice->status != Slice::PENDING) continue; + if (entry.ok) + slice->status = Slice::POSTED; + else + slice->markFailed(); + } + } + } + + if (original_device >= 0) cudaSetDevice(original_device); + return first_error; } int MacaTransport::registerLocalMemory(void *addr, size_t length, @@ -360,6 +828,9 @@ int MacaTransport::registerLocalMemory(void *addr, size_t length, desc.length = alloc_size; desc.name = location; desc.shm_name = serializeBinaryData(&handle, sizeof(cudaIpcMemHandle_t)); +#ifdef ENABLE_MULTI_PROTOCOL + desc.protocol = "maca"; +#endif int rc = metadata_->addLocalMemoryBuffer(desc, true); if (rc == 0) { registered_base_addrs_.insert((uint64_t)base_ptr); @@ -449,15 +920,23 @@ int MacaTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, int MacaTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int MacaTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, false); - return metadata_->updateLocalSegmentDesc(); + int first_error = 0; + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret && !first_error) first_error = ret; + } + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void *MacaTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/nccl_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/nccl_transport/CMakeLists.txt new file mode 100644 index 0000000000..f4324ef1cc --- /dev/null +++ b/mooncake-transfer-engine/src/transport/nccl_transport/CMakeLists.txt @@ -0,0 +1,8 @@ +add_library(nccl_host_transport OBJECT nccl_transport.cpp) +target_include_directories( + nccl_host_transport + PRIVATE ${CMAKE_CURRENT_LIST_DIR}/../../../include) +target_link_libraries( + nccl_host_transport + PRIVATE NCCL::nccl CUDA::cudart JsonCpp::JsonCpp glog::glog pthread) +target_compile_features(nccl_host_transport PRIVATE cxx_std_17) diff --git a/mooncake-transfer-engine/src/transport/nccl_transport/nccl_transport.cpp b/mooncake-transfer-engine/src/transport/nccl_transport/nccl_transport.cpp new file mode 100644 index 0000000000..a28a7b0248 --- /dev/null +++ b/mooncake-transfer-engine/src/transport/nccl_transport/nccl_transport.cpp @@ -0,0 +1,1287 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "transport/nccl_transport/nccl_transport.h" + +#include +#include +#if __has_include() +#include +#else +#include +#endif +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common.h" +#include "error.h" +#include "transfer_metadata.h" + +#if NCCL_VERSION_CODE < 23004 +#error "Mooncake NCCL host transport requires NCCL 2.30.4 or newer" +#endif + +namespace mooncake { +namespace { + +constexpr char kHandshakeProtocol[] = "nccl"; +constexpr int kSessionTimeoutSeconds = 60; + +struct BufferInfo { + uint64_t addr = 0; + uint64_t length = 0; + int device_id = -1; +}; + +bool containsRange(const BufferInfo& buffer, uint64_t addr, size_t length) { + if (length == 0 || addr < buffer.addr || length > buffer.length) { + return false; + } + return addr - buffer.addr <= buffer.length - length; +} + +std::string ncclError(ncclResult_t result, const char* operation) { + std::ostringstream out; + out << operation << " failed: " << ncclGetErrorString(result); + return out.str(); +} + +std::string cudaError(cudaError_t result, const char* operation) { + std::ostringstream out; + out << operation << " failed: " << cudaGetErrorString(result); + return out.str(); +} + +std::string encodeJson(const Json::Value& value) { + Json::StreamWriterBuilder builder; + builder["indentation"] = ""; + return Json::writeString(builder, value); +} + +bool decodeJson(const std::string& encoded, Json::Value* value, + std::string* error) { + if (!value) return false; + Json::CharReaderBuilder builder; + builder["allowComments"] = false; + builder["allowTrailingCommas"] = false; + builder["failIfExtra"] = true; + builder["rejectDupKeys"] = true; + builder["strictRoot"] = true; + std::istringstream input(encoded); + try { + if (!Json::parseFromStream(builder, input, value, error)) return false; + } catch (const Json::Exception& exception) { + if (error) *error = exception.what(); + return false; + } + if (!value->isObject()) { + if (error) *error = "NCCL handshake payload is not an object"; + return false; + } + return true; +} + +bool hasStringField(const Json::Value& value, const char* name) { + return value.isObject() && value.isMember(name) && value[name].isString(); +} + +bool hasIntField(const Json::Value& value, const char* name) { + return value.isObject() && value.isMember(name) && value[name].isInt(); +} + +bool hasArrayField(const Json::Value& value, const char* name) { + return value.isObject() && value.isMember(name) && value[name].isArray(); +} + +Json::Value encodeBuffers(const std::vector& buffers) { + Json::Value result(Json::arrayValue); + for (const auto& buffer : buffers) { + Json::Value item; + item["addr"] = static_cast(buffer.addr); + item["length"] = static_cast(buffer.length); + item["device_id"] = buffer.device_id; + result.append(std::move(item)); + } + return result; +} + +bool decodeBuffers(const Json::Value& value, int expected_device, + std::vector* buffers, std::string* error) { + if (!buffers || expected_device < 0 || !value.isArray()) { + if (error) *error = "NCCL buffer catalog is not an array"; + return false; + } + + std::vector decoded; + decoded.reserve(value.size()); + for (const auto& item : value) { + if (!item.isObject() || !item.isMember("addr") || + !item["addr"].isUInt64() || !item.isMember("length") || + !item["length"].isUInt64() || !item.isMember("device_id") || + !item["device_id"].isInt()) { + if (error) *error = "Malformed NCCL buffer catalog entry"; + return false; + } + BufferInfo buffer; + buffer.addr = item["addr"].asUInt64(); + buffer.length = item["length"].asUInt64(); + buffer.device_id = item["device_id"].asInt(); + if (!buffer.addr || !buffer.length || + buffer.length > + std::numeric_limits::max() - buffer.addr || + buffer.device_id != expected_device) { + if (error) *error = "Invalid NCCL buffer catalog entry"; + return false; + } + decoded.push_back(buffer); + } + if (decoded.empty()) { + if (error) *error = "NCCL buffer catalog is empty"; + return false; + } + // Preserve the sender's registration order: NCCL window registration is + // collective, so both ranks must pair corresponding buffers by the same + // index. Use a separate address-ordered copy only to validate overlap. + auto by_address = decoded; + std::sort(by_address.begin(), by_address.end(), + [](const BufferInfo& lhs, const BufferInfo& rhs) { + return lhs.addr < rhs.addr; + }); + for (size_t i = 1; i < by_address.size(); ++i) { + const auto& previous = by_address[i - 1]; + if (previous.addr + previous.length > by_address[i].addr) { + if (error) *error = "NCCL buffer catalog contains overlap"; + return false; + } + } + *buffers = std::move(decoded); + return true; +} + +std::string encodeUniqueId(const ncclUniqueId& id) { + const auto* bytes = reinterpret_cast(&id); + std::ostringstream out; + out << std::hex << std::setfill('0'); + for (size_t i = 0; i < sizeof(id); ++i) { + out << std::setw(2) << static_cast(bytes[i]); + } + return out.str(); +} + +int decodeHexNibble(char value) { + if (value >= '0' && value <= '9') return value - '0'; + if (value >= 'a' && value <= 'f') return value - 'a' + 10; + if (value >= 'A' && value <= 'F') return value - 'A' + 10; + return -1; +} + +bool decodeUniqueId(const std::string& encoded, ncclUniqueId* id) { + if (!id || encoded.size() != 2 * sizeof(*id)) return false; + auto* bytes = reinterpret_cast(id); + for (size_t i = 0; i < sizeof(*id); ++i) { + const int high = decodeHexNibble(encoded[2 * i]); + const int low = decodeHexNibble(encoded[2 * i + 1]); + if (high < 0 || low < 0) return false; + bytes[i] = static_cast((high << 4) | low); + } + return true; +} + +int getPointerDevice(const void* ptr) { + cudaPointerAttributes attributes{}; + cudaError_t result = cudaPointerGetAttributes(&attributes, ptr); + if (result != cudaSuccess) { + cudaGetLastError(); + return -1; + } + if (attributes.type != cudaMemoryTypeDevice && + attributes.type != cudaMemoryTypeManaged) { + return -1; + } + return attributes.device; +} + +bool endpointLess(const std::string& local_name, int local_device, + const std::string& peer_name, int peer_device) { + return std::tie(local_name, local_device) < + std::tie(peer_name, peer_device); +} + +void destroyNcclSliceEvent(Transport::Slice* slice) { + if (!slice || !slice->nccl.event) return; + + cudaEvent_t event = reinterpret_cast(slice->nccl.event); + const int device_id = slice->nccl.device_id; + // Relinquish slice ownership before calling CUDA so all later cleanup + // paths are idempotent even when CUDA reports an error. + slice->nccl.event = nullptr; + slice->nccl.device_id = -1; + + int saved_device = -1; + cudaGetDevice(&saved_device); + if (device_id >= 0) cudaSetDevice(device_id); + cudaEventDestroy(event); + if (saved_device >= 0) cudaSetDevice(saved_device); +} + +std::string makeSessionKey(const std::string& local_name, int local_device, + const std::string& peer_name, int peer_device) { + std::ostringstream out; + if (endpointLess(local_name, local_device, peer_name, peer_device)) { + out << local_name << '#' << local_device << '|' << peer_name << '#' + << peer_device; + } else { + out << peer_name << '#' << peer_device << '|' << local_name << '#' + << local_device; + } + return out.str(); +} + +class NcclSession { + public: + NcclSession(std::string key, std::string peer_name, int local_device, + int peer_device, int rank, ncclUniqueId unique_id, + std::array, 2> rank_buffers) + : key_(std::move(key)), + peer_name_(std::move(peer_name)), + local_device_(local_device), + peer_device_(peer_device), + rank_(rank), + unique_id_(unique_id), + unique_id_string_(encodeUniqueId(unique_id)), + rank_buffers_(std::move(rank_buffers)) {} + + ~NcclSession() { + if (init_thread_.joinable()) init_thread_.join(); + cleanup(); + } + + const std::string& uniqueIdString() const { return unique_id_string_; } + int rank() const { return rank_; } + + void start() { + std::lock_guard lock(state_mutex_); + if (started_) return; + started_ = true; + init_thread_ = std::thread([this] { initialize(); }); + } + + bool waitReady(std::string* error) { + std::unique_lock lock(state_mutex_); + const bool finished = state_cv_.wait_for( + lock, std::chrono::seconds(kSessionTimeoutSeconds), + [this] { return ready_ || failed_; }); + if (!finished) { + if (error) *error = "Timed out initializing NCCL RMA session"; + return false; + } + if (failed_) { + if (error) *error = error_; + return false; + } + return true; + } + + int enqueuePut(const void* source, int owner_rank, uint64_t dest_addr, + size_t length, cudaEvent_t event, std::string* error) { + if (!waitReady(error)) return -1; + const Window* window = findWindow(owner_rank, dest_addr, length); + if (!window) { + if (error) *error = "Destination is outside an NCCL window"; + return -1; + } + + std::lock_guard lock(submit_mutex_); + int saved_device = -1; + cudaGetDevice(&saved_device); + cudaError_t cuda_result = cudaSetDevice(local_device_); + if (cuda_result != cudaSuccess) { + if (error) *error = cudaError(cuda_result, "cudaSetDevice"); + return -1; + } + + const auto& destination = window->buffers[owner_rank]; + ncclResult_t result = ncclPutSignal( + source, length, ncclUint8, owner_rank, window->handle, + dest_addr - destination.addr, 0, 0, 0, comm_, stream_); + cudaError_t sync_result = cudaSuccess; + if (result == ncclSuccess && event) { + cuda_result = cudaEventRecord(event, stream_); + if (cuda_result != cudaSuccess) { + // The put is already enqueued. Do not let the caller reuse or + // free its source until the stream has drained. + sync_result = cudaStreamSynchronize(stream_); + } + } + if (saved_device >= 0) cudaSetDevice(saved_device); + + if (result != ncclSuccess) { + if (error) *error = ncclError(result, "ncclPutSignal"); + return -1; + } + if (cuda_result != cudaSuccess) { + std::string failure = cudaError(cuda_result, "cudaEventRecord"); + if (sync_result != cudaSuccess) { + failure += "; fallback " + + cudaError(sync_result, "cudaStreamSynchronize"); + setFailure(failure); + } + if (error) *error = std::move(failure); + return -1; + } + return 0; + } + + private: + struct Window { + std::array buffers; + void* local_ptr = nullptr; + ncclWindow_t handle = nullptr; + }; + + void setFailure(const std::string& error) { + LOG(ERROR) << "[Host NCCL] session " << key_ << ": " << error; + std::lock_guard lock(state_mutex_); + error_ = error; + failed_ = true; + state_cv_.notify_all(); + } + + void initialize() { + if (rank_buffers_[0].empty() || + rank_buffers_[0].size() != rank_buffers_[1].size()) { + setFailure( + "NCCL peers must register the same number of CUDA buffers"); + return; + } + for (size_t i = 0; i < rank_buffers_[0].size(); ++i) { + if (rank_buffers_[0][i].length != rank_buffers_[1][i].length) { + setFailure( + "Corresponding NCCL buffers must have identical lengths"); + return; + } + } + + int saved_device = -1; + cudaGetDevice(&saved_device); + cudaError_t cuda_result = cudaSetDevice(local_device_); + if (cuda_result != cudaSuccess) { + setFailure(cudaError(cuda_result, "cudaSetDevice")); + return; + } + + ncclConfig_t config = NCCL_CONFIG_INITIALIZER; + config.numRmaCtx = 1; + ncclResult_t result = + ncclCommInitRankConfig(&comm_, 2, unique_id_, rank_, &config); + if (result != ncclSuccess) { + setFailure(ncclError(result, "ncclCommInitRankConfig")); + if (saved_device >= 0) cudaSetDevice(saved_device); + return; + } + + cuda_result = + cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking); + if (cuda_result != cudaSuccess) { + setFailure(cudaError(cuda_result, "cudaStreamCreate")); + if (saved_device >= 0) cudaSetDevice(saved_device); + return; + } + + for (size_t i = 0; i < rank_buffers_[0].size(); ++i) { + Window window; + window.buffers[0] = rank_buffers_[0][i]; + window.buffers[1] = rank_buffers_[1][i]; + const auto& local_buffer = window.buffers[rank_]; + window.local_ptr = reinterpret_cast(local_buffer.addr); + + result = ncclCommWindowRegister(comm_, window.local_ptr, + local_buffer.length, &window.handle, + NCCL_WIN_COLL_SYMMETRIC); + if (result != ncclSuccess || !window.handle) { + if (result == ncclSuccess) result = ncclInternalError; + setFailure(ncclError(result, "ncclCommWindowRegister")); + if (saved_device >= 0) cudaSetDevice(saved_device); + return; + } + windows_.push_back(window); + } + + // NCCL initializes the host-RMA copy-engine resources collectively on + // the first host RMA submission. Warm them here while both session + // ranks are already participating, so later TE puts are one-sided. + result = ncclSignal(1 - rank_, 0, 0, 0, comm_, stream_); + if (result != ncclSuccess) { + setFailure(ncclError(result, "ncclSignal (RMA warm-up)")); + if (saved_device >= 0) cudaSetDevice(saved_device); + return; + } + cuda_result = cudaStreamSynchronize(stream_); + if (cuda_result != cudaSuccess) { + setFailure( + cudaError(cuda_result, "cudaStreamSynchronize (RMA warm-up)")); + if (saved_device >= 0) cudaSetDevice(saved_device); + return; + } + + if (saved_device >= 0) cudaSetDevice(saved_device); + { + std::lock_guard lock(state_mutex_); + ready_ = true; + } + state_cv_.notify_all(); + LOG(INFO) << "[Host NCCL] session ready peer=" << peer_name_ + << " rank=" << rank_ << " local_device=" << local_device_ + << " peer_device=" << peer_device_ + << " windows=" << windows_.size(); + } + + const Window* findWindow(int owner_rank, uint64_t addr, + size_t length) const { + for (const auto& window : windows_) { + const auto& buffer = window.buffers[owner_rank]; + if (containsRange(buffer, addr, length)) { + return &window; + } + } + return nullptr; + } + + void cleanup() { + if (!comm_) return; + int saved_device = -1; + cudaGetDevice(&saved_device); + cudaSetDevice(local_device_); + if (stream_) cudaStreamSynchronize(stream_); + for (auto it = windows_.rbegin(); it != windows_.rend(); ++it) { + if (it->handle) ncclCommWindowDeregister(comm_, it->handle); + } + windows_.clear(); + if (stream_) { + cudaStreamDestroy(stream_); + stream_ = nullptr; + } + ncclCommDestroy(comm_); + comm_ = nullptr; + if (saved_device >= 0) cudaSetDevice(saved_device); + } + + std::string key_; + std::string peer_name_; + int local_device_; + int peer_device_; + int rank_; + ncclUniqueId unique_id_{}; + std::string unique_id_string_; + std::array, 2> rank_buffers_; + + std::thread init_thread_; + std::mutex state_mutex_; + std::condition_variable state_cv_; + bool started_ = false; + bool ready_ = false; + bool failed_ = false; + std::string error_; + + ncclComm_t comm_ = nullptr; + cudaStream_t stream_ = nullptr; + std::vector windows_; + std::mutex submit_mutex_; +}; + +} // namespace + +class NcclHostTransport::Impl { + public: + explicit Impl(NcclHostTransport* owner) : owner_(owner) {} + + ~Impl() { + std::unordered_map> sessions; + { + std::lock_guard lock(sessions_mutex_); + sessions.swap(sessions_); + bootstrap_ids_.clear(); + } + sessions.clear(); + } + + int install(const std::string& local_server_name, + std::shared_ptr metadata) { + local_server_name_ = local_server_name; + metadata_ = std::move(metadata); + + auto segment = std::make_shared(); + segment->name = local_server_name_; + segment->protocol = kHandshakeProtocol; + int result = metadata_->addLocalSegment( + LOCAL_SEGMENT_ID, local_server_name_, std::move(segment)); + if (result != 0) return result; + + result = metadata_->startHandshakeDaemon( + [this](const HandShakeDesc& peer, HandShakeDesc& local) { + return onHandshake(peer, local); + }, + metadata_->localRpcMeta().rpc_port, + metadata_->localRpcMeta().sockfd); + if (result != 0) return result; + return metadata_->updateLocalSegmentDesc(); + } + + int registerMemory(void* addr, size_t length, const std::string& location, + bool remote_accessible, bool update_metadata) { + std::lock_guard lock(buffers_mutex_); + return registerMemoryLocked(addr, length, location, remote_accessible, + update_metadata); + } + + int unregisterMemory(void* addr, bool update_metadata) { + std::lock_guard lock(buffers_mutex_); + return unregisterMemoryLocked(addr, update_metadata); + } + + int registerMemoryBatch(const std::vector& buffer_list, + const std::string& location) { + std::lock_guard lock(buffers_mutex_); + std::vector registered; + registered.reserve(buffer_list.size()); + for (const auto& buffer : buffer_list) { + int result = registerMemoryLocked(buffer.addr, buffer.length, + location, true, false); + if (result != 0) { + for (auto it = registered.rbegin(); it != registered.rend(); + ++it) { + unregisterMemoryLocked(*it, false); + } + return result; + } + registered.push_back(buffer.addr); + } + int result = metadata_->updateLocalSegmentDesc(); + if (result != 0) { + for (auto it = registered.rbegin(); it != registered.rend(); ++it) { + int rollback_result = unregisterMemoryLocked(*it, false); + if (rollback_result != 0) { + LOG(ERROR) << "[Host NCCL] failed to roll back buffer " + << *it << " after metadata publication failure: " + << rollback_result; + } + } + } + return result; + } + + int unregisterMemoryBatch(const std::vector& addr_list) { + std::lock_guard lock(buffers_mutex_); + if (buffers_frozen_) { + LOG(ERROR) << "[Host NCCL] cannot unregister memory after the " + "session catalog has been frozen"; + return ERR_INVALID_ARGUMENT; + } + std::vector metadata_buffers; + metadata_buffers.reserve(addr_list.size()); + for (size_t index = 0; index < addr_list.size(); ++index) { + void* addr = addr_list[index]; + if (std::find(addr_list.begin(), addr_list.begin() + index, addr) != + addr_list.begin() + index) { + return ERR_INVALID_ARGUMENT; + } + auto it = std::find_if(local_buffers_.begin(), local_buffers_.end(), + [addr](const BufferInfo& buffer) { + return buffer.addr == + reinterpret_cast(addr); + }); + if (it == local_buffers_.end()) return ERR_ADDRESS_NOT_REGISTERED; + BufferDesc metadata_buffer; + if (!findMetadataBuffer(addr, &metadata_buffer)) { + return ERR_ADDRESS_NOT_REGISTERED; + } + metadata_buffers.push_back(std::move(metadata_buffer)); + } + size_t removed = 0; + for (; removed < addr_list.size(); ++removed) { + void* addr = addr_list[removed]; + int result = metadata_->removeLocalMemoryBuffer(addr, false); + if (result != 0) { + restoreMetadataBuffers(metadata_buffers, removed); + return result; + } + } + int result = metadata_->updateLocalSegmentDesc(); + if (result != 0) { + restoreMetadataBuffers(metadata_buffers, metadata_buffers.size()); + return result; + } + for (void* addr : addr_list) { + local_buffers_.erase( + std::remove_if(local_buffers_.begin(), local_buffers_.end(), + [addr](const BufferInfo& buffer) { + return buffer.addr == + reinterpret_cast(addr); + }), + local_buffers_.end()); + } + return 0; + } + + Status submitTasks(const std::vector& task_list) { + Status overall = Status::OK(); + for (TransferTask* task : task_list) { + if (!task || !task->request) { + overall = + Status::InvalidArgument("Missing NCCL transfer request"); + continue; + } + const TransferRequest& request = *task->request; + task->total_bytes = request.length; + + Slice* slice = owner_->getSliceCache().allocate(); + slice->source_addr = request.source; + slice->length = request.length; + slice->opcode = request.opcode; + slice->target_id = request.target_id; + slice->task = task; + slice->status = Slice::PENDING; + slice->ts = getCurrentTimeInNano(); + slice->nccl.event = nullptr; + slice->nccl.device_id = -1; + slice->cleanup_callback = destroyNcclSliceEvent; + task->slice_list.push_back(slice); + __sync_fetch_and_add(&task->slice_count, 1); + + std::string error; + if (request.opcode != TransferRequest::WRITE) { + error = + "NCCL host transport supports WRITE only; NCCL " + "exposes no public host-side Get operation"; + slice->markFailed(); + if (overall.ok()) { + overall = Status::NotSupportedTransport(error); + } + LOG(ERROR) << "[Host NCCL] submit failed: " << error; + continue; + } + + auto target = metadata_->getSegmentDescByID(request.target_id); + if (!target) { + error = "Target segment metadata is unavailable"; + } else if (request.target_id == LOCAL_SEGMENT_ID) { + error = "NCCL host transport requires a remote target"; + } + + BufferInfo remote_buffer; + if (error.empty() && + !findRemoteBuffer(*target, request.target_offset, + request.length, &remote_buffer)) { + error = "Target address is not in a registered NCCL buffer"; + } + + int local_device = -1; + if (error.empty()) { + local_device = getPointerDevice(request.source); + if (local_device < 0) { + error = "NCCL local buffer must be CUDA device memory"; + } else if (!freezeAndValidateLocalBuffer( + reinterpret_cast(request.source), + request.length, local_device)) { + error = "NCCL local buffer is not registered"; + } + } + + std::shared_ptr session; + if (error.empty()) { + session = getOrCreateSession(target->name, local_device, + remote_buffer.device_id, &error); + } + + if (error.empty()) { + int saved_device = -1; + cudaGetDevice(&saved_device); + cudaSetDevice(local_device); + cudaEvent_t event = nullptr; + cudaError_t cuda_result = + cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + if (saved_device >= 0) cudaSetDevice(saved_device); + if (cuda_result != cudaSuccess) { + error = cudaError(cuda_result, "cudaEventCreate"); + } else { + slice->nccl.event = event; + slice->nccl.device_id = local_device; + const int peer_rank = session->rank() == 0 ? 1 : 0; + if (session->enqueuePut( + request.source, peer_rank, request.target_offset, + request.length, event, &error) == 0) { + slice->status = Slice::POSTED; + } + } + } + + if (!error.empty()) { + destroyNcclSliceEvent(slice); + LOG(ERROR) << "[Host NCCL] submit failed: " << error; + slice->markFailed(); + if (overall.ok()) overall = Status::Context(error); + } + } + return overall; + } + + Status poll(BatchID batch_id, size_t task_id, TransferStatus& status) { + auto& batch = Transport::toBatchDesc(batch_id); + if (task_id >= batch.task_list.size()) { + return Status::InvalidArgument("NCCL task ID out of range"); + } + auto& task = batch.task_list[task_id]; + for (Slice* slice : task.slice_list) { + if (!slice || slice->status != Slice::POSTED) continue; + int saved_device = -1; + cudaGetDevice(&saved_device); + cudaSetDevice(slice->nccl.device_id); + cudaEvent_t event = + reinterpret_cast(slice->nccl.event); + cudaError_t result = cudaEventQuery(event); + if (result == cudaSuccess) { + destroyNcclSliceEvent(slice); + slice->markSuccess(); + } else if (result != cudaErrorNotReady) { + cudaGetLastError(); + destroyNcclSliceEvent(slice); + slice->markFailed(); + } + if (saved_device >= 0) cudaSetDevice(saved_device); + } + + status.transferred_bytes = task.transferred_bytes; + if (task.success_slice_count + task.failed_slice_count == + task.slice_count) { + status.s = task.failed_slice_count ? TransferStatusEnum::FAILED + : TransferStatusEnum::COMPLETED; + task.is_finished = true; + } else { + status.s = TransferStatusEnum::WAITING; + } + return Status::OK(); + } + + private: + bool findMetadataBuffer(void* addr, BufferDesc* result) const { + if (!result) return false; + auto segment = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); + if (!segment) return false; + auto it = std::find_if(segment->buffers.begin(), segment->buffers.end(), + [addr](const BufferDesc& buffer) { + return buffer.addr == + reinterpret_cast(addr); + }); + if (it == segment->buffers.end()) return false; + *result = *it; + return true; + } + + void restoreMetadataBuffers(const std::vector& buffers, + size_t count) { + for (size_t index = 0; index < count; ++index) { + int result = metadata_->addLocalMemoryBuffer(buffers[index], false); + if (result != 0) { + LOG(ERROR) << "[Host NCCL] failed to restore metadata buffer " + << reinterpret_cast(buffers[index].addr) + << ": " << result; + } + } + } + + // buffers_mutex_ must be held across each complete catalog mutation so the + // first session cannot snapshot a partially committed registration. + int registerMemoryLocked(void* addr, size_t length, + const std::string& location, + bool remote_accessible, bool update_metadata) { + (void)remote_accessible; + if (!addr || length == 0) return ERR_INVALID_ARGUMENT; + int device_id = getPointerDevice(addr); + if (device_id < 0) { + LOG(ERROR) << "[Host NCCL] only CUDA device or managed memory " + "can be registered"; + return ERR_INVALID_ARGUMENT; + } + + const uint64_t address = reinterpret_cast(addr); + if (length > std::numeric_limits::max() - address) { + LOG(ERROR) << "[Host NCCL] memory registration range overflows"; + return ERR_INVALID_ARGUMENT; + } + + BufferInfo info{address, length, device_id}; + if (buffers_frozen_) { + LOG(ERROR) << "[Host NCCL] register all buffers before the first " + "transfer freezes the session catalog"; + return ERR_INVALID_ARGUMENT; + } + for (const auto& buffer : local_buffers_) { + const uint64_t lhs_end = info.addr + info.length; + const uint64_t rhs_end = buffer.addr + buffer.length; + if (info.addr < rhs_end && buffer.addr < lhs_end) { + LOG(ERROR) << "[Host NCCL] overlapping memory registration"; + return ERR_ADDRESS_OVERLAPPED; + } + } + local_buffers_.push_back(info); + + BufferDesc desc; + desc.name = location; + desc.addr = info.addr; + desc.length = info.length; + desc.device_id = info.device_id; + int result = metadata_->addLocalMemoryBuffer(desc, update_metadata); + if (result != 0) { + int rollback_result = + metadata_->removeLocalMemoryBuffer(addr, false); + if (rollback_result != 0 && + rollback_result != ERR_ADDRESS_NOT_REGISTERED) { + LOG(ERROR) << "[Host NCCL] failed to roll back metadata for " + << addr << ": " << rollback_result; + } + local_buffers_.erase( + std::remove_if(local_buffers_.begin(), local_buffers_.end(), + [addr](const BufferInfo& buffer) { + return buffer.addr == + reinterpret_cast(addr); + }), + local_buffers_.end()); + } + return result; + } + + int unregisterMemoryLocked(void* addr, bool update_metadata) { + if (buffers_frozen_) { + LOG(ERROR) << "[Host NCCL] cannot unregister memory after the " + "session catalog has been frozen"; + return ERR_INVALID_ARGUMENT; + } + + auto it = std::find_if(local_buffers_.begin(), local_buffers_.end(), + [addr](const BufferInfo& buffer) { + return buffer.addr == + reinterpret_cast(addr); + }); + if (it == local_buffers_.end()) return ERR_ADDRESS_NOT_REGISTERED; + + BufferDesc metadata_buffer; + if (!findMetadataBuffer(addr, &metadata_buffer)) { + return ERR_ADDRESS_NOT_REGISTERED; + } + int result = metadata_->removeLocalMemoryBuffer(addr, update_metadata); + if (result != 0) { + if (update_metadata) { + restoreMetadataBuffers({metadata_buffer}, 1); + } + return result; + } + local_buffers_.erase(it); + return 0; + } + + bool freezeBuffersForDevice(int device_id, + std::vector* result) { + if (!result) return false; + std::lock_guard lock(buffers_mutex_); + std::vector snapshot; + for (const auto& buffer : local_buffers_) { + if (buffer.device_id == device_id) snapshot.push_back(buffer); + } + if (snapshot.empty()) return false; + // Registration and unregistration take the same lock and check this + // flag. The first session therefore snapshots a complete, immutable + // catalog instead of racing a buffer mutation between two locks. + buffers_frozen_ = true; + // Keep registration order so collective window index i represents the + // same logical buffer on both ranks even when their virtual addresses + // differ. + *result = std::move(snapshot); + return true; + } + + bool freezeAndValidateLocalBuffer(uint64_t addr, size_t length, + int device_id) { + std::lock_guard lock(buffers_mutex_); + for (const auto& buffer : local_buffers_) { + if (buffer.device_id == device_id && + containsRange(buffer, addr, length)) { + // Once validation succeeds, unregister must not be able to + // remove the source before the session snapshots its windows. + buffers_frozen_ = true; + return true; + } + } + return false; + } + + bool findRemoteBuffer(const SegmentDesc& segment, uint64_t addr, + size_t length, BufferInfo* result) const { + for (const auto& buffer : segment.buffers) { + BufferInfo info{buffer.addr, buffer.length, buffer.device_id}; + if (info.device_id >= 0 && containsRange(info, addr, length)) { + if (result) *result = info; + return true; + } + } + return false; + } + + bool selectBootstrapId(const std::string& key, + const std::string& proposed_id, bool may_create, + ncclUniqueId* unique_id, std::string* encoded_id, + std::string* error) { + std::lock_guard lock(sessions_mutex_); + auto it = bootstrap_ids_.find(key); + if (it != bootstrap_ids_.end()) { + if (!proposed_id.empty() && proposed_id != it->second) { + if (error) *error = "Concurrent NCCL bootstrap ID mismatch"; + return false; + } + *encoded_id = it->second; + if (!decodeUniqueId(*encoded_id, unique_id)) { + if (error) *error = "Stored NCCL bootstrap ID is invalid"; + return false; + } + return true; + } + + if (!proposed_id.empty()) { + *encoded_id = proposed_id; + if (!decodeUniqueId(*encoded_id, unique_id)) { + if (error) *error = "Invalid NCCL unique ID in bootstrap"; + return false; + } + } else { + if (!may_create) { + if (error) *error = "NCCL rank 1 cannot create a unique ID"; + return false; + } + ncclResult_t result = ncclGetUniqueId(unique_id); + if (result != ncclSuccess) { + if (error) *error = ncclError(result, "ncclGetUniqueId"); + return false; + } + *encoded_id = encodeUniqueId(*unique_id); + } + bootstrap_ids_.emplace(key, *encoded_id); + return true; + } + + std::shared_ptr getOrCreateSession( + const std::string& peer_name, int local_device, int peer_device, + std::string* error) { + const std::string key = makeSessionKey(local_server_name_, local_device, + peer_name, peer_device); + std::shared_ptr existing; + { + std::lock_guard lock(sessions_mutex_); + auto it = sessions_.find(key); + if (it != sessions_.end()) { + existing = it->second; + } + } + if (existing) { + // NCCL communicator/window initialization is collective. Retain a + // terminally failed session rather than letting one endpoint retry + // a new collective generation without coordinating its peer. + if (!existing->waitReady(error)) return nullptr; + return existing; + } + + std::vector local_buffers; + if (!freezeBuffersForDevice(local_device, &local_buffers)) { + if (error) *error = "No NCCL buffers registered on local device"; + return nullptr; + } + const int local_rank = endpointLess(local_server_name_, local_device, + peer_name, peer_device) + ? 0 + : 1; + ncclUniqueId unique_id{}; + std::string unique_id_string; + if (local_rank == 0 && !selectBootstrapId(key, "", true, &unique_id, + &unique_id_string, error)) { + return nullptr; + } + + Json::Value request; + request["op"] = "bootstrap"; + request["peer_name"] = local_server_name_; + request["local_device"] = local_device; + request["remote_device"] = peer_device; + request["unique_id"] = unique_id_string; + request["buffers"] = encodeBuffers(local_buffers); + + HandShakeDesc local_desc; + local_desc.payload = encodeJson(request); + HandShakeDesc peer_desc; + int result = metadata_->sendHandshake(peer_name, local_desc, peer_desc); + if (result != 0) { + if (error) *error = "NCCL bootstrap handshake failed"; + return nullptr; + } + + Json::Value response; + std::string parse_error; + if (!decodeJson(peer_desc.payload, &response, &parse_error)) { + if (error) + *error = "Invalid NCCL bootstrap response: " + parse_error; + return nullptr; + } + if (!hasStringField(response, "unique_id") || + !hasArrayField(response, "buffers")) { + if (error) *error = "Malformed NCCL bootstrap response"; + return nullptr; + } + const std::string response_id = response["unique_id"].asString(); + if (!selectBootstrapId(key, response_id, local_rank == 0, &unique_id, + &unique_id_string, error)) { + return nullptr; + } + + std::vector peer_buffers; + if (!decodeBuffers(response["buffers"], peer_device, &peer_buffers, + error)) { + return nullptr; + } + std::array, 2> rank_buffers; + rank_buffers[local_rank] = local_buffers; + rank_buffers[1 - local_rank] = peer_buffers; + + std::shared_ptr session; + { + std::lock_guard lock(sessions_mutex_); + auto it = sessions_.find(key); + if (it != sessions_.end()) { + session = it->second; + if (session->uniqueIdString() != unique_id_string) { + if (error) *error = "Concurrent NCCL bootstrap conflict"; + return nullptr; + } + } else { + session = std::make_shared( + key, peer_name, local_device, peer_device, local_rank, + unique_id, std::move(rank_buffers)); + sessions_.emplace(key, session); + session->start(); + } + } + if (!session->waitReady(error)) return nullptr; + return session; + } + + int onHandshake(const HandShakeDesc& peer_desc, HandShakeDesc& local_desc) { + Json::Value request; + std::string error; + if (!decodeJson(peer_desc.payload, &request, &error)) { + local_desc.reply_msg = "Invalid NCCL handshake payload: " + error; + return 0; + } + if (!hasStringField(request, "op")) { + local_desc.reply_msg = + "Missing or invalid 'op' field in NCCL handshake request"; + return 0; + } + + const std::string op = request["op"].asString(); + Json::Value response; + try { + if (op == "bootstrap") { + if (handleBootstrap(request, &response, &error) != 0) { + local_desc.reply_msg = error; + } + } else { + local_desc.reply_msg = "Unknown NCCL handshake operation"; + } + } catch (const Json::Exception& exception) { + local_desc.reply_msg = "Malformed NCCL handshake request: " + + std::string(exception.what()); + } + local_desc.payload = encodeJson(response); + return 0; + } + + int handleBootstrap(const Json::Value& request, Json::Value* response, + std::string* error) { + if (!response || !hasStringField(request, "peer_name") || + !hasIntField(request, "local_device") || + !hasIntField(request, "remote_device") || + !hasStringField(request, "unique_id") || + !hasArrayField(request, "buffers")) { + if (error) *error = "Malformed NCCL bootstrap request"; + return -1; + } + const std::string peer_name = request["peer_name"].asString(); + const int peer_device = request["local_device"].asInt(); + const int local_device = request["remote_device"].asInt(); + if (peer_name.empty() || local_device < 0 || peer_device < 0) { + if (error) *error = "Invalid NCCL bootstrap endpoint"; + return -1; + } + + std::vector peer_buffers; + if (!decodeBuffers(request["buffers"], peer_device, &peer_buffers, + error)) { + return -1; + } + const int local_rank = endpointLess(local_server_name_, local_device, + peer_name, peer_device) + ? 0 + : 1; + + const std::string key = makeSessionKey(local_server_name_, local_device, + peer_name, peer_device); + ncclUniqueId unique_id{}; + std::string unique_id_string; + const std::string proposed_id = request["unique_id"].asString(); + ncclUniqueId proposed_unique_id{}; + if ((!proposed_id.empty() && + !decodeUniqueId(proposed_id, &proposed_unique_id)) || + (proposed_id.empty() && local_rank != 0)) { + if (error) *error = "Invalid NCCL unique ID in bootstrap"; + return -1; + } + + // Do not let malformed endpoint data permanently freeze registration. + // Once the request is valid, take the immutable catalog snapshot used + // to create the collective windows. + std::vector local_buffers; + if (!freezeBuffersForDevice(local_device, &local_buffers)) { + if (error) *error = "No NCCL buffers registered on local device"; + return -1; + } + + if (!selectBootstrapId(key, proposed_id, local_rank == 0, &unique_id, + &unique_id_string, error)) { + return -1; + } + + std::array, 2> rank_buffers; + rank_buffers[local_rank] = local_buffers; + rank_buffers[1 - local_rank] = peer_buffers; + + { + std::lock_guard lock(sessions_mutex_); + auto it = sessions_.find(key); + if (it != sessions_.end()) { + if (it->second->uniqueIdString() != unique_id_string) { + if (error) *error = "Concurrent NCCL bootstrap conflict"; + return -1; + } + } else { + auto session = std::make_shared( + key, peer_name, local_device, peer_device, local_rank, + unique_id, std::move(rank_buffers)); + sessions_.emplace(key, session); + session->start(); + } + } + + (*response)["unique_id"] = unique_id_string; + (*response)["buffers"] = encodeBuffers(local_buffers); + return 0; + } + + NcclHostTransport* owner_; + std::string local_server_name_; + std::shared_ptr metadata_; + + mutable std::mutex buffers_mutex_; + std::vector local_buffers_; + bool buffers_frozen_ = false; + std::mutex sessions_mutex_; + std::unordered_map> sessions_; + std::unordered_map bootstrap_ids_; +}; + +NcclHostTransport::NcclHostTransport() : impl_(std::make_unique(this)) {} + +NcclHostTransport::~NcclHostTransport() = default; + +int NcclHostTransport::install(std::string& local_server_name, + std::shared_ptr metadata, + std::shared_ptr topology) { + (void)topology; + local_server_name_ = local_server_name; + metadata_ = metadata; + return impl_->install(local_server_name, std::move(metadata)); +} + +Status NcclHostTransport::submitTransfer( + BatchID batch_id, const std::vector& entries) { + auto& batch = toBatchDesc(batch_id); + if (batch.task_list.size() + entries.size() > batch.batch_size) { + return Status::TooManyRequests("NCCL batch capacity exceeded"); + } + const size_t first = batch.task_list.size(); + batch.task_list.resize(first + entries.size()); + std::vector tasks; + tasks.reserve(entries.size()); + for (size_t i = 0; i < entries.size(); ++i) { + auto& task = batch.task_list[first + i]; + task.batch_id = batch_id; + task.transport_ = this; + task.request = &entries[i]; + tasks.push_back(&task); + } + return impl_->submitTasks(tasks); +} + +Status NcclHostTransport::submitTransferTask( + const std::vector& task_list) { + return impl_->submitTasks(task_list); +} + +Status NcclHostTransport::getTransferStatus(BatchID batch_id, size_t task_id, + TransferStatus& status) { + return impl_->poll(batch_id, task_id, status); +} + +int NcclHostTransport::registerLocalMemory(void* addr, size_t length, + const std::string& location, + bool remote_accessible, + bool update_metadata) { + return impl_->registerMemory(addr, length, location, remote_accessible, + update_metadata); +} + +int NcclHostTransport::unregisterLocalMemory(void* addr, bool update_metadata) { + return impl_->unregisterMemory(addr, update_metadata); +} + +int NcclHostTransport::registerLocalMemoryBatch( + const std::vector& buffer_list, const std::string& location) { + return impl_->registerMemoryBatch(buffer_list, location); +} + +int NcclHostTransport::unregisterLocalMemoryBatch( + const std::vector& addr_list) { + return impl_->unregisterMemoryBatch(addr_list); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/nvlink_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/nvlink_transport/CMakeLists.txt index 8846169942..47067f5dd7 100644 --- a/mooncake-transfer-engine/src/transport/nvlink_transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/nvlink_transport/CMakeLists.txt @@ -3,5 +3,5 @@ file(GLOB NVLINK_SOURCES "*.cpp") add_library(nvlink_transport OBJECT ${NVLINK_SOURCES}) if (USE_CUDA) -target_include_directories(nvlink_transport PUBLIC CUDA::cudart "/usr/local/cuda/include") +target_include_directories(nvlink_transport PUBLIC ${CUDAToolkit_INCLUDE_DIRS}) endif() diff --git a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp index 7cf8065274..bac4d625cd 100644 --- a/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp +++ b/mooncake-transfer-engine/src/transport/nvlink_transport/nvlink_transport.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include "common.h" #include "common/serialization.h" @@ -42,6 +43,265 @@ static bool checkCudaErrorReturn(cudaError_t result, const char *message) { } namespace mooncake { + +namespace { + +/// Per-device CUDA stream pool (thread-local). +/// Each thread maintains a map of device_id → stream. +/// The source buffer's device is queried via cudaPointerGetAttributes, +/// and the stream for THAT device is used for DMA copy submission. +struct CudaStreamEntry { + cudaStream_t stream; + int device_id; +}; + +class PerDeviceStreamPool { + public: + CudaStreamEntry getOrCreate(int device_id) { + auto it = pool_.find(device_id); + if (it != pool_.end()) return it->second; + + int saved_device = 0; + cudaGetDevice(&saved_device); + if (cudaSetDevice(device_id) != cudaSuccess) { + LOG(ERROR) << "NvlinkTransport: cudaSetDevice(" << device_id + << ") failed when creating stream"; + return {nullptr, -1}; + } + CudaStreamEntry entry; + entry.device_id = device_id; + cudaError_t err = + cudaStreamCreateWithFlags(&entry.stream, cudaStreamNonBlocking); + if (err != cudaSuccess) + LOG(FATAL) << "Failed to create NVLink CUDA stream on device " + << device_id << ": " << cudaGetErrorString(err); + cudaSetDevice(saved_device); + + cudaDeviceProp prop; + std::string pci = "unknown"; + if (cudaGetDeviceProperties(&prop, device_id) == cudaSuccess) { + pci = std::string(prop.name) + " PCI " + + std::to_string(prop.pciBusID) + ":" + + std::to_string(prop.pciDeviceID); + } + const char *visible = getenv("CUDA_VISIBLE_DEVICES"); + LOG(INFO) << "NvlinkTransport: NVLink CUDA stream created on device " + << device_id << " [physical: " << pci + << "] CUDA_VISIBLE_DEVICES=" + << (visible ? visible : "(not set)") << " pid=" << getpid(); + pool_[device_id] = entry; + return entry; + } + ~PerDeviceStreamPool() { + int saved_device = 0; + cudaGetDevice(&saved_device); + for (auto &kv : pool_) { + cudaSetDevice(kv.first); + if (kv.second.stream) cudaStreamDestroy(kv.second.stream); + } + cudaSetDevice(saved_device); + } + + private: + std::unordered_map pool_; +}; + +static thread_local PerDeviceStreamPool tl_device_stream_pool; + +/// Per-device event pool (thread-local). Caches one event per device to +/// avoid repeated create/destroy on device switches, and ensures proper +/// cleanup when the thread exits. +class PerDeviceEventPool { + public: + cudaEvent_t getOrCreate(int device_id) { + auto it = pool_.find(device_id); + if (it != pool_.end()) return it->second; + int saved_device = 0; + cudaGetDevice(&saved_device); + if (cudaSetDevice(device_id) != cudaSuccess) { + LOG(ERROR) << "NvlinkTransport: cudaSetDevice(" << device_id + << ") failed when creating event"; + return nullptr; + } + cudaEvent_t event = nullptr; + cudaError_t err = + cudaEventCreateWithFlags(&event, cudaEventDisableTiming); + if (err != cudaSuccess) + LOG(FATAL) << "Failed to create NVLink sync event on device " + << device_id << ": " << cudaGetErrorString(err); + cudaSetDevice(saved_device); + pool_[device_id] = event; + return event; + } + ~PerDeviceEventPool() { + int saved_device = 0; + cudaGetDevice(&saved_device); + for (auto &kv : pool_) { + cudaSetDevice(kv.first); + if (kv.second) cudaEventDestroy(kv.second); + } + cudaSetDevice(saved_device); + } + + private: + std::unordered_map pool_; +}; + +static thread_local PerDeviceEventPool tl_device_event_pool; + +static cudaEvent_t getCallerSyncEvent() { + int current_device = 0; + cudaGetDevice(¤t_device); + return tl_device_event_pool.getOrCreate(current_device); +} + +static int getDeviceForPointer(const void *ptr) { + cudaPointerAttributes attr; + if (cudaPointerGetAttributes(&attr, ptr) != cudaSuccess) { + cudaGetLastError(); + return -1; + } + return (attr.type == cudaMemoryTypeDevice) ? attr.device : -1; +} + +static CudaStreamEntry getStreamForRequest(const void *source) { + int device_id = getDeviceForPointer(source); + if (device_id < 0) { + cudaGetDevice(&device_id); + if (device_id < 0) device_id = 0; + } + return tl_device_stream_pool.getOrCreate(device_id); +} + +} // anonymous namespace + +using Slice = Transport::Slice; + +/// Submit batched memcpy operations using cudaMemcpyBatchAsync when available +/// (CUDA 12.8+), falling back to per-slice cudaMemcpyAsync otherwise. +/// Uses cudaMemcpySrcAccessOrderStream to ensure source data visibility for +/// P2P copies. This attribute is REQUIRED — without it, the GPU does not +/// insert the necessary memory barriers for P2P access, causing segfaults. +/// The caller must also establish GPU-level stream synchronization +/// (via cudaEventRecord + cudaStreamWaitEvent) before calling this function. +/// Individual slice errors are tracked so that slices whose memcpy failed +/// are marked as FAILED while successfully submitted ones are POSTED. +static void submitBatchMemcpy(const std::vector &slices, + const std::vector &srcs, + const std::vector &dsts, + const std::vector &sizes, + cudaStream_t stream) { + if (slices.empty()) return; + + const size_t count = slices.size(); + cudaError_t err = cudaSuccess; + + // Log the active memcpy path once per process lifetime + static const bool logged_once = [] { +#if CUDART_VERSION >= 13000 + LOG(INFO) << "NvlinkTransport: using cudaMemcpyBatchAsync " + << "(CUDA >= 13.0 path)"; +#elif CUDART_VERSION >= 12080 + LOG(INFO) << "NvlinkTransport: using cudaMemcpyBatchAsync " + << "(CUDA >= 12.8 path)"; +#else + LOG(INFO) << "NvlinkTransport: using per-slice cudaMemcpyAsync " + << "(CUDA < 12.8 fallback path)"; +#endif + return true; + }(); + (void)logged_once; + +#if CUDART_VERSION >= 12080 + // srcAccessOrderStream is REQUIRED for P2P copies — without it, the GPU + // does not insert necessary memory barriers for cross-device access, + // resulting in segmentation faults. The caller also establishes a + // GPU-level dependency via cudaEventRecord(cudaStreamPerThread) + + // cudaStreamWaitEvent(nvlink_stream) to ensure source data is coherent + // before the memcpy starts; these two mechanisms are complementary. + cudaMemcpyAttributes attr{}; + attr.srcAccessOrder = cudaMemcpySrcAccessOrderStream; + size_t attrs_idx = 0; + // cudaMemcpyBatchAsync in CUDA 12.8 takes non-const size_t* for sizes + std::vector mutable_sizes(sizes); + size_t fail_idx = count; +#endif + +#if CUDART_VERSION >= 13000 + err = cudaMemcpyBatchAsync(const_cast(dsts.data()), + const_cast(srcs.data()), + mutable_sizes.data(), static_cast(count), + &attr, &attrs_idx, 1, stream); + if (err != cudaSuccess) { + LOG(ERROR) << "NvlinkTransport: cudaMemcpyBatchAsync " + << "failed: " << cudaGetErrorString(err); + // CUDA >= 13.0 does not return fail_idx; conservatively mark all + // as FAILED since we cannot determine which copies succeeded. + for (size_t i = 0; i < count; ++i) { + if (slices[i]->status == Slice::PENDING) { + slices[i]->markFailed(); + } + } + } else { + for (size_t i = 0; i < count; ++i) { + slices[i]->status = Slice::POSTED; + slices[i]->local.cuda_stream = (void *)stream; + } + } +#elif CUDART_VERSION >= 12080 + err = cudaMemcpyBatchAsync(const_cast(dsts.data()), + const_cast(srcs.data()), + mutable_sizes.data(), static_cast(count), + &attr, &attrs_idx, 1, &fail_idx, stream); + if (err != cudaSuccess) { + if (fail_idx < count) { + LOG(ERROR) << "NvlinkTransport: cudaMemcpyBatchAsync " + << "failed at index " << fail_idx + << " (src=" << srcs[fail_idx] + << ", dst=" << dsts[fail_idx] + << ", size=" << sizes[fail_idx] + << "): " << cudaGetErrorString(err); + } else { + LOG(ERROR) << "NvlinkTransport: cudaMemcpyBatchAsync " + << "failed: " << cudaGetErrorString(err); + } + // Copies [0, fail_idx) were submitted successfully → POSTED. + // Copy [fail_idx] failed → FAILED. + // Copies (fail_idx, count) were never submitted → FAILED. + for (size_t i = 0; i < fail_idx; ++i) { + slices[i]->status = Slice::POSTED; + slices[i]->local.cuda_stream = (void *)stream; + } + for (size_t i = fail_idx; i < count; ++i) { + if (slices[i]->status == Slice::PENDING) { + slices[i]->markFailed(); + } + } + } else { + for (size_t i = 0; i < count; ++i) { + slices[i]->status = Slice::POSTED; + slices[i]->local.cuda_stream = (void *)stream; + } + } +#else + // Fallback for CUDA < 12.8: submit each memcpy individually + for (size_t i = 0; i < count; ++i) { + auto single_err = cudaMemcpyAsync(dsts[i], srcs[i], sizes[i], + cudaMemcpyDefault, stream); + if (single_err != cudaSuccess) { + LOG(ERROR) << "NvlinkTransport: cudaMemcpyAsync failed at " + << "index " << i << ": " + << cudaGetErrorString(single_err); + slices[i]->markFailed(); + continue; + } + slices[i]->status = Slice::POSTED; + slices[i]->local.cuda_stream = (void *)stream; + } + return; // Slice states already set above +#endif +} + static int getNumDevices() { static int cached_num_devices = -1; if (cached_num_devices == -1) { @@ -203,6 +463,38 @@ Status NvlinkTransport::submitTransfer( size_t task_id = batch_desc.task_list.size(); batch_desc.task_list.resize(task_id + entries.size()); + // Get per-device transfer stream for the source buffer's device. + CudaStreamEntry stream_entry = + getStreamForRequest(entries.empty() ? nullptr : entries[0].source); + cudaStream_t stream = stream_entry.stream; + if (!stream) return Status::Context("Failed to create NVLink CUDA stream"); + + // Synchronize with the caller's GPU work (e.g., PyTorch gather operations) + // that produced the source data. We use cudaEventSynchronize (CPU-blocking) + // instead of cudaStreamWaitEvent to avoid expensive cross-device event + // operations on non-NVIDIA GPUs. The CPU blocking is acceptable because + // the transfer depends on the caller's work anyway — they cannot overlap. + cudaEvent_t sync_event = getCallerSyncEvent(); + cudaError_t sync_err = cudaEventRecord(sync_event, cudaStreamPerThread); + if (sync_err != cudaSuccess) { + LOG(ERROR) << "NvlinkTransport: cudaEventRecord failed: " + << cudaGetErrorString(sync_err); + return Status::Context("cudaEventRecord failed: " + + std::string(cudaGetErrorString(sync_err))); + } + sync_err = cudaEventSynchronize(sync_event); + if (sync_err != cudaSuccess) { + LOG(ERROR) << "NvlinkTransport: cudaEventSynchronize failed: " + << cudaGetErrorString(sync_err); + return Status::Context("cudaEventSynchronize failed: " + + std::string(cudaGetErrorString(sync_err))); + } + + // Phase 1: Prepare slices and collect memcpy parameters + std::vector dsts, srcs; + std::vector sizes; + std::vector slices; + for (auto &request : entries) { TransferTask &task = batch_desc.task_list[task_id]; ++task_id; @@ -221,20 +513,25 @@ Status NvlinkTransport::submitTransfer( slice->task = &task; slice->target_id = request.target_id; slice->status = Slice::PENDING; + slice->ts = getCurrentTimeInNano(); + task.slice_list.push_back(slice); __sync_fetch_and_add(&task.slice_count, 1); - cudaError_t err; - if (slice->opcode == TransferRequest::READ) - err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr, - slice->length, cudaMemcpyDefault); - else - err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr, - slice->length, cudaMemcpyDefault); - if (err != cudaSuccess) - slice->markFailed(); - else - slice->markSuccess(); + + void *src = (request.opcode == TransferRequest::READ) + ? (void *)slice->local.dest_addr + : (void *)slice->source_addr; + void *dst = (request.opcode == TransferRequest::READ) + ? slice->source_addr + : (void *)slice->local.dest_addr; + srcs.push_back(src); + dsts.push_back(dst); + sizes.push_back(slice->length); + slices.push_back(slice); } + // Phase 2: Submit all memcpy operations + submitBatchMemcpy(slices, srcs, dsts, sizes, stream); + return Status::OK(); } @@ -248,6 +545,30 @@ Status NvlinkTransport::getTransferStatus(BatchID batch_id, size_t task_id, std::to_string(batch_id)); } auto &task = batch_desc.task_list[task_id]; + // Poll POSTED slices for async completion via cudaStreamQuery. + // Cache the query result per stream to avoid redundant driver calls. + // With SGLang-side torch.cuda.device(gpu_id), the calling thread's + // active device matches the stream's device, so no device switching + // is needed. + std::unordered_map stream_status_cache; + for (auto *slice : task.slice_list) { + if (slice && slice->status == Slice::POSTED) { + cudaStream_t stream = (cudaStream_t)slice->local.cuda_stream; + auto it = stream_status_cache.find(stream); + cudaError_t cuda_err; + if (it == stream_status_cache.end()) { + cuda_err = cudaStreamQuery(stream); + stream_status_cache[stream] = cuda_err; + } else { + cuda_err = it->second; + } + if (cuda_err == cudaSuccess) { + slice->markSuccess(); + } else if (cuda_err != cudaErrorNotReady) { + slice->markFailed(); + } + } + } status.transferred_bytes = task.transferred_bytes; uint64_t success_slice_count = task.success_slice_count; uint64_t failed_slice_count = task.failed_slice_count; @@ -266,6 +587,33 @@ Status NvlinkTransport::getTransferStatus(BatchID batch_id, size_t task_id, Status NvlinkTransport::submitTransferTask( const std::vector &task_list) { + // Get per-device transfer stream. See submitTransfer() for rationale. + CudaStreamEntry stream_entry = getStreamForRequest( + task_list.empty() ? nullptr : task_list[0]->request->source); + cudaStream_t stream = stream_entry.stream; + if (!stream) return Status::Context("Failed to create NVLink CUDA stream"); + // Synchronize with caller's GPU work via cudaEventSynchronize. + cudaEvent_t sync_event = getCallerSyncEvent(); + cudaError_t sync_err = cudaEventRecord(sync_event, cudaStreamPerThread); + if (sync_err != cudaSuccess) { + LOG(ERROR) << "NvlinkTransport: cudaEventRecord failed: " + << cudaGetErrorString(sync_err); + return Status::Context("cudaEventRecord failed: " + + std::string(cudaGetErrorString(sync_err))); + } + sync_err = cudaEventSynchronize(sync_event); + if (sync_err != cudaSuccess) { + LOG(ERROR) << "NvlinkTransport: cudaEventSynchronize failed: " + << cudaGetErrorString(sync_err); + return Status::Context("cudaEventSynchronize failed: " + + std::string(cudaGetErrorString(sync_err))); + } + + // Phase 1: Prepare slices and collect memcpy parameters + std::vector dsts, srcs; + std::vector sizes; + std::vector slices; + for (size_t index = 0; index < task_list.size(); ++index) { assert(task_list[index]); auto &task = *task_list[index]; @@ -286,20 +634,25 @@ Status NvlinkTransport::submitTransferTask( slice->task = &task; slice->target_id = request.target_id; slice->status = Slice::PENDING; + slice->ts = getCurrentTimeInNano(); task.slice_list.push_back(slice); __sync_fetch_and_add(&task.slice_count, 1); - cudaError_t err; - if (slice->opcode == TransferRequest::READ) - err = cudaMemcpy(slice->source_addr, (void *)slice->local.dest_addr, - slice->length, cudaMemcpyDefault); - else - err = cudaMemcpy((void *)slice->local.dest_addr, slice->source_addr, - slice->length, cudaMemcpyDefault); - if (err != cudaSuccess) - slice->markFailed(); - else - slice->markSuccess(); + + void *src = (request.opcode == TransferRequest::READ) + ? (void *)slice->local.dest_addr + : (void *)slice->source_addr; + void *dst = (request.opcode == TransferRequest::READ) + ? slice->source_addr + : (void *)slice->local.dest_addr; + srcs.push_back(src); + dsts.push_back(dst); + sizes.push_back(slice->length); + slices.push_back(slice); } + + // Phase 2: Submit all memcpy operations + submitBatchMemcpy(slices, srcs, dsts, sizes, stream); + return Status::OK(); } @@ -502,15 +855,23 @@ int NvlinkTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, int NvlinkTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { - for (auto &buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto &buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int NvlinkTransport::unregisterLocalMemoryBatch( const std::vector &addr_list) { - for (auto &addr : addr_list) unregisterLocalMemory(addr, false); - return metadata_->updateLocalSegmentDesc(); + int first_error = 0; + for (auto &addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret && !first_error) first_error = ret; + } + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } void *NvlinkTransport::allocatePinnedLocalMemory(size_t size) { diff --git a/mooncake-transfer-engine/src/transport/nvmeof_transport/CMakeLists.txt b/mooncake-transfer-engine/src/transport/nvmeof_transport/CMakeLists.txt index d64f4bcfa9..51c1982283 100644 --- a/mooncake-transfer-engine/src/transport/nvmeof_transport/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/nvmeof_transport/CMakeLists.txt @@ -1,5 +1,5 @@ file(GLOB NVMEOF_SOURCES "*.cpp") add_library(nvmeof_transport OBJECT ${NVMEOF_SOURCES}) -target_include_directories(nvmeof_transport PUBLIC "/usr/local/cuda/include") +target_include_directories(nvmeof_transport PUBLIC ${CUDAToolkit_INCLUDE_DIRS}) # target_link_libraries(nvmeof_transport PUBLIC transport) \ No newline at end of file diff --git a/mooncake-transfer-engine/src/transport/nvmeof_transport/cufile_desc_pool.cpp b/mooncake-transfer-engine/src/transport/nvmeof_transport/cufile_desc_pool.cpp index b7b4626c50..209284c508 100644 --- a/mooncake-transfer-engine/src/transport/nvmeof_transport/cufile_desc_pool.cpp +++ b/mooncake-transfer-engine/src/transport/nvmeof_transport/cufile_desc_pool.cpp @@ -36,9 +36,7 @@ CUFileDescPool::~CUFileDescPool() { // First, collect and destroy batch_handles from allocated descriptors for (size_t i = 0; i < MAX_NR_DESC; ++i) { if (descs_[i] != nullptr) { - cuFileBatchIODestroy(descs_[i]->batch_handle->handle); - delete descs_[i]->batch_handle; - delete descs_[i]; + destroyDesc(descs_[i]); descs_[i] = nullptr; } } @@ -50,6 +48,10 @@ CUFileDescPool::~CUFileDescPool() { delete batch_handle; } handle_pool_.clear(); + for (auto* desc : quarantined_descs_) { + destroyDesc(desc); + } + quarantined_descs_.clear(); } int CUFileDescPool::allocCUfileDesc(size_t batch_size) { @@ -110,7 +112,10 @@ int CUFileDescPool::allocCUfileDesc(size_t batch_size) { desc->batch_handle = batch_handle; desc->io_params.clear(); desc->io_params.reserve(max_batch_size_); - desc->io_events.resize(max_batch_size_); + desc->io_events.clear(); + desc->io_events.reserve(max_batch_size_); + desc->polled_events.resize(max_batch_size_); + desc->reusable = true; descs_[idx] = desc; return idx; @@ -133,12 +138,18 @@ int CUFileDescPool::pushParams(int idx, const CUfileIOParams_t& io_params) { } auto* desc = descs_[idx]; - if (desc->io_params.size() >= desc->io_params.capacity()) { + if (desc->io_params.size() >= max_batch_size_) { LOG(ERROR) << "Descriptor " << idx << " is full"; return -1; } - desc->io_params.push_back(io_params); + CUfileIOParams_t params = io_params; + const size_t slice_id = desc->io_params.size(); + params.cookie = + reinterpret_cast(static_cast(slice_id + 1)); + desc->io_params.push_back(params); + desc->io_events.push_back(CUfileIOEvents_t{ + .cookie = params.cookie, .status = CUFILE_WAITING, .ret = 0}); return 0; } @@ -163,30 +174,139 @@ int CUFileDescPool::submitBatch(int idx) { } CUfileIOEvents_t CUFileDescPool::getTransferStatus(int idx, int slice_id) { + if (!updateBatchStatus(idx)) { + return failedEvent(); + } + return getCachedTransferStatus(idx, slice_id); +} + +bool CUFileDescPool::updateBatchStatus(int idx) { RWSpinlock::WriteGuard guard(mutex_); if (idx < 0 || idx >= (int)MAX_NR_DESC || descs_[idx] == nullptr) { LOG(ERROR) << "Invalid descriptor index: " << idx; - CUfileIOEvents_t event; - event.status = CUFILE_FAILED; - event.ret = -1; - return event; + return false; + } + + return updateBatchStatus(descs_[idx], idx); +} + +CUfileIOEvents_t CUFileDescPool::getCachedTransferStatus(int idx, + int slice_id) { + RWSpinlock::ReadGuard guard(mutex_); + if (idx < 0 || idx >= (int)MAX_NR_DESC || descs_[idx] == nullptr) { + LOG(ERROR) << "Invalid descriptor index: " << idx; + return failedEvent(); } auto* desc = descs_[idx]; if (slice_id < 0 || slice_id >= (int)desc->io_params.size()) { LOG(ERROR) << "Invalid slice_id " << slice_id << " for descriptor " << idx << " (size: " << desc->io_params.size() << ")"; - CUfileIOEvents_t event; - event.status = CUFILE_FAILED; - event.ret = -1; - return event; + return failedEvent(); } + return desc->io_events[slice_id]; +} + +bool CUFileDescPool::updateBatchStatus(CUFileBatchDesc* desc, int idx) { unsigned nr = desc->io_params.size(); - CUFILE_CHECK(cuFileBatchIOGetStatus(desc->batch_handle->handle, 0, &nr, - desc->io_events.data(), nullptr)); + if (desc->polled_events.size() < nr) { + LOG(ERROR) << "Completion buffer is too small for descriptor " << idx; + return false; + } - return desc->io_events[slice_id]; + CUfileError_t rc = + cuFileBatchIOGetStatus(desc->batch_handle->handle, 0, &nr, + desc->polled_events.data(), nullptr); + if (rc.err != CU_FILE_SUCCESS) { + LOG(WARNING) << "cuFileBatchIOGetStatus failed for descriptor " << idx + << ": " << cuFileGetErrorString(rc); + return false; + } + + for (unsigned i = 0; i < nr; ++i) { + const auto& event = desc->polled_events[i]; + if (!cachePolledEvent(desc->io_events, event)) { + LOG(ERROR) << "Invalid completion cookie " + << reinterpret_cast(event.cookie) + << " for descriptor " << idx; + } + } + + return true; +} + +bool CUFileDescPool::cachePolledEvent(std::vector& io_events, + const CUfileIOEvents_t& event) { + const uintptr_t cookie = reinterpret_cast(event.cookie); + if (cookie == 0 || cookie > io_events.size()) return false; + io_events[cookie - 1] = event; + return true; +} + +bool CUFileDescPool::isTerminalStatus(CUfileStatus_t status) { + return status != CUFILE_WAITING && status != CUFILE_PENDING; +} + +CUfileIOEvents_t CUFileDescPool::failedEvent() { + CUfileIOEvents_t event = {}; + event.status = CUFILE_FAILED; + event.ret = -1; + return event; +} + +void CUFileDescPool::destroyDesc(CUFileBatchDesc* desc) { + cuFileBatchIODestroy(desc->batch_handle->handle); + delete desc->batch_handle; + delete desc; +} + +void CUFileDescPool::cleanupQuarantinedDescs() { + auto it = quarantined_descs_.begin(); + while (it != quarantined_descs_.end()) { + auto* desc = *it; + updateBatchStatus(desc, -1); + + bool all_terminal = true; + for (const auto& event : desc->io_events) { + if (!isTerminalStatus(event.status)) { + all_terminal = false; + break; + } + } + + if (all_terminal) { + destroyDesc(desc); + it = quarantined_descs_.erase(it); + } else { + ++it; + } + } +} + +bool CUFileDescPool::cancelBatch(int idx) { + RWSpinlock::WriteGuard guard(mutex_); + if (idx < 0 || idx >= (int)MAX_NR_DESC || descs_[idx] == nullptr) { + LOG(ERROR) << "Invalid descriptor index: " << idx; + return false; + } + + CUfileError_t rc = cuFileBatchIOCancel(descs_[idx]->batch_handle->handle); + if (rc.err != CU_FILE_SUCCESS) { + LOG(WARNING) << "cuFileBatchIOCancel failed for descriptor " << idx + << ": " << cuFileGetErrorString(rc); + return false; + } + return true; +} + +void CUFileDescPool::markUnreusable(int idx) { + RWSpinlock::WriteGuard guard(mutex_); + if (idx < 0 || idx >= (int)MAX_NR_DESC || descs_[idx] == nullptr) { + LOG(ERROR) << "Invalid descriptor index: " << idx; + return; + } + descs_[idx]->reusable = false; } int CUFileDescPool::getSliceNum(int idx) { @@ -207,23 +327,28 @@ int CUFileDescPool::freeCUfileDesc(int idx) { } auto* desc = descs_[idx]; + const bool reusable = desc->reusable; - // IMPORTANT: Caller should ensure all IOs are completed (via - // getTransferStatus) before calling freeCUfileDesc, as cuFile may still - // access io_params otherwise. This is critical for the handle pooling - // optimization - the handle will be immediately reused and could lead to - // use-after-free bugs if IOs are in-flight. - // + // Reusable descriptors are safe to recycle only after the caller observed + // all IOs complete. Non-reusable descriptors may still be referenced by + // cuFile after a bounded failure cleanup, so keep their handle and params + // quarantined until a later poll observes terminal status for every IO. // Return the handle to pool for reuse (avoid expensive // cuFileBatchIODestroy) { std::lock_guard lock(handle_pool_lock_); - handle_pool_.push_back(desc->batch_handle); + if (reusable) { + handle_pool_.push_back(desc->batch_handle); + } else { + quarantined_descs_.push_back(desc); + } } - // Delete the descriptor (each allocation gets a fresh one) - delete desc; + if (reusable) { + delete desc; + } descs_[idx] = nullptr; + cleanupQuarantinedDescs(); return 0; } @@ -236,4 +361,4 @@ CUFileBatchDesc* CUFileDescPool::getDesc(int idx) { return descs_[idx]; } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/nvmeof_transport/nvmeof_transport.cpp b/mooncake-transfer-engine/src/transport/nvmeof_transport/nvmeof_transport.cpp index d019abd7d0..999e4f3bb8 100644 --- a/mooncake-transfer-engine/src/transport/nvmeof_transport/nvmeof_transport.cpp +++ b/mooncake-transfer-engine/src/transport/nvmeof_transport/nvmeof_transport.cpp @@ -39,6 +39,9 @@ NVMeoFTransport::NVMeoFTransport() { desc_pool_ = std::make_shared(); } +NVMeoFTransport::NVMeoFTransport(std::shared_ptr desc_pool) + : desc_pool_(std::move(desc_pool)) {} + NVMeoFTransport::~NVMeoFTransport() {} Transport::TransferStatusEnum from_cufile_transfer_status( @@ -76,38 +79,139 @@ NVMeoFTransport::BatchID NVMeoFTransport::allocateBatchID(size_t batch_size) { Status NVMeoFTransport::getTransferStatus(BatchID batch_id, size_t task_id, TransferStatus &status) { + if (batch_id == 0) { + return Status::InvalidArgument("NVMeoFTransport: Invalid batch ID"); + } auto &batch_desc = *((BatchDesc *)(batch_id)); + if (task_id >= batch_desc.task_list.size()) { + return Status::InvalidArgument("NVMeoFTransport: Task ID out of range"); + } + if (batch_desc.context == nullptr) { + return Status::InvalidArgument( + "NVMeoFTransport: Batch was not allocated by this transport"); + } auto &task = batch_desc.task_list[task_id]; auto &nvmeof_desc = *((NVMeoFBatchDesc *)(batch_desc.context)); - // LOG(DEBUG) << "get t n " << nr; - // 1. get task -> id map - TransferStatus transfer_status = {.s = Transport::PENDING, - .transferred_bytes = 0}; + if (task_id >= nvmeof_desc.task_to_slices.size()) { + return Status::InvalidArgument( + "NVMeoFTransport: Task has no submitted slices"); + } + + if (task.is_finished && task_id < nvmeof_desc.transfer_status.size()) { + status = nvmeof_desc.transfer_status[task_id]; + return Status::OK(); + } + auto [slice_id, slice_num] = nvmeof_desc.task_to_slices[task_id]; - for (size_t i = slice_id; i < slice_id + slice_num; ++i) { - // LOG(INFO) << "task " << task_id << " i " << i << " upper bound " << - // slice_num; - auto event = desc_pool_->getTransferStatus(nvmeof_desc.desc_idx_, i); - transfer_status.s = from_cufile_transfer_status(event.status); - // TODO(FIXME): what to do if multi slices have different status? - if (transfer_status.s == COMPLETED) { - transfer_status.transferred_bytes += event.ret; - } else { - break; + thread_local std::vector slice_statuses; + collectSliceStatuses(nvmeof_desc.desc_idx_, slice_id, slice_num, + slice_statuses); + + bool is_finished = false; + status = aggregateTransferStatus(slice_statuses, is_finished); + if (!is_finished && isTerminalFailure(status.s)) { + desc_pool_->cancelBatch(nvmeof_desc.desc_idx_); + collectSliceStatuses(nvmeof_desc.desc_idx_, slice_id, slice_num, + slice_statuses); + status = aggregateTransferStatus(slice_statuses, is_finished); + if (!is_finished && isTerminalFailure(status.s)) { + desc_pool_->markUnreusable(nvmeof_desc.desc_idx_); + is_finished = true; } } - if (transfer_status.s == COMPLETED) { + if (is_finished) { + if (task_id < nvmeof_desc.transfer_status.size()) { + nvmeof_desc.transfer_status[task_id] = status; + } task.is_finished = true; } - status = transfer_status; return Status::OK(); } -// Dummy implement for solving build issues, WIP Status NVMeoFTransport::submitTransferTask( const std::vector &task_list) { - /* TBD */ - return Status::OK(); + // MultiTransport owns these generic BatchDesc objects, so this transport + // cannot attach or reclaim the NVMe-specific descriptor required by + // cuFile. No asynchronous work was started; make the tasks releasable. + for (auto *task : task_list) { + if (task != nullptr) task->is_finished = true; + } + return Status::NotImplemented( + "NVMeoFTransport does not support MultiTransport batches"); +} + +Transport::TransferStatus NVMeoFTransport::aggregateTransferStatus( + const std::vector &slice_statuses, bool &is_finished) { + TransferStatus result = {.s = COMPLETED, .transferred_bytes = 0}; + is_finished = true; + bool has_pending = false; + + // Terminal failures use a fixed precedence so the result does not depend + // on the order in which cuFile reports completions. + int failure_priority = 0; + for (const auto &slice_status : slice_statuses) { + switch (slice_status.s) { + case COMPLETED: + result.transferred_bytes += slice_status.transferred_bytes; + break; + case WAITING: + is_finished = false; + break; + case PENDING: + has_pending = true; + is_finished = false; + break; + case INVALID: + if (failure_priority < 1) { + result.s = INVALID; + failure_priority = 1; + } + break; + case CANCELED: + if (failure_priority < 2) { + result.s = CANCELED; + failure_priority = 2; + } + break; + case TIMEOUT: + if (failure_priority < 3) { + result.s = TIMEOUT; + failure_priority = 3; + } + break; + case FAILED: + result.s = FAILED; + failure_priority = 4; + break; + } + } + + if (slice_statuses.empty()) { + result.s = INVALID; + } else if (!is_finished && failure_priority == 0) { + result.s = has_pending ? PENDING : WAITING; + } + return result; +} + +void NVMeoFTransport::collectSliceStatuses( + int desc_idx, size_t slice_id, size_t slice_num, + std::vector &slice_statuses) { + slice_statuses.clear(); + slice_statuses.reserve(slice_num); + desc_pool_->updateBatchStatus(desc_idx); + for (size_t i = slice_id; i < slice_id + slice_num; ++i) { + auto event = desc_pool_->getCachedTransferStatus(desc_idx, i); + auto slice_status = from_cufile_transfer_status(event.status); + slice_statuses.push_back(TransferStatus{ + .s = slice_status, + .transferred_bytes = slice_status == COMPLETED ? event.ret : 0}); + } +} + +bool NVMeoFTransport::isTerminalFailure(TransferStatusEnum status) { + return status == INVALID || status == CANCELED || status == TIMEOUT || + status == FAILED; } Status NVMeoFTransport::submitTransfer( diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp index f6fcfdace9..54db6de43a 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/endpoint_store.cpp @@ -94,13 +94,15 @@ int FIFOEndpointStore::deleteEndpoint(const std::string &peer_nic_path) { return 0; } -int FIFOEndpointStore::deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) { +int FIFOEndpointStore::deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr, + std::string *deleted_peer_nic_path) { RWSpinlock::WriteGuard guard(endpoint_map_lock_); // Find endpoint by pointer for (auto iter = endpoint_map_.begin(); iter != endpoint_map_.end(); ++iter) { if (iter->second.get() == endpoint_ptr) { std::string peer_nic_path = iter->first; + if (deleted_peer_nic_path) *deleted_peer_nic_path = peer_nic_path; waiting_list_len_++; iter->second->beginDestroy(); waiting_list_.insert(iter->second); @@ -142,16 +144,36 @@ void FIFOEndpointStore::reclaimEndpoint() { size_t FIFOEndpointStore::getSize() { return endpoint_map_.size(); } int FIFOEndpointStore::destroyQPs() { + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + int ret = 0; + + // Always transition QPs to ERR before destroy to flush inflight WRs. + for (auto &endpoint : waiting_list_) { + endpoint->beginDestroy(); + } for (auto &kv : endpoint_map_) { - kv.second->destroyQP(); + kv.second->beginDestroy(); } - return 0; + + for (auto &endpoint : waiting_list_) { + if (endpoint->destroyQP()) ret = -1; + } + for (auto &kv : endpoint_map_) { + if (kv.second->destroyQP()) ret = -1; + } + return ret; } int FIFOEndpointStore::disconnectQPs() { + RWSpinlock::WriteGuard guard(endpoint_map_lock_); for (auto &kv : endpoint_map_) { - kv.second->disconnect(); + kv.second->beginDestroy(); + waiting_list_.insert(kv.second); } + waiting_list_len_ += endpoint_map_.size(); + endpoint_map_.clear(); + fifo_list_.clear(); + fifo_map_.clear(); return 0; } @@ -248,13 +270,15 @@ int SIEVEEndpointStore::deleteEndpoint(const std::string &peer_nic_path) { return 0; } -int SIEVEEndpointStore::deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) { +int SIEVEEndpointStore::deleteEndpointByPtr( + const RdmaEndPoint *endpoint_ptr, std::string *deleted_peer_nic_path) { RWSpinlock::WriteGuard guard(endpoint_map_lock_); // Find endpoint by pointer for (auto iter = endpoint_map_.begin(); iter != endpoint_map_.end(); ++iter) { if (iter->second.first.get() == endpoint_ptr) { std::string peer_nic_path = iter->first; + if (deleted_peer_nic_path) *deleted_peer_nic_path = peer_nic_path; iter->second.first->beginDestroy(); waiting_list_len_++; waiting_list_.insert(iter->second.first); @@ -312,14 +336,35 @@ void SIEVEEndpointStore::reclaimEndpoint() { } int SIEVEEndpointStore::destroyQPs() { - for (auto &endpoint : waiting_list_) endpoint->destroyQP(); - for (auto &kv : endpoint_map_) kv.second.first->destroyQP(); - return 0; + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + int ret = 0; + + // Always transition QPs to ERR before destroy to flush inflight WRs. + for (auto &endpoint : waiting_list_) { + endpoint->beginDestroy(); + } + for (auto &kv : endpoint_map_) { + kv.second.first->beginDestroy(); + } + + for (auto &endpoint : waiting_list_) + if (endpoint->destroyQP()) ret = -1; + for (auto &kv : endpoint_map_) + if (kv.second.first->destroyQP()) ret = -1; + return ret; } int SIEVEEndpointStore::disconnectQPs() { - for (auto &endpoint : waiting_list_) endpoint->disconnect(); - for (auto &kv : endpoint_map_) kv.second.first->disconnect(); + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + for (auto &kv : endpoint_map_) { + kv.second.first->beginDestroy(); + waiting_list_.insert(kv.second.first); + } + waiting_list_len_ += endpoint_map_.size(); + endpoint_map_.clear(); + fifo_list_.clear(); + fifo_map_.clear(); + hand_ = std::nullopt; return 0; } diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp index f980797237..a8e395fbd3 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_context.cpp @@ -32,6 +32,7 @@ #include "config.h" #include "cuda_alike.h" #include "environ.h" +#include "hip_device_guard.h" #if defined(USE_HIP_DMABUF) #include @@ -160,6 +161,8 @@ std::string gidBytesToString(const uint8_t *raw) { RdmaContext::RdmaContext(RdmaTransport &engine, const std::string &device_name) : device_name_(device_name), engine_(engine), + connect_pause_( + [] { return static_cast(getCurrentTimeInNano()); }), next_comp_channel_index_(0), next_comp_vector_index_(0), next_cq_list_index_(0), @@ -168,7 +171,9 @@ RdmaContext::RdmaContext(RdmaTransport &engine, const std::string &device_name) static std::once_flag g_once_flag; auto fork_init = []() { int ret = ibv_fork_init(); - if (ret) PLOG(ERROR) << "RDMA context setup failed: fork compatibility"; + if (ret) + LOG(ERROR) << "RDMA context setup failed: fork compatibility: " + << strerror(ret); }; std::call_once(g_once_flag, fork_init); } @@ -180,6 +185,13 @@ RdmaContext::~RdmaContext() { int RdmaContext::construct(size_t num_cq_list, size_t num_comp_channels, uint8_t port, int gid_index, size_t max_cqe, int max_endpoints) { + if (num_cq_list == 0 || num_comp_channels == 0) { + LOG(ERROR) << "Invalid RDMA completion configuration for device " + << device_name_ << ": num_cq_list=" << num_cq_list + << ", num_comp_channels=" << num_comp_channels; + return ERR_INVALID_ARGUMENT; + } + // Create endpoint store based on configuration auto &config = globalConfig(); switch (config.endpoint_store_type) { @@ -201,6 +213,12 @@ int RdmaContext::construct(size_t num_cq_list, size_t num_comp_channels, return ERR_CONTEXT; } + if (context_->num_comp_vectors <= 0) { + LOG(ERROR) << "RDMA device " << device_name_ + << " exposes no completion vectors"; + return ERR_CONTEXT; + } + pd_ = ibv_alloc_pd(context_); if (!pd_) { PLOG(ERROR) << "Failed to allocate new protection domain on device " @@ -293,12 +311,34 @@ int RdmaContext::socketId() { int RdmaContext::deconstruct() { worker_pool_.reset(); - endpoint_store_->destroyQPs(); + // Graceful teardown order: QPs -> MRs. + if (endpoint_store_) { + endpoint_store_->disconnectQPs(); + + // In normal graceful shutdown, reclaim should finish quickly. + constexpr auto kReclaimTimeout = std::chrono::seconds(10); + auto start = std::chrono::steady_clock::now(); + while (endpoint_store_->waitingListSize() > 0) { + endpoint_store_->reclaimEndpoint(); + if (endpoint_store_->waitingListSize() == 0) break; + if (std::chrono::steady_clock::now() - start > kReclaimTimeout) { + LOG(WARNING) << "Endpoint reclaim timed out during graceful " + "shutdown; forcing QP destruction"; + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + if (endpoint_store_->destroyQPs()) { + LOG(ERROR) << "Failed to destroy all QPs before MR deregistration"; + } + } for (auto &[_, entry] : memory_region_map_) { int ret = ibv_dereg_mr(entry.mr); if (ret) { - PLOG(ERROR) << "Failed to unregister memory region"; + LOG(ERROR) << "Failed to unregister memory region: " + << strerror(ret); } } memory_region_map_.clear(); @@ -308,7 +348,8 @@ int RdmaContext::deconstruct() { int ret = ibv_destroy_cq(cq_list_[i].native); if (ret) { - PLOG(ERROR) << "Failed to destroy completion queue"; + LOG(ERROR) << "Failed to destroy completion queue: " + << strerror(ret); } } cq_list_.clear(); @@ -320,59 +361,58 @@ int RdmaContext::deconstruct() { if (comp_channel_) { for (size_t i = 0; i < num_comp_channel_; ++i) - if (comp_channel_[i]) - if (ibv_destroy_comp_channel(comp_channel_[i])) - LOG(ERROR) << "Failed to destroy completion channel"; + if (comp_channel_[i]) { + int ret = ibv_destroy_comp_channel(comp_channel_[i]); + if (ret) + LOG(ERROR) << "Failed to destroy completion channel: " + << strerror(ret); + } delete[] comp_channel_; comp_channel_ = nullptr; } if (pd_) { - if (ibv_dealloc_pd(pd_)) - PLOG(ERROR) << "Failed to deallocate protection domain"; + int ret = ibv_dealloc_pd(pd_); + if (ret) + LOG(ERROR) << "Failed to deallocate protection domain: " + << strerror(ret); pd_ = nullptr; } if (context_) { - if (ibv_close_device(context_)) - PLOG(ERROR) << "Failed to close device context"; + int ret = ibv_close_device(context_); + if (ret) + LOG(ERROR) << "Failed to close device context: " << strerror(ret); context_ = nullptr; } return 0; } -int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, - int access, - MemoryRegionMeta &mrMeta) { - if (length > (size_t)globalConfig().max_mr_size) { - PLOG(WARNING) << "The buffer length exceeds device max_mr_size, " - << "shrink it to " << globalConfig().max_mr_size; - length = (size_t)globalConfig().max_mr_size; - } -#if defined(USE_MLU) || defined(USE_MACA) || defined(USE_CUDA) - // Implement register memory in a way that does not assume the presence of - // nvidia-peermem. If memory is on CPU call ibv_reg_mr() as usual. If memory - // is on GPU then use ibv_reg_dmabuf_mr() instead which does not require - // nvidia-peermem. +int RdmaContext::exportDmabuf(void *addr, DmabufExport &out) { + out = DmabufExport{}; + (void)addr; // unused on the host-only (#else) build +#if defined(USE_MLU) || defined(USE_MACA) || defined(USE_CUDA) || \ + defined(USE_SUPA) + // Decide host vs GPU without assuming the presence of nvidia-peermem. Host + // memory uses the plain ibv_reg_mr() path. GPU memory is exported once as a + // dma_buf fd that every NIC then imports, so the driver keeps a single + // BAR1 window for the buffer instead of one per NIC. CUmemorytype memType; CUresult result = cuPointerGetAttribute( &memType, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, (CUdeviceptr)addr); - // Register memory depending on whether memory is on host or GPU. if (result != CUDA_SUCCESS || memType == CU_MEMORYTYPE_HOST) { - mrMeta.addr = addr; - mrMeta.mr = ibv_reg_mr(pd_, addr, length, access); -#if defined(USE_CUDA) + out.method = DmabufExport::Method::kHostReg; +#if defined(USE_CUDA) || defined(USE_SUPA) } else if (memType == CU_MEMORYTYPE_DEVICE && Environ::Get().GetWithNvidiaPeermem()) { // WITH_NVIDIA_PEERMEM env var is set: use ibv_reg_mr() directly for // GPU memory (requires the nvidia-peermem kernel module to be loaded). - mrMeta.addr = addr; - mrMeta.mr = ibv_reg_mr(pd_, addr, length, access); + out.method = DmabufExport::Method::kHostReg; #endif } else if (memType == CU_MEMORYTYPE_DEVICE) { -#if defined(USE_CUDA) +#if defined(USE_CUDA) || defined(USE_SUPA) // Ensure a CUDA context is current — worker threads or callers // from non-CUDA threads may lack one. unsigned int devOrd = 0; @@ -399,13 +439,15 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, cuGetErrorString(result, &errStr); LOG(ERROR) << "Failed to call cuMemGetAddressRange for " << (uintptr_t)addr << " cuda error=" << errStr; -#if defined(USE_CUDA) +#if defined(USE_CUDA) || defined(USE_SUPA) cuDevicePrimaryCtxRelease(cuDev); #endif return ERR_CONTEXT; } int dmabuf_fd; + // flags must be 0: the PCIE-BAR1 mapping flag is rejected (error 801) + // on some GPU/driver combinations (e.g. B200). result = cuMemGetHandleForAddressRange( &dmabuf_fd, allocBase, allocSize, CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD, 0); @@ -415,23 +457,15 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, LOG(ERROR) << "Failed to retrieve dmabuf for " << (uintptr_t)addr << " base=" << (uintptr_t)allocBase << " size=" << allocSize << " cuda error=" << errStr; -#if defined(USE_CUDA) +#if defined(USE_CUDA) || defined(USE_SUPA) cuDevicePrimaryCtxRelease(cuDev); #endif return ERR_CONTEXT; } - mrMeta.addr = addr; - uint64_t dmabuf_offset = (uintptr_t)addr - (uintptr_t)allocBase; - mrMeta.mr = ibv_reg_dmabuf_mr(pd_, dmabuf_offset, length, - (uintptr_t)addr, dmabuf_fd, access); - const int regErrno = errno; - if (close(dmabuf_fd) != 0) { - PLOG(WARNING) << "Failed to close dmabuf fd"; - } - if (!mrMeta.mr) { - errno = regErrno; - } -#if defined(USE_CUDA) + out.method = DmabufExport::Method::kDmabufReg; + out.fd = dmabuf_fd; + out.offset = (uintptr_t)addr - (uintptr_t)allocBase; +#if defined(USE_CUDA) || defined(USE_SUPA) cuDevicePrimaryCtxRelease(cuDev); #endif } @@ -442,8 +476,7 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, if (hipRes != hipSuccess || hipAttr.type == hipMemoryTypeHost || hipAttr.type == hipMemoryTypeUnregistered) { // Host memory — standard ibv_reg_mr() path. - mrMeta.addr = addr; - mrMeta.mr = ibv_reg_mr(pd_, addr, length, access); + out.method = DmabufExport::Method::kHostReg; } else if (hipAttr.type == hipMemoryTypeManaged) { // Managed (unified) memory pages can migrate between host and device; // hsa_amd_portable_export_dmabuf captures the device-side handle at @@ -452,17 +485,15 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, LOG(WARNING) << "HIP managed memory at " << (uintptr_t)addr << " — dmabuf export skipped (pages may migrate); " "falling back to ibv_reg_mr"; - mrMeta.addr = addr; - mrMeta.mr = ibv_reg_mr(pd_, addr, length, access); + out.method = DmabufExport::Method::kHostReg; } else if (hipAttr.type == hipMemoryTypeDevice && !isKernelDmabufSupported()) { // Kernel lacks CONFIG_PCI_P2PDMA / CONFIG_DMABUF_MOVE_NOTIFY — // ibv_reg_dmabuf_mr may succeed but transfers will silently fail. - // Fail at registration time instead. - mrMeta.addr = addr; - mrMeta.mr = ibv_reg_mr(pd_, addr, length, access); + // Fall back to ibv_reg_mr() instead. + out.method = DmabufExport::Method::kHostReg; } else if (hipAttr.type == hipMemoryTypeDevice) { - // Device memory + kernel support — export dmabuf fd and register. + // Device memory + kernel support — export the dmabuf fd. // Pin to the owning device for the duration of the export calls. struct HipDeviceGuard { int prev_device = 0; @@ -513,23 +544,58 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, return ERR_CONTEXT; } - mrMeta.addr = addr; + out.method = DmabufExport::Method::kDmabufReg; + out.fd = dmabuf_fd; // Offset within the dmabuf-backed region: distance from the // allocation base, plus any offset hsa returned for the export. - uint64_t reg_offset = - (uintptr_t)addr - (uintptr_t)allocBase + hsa_dmabuf_offset; - mrMeta.mr = ibv_reg_dmabuf_mr(pd_, reg_offset, length, (uintptr_t)addr, - dmabuf_fd, access); - const int regErrno = errno; - if (close(dmabuf_fd) != 0) { + out.offset = (uintptr_t)addr - (uintptr_t)allocBase + hsa_dmabuf_offset; + } +#else + out.method = DmabufExport::Method::kHostReg; +#endif + return 0; +} + +void RdmaContext::closeDmabufExport(DmabufExport &exp) { + if (exp.fd >= 0) { + if (close(exp.fd) != 0) { PLOG(WARNING) << "Failed to close dmabuf fd"; } - if (!mrMeta.mr) { - errno = regErrno; - } + exp.fd = -1; + } +} + +int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, + int access, + const DmabufExport &exp, + MemoryRegionMeta &mrMeta) { + if (length > (size_t)globalConfig().max_mr_size) { + // #2017: registerLocalMemory auto-chunks buffers to <= max_mr_size, so + // no larger buffer should reach here. Fail loudly instead of silently + // truncating the MR — a truncated MR advertises bytes past the + // registered region and causes IBV_WC_REM_ACCESS_ERR on RDMA ops whose + // target lands past the boundary. + LOG(ERROR) << "Buffer length " << length + << " exceeds device max_mr_size " + << globalConfig().max_mr_size + << " (should have been chunked before registration, #2017)"; + return ERR_INVALID_ARGUMENT; } -#else mrMeta.addr = addr; +#if defined(USE_MLU) || defined(USE_MACA) || defined(USE_CUDA) || \ + defined(USE_HIP_DMABUF) || defined(USE_SUPA) + if (exp.method == DmabufExport::Method::kDmabufReg) { + // Import the shared dma_buf fd into this NIC's PD. The fd is kept open + // by the caller until every NIC has registered; this MR takes its own + // reference, so all NICs share one dma_buf object (and one BAR1 + // window). + mrMeta.mr = ibv_reg_dmabuf_mr(pd_, exp.offset, length, (uintptr_t)addr, + exp.fd, access); + } else { + mrMeta.mr = ibv_reg_mr(pd_, addr, length, access); + } +#else + (void)exp; mrMeta.mr = ibv_reg_mr(pd_, addr, length, access); #endif if (!mrMeta.mr) { @@ -539,9 +605,12 @@ int RdmaContext::registerMemoryRegionInternal(void *addr, size_t length, return 0; } -int RdmaContext::registerMemoryRegion(void *addr, size_t length, int access) { +int RdmaContext::registerMemoryRegion(void *addr, size_t length, int access, + const DmabufExport &exp) { + // Placeholder context for a failed RNIC: no PD to register against. + if (!pd_) return 0; MemoryRegionMeta mrMeta; - int ret = registerMemoryRegionInternal(addr, length, access, mrMeta); + int ret = registerMemoryRegionInternal(addr, length, access, exp, mrMeta); if (ret != 0) { return ret; } @@ -550,6 +619,21 @@ int RdmaContext::registerMemoryRegion(void *addr, size_t length, int access) { return 0; } +int RdmaContext::registerMemoryRegion(void *addr, size_t length, int access) { + if (!pd_) return 0; // placeholder context: skip the dma_buf export too + // Single-NIC convenience path: export, register, and close the fd here. + // The shared-fd benefit only matters when a buffer is registered against + // multiple NICs (see RdmaTransport::registerLocalMemoryInternal). + DmabufExport exp; + int ret = exportDmabuf(addr, exp); + if (ret != 0) { + return ret; + } + ret = registerMemoryRegion(addr, length, access, exp); + closeDmabufExport(exp); + return ret; +} + int RdmaContext::unregisterMemoryRegion(void *addr) { RWSpinlock::WriteGuard guard(memory_regions_lock_); auto iter = findMemoryRegionContaining(reinterpret_cast(addr)); @@ -565,9 +649,18 @@ int RdmaContext::unregisterMemoryRegion(void *addr) { } int RdmaContext::preTouchMemory(void *addr, size_t length) { + if (!pd_) return 0; // placeholder context + DmabufExport exp; + int ret = exportDmabuf(addr, exp); + if (ret != 0) { + return ret; + } MemoryRegionMeta mrMeta; - int ret = registerMemoryRegionInternal(addr, length, IBV_ACCESS_LOCAL_WRITE, - mrMeta); + ret = registerMemoryRegionInternal(addr, length, IBV_ACCESS_LOCAL_WRITE, + exp, mrMeta); + // The MR (if created) holds its own reference, so closing the fd now is + // safe and does not affect the subsequent dereg. + closeDmabufExport(exp); if (ret != 0) { return ret; } @@ -575,6 +668,8 @@ int RdmaContext::preTouchMemory(void *addr, size_t length) { } uint32_t RdmaContext::rkey(void *addr) { + // Placeholder context holds no MRs; its zero key is never selected. + if (!pd_) return 0; RWSpinlock::ReadGuard guard(memory_regions_lock_); auto iter = findMemoryRegionContaining(reinterpret_cast(addr)); if (iter != memory_region_map_.end()) return iter->second.mr->rkey; @@ -584,6 +679,7 @@ uint32_t RdmaContext::rkey(void *addr) { } uint32_t RdmaContext::lkey(void *addr) { + if (!pd_) return 0; // see rkey() RWSpinlock::ReadGuard guard(memory_regions_lock_); auto iter = findMemoryRegionContaining(reinterpret_cast(addr)); if (iter != memory_region_map_.end()) return iter->second.mr->lkey; @@ -616,7 +712,7 @@ RdmaContext::findMemoryRegionContaining(uintptr_t addr) const { std::shared_ptr RdmaContext::endpoint( const std::string &peer_nic_path) { - if (!active_) { + if (!active_.load(std::memory_order_acquire)) { LOG(ERROR) << "Context is not active: " << deviceName(); return nullptr; } @@ -650,9 +746,40 @@ int RdmaContext::deleteEndpoint(const std::string &peer_nic_path) { } int RdmaContext::deleteEndpointByPtr(const RdmaEndPoint *endpoint_ptr) { - return endpoint_store_->deleteEndpointByPtr(endpoint_ptr); + // Tearing an endpoint down (path failure / QP fatal) means this peer is + // failing; pause active reconnection to its address so the CQ poller isn't + // blocked re-handshaking a likely-gone peer. + // + // Resolve the peer path *inside* the store, under its lock: the raw pointer + // may already be freed (e.g. an IBV_EVENT_QP_FATAL racing endpoint + // destruction), so we must not dereference it here. The store only compares + // pointer identity and returns the path from the live map key, and we arm + // the pause only if the endpoint was actually found. No-op when TTL is 0. + std::string deleted_peer_nic_path; + int ret = endpoint_store_->deleteEndpointByPtr(endpoint_ptr, + &deleted_peer_nic_path); + if (!deleted_peer_nic_path.empty()) pauseConnect(deleted_peer_nic_path); + return ret; +} + +void RdmaContext::pauseConnect(const std::string &peer_nic_path) { + int ttl_ms = globalConfig().conn_pause_ttl_ms; + if (ttl_ms <= 0) return; // disabled + auto server_name = getServerNameFromNicPath(peer_nic_path); + if (server_name.empty()) return; + connect_pause_.pauseFor(server_name, + static_cast(ttl_ms) * 1000000ull); } +bool RdmaContext::isConnectPaused(const std::string &peer_nic_path) { + if (globalConfig().conn_pause_ttl_ms <= 0) return false; // disabled + auto server_name = getServerNameFromNicPath(peer_nic_path); + if (server_name.empty()) return false; + return connect_pause_.isPaused(server_name); +} + +void RdmaContext::pruneConnectPause() { connect_pause_.prune(); } + void RdmaContext::reclaimEndpoints() { endpoint_store_->reclaimEndpoint(); } size_t RdmaContext::waitingListSize() const { @@ -760,9 +887,11 @@ static GidNetworkState autoGidStateFromSelection( const AutoGidSelection &selection) { switch (selection.candidate_class) { case AutoGidCandidateClass::kNetworkRoutable: + case AutoGidCandidateClass::kNetworkPrivateV4: case AutoGidCandidateClass::kNetworkDegraded: return GidNetworkState::GID_WITH_NETWORK; case AutoGidCandidateClass::kNoNetworkRoutable: + case AutoGidCandidateClass::kNoNetworkPrivateV4: case AutoGidCandidateClass::kNoNetworkDegraded: case AutoGidCandidateClass::kFallbackNonzero: return GidNetworkState::GID_WITHOUT_NETWORK; @@ -849,9 +978,11 @@ bool RdmaContext::reprobeAutoGid( } ibv_port_attr port_attr; - if (ibv_query_port(current_context, current_port, &port_attr)) { - PLOG(WARNING) << "Failed to reprobe port attributes on " << device_name_ - << "/" << static_cast(current_port); + int ret = ibv_query_port(current_context, current_port, &port_attr); + if (ret) { + LOG(WARNING) << "Failed to reprobe port attributes on " << device_name_ + << "/" << static_cast(current_port) << ": " + << strerror(ret); return false; } @@ -938,6 +1069,122 @@ bool RdmaContext::reprobeAutoGid( return true; } +GidRefreshResult RdmaContext::refreshCurrentGid(std::string *previous_gid, + std::string *next_gid) { + std::lock_guard reprobe_guard(gid_reprobe_lock_); + std::string current_gid_string; + int current_gid_index = -1; + int next_gid_index = -1; + uint16_t current_lid = 0; + ibv_context *current_context = nullptr; + uint8_t current_port = 0; + bool auto_gid_selection_enabled = false; + { + std::lock_guard guard(gid_lock_); + if (!context_) { + return GidRefreshResult::FAILED; + } + current_gid_index = gid_index_; + current_gid_string = gidBytesToString(gid_.raw); + current_lid = lid_; + current_context = context_; + current_port = port_; + auto_gid_selection_enabled = auto_gid_selection_enabled_; + } + + if (auto_gid_selection_enabled) { + ibv_port_attr port_attr; + int ret = ibv_query_port(current_context, current_port, &port_attr); + if (ret) { + LOG(WARNING) << "Failed to refresh port attributes on " + << device_name_ << "/" + << static_cast(current_port) << ": " + << strerror(ret); + return GidRefreshResult::FAILED; + } + + std::vector candidates; + candidates.reserve(port_attr.gid_tbl_len); + for (int i = 0; i < port_attr.gid_tbl_len; ++i) { + AutoGidCandidate candidate; + candidate.gid_index = i; + + struct ibv_gid_entry gid_entry; + if (ibv_query_gid_ex(current_context, current_port, i, &gid_entry, + 0)) { + candidate.query_succeeded = false; + candidates.push_back(candidate); + continue; + } + + const auto *gid_addr = + reinterpret_cast(gid_entry.gid.raw); + std::string ndev = readGidNdev(device_name_, current_port, i); + candidate.gid = gidBytesToString(gid_entry.gid.raw); + candidate.gid_type = gid_entry.gid_type; + candidate.has_network_device = !ndev.empty(); + candidate.is_ipv4_mapped = ipv6_addr_v4mapped(gid_addr); + candidate.is_link_local_ipv6 = isLinkLocalIpv6(gid_addr); + candidate.is_overlay_network = + candidate.has_network_device && isOverlayNetwork(ndev); + candidate.is_overlay_ipv4 = + candidate.is_ipv4_mapped && isOverlayIPv4(gid_addr); + candidate.is_null_gid = isNullGid(&gid_entry.gid); + candidates.push_back(candidate); + } + + auto selection = selectBestAutoGidCandidate(candidates); + if (!selection.has_value()) { + LOG(WARNING) << "No suitable GID found while refreshing " + << device_name_ << "/" + << static_cast(current_port); + return GidRefreshResult::FAILED; + } + next_gid_index = selection->gid_index; + } else { + next_gid_index = current_gid_index; + } + + ibv_gid new_gid = {}; + std::string next_gid_string; + if (ibv_query_gid(current_context, current_port, next_gid_index, + &new_gid)) { + return GidRefreshResult::FAILED; + } + if (isNullGid(&new_gid)) { + return GidRefreshResult::FAILED; + } + next_gid_string = gidBytesToString(new_gid.raw); + + if (next_gid_index == current_gid_index && + next_gid_string == current_gid_string) { + if (next_gid) *next_gid = current_gid_string; + return GidRefreshResult::UNCHANGED; + } + + int publish_ret = engine_.refreshLocalDeviceDesc(device_name_, current_lid, + next_gid_string); + if (publish_ret) { + LOG(ERROR) << "Failed to refresh local device descriptor for " + << device_name_ << ": " << publish_ret; + return GidRefreshResult::FAILED; + } + + { + std::lock_guard guard(gid_lock_); + gid_ = new_gid; + gid_index_ = next_gid_index; + } + if (previous_gid) *previous_gid = current_gid_string; + if (next_gid) *next_gid = next_gid_string; + + LOG(WARNING) << "Refreshed GID on " << device_name_ << "/" + << static_cast(port_) << ": index " << current_gid_index + << " (" << current_gid_string << ") -> " << next_gid_index + << " (" << next_gid_string << ")"; + return GidRefreshResult::CHANGED; +} + int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, int gid_index) { int num_devices = 0; @@ -966,10 +1213,12 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, ibv_port_attr attr; int ret = ibv_query_port(context, port, &attr); if (ret) { - PLOG(ERROR) << "Failed to query port " << port << " on " - << device_name; - if (ibv_close_device(context)) { - PLOG(ERROR) << "ibv_close_device(" << device_name << ") failed"; + LOG(ERROR) << "Failed to query port " << port << " on " + << device_name << ": " << strerror(ret); + int close_ret = ibv_close_device(context); + if (close_ret) { + LOG(ERROR) << "ibv_close_device(" << device_name + << ") failed: " << strerror(close_ret); } ibv_free_device_list(devices); return ERR_CONTEXT; @@ -977,8 +1226,10 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, if (attr.state != IBV_PORT_ACTIVE) { LOG(WARNING) << "Device " << device_name << " port not active"; - if (ibv_close_device(context)) { - PLOG(ERROR) << "ibv_close_device(" << device_name << ") failed"; + int close_ret = ibv_close_device(context); + if (close_ret) { + LOG(ERROR) << "ibv_close_device(" << device_name + << ") failed: " << strerror(close_ret); } ibv_free_device_list(devices); return ERR_CONTEXT; @@ -987,15 +1238,18 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, ibv_device_attr device_attr; ret = ibv_query_device(context, &device_attr); if (ret) { - PLOG(WARNING) << "Failed to query attributes on " << device_name; - if (ibv_close_device(context)) { - PLOG(ERROR) << "ibv_close_device(" << device_name << ") failed"; + LOG(WARNING) << "Failed to query attributes on " << device_name + << ": " << strerror(ret); + int close_ret = ibv_close_device(context); + if (close_ret) { + LOG(ERROR) << "ibv_close_device(" << device_name + << ") failed: " << strerror(close_ret); } ibv_free_device_list(devices); return ERR_CONTEXT; } -#if defined(USE_MACA) || defined(USE_CUDA) +#if defined(USE_MACA) || defined(USE_CUDA) || defined(USE_SUPA) // Verify DMA-BUF support against the GPU device(s) that the local // topology explicitly maps to this RNIC, rather than assuming the // verbs enumeration order matches GPU enumeration. @@ -1043,7 +1297,7 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, } else { // cuInit is process-global and idempotent; call it once before // the per-device loop, not per cuDeviceGet. -#if defined(USE_CUDA) +#if defined(USE_CUDA) || defined(USE_SUPA) CUresult result = cuInit(0); if (result != CUDA_SUCCESS) { LOG(ERROR) << "Failed to initialize CUDA driver for RNIC " @@ -1086,10 +1340,12 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, ibv_port_attr port_attr; ret = ibv_query_port(context, port, &port_attr); if (ret) { - PLOG(WARNING) << "Failed to query port attributes on " - << device_name << "/" << port; - if (ibv_close_device(context)) { - PLOG(ERROR) << "ibv_close_device(" << device_name << ") failed"; + LOG(WARNING) << "Failed to query port attributes on " << device_name + << "/" << port << ": " << strerror(ret); + int close_ret = ibv_close_device(context); + if (close_ret) { + LOG(ERROR) << "ibv_close_device(" << device_name + << ") failed: " << strerror(close_ret); } ibv_free_device_list(devices); return ERR_CONTEXT; @@ -1130,8 +1386,8 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, // Continue with GID validation ret = ibv_query_gid(context, port, gid_index, &gid_); if (ret) { - PLOG(ERROR) << "Failed to query GID " << gid_index << " on " - << device_name << "/" << port; + LOG(ERROR) << "Failed to query GID " << gid_index << " on " + << device_name << "/" << port << ": " << strerror(ret); goto cleanup_context_and_devices; } @@ -1149,6 +1405,7 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, lid_ = attr.lid; active_mtu_ = attr.active_mtu; active_speed_ = attr.active_speed; + active_width_ = attr.active_width; { std::lock_guard guard(gid_lock_); gid_index_ = gid_index; @@ -1157,10 +1414,13 @@ int RdmaContext::openRdmaDevice(const std::string &device_name, uint8_t port, ibv_free_device_list(devices); return 0; - cleanup_context_and_devices: - if (ibv_close_device(context)) { - PLOG(ERROR) << "ibv_close_device(" << device_name << ") failed"; + cleanup_context_and_devices: { + int close_ret = ibv_close_device(context); + if (close_ret) { + LOG(ERROR) << "ibv_close_device(" << device_name + << ") failed: " << strerror(close_ret); } + } ibv_free_device_list(devices); return ERR_CONTEXT; } @@ -1208,4 +1468,16 @@ int RdmaContext::submitPostSend( const std::vector &slice_list) { return worker_pool_->submitPostSend(slice_list); } + +void RdmaContext::trackPostedSlices( + const std::vector &slice_list, size_t first, + size_t count) { + worker_pool_->trackPostedSlices(slice_list, first, count); +} + +void RdmaContext::untrackPostedSlices( + const std::vector &slice_list, size_t first, + size_t count) { + worker_pool_->untrackPostedSlices(slice_list, first, count); +} } // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp index 29882865f0..04baa0be2a 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_endpoint.cpp @@ -46,6 +46,8 @@ static GidSelectionSnapshot fillLocalHandshakeDesc( local_desc.local_gid = gid_selection.gid; local_desc.peer_nic_path = peer_nic; local_desc.qp_num = qp_num; + local_desc.ready_ack = false; + local_desc.ready_ack_supported = true; local_desc.reply_msg.clear(); return gid_selection; } @@ -64,9 +66,34 @@ static void rememberAutoGidSelection( } } +static std::string qpListToString(const std::vector &qp_num) { + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < qp_num.size(); ++i) { + if (i) oss << ","; + oss << qp_num[i]; + } + oss << "]"; + return oss.str(); +} + +static std::string gidToString(const ibv_gid &gid) { + char buffer[48]; + snprintf(buffer, sizeof(buffer), + "%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x:" + "%02x:%02x:%02x:%02x:%02x:%02x:%02x:%02x", + gid.raw[0], gid.raw[1], gid.raw[2], gid.raw[3], gid.raw[4], + gid.raw[5], gid.raw[6], gid.raw[7], gid.raw[8], gid.raw[9], + gid.raw[10], gid.raw[11], gid.raw[12], gid.raw[13], gid.raw[14], + gid.raw[15]); + return buffer; +} + RdmaEndPoint::RdmaEndPoint(RdmaContext &context) : context_(context), status_(INITIALIZING), + has_connected_(false), + ready_wait_start_ts_(0), wr_depth_list_(nullptr), active_(true), cq_outstanding_(nullptr) {} @@ -90,19 +117,19 @@ int RdmaEndPoint::construct(ibv_cq *cq, size_t num_qp_list, } qp_list_.resize(num_qp_list); - cq_outstanding_ = (volatile int *)cq->cq_context; + cq_outstanding_ = static_cast *>(cq->cq_context); max_wr_depth_ = (int)max_wr_depth; max_sge_per_wr_ = max_sge_per_wr; max_inline_bytes_ = max_inline_bytes; - wr_depth_list_ = new volatile int[num_qp_list](); + wr_depth_list_ = new std::atomic[num_qp_list](); if (!wr_depth_list_) { LOG(ERROR) << "Failed to allocate memory for work request depth list"; return ERR_MEMORY; } for (size_t i = 0; i < num_qp_list; ++i) { - wr_depth_list_[i] = 0; + wr_depth_list_[i].store(0, std::memory_order_relaxed); ibv_qp_init_attr attr; memset(&attr, 0, sizeof(attr)); attr.send_cq = cq; @@ -120,6 +147,7 @@ int RdmaEndPoint::construct(ibv_cq *cq, size_t num_qp_list, } } + ready_wait_start_ts_.store(0, std::memory_order_relaxed); status_.store(UNCONNECTED, std::memory_order_relaxed); return 0; } @@ -148,7 +176,8 @@ int RdmaEndPoint::reconstruct() { // Reconstruct with same parameters as original construction status_.store(INITIALIZING, std::memory_order_relaxed); - active_ = true; + ready_wait_start_ts_.store(0, std::memory_order_relaxed); + active_.store(true, std::memory_order_release); return construct(cq, num_qp, max_sge_per_wr, max_wr_depth, max_inline_bytes); @@ -165,15 +194,16 @@ int RdmaEndPoint::deconstructLocked() { bool displayed = false; if (wr_depth_list_) { for (size_t i = 0; i < qp_list_.size(); ++i) { - if (wr_depth_list_[i] != 0) { + int wr_depth = wr_depth_list_[i].load(std::memory_order_relaxed); + if (wr_depth != 0) { if (!displayed) { LOG(WARNING) << "Outstanding work requests found, CQ will not " "be generated"; displayed = true; } - __sync_fetch_and_sub(cq_outstanding_, wr_depth_list_[i]); - wr_depth_list_[i] = 0; + cq_outstanding_->fetch_sub(wr_depth, std::memory_order_acq_rel); + wr_depth_list_[i].store(0, std::memory_order_relaxed); } } } @@ -181,8 +211,9 @@ int RdmaEndPoint::deconstructLocked() { int result = 0; for (size_t i = 0; i < qp_list_.size(); ++i) { if (!qp_list_[i]) continue; // already destroyed in a previous call - if (ibv_destroy_qp(qp_list_[i])) { - PLOG(ERROR) << "Failed to destroy QP[" << i << "]"; + int ret = ibv_destroy_qp(qp_list_[i]); + if (ret) { + LOG(ERROR) << "Failed to destroy QP[" << i << "]: " << strerror(ret); result = ERR_ENDPOINT; } else { qp_list_[i] = nullptr; @@ -202,21 +233,35 @@ int RdmaEndPoint::destroyQP() { return deconstruct(); } void RdmaEndPoint::beginDestroy() { RWSpinlock::WriteGuard guard(lock_); + beginDestroyLocked(); +} + +void RdmaEndPoint::beginDestroyLocked() { auto current_status = status_.load(std::memory_order_relaxed); if (current_status == DESTROYING || current_status == DESTROYED) return; - active_ = false; - inactive_time_ = getCurrentTimeInNano(); + inactive_time_.store(getCurrentTimeInNano(), std::memory_order_relaxed); + active_.store(false, std::memory_order_release); status_.store(DESTROYING, std::memory_order_release); + ready_wait_start_ts_.store(0, std::memory_order_relaxed); - // Transition QPs to ERR state so hardware flushes all inflight WRs to CQ. - // This allows performPollCq to drain them naturally. - ibv_qp_attr attr; - memset(&attr, 0, sizeof(attr)); - attr.qp_state = IBV_QPS_ERR; + // Only endpoints that reached CONNECTED can have user WRs to flush. For + // pre-connected endpoints, skip RESET/INIT -> ERR because there are no WRs + // and some providers reject that state transition. + if (!has_connected_) return; + + // Transition connected QPs to ERR state so hardware flushes inflight WRs to + // CQ. This allows performPollCq to drain them naturally. for (size_t i = 0; i < qp_list_.size(); ++i) { - if (ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE)) { - PLOG(WARNING) << "Failed to modify QP to ERR during beginDestroy"; + if (!qp_list_[i]) continue; + + ibv_qp_attr attr; + memset(&attr, 0, sizeof(attr)); + attr.qp_state = IBV_QPS_ERR; + int ret = ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE); + if (ret) { + LOG(WARNING) << "Failed to modify QP[" << i + << "] to ERR during beginDestroy: " << strerror(ret); } } } @@ -235,7 +280,7 @@ bool RdmaEndPoint::finishDestroy() { // pre-two-phase predicate (!hasOutstandingSlice == !active_): only // inactive endpoints are eligible for reclaim; active ones must stay. if (current_status != DESTROYING) { - if (active_) return false; + if (active_.load(std::memory_order_acquire)) return false; // Endpoints that never reached construct() own no RDMA resources // and have wr_depth_list_ uninitialized; deconstructLocked() would // delete[] a wild pointer. Drop them directly. @@ -253,13 +298,15 @@ bool RdmaEndPoint::finishDestroy() { // never be flushed; enforce a timeout to avoid leaking forever. bool has_outstanding = false; for (size_t i = 0; i < qp_list_.size(); ++i) { - if (wr_depth_list_[i] != 0) { + if (wr_depth_list_[i].load(std::memory_order_relaxed) != 0) { has_outstanding = true; break; } } if (has_outstanding) { - double elapsed = (getCurrentTimeInNano() - inactive_time_) / 1e9; + double elapsed = (getCurrentTimeInNano() - + inactive_time_.load(std::memory_order_relaxed)) / + 1e9; if (elapsed < kFinishDestroyTimeoutSec) return false; LOG(WARNING) << "finishDestroy timed out after " << elapsed << "s with outstanding WRs, forcing destruction"; @@ -285,9 +332,12 @@ bool RdmaEndPoint::finishDestroy() { void RdmaEndPoint::setPeerNicPath(const std::string &peer_nic_path) { RWSpinlock::WriteGuard guard(lock_); - if (connected()) { - LOG(WARNING) << "Previous connection will be discarded"; - disconnectUnlocked(); + auto curr_status = status_.load(std::memory_order_relaxed); + if (curr_status != INITIALIZING && curr_status != UNCONNECTED) { + LOG(ERROR) << "Cannot change peer NIC path after endpoint lifecycle " + "has started: " + << toString(); + return; } peer_nic_path_ = peer_nic_path; } @@ -301,19 +351,22 @@ int RdmaEndPoint::setupConnectionsByActive() { { RWSpinlock::WriteGuard guard(lock_); - if (connected()) { + if (readyToSend()) { LOG(INFO) << "Connection has been established"; return 0; } // loopback mode if (context_.nicPath() == peer_nic_path_) { - return doSetupConnection(context_.gid(), context_.lid(), qpNum()); + int ret = + doSetupConnection(context_.gid(), context_.lid(), qpNum()); + if (ret == 0) { + ready_wait_start_ts_.store(0, std::memory_order_relaxed); + } + return ret; } - // Only proceed with RPC if we are the first to transition from - // UNCONNECTED. This prevents duplicate concurrent handshake attempts - // from the same endpoint. + // Only the first UNCONNECTED caller transitions to CONNECTING. auto current_status = status_.load(std::memory_order_relaxed); if (current_status == UNCONNECTED) { status_.store(CONNECTING, std::memory_order_relaxed); @@ -359,6 +412,7 @@ int RdmaEndPoint::setupConnectionsByActive() { } } RWSpinlock::ReadGuard guard(lock_); + if (readyToSend()) return 0; return connected() ? 0 : ERR_ENDPOINT; } @@ -383,139 +437,209 @@ int RdmaEndPoint::setupConnectionsByActive() { // `peer_qp_num_list_` with `peer_desc.qp_num`, since a failed RPC may // result in an invalid `peer_desc.qp_num`. // - // If the RPC is failed, even if the state is CONNECTED, (which means - // it is handled by setupConnectionsByPassive in another thread during - // the RPC, or "simultaneous open"), we should resetConnection to be - // safe. Because we're not sure whether the peer needs a connection - // re-establishment. (We don't know `peer_desc.qp_num`) + // If the RPC failed but simultaneous-open passive setup has already + // made this endpoint CONNECTED with the same local QPs used by this + // RPC, reuse that connection. Otherwise reset because + // `peer_desc.qp_num` is invalid and we cannot safely infer the peer + // state. if (rc) { RWSpinlock::WriteGuard write_guard(lock_); + if (connected()) { + auto current_qp_num = qpNum(); + if (current_qp_num == local_desc.qp_num) { + LOG(WARNING) + << "Active handshake RPC failed, but simultaneous-open " + "passive setup already connected this endpoint. " + "Reusing existing connection. rc=" + << rc << ", local_desc.qp_num=" + << qpListToString(local_desc.qp_num) + << ", current_qp_num=" << qpListToString(current_qp_num) + << ", endpoint=" << toString(); + return 0; + } + } + + LOG(ERROR) << "Active handshake RPC failed; resetting endpoint. rc=" + << rc << ", local_desc.qp_num=" + << qpListToString(local_desc.qp_num) + << ", endpoint=" << toString(); resetConnection("handshake RPC failure"); return rc; } bool retry_with_new_gid = false; + bool should_send_ready_ack = false; + HandShakeDesc ready_ack_desc; { // Re-acquire lock after RPC to finalize state transition RWSpinlock::WriteGuard guard(lock_); // Handle simultaneous open: if the peer initiates a connection // during our RPC and it is passively established in - // setupConnectionsByPassive, simply reuse the existing endpoint. + // setupConnectionsByPassive, send an explicit ready ACK after this + // active RPC confirms that the peer's passive QPs are ready. if (connected()) { if (peer_qp_num_list_ == peer_desc.qp_num) { - LOG(INFO) - << "Received same peer QP numbers, reusing connection."; - return 0; - } + if (peer_desc.ready_ack_supported) { + should_send_ready_ack = true; + ready_ack_desc = local_desc; + LOG(INFO) << "Received same peer QP numbers, sending " + "RDMA ready ACK."; + } else { + ready_wait_start_ts_.store(0, + std::memory_order_relaxed); + status_.store(CONNECTED, std::memory_order_relaxed); + LOG(INFO) << "Peer does not support RDMA ready ACK, " + "reusing connection."; + return 0; + } + } else { + // This mismatch scenario should be rare. It may occur when + // a peer first sends us an Active RPC and establishes a + // connection, then restarts, and eventually accepts and + // responds to our Active RPC. + LOG(WARNING) + << "Peer QP list mismatch on connected endpoint, " + "re-establishing connection: " + << toString(); - // This mismatch scenario should be rare. It may occur when a - // peer first sends us an Active RPC and establishes a - // connection, then restarts, and eventually accepts and - // responds to our Active RPC. - LOG(WARNING) << "Peer QP list mismatch on connected endpoint, " - "re-establishing connection: " - << toString(); - - int ret = - resetConnection("re-establishing connection (active)"); - if (ret) return ret; + int ret = + resetConnection("re-establishing connection (active)"); + if (ret) return ret; + } } - if (!peer_desc.reply_msg.empty()) { - LOG(ERROR) << "Rejected handshake request by peer " - << local_desc.peer_nic_path; - disconnectUnlocked(); - return ERR_REJECT_HANDSHAKE; - } + if (!should_send_ready_ack) { + if (!peer_desc.reply_msg.empty()) { + LOG(ERROR) << "Rejected handshake request by peer " + << local_desc.peer_nic_path; + disconnectUnlocked(); + return ERR_REJECT_HANDSHAKE; + } - if (peer_desc.local_nic_path != peer_nic_path_ || - peer_desc.peer_nic_path != local_desc.local_nic_path) { - LOG(ERROR) << "Invalid argument: received packet mismatch, " - "local.local_nic_path: " - << local_desc.local_nic_path - << ", local.peer_nic_path: " - << local_desc.peer_nic_path - << ", peer.local_nic_path: " - << peer_desc.local_nic_path - << ", peer.peer_nic_path: " - << peer_desc.peer_nic_path; - disconnectUnlocked(); - return ERR_REJECT_HANDSHAKE; - } + if (peer_desc.local_nic_path != peer_nic_path_ || + peer_desc.peer_nic_path != local_desc.local_nic_path) { + LOG(ERROR) + << "Invalid argument: received packet mismatch, " + "local.local_nic_path: " + << local_desc.local_nic_path + << ", local.peer_nic_path: " << local_desc.peer_nic_path + << ", peer.local_nic_path: " << peer_desc.local_nic_path + << ", peer.peer_nic_path: " << peer_desc.peer_nic_path; + disconnectUnlocked(); + return ERR_REJECT_HANDSHAKE; + } - int ret = ERR_DEVICE_NOT_FOUND; - std::string failure_message; - SetupConnectionFailureInfo failure_info; - if (!peer_desc.local_gid.empty()) { - ret = doSetupConnection(peer_desc.local_gid, - peer_desc.local_lid, peer_desc.qp_num, - &failure_message, &failure_info); - } else { - auto segment_desc = - context_.engine().meta()->getSegmentDescByName( - peer_server_name); - if (segment_desc) { - for (auto &nic : segment_desc->devices) { - if (nic.name == peer_nic_name) { - ret = doSetupConnection( - nic.gid, nic.lid, peer_desc.qp_num, - &failure_message, &failure_info); - break; + int ret = ERR_DEVICE_NOT_FOUND; + std::string failure_message; + SetupConnectionFailureInfo failure_info; + auto connected_status = peer_desc.ready_ack_supported + ? CONNECTED_WAIT_READY_ACK + : CONNECTED; + if (!peer_desc.local_gid.empty()) { + ret = doSetupConnection(peer_desc.local_gid, + peer_desc.local_lid, + peer_desc.qp_num, connected_status, + &failure_message, &failure_info); + } else { + auto segment_desc = + context_.engine().meta()->getSegmentDescByName( + peer_server_name); + if (segment_desc) { + for (auto &nic : segment_desc->devices) { + if (nic.name == peer_nic_name) { + ret = doSetupConnection( + nic.gid, nic.lid, peer_desc.qp_num, + connected_status, &failure_message, + &failure_info); + break; + } } } } - } - if (ret == 0) { - return 0; - } + if (ret == 0) { + if (peer_desc.ready_ack_supported) { + should_send_ready_ack = true; + ready_ack_desc = local_desc; + } else { + ready_wait_start_ts_.store(0, + std::memory_order_relaxed); + status_.store(CONNECTED, std::memory_order_relaxed); + return 0; + } + } else { + if (shouldAttemptAutoGidHandshakeRetry( + context_.autoGidSelectionEnabled(), + auto_gid_retry_count, + globalConfig().auto_gid_max_retries, + failure_info.stage == + SetupConnectionFailureStage::kRtr, + failure_info.sys_errno)) { + std::string previous_gid; + std::string next_gid; + bool reprobe_changed = context_.reprobeAutoGid( + local_gid_selection, attempted_auto_gid_selections, + &previous_gid, &next_gid); + auto current_gid_selection = context_.gidSelection(); + auto retry_action = decideAutoGidRetryAction( + reprobe_changed, local_gid_selection.gid_index, + local_gid_selection.gid, + current_gid_selection.gid_index, + current_gid_selection.gid); + if (retry_action != AutoGidRetryAction::kDoNotRetry) { + int reset_ret = resetConnection( + retry_action == AutoGidRetryAction:: + kRetryWithReprobedGid + ? "retry after auto GID reprobe (active)" + : "retry with externally reprobed GID " + "(active)"); + if (reset_ret) return reset_ret; + status_.store(CONNECTING, + std::memory_order_relaxed); + ++auto_gid_retry_count; + retry_with_new_gid = true; + LOG(WARNING) + << "Retry active handshake with updated local " + "GID on " + << context_.deviceName() << ": " + << local_gid_selection.gid << " -> " + << current_gid_selection.gid << " (attempt " + << auto_gid_retry_count << "/" + << globalConfig().auto_gid_max_retries << ")"; + } + } - if (shouldAttemptAutoGidHandshakeRetry( - context_.autoGidSelectionEnabled(), auto_gid_retry_count, - globalConfig().auto_gid_max_retries, - failure_info.stage == SetupConnectionFailureStage::kRtr, - failure_info.sys_errno)) { - std::string previous_gid; - std::string next_gid; - bool reprobe_changed = context_.reprobeAutoGid( - local_gid_selection, attempted_auto_gid_selections, - &previous_gid, &next_gid); - auto current_gid_selection = context_.gidSelection(); - auto retry_action = decideAutoGidRetryAction( - reprobe_changed, local_gid_selection.gid_index, - local_gid_selection.gid, current_gid_selection.gid_index, - current_gid_selection.gid); - if (retry_action != AutoGidRetryAction::kDoNotRetry) { - int reset_ret = resetConnection( - retry_action == - AutoGidRetryAction::kRetryWithReprobedGid - ? "retry after auto GID reprobe (active)" - : "retry with externally reprobed GID (active)"); - if (reset_ret) return reset_ret; - status_.store(CONNECTING, std::memory_order_relaxed); - ++auto_gid_retry_count; - retry_with_new_gid = true; - LOG(WARNING) - << "Retry active handshake with updated local GID on " - << context_.deviceName() << ": " - << local_gid_selection.gid << " -> " - << current_gid_selection.gid << " (attempt " - << auto_gid_retry_count << "/" - << globalConfig().auto_gid_max_retries << ")"; + if (!retry_with_new_gid) { + if (ret == ERR_DEVICE_NOT_FOUND) { + LOG(ERROR) << "Peer NIC " << peer_nic_name + << " not found in " << peer_server_name; + disconnectUnlocked(); + } else { + resetConnection("failed connection setup (active)"); + } + return ret; + } } } + } - if (!retry_with_new_gid) { - if (ret == ERR_DEVICE_NOT_FOUND) { - LOG(ERROR) << "Peer NIC " << peer_nic_name - << " not found in " << peer_server_name; - disconnectUnlocked(); - } else { - resetConnection("failed connection setup (active)"); - } - return ret; + if (should_send_ready_ack) { + int ack_ret = sendReadyAck(peer_server_name, ready_ack_desc); + RWSpinlock::WriteGuard guard(lock_); + if (ack_ret) { + resetConnection("failed to send ready ACK"); + return ack_ret; + } + if (!connected()) { + LOG(WARNING) << "Discarding RDMA ready ACK because endpoint " + << "is no longer connected: " << toString(); + ready_wait_start_ts_.store(0, std::memory_order_relaxed); + return ERR_ENDPOINT; } + ready_wait_start_ts_.store(0, std::memory_order_relaxed); + status_.store(CONNECTED, std::memory_order_relaxed); + return 0; } } } @@ -523,12 +647,45 @@ int RdmaEndPoint::setupConnectionsByActive() { int RdmaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc, HandShakeDesc &local_desc) { RWSpinlock::WriteGuard guard(lock_); + if (peer_desc.ready_ack) { + if (!connected()) { + local_desc.reply_msg = + "Received RDMA ready ACK for unconnected endpoint"; + LOG(ERROR) << local_desc.reply_msg << ": " << toString(); + return ERR_REJECT_HANDSHAKE; + } + + if (peer_qp_num_list_ != peer_desc.qp_num) { + local_desc.reply_msg = + "Received stale RDMA ready ACK with mismatched peer QP numbers"; + LOG(WARNING) << local_desc.reply_msg << ", ack_peer_qp_num=" + << qpListToString(peer_desc.qp_num) + << ", current_peer_qp_num=" + << qpListToString(peer_qp_num_list_) << ": " + << toString(); + return ERR_REJECT_HANDSHAKE; + } + + ready_wait_start_ts_.store(0, std::memory_order_relaxed); + status_.store(CONNECTED, std::memory_order_relaxed); + LOG(INFO) << "Received RDMA ready ACK."; + return 0; + } + if (connected()) { // If already connected with the same peer QP info, return success if (peer_qp_num_list_ == peer_desc.qp_num) { fillLocalHandshakeDesc(context_, peer_nic_path_, qpNum(), local_desc); - LOG(INFO) << "Received same peer QP numbers, reusing connection."; + if (!peer_desc.ready_ack_supported) { + ready_wait_start_ts_.store(0, std::memory_order_relaxed); + status_.store(CONNECTED, std::memory_order_relaxed); + LOG(INFO) << "Peer does not support RDMA ready ACK, " + "reusing connection."; + } else { + LOG(INFO) << "Received same peer QP numbers, reusing " + "connection while waiting for ready ACK."; + } return 0; } // Different peer (e.g., peer restarted) @@ -543,9 +700,8 @@ int RdmaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc, // establish the connection on this same endpoint. Because we're holding // the lock, even if there are already Active RPCs sent to the same // peer nic path by setupConnectionsByActive on another thread, it will - // be blocked after the RPC return. Once the lock is released, - // they will simply observe the CONNECTED state and safely reuse the QP. - // This inherently handles simultaneous open. + // be blocked after the RPC return. Once the lock is released, active + // callers will confirm readiness before posting WRs. if (peer_desc.peer_nic_path != context_.nicPath() || peer_desc.local_nic_path != peer_nic_path_) { @@ -579,9 +735,22 @@ int RdmaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc, local_gid_selection); SetupConnectionFailureInfo failure_info; + auto connected_status = peer_desc.ready_ack_supported + ? CONNECTED_WAIT_READY_ACK + : CONNECTED; int ret = doSetupConnection(peer_gid, peer_lid, peer_desc.qp_num, - &local_desc.reply_msg, &failure_info); + connected_status, &local_desc.reply_msg, + &failure_info); if (ret == 0) { + if (peer_desc.ready_ack_supported) { + ready_wait_start_ts_.store(getCurrentTimeInNano(), + std::memory_order_relaxed); + status_.store(CONNECTED_WAIT_READY_ACK, + std::memory_order_relaxed); + } else { + ready_wait_start_ts_.store(0, std::memory_order_relaxed); + status_.store(CONNECTED, std::memory_order_relaxed); + } return 0; } @@ -641,6 +810,7 @@ int RdmaEndPoint::setupConnectionsByPassive(const HandShakeDesc &peer_desc, } local_desc.reply_msg = "Peer nic not found in that server: " + peer_nic_path_; + ready_wait_start_ts_.store(0, std::memory_order_relaxed); status_.store(UNCONNECTED, std::memory_order_relaxed); LOG(ERROR) << local_desc.reply_msg; return ERR_DEVICE_NOT_FOUND; @@ -653,53 +823,76 @@ void RdmaEndPoint::disconnect() { int RdmaEndPoint::disconnectUnlocked() { auto curr_status = status_.load(std::memory_order_acquire); - if (curr_status != CONNECTED && curr_status != CONNECTING) return 0; - - ibv_qp_attr attr; - memset(&attr, 0, sizeof(attr)); - attr.qp_state = IBV_QPS_RESET; - int ret = 0; - for (size_t i = 0; i < qp_list_.size(); ++i) { - int curr_ret = ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE); - if (curr_ret) { - PLOG(ERROR) << "Failed to modify QP to RESET"; - ret = ERR_ENDPOINT; - } - // After resetting QP, the wr_depth_list_ won't change - bool displayed = false; - if (wr_depth_list_[i] != 0) { - if (!displayed) { - LOG(WARNING) << "Outstanding work requests found, CQ will not " - "be generated"; - displayed = true; - } - __sync_fetch_and_sub(cq_outstanding_, wr_depth_list_[i]); - wr_depth_list_[i] = 0; + if (!isConnectedStatus(curr_status) && curr_status != CONNECTING) return 0; + ready_wait_start_ts_.store(0, std::memory_order_relaxed); + + if (!has_connected_) { + // This endpoint has not posted user WRs yet, but its handshake may + // already have reached the peer. The peer can cache our QP numbers and + // rebuild its passive endpoint around them before the active side sees + // a setup failure. Reusing the same local QPs after RESET would make a + // later retry ambiguous with that partially processed handshake, so + // retry with fresh QP numbers instead. + for (size_t i = 0; i < qp_list_.size(); ++i) { + CHECK_EQ(wr_depth_list_[i].load(std::memory_order_relaxed), 0) + << "Pre-connected endpoint must not have outstanding WRs"; } + return reconstruct(); } - peer_qp_num_list_.clear(); - status_.store(UNCONNECTED, std::memory_order_release); - return ret; + + beginDestroyLocked(); + return 0; } int RdmaEndPoint::resetConnection(const std::string &reason) { auto curr_status = status_.load(std::memory_order_acquire); - if (curr_status != CONNECTING && curr_status != CONNECTED) return 0; + if (curr_status != CONNECTING && !isConnectedStatus(curr_status)) return 0; + ready_wait_start_ts_.store(0, std::memory_order_relaxed); + + if (!has_connected_) { + int ret = disconnectUnlocked(); + if (ret) { + LOG(ERROR) << "Failed to reset pre-connected endpoint " + << "(triggered by: " << reason << "): error=" << ret; + } else { + LOG(INFO) << "Successfully reset pre-connected endpoint " + << "(triggered by: " << reason << ")."; + } + return ret; + } -#ifdef CONFIG_ERDMA - int ret = reconstruct(); -#else - int ret = disconnectUnlocked(); -#endif + LOG(WARNING) << "Retiring endpoint instead of resetting it (triggered by: " + << reason << "): " << toString(); + beginDestroyLocked(); + return ERR_ENDPOINT; +} - if (ret) { - LOG(ERROR) << "Failed to reset the endpoint (triggered by: " << reason - << "): error=" << ret; - } else { - LOG(INFO) << "Successfully reset the endpoint (triggered by: " << reason - << ")."; +bool RdmaEndPoint::readyAckTimedOut() const { + if (status() != CONNECTED_WAIT_READY_ACK) return false; + uint64_t start_ts = ready_wait_start_ts_.load(std::memory_order_relaxed); + return start_ts != 0 && + getCurrentTimeInNano() - start_ts > kReadyAckTimeoutNano; +} + +int RdmaEndPoint::sendReadyAck(const std::string &peer_server_name, + const HandShakeDesc &local_desc) { + HandShakeDesc ready_ack_desc = local_desc; + ready_ack_desc.ready_ack = true; + + HandShakeDesc peer_desc; + int rc = context_.engine().sendHandshake(peer_server_name, ready_ack_desc, + peer_desc); + if (rc) { + LOG(ERROR) << "Failed to send RDMA ready ACK to " << peer_server_name + << ": " << rc; + return rc; } - return ret; + if (!peer_desc.reply_msg.empty()) { + LOG(ERROR) << "RDMA ready ACK rejected by " << peer_server_name << ": " + << peer_desc.reply_msg; + return ERR_REJECT_HANDSHAKE; + } + return 0; } const std::string RdmaEndPoint::toString() const { @@ -707,6 +900,9 @@ const std::string RdmaEndPoint::toString() const { if (status == CONNECTED) return "EndPoint: local " + context_.nicPath() + ", peer " + peer_nic_path_; + else if (status == CONNECTED_WAIT_READY_ACK) + return "EndPoint: local " + context_.nicPath() + ", peer " + + peer_nic_path_ + " (waiting ready ACK)"; else if (status == DESTROYING) return "EndPoint: local " + context_.nicPath() + ", peer " + peer_nic_path_ + " (destroying)"; @@ -720,7 +916,8 @@ int RdmaEndPoint::submitPostSend( std::vector &slice_list, std::vector &failed_slice_list) { RWSpinlock::WriteGuard guard(lock_); - if (!active_ || status_.load(std::memory_order_relaxed) != CONNECTED) { + if (!active_.load(std::memory_order_acquire) || + status_.load(std::memory_order_relaxed) != CONNECTED) { for (auto &slice : slice_list) failed_slice_list.push_back(slice); slice_list.clear(); return 0; @@ -729,16 +926,25 @@ int RdmaEndPoint::submitPostSend( const size_t num_qp = qp_list_.size(); if (slice_list.empty()) return 0; const size_t requested = slice_list.size(); - int cq_remaining = int(globalConfig().max_cqe) - *cq_outstanding_; - std::vector wr_list(requested, ibv_send_wr{}); - std::vector sge_list(requested); + int cq_remaining = int(globalConfig().max_cqe) - + cq_outstanding_->load(std::memory_order_relaxed); + if (cq_remaining <= 0) return 0; + + // Only allocate for the max number of WRs we can actually post per QP, + // not the entire requested slice count. Each QP iteration reuses the + // wr_list/sge_list from index 0, so we only need max_wr_depth_ entries. + size_t max_postable_per_qp = + std::min({(size_t)max_wr_depth_, (size_t)cq_remaining, requested}); + std::vector wr_list(max_postable_per_qp, ibv_send_wr{}); + std::vector sge_list(max_postable_per_qp); size_t total_posted = 0; size_t cursor = 0; for (size_t qp_index = 0; qp_index < num_qp && cq_remaining > 0 && cursor < requested; ++qp_index) { - int qp_avail = max_wr_depth_ - wr_depth_list_[qp_index]; + int qp_avail = max_wr_depth_ - + wr_depth_list_[qp_index].load(std::memory_order_relaxed); if (qp_avail <= 0) continue; size_t remaining_qps = num_qp - qp_index; @@ -776,16 +982,24 @@ int RdmaEndPoint::submitPostSend( } ibv_send_wr *bad_wr = nullptr; - __sync_fetch_and_add(&wr_depth_list_[qp_index], wr_count); - __sync_fetch_and_add(cq_outstanding_, wr_count); + wr_depth_list_[qp_index].fetch_add(wr_count, std::memory_order_acq_rel); + cq_outstanding_->fetch_add(wr_count, std::memory_order_acq_rel); + // Register before ringing the doorbell. A fast completion may otherwise + // be polled before the diagnostic registry sees the slice. + context_.trackPostedSlices(slice_list, start, wr_count); int rc = ibv_post_send(qp_list_[qp_index], wr_list.data(), &bad_wr); if (rc) { - PLOG(ERROR) << "Failed to ibv_post_send"; + LOG(ERROR) << "Failed to ibv_post_send: " << strerror(rc); + const size_t first_failed = + bad_wr ? static_cast(bad_wr - wr_list.data()) : 0; + context_.untrackPostedSlices(slice_list, start + first_failed, + wr_count - first_failed); while (bad_wr) { int i = bad_wr - wr_list.data(); failed_slice_list.push_back(slice_list[start + i]); - __sync_fetch_and_sub(&wr_depth_list_[qp_index], 1); - __sync_fetch_and_sub(cq_outstanding_, 1); + wr_depth_list_[qp_index].fetch_sub(1, + std::memory_order_acq_rel); + cq_outstanding_->fetch_sub(1, std::memory_order_acq_rel); bad_wr = bad_wr->next; } total_posted += wr_count; @@ -851,6 +1065,7 @@ static int parseGidString(const std::string &gid_str, ibv_gid &gid_out) { int RdmaEndPoint::doSetupConnection(const std::string &peer_gid, uint16_t peer_lid, std::vector peer_qp_num_list, + Status connected_status, std::string *reply_msg, SetupConnectionFailureInfo *failure_info) { if (qp_list_.size() != peer_qp_num_list.size()) { @@ -889,7 +1104,8 @@ int RdmaEndPoint::doSetupConnection(const std::string &peer_gid, } peer_qp_num_list_ = std::move(peer_qp_num_list); - status_.store(CONNECTED, std::memory_order_relaxed); + has_connected_ = true; + status_.store(connected_status, std::memory_order_relaxed); return 0; } @@ -908,11 +1124,11 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const ibv_gid &peer_gid, int ret = ibv_modify_qp(qp, &attr, IBV_QP_STATE); if (ret) { std::string message = "Failed to modify QP to RESET"; - PLOG(ERROR) << "[Handshake] " << message; - if (reply_msg) *reply_msg = message + ": " + strerror(errno); + LOG(ERROR) << "[Handshake] " << message << ": " << strerror(ret); + if (reply_msg) *reply_msg = message + ": " + strerror(ret); if (failure_info) { failure_info->stage = SetupConnectionFailureStage::kReset; - failure_info->sys_errno = errno; + failure_info->sys_errno = ret; } return ERR_ENDPOINT; } @@ -930,11 +1146,11 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const ibv_gid &peer_gid, if (ret) { std::string message = "Failed to modify QP to INIT, check local context port num"; - PLOG(ERROR) << "[Handshake] " << message; - if (reply_msg) *reply_msg = message + ": " + strerror(errno); + LOG(ERROR) << "[Handshake] " << message << ": " << strerror(ret); + if (reply_msg) *reply_msg = message + ": " + strerror(ret); if (failure_info) { failure_info->stage = SetupConnectionFailureStage::kInit; - failure_info->sys_errno = errno; + failure_info->sys_errno = ret; } return ERR_ENDPOINT; } @@ -955,7 +1171,11 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const ibv_gid &peer_gid, static_cast(globalConfig().ib_traffic_class); } attr.ah_attr.dlid = peer_lid; + // Set service level if configured (-1 means use default) attr.ah_attr.sl = 0; + if (globalConfig().ib_service_level >= 0) { + attr.ah_attr.sl = static_cast(globalConfig().ib_service_level); + } attr.ah_attr.src_path_bits = 0; attr.ah_attr.static_rate = 0; attr.ah_attr.is_global = 1; @@ -971,11 +1191,22 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const ibv_gid &peer_gid, if (ret) { std::string message = "Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num"; - PLOG(ERROR) << "[Handshake] " << message; - if (reply_msg) *reply_msg = message + ": " + strerror(errno); + LOG(ERROR) << "[Handshake] " << message + << ": local=" << context_.nicPath() + << ", peer=" << peer_nic_path_ << ", qp_index=" << qp_index + << ", local_qp=" << qp->qp_num + << ", peer_qp=" << peer_qp_num + << ", local_gid=" << context_.gid() + << ", local_gid_index=" << local_gid_index + << ", peer_gid=" << gidToString(peer_gid) + << ", peer_lid=" << peer_lid + << ", path_mtu=" << attr.path_mtu + << ", port_num=" << static_cast(context_.portNum()) + << ": " << strerror(ret); + if (reply_msg) *reply_msg = message + ": " + strerror(ret); if (failure_info) { failure_info->stage = SetupConnectionFailureStage::kRtr; - failure_info->sys_errno = errno; + failure_info->sys_errno = ret; } return ERR_ENDPOINT; } @@ -994,11 +1225,11 @@ int RdmaEndPoint::doSetupConnection(int qp_index, const ibv_gid &peer_gid, IBV_QP_MAX_QP_RD_ATOMIC); if (ret) { std::string message = "Failed to modify QP to RTS"; - PLOG(ERROR) << "[Handshake] " << message; - if (reply_msg) *reply_msg = message + ": " + strerror(errno); + LOG(ERROR) << "[Handshake] " << message << ": " << strerror(ret); + if (reply_msg) *reply_msg = message + ": " + strerror(ret); if (failure_info) { failure_info->stage = SetupConnectionFailureStage::kRts; - failure_info->sys_errno = errno; + failure_info->sys_errno = ret; } return ERR_ENDPOINT; } diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp index 6273a26218..e444b38b29 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -25,6 +26,7 @@ #include #include #include +#include #include @@ -41,8 +43,91 @@ namespace mooncake { static bool MCIbRelaxedOrderingEnabled = false; static int MCIbRelaxedOrderingMode = 2; +static std::string resolveBufferLocation( + const TransferMetadata::BufferDesc &buffer, uint64_t offset) { + std::string location = buffer.name; + SegmentsLocationInfo seg_info; + if (parseSegmentsLocation(buffer.name, seg_info)) { + location = resolveSegmentsLocation(seg_info, buffer.length, + offset - buffer.addr); + } + return location; +} + +// Calculate remaining bytes in buffer for 'address'. +// Caches 'last_buffer_idx' across sequential slices to avoid rescanning +// buffers. +static uint64_t bytesUntilBufferEnd( + const TransferMetadata::SegmentDesc *segment_desc, uint64_t address, + bool require_remote_key, size_t *last_buffer_idx = nullptr) { + if (!segment_desc || segment_desc->buffers.empty()) return 0; + + auto is_valid_buffer = [&](const TransferMetadata::BufferDesc &buffer) { +#ifdef ENABLE_MULTI_PROTOCOL + if (!buffer.protocol.empty() && buffer.protocol != "rdma") return false; +#endif + if (require_remote_key ? buffer.rkey.empty() : buffer.lkey.empty()) + return false; + if (address < buffer.addr || address - buffer.addr >= buffer.length) + return false; + return true; + }; + + if (last_buffer_idx && *last_buffer_idx < segment_desc->buffers.size()) { + const auto &buffer = segment_desc->buffers[*last_buffer_idx]; + if (is_valid_buffer(buffer)) { + return buffer.length - (address - buffer.addr); + } + } + + uint64_t max_remaining = 0; + for (size_t i = 0; i < segment_desc->buffers.size(); ++i) { + const auto &buffer = segment_desc->buffers[i]; + if (is_valid_buffer(buffer)) { + uint64_t remaining = buffer.length - (address - buffer.addr); + if (remaining > max_remaining) { + max_remaining = remaining; + if (last_buffer_idx) *last_buffer_idx = i; + } + } + } + return max_remaining; +} + +// Calculate slice length capped at MR boundary to avoid multi-MR WRs. +// Caches source and target buffer indices across slices for a single request. +struct SliceLengthCalculator { + const Transport::TransferRequest &request; + size_t block_size; + size_t fragment_size; + const TransferMetadata::SegmentDesc *local_desc; + const TransferMetadata::SegmentDesc *target_desc; + + size_t src_buffer_idx = 0; + size_t tgt_buffer_idx = 0; + + inline size_t calculate(uint64_t offset) { + const size_t remaining = request.length - offset; + size_t slice_length = + remaining <= block_size + fragment_size ? remaining : block_size; + + const uint64_t src_rem = bytesUntilBufferEnd( + local_desc, reinterpret_cast(request.source) + offset, + false, &src_buffer_idx); + if (src_rem > 0) + slice_length = std::min(slice_length, src_rem); + + const uint64_t tgt_rem = bytesUntilBufferEnd( + target_desc, request.target_offset + offset, true, &tgt_buffer_idx); + if (tgt_rem > 0) + slice_length = std::min(slice_length, tgt_rem); + + return slice_length; + } +}; + // Mode definition for MC_IB_PCI_RELAXED_ORDERING env. -// 0 - disabled, 1 - enabled if supported, 2 - auto (default, same as 1 today). +// 0 - disabled, 1 - enabled if supported (default), 2 - auto (same as 1 today). static int getIbRelaxedOrderingMode() { int val = globalConfig().ib_pci_relaxed_ordering_mode; if (val < 0 || val > 2) { @@ -201,121 +286,252 @@ int RdmaTransport::registerLocalMemoryInternal(void *addr, size_t length, bool remote_accessible, bool update_metadata, bool force_sequential) { - (void)remote_accessible; - BufferDesc buffer_desc; - const int kBaseAccessRights = IBV_ACCESS_LOCAL_WRITE | - IBV_ACCESS_REMOTE_WRITE | - IBV_ACCESS_REMOTE_READ; - - int access_rights = kBaseAccessRights; + int access_rights = IBV_ACCESS_LOCAL_WRITE; + if (remote_accessible) + access_rights |= IBV_ACCESS_REMOTE_WRITE | IBV_ACCESS_REMOTE_READ; if (MCIbRelaxedOrderingEnabled) { access_rights |= IBV_ACCESS_RELAXED_ORDERING; } - bool do_pre_touch = context_list_.size() > 0 && - std::thread::hardware_concurrency() >= 4 && - length >= (size_t)4 * 1024 * 1024 * 1024; - if (do_pre_touch) { - // Parallel Pre-touch the memory to speedup the registration process. - int ret = preTouchMemory(addr, length); - if (ret != 0) { - return ret; - } - } - - /* Parallel register when: - 1. parallel_reg_mr is enabled via MC_ENABLE_PARALLEL_REG_MR; - 2. parallel_reg_mr not set and multiple contexts exist and memory has been - pre-touched - Note: If memory hasn't been touched, parallel register can be - slower. Details in: https://github.com/kvcache-ai/Mooncake/issues/848 - Note: force_sequential is used by batch operations to avoid nested - parallelism. - */ - int use_parallel_reg = 0; - if (!force_sequential) { - use_parallel_reg = globalConfig().parallel_reg_mr; - if (use_parallel_reg == -1) { - use_parallel_reg = context_list_.size() > 1 && do_pre_touch; + + // Mooncake#2017: ibv_reg_mr silently truncates a registration to the device + // max_mr_size, but the metadata would still advertise the full BufferDesc + // length, so any remote RDMA op past the boundary fails with + // IBV_WC_REM_ACCESS_ERR (ionic CQE error 10). Split buffers larger than + // max_mr_size into chunks of <= max_mr_size, register each as its own MR, + // and publish one BufferDesc per chunk (the per-context rkey/lkey lookups + // are address-range based, so each chunk gets the correct key). + size_t chunk_limit = (size_t)globalConfig().max_mr_size; + std::vector> chunks; + if (chunk_limit > 0 && length > chunk_limit) { + for (size_t offset = 0; offset < length;) { + size_t chunk_len = std::min(chunk_limit, length - offset); + chunks.emplace_back(static_cast(addr) + offset, chunk_len); + offset += chunk_len; } + LOG(WARNING) << "Auto-splitting buffer " << addr << " (" << length + << " bytes) into " << chunks.size() + << " chunks of <= " << chunk_limit + << " bytes each (device max_mr_size; Mooncake#2017)"; + } else { + chunks.emplace_back(addr, length); } - auto reg_start = std::chrono::steady_clock::now(); - - if (use_parallel_reg) { - std::vector reg_threads; - reg_threads.reserve(context_list_.size()); - std::vector ret_codes(context_list_.size(), 0); - const int ar = access_rights; // Local copy for lambda capture + // Resolve the location name once, from the original buffer. + std::string resolved_name; + if (name == kWildcardLocation) { + bool only_first_page = true; + const std::vector entries = + getMemoryLocation(addr, length, only_first_page); + if (entries.empty()) return -1; + resolved_name = entries[0].location; + } else { + resolved_name = name; + } - for (size_t i = 0; i < context_list_.size(); ++i) { - reg_threads.emplace_back([this, &ret_codes, i, addr, length, ar]() { - ret_codes[i] = - context_list_[i]->registerMemoryRegion(addr, length, ar); - }); + // Export a single dma_buf fd for the whole buffer and import it into every + // NIC's PD during each chunk's registration below (one dma_buf object + // shared across NICs keeps a single BAR1 window for the buffer instead of + // one per NIC). Host memory yields an empty export (plain ibv_reg_mr). The + // fd must stay open across every registration that consumes it (each MR + // takes its own reference), so it is closed once, on any return path, by + // the RAII guard below. + DmabufExport dmabuf_exp; + if (!context_list_.empty()) { + int eret = RdmaContext::exportDmabuf(addr, dmabuf_exp); + if (eret != 0) { + LOG(ERROR) << "Failed to export dma_buf for addr=" << addr; + return eret; } - - for (auto &thread : reg_threads) { - thread.join(); + } + struct DmabufCloser { + DmabufExport &exp; + ~DmabufCloser() { RdmaContext::closeDmabufExport(exp); } + } dmabuf_closer{dmabuf_exp}; + + // Best-effort unregister of ONE chunk's MRs across all contexts. Used to + // clean up a chunk whose registration failed part-way (some contexts + // succeeded, one failed) BEFORE its metadata was committed — that chunk is + // not a "committed" chunk, so rollbackChunks() below must not touch its + // (never-added) metadata, but its partial MRs still need releasing. + auto unregisterChunkMRs = [&](void *chunk_addr) { + for (auto &context : context_list_) { + int ret = context->unregisterMemoryRegion(chunk_addr); + if (ret) + LOG(WARNING) << "Rollback: failed to unregister chunk MR at " + << chunk_addr << " (ret=" << ret << ")"; + } + }; + + // Best-effort rollback of the first `committed` FULLY-committed chunks + // (metadata added to the local segment desc AND MRs registered), i.e. + // chunks [0, committed). Pass the count of chunks whose + // addLocalMemoryBuffer has succeeded — `committed == 0` is a no-op (nothing + // to undo). Metadata is removed WITHOUT a per-chunk publish; one + // updateLocalSegmentDesc() at the end republishes the cleaned desc. + auto rollbackChunks = [&](size_t committed) { + size_t n = std::min(committed, chunks.size()); + for (size_t ri = 0; ri < n; ++ri) { + int rc = metadata_->removeLocalMemoryBuffer( + chunks[ri].first, /*update_metadata=*/false); + if (rc) + LOG(WARNING) << "Rollback: failed to remove metadata for chunk " + "at " + << chunks[ri].first << " (ret=" << rc << ")"; + unregisterChunkMRs(chunks[ri].first); + } + if (n > 0 && update_metadata) metadata_->updateLocalSegmentDesc(); + }; + + // Pre-touch decision is loop-invariant: it depends only on context_list_, + // hardware_concurrency(), and the ORIGINAL buffer length (never chunk_len, + // which is capped at max_mr_size and would silently disable pre-touch for a + // >=4GiB buffer). Compute once above the loop to avoid repeated + // hardware_concurrency() OS queries per chunk. + const bool do_pre_touch = context_list_.size() > 0 && + std::thread::hardware_concurrency() >= 4 && + length >= (size_t)4 * 1024 * 1024 * 1024; + + for (size_t ci = 0; ci < chunks.size(); ++ci) { + void *chunk_addr = chunks[ci].first; + size_t chunk_len = chunks[ci].second; + const uint64_t chunk_offset = reinterpret_cast(chunk_addr) - + reinterpret_cast(addr); + DmabufExport chunk_dmabuf_exp = dmabuf_exp; + if (chunk_dmabuf_exp.method == DmabufExport::Method::kDmabufReg) + chunk_dmabuf_exp.offset += chunk_offset; + + if (do_pre_touch) { + // Parallel pre-touch the memory to speed up registration. + int ret = preTouchMemory(chunk_addr, chunk_len); + if (ret != 0) { + // pre-touch is before MR registration for chunk ci, so ci has + // no MR/metadata yet: roll back only committed chunks [0, ci). + rollbackChunks(ci); + return ret; + } } - for (size_t i = 0; i < ret_codes.size(); ++i) { - if (ret_codes[i] != 0) { - LOG(ERROR) << "Failed to register memory region with context " - << i; - return ret_codes[i]; + /* Parallel register when: + 1. parallel_reg_mr is enabled via MC_ENABLE_PARALLEL_REG_MR; + 2. parallel_reg_mr not set, multiple contexts exist, memory pre-touched. + force_sequential is used by batch operations to avoid nested + parallelism. + */ + int use_parallel_reg = 0; + if (!force_sequential) { + use_parallel_reg = globalConfig().parallel_reg_mr; + if (use_parallel_reg == -1) { + use_parallel_reg = context_list_.size() > 1 && do_pre_touch; } } - } else { - for (size_t i = 0; i < context_list_.size(); ++i) { - int ret = context_list_[i]->registerMemoryRegion(addr, length, - access_rights); - if (ret) { - LOG(ERROR) << "Failed to register memory region with context " - << i; - return ret; + + auto reg_start = std::chrono::steady_clock::now(); + + if (use_parallel_reg) { + std::vector reg_threads; + reg_threads.reserve(context_list_.size()); + std::vector ret_codes(context_list_.size(), 0); + const int ar = access_rights; // Local copy for lambda capture + + for (size_t i = 0; i < context_list_.size(); ++i) { + reg_threads.emplace_back([this, &ret_codes, chunk_dmabuf_exp, i, + chunk_addr, chunk_len, ar]() { + ret_codes[i] = context_list_[i]->registerMemoryRegion( + chunk_addr, chunk_len, ar, chunk_dmabuf_exp); + }); + } + + for (auto &thread : reg_threads) thread.join(); + + for (size_t i = 0; i < ret_codes.size(); ++i) { + if (ret_codes[i] != 0) { + LOG(ERROR) << "Failed to register memory region (chunk " + << ci << ") with context " << i; + // chunk ci's MRs are partially registered but its metadata + // was never added; release ci's MRs, then roll back the + // committed chunks [0, ci). + unregisterChunkMRs(chunk_addr); + rollbackChunks(ci); + return ret_codes[i]; + } + } + } else { + for (size_t i = 0; i < context_list_.size(); ++i) { + int ret = context_list_[i]->registerMemoryRegion( + chunk_addr, chunk_len, access_rights, chunk_dmabuf_exp); + if (ret) { + LOG(ERROR) << "Failed to register memory region (chunk " + << ci << ") with context " << i; + // chunk ci's MRs are partially registered but its metadata + // was never added; release ci's MRs, then roll back [0, + // ci). + unregisterChunkMRs(chunk_addr); + rollbackChunks(ci); + return ret; + } } } - } - auto reg_end = std::chrono::steady_clock::now(); - auto reg_duration_ms = - std::chrono::duration_cast(reg_end - - reg_start) - .count(); + auto reg_end = std::chrono::steady_clock::now(); + auto reg_duration_ms = + std::chrono::duration_cast(reg_end - + reg_start) + .count(); + if (globalConfig().trace) { + LOG(INFO) << "registerMemoryRegion: chunk " << ci << "/" + << chunks.size() << ", addr=" << chunk_addr + << ", length=" << chunk_len + << ", contexts=" << context_list_.size() + << ", parallel=" << (use_parallel_reg ? "true" : "false") + << ", duration=" << reg_duration_ms << "ms"; + } - if (globalConfig().trace) { - LOG(INFO) << "registerMemoryRegion: addr=" << addr - << ", length=" << length - << ", contexts=" << context_list_.size() - << ", parallel=" << (use_parallel_reg ? "true" : "false") - << ", duration=" << reg_duration_ms << "ms"; + // Collect per-context keys for THIS chunk (address-range lookup). + BufferDesc buffer_desc; + for (auto &context : context_list_) { + buffer_desc.lkey.push_back(context->lkey(chunk_addr)); + if (remote_accessible) + buffer_desc.rkey.push_back(context->rkey(chunk_addr)); + } + buffer_desc.name = resolved_name; + buffer_desc.addr = (uint64_t)chunk_addr; + buffer_desc.length = chunk_len; +#ifdef ENABLE_MULTI_PROTOCOL + buffer_desc.protocol = "rdma"; +#endif + // Add to the LOCAL segment desc only (update_metadata=false); a chunked + // buffer otherwise publishes to the metadata server once PER CHUNK. We + // publish once, below, after every chunk has been added. + int rc = metadata_->addLocalMemoryBuffer(buffer_desc, + /*update_metadata=*/false); + if (rc) { + // ci's MRs are registered but its metadata add failed; release ci's + // MRs, then roll back the committed chunks [0, ci). + unregisterChunkMRs(chunk_addr); + rollbackChunks(ci); + return rc; + } } - // Collect keys from all contexts - for (auto &context : context_list_) { - buffer_desc.lkey.push_back(context->lkey(addr)); - buffer_desc.rkey.push_back(context->rkey(addr)); + // Publish the accumulated per-chunk BufferDescs in a SINGLE metadata update + // (a chunked buffer otherwise publishes once per chunk). + if (update_metadata) { + int rc = metadata_->updateLocalSegmentDesc(); + if (rc) { + rollbackChunks(chunks.size()); + return rc; + } } - // Get the memory location automatically after registered MR(pinned), - // when the name is kWildcardLocation("*"). - if (name == kWildcardLocation) { - bool only_first_page = true; - const std::vector entries = - getMemoryLocation(addr, length, only_first_page); - if (entries.empty()) return -1; - buffer_desc.name = entries[0].location; - } else { - buffer_desc.name = name; + // Remember chunk start-addresses so unregisterLocalMemory(addr) (which only + // gets the base addr) can clean up every chunk. + if (chunks.size() > 1) { + std::lock_guard lock(chunk_map_mutex_); + std::vector chunk_addrs; + chunk_addrs.reserve(chunks.size()); + for (auto &c : chunks) chunk_addrs.push_back((uint64_t)c.first); + chunk_map_[(uint64_t)addr] = std::move(chunk_addrs); } - - buffer_desc.addr = (uint64_t)addr; - buffer_desc.length = length; -#ifdef ENABLE_MULTI_PROTOCOL - buffer_desc.protocol = "rdma"; -#endif - int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata); - if (rc) return rc; return 0; } @@ -326,6 +542,80 @@ int RdmaTransport::unregisterLocalMemory(void *addr, bool update_metadata) { int RdmaTransport::unregisterLocalMemoryInternal(void *addr, bool update_metadata, bool force_sequential) { + // Mooncake#2017: if this base buffer was split into chunks at registration, + // unregister each chunk's MR + metadata entry (unregisterLocalMemory only + // receives the base addr). + std::vector chunk_addrs; + { + std::lock_guard lock(chunk_map_mutex_); + auto it = chunk_map_.find((uint64_t)addr); + if (it != chunk_map_.end()) { + chunk_addrs = std::move(it->second); + chunk_map_.erase(it); + } + } + if (!chunk_addrs.empty()) { + // Unregister EVERY chunk even if one fails; chunk_map_ was already + // erased, so an early return would leak the remaining chunks' MRs + + // metadata. Remember the first error and report it at the end. + int first_err = 0; + + // Metadata: remove each chunk from the local desc WITHOUT publishing (a + // chunked buffer otherwise publishes to the metadata server once per + // chunk); publish once after all removals, below. + for (uint64_t ca : chunk_addrs) { + int rc = metadata_->removeLocalMemoryBuffer( + reinterpret_cast(ca), /*update_metadata=*/false); + if (rc && !first_err) first_err = rc; + } + + // MRs: unregister across contexts in PARALLEL (one thread per context, + // each releasing all chunks) — the previous code did chunks × contexts + // fully sequentially (e.g. 4 chunks × 8 NICs = 32 sequential + // ibv_dereg_mr). Mirrors the parallel path used for single buffers. + int use_parallel_unreg = 0; + if (!force_sequential) { + use_parallel_unreg = globalConfig().parallel_reg_mr; + if (use_parallel_unreg == -1) + use_parallel_unreg = context_list_.size() > 1; + } + if (use_parallel_unreg) { + std::vector threads; + threads.reserve(context_list_.size()); + std::vector ret_codes(context_list_.size(), 0); + for (size_t i = 0; i < context_list_.size(); ++i) { + threads.emplace_back([this, &ret_codes, i, &chunk_addrs]() { + for (uint64_t ca : chunk_addrs) { + int ret = context_list_[i]->unregisterMemoryRegion( + reinterpret_cast(ca)); + if (ret && !ret_codes[i]) ret_codes[i] = ret; + } + }); + } + for (auto &t : threads) t.join(); + for (int rc : ret_codes) + if (rc && !first_err) first_err = rc; + } else { + for (uint64_t ca : chunk_addrs) + for (auto &context : context_list_) { + int ret = context->unregisterMemoryRegion( + reinterpret_cast(ca)); + if (ret) { + LOG(ERROR) << "Failed to unregister chunk MR at " + << reinterpret_cast(ca); + if (!first_err) first_err = ret; + } + } + } + + // Single metadata publish covering all removed chunks. + if (update_metadata) { + int rc = metadata_->updateLocalSegmentDesc(); + if (rc && !first_err) first_err = rc; + } + return first_err; + } + int rc = metadata_->removeLocalMemoryBuffer(addr, update_metadata); if (rc) return rc; @@ -433,7 +723,7 @@ int RdmaTransport::refreshLocalDeviceDesc(const std::string &device_name, int RdmaTransport::registerLocalMemoryBatch( const std::vector &buffer_list, const std::string &location) { -#if defined(USE_CUDA) +#if defined(USE_CUDA) || defined(USE_SUPA) if (!Environ::Get().GetWithNvidiaPeermem()) { for (auto &buffer : buffer_list) { int ret = registerLocalMemory(buffer.addr, buffer.length, location, @@ -442,6 +732,7 @@ int RdmaTransport::registerLocalMemoryBatch( LOG(WARNING) << "RdmaTransport: Failed to register memory: addr " << buffer.addr << " length " << buffer.length; + return ret; } } } else { @@ -457,15 +748,19 @@ int RdmaTransport::registerLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < buffer_list.size(); ++i) { - if (results[i].get()) { + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "RdmaTransport: Failed to register memory: addr " << buffer_list[i].addr << " length " << buffer_list[i].length; + if (!first_error) first_error = ret; } } -#if defined(USE_CUDA) + if (first_error) return first_error; +#if defined(USE_CUDA) || defined(USE_SUPA) } // Environ::Get().GetWithNvidiaPeermem() #endif @@ -483,13 +778,17 @@ int RdmaTransport::unregisterLocalMemoryBatch( })); } + int first_error = 0; for (size_t i = 0; i < addr_list.size(); ++i) { - if (results[i].get()) + int ret = results[i].get(); + if (ret) { LOG(WARNING) << "RdmaTransport: Failed to unregister memory: addr " << addr_list[i]; + if (!first_error) first_error = ret; + } } - - return metadata_->updateLocalSegmentDesc(); + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status RdmaTransport::submitTransfer( @@ -506,7 +805,17 @@ Status RdmaTransport::submitTransfer( size_t task_id = batch_desc.task_list.size(); batch_desc.task_list.resize(task_id + entries.size()); std::vector task_list; - for (auto &task : batch_desc.task_list) task_list.push_back(&task); + for (auto &request : entries) { + auto &task = batch_desc.task_list[task_id]; + ++task_id; + task.batch_id = batch_id; +#ifdef USE_ASCEND_HETEROGENEOUS + task.request = const_cast(&request); +#else + task.request = &request; +#endif + task_list.push_back(&task); + } return submitTransferTask(task_list); } @@ -514,6 +823,8 @@ Status RdmaTransport::submitTransferTask( const std::vector &task_list) { std::unordered_map, std::vector> slices_to_post; + std::unordered_map> + target_segment_descs; auto local_segment_desc = metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); assert(local_segment_desc.get()); const size_t kBlockSize = globalConfig().slice_size; @@ -521,6 +832,38 @@ Status RdmaTransport::submitTransferTask( const size_t kFragmentSize = globalConfig().fragment_limit; const size_t kSubmitWatermark = globalConfig().max_wr * globalConfig().num_qp_per_ep; + auto fail_unposted_slices = [&]() { + for (auto &entry : slices_to_post) + for (auto *slice : entry.second) slice->markFailed(); + slices_to_post.clear(); + }; + + // Fabricate a zero-length failed slice for unstarted tasks so that the + // existing success + failed == slice_count terminal check can drive the + // task to FAILED without special-casing, and ~TransferTask reclaims the + // slice. + auto fail_unstarted_tasks = [&](size_t first_task_index) { + for (size_t i = first_task_index; i < task_list.size(); ++i) { + auto &task = *task_list[i]; + Slice *slice = getSliceCache().allocate(); + assert(slice); + slice->source_addr = nullptr; + slice->length = 0; + slice->task = &task; + slice->status = Slice::PENDING; + task.slice_list.push_back(slice); + __sync_fetch_and_add(&task.slice_count, 1); + slice->markFailed(); + } + }; + auto fail_task_and_cleanup = [&](TransferTask &task, Slice *slice, + size_t task_index) { + task.total_bytes += slice->length; + __sync_fetch_and_add(&task.slice_count, 1); + slice->markFailed(); + fail_unposted_slices(); + fail_unstarted_tasks(task_index + 1); + }; uint64_t nr_slices; for (size_t index = 0; index < task_list.size(); ++index) { assert(task_list[index]); @@ -528,6 +871,15 @@ Status RdmaTransport::submitTransferTask( nr_slices = 0; assert(task.request); auto &request = *task.request; + auto target_desc_it = target_segment_descs.find(request.target_id); + if (target_desc_it == target_segment_descs.end()) { + target_desc_it = + target_segment_descs + .emplace(request.target_id, + metadata_->getSegmentDescByID(request.target_id)) + .first; + } + const auto &target_segment_desc = target_desc_it->second; auto request_buffer_id = -1, request_device_id = -1; if (selectDevice(local_segment_desc.get(), (uint64_t)request.source, @@ -537,20 +889,21 @@ Status RdmaTransport::submitTransferTask( request_device_id = -1; } - for (uint64_t offset = 0; offset < request.length; - offset += kBlockSize) { + SliceLengthCalculator slice_calc{request, kBlockSize, kFragmentSize, + local_segment_desc.get(), + target_segment_desc.get()}; + for (uint64_t offset = 0; offset < request.length;) { + size_t slice_length = slice_calc.calculate(offset); + Slice *slice = getSliceCache().allocate(); assert(slice); if (!slice->from_cache) { nr_slices++; } - bool merge_final_slice = - request.length - offset <= kBlockSize + kFragmentSize; - slice->source_addr = (char *)request.source + offset; - slice->length = - merge_final_slice ? request.length - offset : kBlockSize; + slice->length = slice_length; + slice->source_location.clear(); slice->opcode = request.opcode; slice->rdma.dest_addr = request.target_offset + offset; slice->rdma.retry_cnt = request.advise_retry_cnt; @@ -565,9 +918,12 @@ Status RdmaTransport::submitTransferTask( retry_cnt = request.advise_retry_cnt; bool found_device = false; if (request_buffer_id >= 0 && request_device_id >= 0) { - found_device = true; - buffer_id = request_buffer_id; - device_id = request_device_id; + auto &request_context = context_list_[request_device_id]; + if (request_context && request_context->active()) { + found_device = true; + buffer_id = request_buffer_id; + device_id = request_device_id; + } } while (retry_cnt < kMaxRetryCount && !found_device) { if (selectDevice(local_segment_desc.get(), @@ -589,8 +945,7 @@ Status RdmaTransport::submitTransferTask( } if (!found_device) { auto source_addr = slice->source_addr; - for (auto &entry : slices_to_post) - for (auto s : entry.second) getSliceCache().deallocate(s); + fail_task_and_cleanup(task, slice, index); LOG(ERROR) << "Memory region not registered by any active device(s): " << source_addr; @@ -600,6 +955,7 @@ Status RdmaTransport::submitTransferTask( } else { auto &context = context_list_[device_id]; if (!context->active()) { + fail_task_and_cleanup(task, slice, index); LOG(ERROR) << "Device " << device_id << " is not active"; return Status::InvalidArgument("Device " + std::to_string(device_id) + @@ -607,6 +963,11 @@ Status RdmaTransport::submitTransferTask( } slice->rdma.source_lkey = local_segment_desc->buffers[buffer_id].lkey[device_id]; + if (globalConfig().log_rdma_slice_affinity) { + slice->source_location = resolveBufferLocation( + local_segment_desc->buffers[buffer_id], + reinterpret_cast(slice->source_addr)); + } slices_to_post[context].push_back(slice); task.total_bytes += slice->length; __sync_fetch_and_add(&task.slice_count, 1); @@ -619,9 +980,7 @@ Status RdmaTransport::submitTransferTask( nr_slices = 0; } - if (merge_final_slice) { - break; - } + offset += slice->length; } } @@ -686,7 +1045,11 @@ RdmaTransport::SegmentID RdmaTransport::getSegmentID( int RdmaTransport::onSetupRdmaConnections(const HandShakeDesc &peer_desc, HandShakeDesc &local_desc) { auto local_nic_name = getNicNameFromNicPath(peer_desc.peer_nic_path); - if (local_nic_name.empty()) return ERR_INVALID_ARGUMENT; + if (local_nic_name.empty()) { + local_desc.reply_msg = + "Invalid peer_nic_path in handshake: " + peer_desc.peer_nic_path; + return ERR_INVALID_ARGUMENT; + } std::shared_ptr context; int index = 0; @@ -697,12 +1060,43 @@ int RdmaTransport::onSetupRdmaConnections(const HandShakeDesc &peer_desc, } index++; } - if (!context) return ERR_INVALID_ARGUMENT; + if (!context) { + local_desc.reply_msg = + "Local RDMA context not found for handshake NIC: " + local_nic_name; + return ERR_INVALID_ARGUMENT; + } // Use existing endpoint or create new one. auto endpoint = context->endpoint(peer_desc.local_nic_path); - if (!endpoint) return ERR_ENDPOINT; - return endpoint->setupConnectionsByPassive(peer_desc, local_desc); + if (!endpoint) { + local_desc.reply_msg = "Local RDMA endpoint unavailable for " + + local_nic_name + " <- " + + peer_desc.local_nic_path; + return ERR_ENDPOINT; + } + int ret = endpoint->setupConnectionsByPassive(peer_desc, local_desc); + if (endpoint->retired()) { + context->deleteEndpointByPtr(endpoint.get()); + if (ret == ERR_ENDPOINT) { + // setupConnectionsByPassive() can retire a stale endpoint before + // creating a usable passive connection for this incoming handshake. + // That is a local endpoint-store race, not necessarily a peer + // handshake failure, so absorb it once with a fresh endpoint. + local_desc = HandShakeDesc(); + endpoint = context->endpoint(peer_desc.local_nic_path); + if (!endpoint) { + local_desc.reply_msg = + "Fresh local RDMA endpoint unavailable after retiring " + "stale endpoint for " + + local_nic_name + " <- " + peer_desc.local_nic_path; + return ERR_ENDPOINT; + } + ret = endpoint->setupConnectionsByPassive(peer_desc, local_desc); + if (endpoint->retired()) + context->deleteEndpointByPtr(endpoint.get()); + } + } + return ret; } int RdmaTransport::initializeRdmaResources() { @@ -717,6 +1111,16 @@ int RdmaTransport::initializeRdmaResources() { if (ret) { local_topology_->disableDevice(device_name); LOG(WARNING) << "Disable device " << device_name; + // Keep context_list_ index-aligned with getHcaList(): both it and + // BufferDesc::lkey are subscripted by the HCA index, which + // disableDevice() leaves in place. Dropping a slot would make a + // later device_id name the wrong RNIC or run off the end. A + // never-constructed context is an inert placeholder; the partially + // built one is released so it does not pin an open uverbs fd. + auto placeholder = + std::make_shared(*this, device_name); + placeholder->set_active(false); + context_list_.push_back(std::move(placeholder)); } else { context_list_.push_back(context); } @@ -748,6 +1152,17 @@ int RdmaTransport::selectDevice(SegmentDesc *desc, uint64_t offset, ++buffer_id) { const auto &buffer = buffers[buffer_id]; +#ifdef ENABLE_MULTI_PROTOCOL + // The RDMA transport must only bind buffers registered under the rdma + // protocol. Device (hip) buffers alias the same GPU addresses but carry + // no lkey/rkey, so picking one yields an empty-lkey out-of-bounds read + // in the submit path. The !empty() guard leaves legacy single-protocol + // descriptors (empty protocol field) unaffected. + if (!buffer.protocol.empty() && buffer.protocol != "rdma") { + continue; + } +#endif + // Check if offset is within buffer range if (offset < buffer.addr || length > buffer.length || offset - buffer.addr > buffer.length - length) { @@ -776,6 +1191,39 @@ int RdmaTransport::selectDevice(SegmentDesc *desc, uint64_t offset, return ERR_ADDRESS_NOT_REGISTERED; } +int RdmaTransport::selectDeviceByLocalHca(SegmentDesc *desc, uint64_t offset, + size_t length, + std::string_view local_hca, + int &buffer_id, int &device_id, + int retry_count) { + if (desc == nullptr) return ERR_ADDRESS_NOT_REGISTERED; + const auto &buffers = desc->buffers; + for (buffer_id = 0; buffer_id < static_cast(buffers.size()); + ++buffer_id) { + const auto &buffer = buffers[buffer_id]; + +#ifdef ENABLE_MULTI_PROTOCOL + if (!buffer.protocol.empty() && buffer.protocol != "rdma") { + continue; + } +#endif + + if (offset < buffer.addr || length > buffer.length || + offset - buffer.addr > buffer.length - length) { + continue; + } + + const auto location = resolveBufferLocation(buffer, offset); + device_id = desc->topology.selectDeviceByLocalHca(location, local_hca, + retry_count); + if (device_id >= 0) return 0; + device_id = desc->topology.selectDeviceByLocalHca( + kWildcardLocation, local_hca, retry_count); + if (device_id >= 0) return 0; + } + return ERR_ADDRESS_NOT_REGISTERED; +} + int RdmaTransport::selectDevice(SegmentDesc *desc, uint64_t offset, size_t length, int &buffer_id, int &device_id, int retry_count) { diff --git a/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp b/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp index 9ef0dd1282..abd24be303 100644 --- a/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp +++ b/mooncake-transfer-engine/src/transport/rdma_transport/worker_pool.cpp @@ -17,8 +17,10 @@ #include #include +#include #include "config.h" +#include "memory_location.h" #include "transport/rdma_transport/rdma_context.h" #include "transport/rdma_transport/rdma_endpoint.h" #include "transport/rdma_transport/rdma_transport.h" @@ -31,11 +33,93 @@ namespace mooncake { const static int kTransferWorkerCount = globalConfig().workers_per_ctx; +static std::string resolveBufferLocation( + const TransferMetadata::BufferDesc &buffer, uint64_t offset) { + std::string location = buffer.name; + SegmentsLocationInfo seg_info; + if (parseSegmentsLocation(buffer.name, seg_info)) { + location = resolveSegmentsLocation(seg_info, buffer.length, + offset - buffer.addr); + } + return location; +} + +static const std::string &sourceLocationOrUnknown(Transport::Slice *slice) { + static const std::string kUnknown = ""; + return slice->source_location.empty() ? kUnknown : slice->source_location; +} + +static int selectPeerDevice(RdmaTransport::SegmentDesc *peer_segment_desc, + uint64_t offset, size_t length, + const std::string &local_hca, int &buffer_id, + int &device_id, int retry_count = 0) { + const auto &config = globalConfig(); + int ret = 0; + if (config.enable_hca_peer_affinity) { + ret = RdmaTransport::selectDeviceByLocalHca( + peer_segment_desc, offset, length, local_hca, buffer_id, device_id, + retry_count); + } else { + auto hint = config.enable_dest_device_affinity + ? std::string_view(local_hca) + : std::string_view(); + ret = + RdmaTransport::selectDevice(peer_segment_desc, offset, length, hint, + buffer_id, device_id, retry_count); + } + if (ret) return ret; + + if (buffer_id < 0 || + static_cast(buffer_id) >= peer_segment_desc->buffers.size() || + device_id < 0 || + static_cast(device_id) >= + peer_segment_desc->buffers[buffer_id].rkey.size()) { + LOG(ERROR) << "[RDMA] No rkey for MR access: seg=" + << (peer_segment_desc ? peer_segment_desc->name : "null") + << " addr=" << (void *)offset << " len=" << length; + return ERR_ADDRESS_NOT_REGISTERED; + } + + // device_id comes from the peer-supplied topology, whose HCA list is + // independent of the peer 'devices' array, so bound it against that array + // too before the devices[device_id] accesses below. decodeSegmentDesc() + // now rejects a descriptor whose key count and device count disagree, which + // makes this check redundant for descriptors that arrived through the + // metadata path; it is kept as a local bound on the value actually used. + if (static_cast(device_id) >= peer_segment_desc->devices.size()) { + LOG(ERROR) << "[RDMA] Peer device index out of range: seg=" + << peer_segment_desc->name << " device_id=" << device_id + << " devices=" << peer_segment_desc->devices.size(); + return ERR_ADDRESS_NOT_REGISTERED; + } + return 0; +} + +static bool workerCanPost(int thread_id) { + return kTransferWorkerCount == 1 || thread_id != 0; +} + +static bool workerCanPoll(int thread_id) { + return kTransferWorkerCount == 1 || thread_id == 0; +} + +static void getPostingShardAssignment(int thread_id, int &post_tid, + int &post_count) { + assert(workerCanPost(thread_id)); + if (kTransferWorkerCount > 1) { + post_tid = thread_id - 1; + post_count = kTransferWorkerCount - 1; + } else { + post_tid = thread_id; + post_count = kTransferWorkerCount; + } +} + WorkerPool::WorkerPool(RdmaContext &context, int numa_socket_id) : context_(context), numa_socket_id_(numa_socket_id), workers_running_(true), - suspended_flag_(0), + parked_worker_count_(0), redispatch_counter_(0), submitted_slice_count_(0), processed_slice_count_(0) { @@ -111,12 +195,9 @@ int WorkerPool::submitPostSend( } auto &peer_segment_desc = segment_desc_map[slice->target_id]; int buffer_id, device_id; - auto hint = globalConfig().enable_dest_device_affinity - ? context_.deviceName() - : ""; - if (RdmaTransport::selectDevice(peer_segment_desc.get(), - slice->rdma.dest_addr, slice->length, - hint, buffer_id, device_id)) { + if (selectPeerDevice(peer_segment_desc.get(), slice->rdma.dest_addr, + slice->length, context_.deviceName(), buffer_id, + device_id)) { peer_segment_desc = context_.engine().meta()->getSegmentDescByID( slice->target_id, true); if (!peer_segment_desc) { @@ -127,9 +208,9 @@ int WorkerPool::submitPostSend( continue; } - if (RdmaTransport::selectDevice( - peer_segment_desc.get(), slice->rdma.dest_addr, - slice->length, hint, buffer_id, device_id)) { + if (selectPeerDevice(peer_segment_desc.get(), slice->rdma.dest_addr, + slice->length, context_.deviceName(), + buffer_id, device_id)) { slice->markFailed(); context_.engine().meta()->dumpMetadataContent( peer_segment_desc->name, slice->rdma.dest_addr, @@ -152,9 +233,13 @@ int WorkerPool::submitPostSend( bool found = false; for (size_t alt_dev_id = 0; alt_dev_id < peer_segment_desc->devices.size(); ++alt_dev_id) { - if (alt_dev_id == (size_t)device_id) continue; + if (alt_dev_id == (size_t)device_id || + alt_dev_id >= + peer_segment_desc->buffers[buffer_id].rkey.size()) { + continue; + } auto alt_path = - MakeNicPath(peer_segment_desc->name, + MakeNicPath(peer_segment_desc->nicPathServerName(), peer_segment_desc->devices[alt_dev_id].name); if (isRailAvailable(alt_path)) { device_id = alt_dev_id; @@ -173,11 +258,40 @@ int WorkerPool::submitPostSend( } slice->peer_nic_path = peer_nic_path; + if (globalConfig().log_rdma_slice_affinity) { + VLOG(1) << "RDMA slice affinity: source_location=" + << sourceLocationOrUnknown(slice) << ", target_location=" + << resolveBufferLocation( + peer_segment_desc->buffers[buffer_id], + slice->rdma.dest_addr) + << ", local_device_name=" << context_.deviceName() + << ", peer_device_name=" + << peer_segment_desc->devices[device_id].name + << ", target_id=" << slice->target_id + << ", source_addr=" << slice->source_addr << ", dest_addr=" + << reinterpret_cast(slice->rdma.dest_addr) + << ", length=" << slice->length; + } int shard_id = (slice->target_id * 10007 + device_id) % kShardCount; slice_list_map[shard_id].push_back(slice); submitted_slice_count++; } + enqueuePreparedSlices(slice_list_map, submitted_slice_count); + + // Context-level health tracking: if all slices failed due to no available + // rails, increment the context failure counter. This detects catastrophic + // local RNIC hardware failure where all paths through the RNIC are down. + if (submitted_slice_count == 0 && + all_rails_failed_count == (int)slice_list.size()) { + if (markContextFailure()) refreshPublishedLocalTopology(); + } + + return 0; +} + +void WorkerPool::enqueuePreparedSlices(SliceList (&slice_list_map)[kShardCount], + uint64_t submitted_slice_count) { for (int shard_id = 0; shard_id < kShardCount; ++shard_id) { if (slice_list_map[shard_id].empty()) continue; slice_queue_lock_[shard_id].lock(); @@ -189,46 +303,90 @@ int WorkerPool::submitPostSend( } submitted_slice_count_.fetch_add(submitted_slice_count); - if (suspended_flag_.load()) { + if (submitted_slice_count && + parked_worker_count_.load(std::memory_order_acquire) > 0) { std::lock_guard lock(cond_mutex_); cond_var_.notify_all(); } +} - // Context-level health tracking: if all slices failed due to no available - // rails, increment the context failure counter. This detects catastrophic - // local RNIC hardware failure where all paths through the RNIC are down. - if (submitted_slice_count == 0 && - all_rails_failed_count == (int)slice_list.size()) { - markContextFailure(); +int WorkerPool::submitPreparedPostSend( + const std::vector &slice_list) { + // Called by a different local RNIC's worker during local failover. The + // slice already carries the chosen peer_nic_path and refreshed local lkey, + // so enqueue it directly instead of running remote-path selection again. + SliceList slice_list_map[kShardCount]; + uint64_t submitted_slice_count = 0; + + for (auto &slice : slice_list) { + if (slice->peer_nic_path.empty()) { + slice->markFailed(); + continue; + } + auto shard_id = static_cast( + std::hash{}(slice->peer_nic_path) % kShardCount); + slice_list_map[shard_id].push_back(slice); + submitted_slice_count++; } + enqueuePreparedSlices(slice_list_map, submitted_slice_count); + return 0; } +void WorkerPool::trackPostedSlices( + const std::vector &slice_list, size_t first, + size_t count) { + if (!globalConfig().track_rdma_posted_slices) return; + + std::lock_guard lock(posted_slices_mutex_); + for (size_t i = first; i < first + count; ++i) + posted_slices_.insert(slice_list[i]); +} + +void WorkerPool::untrackPostedSlices( + const std::vector &slice_list, size_t first, + size_t count) { + if (!globalConfig().track_rdma_posted_slices) return; + + std::lock_guard lock(posted_slices_mutex_); + for (size_t i = first; i < first + count; ++i) + posted_slices_.erase(slice_list[i]); +} + void WorkerPool::performPostSend(int thread_id) { - // Fast-fail if context is unhealthy due to catastrophic hardware failure - if (!contextHealthy()) { - auto &local_slice_queue = collective_slice_queue_[thread_id]; - for (int shard_id = thread_id; shard_id < kShardCount; - shard_id += kTransferWorkerCount) { + int post_tid = 0; + int post_count = 0; + getPostingShardAssignment(thread_id, post_tid, post_count); + auto &local_slice_queue = collective_slice_queue_[thread_id]; + + // If this local RNIC is inactive/unhealthy, the remote rail is not the + // problem. Move queued work to another local RNIC while preserving the + // already selected peer rail. + if (!context_.active() || !contextHealthy()) { + auto local_slice_queue_clone = local_slice_queue; + local_slice_queue.clear(); + for (auto &entry : local_slice_queue_clone) + redispatch(entry.second, thread_id, true); + + for (int shard_id = post_tid; shard_id < kShardCount; + shard_id += post_count) { if (slice_queue_count_[shard_id].load(std::memory_order_relaxed) == 0) continue; slice_queue_lock_[shard_id].lock(); - for (auto &entry : slice_queue_[shard_id]) { - for (auto &slice : entry.second) slice->markFailed(); - processed_slice_count_ += entry.second.size(); - } + auto slice_queue_clone = slice_queue_[shard_id]; slice_queue_[shard_id].clear(); slice_queue_count_[shard_id].store(0, std::memory_order_relaxed); slice_queue_lock_[shard_id].unlock(); + for (auto &entry : slice_queue_clone) + redispatch(entry.second, thread_id, true); } return; } - auto &local_slice_queue = collective_slice_queue_[thread_id]; - for (int shard_id = thread_id; shard_id < kShardCount; - shard_id += kTransferWorkerCount) { + for (int shard_id = post_tid; shard_id < kShardCount; + shard_id += post_count) { if (slice_queue_count_[shard_id].load(std::memory_order_relaxed) == 0) continue; @@ -250,8 +408,9 @@ void WorkerPool::performPostSend(int thread_id) { redispatch_counter_.load(std::memory_order_relaxed); auto local_slice_queue_clone = local_slice_queue; local_slice_queue.clear(); + bool handoff_to_local_worker = !context_.active() || !contextHealthy(); for (auto &entry : local_slice_queue_clone) - redispatch(entry.second, thread_id); + redispatch(entry.second, thread_id, handoff_to_local_worker); return; } @@ -275,6 +434,18 @@ void WorkerPool::performPostSend(int thread_id) { processed_slice_count_.fetch_add(entry.second.size()); entry.second.clear(); #else + // Check the connection pause before looking up or creating an endpoint. + // A paused peer is a policy decision, not a new path failure: routing + // it through setupConnectionsByActive() would return ERR_ENDPOINT, + // delete the endpoint, and refresh the pause without attempting a + // connection. Keeping the check here gives the pause a hard retry + // deadline; only a genuine connection/QP failure can arm or extend it. + if (!isRailAvailable(entry.first) || + context_.isConnectPaused(entry.first)) { + for (auto &slice : entry.second) failed_slice_list.push_back(slice); + entry.second.clear(); + continue; + } #ifdef CONFIG_CACHE_ENDPOINT auto &endpoint = endpoint_map[entry.first]; if (endpoint == nullptr || !endpoint->active()) @@ -287,13 +458,65 @@ void WorkerPool::performPostSend(int thread_id) { entry.second.clear(); continue; } - if (!endpoint->connected() && endpoint->setupConnectionsByActive()) { - LOG(ERROR) << "Worker: Cannot make connection for endpoint: " - << entry.first << ", deleting endpoint"; - // Unified path failure handling - handlePathFailure(entry.first, endpoint.get()); - for (auto &slice : entry.second) failed_slice_list.push_back(slice); - entry.second.clear(); + if (!endpoint->connected()) { + int setup_ret = endpoint->setupConnectionsByActive(); + if (setup_ret) { + // Active handshake setup failures are ambiguous: the failed + // side may be the peer rail, or this local RNIC may have just + // gone inactive. Prefer switching peer rails when one is + // available; otherwise hand off to another local RNIC only when + // this context is already known inactive. + bool local_context_inactive = !context_.active(); + bool has_peer_alternative = false; + for (auto &slice : entry.second) { + if (hasAvailablePeerRailAlternative(slice, entry.first)) { + has_peer_alternative = true; + break; + } + } + LOG(WARNING) << "Worker: Cannot make connection for endpoint: " + << entry.first + << (has_peer_alternative + ? ", pausing peer rail and retrying " + "through an alternate peer RNIC" + : local_context_inactive + ? ", local RNIC is inactive; " + "trying another local RNIC" + : ", no alternate peer RNIC is available; " + "retrying without pausing peer rail"); + if (has_peer_alternative) { + markRailFailed(entry.first, true); + redispatch_counter_++; + } else if (local_context_inactive) { + context_.set_active(false); + refreshPublishedLocalTopology(); + redispatch_counter_++; + } + context_.deleteEndpointByPtr(endpoint.get()); + for (auto &slice : entry.second) { + if (!has_peer_alternative && local_context_inactive && + tryHandoffToAnotherLocalWorker(slice)) { + processed_slice_count_++; + } else { + failed_slice_list.push_back(slice); + } + } + entry.second.clear(); + continue; + } + } + if (!endpoint->readyToSend()) { + if (endpoint->readyAckTimedOut()) { + LOG(ERROR) << "Worker: Timed out waiting for RDMA ready ACK " + << "for endpoint: " << entry.first + << ", deleting endpoint"; + markRailFailed(entry.first, true); + redispatch_counter_++; + context_.deleteEndpointByPtr(endpoint.get()); + for (auto &slice : entry.second) + failed_slice_list.push_back(slice); + entry.second.clear(); + } continue; } // Set endpoint pointer for each slice before submitting @@ -306,9 +529,13 @@ void WorkerPool::performPostSend(int thread_id) { if (!failed_slice_list.empty()) { SliceList retry_list; + SliceList local_retry_list; for (auto &slice : failed_slice_list) { if (shouldRetrySlice(slice)) { - retry_list.push_back(slice); + if (!context_.active()) + local_retry_list.push_back(slice); + else + retry_list.push_back(slice); } else { slice->markFailed(); processed_slice_count_++; @@ -317,16 +544,53 @@ void WorkerPool::performPostSend(int thread_id) { if (!retry_list.empty()) { redispatch(retry_list, thread_id); } + if (!local_retry_list.empty()) + redispatch(local_retry_list, thread_id, true); } } void WorkerPool::performPollCq(int thread_id) { + const uint64_t poll_ts = getCurrentTimeInNano(); + const uint64_t previous_poll_ts = + last_poll_ts_ns_.exchange(poll_ts, std::memory_order_relaxed); + if (previous_poll_ts > 0 && poll_ts > previous_poll_ts) { + const uint64_t interval = poll_ts - previous_poll_ts; + last_poll_interval_ns_.store(interval, std::memory_order_relaxed); + uint64_t previous_max = + max_poll_interval_ns_.load(std::memory_order_relaxed); + while (interval > previous_max && + !max_poll_interval_ns_.compare_exchange_weak( + previous_max, interval, std::memory_order_relaxed)) { + } + } + + // Slices this loop drove to a terminal state, successes and + // retry-exhausted failures alike; folded into processed_slice_count_, + // which gates worker parking. Every terminal outcome below goes through + // finalize_slice() so it is counted exactly once. Slices handed to + // redispatch() are not terminal here and are accounted for there. int processed_slice_count = 0; + // Successful completions only, kept apart because it clears the context + // health counter. A completion error is no evidence that this RNIC can + // still move data -- including IBV_WC_WR_FLUSH_ERR, which only reports + // WRs the hardware discarded after the QP had already entered ERR. + int succeeded_slice_count = 0; + auto finalize_slice = [&](Transport::Slice *slice, bool success) { + if (success) { + slice->markSuccess(); + succeeded_slice_count++; + } else { + slice->markFailed(); + } + processed_slice_count++; + }; const static size_t kPollCount = 64; - std::unordered_map qp_depth_set; - SliceList failed_slice_list; // Unified: collect all slices for redispatch - for (int cq_index = thread_id; cq_index < context_.cqCount(); - cq_index += kTransferWorkerCount) { + std::unordered_map *, int> qp_depth_set; + std::unordered_set local_failed_endpoints; + bool recorded_local_context_failure = false; + SliceList failed_slice_list; + SliceList local_failed_slice_list; + for (int cq_index = 0; cq_index < context_.cqCount(); cq_index++) { ibv_wc wc[kPollCount]; int nr_poll = context_.poll(kPollCount, wc, cq_index); if (nr_poll < 0) { @@ -334,6 +598,14 @@ void WorkerPool::performPollCq(int thread_id) { continue; } + if (nr_poll > 0 && globalConfig().track_rdma_posted_slices) { + std::lock_guard lock(posted_slices_mutex_); + for (int i = 0; i < nr_poll; ++i) { + auto *slice = reinterpret_cast(wc[i].wr_id); + posted_slices_.erase(slice); + } + } + for (int i = 0; i < nr_poll; ++i) { Transport::Slice *slice = (Transport::Slice *)wc[i].wr_id; assert(slice); @@ -348,18 +620,33 @@ void WorkerPool::performPollCq(int thread_id) { // not real network errors and should not trigger rail failure // handling or endpoint deletion. if (wc[i].status == IBV_WC_WR_FLUSH_ERR) { - if (globalConfig().trace) - LOG(INFO) << "Worker: WR flush error (peer_nic: " - << slice->peer_nic_path - << "), marking failed without retry"; - slice->markFailed(); - processed_slice_count++; + if (!context_.active()) { + if (globalConfig().trace) + LOG(INFO) + << "Worker: WR flush error on inactive " + << "local context " << context_.deviceName() + << " (peer_nic: " << slice->peer_nic_path + << "), handing off if retry allows"; + if (shouldRetrySlice(slice)) + local_failed_slice_list.push_back(slice); + else + finalize_slice(slice, false); + } else { + if (globalConfig().trace) + LOG(INFO) << "Worker: WR flush error (peer_nic: " + << slice->peer_nic_path + << "), redispatching if retry allows"; + if (shouldRetrySlice(slice)) + failed_slice_list.push_back(slice); + else + finalize_slice(slice, false); + } continue; } - // All other WC errors indicate real path/network failures and - // should trigger redispatch to an alternate path (or fail if - // retry exhausted) + // Completion errors are split by local context health. Local + // faults hand off to another local RNIC; remote/default faults + // keep this local context and switch peer rails. LOG(ERROR) << "Worker: Process failed for slice (opcode: " << slice->opcode << ", source_addr: " << slice->source_addr @@ -369,47 +656,82 @@ void WorkerPool::performPollCq(int thread_id) { << ", peer_nic: " << slice->peer_nic_path << ", dest_rkey: " << slice->rdma.dest_rkey << ", retry_cnt: " << slice->rdma.retry_cnt + << ", max_retry_cnt: " << slice->rdma.max_retry_cnt << "): " << ibv_wc_status_str(wc[i].status); - // Unified path failure handling - handlePathFailure(slice->peer_nic_path, slice->rdma.endpoint); + auto *retry_list = &failed_slice_list; + if (!context_.active() || isLocalWcFailure(wc[i])) { + if (!recorded_local_context_failure) { + handleLocalFailure(slice->peer_nic_path, + slice->rdma.endpoint); + recorded_local_context_failure = true; + if (slice->rdma.endpoint) + local_failed_endpoints.insert(slice->rdma.endpoint); + } else if (slice->rdma.endpoint && + !local_failed_endpoints.count( + slice->rdma.endpoint)) { + context_.deleteEndpointByPtr(slice->rdma.endpoint); + local_failed_endpoints.insert(slice->rdma.endpoint); + } + retry_list = &local_failed_slice_list; + } else { + if (hasAvailablePeerRailAlternative(slice, + slice->peer_nic_path)) { + markRailFailed(slice->peer_nic_path, true); + redispatch_counter_++; + } + if (slice->rdma.endpoint) { + context_.deleteEndpointByPtr(slice->rdma.endpoint); + } + } if (shouldRetrySlice(slice)) { - failed_slice_list.push_back(slice); + retry_list->push_back(slice); } else { - slice->markFailed(); - processed_slice_count_++; + finalize_slice(slice, false); } } else { - slice->markSuccess(); - processed_slice_count++; + finalize_slice(slice, true); } } if (nr_poll) - __sync_fetch_and_sub(context_.cqOutstandingCount(cq_index), - nr_poll); + context_.cqOutstandingCount(cq_index)->fetch_sub( + nr_poll, std::memory_order_acq_rel); } for (auto &entry : qp_depth_set) - __sync_fetch_and_sub(entry.first, entry.second); + entry.first->fetch_sub(entry.second, std::memory_order_acq_rel); - if (processed_slice_count) { + if (processed_slice_count) processed_slice_count_.fetch_add(processed_slice_count); - markContextSuccess(); - } + // Clear the consecutive-failure counter only on proven data movement, so + // that repeated local completion failures can still reach the threshold + // and retire this RNIC (see handleLocalFailure()). + if (succeeded_slice_count) markContextSuccess(); + if (!local_failed_slice_list.empty()) { + redispatch(local_failed_slice_list, thread_id, true); + } if (!failed_slice_list.empty()) { redispatch(failed_slice_list, thread_id); } } void WorkerPool::redispatch(std::vector &slice_list, - int thread_id) { + int thread_id, bool handoff_to_local_worker) { std::unordered_map> segment_desc_map; - for (auto &slice : slice_list) { - auto target_id = slice->target_id; - if (!segment_desc_map.count(target_id)) { - segment_desc_map[target_id] = - context_.engine().meta()->getSegmentDescByID(target_id, true); + const bool use_local_queue = workerCanPost(thread_id); + int shared_redispatch_count = 0; + // Remote redispatch needs target metadata to choose a new peer RNIC. + // Local handoff keeps the peer RNIC fixed and only switches source RNIC, so + // it can skip this lookup. + if (!handoff_to_local_worker) { + for (auto &slice : slice_list) { + auto target_id = slice->target_id; + if (!segment_desc_map.count(target_id)) { + segment_desc_map[target_id] = + context_.engine().meta()->getSegmentDescByID(target_id, + true); + } } } @@ -418,13 +740,33 @@ void WorkerPool::redispatch(std::vector &slice_list, slice->markFailed(); processed_slice_count_++; } else { + if (handoff_to_local_worker) { + if (tryHandoffToAnotherLocalWorker(slice)) { + processed_slice_count_++; + continue; + } + // A local RNIC failure cannot be repaired by keeping this + // worker/context and changing the remote rail. If no other + // local worker can take the slice, fail it immediately. + slice->markFailed(); + processed_slice_count_++; + continue; + } + + // Remote-side/default policy: keep local context fixed and switch + // remote path. auto &peer_segment_desc = segment_desc_map[slice->target_id]; int buffer_id, device_id; if (!peer_segment_desc || - RdmaTransport::selectDevice(peer_segment_desc.get(), - slice->rdma.dest_addr, - slice->length, buffer_id, device_id, - slice->rdma.retry_cnt)) { + selectPeerDevice(peer_segment_desc.get(), slice->rdma.dest_addr, + slice->length, context_.deviceName(), + buffer_id, device_id, slice->rdma.retry_cnt)) { + LOG(ERROR) << "Worker: Cannot redispatch slice for target " + << slice->target_id + << ", peer segment unavailable or no target RNIC, " + << "dest_addr=" << (void *)slice->rdma.dest_addr + << ", length=" << slice->length + << ", retry_cnt=" << slice->rdma.retry_cnt; slice->markFailed(); processed_slice_count_++; continue; @@ -434,84 +776,277 @@ void WorkerPool::redispatch(std::vector &slice_list, auto peer_nic_path = MakeNicPath(peer_segment_desc->nicPathServerName(), peer_segment_desc->devices[device_id].name); + if (!isRailAvailable(peer_nic_path)) { + bool found = false; + for (size_t alt_dev_id = 0; + alt_dev_id < peer_segment_desc->devices.size(); + ++alt_dev_id) { + if (alt_dev_id == (size_t)device_id || + alt_dev_id >= + peer_segment_desc->buffers[buffer_id].rkey.size()) { + continue; + } + auto alt_path = MakeNicPath( + peer_segment_desc->nicPathServerName(), + peer_segment_desc->devices[alt_dev_id].name); + if (isRailAvailable(alt_path)) { + device_id = alt_dev_id; + slice->rdma.dest_rkey = + peer_segment_desc->buffers[buffer_id] + .rkey[device_id]; + peer_nic_path = alt_path; + found = true; + break; + } + } + if (!found) { + LOG(ERROR) + << "Worker: Cannot redispatch slice because all peer " + "rails are paused for target " + << slice->target_id + << ", selected peer=" << peer_nic_path + << ", retry_cnt=" << slice->rdma.retry_cnt; + slice->markFailed(); + processed_slice_count_++; + continue; + } + } slice->peer_nic_path = peer_nic_path; - collective_slice_queue_[thread_id][peer_nic_path].push_back(slice); + if (globalConfig().log_rdma_slice_affinity) { + VLOG(1) << "RDMA slice affinity: source_location=" + << sourceLocationOrUnknown(slice) + << ", target_location=" + << resolveBufferLocation( + peer_segment_desc->buffers[buffer_id], + slice->rdma.dest_addr) + << ", local_device_name=" << context_.deviceName() + << ", peer_device_name=" + << peer_segment_desc->devices[device_id].name + << ", target_id=" << slice->target_id + << ", source_addr=" << slice->source_addr + << ", dest_addr=" + << reinterpret_cast(slice->rdma.dest_addr) + << ", length=" << slice->length + << ", retry_cnt=" << slice->rdma.retry_cnt; + } + slice->ts = 0; + if (use_local_queue) { + collective_slice_queue_[thread_id][peer_nic_path].push_back( + slice); + } else { + int shard_id = + (slice->target_id * 10007 + device_id) % kShardCount; + slice_queue_lock_[shard_id].lock(); + slice_queue_[shard_id][peer_nic_path].push_back(slice); + slice_queue_count_[shard_id].fetch_add( + 1, std::memory_order_relaxed); + slice_queue_lock_[shard_id].unlock(); + shared_redispatch_count++; + } } } + + if (shared_redispatch_count && + parked_worker_count_.load(std::memory_order_acquire) > 0) { + std::lock_guard lock(cond_mutex_); + cond_var_.notify_all(); + } +} + +bool WorkerPool::tryHandoffToAnotherLocalWorker(Transport::Slice *slice) { + // Local failover changes only the source RNIC. Keep target_id, + // peer_nic_path, dest_addr, and dest_rkey intact; only replace source_lkey + // for the selected alternate local context. + auto local_segment_desc = + context_.engine().meta()->getSegmentDescByID(LOCAL_SEGMENT_ID); + auto &contexts = context_.engine().context_list_; + if (!local_segment_desc || contexts.size() <= 1) { + return false; + } + + int current_ctx_id = -1; + for (size_t i = 0; i < contexts.size(); ++i) { + if (contexts[i] && contexts[i].get() == &context_) { + current_ctx_id = static_cast(i); + break; + } + } + if (current_ctx_id < 0) { + return false; + } + + int start_ctx = static_cast(slice->rdma.retry_cnt % contexts.size()); + for (size_t offset = 0; offset < contexts.size(); ++offset) { + int device_id = (start_ctx + static_cast(offset)) % + static_cast(contexts.size()); + if (device_id == current_ctx_id) continue; + + auto &alt_ctx = contexts[device_id]; + if (!alt_ctx || !alt_ctx->active()) continue; + + int buffer_id = -1; + for (size_t idx = 0; idx < local_segment_desc->buffers.size(); ++idx) { + auto &buffer = local_segment_desc->buffers[idx]; + auto source = reinterpret_cast(slice->source_addr); + auto buffer_start = reinterpret_cast(buffer.addr); + auto buffer_end = buffer_start + buffer.length; + if (buffer_start <= source && + source + slice->length <= buffer_end) { + buffer_id = static_cast(idx); + break; + } + } + if (buffer_id < 0) { + continue; + } + if (device_id >= + static_cast( + local_segment_desc->buffers[buffer_id].lkey.size())) { + continue; + } + + slice->rdma.source_lkey = + local_segment_desc->buffers[buffer_id].lkey[device_id]; + slice->rdma.endpoint = nullptr; + slice->ts = 0; + + std::vector handoff{slice}; + alt_ctx->worker_pool_->submitPreparedPostSend(handoff); + + VLOG(1) << "Local-side retry handed slice from worker pool on " + << context_.deviceName() << " to worker pool on " + << alt_ctx->deviceName() << " while keeping remote peer " + << slice->peer_nic_path; + return true; + } + + return false; +} + +bool WorkerPool::hasOutstandingCq(int thread_id) { + if (!workerCanPoll(thread_id)) return false; + for (int cq_index = 0; cq_index < context_.cqCount(); ++cq_index) { + if (context_.cqOutstandingCount(cq_index)->load( + std::memory_order_relaxed) > 0) + return true; + } + return false; } void WorkerPool::transferWorker(int thread_id) { bindToSocket(numa_socket_id_); const static uint64_t kWaitPeriodInNano = 100000000; // 100ms uint64_t last_wait_ts = getCurrentTimeInNano(); + const bool can_post = workerCanPost(thread_id); + const bool can_poll = workerCanPoll(thread_id); while (workers_running_.load(std::memory_order_relaxed)) { auto processed_slice_count = processed_slice_count_.load(std::memory_order_relaxed); auto submitted_slice_count = submitted_slice_count_.load(std::memory_order_relaxed); - if (processed_slice_count == submitted_slice_count) { + if (processed_slice_count == submitted_slice_count && + !hasOutstandingCq(thread_id)) { uint64_t curr_wait_ts = getCurrentTimeInNano(); if (curr_wait_ts - last_wait_ts > kWaitPeriodInNano) { std::unique_lock lock(cond_mutex_); - suspended_flag_.fetch_add(1); + parked_worker_count_.fetch_add(1, std::memory_order_acq_rel); // Double-check condition after acquiring lock to avoid lost - // wakeup + // wakeup. parked_worker_count_ is set before this check so + // producers that submit after it will notify this worker. if (processed_slice_count_.load(std::memory_order_relaxed) == - submitted_slice_count_.load()) { + submitted_slice_count_.load() && + !hasOutstandingCq(thread_id)) { cond_var_.wait_for(lock, std::chrono::seconds(1)); } - suspended_flag_.fetch_sub(1); + parked_worker_count_.fetch_sub(1, std::memory_order_acq_rel); last_wait_ts = curr_wait_ts; } continue; } - performPostSend(thread_id); + if (can_post) { + performPostSend(thread_id); + } #ifndef USE_FAKE_POST_SEND - performPollCq(thread_id); + if (can_poll) { + performPollCq(thread_id); + } #endif last_wait_ts = getCurrentTimeInNano(); } } int WorkerPool::doProcessContextEvents() { - ibv_async_event event; - bool event_acked = false; - if (ibv_get_async_event(context_.context(), &event) < 0) return ERR_CONTEXT; - LOG(WARNING) << "Worker: Received context async event " - << ibv_event_type_str(event.event_type) << " for context " + // The async fd is edge-triggered (joinNonblockingPollList) and + // ibv_get_async_event() returns one record per read, so anything left + // queued here waits for an unrelated later event to release it. Bursts + // are routine -- IBV_EVENT_COMM_EST fires once per connection -- and a + // backlog delays every event behind it, port and device errors included. + while (true) { + ibv_async_event event; + bool event_acked = false; + errno = 0; + if (ibv_get_async_event(context_.context(), &event) < 0) { + if (errno == EAGAIN || errno == EWOULDBLOCK) return 0; // drained + if (errno == EINTR) continue; + return ERR_CONTEXT; + } + LOG(WARNING) << "Worker: Received context async event " + << ibv_event_type_str(event.event_type) << " for context " + << context_.deviceName(); + if (event.event_type == IBV_EVENT_QP_FATAL) { + auto endpoint_ptr = (RdmaEndPoint *)event.element.qp->qp_context; + + /** + * There might be a deadlock if we call endpoint->set_active(false) + * before ack the event: + * + * Thread A: + * Holding endpoint->lock_ and calling ibv_destroy_qp (if using + * eRDMA), ibv_destroy_qp will block until the event is acked. + * + * Thread B (this thread): + * Calling endpoint->set_active(false), which blocks as + * endpoint->lock_ is held by Thread A. + */ + ibv_ack_async_event(&event); + event_acked = true; + + /** + * After ack the event, the endpoint might be destroyed if it + * happened to be destroying event.element.qp. Therefore, we cannot + * just dereference endpoint_ptr. Instead, we need to get the + * shared_ptr of the endpoint from context_ and use that shared_ptr + * to access the endpoint. + */ + context_.deleteEndpointByPtr(endpoint_ptr); + } else if (handleContextEvent(event.event_type, false, &event)) { + event_acked = true; + } + + if (!event_acked) { + ibv_ack_async_event(&event); + } + } +} + +void WorkerPool::processContextEventForTest(ibv_event_type event_type) { + LOG(WARNING) << "Worker: Injected context async event " + << ibv_event_type_str(event_type) << " for context " << context_.deviceName(); - if (event.event_type == IBV_EVENT_QP_FATAL) { - auto endpoint_ptr = (RdmaEndPoint *)event.element.qp->qp_context; - /** - * There might be a deadlock if we call endpoint->set_active(false) - * before ack the event: - * - * Thread A: - * Holding endpoint->lock_ and calling ibv_destroy_qp (if using - * eRDMA), ibv_destroy_qp will block until the event is acked. - * - * Thread B (this thread): - * Calling endpoint->set_active(false), which blocks as - * endpoint->lock_ is held by Thread A. - */ - ibv_ack_async_event(&event); - event_acked = true; + handleContextEvent(event_type, true); +} - /** - * After ack the event, the endpoint might be destroyed if it happened - * to be destroying event.element.qp. Therefore, we cannot just - * dereference endpoint_ptr. Instead, we need to get the shared_ptr of - * the endpoint from context_ and use that shared_ptr to access the - * endpoint. - */ - context_.deleteEndpointByPtr(endpoint_ptr); - } else if (event.event_type == IBV_EVENT_DEVICE_FATAL || - event.event_type == IBV_EVENT_CQ_ERR || - event.event_type == IBV_EVENT_WQ_FATAL || - event.event_type == IBV_EVENT_PORT_ERR || - event.event_type == IBV_EVENT_LID_CHANGE) { +bool WorkerPool::handleContextEvent(ibv_event_type event_type, + bool injected_for_test, + struct ibv_async_event *event) { + if (event_type == IBV_EVENT_DEVICE_FATAL || + event_type == IBV_EVENT_CQ_ERR || event_type == IBV_EVENT_WQ_FATAL || + event_type == IBV_EVENT_PORT_ERR || + event_type == IBV_EVENT_LID_CHANGE) { + recovery_activate_after_ns_.store(0, std::memory_order_relaxed); context_.set_active(false); + refreshPublishedLocalTopology(); /** * Similar deadlock might happen if we call @@ -525,39 +1060,281 @@ int WorkerPool::doProcessContextEvents() { * Calling endpoint->disconnect(), which blocks as endpoint->lock_ * is held by Thread A. */ - ibv_ack_async_event(&event); - event_acked = true; + if (event != nullptr) ibv_ack_async_event(event); context_.disconnectAllEndpoints(); LOG(INFO) << "Worker: Context " << context_.deviceName() - << " is now inactive"; - } else if (event.event_type == IBV_EVENT_PORT_ACTIVE) { - context_.set_active(true); - markContextSuccess(); // Reset failure counter on port recovery + << " is now inactive due to " + << (injected_for_test ? "injected fatal event: " + : "fatal event: ") + << event_type; + } else if (event_type == IBV_EVENT_GID_CHANGE) { + auto gid_refresh_result = refreshPublishedLocalGid(); + if (event != nullptr) ibv_ack_async_event(event); + + if (gid_refresh_result != GidRefreshResult::UNCHANGED) { + context_.disconnectAllEndpoints(); + LOG(INFO) << "Worker: Context " << context_.deviceName() + << (injected_for_test ? " injected GID refresh result=" + : " GID refresh result=") + << static_cast(gid_refresh_result) + << ", disconnected all endpoints"; + } + } else if (event_type == IBV_EVENT_PORT_ACTIVE) { + // PORT_ACTIVE only means the link started coming back. Real mlx5/RoCE + // data path can still reject RTR for a while after link-up, so delay + // publishing this local RNIC back to metadata. Injected tests follow + // the same path. + scheduleContextRecovery(); + if (event != nullptr) ibv_ack_async_event(event); + } else { + return false; + } + + return true; +} + +void WorkerPool::scheduleContextRecovery(uint64_t delay_ns) { + uint64_t activate_after = getCurrentTimeInNano() + delay_ns; + recovery_activate_after_ns_.store(activate_after, + std::memory_order_relaxed); + LOG(INFO) << "Worker: Context " << context_.deviceName() + << " scheduled recovery probe after " << delay_ns / 1000000000ull + << " seconds"; +} + +void WorkerPool::maybeActivateRecoveredContext() { + uint64_t activate_after = + recovery_activate_after_ns_.load(std::memory_order_relaxed); + if (activate_after == 0 || + static_cast(getCurrentTimeInNano()) < activate_after) + return; + + uint64_t expected = activate_after; + if (!recovery_activate_after_ns_.compare_exchange_strong( + expected, 0, std::memory_order_relaxed)) { + return; + } + + auto gid_refresh_result = refreshPublishedLocalGid(); + if (gid_refresh_result == GidRefreshResult::FAILED) { + context_.set_active(false); + refreshPublishedLocalTopology(); + scheduleContextRecovery(); + LOG(WARNING) << "Worker: Context " << context_.deviceName() + << " failed to refresh GID during recovery; " + "keeping inactive"; + return; + } + if (gid_refresh_result == GidRefreshResult::CHANGED) { + context_.disconnectAllEndpoints(); LOG(INFO) << "Worker: Context " << context_.deviceName() - << " is now active"; + << " GID changed during recovery, disconnected all endpoints"; } - if (!event_acked) { - ibv_ack_async_event(&event); + context_.set_active(true); + refreshPublishedLocalTopology(); + context_failure_count_.store(0, std::memory_order_relaxed); + LOG(INFO) << "Worker: Context " << context_.deviceName() + << " is now active after recovery delay"; +} + +bool WorkerPool::hasAvailablePeerRailAlternative( + Transport::Slice *slice, const std::string &failed_peer_path) { + auto peer_segment_desc = + context_.engine().meta()->getSegmentDescByID(slice->target_id, false); + if (!peer_segment_desc) return false; + + int buffer_id = -1; + for (size_t idx = 0; idx < peer_segment_desc->buffers.size(); ++idx) { + auto &buffer = peer_segment_desc->buffers[idx]; + uint64_t buffer_start = reinterpret_cast(buffer.addr); + uint64_t buffer_end = buffer_start + buffer.length; + if (buffer_start <= slice->rdma.dest_addr && + slice->rdma.dest_addr + slice->length <= buffer_end) { + buffer_id = static_cast(idx); + break; + } } + if (buffer_id < 0) return false; - return 0; + auto server_name = peer_segment_desc->nicPathServerName(); + for (size_t dev_id = 0; dev_id < peer_segment_desc->devices.size(); + ++dev_id) { + if (dev_id >= peer_segment_desc->buffers[buffer_id].rkey.size()) { + continue; + } + auto peer_path = + MakeNicPath(server_name, peer_segment_desc->devices[dev_id].name); + if (peer_path != failed_peer_path && isRailAvailable(peer_path)) { + return true; + } + } + return false; +} + +void WorkerPool::refreshPublishedLocalTopology() { + std::lock_guard guard(context_.engine().local_desc_lock_); + auto desc = + context_.engine().metadata_->getSegmentDescByID(LOCAL_SEGMENT_ID); + if (!desc || !context_.engine().local_topology_) return; + + auto updated_desc = std::make_shared(*desc); + updated_desc->topology = *context_.engine().local_topology_; + for (const auto &context : context_.engine().context_list_) { + if (context->active()) continue; + updated_desc->topology.disableDevice(context->deviceName()); + } + + context_.engine().metadata_->addLocalSegment( + LOCAL_SEGMENT_ID, updated_desc->name, std::move(updated_desc)); + int ret = context_.engine().metadata_->updateLocalSegmentDesc(); + if (ret) { + LOG(WARNING) << "Failed to publish RDMA topology update for " + << context_.deviceName() << ", ret=" << ret; + } +} + +GidRefreshResult WorkerPool::refreshPublishedLocalGid() { + std::string previous_gid; + std::string next_gid; + auto result = context_.refreshCurrentGid(&previous_gid, &next_gid); + if (result == GidRefreshResult::CHANGED) { + LOG(WARNING) << "Worker: refreshed published GID for " + << context_.deviceName() << ": " << previous_gid << " -> " + << next_gid; + } else if (result == GidRefreshResult::UNCHANGED) { + LOG(INFO) << "Worker: received GID change event for " + << context_.deviceName() << ", current GID is unchanged"; + } else { + LOG(ERROR) << "Worker: failed to refresh published GID for " + << context_.deviceName() + << ", disconnecting endpoints to avoid stale GID reuse"; + } + return result; } void WorkerPool::monitorWorker() { bindToSocket(numa_socket_id_); auto last_reset_ts = getCurrentTimeInNano(); + uint64_t outstanding_since_ns = 0; + uint64_t last_timeout_log_ns = 0; + uint64_t last_processed_count = + processed_slice_count_.load(std::memory_order_relaxed); while (workers_running_) { - auto current_ts = getCurrentTimeInNano(); + const uint64_t current_ts = + static_cast(getCurrentTimeInNano()); + maybeActivateRecoveredContext(); if (current_ts - last_reset_ts > 1000000000ll) { // Drain endpoint_store_->waiting_list_ even when no new // insertions are happening. Without this, reclaim only runs // from RdmaContext::endpoint() and the waiting list grows // unboundedly under failure load. See issue #1845. context_.reclaimEndpoints(); + // Drop expired active-connect pause entries so the map doesn't grow + // for peers that are never re-attempted after their pause lapses. + context_.pruneConnectPause(); last_reset_ts = current_ts; } + + int64_t cq_outstanding = 0; + for (int cq_index = 0; cq_index < context_.cqCount(); ++cq_index) { + cq_outstanding += context_.cqOutstandingCount(cq_index)->load( + std::memory_order_relaxed); + } + const uint64_t processed_count = + processed_slice_count_.load(std::memory_order_relaxed); + if (processed_count != last_processed_count) { + last_processed_count = processed_count; + outstanding_since_ns = + cq_outstanding > 0 ? current_ts : static_cast(0); + } + if (cq_outstanding > 0) { + if (outstanding_since_ns == 0) outstanding_since_ns = current_ts; + + const uint64_t outstanding_age_ns = + current_ts - outstanding_since_ns; + const uint64_t last_poll_ts = + last_poll_ts_ns_.load(std::memory_order_relaxed); + const uint64_t poll_gap_ns = + last_poll_ts > 0 && current_ts > last_poll_ts + ? current_ts - last_poll_ts + : 0; + + // Log a stalled poller quickly, and also log at the same 30-second + // boundary used by TransferEnginePy when polling continues. + const bool poll_stalled = poll_gap_ns >= 5ULL * 1000 * 1000 * 1000; + const bool transfer_timed_out = + outstanding_age_ns >= 30ULL * 1000 * 1000 * 1000; + if ((poll_stalled || transfer_timed_out) && + current_ts - last_timeout_log_ns >= 5ULL * 1000 * 1000 * 1000) { + LOG(ERROR) + << "CQ completion timeout diagnostic: context=" + << context_.deviceName() + << ", outstanding=" << cq_outstanding + << ", outstanding_age_ms=" << outstanding_age_ns / 1000000 + << ", poll_gap_ms=" << poll_gap_ns / 1000000 + << ", last_poll_interval_ms=" + << last_poll_interval_ns_.load(std::memory_order_relaxed) / + 1000000 + << ", max_poll_interval_ms=" + << max_poll_interval_ns_.load(std::memory_order_relaxed) / + 1000000 + << ", submitted=" + << submitted_slice_count_.load(std::memory_order_relaxed) + << ", processed=" + << processed_slice_count_.load(std::memory_order_relaxed); + + if (globalConfig().track_rdma_posted_slices) { + struct StuckGroup { + size_t slice_count = 0; + uint64_t total_bytes = 0; + uint64_t oldest_post_ts = 0; + void *sample_source_addr = nullptr; + uint64_t sample_dest_addr = 0; + }; + std::unordered_map stuck_groups; + { + std::lock_guard lock(posted_slices_mutex_); + for (auto *slice : posted_slices_) { + auto &group = stuck_groups[slice->peer_nic_path]; + group.slice_count++; + group.total_bytes += slice->length; + if (group.oldest_post_ts == 0 || + static_cast(slice->ts) < + group.oldest_post_ts) { + group.oldest_post_ts = + static_cast(slice->ts); + group.sample_source_addr = slice->source_addr; + group.sample_dest_addr = slice->rdma.dest_addr; + } + } + } + for (const auto &entry : stuck_groups) { + const auto &group = entry.second; + const uint64_t oldest_age_ms = + group.oldest_post_ts > 0 && + current_ts > group.oldest_post_ts + ? (current_ts - group.oldest_post_ts) / 1000000 + : 0; + LOG(ERROR) + << "CQ stuck transfer group: context=" + << context_.deviceName() + << ", peer_nic=" << entry.first + << ", slices=" << group.slice_count + << ", bytes=" << group.total_bytes + << ", oldest_post_age_ms=" << oldest_age_ms + << ", sample_source_addr=" + << group.sample_source_addr << ", sample_dest_addr=" + << reinterpret_cast(group.sample_dest_addr); + } + } + last_timeout_log_ns = current_ts; + } + } else { + outstanding_since_ns = 0; + } + struct epoll_event event; int num_events = epoll_wait(context_.eventFd(), &event, 1, 100); if (num_events < 0) { @@ -575,15 +1352,22 @@ void WorkerPool::monitorWorker() { } } -void WorkerPool::markRailFailed(const std::string &peer_nic_path) { +void WorkerPool::markRailFailed(const std::string &peer_nic_path, + bool immediate_pause) { std::lock_guard lock(rail_state_lock_); auto &state = rail_states_[peer_nic_path]; + uint64_t now = getCurrentTimeInNano(); state.error_count++; + if (immediate_pause && state.error_count < kRailErrorThreshold) { + state.error_count = kRailErrorThreshold; + } if (state.error_count >= kRailErrorThreshold) { - uint64_t now = getCurrentTimeInNano(); - state.pause_until_ns = now + kRailPauseNs; + const uint64_t rail_pause_ns = + globalConfig().rdma_rail_pause_seconds * 1000000000ull; + state.pause_until_ns = now + rail_pause_ns; LOG(WARNING) << "Rail paused: peer=" << peer_nic_path - << " error_count=" << state.error_count; + << " error_count=" << state.error_count + << " pause_ms=" << rail_pause_ns / 1000000ull; } } @@ -610,14 +1394,49 @@ bool WorkerPool::shouldRetrySlice(Transport::Slice *slice) { return slice->rdma.retry_cnt < slice->rdma.max_retry_cnt; } -// Unified path failure handler -void WorkerPool::handlePathFailure(const std::string &peer_nic_path, - RdmaEndPoint *endpoint) { - markRailFailed(peer_nic_path); - redispatch_counter_++; // Notify all workers to redispatch their queues +bool WorkerPool::isLocalWcFailure(const ibv_wc &wc) { + // IBV_WC_GENERAL_ERR is intentionally not treated as a local RNIC failure. + // Providers use it for broad connection/path failures too, and disabling + // the local context here can mask endpoint GID reprobe and remote rail + // recovery paths. + switch (wc.status) { + case IBV_WC_LOC_LEN_ERR: + case IBV_WC_LOC_QP_OP_ERR: + case IBV_WC_LOC_PROT_ERR: + case IBV_WC_MW_BIND_ERR: + case IBV_WC_LOC_ACCESS_ERR: +#ifdef IBV_WC_LOC_RDD_VIOL_ERR + case IBV_WC_LOC_RDD_VIOL_ERR: +#endif +#ifdef IBV_WC_LOC_EEC_OP_ERR + case IBV_WC_LOC_EEC_OP_ERR: +#endif +#ifdef IBV_WC_LOC_EEC_STATE_ERR + case IBV_WC_LOC_EEC_STATE_ERR: +#endif + return true; + + default: + return false; + } +} + +void WorkerPool::handleLocalFailure(const std::string &peer_nic_path, + RdmaEndPoint *endpoint) { + // Local completion faults can be caused by a poisoned QP/MR as well as a + // bad RNIC. Retry this slice elsewhere, but only disable the whole context + // after repeated local failures or an async port/device event. + bool context_disabled = markContextFailure(); + if (context_disabled) refreshPublishedLocalTopology(); + redispatch_counter_++; + + // Endpoint may also be poisoned; retire it for safety. if (endpoint) { context_.deleteEndpointByPtr(endpoint); } + + LOG(WARNING) << "Local-side RDMA failure detected on context " + << context_.deviceName() << ", peer=" << peer_nic_path; } } // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt b/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt index 841a36aa74..3b6bf0cd25 100644 --- a/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt @@ -1,6 +1,6 @@ file(GLOB CORO_RPC_SOURCES "*.cpp") -find_package(Python3 COMPONENTS Interpreter Development REQUIRED) +find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED) add_library(rpc_communicator OBJECT ${CORO_RPC_SOURCES}) @@ -10,7 +10,6 @@ target_link_libraries(rpc_communicator yalantinglibs::yalantinglibs glog::glog pthread - ${Python3_LIBRARIES} ) target_include_directories(rpc_communicator diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp index 158725f0ed..18d3c2ad03 100644 --- a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp @@ -1,4 +1,5 @@ #include "transport/rpc_communicator/rpc_communicator.h" +#include #include #include #include @@ -11,11 +12,12 @@ #include #include "async_simple/coro/SyncAwait.h" #include "default_config.h" +#include "transfer_engine_rpc_client_io_context.h" namespace mooncake { namespace py = pybind11; -class py_rpc_context { +class __attribute__((visibility("hidden"))) py_rpc_context { public: void response_msg(py::buffer msg, py::object done) { py::buffer_info info = msg.request(); @@ -33,7 +35,12 @@ class py_rpc_context { coro_rpc::context context_; }; -RpcCommunicator::RpcCommunicator() {} +struct __attribute__((visibility("hidden"))) RpcCommunicator::PyCallbackHolder { + py::handle callback; +}; + +RpcCommunicator::RpcCommunicator() + : py_callback_(std::make_unique()) {} RpcCommunicator::~RpcCommunicator() { stopServer(); } @@ -55,12 +62,20 @@ bool RpcCommunicator::initialize(const RpcCommunicatorConfig& config) { pool_conf.client_config.socket_config = coro_io::ib_socket_t::config_t{}; } + if (config.pool_size > 0 && + config.pool_size <= std::numeric_limits::max()) { + pool_conf.max_connection = static_cast(config.pool_size); + } else { + LOG(WARNING) << "Invalid RPC client per-target pool_size " + << config.pool_size << "; using default " + << pool_conf.max_connection; + } client_pools_ = std::make_shared>( - pool_conf); + pool_conf, GetTransferEngineRpcClientIoContextPool()); - LOG(INFO) << "create coro_rpc_client_pool with " << config.pool_size - << " threads"; + LOG(INFO) << "Created coro_rpc client pools with up to " + << pool_conf.max_connection << " cached connections per target"; if (!config.listen_address.empty()) { LOG(INFO) << "Initializing server on " << config.listen_address; @@ -395,7 +410,7 @@ void RpcCommunicator::handleDataTransferWithAttachment( auto view = py::memoryview::from_buffer(data.data(), {data.size()}, {sizeof(char)}); - py_callback_(std::move(t), view); + py_callback_->callback(std::move(t), view); } void RpcCommunicator::handleTensorTransferWithAttachment( @@ -412,7 +427,7 @@ void RpcCommunicator::handleTensorTransferWithAttachment( auto view = py::memoryview::from_buffer( attachment.data(), {attachment.size()}, {sizeof(int8_t)}); - py_callback_(std::move(t), view); + py_callback_->callback(std::move(t), view); } -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp index 13d44f6ff7..2a3b712c1b 100644 --- a/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp +++ b/mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp @@ -16,7 +16,7 @@ static constexpr size_t MAX_TENSOR_DIMS = 4; static constexpr size_t TENSOR_METADATA_SIZE = 4 + 4 + MAX_TENSOR_DIMS * 8; // Implementation class -class RpcInterface::Impl { +class __attribute__((visibility("hidden"))) RpcInterface::Impl { public: std::unique_ptr communicator; pybind11::function data_receive_callback; @@ -52,7 +52,7 @@ bool RpcInterface::initializeClient(size_t pool_size, size_t timeout_seconds) { bool RpcInterface::initializeServer(const std::string& listen_address, size_t thread_count, size_t timeout_seconds) { - return initialize(listen_address, thread_count, timeout_seconds, 4); + return initialize(listen_address, thread_count, timeout_seconds, 100); } bool RpcInterface::startServer() { @@ -558,9 +558,9 @@ void bind_rpc_interface(pybind11::module_& m) { .def(py::init<>()) .def("initialize", &RpcInterface::initialize, py::arg("listen_address") = "", py::arg("thread_count") = 0, - py::arg("timeout_seconds") = 30, py::arg("pool_size") = 10) + py::arg("timeout_seconds") = 30, py::arg("pool_size") = 100) .def("initialize_client", &RpcInterface::initializeClient, - py::arg("pool_size") = 10, py::arg("timeout_seconds") = 30) + py::arg("pool_size") = 100, py::arg("timeout_seconds") = 30) .def("initialize_server", &RpcInterface::initializeServer, py::arg("listen_address"), py::arg("thread_count") = 8, py::arg("timeout_seconds") = 30) diff --git a/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp b/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp index f3c3d7eada..ebf3f654b7 100644 --- a/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp +++ b/mooncake-transfer-engine/src/transport/sunrise_link/sunrise_link_transport.cpp @@ -1,6 +1,7 @@ // Copyright 2024 KVCache.AI #include "transport/sunrise_link_transport/sunrise_link_transport.h" +#include "sunrise_allocator.h" #include #include @@ -139,6 +140,13 @@ static tangError_t QueryPointerAttrsBestEffort(void* ptr, tangPointerAttributes* attr, int preferred_dev) { if (!ptr || !attr) return tangErrorInvalidValue; + + if (!sunrise_is_device_memory_range(ptr)) { + attr->type = tangMemoryTypeHost; + attr->device = -1; + return tangSuccess; + } + SavedTangDevice saved; if (preferred_dev >= 0) { @@ -148,14 +156,22 @@ static tangError_t QueryPointerAttrsBestEffort(void* ptr, tangError_t ret = tangPointerGetAttributes(attr, ptr); if (ret == tangSuccess) return ret; + if (preferred_dev < 0) { + return ret; + } + int dev_count = 0; if (tangGetDeviceCount(&dev_count) != tangSuccess || dev_count <= 0) { return ret; } for (int d = 0; d < dev_count; ++d) { if (d == preferred_dev) continue; - tangSetDevice(d); - ret = tangPointerGetAttributes(attr, ptr); + { + SavedTangDevice loop_saved; + tangError_t sd = tangSetDevice(d); + if (sd != tangSuccess) continue; + ret = tangPointerGetAttributes(attr, ptr); + } if (ret == tangSuccess) return ret; } return ret; @@ -336,39 +352,42 @@ int SunriseLinkTransport::registerLocalMemory(void* addr, size_t length, int preferred_dev = ParseSunriseDeviceId(location); tangPointerAttributes attr = {}; tangError_t err = QueryPointerAttrsBestEffort(addr, &attr, preferred_dev); - if (err != tangSuccess || attr.type != tangMemoryTypeDevice) { - LOG(ERROR) << "SunriseLinkTransport: unsupported memory for register" - << " addr=" << addr << " location=" << location - << " preferred_dev=" << preferred_dev << " err=" << err - << " type=" << attr.type << " device=" << attr.device; - return -1; - } - tangIpcMemHandle_t handle; - int saved_dev = -1; - tangGetDevice(&saved_dev); - int handle_dev = preferred_dev >= 0 ? preferred_dev : attr.device; - if (handle_dev >= 0) { - tangError_t sd = tangSetDevice(handle_dev); - if (sd != tangSuccess) { - LOG(ERROR) << "SunriseLinkTransport: tangSetDevice before " - << "tangIpcGetMemHandle failed: " << sd << " " - << tangGetErrorString(sd) << " device=" << handle_dev; - if (saved_dev >= 0) tangSetDevice(saved_dev); - return -1; - } - } - err = tangIpcGetMemHandle(&handle, addr); - if (saved_dev >= 0) tangSetDevice(saved_dev); + bool is_device_mem = + (err == tangSuccess && attr.type == tangMemoryTypeDevice); + int mem_dev = + is_device_mem ? attr.device : (preferred_dev >= 0 ? preferred_dev : 0); std::string shm_name; - if (err == tangSuccess) { - shm_name = serializeBinaryData(&handle, sizeof(handle)); + if (is_device_mem) { + tangIpcMemHandle_t handle; + int saved_dev = -1; + tangGetDevice(&saved_dev); + int handle_dev = preferred_dev >= 0 ? preferred_dev : attr.device; + if (handle_dev >= 0) { + tangError_t sd = tangSetDevice(handle_dev); + if (sd != tangSuccess) { + LOG(ERROR) << "SunriseLinkTransport: tangSetDevice before " + << "tangIpcGetMemHandle failed: " << sd << " " + << tangGetErrorString(sd) + << " device=" << handle_dev; + if (saved_dev >= 0) tangSetDevice(saved_dev); + return -1; + } + } + err = tangIpcGetMemHandle(&handle, addr); + if (saved_dev >= 0) tangSetDevice(saved_dev); + + if (err == tangSuccess) { + shm_name = serializeBinaryData(&handle, sizeof(handle)); + } else { + LOG(WARNING) << "SunriseLinkTransport: tangIpcGetMemHandle failed: " + << err << " " << tangGetErrorString(err) + << ", falling back to RAW_ADDR"; + shm_name = kRawAddrPrefix + std::to_string(attr.device); + } } else { - LOG(WARNING) << "SunriseLinkTransport: tangIpcGetMemHandle failed: " - << err << " " << tangGetErrorString(err) - << ", falling back to RAW_ADDR"; - shm_name = kRawAddrPrefix + std::to_string(attr.device); + shm_name = kRawAddrPrefix + std::to_string(mem_dev); } { @@ -376,7 +395,7 @@ int SunriseLinkTransport::registerLocalMemory(void* addr, size_t length, if (registered_regions_.count(addr)) { return 0; } - registered_regions_[addr] = RegisteredRegion{length, attr.device}; + registered_regions_[addr] = RegisteredRegion{length, mem_dev}; } BufferDesc desc; @@ -399,15 +418,22 @@ int SunriseLinkTransport::unregisterLocalMemory(void* addr, int SunriseLinkTransport::registerLocalMemoryBatch( const std::vector& buffer_list, const std::string& location) { for (const auto& buffer : buffer_list) { - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; } return metadata_->updateLocalSegmentDesc(); } int SunriseLinkTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { - for (auto* addr : addr_list) unregisterLocalMemory(addr, false); - return metadata_->updateLocalSegmentDesc(); + int first_error = 0; + for (auto* addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret && !first_error) first_error = ret; + } + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } int SunriseLinkTransport::relocateSharedMemoryAddress(uint64_t& dest_addr, @@ -501,35 +527,86 @@ int SunriseLinkTransport::relocateSharedMemoryAddress(uint64_t& dest_addr, int SunriseLinkTransport::startCopy(void* src, void* dst, size_t length, int remote_dev, int local_dev) { - tangPointerAttributes src_attr = {}; - tangPointerAttributes dst_attr = {}; - tangError_t src_ret = - QueryPointerAttrsBestEffort(src, &src_attr, local_dev); - tangError_t dst_ret = - QueryPointerAttrsBestEffort(dst, &dst_attr, remote_dev); - NormalizeMappedHostPointer(&src, &src_attr); - NormalizeMappedHostPointer(&dst, &dst_attr); - - int src_dev = - (src_ret == tangSuccess && src_attr.type == tangMemoryTypeDevice) - ? src_attr.device - : -1; - int dst_dev = - (dst_ret == tangSuccess && dst_attr.type == tangMemoryTypeDevice) - ? dst_attr.device - : -1; - - tangMemoryType src_type = - src_ret == tangSuccess ? src_attr.type : tangMemoryTypeHost; - tangMemoryType dst_type = - dst_ret == tangSuccess ? dst_attr.type : tangMemoryTypeHost; + bool src_is_host_alloc = sunrise_is_host_allocated(src); + bool dst_is_host_alloc = sunrise_is_host_allocated(dst); + bool src_is_dev = sunrise_is_device_memory_range(src); + bool dst_is_dev = sunrise_is_device_memory_range(dst); + + VLOG(1) << "startCopy: src=" << src << " dst=" << dst << " len=" << length + << " src_is_dev=" << src_is_dev << " dst_is_dev=" << dst_is_dev + << " src_is_host_alloc=" << src_is_host_alloc + << " dst_is_host_alloc=" << dst_is_host_alloc; + + bool any_tang = + src_is_dev || dst_is_dev || src_is_host_alloc || dst_is_host_alloc; + if (!any_tang) { + memcpy(dst, src, length); + return 0; + } + + if (src_is_dev && dst_is_host_alloc && !dst_is_dev) { + SavedTangDevice saved; + tangSetDevice(0); + std::vector staging(length); + tangError_t ret = + tangMemcpy(staging.data(), src, length, tangMemcpyDeviceToHost); + if (ret != tangSuccess) { + LOG(ERROR) << "startCopy: tangMemcpy D2H staging failed"; + return -1; + } + memcpy(dst, staging.data(), length); + return 0; + } - tangError_t ret = - doTangCopy(dst, dst_dev, src, src_dev, length, src_type, dst_type); + if (src_is_host_alloc && !src_is_dev && dst_is_dev) { + SavedTangDevice saved; + tangSetDevice(0); + std::vector staging(length); + memcpy(staging.data(), src, length); + tangError_t ret = + tangMemcpy(dst, staging.data(), length, tangMemcpyHostToDevice); + if (ret != tangSuccess) { + LOG(ERROR) << "startCopy: tangMemcpy H2D staging failed"; + return -1; + } + return 0; + } + + if ((src_is_host_alloc || dst_is_host_alloc) && !src_is_dev && + !dst_is_dev) { + memcpy(dst, src, length); + return 0; + } + + SavedTangDevice saved; + tangSetDevice(0); + + tangMemcpyKind kind = tangMemcpyHostToHost; + if ((src_is_dev || src_is_host_alloc) && + (dst_is_dev || dst_is_host_alloc)) { + kind = tangMemcpyDeviceToDevice; + } else if (src_is_dev || src_is_host_alloc) { + kind = tangMemcpyDeviceToHost; + } else if (dst_is_dev || dst_is_host_alloc) { + kind = tangMemcpyHostToDevice; + } + + if (kind == tangMemcpyHostToHost) { + memcpy(dst, src, length); + return 0; + } + + tangError_t ret = tangMemcpy(dst, src, length, kind); + if (dst_is_host_alloc || (kind == tangMemcpyDeviceToHost)) { + tangDeviceSynchronize(); + } if (ret != tangSuccess) { - LOG(ERROR) << "SunriseLinkTransport: tang copy failed: " << ret << " " - << tangGetErrorString(ret) << " src_dev=" << src_dev - << " dst_dev=" << dst_dev; + LOG(ERROR) << "SunriseLinkTransport: tangMemcpy failed: " << ret << " " + << tangGetErrorString(ret) << " kind=" << kind + << " src_is_dev=" << src_is_dev + << " dst_is_dev=" << dst_is_dev + << " src_is_host_alloc=" << src_is_host_alloc + << " dst_is_host_alloc=" << dst_is_host_alloc; return -1; } return 0; diff --git a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp index dcb0e55682..c2aa41bbcd 100644 --- a/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp +++ b/mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -26,6 +27,7 @@ #include #include #include +#include #include #include "common.h" @@ -42,8 +44,21 @@ static size_t getChunkSize() { static const size_t val = [] { const char* env = std::getenv("MC_TCP_SLICE_SIZE"); if (env) { - size_t v = std::stoull(env); - if (v > 0) return v; + try { + size_t v = std::stoull(env); + if (v > 0) return v; + LOG(WARNING) + << "Ignore non-positive MC_TCP_SLICE_SIZE value: " << env + << ", using default 65536"; + } catch (const std::exception& e) { + // A non-numeric or out-of-range value makes std::stoull throw; + // fall through to the default instead of letting the exception + // propagate out of this static initializer and abort the + // transfer that first reads the chunk size. + LOG(WARNING) + << "Invalid MC_TCP_SLICE_SIZE value: " << env + << ". Error: " << e.what() << ", using default 65536"; + } } return size_t(65536); // 64KB default }(); @@ -100,7 +115,45 @@ class TcpTransport; using ValidateAddrFn = std::function; -// Server-side session: handles one transfer request on a persistent connection +// --- Acknowledged framing (protocol v2, #2086) ------------------------------ +// v1 framing gives the initiator no channel to learn whether the receiver +// applied (or even accepted) a WRITE: COMPLETED fires when the final chunk +// enters the initiator's kernel socket buffer, while megabytes may still be +// in flight toward destination memory, and a rejected request is silently +// "successful". v2 requests set the high bit of the opcode; the server then +// (a) prefixes every READ response with an 8-byte status frame and (b) sends +// an 8-byte status frame for WRITE only after the final chunk has been +// applied to destination memory. Initiators enable v2 only when the target +// segment advertises tcp_proto_version >= 2, so old servers never see +// flagged opcodes and old initiators keep receiving v1 framing. +static constexpr uint8_t kOpcodeV2Flag = 0x80; +// Status frames carry a magic in the high 32 bits so that a v2 initiator +// which reaches a v1 server through a stale descriptor (v1 treats unknown +// opcodes as READ and immediately streams payload bytes) fails fast on the +// first frame instead of misinterpreting the stream. Residual risk: payload +// bytes that happen to equal a valid frame (2^-64 per request, data +// dependent) are indistinguishable in-band; eliminating that would need a +// nonce/checksum handshake, which this deliberately avoids. +static constexpr uint64_t kStatusMagic = 0x4D435456ull << 32; // "MCTV" +static constexpr uint64_t kStatusOk = kStatusMagic | 0; +static constexpr uint64_t kStatusAddrRejected = kStatusMagic | 1; +static inline bool statusFrameValid(uint64_t frame) { + return (frame & 0xFFFFFFFF00000000ull) == kStatusMagic; +} + +// Operational escape hatch: MC_TCP_PROTO=1 forces initiators to speak the +// legacy unacknowledged framing even to v2-capable servers. Also used by +// tests to cover the mixed-version matrix in one process. +static bool forceLegacyTcpProto() { + // Read per call (startTransfer already does metadata lookups; getenv is + // noise) so tests can cover both protocol modes in one process. + const char* env = std::getenv("MC_TCP_PROTO"); + return env && env[0] == '1' && env[1] == '\0'; +} + +// Server-side session: handles transfer requests on a persistent connection. +// The session owns the socket; ending the callback chain without rearming +// (start()/next handler) drops the last reference and closes the connection. struct ServerSession : public std::enable_shared_from_this { explicit ServerSession(std::shared_ptr socket, ValidateAddrFn validate_addr) @@ -112,16 +165,30 @@ struct ServerSession : public std::enable_shared_from_this { SessionHeader header_; uint64_t total_transferred_bytes_; char* local_buffer_; - std::function on_finalize_; - std::mutex session_mutex_; + bool v2_ = false; + uint64_t status_frame_; void start() { - session_mutex_.lock(); total_transferred_bytes_ = 0; readHeader(); } private: + // Send an 8-byte status frame, then run `next` (or end the session — + // closing the connection — when `next` is empty or the send fails). + void sendStatus(uint64_t status, std::function next) { + status_frame_ = htole64(status); + auto self(shared_from_this()); + asio::async_write(*socket_, + asio::buffer(&status_frame_, sizeof(status_frame_)), + [this, self, next = std::move(next)]( + const asio::error_code& ec, std::size_t) { + if (ec) + return; // connection closes with the session + if (next) next(); + }); + } + void readHeader() { auto self(shared_from_this()); asio::async_read( @@ -134,10 +201,11 @@ struct ServerSession : public std::enable_shared_from_this { << ec.message() << " (value: " << ec.value() << ")" << ", bytes read: " << len; } - session_mutex_.unlock(); return; } + v2_ = (header_.opcode & kOpcodeV2Flag) != 0; + const uint8_t opcode = header_.opcode & ~kOpcodeV2Flag; local_buffer_ = (char*)(le64toh(header_.addr)); uint64_t size = le64toh(header_.size); if (validate_addr_ && @@ -146,13 +214,21 @@ struct ServerSession : public std::enable_shared_from_this { << std::hex << (uint64_t)local_buffer_ << std::dec << " with size " << size << " is not within any registered buffer"; - session_mutex_.unlock(); + // v2 initiators learn of the rejection; v1 initiators + // only see the connection close (and, for small WRITEs, + // may have already reported success — the defect v2 + // exists to fix). + if (v2_) sendStatus(kStatusAddrRejected, nullptr); return; } - if (header_.opcode == (uint8_t)TransferRequest::WRITE) + if (opcode == (uint8_t)TransferRequest::WRITE) { readBody(); - else + } else if (v2_) { + // READ, v2: status frame precedes the data. + sendStatus(kStatusOk, [this] { writeBody(); }); + } else { writeBody(); + } }); } @@ -164,7 +240,6 @@ struct ServerSession : public std::enable_shared_from_this { size_t buffer_size = std::min(getChunkSize(), size - total_transferred_bytes_); if (buffer_size == 0) { - session_mutex_.unlock(); // Transfer complete, wait for next request on this connection start(); return; @@ -192,7 +267,6 @@ struct ServerSession : public std::enable_shared_from_this { LOG(ERROR) << "ServerSession::writeBody failed to copy from " "CUDA memory. " << "Error: " << cudaGetErrorString(cuda_status); - session_mutex_.unlock(); delete[] dram_buffer; return; // Connection will be closed } @@ -217,7 +291,6 @@ struct ServerSession : public std::enable_shared_from_this { << " using buffer " << static_cast(dram_buffer) << ". Error: " << ec.message() << " (value: " << ec.value() << ")"; - session_mutex_.unlock(); return; // Connection will be closed } total_transferred_bytes_ += transferred_bytes; @@ -233,9 +306,15 @@ struct ServerSession : public std::enable_shared_from_this { size_t buffer_size = std::min(getChunkSize(), size - total_transferred_bytes_); if (buffer_size == 0) { - session_mutex_.unlock(); - // Transfer complete, wait for next request on this connection - start(); + // Destination memory now holds the complete payload. Under v2, + // acknowledge before accepting the next request — this is what + // makes the initiator's COMPLETED mean "applied at the + // destination" rather than "left my socket buffer". + if (v2_) { + sendStatus(kStatusOk, [this] { start(); }); + } else { + start(); + } return; } @@ -267,7 +346,6 @@ struct ServerSession : public std::enable_shared_from_this { << ". Error: " << ec.message() << " (value: " << ec.value() << ")"; } - session_mutex_.unlock(); if (cuda_device >= 0) delete[] dram_buffer; return; // Connection will be closed } @@ -292,7 +370,6 @@ struct ServerSession : public std::enable_shared_from_this { "memory. " << "Error: " << cudaGetErrorString(cuda_status); delete[] dram_buffer; - session_mutex_.unlock(); return; // Connection will be closed } delete[] dram_buffer; @@ -306,30 +383,133 @@ struct ServerSession : public std::enable_shared_from_this { // Client-side session: initiates one transfer request struct ClientSession : public std::enable_shared_from_this { - explicit ClientSession(std::shared_ptr socket, - std::function on_complete = nullptr) - : socket_(std::move(socket)), on_complete_(std::move(on_complete)) {} + explicit ClientSession(std::shared_ptr socket, bool use_v2, + std::function on_complete = nullptr) + : socket_(std::move(socket)), + v2_(use_v2), + on_complete_(std::move(on_complete)) {} std::shared_ptr socket_; SessionHeader header_; uint64_t total_transferred_bytes_; char* local_buffer_; + bool v2_; + uint64_t status_frame_; + // v2 WRITE runs the body stream and the ack read concurrently (one + // async op per direction; handlers serialize on the io thread). The + // concurrent read lets a rejection — or a v1 server's bogus payload — + // abort a large in-flight WRITE instead of deadlocking on mutually + // full socket buffers, and delivers rejection frames before the close. + bool write_body_done_ = false; + bool write_acked_ok_ = false; + // An early negative/malformed ack can arrive while asio::async_write still + // owns a buffer pointing into the caller's source memory. Do not publish a + // terminal status until that body operation has completed or been + // cancelled: callers are allowed to release the source buffer as soon as + // the transfer becomes terminal. + bool write_body_in_flight_ = false; + bool write_abort_requested_ = false; + // A v2 status frame is prompt by construction: a READ status precedes + // any payload, and a WRITE ack follows at most one chunk's apply after + // the body is done. The only peer that never sends one is a legacy + // server reached through a stale v2 descriptor — and for requests + // shorter than a frame it also keeps the connection open (it streamed + // size < 8 bytes of "READ payload" and is waiting for our next header), + // so without a deadline both sides wait forever. Bound that wait; the + // default is generous so no healthy slow path can trip it, since it only + // covers the frame itself, never payload streaming. + static int statusFrameTimeoutSec() { + const char* env = std::getenv("MC_TCP_STATUS_TIMEOUT_SEC"); + if (env) { + int v = std::atoi(env); + if (v > 0) return v; + } + return 30; + } + std::optional status_timer_; + bool status_deadline_disarmed_ = false; std::function on_finalize_; - std::function on_complete_; // Callback when transfer completes - std::mutex session_mutex_; + // Invoked exactly once per request with clean=true iff the protocol + // exchange terminated in a well-defined connection state. A socket whose + // request did not end cleanly must not be reused: the server-side session + // may be mid-frame, and the next request's header would be consumed as + // body bytes. + std::function on_complete_; void initiate(void* buffer, uint64_t dest_addr, size_t size, TransferRequest::OpCode opcode) { - session_mutex_.lock(); local_buffer_ = (char*)buffer; header_.addr = htole64(dest_addr); header_.size = htole64(size); - header_.opcode = (uint8_t)opcode; + header_.opcode = (uint8_t)opcode | (v2_ ? kOpcodeV2Flag : 0); total_transferred_bytes_ = 0; writeHeader(); } private: + // All handlers run on the transport's single io thread, so arm/cancel + // and the expiry handler never race. Expiry only closes the socket: the + // pending status read then completes with an error and its handler owns + // the failure path (including source-buffer quiescence for WRITE). + void armStatusDeadline() { + auto self(shared_from_this()); + status_deadline_disarmed_ = false; + status_timer_.emplace(socket_->get_executor()); + status_timer_->expires_after( + std::chrono::seconds(statusFrameTimeoutSec())); + status_timer_->async_wait([this, self](const asio::error_code& ec) { + // The disarmed flag also covers an expiry that was already + // queued when cancel() ran (cancel cannot revoke those, and by + // then the socket may have been re-pooled). + if (ec == asio::error::operation_aborted || + status_deadline_disarmed_) { + return; + } + LOG(ERROR) << "ClientSession: no status frame within " + << statusFrameTimeoutSec() + << "s (peer likely speaks the legacy protocol); " + "dropping connection"; + if (socket_ && socket_->is_open()) { + asio::error_code cec; + socket_->close(cec); + } + }); + } + + void cancelStatusDeadline() { + status_deadline_disarmed_ = true; + if (status_timer_) status_timer_->cancel(); + } + + // Single terminal path: finish connection ownership, then report the + // outcome. Posted so it runs after the invoking callback returns. + void finalize(TransferStatusEnum status, bool clean) { + cancelStatusDeadline(); + auto self(shared_from_this()); + asio::post( + socket_->get_executor(), + [this, self, status, clean, on_finalize = std::move(on_finalize_), + on_complete = std::move(on_complete_)]() { + // Finish connection ownership before publishing terminal + // status. Once on_finalize marks the slice, the caller may + // immediately free the batch or destroy the transport. + if (on_complete) on_complete(clean); + if (on_finalize) on_finalize(status); + }); + } + + // Abort a v2 WRITE and cancel any body operation. If asio still owns the + // current source buffer, its completion handler is responsible for + // finalizing after the buffer is quiescent. + void abortWrite() { + write_abort_requested_ = true; + if (socket_ && socket_->is_open()) { + asio::error_code ec; + socket_->close(ec); + } + if (!write_body_in_flight_) finalize(TransferStatusEnum::FAILED, false); + } + void writeHeader() { auto self(shared_from_this()); asio::async_write( @@ -340,21 +520,105 @@ struct ClientSession : public std::enable_shared_from_this { << "ClientSession::writeHeader failed. Error: " << ec.message() << " (value: " << ec.value() << ")" << ", bytes written: " << len; - asio::post( - socket_->get_executor(), - [this, self, on_finalize = std::move(on_finalize_), - on_complete = std::move(on_complete_)]() { - if (on_finalize) - on_finalize(TransferStatusEnum::FAILED); - session_mutex_.unlock(); - if (on_complete) on_complete(); - }); + finalize(TransferStatusEnum::FAILED, false); return; } - if (header_.opcode == (uint8_t)TransferRequest::WRITE) + if ((header_.opcode & ~kOpcodeV2Flag) == + (uint8_t)TransferRequest::WRITE) { + if (v2_) readWriteAck(); // concurrent with the body writeBody(); - else + } else if (v2_) { + readReadStatus(); + } else { readBody(); + } + }); + } + + // v2 READ: the server prefixes the data with a status frame. + void readReadStatus() { + auto self(shared_from_this()); + armStatusDeadline(); + asio::async_read( + *socket_, asio::buffer(&status_frame_, sizeof(status_frame_)), + [this, self](const asio::error_code& ec, std::size_t len) { + cancelStatusDeadline(); + if (ec || len != sizeof(status_frame_)) { + LOG(ERROR) + << "ClientSession: failed to read READ status " + "frame. Error: " + << ec.message() << " (value: " << ec.value() << ")"; + finalize(TransferStatusEnum::FAILED, false); + return; + } + uint64_t frame = le64toh(status_frame_); + if (!statusFrameValid(frame)) { + LOG(ERROR) << "ClientSession: malformed READ status " + "frame (peer likely speaks the legacy " + "protocol); dropping connection"; + finalize(TransferStatusEnum::FAILED, false); + return; + } + if (frame != kStatusOk) { + LOG(ERROR) << "ClientSession: READ rejected by server, " + "status " + << (frame & 0xFFFFFFFFull); + finalize(TransferStatusEnum::FAILED, false); + return; + } + readBody(); + }); + } + + // v2 WRITE: completion is the server's acknowledgment that the payload + // has been applied to destination memory. Armed concurrently with the + // body stream; a well-behaved v2 server only sends the frame after the + // final chunk, so a frame arriving before the body is done is either a + // rejection or a legacy peer's payload — both close the socket immediately + // and publish failure only after the outstanding body write is quiescent. + void readWriteAck() { + auto self(shared_from_this()); + asio::async_read( + *socket_, asio::buffer(&status_frame_, sizeof(status_frame_)), + [this, self](const asio::error_code& ec, std::size_t len) { + cancelStatusDeadline(); + if (ec || len != sizeof(status_frame_)) { + // The body path may have already finalized a failure and + // closed the socket; finalize() is idempotent (moved-from + // callbacks are null-checked). + if (ec != asio::error::operation_aborted) { + LOG(ERROR) + << "ClientSession: failed to read WRITE " + "ack frame. Error: " + << ec.message() << " (value: " << ec.value() << ")"; + } + abortWrite(); + return; + } + uint64_t frame = le64toh(status_frame_); + if (!statusFrameValid(frame)) { + LOG(ERROR) << "ClientSession: malformed WRITE ack frame " + "(peer likely speaks the legacy protocol); " + "dropping connection"; + abortWrite(); + return; + } + if (frame != kStatusOk) { + LOG(ERROR) << "ClientSession: WRITE rejected by server, " + "status " + << (frame & 0xFFFFFFFFull); + abortWrite(); + return; + } + if (!write_body_done_) { + // The server's ack can legitimately overtake the final + // local write-completion handler (both become ready + // together for small writes; the io thread may run this + // handler first). Record it; the body path finalizes. + write_acked_ok_ = true; + return; + } + finalize(TransferStatusEnum::COMPLETED, true); }); } @@ -366,14 +630,7 @@ struct ClientSession : public std::enable_shared_from_this { size_t buffer_size = std::min(getChunkSize(), size - total_transferred_bytes_); if (buffer_size == 0) { - asio::post(socket_->get_executor(), - [this, self, on_finalize = std::move(on_finalize_), - on_complete = std::move(on_complete_)]() { - if (on_finalize) - on_finalize(TransferStatusEnum::COMPLETED); - session_mutex_.unlock(); - if (on_complete) on_complete(); - }); + finalize(TransferStatusEnum::COMPLETED, true); return; } @@ -400,22 +657,12 @@ struct ClientSession : public std::enable_shared_from_this { << " using buffer " << static_cast(dram_buffer) << ". Error: " << ec.message() << " (value: " << ec.value() << ")"; - // Post entire cleanup to ensure it runs after callback - // returns - asio::post(socket_->get_executor(), - [this, self, dram_buffer, cuda_device, - on_finalize = std::move(on_finalize_), - on_complete = std::move(on_complete_)]() { - if (on_finalize) - on_finalize(TransferStatusEnum::FAILED); #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ defined(USE_COREX) - if (cuda_device >= 0) delete[] dram_buffer; + if (cuda_device >= 0) delete[] dram_buffer; #endif - session_mutex_.unlock(); - if (on_complete) on_complete(); - }); + finalize(TransferStatusEnum::FAILED, false); return; } @@ -438,19 +685,8 @@ struct ClientSession : public std::enable_shared_from_this { << "ClientSession::readBody failed to copy to CUDA " "memory. " << "Error: " << cudaGetErrorString(cuda_status); - // Post entire cleanup to ensure it runs after callback - // returns - asio::post( - socket_->get_executor(), - [this, self, dram_buffer, - on_finalize = std::move(on_finalize_), - on_complete = std::move(on_complete_)]() { - if (on_finalize) - on_finalize(TransferStatusEnum::FAILED); - delete[] dram_buffer; - session_mutex_.unlock(); - if (on_complete) on_complete(); - }); + delete[] dram_buffer; + finalize(TransferStatusEnum::FAILED, false); return; } delete[] dram_buffer; @@ -469,15 +705,28 @@ struct ClientSession : public std::enable_shared_from_this { size_t buffer_size = std::min(getChunkSize(), size - total_transferred_bytes_); if (buffer_size == 0) { - // Post cleanup to ensure it runs after callback returns - asio::post(socket_->get_executor(), - [this, self, on_finalize = std::move(on_finalize_), - on_complete = std::move(on_complete_)]() { - if (on_finalize) - on_finalize(TransferStatusEnum::COMPLETED); - session_mutex_.unlock(); - if (on_complete) on_complete(); - }); + if (v2_) { + if (write_abort_requested_) { + finalize(TransferStatusEnum::FAILED, false); + return; + } + // Completion comes from the server's acknowledgment, whose + // read is already in flight (armed in writeHeader) and may + // have finished first. + write_body_done_ = true; + if (write_acked_ok_) { + finalize(TransferStatusEnum::COMPLETED, true); + } else { + // From here a well-behaved server owes at most one + // chunk's apply plus the frame; a legacy peer behind a + // stale descriptor may owe nothing, ever. + armStatusDeadline(); + } + } else { + // v1: no acknowledgment exists in the protocol; this only + // means the payload left the initiator (#2086). + finalize(TransferStatusEnum::COMPLETED, true); + } return; } @@ -503,26 +752,19 @@ struct ClientSession : public std::enable_shared_from_this { LOG(ERROR) << "ClientSession::writeBody failed to copy from " "CUDA memory. " << "Error: " << cudaGetErrorString(cuda_status); - // Post entire cleanup to ensure it runs after callback returns - asio::post(socket_->get_executor(), - [this, self, dram_buffer, - on_finalize = std::move(on_finalize_), - on_complete = std::move(on_complete_)]() { - if (on_finalize) - on_finalize(TransferStatusEnum::FAILED); - delete[] dram_buffer; - session_mutex_.unlock(); - if (on_complete) on_complete(); - }); + delete[] dram_buffer; + abortWrite(); return; } } #endif + write_body_in_flight_ = true; asio::async_write( *socket_, asio::buffer(dram_buffer, buffer_size), [this, addr, dram_buffer, cuda_device, self]( const asio::error_code& ec, std::size_t transferred_bytes) { + write_body_in_flight_ = false; if (cuda_device >= 0) { delete[] dram_buffer; } @@ -533,17 +775,14 @@ struct ClientSession : public std::enable_shared_from_this { << " using buffer " << static_cast(dram_buffer) << ". Error: " << ec.message() << " (value: " << ec.value() << ")"; - // Post entire cleanup to ensure it runs after callback - // returns - asio::post( - socket_->get_executor(), - [this, self, on_finalize = std::move(on_finalize_), - on_complete = std::move(on_complete_)]() { - if (on_finalize) - on_finalize(TransferStatusEnum::FAILED); - session_mutex_.unlock(); - if (on_complete) on_complete(); - }); + abortWrite(); + return; + } + if (write_abort_requested_) { + // The early ack path closed the socket while this + // operation still owned the caller's source buffer. It is + // safe to publish failure now that the handler has run. + finalize(TransferStatusEnum::FAILED, false); return; } total_transferred_bytes_ += transferred_bytes; @@ -695,6 +934,9 @@ int TcpTransport::allocateLocalSegmentID(int tcp_data_port) { desc->protocol = "tcp"; #endif desc->tcp_data_port = tcp_data_port; + // Advertise acknowledged framing (#2086); initiators fall back to v1 + // against descriptors that do not carry the field. + desc->tcp_proto_version = 2; metadata_->addLocalSegment(LOCAL_SEGMENT_ID, local_server_name_, std::move(desc)); return 0; @@ -722,15 +964,23 @@ int TcpTransport::unregisterLocalMemory(void* addr, bool update_metadata) { int TcpTransport::registerLocalMemoryBatch( const std::vector& buffer_list, const std::string& location) { - for (auto& buffer : buffer_list) - registerLocalMemory(buffer.addr, buffer.length, location, true, false); + for (auto& buffer : buffer_list) { + int ret = registerLocalMemory(buffer.addr, buffer.length, location, + true, false); + if (ret) return ret; + } return metadata_->updateLocalSegmentDesc(); } int TcpTransport::unregisterLocalMemoryBatch( const std::vector& addr_list) { - for (auto& addr : addr_list) unregisterLocalMemory(addr, false); - return metadata_->updateLocalSegmentDesc(); + int first_error = 0; + for (auto& addr : addr_list) { + int ret = unregisterLocalMemory(addr, false); + if (ret && !first_error) first_error = ret; + } + int metadata_ret = metadata_->updateLocalSegmentDesc(); + return first_error ? first_error : metadata_ret; } Status TcpTransport::getTransferStatus(BatchID batch_id, size_t task_id, @@ -776,19 +1026,7 @@ Status TcpTransport::submitTransfer( for (auto& request : entries) { TransferTask& task = batch_desc.task_list[task_id]; ++task_id; - task.total_bytes = request.length; - Slice* slice = getSliceCache().allocate(); - slice->source_addr = (char*)request.source; - slice->length = request.length; - slice->opcode = request.opcode; - slice->tcp.dest_addr = request.target_offset; - slice->task = &task; - slice->target_id = request.target_id; - slice->status = Slice::PENDING; - slice->ts = 0; - task.slice_list.push_back(slice); - __sync_fetch_and_add(&task.slice_count, 1); - startTransfer(slice); + startTransfer(prepareTransfer(&task, request)); } return Status::OK(); @@ -796,28 +1034,63 @@ Status TcpTransport::submitTransfer( Status TcpTransport::submitTransferTask( const std::vector& task_list) { - for (size_t index = 0; index < task_list.size(); ++index) { - assert(task_list[index]); - auto& task = *task_list[index]; - assert(task.request); - auto& request = *task.request; - task.total_bytes = request.length; - Slice* slice = getSliceCache().allocate(); - slice->source_addr = (char*)request.source; - slice->length = request.length; - slice->opcode = request.opcode; - slice->tcp.dest_addr = request.target_offset; - slice->task = &task; - slice->target_id = request.target_id; - slice->status = Slice::PENDING; - slice->ts = 0; - task.slice_list.push_back(slice); - __sync_fetch_and_add(&task.slice_count, 1); - startTransfer(slice); + for (auto* task : task_list) { + assert(task && task->request); + startTransfer(prepareTransfer(task, *task->request)); } return Status::OK(); } +Status TcpTransport::submitTransferTaskGroup( + const std::vector& task_list) { + std::vector slices; + slices.reserve(task_list.size()); + for (auto* task : task_list) { + assert(task && task->request); + slices.push_back(prepareTransfer(task, *task->request)); + } + startTransferSequence(std::move(slices)); + return Status::OK(); +} + +Transport::Slice* TcpTransport::prepareTransfer( + TransferTask* task, const TransferRequest& request) { + task->total_bytes = request.length; + Slice* slice = getSliceCache().allocate(); + slice->source_addr = static_cast(request.source); + slice->length = request.length; + slice->opcode = request.opcode; + slice->tcp.dest_addr = request.target_offset; + slice->task = task; + slice->target_id = request.target_id; + slice->status = Slice::PENDING; + slice->ts = 0; + task->slice_list.push_back(slice); + __sync_fetch_and_add(&task->slice_count, 1); + return slice; +} + +void TcpTransport::startTransferSequence(std::vector slices) { + struct Sequence { + std::vector slices; + size_t next = 0; + }; + auto sequence = std::make_shared(); + sequence->slices = std::move(slices); + auto advance = std::make_shared>(); + std::weak_ptr> weak_advance = advance; + *advance = [this, sequence, weak_advance]() { + auto advance = weak_advance.lock(); + if (!advance || sequence->next == sequence->slices.size()) return; + auto* slice = sequence->slices[sequence->next++]; + std::function continuation; + if (sequence->next < sequence->slices.size()) + continuation = [advance]() { (*advance)(); }; + startTransfer(slice, std::move(continuation), true); + }; + (*advance)(); +} + void TcpTransport::worker() { while (running_) { try { @@ -833,9 +1106,9 @@ void TcpTransport::worker() { } std::shared_ptr TcpTransport::getConnection( - const std::string& host, uint16_t port) { - // If connection pool is disabled, always create a new connection - if (!enable_connection_pool_) { + const std::string& host, uint16_t port, bool use_pool) { + // Ungrouped transfers keep the configured connection-pool behavior. + if (!use_pool) { try { asio::ip::tcp::resolver resolver(context_->io_context); auto endpoint_iterator = @@ -1013,13 +1286,46 @@ bool TcpTransport::validateAddress(uint64_t addr, uint64_t size) const { return false; } -void TcpTransport::startTransfer(Slice* slice) { +void TcpTransport::discardConnection( + const std::string& host, uint16_t port, + std::shared_ptr socket) { + if (socket && socket->is_open()) { + asio::error_code ec; + socket->close(ec); + } + std::lock_guard lock(pool_mutex_); + auto it = connection_pool_.find(ConnectionKey{host, port}); + if (it != connection_pool_.end()) { + auto& queue = it->second; + for (auto queue_it = queue.begin(); queue_it != queue.end(); + ++queue_it) { + if ((*queue_it)->socket == socket) { + queue.erase(queue_it); + break; + } + } + if (queue.empty()) connection_pool_.erase(it); + } +} + +void TcpTransport::startTransfer(Slice* slice, + std::function continuation, + bool reuse_connection) { + auto finish = [this, slice, continuation = std::move(continuation)]( + TransferStatusEnum status) mutable { + if (status == TransferStatusEnum::COMPLETED) + slice->markSuccess(); + else + slice->markFailed(); + if (continuation) + asio::post(context_->io_context, std::move(continuation)); + }; auto desc = metadata_->getSegmentDescByID(slice->target_id); if (!desc) { LOG(ERROR) << "TcpTransport::startTransfer failed to get segment " "description for target_id: " << slice->target_id; - slice->markFailed(); + finish(TransferStatusEnum::FAILED); return; } @@ -1028,39 +1334,51 @@ void TcpTransport::startTransfer(Slice* slice) { LOG(ERROR) << "TcpTransport::startTransfer failed to get RPC meta " "entry for segment name: " << desc->name; - slice->markFailed(); + finish(TransferStatusEnum::FAILED); + return; + } + + // Zero-length requests are complete by definition. v1 reported them + // COMPLETED while the server silently rejected size==0 in address + // validation; short-circuiting keeps that outcome (rather than turning + // no-ops into v2 rejection failures) without the pointless round trip. + if (slice->length == 0) { + finish(TransferStatusEnum::COMPLETED); return; } - // Get connection from pool - auto socket = - getConnection(meta_entry.ip_or_host_name, desc->tcp_data_port); + const bool use_pool = enable_connection_pool_ || reuse_connection; + auto socket = getConnection(meta_entry.ip_or_host_name, desc->tcp_data_port, + use_pool); if (!socket) { LOG(ERROR) << "TcpTransport::startTransfer failed to get connection to " << meta_entry.ip_or_host_name << ":" << desc->tcp_data_port; - slice->markFailed(); + finish(TransferStatusEnum::FAILED); return; } try { - auto session = std::make_shared(socket); - - session->on_finalize_ = [slice](TransferStatusEnum status) { - if (status == TransferStatusEnum::COMPLETED) - slice->markSuccess(); - else - slice->markFailed(); - }; - - // Return connection to pool when transfer completes, or close if - // disabled - if (enable_connection_pool_) { + const bool use_v2 = + desc->tcp_proto_version >= 2 && !forceLegacyTcpProto(); + auto session = std::make_shared(socket, use_v2); + + session->on_finalize_ = finish; + + // Return connection to pool when the request terminated cleanly; + // otherwise the server-side session state is unknown (it may be + // mid-frame), so reusing the socket would desynchronize the next + // request. Discard it instead. + if (use_pool) { session->on_complete_ = [this, host = meta_entry.ip_or_host_name, - port = desc->tcp_data_port, socket]() { - returnConnection(host, port, socket); + port = desc->tcp_data_port, + socket](bool clean) { + if (clean) + returnConnection(host, port, socket); + else + discardConnection(host, port, socket); }; } else { - session->on_complete_ = [socket]() { + session->on_complete_ = [socket](bool) { // Close connection immediately after transfer if (socket && socket->is_open()) { asio::error_code ec; @@ -1078,33 +1396,12 @@ void TcpTransport::startTransfer(Slice* slice) { << ", opcode: " << (int)slice->opcode << ", target_id: " << slice->target_id << ". Exception: " << e.what(); - // On exception, always close the socket and remove from pool if present - // Don't return it to the pool as it may be in an inconsistent state - if (socket && socket->is_open()) { - asio::error_code ec; - socket->close(ec); - } - if (enable_connection_pool_) { - // Remove the connection from pool if it was pooled - ConnectionKey key{meta_entry.ip_or_host_name, - static_cast(desc->tcp_data_port)}; - std::lock_guard lock(pool_mutex_); - auto it = connection_pool_.find(key); - if (it != connection_pool_.end()) { - auto& queue = it->second; - for (auto queue_it = queue.begin(); queue_it != queue.end(); - ++queue_it) { - if ((*queue_it)->socket == socket) { - queue.erase(queue_it); - break; - } - } - if (queue.empty()) { - connection_pool_.erase(it); - } - } - } - slice->markFailed(); + // On exception, always close the socket and remove from pool if + // present. Don't return it to the pool as it may be in an + // inconsistent state. + discardConnection(meta_entry.ip_or_host_name, + static_cast(desc->tcp_data_port), socket); + finish(TransferStatusEnum::FAILED); } } diff --git a/mooncake-transfer-engine/tent/CMakeLists.txt b/mooncake-transfer-engine/tent/CMakeLists.txt index b5bcc3795e..f7970376a4 100644 --- a/mooncake-transfer-engine/tent/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/CMakeLists.txt @@ -1,13 +1,22 @@ cmake_minimum_required(VERSION 3.16) -if (CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) - project(mc_tent LANGUAGES CXX) +if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) + project(mc_tent LANGUAGES CXX) +endif() + +if(BUILD_UNIT_TESTS AND NOT TARGET gtest) + include(${CMAKE_CURRENT_SOURCE_DIR}/../../mooncake-common/FindGTest.cmake) endif() set(CMAKE_POSITION_INDEPENDENT_CODE ON) +# Compile-time metrics switch (default OFF for zero overhead on the transfer hot +# path). Declared at the top level so it is discoverable; the actual compile +# definition is applied in tent/src/metrics/CMakeLists.txt. +option(TENT_METRICS_ENABLED "Enable TENT metrics collection" OFF) + add_subdirectory(src) add_subdirectory(plugins) -if (BUILD_UNIT_TESTS) - add_subdirectory(tests) +if(BUILD_UNIT_TESTS) + add_subdirectory(tests) endif() diff --git a/mooncake-transfer-engine/tent/config/transfer-engine.json b/mooncake-transfer-engine/tent/config/transfer-engine.json index bfe46b2ed7..632346afb5 100644 --- a/mooncake-transfer-engine/tent/config/transfer-engine.json +++ b/mooncake-transfer-engine/tent/config/transfer-engine.json @@ -17,18 +17,7 @@ "http_port": 9100, "http_host": "0.0.0.0", "http_server_threads": 2, - "report_interval_seconds": 30, - "enable_prometheus": true, - "enable_json": true, - "latency_buckets": [ - 0.000125, 0.00015, 0.0002, 0.00025, 0.0003, 0.0004, 0.0005, - 0.00075, 0.001, 0.0015, 0.002, 0.003, 0.005, 0.007, - 0.015, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0 - ], - "size_buckets": [ - 1024, 4096, 16384, 65536, 262144, 1048576, 4194304, - 16777216, 67108864, 268435456, 1073741824 - ] + "report_interval_seconds": 30 }, "transports": { "rdma": { @@ -71,6 +60,21 @@ "rail_topo_path": "/path/to/rail_topo.json" } }, + "ub": { + "enable": false, + "device_filter": [], + "worker_count": 4, + "poller_count": 1, + "jfc_per_context": 1, + "jetty_per_endpoint": 4, + "max_endpoints": 4096, + "slice_size": 65536, + "max_retries": 3, + "slice_timeout_ms": 10000, + "endpoint_cooldown_ms": 30000, + "enable_bandwidth_estimation": true, + "enable_notifications": false + }, "tcp": { "enable": true, "max_retry_count": 3, @@ -83,8 +87,7 @@ }, "shm": { "enable" : true, - "cxl_mount_path": "", - "async_memcpy_threshold": 4 + "cxl_mount_path": "" }, "mnnvl": { "enable" : false diff --git a/mooncake-transfer-engine/tent/include/tent/common/concurrent/rw_spinlock.h b/mooncake-transfer-engine/tent/include/tent/common/concurrent/rw_spinlock.h index f946a7cfbb..4eb8d8af4b 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/concurrent/rw_spinlock.h +++ b/mooncake-transfer-engine/tent/include/tent/common/concurrent/rw_spinlock.h @@ -26,6 +26,7 @@ namespace tent { class RWSpinlock { union RWTicket { constexpr RWTicket() : whole(0) {} + constexpr RWTicket(uint64_t v) : whole(v) {} uint64_t whole; uint32_t readWrite; struct { @@ -33,26 +34,12 @@ class RWSpinlock { uint16_t read; uint16_t users; }; - } ticket; - - private: - static void asm_volatile_memory() { asm volatile("" ::: "memory"); } - - template - static T load_acquire(T *addr) { - T t = *addr; - asm_volatile_memory(); - return t; - } + }; - template - static void store_release(T *addr, T v) { - asm_volatile_memory(); - *addr = v; - } + std::atomic ticket; public: - RWSpinlock() {} + RWSpinlock() : ticket(0) {} RWSpinlock(RWSpinlock const &) = delete; RWSpinlock &operator=(RWSpinlock const &) = delete; @@ -60,17 +47,21 @@ class RWSpinlock { void lock() { writeLockNice(); } bool tryLock() { - RWTicket t; - uint64_t old = t.whole = load_acquire(&ticket.whole); + RWTicket t, expected; + expected.whole = ticket.load(std::memory_order_acquire); + t.whole = expected.whole; if (t.users != t.write) return false; ++t.users; - return __sync_bool_compare_and_swap(&ticket.whole, old, t.whole); + return ticket.compare_exchange_weak(expected.whole, t.whole, + std::memory_order_acquire); } void writeLockAggressive() { uint32_t count = 0; - uint16_t val = __sync_fetch_and_add(&ticket.users, 1); - while (val != load_acquire(&ticket.write)) { + uint16_t val = fetch_add_users(1); + RWTicket t; + while (val != + (t.whole = ticket.load(std::memory_order_acquire), t.write)) { PAUSE(); if (++count > 1000) std::this_thread::yield(); } @@ -85,16 +76,22 @@ class RWSpinlock { } void unlockAndLockShared() { - uint16_t val = __sync_fetch_and_add(&ticket.read, 1); + uint16_t val = fetch_add_read(1); (void)val; } void unlock() { + uint64_t expected = ticket.load(std::memory_order_relaxed); + uint64_t new_val; RWTicket t; - t.whole = load_acquire(&ticket.whole); - ++t.read; - ++t.write; - store_release(&ticket.readWrite, t.readWrite); + do { + t.whole = expected; + ++t.read; + ++t.write; + new_val = t.whole; + } while (!ticket.compare_exchange_weak(expected, new_val, + std::memory_order_release, + std::memory_order_relaxed)); } void lockShared() { @@ -106,15 +103,58 @@ class RWSpinlock { } bool tryLockShared() { - RWTicket t, old; - old.whole = t.whole = load_acquire(&ticket.whole); - old.users = old.read; + RWTicket t, expected; + expected.whole = ticket.load(std::memory_order_acquire); + t.whole = expected.whole; + expected.users = expected.read; ++t.read; ++t.users; - return __sync_bool_compare_and_swap(&ticket.whole, old.whole, t.whole); + return ticket.compare_exchange_weak(expected.whole, t.whole, + std::memory_order_acquire); } - void unlockShared() { __sync_fetch_and_add(&ticket.write, 1); } + void unlockShared() { fetch_add_write(1); } + + private: + uint16_t fetch_add_users(uint16_t delta) { + uint64_t expected = ticket.load(std::memory_order_relaxed); + uint64_t new_val; + RWTicket t; + do { + t.whole = expected; + t.users += delta; + new_val = t.whole; + } while (!ticket.compare_exchange_weak(expected, new_val, + std::memory_order_acquire, + std::memory_order_relaxed)); + return static_cast(t.users - delta); + } + + uint16_t fetch_add_read(uint16_t delta) { + uint64_t expected = ticket.load(std::memory_order_relaxed); + uint64_t new_val; + RWTicket t; + do { + t.whole = expected; + t.read += delta; + new_val = t.whole; + } while (!ticket.compare_exchange_weak(expected, new_val, + std::memory_order_release)); + return static_cast(t.read - delta); + } + + uint16_t fetch_add_write(uint16_t delta) { + uint64_t expected = ticket.load(std::memory_order_relaxed); + uint64_t new_val; + RWTicket t; + do { + t.whole = expected; + t.write += delta; + new_val = t.whole; + } while (!ticket.compare_exchange_weak(expected, new_val, + std::memory_order_release)); + return static_cast(t.write - delta); + } public: struct WriteGuard { @@ -150,4 +190,4 @@ class RWSpinlock { } // namespace tent } // namespace mooncake -#endif // TENT_RW_SPINLOCK_H \ No newline at end of file +#endif // TENT_RW_SPINLOCK_H diff --git a/mooncake-transfer-engine/tent/include/tent/common/concurrent/thread_local_storage.h b/mooncake-transfer-engine/tent/include/tent/common/concurrent/thread_local_storage.h index ab2fc36937..63b28accd1 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/concurrent/thread_local_storage.h +++ b/mooncake-transfer-engine/tent/include/tent/common/concurrent/thread_local_storage.h @@ -15,70 +15,183 @@ #ifndef TENT_THREAD_LOCAL_STORAGE_H #define TENT_THREAD_LOCAL_STORAGE_H +#include +#include #include +#include #include -#include +#include #include namespace mooncake { namespace tent { +namespace detail { +// Hoisted out of the class template so ids are unique across ALL +// ThreadLocalStorage instantiations, not just within one T. The per-thread +// maps are per-T, so per-T uniqueness would suffice today — global +// uniqueness removes the aliasing hazard if the per-thread state is ever +// shared across instantiations. +inline uint64_t nextThreadLocalStorageId() { + static std::atomic counter{1}; + return counter.fetch_add(1, std::memory_order_relaxed); +} +} // namespace detail + +// Per-instance thread-local storage (#2717). +// +// Each (ThreadLocalStorage instance, thread) pair owns a distinct T. The +// previous implementation kept one `thread_local` slot per template +// instantiation, so all instances of ThreadLocalStorage in a process +// aliased each other's per-thread state, and its heap-allocated holders were +// never destroyed (the deregistration path was unreachable). +// +// Lifetime rules: +// - A thread's T values are destroyed when the thread exits. +// - Destroying the storage does not destroy other threads' values; they are +// orphaned (at most one T per thread per destroyed storage) and reclaimed +// at those threads' exit. Instance ids are process-monotonic and never +// reused, so an orphaned value can never be served to a new instance that +// happens to reuse the same address. +// - Both teardown orders are safe: a thread exiting while the owner is alive +// deregisters its value from the owner's registry; an owner destroyed +// while threads are alive marks the jointly-owned control block dead, and +// those threads skip deregistration at exit. Orphans are additionally +// swept opportunistically: the next first-use get() of ANY instance on +// that thread reclaims all of the thread's dead-owner values, so orphan +// count stays bounded by live instances even under storage churn. +// +// Precondition: callers must keep the storage alive across every get() / +// forEach() call (the usual member-of-owner pattern satisfies this); the +// destructor may run concurrently only with other threads' exits, not with +// their accesses. +// +// Concurrency: +// - get() is lock-free after the first call per (instance, thread): one +// thread_local access plus an id compare on the hot path. +// - forEach() runs under the registry mutex and visits exactly the values of +// live registered threads. It synchronizes registry membership only — if +// owning threads mutate their T concurrently, the callback observes those +// fields with whatever synchronization T itself provides (unchanged from +// the previous implementation). template class ThreadLocalStorage { public: ThreadLocalStorage() = default; - ~ThreadLocalStorage() = default; - // Disable copy/move + ~ThreadLocalStorage() { + std::lock_guard lock(control_->mutex); + control_->owner_alive.store(false, std::memory_order_release); + control_->values.clear(); + } + ThreadLocalStorage(const ThreadLocalStorage&) = delete; ThreadLocalStorage& operator=(const ThreadLocalStorage&) = delete; - // Access the thread-local instance (lock-free) + // Access the calling thread's instance, constructing it on first use. T& get() { - if (!instance_) { - instance_ = new InstanceHolder(this); + ThreadState& state = threadState(); + if (state.cached_id == id_) return *state.cached_value; + auto it = state.nodes.find(id_); + if (it == state.nodes.end()) { + // First use of this instance on this thread — the cold path. + // Piggyback a sweep of values whose owners are gone, so a + // long-lived thread does not accumulate one orphan per + // destroyed storage it ever touched (their T contents would + // otherwise stay pinned until thread exit). + sweepDeadNodes(state); + it = state.nodes.try_emplace(id_).first; + ThreadNode& node = it->second; + node.control = control_; + // The caller holds a reference to *this, so the owner is alive + // and registration cannot race ~ThreadLocalStorage's clear(). + std::lock_guard lock(control_->mutex); + control_->values.insert(&node.value); } - return instance_->value; + state.cached_id = id_; + state.cached_value = &it->second.value; + return *state.cached_value; } - // Safe iteration over all instances (with locking) + // Safe iteration over the values of all live threads that have called + // get() on this instance. void forEach(const std::function& fn) { - std::lock_guard lock(global_mutex_); - for (auto* inst : instances_) { - fn(inst->value); + std::lock_guard lock(control_->mutex); + for (T* value : control_->values) { + fn(*value); } } private: - struct InstanceHolder { - T value; - ThreadLocalStorage* owner; + // Shared between the owner and every thread node so that whichever side + // is torn down last still has a valid registry (or a dead flag) to look + // at. + struct ControlBlock { + std::mutex mutex; + std::unordered_set values; + // Atomic so the orphan sweep can test liveness without taking the + // mutex of every node it scans. + std::atomic owner_alive{true}; + }; - InstanceHolder(ThreadLocalStorage* owner) : owner(owner) { - std::lock_guard lock(owner->global_mutex_); - owner->instances_.insert(this); - } + struct ThreadNode { + T value{}; + std::shared_ptr control; - ~InstanceHolder() { - std::lock_guard lock(owner->global_mutex_); - owner->instances_.erase(this); + ThreadNode() = default; + ThreadNode(const ThreadNode&) = delete; + ThreadNode& operator=(const ThreadNode&) = delete; + + ~ThreadNode() { + if (!control) return; + std::lock_guard lock(control->mutex); + if (control->owner_alive.load(std::memory_order_acquire)) + control->values.erase(&value); } }; - // Thread-local pointer to the per-thread instance - thread_local static InstanceHolder* instance_; + struct ThreadState { + // Node-based map: ThreadNode addresses are stable across rehash, + // which the registry and the one-entry cache below rely on. + std::unordered_map nodes; + // One-entry cache so the common get() is a single thread_local + // access plus a compare. Entries live until thread exit, so the + // cached pointer cannot dangle while the id matches. + uint64_t cached_id = 0; // ids start at 1; 0 never matches + T* cached_value = nullptr; + }; - // Global list of all thread instances - std::unordered_set instances_; - std::mutex global_mutex_; -}; + static void sweepDeadNodes(ThreadState& state) { + for (auto it = state.nodes.begin(); it != state.nodes.end();) { + // A null control cannot be observed today (the sweep runs before + // try_emplace on the same thread, and control is assigned + // immediately after emplace by a nothrow shared_ptr copy), but + // check it for consistency with ~ThreadNode and robustness to + // reordering. + if (!it->second.control || !it->second.control->owner_alive.load( + std::memory_order_acquire)) { + if (state.cached_value == &it->second.value) { + state.cached_id = 0; + state.cached_value = nullptr; + } + // ~ThreadNode sees the dead owner and skips the registry. + it = state.nodes.erase(it); + } else { + ++it; + } + } + } -// Definition of thread_local variable (must be outside the class) -template -thread_local typename ThreadLocalStorage::InstanceHolder* - ThreadLocalStorage::instance_ = nullptr; + static ThreadState& threadState() { + thread_local ThreadState state; + return state; + } + + const uint64_t id_ = detail::nextThreadLocalStorageId(); + std::shared_ptr control_ = std::make_shared(); +}; } // namespace tent } // namespace mooncake -#endif // TENT_THREAD_LOCAL_STORAGE_H \ No newline at end of file +#endif // TENT_THREAD_LOCAL_STORAGE_H diff --git a/mooncake-transfer-engine/tent/include/tent/common/qos_metrics.h b/mooncake-transfer-engine/tent/include/tent/common/qos_metrics.h new file mode 100644 index 0000000000..e182a9ae28 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/common/qos_metrics.h @@ -0,0 +1,103 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_COMMON_QOS_METRICS_H +#define TENT_COMMON_QOS_METRICS_H + +#include +#include +#include +#include +#include + +namespace mooncake { +namespace tent { + +// One entry in --qos_classes: +// name:threads:slo_us:weight[:isolated_gbps] +// isolated_gbps is an optional result from a matching isolated run. It is +// deliberately an input: isolation loss cannot be inferred from a mixed run. +struct QosClassConfig { + std::string name; + int threads = 0; + uint64_t slo_us = 0; + double weight = 1.0; + std::optional isolated_throughput_gbps; +}; + +struct QosClassMetrics { + std::string name; + int threads = 0; + uint64_t slo_us = 0; + double weight = 1.0; + size_t operations = 0; + double throughput_gbps = 0.0; + double p99_us = 0.0; + std::optional slo_attainment; + double goodput_gbps = 0.0; + double weighted_goodput_gbps = 0.0; + std::optional isolated_throughput_gbps; + std::optional isolation_leakage; + uint64_t transferred_bytes = 0; +}; + +struct QosClassSample { + size_t operations = 0; + double total_duration_us = 0.0; + double p99_us = 0.0; + std::optional slo_attainment; + std::optional transferred_bytes; +}; + +struct QosMetricsReport { + size_t block_size = 0; + size_t batch_size = 0; + int num_threads = 0; + double aggregate_throughput_gbps = 0.0; + double weighted_goodput_gbps = 0.0; + double jain_fairness = 0.0; + std::optional max_isolation_leakage; + std::optional link_capacity_gbps; + std::optional total_utilization; + std::vector classes; +}; + +bool parseQosClasses(const std::string& spec, + std::vector* classes, std::string* error); + +bool parseQosClassesJson(const std::string& spec, + std::vector* classes, + std::string* error); + +bool validateQosClasses(const std::vector& classes, + int num_threads, std::string* error); + +size_t qosClassForThread(const std::vector& classes, + int thread_id); + +QosMetricsReport calculateQosMetrics(size_t block_size, size_t batch_size, + int num_threads, + const std::vector& classes, + const std::vector& samples, + double link_capacity_gbps); + +void printQosMetrics(const QosMetricsReport& report); + +bool appendQosMetricsJsonl(const std::string& path, + const QosMetricsReport& report, std::string* error); + +} // namespace tent +} // namespace mooncake + +#endif // TENT_COMMON_QOS_METRICS_H diff --git a/mooncake-transfer-engine/tent/include/tent/common/types.h b/mooncake-transfer-engine/tent/include/tent/common/types.h index 974ed4d6dd..088f0b3128 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/types.h +++ b/mooncake-transfer-engine/tent/include/tent/common/types.h @@ -54,6 +54,8 @@ enum TransportType : int { TCP, AscendDirect, SUNRISE_LINK, + TPU, + UB, // Sentinel: must remain the last enumerator. kNumTransportTypes, }; @@ -65,6 +67,64 @@ inline TransportType c_to_transport_hint(int v) { return static_cast(v); } +inline const char* transportTypeName(TransportType type) { + switch (type) { + case UNSPEC: + return "unspec"; + case RDMA: + return "rdma"; + case MNNVL: + return "mnnvl"; + case SHM: + return "shm"; + case NVLINK: + return "nvlink"; + case GDS: + return "gds"; + case IOURING: + return "io_uring"; + case TCP: + return "tcp"; + case AscendDirect: + return "ascend"; + case SUNRISE_LINK: + return "sunrise_link"; + case TPU: + return "tpu"; + case UB: + return "ub"; + case kNumTransportTypes: + return "unknown"; + } + return "unknown"; +} + +inline TransportType parseTransportType(const std::string& str) { + if (str == "unspec") return UNSPEC; + if (str == "rdma") return RDMA; + if (str == "mnnvl") return MNNVL; + if (str == "shm") return SHM; + if (str == "nvlink") return NVLINK; + if (str == "gds") return GDS; + if (str == "io_uring") return IOURING; + if (str == "tcp") return TCP; + if (str == "ascend") return AscendDirect; + if (str == "sunrise_link") return SUNRISE_LINK; + if (str == "tpu") return TPU; + if (str == "ub") return UB; + return UNSPEC; +} + +enum class IntentType : int { + INTENT_UNSPEC = 0, + FOREGROUND_GET, + BACKGROUND_PREFETCH, + MIGRATION, + CHECKPOINT, + WEIGHT_LOADING, + STAGING_INTERNAL, +}; + struct Request { enum OpCode { READ, WRITE }; OpCode opcode; @@ -79,6 +139,13 @@ struct Request { TransportType transport_hint = UNSPEC; // UNSPEC = follow policy; otherwise pin this request to the // name transport. + // Optional SLO deadline as an absolute steady_clock timestamp in + // nanoseconds. 0 = no deadline (default), behaves exactly as today. + // When set, the engine emits an observability-only feasibility metric + // (MLU = actual transfer time / available window) on completion; it does + // not yet drive any admission or scheduling decision. See RFC #2519. + uint64_t deadline_ns = 0; + IntentType intent_type = IntentType::INTENT_UNSPEC; }; enum TransferStatusEnum { @@ -96,6 +163,12 @@ struct TransferStatus { size_t transferred_bytes; }; +struct NicLoadStats { + std::string device_name; + uint64_t inflight_bytes{0}; + double ewma_bandwidth_bps{0.0}; +}; + enum Permission { kLocalReadWrite, kGlobalReadOnly, diff --git a/mooncake-transfer-engine/tent/include/tent/common/utils/os.h b/mooncake-transfer-engine/tent/include/tent/common/utils/os.h index 1c31915b3a..c174c0ad71 100644 --- a/mooncake-transfer-engine/tent/include/tent/common/utils/os.h +++ b/mooncake-transfer-engine/tent/include/tent/common/utils/os.h @@ -27,6 +27,7 @@ #include #include #include +#include #include #include @@ -41,6 +42,16 @@ namespace mooncake { namespace tent { +// libnuma fills the cache numa_node_to_cpus() reads lazily and without locking, +// so concurrent first callers each allocate it and all but one are orphaned -- +// a leak LeakSanitizer reports. Workers bind every thread at startup, so they +// hit that window. An inline function, not a static local: bindToSocket() has +// internal linkage, so a static local would be per-TU. +inline std::mutex &numaNodeCpuCacheMutex() { + static std::mutex mutex; + return mutex; +} + static inline int bindToSocket(int socket_id) { if (numa_available() < 0) { LOG(WARNING) << "The platform does not support NUMA"; @@ -51,7 +62,10 @@ static inline int bindToSocket(int socket_id) { if (socket_id < 0 || socket_id >= numa_num_configured_nodes()) socket_id = 0; struct bitmask *cpu_list = numa_allocate_cpumask(); - numa_node_to_cpus(socket_id, cpu_list); + { + std::lock_guard guard(numaNodeCpuCacheMutex()); + numa_node_to_cpus(socket_id, cpu_list); + } int nr_possible_cpus = numa_num_possible_cpus(); int nr_cpus = 0; for (int cpu = 0; cpu < nr_possible_cpus; ++cpu) { @@ -96,4 +110,4 @@ static inline std::string getCurrentDateTime() { } // namespace tent } // namespace mooncake -#endif // TENT_OS_H \ No newline at end of file +#endif // TENT_OS_H diff --git a/mooncake-transfer-engine/tent/include/tent/metastore/http.h b/mooncake-transfer-engine/tent/include/tent/metastore/http.h index a9b5532f29..8d764dd4c4 100644 --- a/mooncake-transfer-engine/tent/include/tent/metastore/http.h +++ b/mooncake-transfer-engine/tent/include/tent/metastore/http.h @@ -45,17 +45,16 @@ class HttpMetaStore : public MetaStore { return size * nmemb; } - std::string encodeUrl(const std::string &key) { - char *newkey = curl_easy_escape(client_, key.c_str(), key.size()); - std::string encodedKey(newkey); + std::string encodeUrl(CURL *curl, const std::string &key) { + char *newkey = curl_easy_escape(curl, key.c_str(), key.size()); + std::string encodedKey(newkey ? newkey : ""); std::string url = endpoint_ + "?key=" + encodedKey; - curl_free(newkey); + if (newkey) curl_free(newkey); return url; } private: std::atomic connected_; - CURL *client_; std::string endpoint_; }; } // namespace tent diff --git a/mooncake-transfer-engine/tent/include/tent/metrics/config_loader.h b/mooncake-transfer-engine/tent/include/tent/metrics/config_loader.h index a61a7c10fd..88073c51a0 100644 --- a/mooncake-transfer-engine/tent/include/tent/metrics/config_loader.h +++ b/mooncake-transfer-engine/tent/include/tent/metrics/config_loader.h @@ -29,10 +29,6 @@ struct MetricsConfig { uint16_t http_port = 9100; uint16_t http_server_threads = 2; // HTTP server thread count uint32_t report_interval_seconds = 30; // 0 means disabled - bool enable_prometheus = true; - bool enable_json = true; - std::vector latency_buckets; - std::vector size_buckets; }; // Helper class to load metrics configuration from various sources @@ -71,12 +67,6 @@ constexpr const char* METRICS_HTTP_SERVER_THREADS = "metrics/http_server_threads"; constexpr const char* METRICS_REPORT_INTERVAL = "metrics/report_interval_seconds"; -constexpr const char* METRICS_ENABLE_PROMETHEUS = "metrics/enable_prometheus"; -constexpr const char* METRICS_ENABLE_JSON = "metrics/enable_json"; - -// Bucket configurations -constexpr const char* METRICS_LATENCY_BUCKETS = "metrics/latency_buckets"; -constexpr const char* METRICS_SIZE_BUCKETS = "metrics/size_buckets"; // Environment variable names (with TENT_ prefix) constexpr const char* ENV_METRICS_ENABLED = "TENT_METRICS_ENABLED"; @@ -86,15 +76,9 @@ constexpr const char* ENV_METRICS_HTTP_SERVER_THREADS = "TENT_METRICS_HTTP_SERVER_THREADS"; constexpr const char* ENV_METRICS_REPORT_INTERVAL = "TENT_METRICS_REPORT_INTERVAL"; -constexpr const char* ENV_METRICS_ENABLE_PROMETHEUS = - "TENT_METRICS_ENABLE_PROMETHEUS"; -constexpr const char* ENV_METRICS_ENABLE_JSON = "TENT_METRICS_ENABLE_JSON"; -constexpr const char* ENV_METRICS_LATENCY_BUCKETS = - "TENT_METRICS_LATENCY_BUCKETS"; -constexpr const char* ENV_METRICS_SIZE_BUCKETS = "TENT_METRICS_SIZE_BUCKETS"; } // namespace config_keys } // namespace tent } // namespace mooncake -#endif // TENT_METRICS_CONFIG_LOADER_H \ No newline at end of file +#endif // TENT_METRICS_CONFIG_LOADER_H diff --git a/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h b/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h index 90162e11d6..b23e9d4a8f 100644 --- a/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h +++ b/mooncake-transfer-engine/tent/include/tent/metrics/tent_metrics.h @@ -14,6 +14,7 @@ #pragma once +#include #include #include #include @@ -23,6 +24,7 @@ #include #include "tent/common/status.h" +#include "tent/common/types.h" #include "tent/metrics/config_loader.h" // Compile-time metrics enable/disable switch @@ -74,12 +76,50 @@ class TentMetrics { return runtime_enabled_.load(std::memory_order_relaxed); } - // Record transfer operations - void recordReadCompleted(size_t bytes, double latency_seconds = 0.0); - void recordWriteCompleted(size_t bytes, double latency_seconds = 0.0); - void recordReadFailed(size_t bytes); - void recordWriteFailed(size_t bytes); - void recordTransportFailover(); + // Record logical transfer outcomes. The TransportType argument is the + // terminal transport: after failover, a recovered request is attributed to + // the transport that completed it. + void recordReadCompleted(TransportType tp, size_t bytes, + double latency_seconds = 0.0); + void recordWriteCompleted(TransportType tp, size_t bytes, + double latency_seconds = 0.0); + void recordReadFailed(TransportType tp); + void recordWriteFailed(TransportType tp); + // Failover counter is labeled with both the source and destination + // transport types so failover flows (e.g. rdma->tcp) are queryable. + void recordTransportFailover(TransportType from, TransportType to); + + // Record physical transport attempts separately from logical requests. + // "Started" means the engine is about to call submitTransferTasks (or the + // staging equivalent). "Finished" records the observed attempt duration + // and, for a FAILED status (including synchronous submit failure), + // increments the attempt-failure counter. + void recordTransportAttemptStarted(TransportType tp, + Request::OpCode operation); + void recordTransportAttemptFinished(TransportType tp, + Request::OpCode operation, + TransferStatusEnum status, + double latency_us); + + // Record the deadline feasibility ratio (MLU) for a completed transfer + // that carried a deadline. mlu = actual_transfer_seconds / window_seconds, + // where window_seconds is (deadline - submit_time). mlu < 1 means the + // transfer met its deadline; mlu >= 1 means it missed. Observability only. + void recordDeadlineMLU(TransportType tp, double mlu); + + // Record a transfer whose deadline was already in the past at submit time + // (infeasible window). Recorded into a dedicated counter so it is + // distinguishable from genuine MLU samples in the histogram above. + void recordDeadlineInfeasible(TransportType tp); + + enum class Stage { + QueueWait, + Dispatch, + Transport, + }; + + // Causal chain: record per-stage latency breakdown (microseconds). + void recordStageLatency(Stage stage, TransportType tp, double latency_us); // Get metrics for HTTP server std::string getPrometheusMetrics(); @@ -89,6 +129,14 @@ class TentMetrics { // Check if initialized bool isInitialized() const { return initialized_; } + // Port the HTTP metrics server is bound to, or 0 when the endpoint is not + // running (log-only mode, or metrics disabled at compile time). Backed by + // an atomic so it is safe to read from other threads while initialize() is + // still running. + uint16_t httpPort() const { + return bound_http_port_.load(std::memory_order_relaxed); + } + private: TentMetrics() = default; ~TentMetrics(); @@ -100,10 +148,22 @@ class TentMetrics { std::atomic initialized_{false}; MetricsConfig config_; + // Port the HTTP server actually bound to, 0 until a successful bind. Kept + // separate from config_.http_port and atomic because httpPort() may be read + // by other threads while the initializing thread is still binding a port. + std::atomic bound_http_port_{0}; #if TENT_METRICS_ENABLED - // Initialize HTTP server with endpoints - void initHttpServer(); + // Initialize and start the HTTP server on the configured port. Port + // assignment is deterministic: co-located ranks are expected to be given + // distinct ports explicitly (e.g. base_port + local_rank) rather than + // auto-scanned. Returns an error if the port cannot be bound (the caller + // then degrades to log-only metrics); on success bound_http_port_ is set. + Status initHttpServer(); + + // Register the /metrics, /metrics/summary, /metrics/json and /health + // endpoints on the current http_server_ instance. + void registerHttpHandlers(); // HTTP server for metrics endpoint std::unique_ptr http_server_; @@ -114,55 +174,193 @@ class TentMetrics { std::mutex metric_report_mutex_; std::condition_variable metric_report_cv_; - // Counters - stored as pointers for unified management - std::vector counters_; - ylt::metric::counter_t read_bytes_total_{"tent_read_bytes_total", - "Total bytes read via TENT"}; - ylt::metric::counter_t write_bytes_total_{"tent_write_bytes_total", - "Total bytes written via TENT"}; - ylt::metric::counter_t read_requests_total_{"tent_read_requests_total", - "Total read requests via TENT"}; - ylt::metric::counter_t write_requests_total_{ - "tent_write_requests_total", "Total write requests via TENT"}; - ylt::metric::counter_t read_failures_total_{"tent_read_failures_total", - "Total read failures via TENT"}; - ylt::metric::counter_t write_failures_total_{ - "tent_write_failures_total", "Total write failures via TENT"}; - ylt::metric::counter_t failover_total_{ + // Counters — stored as base pointers (metric_t*) so that counters with + // different label arities (N=1 for per-transport, N=2 for failover and + // transport-attempt operation labels) share one vector for Prometheus + // serialize(). The concrete typed members below are used directly for + // JSON/summary aggregation (which need to iterate label values via copy()). + std::vector counters_; + + // Label name arrays for dynamic metric construction. + static inline const std::array kTransportLabel{"transport"}; + static inline const std::array kFailoverLabels{"from", + "to"}; + static inline const std::array kAttemptLabels{"transport", + "operation"}; + + // Per-transport counters (label: transport). Values are int64_t. + ylt::metric::basic_dynamic_counter read_bytes_total_{ + "tent_read_bytes_total", "Total bytes read via TENT", kTransportLabel}; + ylt::metric::basic_dynamic_counter write_bytes_total_{ + "tent_write_bytes_total", "Total bytes written via TENT", + kTransportLabel}; + ylt::metric::basic_dynamic_counter read_requests_total_{ + "tent_read_requests_total", "Total read requests via TENT", + kTransportLabel}; + ylt::metric::basic_dynamic_counter write_requests_total_{ + "tent_write_requests_total", "Total write requests via TENT", + kTransportLabel}; + ylt::metric::basic_dynamic_counter read_failures_total_{ + "tent_read_failures_total", "Total read failures via TENT", + kTransportLabel}; + ylt::metric::basic_dynamic_counter write_failures_total_{ + "tent_write_failures_total", "Total write failures via TENT", + kTransportLabel}; + // Failover counter has two labels (from, to) so failover flows are + // queryable as e.g. tent_transport_failover_total{from="rdma",to="tcp"}. + ylt::metric::basic_dynamic_counter failover_total_{ "tent_transport_failover_total", - "Total cross-transport failover events"}; - - // Histograms - stored as pointers for unified management - std::vector histograms_; - // Store bucket boundaries separately since ylt histogram doesn't expose - // them publicly - std::vector> histogram_boundaries_; + "Total cross-transport failover events", kFailoverLabels}; + ylt::metric::basic_dynamic_counter transport_attempts_total_{ + "tent_transport_attempts_total", + "Total physical transport attempts submitted for execution", + kAttemptLabels}; + ylt::metric::basic_dynamic_counter + transport_attempt_failures_total_{ + "tent_transport_attempt_failures_total", + "Physical transport attempts that terminated with FAILED", + kAttemptLabels}; + ylt::metric::basic_dynamic_counter deadline_infeasible_total_{ + "tent_deadline_infeasible_total", + "Transfers whose deadline was already in the past at submit", + kTransportLabel}; + + // Histograms - paired with their bucket boundaries in a single vector so + // the two cannot drift out of sync (ylt histogram doesn't expose its + // boundaries publicly, so we hold them alongside the pointer). + // Each entry also carries a parallel counter used to track the per-label + // sum, because ylt's basic_dynamic_histogram keeps sum_ private with no + // public accessor — without this counter, _sum could only be emitted as 0 + // in the Prometheus endpoint, breaking _sum / _count queries. + // Only N=1 (per-transport) histograms live here; the N=2 + // transport_attempt_latency_ histogram is serialized separately (see + // getPrometheusMetrics / getJsonMetrics) via the same templated helpers. + struct HistogramEntry { + ylt::metric::basic_dynamic_histogram* h; + const std::vector* boundaries; + ylt::metric::basic_dynamic_counter* sum; + }; + std::vector histograms_; // Latency histograms use microseconds (us) as unit // Default buckets: 100us, 500us, 1ms, 5ms, 10ms, 50ms, 100ms, 500ms, 1s static inline const std::vector kLatencyBuckets{ 100, 500, 1000, 5000, 10000, 50000, 100000, 500000, 1000000}; - ylt::metric::histogram_t read_latency_{ + ylt::metric::basic_dynamic_histogram read_latency_{ "tent_read_latency_us", "Read latency distribution in microseconds", - kLatencyBuckets}; - ylt::metric::histogram_t write_latency_{ + kLatencyBuckets, kTransportLabel}; + ylt::metric::basic_dynamic_histogram write_latency_{ "tent_write_latency_us", "Write latency distribution in microseconds", - kLatencyBuckets}; + kLatencyBuckets, kTransportLabel}; // Size histograms for request size distribution (in bytes) // Default buckets: 1KB, 4KB, 16KB, 64KB, 256KB, 1MB, 4MB, 16MB, 64MB, // 256MB, 1GB static inline const std::vector kSizeBuckets{ 1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, 67108864, 268435456, 1073741824}; - ylt::metric::histogram_t read_size_{ + ylt::metric::basic_dynamic_histogram read_size_{ "tent_read_size_bytes", "Read request size distribution in bytes", - kSizeBuckets}; - ylt::metric::histogram_t write_size_{ + kSizeBuckets, kTransportLabel}; + ylt::metric::basic_dynamic_histogram write_size_{ "tent_write_size_bytes", "Write request size distribution in bytes", - kSizeBuckets}; + kSizeBuckets, kTransportLabel}; + + // Deadline feasibility ratio (MLU) distribution for transfers that carried + // a deadline. Stored in per-mille (MLU x 1000) so the histogram can use + // integer observe() like the others; the 1000 boundary is MLU == 1.0, the + // feasible/infeasible line (< 1000 met the deadline, >= 1000 missed it). + static inline const std::vector kMluPerMilleBuckets{ + 100, 250, 500, 750, 900, 1000, 1250, 1500, 2000, 5000}; + ylt::metric::basic_dynamic_histogram deadline_mlu_{ + "tent_deadline_mlu_permille", + "Deadline feasibility ratio (MLU x 1000) distribution", + kMluPerMilleBuckets, kTransportLabel}; + + // Causal chain stage latency histograms (microseconds) + // Buckets span 10us to 500ms to capture both fast RDMA and slower TCP. + static inline const std::vector kStageBuckets{ + 10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000, 500000}; + ylt::metric::basic_dynamic_histogram stage_queue_wait_{ + "tent_stage_queue_wait_us", + "Causal chain: queue wait latency in microseconds", kStageBuckets, + kTransportLabel}; + ylt::metric::basic_dynamic_histogram stage_dispatch_{ + "tent_stage_dispatch_us", + "Causal chain: dispatch latency in microseconds", kStageBuckets, + kTransportLabel}; + ylt::metric::basic_dynamic_histogram stage_transport_{ + "tent_stage_transport_us", + "Causal chain: transport execution latency in microseconds", + kStageBuckets, kTransportLabel}; + ylt::metric::basic_dynamic_histogram transport_attempt_latency_{ + "tent_transport_attempt_latency_us", + "Observed physical transport attempt latency in microseconds", + kLatencyBuckets, kAttemptLabels}; + + // Parallel counters tracking per-label sum for each histogram. ylt's + // basic_dynamic_histogram keeps sum_ private with no public accessor, so + // we maintain these alongside observe() calls and read them back via + // copy() in the Prometheus serializer. Each counter's str_name is the + // histogram name suffixed with "_sum" so it emits as _sum{...} in + // Prometheus output. + ylt::metric::basic_dynamic_counter read_latency_sum_{ + "tent_read_latency_us_sum", + "Sum of read latency observations (microseconds)", kTransportLabel}; + ylt::metric::basic_dynamic_counter write_latency_sum_{ + "tent_write_latency_us_sum", + "Sum of write latency observations (microseconds)", kTransportLabel}; + ylt::metric::basic_dynamic_counter read_size_sum_{ + "tent_read_size_bytes_sum", "Sum of read request sizes (bytes)", + kTransportLabel}; + ylt::metric::basic_dynamic_counter write_size_sum_{ + "tent_write_size_bytes_sum", "Sum of write request sizes (bytes)", + kTransportLabel}; + ylt::metric::basic_dynamic_counter deadline_mlu_sum_{ + "tent_deadline_mlu_permille_sum", + "Sum of deadline MLU (permille) observations", kTransportLabel}; + ylt::metric::basic_dynamic_counter stage_queue_wait_sum_{ + "tent_stage_queue_wait_us_sum", + "Sum of queue wait latency observations (microseconds)", + kTransportLabel}; + ylt::metric::basic_dynamic_counter stage_dispatch_sum_{ + "tent_stage_dispatch_us_sum", + "Sum of dispatch latency observations (microseconds)", kTransportLabel}; + ylt::metric::basic_dynamic_counter stage_transport_sum_{ + "tent_stage_transport_us_sum", + "Sum of transport execution latency observations (microseconds)", + kTransportLabel}; + // Parallel sum counter for the N=2 transport-attempt latency histogram + // (labels: transport, operation). Same rationale as the N=1 sum counters + // above — ylt keeps the histogram's sum_ private. + ylt::metric::basic_dynamic_counter + transport_attempt_latency_sum_{"tent_transport_attempt_latency_us_sum", + "Sum of physical transport attempt " + "latency observations (microseconds)", + kAttemptLabels}; // Helper to register all metrics to the vectors void registerMetrics(); + + // Serialize a single histogram in Prometheus text format. Walks the + // same get_bucket_counts() / copy() data the JSON path uses, so the two + // endpoints can never drift. Unlike ylt's serialize(), this never + // silently drops a histogram that has observed >=1 sample — even when + // every observation landed in the first bucket (which makes sum_==0 + // and causes ylt's serialize() to clear() its output string, taking the + // # HELP / # TYPE header with it). Reachable in production when + // sub-microsecond latencies truncate to 0 under int64_t observation. + // + // `boundaries` is the compile-time bucket boundary vector paired with + // the histogram in HistogramEntry. The caller (getPrometheusMetrics) + // emits the # HELP / # TYPE header so the format stays identical to ylt. + // Templated over label arity N so both the per-transport (N=1) histograms + // and the transport-attempt (N=2) histogram share one implementation. + template + void serializeHistogramPrometheus( + ylt::metric::basic_dynamic_histogram* hist, + const std::vector& boundaries, + ylt::metric::basic_dynamic_counter& sum, + std::string& out) const; #endif // TENT_METRICS_ENABLED }; @@ -178,93 +376,104 @@ class ScopedLatencyRecorder { public: enum class OperationType { Read, Write }; - ScopedLatencyRecorder(OperationType type, size_t bytes) - : type_(type), bytes_(bytes), enabled_(TentMetrics::isEnabled()) { - // Only record start time if metrics are enabled (avoid clock overhead - // when disabled) + ScopedLatencyRecorder(OperationType type, TransportType tp, size_t bytes) + : type_(type), + tp_(tp), + bytes_(bytes), + enabled_(TentMetrics::isEnabled()) { if (enabled_) { start_ = std::chrono::steady_clock::now(); } } ~ScopedLatencyRecorder() { - if (!enabled_ || failed_) - return; // Skip if disabled or already marked as failed + if (!enabled_ || failed_) return; auto end = std::chrono::steady_clock::now(); double latency = std::chrono::duration(end - start_).count(); if (type_ == OperationType::Read) { - TentMetrics::instance().recordReadCompleted(bytes_, latency); + TentMetrics::instance().recordReadCompleted(tp_, bytes_, latency); } else { - TentMetrics::instance().recordWriteCompleted(bytes_, latency); + TentMetrics::instance().recordWriteCompleted(tp_, bytes_, latency); } } void markFailed() { - if (!enabled_) return; // Skip if disabled + if (!enabled_) return; failed_ = true; if (type_ == OperationType::Read) { - TentMetrics::instance().recordReadFailed(bytes_); + TentMetrics::instance().recordReadFailed(tp_); } else { - TentMetrics::instance().recordWriteFailed(bytes_); + TentMetrics::instance().recordWriteFailed(tp_); } } private: OperationType type_; + TransportType tp_; size_t bytes_; std::chrono::steady_clock::time_point start_; - bool enabled_; // Captured at construction time for consistent behavior + bool enabled_; bool failed_ = false; }; -// Convenience macros for recording metrics (enabled version) -#define TENT_RECORD_READ_COMPLETED(bytes, latency) \ +// Convenience macros for recording metrics (enabled version). +// The TransportType argument labels each metric so Prometheus queries can +// break down traffic by transport. For failover, from→to labels the flow. +#define TENT_RECORD_READ_COMPLETED(tp, bytes, latency) \ do { \ if (::mooncake::tent::TentMetrics::isEnabled()) { \ ::mooncake::tent::TentMetrics::instance().recordReadCompleted( \ - bytes, latency); \ + tp, bytes, latency); \ } \ } while (0) -#define TENT_RECORD_WRITE_COMPLETED(bytes, latency) \ +#define TENT_RECORD_WRITE_COMPLETED(tp, bytes, latency) \ do { \ if (::mooncake::tent::TentMetrics::isEnabled()) { \ ::mooncake::tent::TentMetrics::instance().recordWriteCompleted( \ - bytes, latency); \ + tp, bytes, latency); \ } \ } while (0) -#define TENT_RECORD_READ_FAILED(bytes) \ - do { \ - if (::mooncake::tent::TentMetrics::isEnabled()) { \ - ::mooncake::tent::TentMetrics::instance().recordReadFailed(bytes); \ - } \ +#define TENT_RECORD_READ_FAILED(tp) \ + do { \ + if (::mooncake::tent::TentMetrics::isEnabled()) { \ + ::mooncake::tent::TentMetrics::instance().recordReadFailed(tp); \ + } \ } while (0) -#define TENT_RECORD_WRITE_FAILED(bytes) \ - do { \ - if (::mooncake::tent::TentMetrics::isEnabled()) { \ - ::mooncake::tent::TentMetrics::instance().recordWriteFailed( \ - bytes); \ - } \ +#define TENT_RECORD_WRITE_FAILED(tp) \ + do { \ + if (::mooncake::tent::TentMetrics::isEnabled()) { \ + ::mooncake::tent::TentMetrics::instance().recordWriteFailed(tp); \ + } \ } while (0) -#define TENT_RECORD_TRANSPORT_FAILOVER() \ - do { \ - if (::mooncake::tent::TentMetrics::isEnabled()) { \ - ::mooncake::tent::TentMetrics::instance() \ - .recordTransportFailover(); \ - } \ +#define TENT_RECORD_TRANSPORT_FAILOVER(from, to) \ + do { \ + if (::mooncake::tent::TentMetrics::isEnabled()) { \ + ::mooncake::tent::TentMetrics::instance().recordTransportFailover( \ + from, to); \ + } \ } while (0) -// RAII macro for automatic latency measurement -#define TENT_SCOPED_READ_LATENCY(bytes) \ - ::mooncake::tent::ScopedLatencyRecorder _tent_latency_recorder_( \ - ::mooncake::tent::ScopedLatencyRecorder::OperationType::Read, bytes) - -#define TENT_SCOPED_WRITE_LATENCY(bytes) \ - ::mooncake::tent::ScopedLatencyRecorder _tent_latency_recorder_( \ - ::mooncake::tent::ScopedLatencyRecorder::OperationType::Write, bytes) +#define TENT_SCOPED_READ_LATENCY(tp, bytes) \ + ::mooncake::tent::ScopedLatencyRecorder _tent_latency_recorder_( \ + ::mooncake::tent::ScopedLatencyRecorder::OperationType::Read, tp, \ + bytes) + +#define TENT_SCOPED_WRITE_LATENCY(tp, bytes) \ + ::mooncake::tent::ScopedLatencyRecorder _tent_latency_recorder_( \ + ::mooncake::tent::ScopedLatencyRecorder::OperationType::Write, tp, \ + bytes) + +#define TENT_RECORD_STAGE_LATENCY(stage, tp, latency_us) \ + do { \ + if (::mooncake::tent::TentMetrics::isEnabled()) { \ + ::mooncake::tent::TentMetrics::instance().recordStageLatency( \ + stage, tp, latency_us); \ + } \ + } while (0) #else // !TENT_METRICS_ENABLED @@ -272,18 +481,19 @@ class ScopedLatencyRecorder { class ScopedLatencyRecorder { public: enum class OperationType { Read, Write }; - ScopedLatencyRecorder(OperationType, size_t) {} + ScopedLatencyRecorder(OperationType, TransportType, size_t) {} void markFailed() {} }; // Zero-overhead macros when metrics are disabled at compile time -#define TENT_RECORD_READ_COMPLETED(bytes, latency) ((void)0) -#define TENT_RECORD_WRITE_COMPLETED(bytes, latency) ((void)0) -#define TENT_RECORD_READ_FAILED(bytes) ((void)0) -#define TENT_RECORD_WRITE_FAILED(bytes) ((void)0) -#define TENT_RECORD_TRANSPORT_FAILOVER() ((void)0) -#define TENT_SCOPED_READ_LATENCY(bytes) ((void)0) -#define TENT_SCOPED_WRITE_LATENCY(bytes) ((void)0) +#define TENT_RECORD_READ_COMPLETED(tp, bytes, latency) ((void)0) +#define TENT_RECORD_WRITE_COMPLETED(tp, bytes, latency) ((void)0) +#define TENT_RECORD_READ_FAILED(tp) ((void)0) +#define TENT_RECORD_WRITE_FAILED(tp) ((void)0) +#define TENT_RECORD_TRANSPORT_FAILOVER(from, to) ((void)0) +#define TENT_SCOPED_READ_LATENCY(tp, bytes) ((void)0) +#define TENT_SCOPED_WRITE_LATENCY(tp, bytes) ((void)0) +#define TENT_RECORD_STAGE_LATENCY(stage, tp, latency_us) ((void)0) #endif // TENT_METRICS_ENABLED diff --git a/mooncake-transfer-engine/tent/include/tent/platform/tpu.h b/mooncake-transfer-engine/tent/include/tent/platform/tpu.h new file mode 100644 index 0000000000..1bca92a70d --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/platform/tpu.h @@ -0,0 +1,66 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_PLATFORM_TPU_H_ +#define TENT_PLATFORM_TPU_H_ + +#include "tent/common/config.h" +#include "tent/platform/cpu.h" +#include "tent/runtime/platform.h" + +namespace mooncake { +namespace tent { + +// TpuPlatform models a TPU host: the host DRAM and RDMA topology are identical +// to CpuPlatform, so those paths (host allocation, NUMA probing, host<->host +// copy) are inherited unchanged. Only the TPU-device-aware operations are +// overridden, and every one of them is delegated to TpuPjrtShim so that this +// class carries no direct PJRT dependency. +// +// TPU HBM cannot be reached by the NIC, so there is no direct device transport: +// the HBM<->host hop is performed here via copy(), and the host<->host hop is +// carried by RDMA/TCP. ProxyManager chains the two (see findStagingPolicy). +class TpuPlatform : public CpuPlatform { + public: + explicit TpuPlatform(std::shared_ptr config) + : CpuPlatform(config) {} + + ~TpuPlatform() override {} + + // Host + RDMA discovery from CpuPlatform, plus one MemEntry per TPU device + // ("tpu:N") so findNearMem() can resolve a device to its nearest host node. + Status probe(std::vector &nic_list, + std::vector &mem_list) override; + + // Device allocation is owned by the serving framework (JAX / torch-XLA); + // host allocation is inherited. "tpu" locations therefore return + // NotImplemented here. + Status allocate(void **pptr, size_t size, MemoryOptions &options) override; + + // HBM<->host copy via the adapter when either side is TPU memory; otherwise + // the inherited host memcpy. + Status copy(void *dst, void *src, size_t length) override; + + MemoryType getMemoryType(void *addr) override; + + const std::vector getLocation( + void *start, size_t len, bool skip_prefault = false) override; + + const std::string type() const override { return "tpu"; } +}; + +} // namespace tent +} // namespace mooncake + +#endif // TENT_PLATFORM_TPU_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h new file mode 100644 index 0000000000..3cb2be6fc8 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_abi.h @@ -0,0 +1,86 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// C ABI contract between TENT and the TPU/PJRT adapter shared library. +// +// The adapter is built separately against the PJRT runtime and exports these +// symbols with C linkage. TENT resolves them at runtime via dlopen()/dlsym() +// (see TpuPjrtShim); it never links the adapter or the PJRT runtime directly. +// +// Pointer classification note: the "device pointer" passed across this ABI is +// the stable token the serving-engine integration registers with TENT for a TPU +// buffer (see registerLocalMemory with a "tpu:N" location). The adapter owns +// the mapping from that token to the underlying PJRT buffer; mc_tpu_pjrt_* copy +// and classification calls resolve the token through the adapter's own +// registry. +// +// INTERIOR POINTERS (load-bearing): TENT stages transfers through host DRAM in +// chunks (see ProxyManager, chunk_size defaults to 4 MiB), and hands the +// adapter `token + chunk_offset` for every chunk after the first. Every +// entrypoint below therefore takes an address that may point into the MIDDLE of +// a registered buffer, not only at its base. An adapter whose registry only +// matches base addresses will report `is_device_ptr(token + off) == 0`, TENT +// will classify TPU HBM as host memory, and the staging copy silently degrades +// into a memcpy from a non-data address -- corrupting every transfer larger +// than one chunk without raising an error. Adapters MUST resolve an address to +// the registered buffer whose range [base, base + size) contains it. +// +// The token is NOT required to be host-dereferenceable, and on real PJRT/TPU it +// is not: PJRT_Buffer_UnsafePointer returns an internal handle that reads as +// garbage rather than buffer contents. Never dereference it. + +#ifndef TENT_PLATFORM_TPU_PJRT_ABI_H_ +#define TENT_PLATFORM_TPU_PJRT_ABI_H_ + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Initialize the adapter (creates/attaches the PJRT client). Returns 0 on +// success, non-zero on failure. Idempotent; safe to call more than once. +int mc_tpu_pjrt_init(void); + +// Returns 1 if `addr` falls inside any TPU device buffer known to the adapter +// (base address or interior), else 0. +int mc_tpu_pjrt_is_device_ptr(const void *addr); + +// Returns the device ordinal of the buffer containing `addr` (base address or +// interior), or -1 if `addr` is not inside a known TPU device buffer. +int mc_tpu_pjrt_device_index(const void *addr); + +// Synchronous device->host copy. `device_src` may be an interior address; the +// whole range [device_src, device_src + len) must lie within a single +// registered buffer, otherwise the adapter must fail rather than copy short. +// Returns 0 on success, non-zero on failure. +int mc_tpu_pjrt_copy_d2h(void *host_dst, const void *device_src, size_t len); + +// Synchronous host->device copy. `device_dst` may be an interior address; the +// whole range [device_dst, device_dst + len) must lie within a single +// registered buffer, otherwise the adapter must fail rather than copy short. +// Returns 0 on success, non-zero on failure. +int mc_tpu_pjrt_copy_h2d(void *device_dst, const void *host_src, size_t len); + +// Number of visible TPU devices. +int mc_tpu_pjrt_device_count(void); + +// NUMA node closest to device `index`, or -1 if unknown. +int mc_tpu_pjrt_device_numa(int index); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // TENT_PLATFORM_TPU_PJRT_ABI_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h new file mode 100644 index 0000000000..e2dd2d0a3d --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/platform/tpu_pjrt_shim.h @@ -0,0 +1,97 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_PLATFORM_TPU_PJRT_SHIM_H_ +#define TENT_PLATFORM_TPU_PJRT_SHIM_H_ + +#include +#include + +#include "tent/common/status.h" + +namespace mooncake { +namespace tent { + +// TpuPjrtShim isolates every TPU/PJRT dependency behind a narrow interface. +// +// The TPU device I/O primitives (HBM<->host DMA, device-pointer classification, +// device topology) are provided by a separate adapter shared library built +// against the PJRT runtime. TENT itself carries no build-time dependency on +// that runtime: the adapter is resolved at runtime via dlopen(), matching the +// way the other accelerator backends keep vendor SDKs out of the core build. +// +// The adapter library must export the C ABI declared in tpu_pjrt_abi.h. Its +// path defaults to "libmooncake_tpu_pjrt.so" and can be overridden with the +// MC_TPU_PJRT_LIB environment variable. When the adapter cannot be loaded, +// available() returns false and every operation returns a non-OK Status; this +// keeps a USE_TPU build functional (and unit-testable with a mock adapter) +// without the real runtime present. +class TpuPjrtShim { + public: + // Process-wide singleton. The adapter is loaded (and initialized) lazily on + // first use; loading is attempted at most once. + static TpuPjrtShim &instance(); + + // True when the adapter library was loaded and initialized successfully. + bool available() const { return available_; } + + // Returns true if `addr` refers to memory owned by the TPU runtime (HBM), + // including addresses interior to a registered buffer. Returns false when + // the adapter is unavailable or the pointer is host memory, so a caller can + // safely treat "not TPU" as host memory. + bool isDevicePtr(const void *addr) const; + + // Device ordinal of the buffer containing `addr` (base or interior), or -1 + // if `addr` is not TPU device memory. + int deviceIndex(const void *addr) const; + + // Synchronous HBM -> host DMA copy of `length` bytes. `device_src` may be + // interior to a registered buffer; the range must not run past its end. + Status copyD2H(void *host_dst, const void *device_src, size_t length) const; + + // Synchronous host -> HBM DMA copy of `length` bytes. `device_dst` may be + // interior to a registered buffer; the range must not run past its end. + Status copyH2D(void *device_dst, const void *host_src, size_t length) const; + + // Number of visible TPU devices (0 when the adapter is unavailable). + int deviceCount() const; + + // NUMA node closest to TPU device `index`, or -1 if unknown. + int deviceNumaNode(int index) const; + + private: + TpuPjrtShim(); + ~TpuPjrtShim(); + TpuPjrtShim(const TpuPjrtShim &) = delete; + TpuPjrtShim &operator=(const TpuPjrtShim &) = delete; + + void load(); + + void *handle_ = nullptr; + bool available_ = false; + + // Resolved adapter entrypoints (see tpu_pjrt_abi.h for the contract). + int (*fn_init_)() = nullptr; + int (*fn_is_device_ptr_)(const void *) = nullptr; + int (*fn_device_index_)(const void *) = nullptr; + int (*fn_copy_d2h_)(void *, const void *, size_t) = nullptr; + int (*fn_copy_h2d_)(void *, const void *, size_t) = nullptr; + int (*fn_device_count_)() = nullptr; + int (*fn_device_numa_)(int) = nullptr; +}; + +} // namespace tent +} // namespace mooncake + +#endif // TENT_PLATFORM_TPU_PJRT_SHIM_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h index 1897454514..4253adb191 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/admission_queue.h @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -42,6 +43,29 @@ struct QueueLimits { size_t max_outstanding_bytes{0}; size_t staging_owner_reserve{0}; size_t staging_byte_reserve{0}; + // Opt-in deadline-aware dispatch (RFC #2519 step 2). When false (default), + // pickForDispatch keeps strict FIFO order — unchanged behavior. When true, + // owners carrying a deadline (request.deadline_ns != 0) are dispatched + // earliest-deadline-first; owners without a deadline keep FIFO order behind + // them. This only reorders selection within the existing capacity limits; + // it does not admit/reject or otherwise change what gets dispatched. + bool deadline_aware{false}; + // Opt-in deadline-infeasible drop (RFC #2519 step 3). Local-decode MLU + // threshold θ_local. 0 (default) disables drop entirely — behavior is the + // step-2 EDF ordering (or FIFO). When > 0 (e.g. 1.5) and a bandwidth + // provider is set, an owner whose predicted MLU + // (= predicted_transfer_time / remaining_window) reaches this threshold is + // dropped instead of dispatched, and on_local_decode_suggested is raised so + // the caller can recompute locally. Requires deadline_aware = true. + double mlu_local_threshold{0.0}; + // Opt-in deadline proximity promotion. When > 0, pickForDispatch promotes + // queued owners whose remaining slack (deadline_ns - now) is below this + // threshold to the front of the dispatch queue, ahead of owners with more + // slack or no deadline. This dynamically boosts urgency as a deadline + // approaches, regardless of original admission order. Requires a + // NowProvider (via setDegradationPolicy) or defaults to steady_clock. + // 0 (default) disables promotion entirely. + uint64_t promotion_slack_ns{0}; }; struct QueueOwnerInput { @@ -50,6 +74,11 @@ struct QueueOwnerInput { std::vector derived_task_ids; Request request{}; QueueOwnerKind kind{QueueOwnerKind::User}; + // True only when the caller has established that this owner's transfer + // time is governed by the installed bandwidth provider. Default false + // keeps degradation explicitly opt-in so a new enqueue path cannot + // accidentally apply an RDMA EWMA to MNNVL/TCP/staging paths. + bool degradation_eligible{false}; }; struct QueueSubmit { @@ -59,6 +88,23 @@ struct QueueSubmit { std::vector owners; }; +// RFC #2519 step 3: degradation signal raised when a transfer is predicted to +// miss its deadline and is dropped from dispatch. The bodies (compression / +// local recompute) live in the upper layer (vLLM/SGLang); TENT only raises the +// signal. No hook registered ⇒ the drop still happens but nothing is notified. +struct DegradationHooks { + std::function on_local_decode_suggested; +}; + +// Returns the predicted transfer bandwidth in bytes/second, or <= 0 if unknown +// (in which case the drop decision is skipped). Injected by the owner so the +// admission queue does not depend on the device-selection layer directly. +using BandwidthProvider = std::function; + +// Returns "now" as a steady-clock timestamp in nanoseconds, matching the units +// of Request.deadline_ns. Injectable so tests are deterministic. +using NowProvider = std::function; + // Runtime-private admission model. It is intentionally single-threaded; the // eventual TransferEngineImpl integration owns synchronization. class LocalTransferAdmissionQueue { @@ -75,11 +121,29 @@ class LocalTransferAdmissionQueue { Status tryAdmit(const QueueSubmit& submit, std::vector& admitted_owner_ids); - std::vector pickForDispatch(size_t max_owners, - size_t max_bytes); + // Returns the owners to dispatch. When step-3 drop is enabled + // (mlu_local_threshold > 0, deadline_aware, and a bandwidth provider set), + // owners predicted to miss their deadline are dropped: charged out of the + // outstanding accounting, marked terminal (CANCELED), appended to + // `dropped_owner_ids` (if non-null), and on_local_decode_suggested is + // raised. `dropped_owner_ids` is cleared on entry. + std::vector pickForDispatch( + size_t max_owners, size_t max_bytes, + std::vector* dropped_owner_ids = nullptr); + + // Install the step-3 degradation policy inputs. Optional; without it the + // queue never drops (default behavior). now defaults to steady_clock. + void setDegradationPolicy(BandwidthProvider bandwidth_provider, + DegradationHooks hooks, + NowProvider now_provider = nullptr); Status complete(QueueOwnerId owner_id, TransferStatusEnum terminal_status); + // Cancel an owner that has not entered the dispatch window. Idempotent for + // an owner already canceled; dispatching owners must be canceled through + // their selected transport instead. + Status cancel(QueueOwnerId owner_id); + Status retireBatch(uint64_t batch_token); Status resolveOwner(uint64_t batch_token, size_t public_task_id, @@ -96,15 +160,16 @@ class LocalTransferAdmissionQueue { enum class QueueState { Queued, Dispatching, - Completed, - Failed, + Terminal, }; struct QueueOwner { uint64_t batch_token{0}; Request request{}; QueueOwnerKind kind{QueueOwnerKind::User}; + bool degradation_eligible{false}; QueueState state{QueueState::Queued}; + TransferStatusEnum terminal_status{TransferStatusEnum::PENDING}; }; QueueLimits limits_; @@ -117,6 +182,11 @@ class LocalTransferAdmissionQueue { size_t outstanding_bytes_{0}; size_t outstanding_user_owners_{0}; size_t outstanding_user_bytes_{0}; + + // RFC #2519 step 3 degradation policy (all optional / opt-in). + BandwidthProvider bandwidth_provider_; + DegradationHooks degradation_hooks_; + NowProvider now_provider_; }; } // namespace tent diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/platform.h b/mooncake-transfer-engine/tent/include/tent/runtime/platform.h index 95a4c7aefa..e53265adb0 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/platform.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/platform.h @@ -20,7 +20,10 @@ namespace mooncake { namespace tent { -enum MemoryType { MTYPE_UNKNOWN, MTYPE_CPU, MTYPE_CUDA, MTYPE_ROCM }; +// MTYPE_TPU is appended last so the numeric values of the existing entries are +// preserved. TPU HBM is not NIC-addressable, so transfers touching it are +// staged through host DRAM by ProxyManager (see findStagingPolicy). +enum MemoryType { MTYPE_UNKNOWN, MTYPE_CPU, MTYPE_CUDA, MTYPE_ROCM, MTYPE_TPU }; class Platform { public: diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/progress_worker.h b/mooncake-transfer-engine/tent/include/tent/runtime/progress_worker.h index de8ed0666c..30fc01a098 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/progress_worker.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/progress_worker.h @@ -16,6 +16,7 @@ #define PROGRESS_WORKER_H_ #include +#include #include #include #include @@ -30,15 +31,16 @@ namespace tent { class TransferEngineImpl; // Event-driven progress worker for issue #2116. When the engine is configured -// with enable_progress_worker=true, transports (or test hooks) call +// with enable_progress_worker=true, transports call // notifyBatchMaybeReady to wake this worker, which then drives one -// progressBatch step per notification. This decouples failover/resubmit from -// the caller polling loop, so integrators that turn off -// enable_auto_failover_on_poll do not need to spin a polling thread of their -// own to keep failover progressing. +// progressBatch step per notification. When the runtime queue has active work, +// the worker also uses a low-frequency fallback tick so transports that do not +// yet emit completion wakes can still drain the queue. class ProgressWorker { public: - explicit ProgressWorker(TransferEngineImpl* impl); + explicit ProgressWorker(TransferEngineImpl* impl, + std::chrono::microseconds fallback_interval = + std::chrono::microseconds(0)); ~ProgressWorker(); ProgressWorker(const ProgressWorker&) = delete; @@ -55,10 +57,15 @@ class ProgressWorker { // started. void notifyBatchMaybeReady(BatchID batch_id); + // Safe from any thread. Coalesces multiple runtime queue wakes into one + // bounded refill step. + void notifyRuntimeQueueReady(); + private: void runner(); TransferEngineImpl* impl_; + std::chrono::microseconds fallback_interval_; std::atomic running_{false}; std::thread thread_; @@ -66,6 +73,7 @@ class ProgressWorker { std::condition_variable cv_; std::unordered_set queued_; std::deque order_; + bool queue_ready_{false}; }; } // namespace tent diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/proxy_manager.h b/mooncake-transfer-engine/tent/include/tent/runtime/proxy_manager.h index 0620810d9f..959197cac3 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/proxy_manager.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/proxy_manager.h @@ -31,6 +31,7 @@ struct StageBufferCache; struct StagingTask { TaskInfo* native{nullptr}; + BatchID batch{0}; std::vector params; }; @@ -51,7 +52,8 @@ class ProxyManager { Status deconstruct(); - Status submit(TaskInfo* task, const std::vector& params); + Status submit(TaskInfo* task, BatchID batch, + const std::vector& params); Status getStatus(TaskInfo* task, TransferStatus& task_status); @@ -113,4 +115,4 @@ class ProxyManager { } // namespace tent } // namespace mooncake -#endif // PROXY_MANAGER_H_ \ No newline at end of file +#endif // PROXY_MANAGER_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/qos_contract.h b/mooncake-transfer-engine/tent/include/tent/runtime/qos_contract.h new file mode 100644 index 0000000000..334e240e33 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/runtime/qos_contract.h @@ -0,0 +1,110 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_RUNTIME_QOS_CONTRACT_H +#define TENT_RUNTIME_QOS_CONTRACT_H + +#include "tent/common/config.h" +#include "tent/common/status.h" +#include "tent/common/types.h" +#include "tent/thirdparty/nlohmann/json.h" + +#include +#include +#include +#include +#include + +namespace mooncake { +namespace tent { + +struct QosPolicyFields { + std::optional priority; + std::optional max_inflight_bytes; + std::optional max_inflight_requests; + std::optional> allowed_degraded_actions; +}; + +struct QosRequestContext { + // Opaque administrative isolation identifier supplied by the caller. A + // Store integration should pass its canonical Store tenant ID here rather + // than create a second TENT-specific tenant namespace. Empty values + // normalize to "default". + std::string tenant_id = "default"; + // Intent is the normalized business meaning of the transfer + // (foreground_get, background_prefetch, checkpoint, ...). When no explicit + // selects the tenant-local intent override. Empty values normalize to + // "unspec". + std::string intent = "unspec"; + // Legacy caller priority. It is preserved when no resolved contract sets a + // priority, and is overridden by the effective QoS contract priority when + // present. + int requested_priority = PRIO_HIGH; +}; + +struct EffectiveQosPolicy : public QosPolicyFields { + bool enabled = false; + bool matched = false; + // Normalized request identity used for resolution and explain output. + std::string tenant_id = "default"; + std::string intent = "unspec"; + // Name of the contract layer that supplied the most specific match: + // compatibility_default, global_default, .default, or + // .. + std::string matched_contract = "compatibility_default"; + int requested_priority = PRIO_HIGH; + int effective_priority = PRIO_HIGH; +}; + +class QosContractResolver { + public: + Status loadFromConfig(const Config& config); + + bool enabled() const { return enabled_; } + + Status resolve(const QosRequestContext& context, + EffectiveQosPolicy* out) const; + + std::string explainJson(const EffectiveQosPolicy& policy, + int indent = 2) const; + + static std::string intentTypeName(IntentType intent); + + private: + struct TenantContract { + QosPolicyFields defaults; + std::unordered_map intents; + }; + + static Status parsePolicyFields(const json& node, QosPolicyFields* out, + const std::string& path); + static Status parsePriority(const json& node, int* out, + const std::string& path); + static Status parseBytes(const json& node, uint64_t* out, + const std::string& path); + static Status parseUint64(const json& node, uint64_t* out, + const std::string& path); + static std::string normalizeKey(const std::string& value); + static void mergeFields(QosPolicyFields* dst, const QosPolicyFields& src); + static void fieldsToJson(json* out, const QosPolicyFields& fields); + + bool enabled_{false}; + QosPolicyFields global_defaults_; + std::unordered_map tenants_; +}; + +} // namespace tent +} // namespace mooncake + +#endif // TENT_RUNTIME_QOS_CONTRACT_H diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h new file mode 100644 index 0000000000..09b243ab13 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/runtime/receiver_credit.h @@ -0,0 +1,98 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_RUNTIME_RECEIVER_CREDIT_H +#define TENT_RUNTIME_RECEIVER_CREDIT_H + +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" + +namespace mooncake::tent { + +struct ReceiverSessionId { + uint64_t high{0}, low{0}; + bool operator==(const ReceiverSessionId& o) const { + return high == o.high && low == o.low; + } +}; +struct CreditKey { + ReceiverSessionId receiver_session; + uint64_t sender_peer{0}; + uint32_t qos_class{0}; + bool operator==(const CreditKey& o) const { + return receiver_session == o.receiver_session && + sender_peer == o.sender_peer && qos_class == o.qos_class; + } +}; +struct CreditKeyHash { + size_t operator()(const CreditKey&) const noexcept; +}; +enum class CreditResource : uint16_t { + DataBytes = 1, + RequestSlots, + StagingSlots, + ConsumerSlots +}; +constexpr size_t kCreditResourceCount = 4; +struct CreditAmount { + CreditResource resource; + uint64_t grant_total{0}; +}; +struct CreditCharge { + std::vector> resources; +}; +struct ReceiverCreditUpdateV1 { + uint16_t schema_version{1}, flags{0}; + uint32_t qos_class{0}; + ReceiverSessionId receiver_session_id; + uint64_t epoch{0}, sequence{0}; + uint32_t freshness_ttl_ms{0}; + std::vector grants; +}; +enum class CreditUpdateDisposition : uint8_t { + Applied, + DuplicateOrOld, + SequenceGap +}; + +// Private, sender-side state model. It has no network or Admission integration. +class SenderCreditLedger { + public: + explicit SenderCreditLedger(size_t max_entries = 1024) + : max_entries_(max_entries) {} + Status activate(const CreditKey&, uint64_t epoch); + // Removes an exactly matched epoch after the caller has fenced and drained + // (or failed) its transport-owned work. An old cleanup cannot erase a + // reactivated, newer epoch. + Status deactivate(const CreditKey&, uint64_t epoch); + Status applyUpdate(const CreditKey&, const ReceiverCreditUpdateV1&, + CreditUpdateDisposition&); + Status tryReserve(const CreditKey&, const CreditCharge&); + // Only for work not yet handed to a transport; completions need a new + // grant. + Status rollbackReservation(const CreditKey&, const CreditCharge&); + Status available(const CreditKey&, CreditResource, uint64_t&) const; + Status consumed(const CreditKey&, CreditResource, uint64_t&) const; + + private: + struct Entry { + uint64_t epoch{0}, last_sequence{0}; + bool has_update{false}; + std::array grants{}, consumed{}; + }; + static Status resourceIndex(CreditResource, size_t&); + static Status normalize(const CreditCharge&, + std::array&); + mutable std::mutex mutex_; + const size_t max_entries_; + std::unordered_map entries_; +}; + +} // namespace mooncake::tent +#endif diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h index 98be76bcf4..26fccd234b 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment_manager.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -58,6 +59,11 @@ class SegmentManager { // cache invalidation and retry on stale segment cache. Status getRemoteCached(SegmentDesc *&desc, SegmentID handle); + // Owning-reference variant: the returned SegmentDescRef keeps the desc + // (and raw pointers into it) alive independently of the thread-local + // cache's lifetime. + Status getRemoteCached(SegmentDescRef &desc, SegmentID handle); + Status getRemote(SegmentDescRef &desc, const std::string &segment_name); // Invalidates the thread-local cache. @@ -77,18 +83,33 @@ class SegmentManager { // invalidated, the segment is refetched and the operation is retried. template Status withCachedSegment(SegmentID segment_id, Func operation) { + SegmentDescRef pin; + return withCachedSegment(segment_id, pin, operation); + } + + // Same as above, but additionally hands the caller an owning reference + // to the segment snapshot in `pin`. Use this variant whenever a raw + // pointer obtained inside `operation` (e.g. a BufferDesc* from + // findBuffer()) must remain valid after `operation` returns: the pointer + // is guaranteed valid only for as long as `pin` is held. + template + Status withCachedSegment(SegmentID segment_id, SegmentDescRef &pin, + Func operation) { static_assert( std::is_same_v, Status>, "operation must return Status"); - // Local segment: no cache lookup or retry required. + // Local segment: no cache lookup or retry required. Pin the current + // snapshot so pointers into it survive concurrent updateLocal(). if (segment_id == LOCAL_SEGMENT_ID) { - return operation(getLocal().get()); + pin = getLocal(); + return operation(pin.get()); } // First get a cached version. SegmentDesc *desc = nullptr; - CHECK_STATUS(getRemoteCached(desc, segment_id)); + CHECK_STATUS(getRemoteCached(pin, segment_id)); + desc = pin.get(); // Do operation under the cached segment. Status res = operation(desc); @@ -98,7 +119,8 @@ class SegmentManager { // If result status is IsNeedsRefreshCache, invalidate cache and retry invalidateRemote(segment_id); - CHECK_STATUS(getRemoteCached(desc, segment_id)); + CHECK_STATUS(getRemoteCached(pin, segment_id)); + desc = pin.get(); // Do operation again res = operation(desc); @@ -114,14 +136,60 @@ class SegmentManager { } public: - SegmentDescRef getLocal() { return local_desc_; } + // Returns the current immutable snapshot of the local SegmentDesc. + // + // Snapshots are copy-on-write: once published they are never mutated, so + // the returned SegmentDescRef (and raw pointers into it, e.g. from + // findBuffer()) stays valid and self-consistent for as long as the caller + // holds the reference. A reader may observe a snapshot that is stale with + // respect to a concurrent register/unregister, never a torn one — + // staleness is already a designed-in property of segment metadata (remote + // peers cache descs with a TTL plus best-effort invalidation pushes). + // + // Do NOT write through the returned pointer; all mutations must go + // through updateLocal(). + SegmentDescRef getLocal() { + // Per-thread snapshot cache with version-based invalidation — the + // same pattern getRemoteCached() uses for remote descs. The fast + // path is one relaxed-acquire load; the shared_mutex is only taken + // after a publication. Keyed by a monotonic manager id (not `this`) + // so a recycled allocation can never satisfy a stale cache entry. + struct Cache { + uint64_t manager_id = UINT64_MAX; + uint64_t version = 0; + SegmentDescRef ref; + }; + thread_local Cache cache; + auto version = local_desc_version_.load(std::memory_order_acquire); + if (cache.manager_id != manager_id_ || cache.version != version || + !cache.ref) { + std::shared_lock guard(local_desc_lock_); + cache.ref = local_desc_; + cache.manager_id = manager_id_; + // Tag with the pre-lock version: if a publication raced in + // between, the tag mismatches on the next call and we simply + // refresh again — the cache can serve a fresher snapshot than + // its tag, never a staler one. + cache.version = version; + } + return cache.ref; + } + + // Applies `mutator` to a private clone of the local SegmentDesc and + // atomically publishes the result as the new snapshot. Mutations are + // serialized by an internal writer mutex. If `mutator` returns a non-OK + // status, nothing is published. + // + // This is the only way to modify the local SegmentDesc; see getLocal() + // for the snapshot semantics it guarantees. + Status updateLocal(const std::function &mutator); // Returns a serialized JSON snapshot of local_desc_. The result is cached // and shared across concurrent GetSegmentDesc RPC handlers; the cache is - // invalidated by synchronizeLocal() (which is called on every register / - // unregisterLocalMemory). This avoids re-dumping the full segment desc on - // every peer fetch — the previous behavior multiplied dump cost by the - // number of concurrent peer RPCs and dominated remote getRemote latency. + // invalidated by updateLocal() on every publication. This avoids + // re-dumping the full segment desc on every peer fetch — the previous + // behavior multiplied dump cost by the number of concurrent peer RPCs and + // dominated remote getRemote latency. std::shared_ptr getLocalDumpedJson(); Status synchronizeLocal(); @@ -153,7 +221,25 @@ class SegmentManager { std::atomic version_; + // Current published snapshot of the local SegmentDesc. Replaced wholesale + // by updateLocal(); never mutated in place. local_desc_lock_ only guards + // the pointer swap, not the pointee. std::shared_mutex (rather than the + // in-tree RWSpinlock) keeps the synchronization visible to + // ThreadSanitizer. + std::shared_mutex local_desc_lock_; SegmentDescRef local_desc_; + // Serializes clone-mutate-publish cycles in updateLocal(). + std::mutex local_update_mu_; + // Serializes {snapshot, putSegmentDesc} pairs in synchronizeLocal() so a + // stale in-flight put cannot overwrite a newer snapshot in the registry. + std::mutex local_sync_mu_; + // Publication counter; invalidates the per-thread snapshot caches in + // getLocal() and tags local_json_cache_ so that a slow JSON dump of an + // old snapshot can never overwrite the cache entry of a newer one. + std::atomic local_desc_version_{0}; + // Process-unique id for the thread-local cache key in getLocal(). + const uint64_t manager_id_; + ThreadLocalStorage tl_remote_cache_; std::unique_ptr registry_; @@ -164,9 +250,12 @@ class SegmentManager { std::shared_ptr subscribers_lock_; std::shared_ptr> subscribers_; - // Cache for the serialized JSON of local_desc_. Reset by synchronizeLocal. + // Cache for the serialized JSON of local_desc_. Invalidated by + // updateLocal(); local_json_cache_version_ records which publication the + // cached string was computed from. std::mutex local_json_cache_mu_; std::shared_ptr local_json_cache_; + uint64_t local_json_cache_version_{0}; }; } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/segment_tracker.h b/mooncake-transfer-engine/tent/include/tent/runtime/segment_tracker.h index 467a997565..a32aec8d27 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/segment_tracker.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/segment_tracker.h @@ -29,13 +29,18 @@ #include #include "tent/runtime/segment.h" +#include "tent/runtime/segment_manager.h" namespace mooncake { namespace tent { +// Maintains the buffer list of the local SegmentDesc (ref-counted +// registration / deregistration). All mutations are applied through +// SegmentManager::updateLocal(), so concurrent readers of getLocal() always +// observe consistent snapshots; SegmentTracker itself holds no state besides +// the manager reference. class SegmentTracker { public: - SegmentTracker(const SegmentDescRef& local_desc) - : local_desc_(local_desc) {} + explicit SegmentTracker(SegmentManager& manager) : manager_(manager) {} ~SegmentTracker() {} @@ -43,9 +48,6 @@ class SegmentTracker { SegmentTracker& operator==(const SegmentTracker&) = delete; public: - Status query(uint64_t base, size_t length, - std::vector& result); - Status addInBatch(std::vector& desc_list, std::function&)> callback); @@ -55,13 +57,15 @@ class SegmentTracker { Status remove(uint64_t base, size_t length, std::function callback); - Status forEach(std::function callback); + // Iterates over the current snapshot; entries are immutable. Callers + // needing a mutable copy (e.g. transports scrubbing keys during + // deregistration) must copy explicitly. + Status forEach(std::function callback); private: - SegmentDescRef local_desc_; - std::mutex mutex_; + SegmentManager& manager_; }; } // namespace tent } // namespace mooncake -#endif // SEGMENT_TRACKER_H \ No newline at end of file +#endif // SEGMENT_TRACKER_H diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/topology.h b/mooncake-transfer-engine/tent/include/tent/runtime/topology.h index 9360b5e2f4..d0c8d0524a 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/topology.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/topology.h @@ -36,17 +36,29 @@ namespace tent { class Platform; class Topology { public: - const static size_t DevicePriorityRanks = 3; - - enum NicType { NIC_RDMA, NIC_TCP, NIC_UNKNOWN }; + inline static constexpr size_t DevicePriorityRanks = 3; + + // Keep the existing RDMA/TCP numeric values stable for serialized + // topologies. UB is a distinct link type and must never be selected as an + // RDMA verbs device. + enum NicType { + NIC_RDMA = 0, + NIC_TCP = 1, + NIC_UNKNOWN = 2, + NIC_UB = 3, + }; enum MemType { MEM_HOST, MEM_CUDA, MEM_ROCM, MEM_ASCEND, MEM_UNKNOWN }; using NicID = int; struct NicEntry { std::string name; std::string pci_bus_id; - NicType type; - int numa_node; + NicType type{NIC_UNKNOWN}; + int numa_node{-1}; + + // Hardware-specific discovery metadata is opaque to the common + // topology layer. Transports own namespaced keys and interpretation. + std::unordered_map device_attrs{}; }; using MemID = int; @@ -67,8 +79,12 @@ class Topology { void clear(); + // Preserve the original one-argument symbol for source and binary + // compatibility with callers that do not opt into UB discovery. Status discover(const std::vector& platforms); + Status discover(const std::vector& platforms, bool discover_ub); + Status parse(const std::string& json_content); std::string toString() const; @@ -140,4 +156,4 @@ struct RangeLocation { } // namespace tent } // namespace mooncake -#endif // TENT_TOPOLOGY_H \ No newline at end of file +#endif // TENT_TOPOLOGY_H diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h index 3eae2e4198..38b1189987 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h @@ -30,14 +30,14 @@ #include "tent/common/config.h" #include "tent/common/status.h" #include "tent/common/types.h" -#include "tent/common/concurrent/thread_local_storage.h" +#include "tent/runtime/admission_queue.h" +#include "tent/runtime/transport.h" #include "tent/runtime/transport_selector.h" namespace mooncake { namespace tent { class Batch; -class BatchSet; class Topology; class Transport; class SegmentDesc; @@ -55,11 +55,23 @@ struct TaskInfo { int xport_priority{0}; // transport priority (for fallback) int failover_count{0}; // number of failover attempts uint64_t device_mask{~0ULL}; // Device mask for quota allocation + std::string qp_pool; // Named QP pool (RFC #2568 step 3), "" = none Request request; bool staging{false}; + bool cancel_requested{false}; TransferStatusEnum status{TransferStatusEnum::PENDING}; volatile TransferStatusEnum staging_status{TransferStatusEnum::PENDING}; - std::chrono::steady_clock::time_point start_time{}; // For latency tracking + std::chrono::steady_clock::time_point start_time{}; // Request submit + std::chrono::steady_clock::time_point dispatch_time{}; // Initial dispatch + std::chrono::steady_clock::time_point post_time{}; // Initial post + // Current physical attempt. Replaced before each transport submission; + // post_time above intentionally remains the logical request's first post. + // attempt_type is captured at attempt start so the attempt is attributed to + // the transport that actually ran it, even if task.type is later + // overwritten by failover before the attempt is finished. + std::chrono::steady_clock::time_point attempt_post_time{}; + TransportType attempt_type{UNSPEC}; + bool attempt_active{false}; }; class TransferEngineImpl { @@ -139,6 +151,8 @@ class TransferEngineImpl { const std::vector& request_list, const Notification& notifi); + Status cancelTransfer(BatchID batch_id, size_t task_id); + Status sendNotification(SegmentID target_id, const Notification& notifi); Status receiveNotification(std::vector& notifi_list); @@ -155,6 +169,8 @@ class TransferEngineImpl { Status progressBatch(BatchID batch_id, TransferStatus& overall_status); + Status getNicLoadStats(std::vector& stats) const; + Status waitTransferCompletion(BatchID batch_id); Status transferSync(const std::vector& request_list); @@ -176,11 +192,13 @@ class TransferEngineImpl { } // Wake the optional event-driven progress worker for `batch_id`. No-op if - // enable_progress_worker is false. Currently used by test/integration - // hooks; transports will be migrated to call this in a follow-up PR. + // enable_progress_worker is false. Transport completion paths use this as + // an idempotent "maybe ready" signal. void notifyBatchMaybeReady(BatchID batch_id); private: + friend class ProgressWorker; + Status construct(); Status deconstruct(); @@ -197,6 +215,56 @@ class TransferEngineImpl { Status resubmitTransferTask(Batch* batch, size_t task_id); + Status retainBatch(BatchID batch_id, Batch*& batch); + + Status releaseBatch(Batch* batch); + + class BatchRef; + + struct PreparedSubmit; + + Status submitTransfer(BatchID batch_id, + const std::vector& request_list, + const Notification* notifi, + QueueOwnerKind owner_kind); + + Status submitStagingTransfer(BatchID batch_id, + const std::vector& request_list); + + Status enqueuePreparedSubmit(Batch* batch, const PreparedSubmit& prepared, + QueueOwnerKind owner_kind); + + bool shouldQueueSubmit(const PreparedSubmit& prepared, + QueueOwnerKind owner_kind) const; + + Status prepareSubmit(Batch* batch, const std::vector& request_list, + PreparedSubmit& prepared); + + Status commitPreparedSubmit(Batch* batch, const PreparedSubmit& prepared); + + void attachProgressNotifier(Batch* batch, Transport::SubBatchRef sub_batch); + + uint64_t nextBatchToken(); + + Status refillDispatchWindow(); + + Status progressRuntimeQueue(); + + bool hasActiveRuntimeQueue(); + + void notifyRuntimeQueueReady(); + + Status dispatchQueuedOwner(QueueOwnerId owner_id); + + Status markQueuedOwnerSubmitted(QueueOwnerId owner_id); + + Status finishQueuedOwner(QueueOwnerId owner_id, + TransferStatusEnum terminal_status); + + Status cancelQueuedOwner(QueueOwnerId owner_id); + + Status retireQueueForBatch(Batch* batch); + Status pollTaskStatus(Batch* batch, size_t task_id, TransferStatus& task_status); @@ -220,10 +288,20 @@ class TransferEngineImpl { Status maybeFireSubmitHooks(Batch* batch, bool check = true); + void addSubmitHook(Batch* batch, size_t start_task_id, + const std::vector& request_list, + const Notification& notifi); + void recordTaskCompletionMetrics(TaskInfo& task, TransferStatusEnum prev_status, TransferStatusEnum new_status); + void startTransportAttempt(TaskInfo& task, TransportType type, + std::chrono::steady_clock::time_point post_time); + + void finishTransportAttempt(TaskInfo& task, TransferStatusEnum status, + std::chrono::steady_clock::time_point end_time); + private: struct AllocatedMemory { void* addr; @@ -237,6 +315,22 @@ class TransferEngineImpl { std::vector freelist; }; + struct RuntimeQueueConfig { + bool enabled{false}; + QueueLimits limits{}; + size_t max_dispatch_owners{0}; + size_t max_dispatch_bytes{0}; + std::chrono::microseconds progress_fallback_interval{50000}; + }; + + struct QueuedOwnerState { + Batch* batch{nullptr}; + size_t owner_task_id{0}; + std::vector public_task_ids; + size_t byte_charge{0}; + bool in_dispatch_window{false}; + }; + private: std::shared_ptr conf_; std::shared_ptr metadata_; @@ -248,7 +342,7 @@ class TransferEngineImpl { transport_list_; std::unique_ptr local_segment_tracker_; - ThreadLocalStorage batch_set_; + BatchSet batch_set_; std::vector allocated_memory_; std::mutex mutex_; @@ -263,6 +357,12 @@ class TransferEngineImpl { int max_failover_attempts_{3}; bool enable_auto_failover_on_poll_{true}; bool enable_progress_worker_{false}; + RuntimeQueueConfig runtime_queue_config_; + std::unique_ptr runtime_queue_; + std::unordered_map queued_owners_; + size_t dispatch_inflight_owners_{0}; + size_t dispatch_inflight_bytes_{0}; + uint64_t next_batch_token_{1}; // Guards alive_batches_ and serializes pollTaskStatus / // updateTaskStatusAfterPoll / lazyFreeBatch against the optional @@ -275,4 +375,4 @@ class TransferEngineImpl { } // namespace tent } // namespace mooncake -#endif \ No newline at end of file +#endif diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h index 39925c1c53..93bf585a9c 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport.h @@ -22,12 +22,14 @@ #include #include +#include #include #include #include #include #include "tent/common/status.h" +#include "tent/common/types.h" #include "tent/runtime/platform.h" #include "tent/runtime/control_plane.h" @@ -48,17 +50,27 @@ class Transport { SubBatch() : device_mask(~0ULL) {} virtual ~SubBatch() {} virtual size_t size() const = 0; + void notifyProgress() { + if (notify_progress) notify_progress(progress_batch_id); + } + uint64_t device_mask; // Device mask for transport selection + // Named QP pool for this batch's transfers (RFC #2568 step 3). Empty = + // no pool (default spray). Carried like device_mask, from the matched + // SelectionPolicy down to each RdmaTask. + std::string qp_pool; + BatchID progress_batch_id{0}; + std::function notify_progress; }; - using SubBatchRef = SubBatch *; + using SubBatchRef = SubBatch*; public: Transport() = default; virtual ~Transport() = default; - virtual Status install(std::string &local_segment_name, + virtual Status install(std::string& local_segment_name, std::shared_ptr metadata, std::shared_ptr local_topology, std::shared_ptr conf = nullptr) { @@ -69,56 +81,67 @@ class Transport { virtual const Capabilities capabilities() const { return caps; } - virtual Status allocateSubBatch(SubBatchRef &batch, size_t max_size) { + virtual Status allocateSubBatch(SubBatchRef& batch, size_t max_size) { return Status::NotImplemented( "allocateSubBatch not implemented" LOC_MARK); } - virtual Status freeSubBatch(SubBatchRef &batch) { + virtual Status freeSubBatch(SubBatchRef& batch) { return Status::NotImplemented("freeSubBatch not implemented" LOC_MARK); } virtual Status submitTransferTasks( - SubBatchRef batch, const std::vector &request_list) { + SubBatchRef batch, const std::vector& request_list) { return Status::NotImplemented( "submitTransferTasks not implemented" LOC_MARK); } virtual Status getTransferStatus(SubBatchRef batch, int task_id, - TransferStatus &status) { + TransferStatus& status) { return Status::NotImplemented( "getTransferStatus not implemented" LOC_MARK); } - virtual Status allocateLocalMemory(void **addr, size_t size, - MemoryOptions &options) { + // Cancellation is best effort: implementations must prevent work that has + // not reached the device from being submitted, but work already posted to + // a device may still complete. Callers must continue polling until the + // task reaches a terminal state. + virtual bool supportsCancellation() const { return false; } + + virtual Status cancelTransferTask(SubBatchRef batch, int task_id) { + return Status::NotImplemented( + "cancelTransferTask not implemented" LOC_MARK); + } + + virtual Status allocateLocalMemory(void** addr, size_t size, + MemoryOptions& options) { return Platform::getLoader().allocate(addr, size, options); } - virtual Status freeLocalMemory(void *addr, size_t size) { + virtual Status freeLocalMemory(void* addr, size_t size) { return Platform::getLoader().free(addr, size); } // Pre-registration warm-up that pins pages before NUMA probing. // Returns true if pages were successfully pinned (caller may skip // prefault). Default: no-op, returns false. - virtual bool warmupMemory(void *addr, size_t length) { return false; } + virtual bool warmupMemory(void* addr, size_t length) { return false; } - virtual Status addMemoryBuffer(BufferDesc &desc, - const MemoryOptions &options) { + virtual Status addMemoryBuffer(BufferDesc& desc, + const MemoryOptions& options) { return Status::NotImplemented( "addMemoryBuffer not implemented" LOC_MARK); } - virtual Status addMemoryBuffer(std::vector &desc_list, - const MemoryOptions &options) { - for (auto &desc : desc_list) { + virtual Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) { + for (auto& desc : desc_list) { CHECK_STATUS(addMemoryBuffer(desc, options)); } return Status::OK(); } - virtual Status removeMemoryBuffer(BufferDesc &desc) { + virtual Status removeMemoryBuffer(BufferDesc& desc) { return Status::NotImplemented( "removeMemoryBuffer not implemented" LOC_MARK); } @@ -126,17 +149,23 @@ class Transport { virtual bool supportNotification() const { return false; } virtual Status sendNotification(SegmentID target_id, - const Notification ¬ify) { + const Notification& notify) { return Status::NotImplemented( "sendNotification not implemented" LOC_MARK); } - virtual Status receiveNotification(std::vector ¬ify_list) { + virtual Status receiveNotification(std::vector& notify_list) { return Status::NotImplemented( "receiveNotification not implemented" LOC_MARK); } - virtual const char *getName() const { return ""; } + virtual const char* getName() const { return ""; } + + virtual double getEstimatedBandwidth() const { return -1.0; } + + virtual Status getNicLoadStats(std::vector&) const { + return Status::OK(); + } protected: Capabilities caps; diff --git a/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h b/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h index eaa2341f33..63773ceeb2 100644 --- a/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h +++ b/mooncake-transfer-engine/tent/include/tent/runtime/transport_selector.h @@ -82,6 +82,8 @@ struct SelectionContext { int priority_level; // Request priority level (lower = more urgent) std::optional policy_name; // Optional: bind to specific policy by name + IntentType intent_type{ + IntentType::INTENT_UNSPEC}; // Business intent for policy matching }; /** @@ -117,6 +119,18 @@ struct SelectionPolicy { // Transport preference list (evaluated in order) std::vector transports; + + // Per-policy link-layer QoS. nullopt = fall back to the global RdmaParams + // value. InfiniBand Service Level (0-15) and Traffic Class / DSCP (0-255). + std::optional service_level; + std::optional traffic_class; + // Named QP pool this policy's traffic should land on; parsed and stored for + // now, routing to be wired later. Unset = the current single "data QP". + std::optional qp_pool; + + // Optional business-intent filter. nullopt preserves the historical + // catch-all behavior; otherwise the request intent must match exactly. + std::optional intent_type; }; /** @@ -125,6 +139,10 @@ struct SelectionPolicy { struct SelectionResult { TransportType transport = UNSPEC; uint64_t device_mask = ~0ULL; // Bitmask of allowed devices (~0 = all) + // Resolved link-layer QoS from the matched policy (nullopt = default). + std::optional service_level; + std::optional traffic_class; + std::optional qp_pool; }; /** @@ -167,8 +185,6 @@ class TransportSelector { */ bool isLegacyMode() const { return legacy_mode_; } - static std::string transportTypeName(TransportType type); - static TransportType parseTransportType(const std::string& str); static std::optional> reorderWithHint( const std::vector& raw, TransportType hint); diff --git a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h index ff93e0ed02..c13c74564c 100644 --- a/mooncake-transfer-engine/tent/include/tent/transfer_engine.h +++ b/mooncake-transfer-engine/tent/include/tent/transfer_engine.h @@ -109,6 +109,8 @@ typedef struct tent_notifi_info tent_notifi_info; #define TRANSPORT_TCP (7) #define TRANSPORT_ASCEND_DIRECT (8) #define TRANSPORT_SUNRISE_LINK (9) +#define TRANSPORT_TPU (10) +#define TRANSPORT_UB (11) struct tent_memory_options { char location[64]; @@ -174,6 +176,9 @@ void tent_free_notifs(tent_notifi_info* info); int tent_task_status(tent_engine_t engine, tent_batch_id_t batch_id, size_t task_id, tent_status_t* status); +int tent_cancel_task(tent_engine_t engine, tent_batch_id_t batch_id, + size_t task_id); + int tent_overall_status(tent_engine_t engine, tent_batch_id_t batch_id, tent_status_t* status); @@ -201,6 +206,17 @@ int tent_register_memory_batch_ex(tent_engine_t engine, void** addrs, int tent_task_status_list(tent_engine_t engine, tent_batch_id_t batch_id, tent_status_t* statuses, size_t* count); +struct tent_nic_load_stat { + char device_name[64]; + uint64_t inflight_bytes; + double ewma_bandwidth_bps; +}; + +typedef struct tent_nic_load_stat tent_nic_load_stat_t; + +int tent_get_nic_load_stats(tent_engine_t engine, tent_nic_load_stat_t* stats, + size_t* count); + #ifdef __cplusplus } #endif // __cplusplus @@ -299,6 +315,11 @@ class TransferEngine { const std::vector& request_list, const Notification& notifi); + // Best-effort task cancellation. Work that has not reached the transport + // is prevented from being submitted. Device work already posted may still + // complete, so callers must continue polling for a terminal status. + Status cancelTransfer(BatchID batch_id, size_t task_id); + Status sendNotification(SegmentID target_id, const Notification& notifi); Status receiveNotification(std::vector& notifi_list); @@ -321,6 +342,8 @@ class TransferEngine { // progress later"; terminal states (COMPLETED/FAILED) will not be revived. Status progressBatch(BatchID batch_id, TransferStatus& overall_status); + Status getNicLoadStats(std::vector& stats) const; + private: std::unique_ptr impl_; }; @@ -328,4 +351,4 @@ class TransferEngine { } // namespace mooncake #endif -#endif \ No newline at end of file +#endif diff --git a/mooncake-transfer-engine/tent/include/tent/transport/gds/gds_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/gds/gds_transport.h index a184ea9528..d63c42bfe6 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/gds/gds_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/gds/gds_transport.h @@ -36,8 +36,11 @@ namespace tent { class GdsFileContext; struct IOParamRange { - size_t base; - size_t count; + size_t base = 0; + size_t count = 0; + size_t complete_count = 0; + size_t transferred_bytes = 0; + TransferStatusEnum status = TransferStatusEnum::PENDING; }; // Wrapper for reusable CUfileBatchHandle_t diff --git a/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h index 8986038bba..762a0dc9a5 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/mnnvl/mnnvl_transport.h @@ -20,6 +20,7 @@ #include #include +#include #include #include #include @@ -37,7 +38,6 @@ struct MnnvlTask { volatile TransferStatusEnum status_word; volatile size_t transferred_bytes; uint64_t target_addr = 0; - int cuda_id = 0; cudaEvent_t completion_event = nullptr; }; @@ -46,6 +46,7 @@ struct MnnvlSubBatch : public Transport::SubBatch { size_t max_size; CUDAStreamHandle sync_stream; CUDAStreamHandle async_stream; + int stream_device_id = -1; // Completion events created in startTransfer (one per submit). Destroyed by // the destructor (RAII); Slab::deallocate() invokes ~MnnvlSubBatch() // before reusing the storage, so this runs on every free. @@ -135,4 +136,4 @@ class MnnvlTransport : public Transport { } // namespace tent } // namespace mooncake -#endif // MNNVL_TRANSPORT_H_ \ No newline at end of file +#endif // MNNVL_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h index 38446c6456..fd786f08b2 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/nvlink/nvlink_transport.h @@ -40,7 +40,6 @@ struct NVLinkTask { volatile size_t transferred_bytes; uint64_t target_addr = 0; bool is_cuda_ipc; - int cuda_id = 0; cudaEvent_t completion_event = nullptr; }; @@ -49,6 +48,7 @@ struct NVLinkSubBatch : public Transport::SubBatch { size_t max_size; CUDAStreamHandle sync_stream; CUDAStreamHandle async_stream; + int stream_device_id = -1; // Completion events created in startTransfer (one per submit). Destroyed by // the destructor (RAII); Slab::deallocate() invokes ~NVLinkSubBatch() // before reusing the storage, so this runs on every free. @@ -124,10 +124,10 @@ class NVLinkTransport : public Transport { uint64_t async_memcpy_threshold_; bool host_register_; - std::mutex register_mutex_; + mutable std::mutex register_mutex_; std::unordered_set registered_base_addrs_; }; } // namespace tent } // namespace mooncake -#endif // NVLINK_TRANSPORT_H_ \ No newline at end of file +#endif // NVLINK_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/bw_arbitration.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/bw_arbitration.h new file mode 100644 index 0000000000..046ff6a341 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/bw_arbitration.h @@ -0,0 +1,83 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Deadline-aware NIC bandwidth arbitration WITHIN a priority tier (RFC #2792). +// +// TENT's QoS is otherwise vertical: SL/TC and priority tiers separate business +// classes. But when several flows in the SAME tier contend for one NIC, the +// NIC's bandwidth is split blindly/equally (measured: a ~388 Gb/s NIC gives +// ~97 Gb/s to each of 4 contending flows). There is no way to let a flow that +// is about to miss its deadline claim a larger share. +// +// This header isolates the pure ordering decision so it can be unit-tested +// without the RDMA stack: given the contending slices' (deadline_ns, length) +// and a bandwidth estimate, order them most-urgent-first by predicted MLU +// (predicted transfer time / remaining deadline window) — the same MLU used by +// the admission layer (#2618/#2764). Opt-in: when disabled, the original order +// is preserved exactly (byte-identical to today's equal split). + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace mooncake { +namespace tent { + +// Minimal view of a contending slice for arbitration. Kept free of RdmaSlice +// so the policy is unit-testable in isolation. +struct ArbFlow { + uint64_t deadline_ns; // 0 == no deadline + size_t length; // bytes to transfer +}; + +// Predicted Missed-Latency-per-Unit: predicted transfer time / remaining +// deadline window. Higher == more urgent (closer to / past its deadline). +// A flow with no deadline (deadline_ns == 0) is least urgent (MLU 0). A flow +// already past its deadline, or that cannot fit its window at the given +// bandwidth, gets a very high MLU so it sorts first. bw_bps <= 0 disables +// prediction (returns 0 for everyone == no reordering). +inline double PredictedMlu(const ArbFlow& f, uint64_t now_ns, double bw_bps) { + if (f.deadline_ns == 0 || bw_bps <= 0.0) return 0.0; + if (f.deadline_ns <= now_ns) return std::numeric_limits::max(); + const double window_s = (f.deadline_ns - now_ns) / 1e9; + const double predicted_time_s = static_cast(f.length) / bw_bps; + return predicted_time_s / window_s; +} + +// Return the indices of `flows` ordered most-urgent-first (highest predicted +// MLU first). Ties and the no-deadline case keep original (FIFO) order — a +// stable sort — so a symmetric set degrades to today's behavior. This does not +// drop or admit anything; it only reorders selection among already-eligible, +// same-tier flows. +inline std::vector OrderByUrgency(const std::vector& flows, + uint64_t now_ns, double bw_bps) { + std::vector idx(flows.size()); + std::iota(idx.begin(), idx.end(), size_t{0}); + std::vector mlu(flows.size()); + for (size_t i = 0; i < flows.size(); ++i) { + mlu[i] = PredictedMlu(flows[i], now_ns, bw_bps); + } + std::stable_sort(idx.begin(), idx.end(), [&](size_t a, size_t b) { + return mlu[a] > mlu[b]; // higher MLU first; stable keeps FIFO on ties + }); + return idx; +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/cq.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/cq.h index 78cacec481..f80f060391 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/cq.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/cq.h @@ -19,6 +19,8 @@ #include #include +#include + #include "tent/common/status.h" namespace mooncake { @@ -29,6 +31,10 @@ class RdmaContext; class RdmaCQ { public: RdmaCQ() : cq_(nullptr), cqe_now_(0), cqe_limit_(-1), context_(nullptr) {} + RdmaCQ(const RdmaCQ &) = delete; + RdmaCQ &operator=(const RdmaCQ &) = delete; + RdmaCQ(RdmaCQ &&) = delete; + RdmaCQ &operator=(RdmaCQ &&) = delete; ~RdmaCQ(); @@ -43,11 +49,11 @@ class RdmaCQ { void cancelQuota(int num_entries); - int getQuota() const { return cqe_now_; } + int getQuota() const { return cqe_now_.load(std::memory_order_relaxed); } int maxCqe() const { return cqe_limit_; } - bool empty() const { return cqe_now_ == 0; } + bool empty() const { return cqe_now_.load(std::memory_order_relaxed) == 0; } int poll(int num_entries, ibv_wc *wc); @@ -57,7 +63,7 @@ class RdmaCQ { private: ibv_cq *cq_; - volatile int cqe_now_; + std::atomic cqe_now_; int cqe_limit_; RdmaContext *context_; }; diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h index 7ef3588b89..79e3bafa48 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint.h @@ -15,6 +15,7 @@ #ifndef TENT_ENDPOINT_H #define TENT_ENDPOINT_H +#include #include #include #include @@ -26,7 +27,7 @@ namespace mooncake { namespace tent { class RdmaEndPoint : public std::enable_shared_from_this { struct WrDepthBlock { - volatile int value; + std::atomic value; uint64_t padding[7]; }; @@ -93,8 +94,7 @@ class RdmaEndPoint : public std::enable_shared_from_this { }; int construct(RdmaContext* context, EndPointParams* params, - const std::string& endpoint_name, - std::atomic* endpoints_count = nullptr); + const std::string& endpoint_name); int deconstruct(); @@ -107,7 +107,6 @@ class RdmaEndPoint : public std::enable_shared_from_this { // WRs have been drained, actually destroys QPs and frees resources. // Returns true if destruction is complete, false if outstanding WRs remain. void beginDestroy(); - void beginDestroyNoLock(); // Internal version without locking bool finishDestroy(); Status connect(const std::string& peer_server_name, @@ -162,7 +161,7 @@ class RdmaEndPoint : public std::enable_shared_from_this { size_t acknowledge(RdmaSlice* slice, TransferStatusEnum status); - volatile int* getQuotaCounter(int qp_index) const { + std::atomic* getQuotaCounter(int qp_index) const { return &wr_depth_list_[qp_index].value; } @@ -174,12 +173,17 @@ class RdmaEndPoint : public std::enable_shared_from_this { int setupOneQP(int qp_index, const std::string& peer_gid, uint16_t peer_lid, uint32_t peer_qp_num, std::string* reply_msg = nullptr); + // Returns the pool segment owning qp_index, or nullptr when no pools are + // configured (the default single-pool case). Read-only after construct(). + const QpPoolSegment* poolForQp(int qp_index) const; + bool reserveQuota(int qp_index, int num_entries); void cancelQuota(int qp_index, int num_entries); private: // Caller must hold lock_ in write mode. + void beginDestroyNoLock(); int deconstructUnlocked(); void resetInflightSlices(); @@ -188,17 +192,23 @@ class RdmaEndPoint : public std::enable_shared_from_this { void repostAllNotifyRecvs(); private: + friend class EndpointTestAccess; + std::atomic status_; RdmaContext* context_; EndPointParams* params_; std::string endpoint_name_; std::vector qp_list_; + // Per-pool QP layout, resolved once in construct() from params_->qp_pools. + // Empty = default single pool spanning all of qp_list_. Each segment's + // [begin, begin+num_qp) indexes into qp_list_. Read-only after construct(). + std::vector qp_pool_segments_; // Each data QP queue is owned by exactly one worker lane; reset/deconstruct // are synchronized by the endpoint lifecycle lock. std::vector slice_queue_; WrDepthBlock* wr_depth_list_; - volatile int inflight_slices_; + std::atomic inflight_slices_; uint32_t padding_[7]; RWSpinlock lock_; @@ -208,8 +218,6 @@ class RdmaEndPoint : public std::enable_shared_from_this { std::string peer_server_name_; std::string peer_nic_name_; std::vector peer_qp_num_list_; - std::atomic* endpoints_count_; - // Notification QP (one per endpoint for control plane operations) ibv_qp* notify_qp_ = nullptr; @@ -220,16 +228,18 @@ class RdmaEndPoint : public std::enable_shared_from_this { std::vector notify_recv_mrs_; // Memory regions for recv buffers std::vector notify_send_buffer_; // Single contiguous send buffer ibv_mr* notify_send_mr_ = nullptr; // Single MR for all send slots + // Serializes notification buffer/QP access against deconstruction. + std::mutex notify_resource_mutex_; std::mutex notify_send_mutex_; std::condition_variable notify_send_cv_; - int notify_pending_count_ = 0; // Number of pending sends - uint64_t notify_send_wr_id_ = 0; // Circular counter for wr_id - bool notify_connected_ = false; + size_t notify_pending_count_ = 0; // Number of pending sends + uint64_t notify_send_wr_id_ = 0; // Circular counter for wr_id + std::atomic notify_connected_{false}; // Two-phase destruction constants (matching TE) static constexpr double kFinishDestroyTimeoutSec = 30.0; - static constexpr int kFinishDestroyMaxRetries = 3; - int finish_destroy_retries_ = 0; // Retry counter for finishDestroy + static constexpr int kMaxDestroyErrorLogs = 3; + int destroy_error_count_ = 0; }; } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint_store.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint_store.h index 8655681616..48f656d0f7 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint_store.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/endpoint_store.h @@ -28,6 +28,8 @@ using namespace mooncake; namespace mooncake { namespace tent { +class EndpointStoreTestAccess; + class EndpointStore { public: virtual std::shared_ptr get(const std::string &key) = 0; @@ -37,7 +39,11 @@ class EndpointStore { virtual int remove(RdmaEndPoint *ep) = 0; - virtual void clear() = 0; + // Terminal context-shutdown operation. Unlike remove()/reclaim(), clear() + // synchronously releases every endpoint's verbs resources because workers + // have stopped and can no longer drain CQ flush completions. Callers must + // not use the store again after clear() returns. + virtual int clear() = 0; virtual size_t size() = 0; @@ -57,7 +63,7 @@ class FIFOEndpointStore : public EndpointStore { int remove(RdmaEndPoint *ep) override; - void clear(); + int clear() override; size_t size() override; @@ -65,6 +71,8 @@ class FIFOEndpointStore : public EndpointStore { void reclaim() override; private: + friend class EndpointStoreTestAccess; + RdmaContext &context_; RWSpinlock endpoint_map_lock_; std::unordered_map> @@ -87,7 +95,7 @@ class SIEVEEndpointStore : public EndpointStore { std::shared_ptr getOrInsert(const std::string &key) override; - void clear(); + int clear() override; size_t size() override; @@ -96,6 +104,8 @@ class SIEVEEndpointStore : public EndpointStore { void reclaim() override; private: + friend class EndpointStoreTestAccess; + RdmaContext &context_; RWSpinlock endpoint_map_lock_; // The bool represents visited @@ -111,9 +121,8 @@ class SIEVEEndpointStore : public EndpointStore { std::atomic waiting_list_len_; size_t max_size_; - std::atomic endpoints_count_{0}; }; } // namespace tent } // namespace mooncake -#endif \ No newline at end of file +#endif diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/gdr_reachability.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/gdr_reachability.h new file mode 100644 index 0000000000..54d824f485 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/gdr_reachability.h @@ -0,0 +1,112 @@ +// Copyright 2025 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_GDR_REACHABILITY_H +#define TENT_GDR_REACHABILITY_H + +#include +#include +#include +#include +#include +#include + +namespace mooncake { +namespace tent { + +class Config; + +// GdrReachability tracks whether a NIC can actually GPUDirect-DMA (P2P) to a +// GPU. ibv_reg_mr succeeds on a GPU buffer for every NIC (nvidia-peermem maps +// the GPU pages per-GPU and hands back a BAR address), so registration cannot +// tell whether a given NIC's PCIe path can really reach the GPU -- that depends +// on PCIe topology / ACS and only surfaces on the data plane as +// IBV_WC_LOC_PROT_ERR (local NIC -> local GPU) or IBV_WC_REM_ACCESS_ERR +// (remote NIC -> remote GPU). +// +// This component learns unreachable (NIC, GPU) pairs at runtime (and, when the +// registration probe is enabled, up front) and lets device selection skip +// them. A permissive fabric never records anything and keeps full multi-rail +// aggregation; a restrictive fabric converges onto the reachable NIC(s) instead +// of exhausting retries and failing the whole transfer. +// +// Exclusion uses the same threshold + exponential-cooldown + re-admit scheme as +// RailMonitor, so a transient error self-heals after one cooldown and a +// genuinely-dead path is only re-probed occasionally. +// +// It is a process-wide singleton keyed by stable identifiers (RDMA device name +// + GPU ordinal, plus peer machine id for the remote side), so a single +// instance is correct even across multiple TransferEngine instances in one +// process. All methods are thread-safe. +class GdrReachability { + public: + static GdrReachability &instance(); + + // Cheap global fast path: true only after at least one pair has ever been + // excluded. Hot selection paths skip all work while this is false. + static bool hasAnyExclusion() { + return any_exclusion_.load(std::memory_order_relaxed); + } + + // Override error_threshold / error_window / cooldown from config. Config + // keys mirror RailMonitor: transports/rdma/gdr_error_threshold, + // transports/rdma/gdr_error_window_secs, transports/rdma/gdr_cooldown_secs. + void configure(const Config *conf); + + // --- Local side: this host's NIC -> this host's GPU --- + void reportLocalFailure(const std::string &nic_name, int gpu_ordinal); + void reportLocalSuccess(const std::string &nic_name, int gpu_ordinal); + bool localReachable(const std::string &nic_name, int gpu_ordinal); + + // --- Remote side: peer `machine_id`'s NIC -> that peer's GPU --- + void reportRemoteFailure(const std::string &machine_id, + const std::string &nic_name, int gpu_ordinal); + void reportRemoteSuccess(const std::string &machine_id, + const std::string &nic_name, int gpu_ordinal); + bool remoteReachable(const std::string &machine_id, + const std::string &nic_name, int gpu_ordinal); + + private: + GdrReachability() = default; + + struct State { + uint32_t error_count = 0; + std::chrono::steady_clock::time_point last_error{}; + std::chrono::steady_clock::duration cooldown{std::chrono::seconds(0)}; + std::chrono::steady_clock::time_point resume_time{}; // paused until + }; + + bool reachable(const std::string &key); + void markFailed(const std::string &key); + void markRecovered(const std::string &key); + + static std::string localKey(const std::string &nic_name, int gpu_ordinal); + static std::string remoteKey(const std::string &machine_id, + const std::string &nic_name, int gpu_ordinal); + + static std::atomic any_exclusion_; + + std::shared_mutex mutex_; + std::unordered_map state_; + + int error_threshold_ = 2; + std::chrono::seconds error_window_{10}; + std::chrono::seconds cooldown_{30}; + static constexpr std::chrono::seconds kMaxCooldown{300}; +}; + +} // namespace tent +} // namespace mooncake + +#endif // TENT_GDR_REACHABILITY_H diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/params.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/params.h index 0912be3504..b438281dd6 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/params.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/params.h @@ -15,8 +15,12 @@ #ifndef TENT_PARAMS_H #define TENT_PARAMS_H +#include + #include #include +#include +#include namespace mooncake { namespace tent { @@ -29,6 +33,24 @@ struct DeviceParams { int max_cqe = 4096; }; +// One named QP pool. When SelectionPolicy entries declare a `qp_pool`, each +// distinct pool gets its own contiguous run of data QPs inside every endpoint, +// handshaked with that pool's SL/TC. Transfers routed to a pool only use its +// QPs, giving link-layer isolation between traffic classes (RFC #2568 step 2). +// +// Layering note: the wire handshake is unchanged. Both peers derive the same +// pool layout from the same SelectionPolicy config, so the flat qp_num list is +// still paired positionally — pool P's i-th QP on one side lines up with pool +// P's i-th QP on the other. This keeps BootstrapDesc byte-compatible with peers +// that don't know about pools (they simply run a single default pool). +struct QpPoolSegment { + std::string name; + int num_qp = 0; // Number of data QPs dedicated to this pool. + int begin = 0; // Index into qp_list_ where this pool's QPs start. + int service_level = -1; // -1 = fall back to EndPointParams::service_level. + int traffic_class = -1; // -1 = fall back to EndPointParams::traffic_class. +}; + struct EndPointParams { int endpoint_store_cap = 65536; int qp_mul_factor = 6; // Derived from RdmaParams::num_lanes. @@ -37,6 +59,12 @@ struct EndPointParams { int max_inline_bytes = 64; ibv_mtu path_mtu = IBV_MTU_4096; + // Named QP pools. Empty (default) = today's behavior: a single homogeneous + // run of qp_mul_factor data QPs, all handshaked with the global SL/TC. When + // non-empty, the segments partition the data QPs by pool; total QP count is + // the sum of per-pool num_qp. Both peers must derive the same layout. + std::vector qp_pools; + // Advanced parameters, do not change unless you understand them // INIT State uint16_t pkey_index = 0; @@ -58,6 +86,67 @@ struct EndPointParams { uint8_t max_rd_atomic = 16; }; +// Result of resolving the QP-pool layout: the concrete per-pool segments (each +// with its [begin, begin+num_qp) span filled in) and the total data-QP count. +struct QpPoolLayout { + std::vector segments; + int total_qp = 0; + bool valid = false; // false = invalid config (non-positive total). +}; + +// Pure resolver used by RdmaEndPoint::construct(). Kept free-standing (no RDMA +// handles) so the layout math can be unit-tested. Empty `pools` reproduces the +// historical single homogeneous run of `qp_mul_factor` data QPs (segments left +// empty, meaning "one default pool"); a non-empty config lays out one +// contiguous segment per pool and the total is the sum of per-pool num_qp. +inline QpPoolLayout computeQpPoolSegments( + const std::vector& pools, int qp_mul_factor) { + QpPoolLayout layout; + if (pools.empty()) { + layout.total_qp = qp_mul_factor; + } else { + for (const auto& pool : pools) { + // Every pool must claim at least one QP; a non-positive num_qp + // would produce an empty/negative [begin, begin+num_qp) span and + // break the router. Reject the whole layout so the caller falls + // back to the default single-pool behavior. + if (pool.num_qp <= 0) { + layout.segments.clear(); + layout.total_qp = 0; + layout.valid = false; + return layout; + } + QpPoolSegment seg = pool; + seg.begin = layout.total_qp; + layout.segments.push_back(seg); + layout.total_qp += pool.num_qp; + } + } + layout.valid = layout.total_qp > 0; + return layout; +} + +// Pure QP router used by RdmaEndPoint::submitSlices (RFC #2568 step 3). Given +// the resolved segments, the pool a transfer asked for, and a caller-provided +// candidate index (the worker lane), return the QP index the transfer should +// use. When the pool is named and found, the candidate is folded into that +// pool's [begin, begin+num_qp) span so the transfer only ever touches its +// pool's QPs. When no pool is named, or the name is unknown, or no pools are +// configured, the candidate passes through unchanged (default spray behavior). +// total_qp must be > 0; the result is always in [0, total_qp). +inline int selectQpInPool(const std::vector& segments, + const std::string& pool_name, int candidate, + int total_qp) { + if (candidate < 0) candidate = 0; + if (!pool_name.empty()) { + for (const auto& seg : segments) { + if (seg.name == pool_name && seg.num_qp > 0) + return seg.begin + (candidate % seg.num_qp); + } + } + return candidate % total_qp; +} + struct WorkerParams { int num_workers = 6; // Derived from RdmaParams::num_lanes. int max_retry_count = 8; @@ -72,7 +161,8 @@ struct RdmaParams { DeviceParams device; EndPointParams endpoint; WorkerParams workers; - bool verbose; + bool verbose = false; + bool log_slice_affinity = false; }; } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h new file mode 100644 index 0000000000..a5d3a92549 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/promotion_policy.h @@ -0,0 +1,92 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Starvation-prevention promotion policy for the RDMA worker priority queues, +// factored out of Workers::promoteTimedOutRequests so the decision logic can be +// unit-tested without the full RDMA stack (issue #2528). +// +// A worker keeps three FIFO priority queues (HIGH > MEDIUM > LOW). To keep +// lower-priority requests from starving, every ~1ms a promotion pass looks at +// timed-out entries and moves them up one level. This header isolates *which* +// entries a pass decides to promote, given each entry's enqueue timestamp. + +#pragma once + +#include +#include +#include + +namespace mooncake { +namespace tent { + +// Result of a promotion decision for a single queue: the indices (into the +// drained queue order) that should move up one level. Empty == promote none. +struct PromotionDecision { + std::vector promote_indices; + bool promoted_any() const { return !promote_indices.empty(); } +}; + +// Historical policy (as implemented in Workers::promoteTimedOutRequests today): +// the pass drains the whole queue, inspects ONLY the head entry, and if the +// head has timed out it promotes EVERY entry in the queue. This is the behavior +// #2528 flags as unintended: +// * decision is head-only but applied to the whole queue -> freshly enqueued +// entries that are not starving get promoted alongside the starving head; +// * if the head has NOT timed out, later timed-out entries are not promoted. +// +// `enqueue_ts` is per drained entry in queue order (index 0 == head). A zero +// timestamp means "no timestamp" and never counts as timed out (matches the +// `enqueue_ts > 0` guard in the current code). +inline PromotionDecision DecidePromotionHeadOnly( + const std::vector& enqueue_ts, uint64_t current_ts, + uint64_t promotion_timeout_ns) { + PromotionDecision d; + if (enqueue_ts.empty()) return d; + const uint64_t head = enqueue_ts.front(); + // Guard current_ts >= head before subtracting: both are unsigned, so a + // non-monotonic clock / race where current_ts < head would otherwise + // underflow and spuriously mark the entry timed out. + const bool head_timed_out = head > 0 && current_ts >= head && + (current_ts - head) >= promotion_timeout_ns; + if (head_timed_out) { + d.promote_indices.reserve(enqueue_ts.size()); + for (size_t i = 0; i < enqueue_ts.size(); ++i) { + d.promote_indices.push_back(i); // promote ALL + } + } + return d; +} + +// Per-entry policy: promote exactly the entries that have themselves timed out, +// leaving freshly enqueued entries in place. This is the behavior #2528 +// proposes; kept here next to the historical policy so a test can contrast the +// two and a follow-up fix can switch Workers over to it. +inline PromotionDecision DecidePromotionPerEntry( + const std::vector& enqueue_ts, uint64_t current_ts, + uint64_t promotion_timeout_ns) { + PromotionDecision d; + for (size_t i = 0; i < enqueue_ts.size(); ++i) { + const uint64_t ts = enqueue_ts[i]; + // See DecidePromotionHeadOnly: guard against unsigned underflow when + // current_ts < ts. + if (ts > 0 && current_ts >= ts && + (current_ts - ts) >= promotion_timeout_ns) { + d.promote_indices.push_back(i); + } + } + return d; +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h index b3a857ec44..9273e507d7 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/quota.h @@ -131,6 +131,8 @@ class DeviceSelector { Status release(int dev_id, uint64_t length, double latency); + Status getNicLoadStats(std::vector &stats) const; + void updateTrafficStats(int dev_id, uint64_t length) { auto it = devices_.find(dev_id); if (it != devices_.end()) { @@ -149,6 +151,8 @@ class DeviceSelector { void printTrafficStats(); + double getAggregateEwmaBandwidth() const; + void fillDevicePriorities(); int getDevicePriority(int dev_id) const; @@ -221,4 +225,4 @@ class DeviceSelector { } // namespace tent } // namespace mooncake -#endif // TENT_SELECTOR_H \ No newline at end of file +#endif // TENT_SELECTOR_H diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h index 2a7395799c..6ff05db7b8 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/rdma_transport.h @@ -56,6 +56,7 @@ struct RdmaSubBatch : public Transport::SubBatch { class RdmaTransport : public Transport { friend class Workers; friend class RdmaEndPoint; + friend class RdmaTransportTestPeer; public: RdmaTransport(); @@ -79,6 +80,10 @@ class RdmaTransport : public Transport { virtual Status getTransferStatus(SubBatchRef batch, int task_id, TransferStatus& status); + bool supportsCancellation() const override { return true; } + + Status cancelTransferTask(SubBatchRef batch, int task_id) override; + virtual Status addMemoryBuffer(BufferDesc& desc, const MemoryOptions& options); @@ -91,6 +96,9 @@ class RdmaTransport : public Transport { virtual const char* getName() const { return "rdma"; } + double getEstimatedBandwidth() const override; + Status getNicLoadStats(std::vector& stats) const override; + virtual bool supportNotification() const override { return true; } virtual Status sendNotification(SegmentID target_id, @@ -116,6 +124,11 @@ class RdmaTransport : public Transport { std::shared_ptr config() const { return conf_; } + private: + // Builds context_set_ with one slot per NicID; returns how many RNICs + // came up. Remaining slots hold inert contexts. + size_t initializeContexts(); + private: bool installed_; std::shared_ptr conf_; @@ -140,10 +153,12 @@ class RdmaTransport : public Transport { // Map QP number to Endpoint for notification processing RWSpinlock notify_endpoint_map_lock_; - std::unordered_map notify_qp_to_endpoint_; + std::unordered_map> + notify_qp_to_endpoint_; // Register/unregister notification QP (called by Endpoint) - void registerNotifyQp(uint32_t qp_num, RdmaEndPoint* endpoint); + void registerNotifyQp(uint32_t qp_num, + const std::shared_ptr& endpoint); void unregisterNotifyQp(uint32_t qp_num); std::shared_ptr getEndpoint(SegmentID target_id, int device_id); @@ -159,4 +174,4 @@ class RdmaTransport : public Transport { } // namespace tent } // namespace mooncake -#endif // TENT_RDMA_TRANSPORT_H \ No newline at end of file +#endif // TENT_RDMA_TRANSPORT_H diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h index f681c8d274..17fa125fa9 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/slice.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -50,12 +51,20 @@ using RdmaTaskStorage = Slab; struct RdmaTask { int num_slices; Request request; + // Named QP pool this task's slices should use (RFC #2568 step 3). Empty = + // no pool selected: slices spray across all data QPs as before. Resolved + // from SelectionResult.qp_pool at task creation. + std::string qp_pool; volatile TransferStatusEnum status_word; volatile size_t transferred_bytes; - volatile int success_slices; - volatile int resolved_slices; + std::atomic success_slices{0}; + std::atomic resolved_slices{0}; volatile TransferStatusEnum first_error = PENDING; + // Set by the control thread. Workers observe this flag before posting or + // retrying a slice. Already-posted WRs are allowed to drain normally. + std::atomic cancel_requested{false}; + // Reference counting for UAF protection std::atomic ref_count{0}; @@ -81,12 +90,28 @@ struct RdmaSlice { uint32_t target_rkey = 0; int source_dev_id = -1; int target_dev_id = -1; + // GPUDirect reachability learning (see GdrReachability). Resolved once per + // (re)submit in Workers::generatePostPath. GPU ordinals are -1 for host + // memory; the name pointers alias stable Topology::NicEntry / segment + // storage and stay valid for the slice's lifetime. + int source_gpu_ordinal = -1; + int target_gpu_ordinal = -1; + const char* source_nic_name = nullptr; + const char* target_nic_name = nullptr; + const std::string* target_machine_id = nullptr; std::weak_ptr ep_weak_ptr; TransferStatusEnum word = TransferStatusEnum::INITIAL; int qp_index = 0; int retry_count = 0; + // Flat (source,target) combination index last tried by + // selectFallbackDevice; the next fallback resumes just past it so retries + // rotate through all combinations with wraparound instead of hammering one. + int last_fallback_idx = -1; bool failed = false; + // True while DeviceSelector accounts this slice against source_dev_id. + // The worker clears it exactly once on completion, failure, or cancel. + bool quota_charged = false; uint64_t enqueue_ts = 0; uint64_t submit_ts = 0; // Non-owning pointer to the per-worker RailMonitor for this slice's @@ -106,15 +131,18 @@ static inline void updateSliceStatus(RdmaSlice* slice, if (!__sync_bool_compare_and_swap(&slice->word, PENDING, status)) return; if (status == COMPLETED) { __sync_fetch_and_add(&task->transferred_bytes, slice->length); - __sync_fetch_and_add(&task->success_slices, 1); + task->success_slices.fetch_add(1, std::memory_order_acq_rel); } else { __sync_bool_compare_and_swap(&task->first_error, PENDING, status); } - int resolved = __sync_add_and_fetch(&task->resolved_slices, 1); + int resolved = + task->resolved_slices.fetch_add(1, std::memory_order_acq_rel) + 1; if (resolved >= task->num_slices) { - TransferStatusEnum final_st = (task->success_slices == task->num_slices) - ? COMPLETED - : task->first_error; + TransferStatusEnum final_st = + (task->success_slices.load(std::memory_order_acquire) == + task->num_slices) + ? COMPLETED + : task->first_error; if (final_st == PENDING) final_st = FAILED; __sync_bool_compare_and_swap(&task->status_word, PENDING, final_st); } @@ -124,4 +152,4 @@ static inline void updateSliceStatus(RdmaSlice* slice, } // namespace tent } // namespace mooncake -#endif // TENT_SLICE_H \ No newline at end of file +#endif // TENT_SLICE_H diff --git a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h index 83b3da8ed8..e9578d0843 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/rdma/workers.h @@ -38,12 +38,14 @@ class RdmaTransport; class DeviceSelector; class Workers { + friend class RdmaTransportTestPeer; + public: static constexpr size_t kCapacity = 1024 * 8; using BoundedSliceQueue = BoundedMPSCQueue; public: - Workers(RdmaTransport *transport); + Workers(RdmaTransport* transport); ~Workers(); @@ -51,16 +53,17 @@ class Workers { Status stop(); - Status submit(RdmaSlice *slice); + Status submit(RdmaSlice* slice); - Status submit(RdmaSliceList &slice_list, int worker_id = -1); + Status submit(RdmaSliceList& slice_list, int worker_id = -1); - Status cancel(RdmaSliceList &slice_list); + Status cancel(RdmaTask* task); - DeviceSelector *getDeviceSelector() const { return device_selector_.get(); } + DeviceSelector* getDeviceSelector() const { return device_selector_.get(); } private: using Task = std::function; + struct WorkerContext; void workerThread(int thread_id); @@ -68,37 +71,56 @@ class Workers { void asyncPollCq(); + bool cancelUnpostedSlice(WorkerContext& worker, RdmaSlice* slice); + + void releaseSliceQuota(RdmaSlice* slice, double latency = 0.0); + void monitorThread(); - int handleContextEvents(std::shared_ptr &context); + // 1 Hz heartbeat from monitorThread(): drains every context's retiring + // endpoints so reclaim is not gated on new insertions, which stall under + // failure load. + void reclaimEndpoints(); - Status generatePostPath(RdmaSlice *slice); + int handleContextEvents(std::shared_ptr& context); + + Status generatePostPath(RdmaSlice* slice); private: struct RouteHint { - SegmentDesc *segment; - BufferDesc *buffer; - const Topology::MemEntry *topo_entry; - const Topology *topo; + // Owning reference to the segment snapshot; keeps all raw pointers + // below valid for the lifetime of this hint. + SegmentDescRef pin; + SegmentDesc* segment; + BufferDesc* buffer; + const Topology::MemEntry* topo_entry; + const Topology* topo; + std::string location; }; - Status getRouteHint(RouteHint &hint, SegmentID segment_id, uint64_t addr, + Status getRouteHint(RouteHint& hint, SegmentID segment_id, uint64_t addr, uint64_t length); - Status selectOptimalDevice(RouteHint &source, RouteHint &target, - RdmaSlice *slice); + Status selectOptimalDevice(RouteHint& source, RouteHint& target, + RdmaSlice* slice); + + Status selectFallbackDevice(RouteHint& source, RouteHint& target, + RdmaSlice* slice); - Status selectFallbackDevice(RouteHint &source, RouteHint &target, - RdmaSlice *slice); + int getDeviceByFlatIndex(const RouteHint& hint, size_t flat_idx); - int getDeviceByFlatIndex(const RouteHint &hint, size_t flat_idx); + // True if the (sdev -> tdev) NIC pair is known-unable to GPUDirect-DMA to + // the source/target GPU (learned from prior completion errors). Used to + // steer selection away from dead rails before posting. + bool gdrPairExcluded(const RouteHint& source, const RouteHint& target, + int sdev, int tdev, int src_gpu, int dst_gpu); - int getDeviceRank(const RouteHint &hint, int device_id); + int getDeviceRank(const RouteHint& hint, int device_id); void showLatencyInfo(); private: - RdmaTransport *transport_; + RdmaTransport* transport_; size_t num_workers_; std::thread monitor_; @@ -109,7 +131,7 @@ class Workers { SegmentID remote_segment_id; int remote_device_id; - bool operator==(const PostPath &rhs) const { + bool operator==(const PostPath& rhs) const { return local_device_id == rhs.local_device_id && remote_segment_id == rhs.remote_segment_id && remote_device_id == rhs.remote_device_id; @@ -117,7 +139,7 @@ class Workers { }; struct PostPathHash { - size_t operator()(const PostPath &postPath) const { + size_t operator()(const PostPath& postPath) const { size_t h1 = std::hash{}(postPath.local_device_id); size_t h2 = std::hash{}(postPath.remote_segment_id); size_t h3 = std::hash{}(postPath.remote_device_id); @@ -127,10 +149,10 @@ class Workers { std::shared_ptr getEndpoint(PostPath path); - void disableEndpoint(RdmaSlice *slice); + void disableEndpoint(RdmaSlice* slice); using GroupedRequests = - std::unordered_map, PostPathHash>; + std::unordered_map, PostPathHash>; struct PerfMetric { void add(double val) { samples.push_back(val); } @@ -187,7 +209,7 @@ class Workers { std::thread thread; BoundedSliceQueue queues[kNumPriorityLevels]; // Priority queues GroupedRequests requests; - std::unordered_set inflight_slice_set; + std::unordered_set inflight_slice_set; std::atomic inflight_slices = 0; std::mutex mutex; @@ -206,14 +228,25 @@ class Workers { }; // Promote timed-out low priority requests to higher priority queues - void promoteTimedOutRequests(WorkerContext &worker); + void promoteTimedOutRequests(WorkerContext& worker); - WorkerContext *worker_context_; + WorkerContext* worker_context_; uint64_t slice_timeout_ns_; uint64_t priority_promotion_timeout_ns_; // Timeout for priority promotion + // Opt-in (issue #2528): when true, a promotion pass promotes exactly the + // entries that have themselves timed out, instead of promoting the whole + // queue whenever only the head has timed out. Default false keeps the + // historical "flush the tier" behavior. + bool priority_promotion_per_entry_ = false; std::unique_ptr device_selector_; + // File contents loaded once from workers.rail_topo_path and shared by all + // per-worker/per-peer RailMonitor instances. + std::string rail_topo_json_; bool always_tier1_ = false; + // Opt-in deadline-aware bandwidth arbitration within a priority tier + // (RFC #2792). Default false = original FIFO order (equal bandwidth split). + bool deadline_bw_arbitration_ = false; }; } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/include/tent/transport/shm/shm_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/shm/shm_transport.h index 85fac8282a..b749d3843e 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/shm/shm_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/shm/shm_transport.h @@ -75,6 +75,8 @@ class ShmTransport : public Transport { virtual Status freeLocalMemory(void *addr, size_t size); private: + friend class ShmTransportTestPeer; + void startTransfer(ShmTask *task, ShmSubBatch *batch); void *createSharedMemory(const std::string &path, size_t size); @@ -89,14 +91,15 @@ class ShmTransport : public Transport { std::shared_ptr metadata_; struct OpenedShmEntry { - int shm_fd; void *shm_addr; uint64_t length; }; - using HashMap = - std::unordered_map>; + using RelocateMap = std::unordered_map; + using HashMap = std::unordered_map; + + static bool tryResolve(const RelocateMap &relocate_map, uint64_t &dest_addr, + uint64_t length); RWSpinlock relocate_lock_; HashMap relocate_map_; @@ -108,8 +111,11 @@ class ShmTransport : public Transport { std::unordered_map shm_path_map_; std::string cxl_mount_path_; + + static constexpr int kShmCreateMaxRetries = 8; }; + } // namespace tent } // namespace mooncake -#endif // SHM_TRANSPORT_H_ \ No newline at end of file +#endif // SHM_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h index acdf117dc6..3d5f2ebedd 100644 --- a/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h +++ b/mooncake-transfer-engine/tent/include/tent/transport/tcp/tcp_transport.h @@ -38,6 +38,8 @@ struct TcpParams { struct TcpTask { Request request; + BatchID progress_batch_id{0}; + std::function notify_progress; std::atomic status_word{TransferStatusEnum::PENDING}; std::atomic transferred_bytes{0}; uint64_t target_addr = 0; @@ -45,6 +47,8 @@ struct TcpTask { TcpTask() = default; TcpTask(TcpTask &&other) noexcept : request(std::move(other.request)), + progress_batch_id(other.progress_batch_id), + notify_progress(std::move(other.notify_progress)), status_word(other.status_word.load(std::memory_order_relaxed)), transferred_bytes( other.transferred_bytes.load(std::memory_order_relaxed)), diff --git a/mooncake-transfer-engine/tent/include/tent/transport/tpu/tpu_transport.h b/mooncake-transfer-engine/tent/include/tent/transport/tpu/tpu_transport.h new file mode 100644 index 0000000000..e4d0890d4c --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/tpu/tpu_transport.h @@ -0,0 +1,93 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_TPU_TRANSPORT_H_ +#define TENT_TPU_TRANSPORT_H_ + +#include +#include + +#include "tent/runtime/control_plane.h" +#include "tent/runtime/transport.h" + +namespace mooncake { +namespace tent { + +struct TpuTask { + Request request; + volatile TransferStatusEnum status_word; + volatile size_t transferred_bytes; +}; + +struct TpuSubBatch : public Transport::SubBatch { + std::vector task_list; + size_t max_size; + virtual size_t size() const { return task_list.size(); } +}; + +// TpuTransport is the local staging executor for TPU: it performs the +// HBM<->host-DRAM hop of a staged transfer by delegating to +// Platform::copy() (which routes to the PJRT adapter for TPU memory). It never +// touches the network — the host<->host hop is carried by RDMA/TCP, and +// ProxyManager chains the two stages (see findStagingPolicy). +// +// Because TPU HBM is not NIC-addressable, this transport advertises only the +// device<->host capabilities (gpu_to_dram / dram_to_gpu) and leaves gpu_to_gpu +// false so the engine always stages cross-node traffic through host DRAM. It is +// therefore only ever selected for LOCAL_SEGMENT_ID copies. +class TpuTransport : public Transport { + public: + TpuTransport(); + + ~TpuTransport(); + + virtual Status install(std::string &local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf = nullptr); + + virtual Status uninstall(); + + virtual Status allocateSubBatch(SubBatchRef &batch, size_t max_size); + + virtual Status freeSubBatch(SubBatchRef &batch); + + virtual Status submitTransferTasks( + SubBatchRef batch, const std::vector &request_list); + + virtual Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus &status); + + virtual Status addMemoryBuffer(BufferDesc &desc, + const MemoryOptions &options); + + virtual Status removeMemoryBuffer(BufferDesc &desc); + + virtual const char *getName() const { return "tpu"; } + + private: + void startTransfer(TpuTask *task, TpuSubBatch *batch); + + private: + bool installed_; + std::string local_segment_name_; + std::shared_ptr local_topology_; + std::shared_ptr metadata_; + std::shared_ptr conf_; +}; + +} // namespace tent +} // namespace mooncake + +#endif // TENT_TPU_TRANSPORT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/buffers.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/buffers.h new file mode 100644 index 0000000000..3fcd17799b --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/buffers.h @@ -0,0 +1,158 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_BUFFERS_H_ +#define TENT_TRANSPORT_UB_BUFFERS_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/runtime/segment.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +struct UbBufferSegmentMetadata { + Topology::NicID topology_id{-1}; + std::string device_name; + std::string eid; + int eid_index{-1}; + SegmentDescriptor descriptor; +}; + +struct UbBufferMetadata { + static constexpr uint32_t kSchemaVersion = 1; + + uint32_t schema_version{kSchemaVersion}; + uint64_t generation{0}; + uint64_t base{0}; + uint64_t length{0}; + std::string location; + Permission permission{kLocalReadWrite}; + std::vector segments; +}; + +Status encodeBufferMetadata(const UbBufferMetadata& metadata, + std::string& encoded); +Status decodeBufferMetadata(std::string_view encoded, + UbBufferMetadata& metadata); + +struct LocalSegmentRef { + UbContextPtr context; + LocalSegmentPtr segment; + uint64_t generation{0}; + uint64_t buffer_base{0}; + uint64_t buffer_length{0}; +}; + +struct ImportedSegmentRef { + UbContextPtr context; + RemoteSegmentPtr segment; + uint64_t generation{0}; + uint64_t buffer_base{0}; + uint64_t buffer_length{0}; + Topology::NicID remote_topology_id{-1}; +}; + +// Owns registrations and imports independently from Classic TE. Every +// adapter handle is shared with completion tokens, so removal invalidates the +// metadata immediately while native destruction is deferred until in-flight +// work releases its last reference. +class UbBufferManager final { + public: + UbBufferManager(std::shared_ptr adapter, + std::vector contexts); + ~UbBufferManager(); + + UbBufferManager(const UbBufferManager&) = delete; + UbBufferManager& operator=(const UbBufferManager&) = delete; + + Status addBuffer(BufferDesc& desc, const MemoryOptions& options); + Status addBuffers(std::vector& descs, + const MemoryOptions& options); + Status removeBuffer(BufferDesc& desc); + Status clear(); + + Status findLocal(uint64_t address, size_t length, + Topology::NicID local_topology_id, + LocalSegmentRef& result) const; + + Status importRemote(SegmentID remote_segment_id, + Topology::NicID local_topology_id, + Topology::NicID remote_topology_id, + const BufferDesc& remote_buffer, Request::OpCode opcode, + uint64_t address, size_t length, + ImportedSegmentRef& result); + + [[nodiscard]] size_t localBufferCount() const; + [[nodiscard]] size_t importedSegmentCount() const; + + private: + struct AddressRange { + uint64_t base{0}; + uint64_t length{0}; + + bool operator<(const AddressRange& rhs) const { + return base < rhs.base || (base == rhs.base && length < rhs.length); + } + [[nodiscard]] bool contains(uint64_t address, size_t size) const; + }; + + struct LocalRecord { + MemoryOptions options; + uint64_t generation{0}; + std::unordered_map segments; + }; + + struct ImportKey { + Topology::NicID local_topology_id{-1}; + SegmentID remote_segment_id{LOCAL_SEGMENT_ID}; + Topology::NicID remote_topology_id{-1}; + uint64_t buffer_base{0}; + uint64_t generation{0}; + + bool operator==(const ImportKey&) const = default; + }; + + struct ImportKeyHash { + size_t operator()(const ImportKey& key) const noexcept; + }; + + Status addBufferInternal(BufferDesc& desc, const MemoryOptions& options); + Status unregisterRecord(LocalRecord& record); + void retainPendingRecord(LocalRecord& record); + static uint32_t segmentAccess(Permission permission); + static bool permissionAllows(Permission permission, Request::OpCode opcode); + UbContextPtr findContext(Topology::NicID topology_id) const; + + std::shared_ptr adapter_; + std::vector contexts_; + std::unordered_map context_by_topology_id_; + + mutable std::shared_mutex local_mutex_; + std::map local_buffers_; + // Registrations created by an add/transaction rollback stay owned here + // when the provider refuses the first unregister attempt. clear() retries + // them before the manager can be destroyed. + std::vector pending_local_segments_; + + mutable std::shared_mutex import_mutex_; + std::unordered_map imports_; + // Partial imports returned alongside a provider error have no usable cache + // key, but still need stable ownership until unimport succeeds. + std::vector pending_remote_segments_; + std::atomic next_generation_{1}; +}; + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_BUFFERS_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/context.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/context.h new file mode 100644 index 0000000000..03999a6843 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/context.h @@ -0,0 +1,112 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_CONTEXT_H_ +#define TENT_TRANSPORT_UB_CONTEXT_H_ + +#include +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/runtime/topology.h" +#include "tent/transport/ub/jfc.h" +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +class UbContext final : public std::enable_shared_from_this { + public: + enum class State : uint8_t { + kUninitialized, + kActive, + kFailed, + kDraining, + kClosed, + }; + + UbContext(Topology::NicID topology_id, DeviceInfo device, + std::shared_ptr adapter); + ~UbContext(); + + UbContext(const UbContext&) = delete; + UbContext& operator=(const UbContext&) = delete; + + Status initialize(uint32_t jfc_count, const JfcOptions& options); + Status shutdown(); + // Returns true only for the Active -> Failed transition. Every subsequent + // poll error advances the failure epoch so stale successes from another + // JFC can never reactivate the device. + [[nodiscard]] bool markUnavailable() noexcept; + // The failure owner calls this only after every endpoint backed by this + // device has been unpublished and has reached Destroyed. Pending native + // cleanup remains quarantined and keeps this recovery barrier incomplete. + void completeFailureCleanup() noexcept; + // Pollers continue probing every JFC. A transiently failed context becomes + // active only after all JFCs have succeeded in the current failure epoch, + // the error-free cooldown has elapsed, old WRs have drained, and endpoint + // unpublication has completed. + [[nodiscard]] bool recordPollSuccess(size_t jfc_index, + uint64_t cooldown_ns) noexcept; + + [[nodiscard]] Topology::NicID topologyId() const noexcept { + return topology_id_; + } + [[nodiscard]] const DeviceInfo& deviceInfo() const noexcept { + return device_; + } + [[nodiscard]] const ContextPtr& handle() const noexcept { return handle_; } + [[nodiscard]] State state() const noexcept { + return state_.load(std::memory_order_acquire); + } + [[nodiscard]] bool active() const noexcept { + return state() == State::kActive; + } + [[nodiscard]] const std::vector>& jfcs() const { + return jfcs_; + } + [[nodiscard]] std::shared_ptr jfc(size_t index) const; + + void addInflight(uint64_t bytes) noexcept; + void removeInflight(uint64_t bytes) noexcept; + [[nodiscard]] uint64_t inflightBytes() const noexcept { + return inflight_bytes_.load(std::memory_order_relaxed); + } + [[nodiscard]] uint64_t outstandingWrs() const noexcept { + return outstanding_wrs_.load(std::memory_order_relaxed); + } + [[nodiscard]] uint64_t recoveryCount() const noexcept { + return recovery_count_.load(std::memory_order_relaxed); + } + [[nodiscard]] uint64_t failureStartedNs() const noexcept { + return failure_started_ns_.load(std::memory_order_acquire); + } + + private: + Status shutdownLocked(); + + const Topology::NicID topology_id_; + const DeviceInfo device_; + std::shared_ptr adapter_; + ContextPtr handle_; + std::vector> jfcs_; + mutable std::mutex lifecycle_mutex_; + std::atomic state_{State::kUninitialized}; + std::atomic inflight_bytes_{0}; + std::atomic outstanding_wrs_{0}; + std::atomic failure_started_ns_{0}; + std::atomic last_failure_ns_{0}; + std::atomic recovery_count_{0}; + uint64_t failure_epoch_{0}; + std::vector jfc_success_epochs_; + bool failure_cleanup_complete_{false}; +}; + +using UbContextPtr = std::shared_ptr; + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_CONTEXT_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/jfc.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/jfc.h new file mode 100644 index 0000000000..ec7dd14a15 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/jfc.h @@ -0,0 +1,58 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_JFC_H_ +#define TENT_TRANSPORT_UB_JFC_H_ + +#include +#include +#include +#include +#include + +#include "tent/common/status.h" +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +// TENT-facing completion queue. Native JFC handles stay behind the adapter; +// this object only adds stable ownership and transport telemetry. +class UbJfc final { + public: + UbJfc(size_t index, std::shared_ptr adapter, JfcPtr handle) + : index_(index), + adapter_(std::move(adapter)), + handle_(std::move(handle)) {} + + UbJfc(const UbJfc&) = delete; + UbJfc& operator=(const UbJfc&) = delete; + + ~UbJfc(); + + [[nodiscard]] size_t index() const noexcept { return index_; } + [[nodiscard]] const JfcPtr& handle() const noexcept { return handle_; } + [[nodiscard]] bool valid() const noexcept { + return handle_ && handle_->valid(); + } + + Status poll(size_t max_completions, std::vector& completions); + Status close(); + + [[nodiscard]] uint64_t completionCount() const noexcept { + return completion_count_.load(std::memory_order_relaxed); + } + [[nodiscard]] uint64_t pollErrorCount() const noexcept { + return poll_error_count_.load(std::memory_order_relaxed); + } + + private: + const size_t index_; + std::shared_ptr adapter_; + JfcPtr handle_; + std::atomic completion_count_{0}; + std::atomic poll_error_count_{0}; +}; + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_JFC_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/params.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/params.h new file mode 100644 index 0000000000..175dc8739c --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/params.h @@ -0,0 +1,219 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_TRANSPORT_UB_PARAMS_H_ +#define TENT_TRANSPORT_UB_PARAMS_H_ + +#include +#include +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/common/status.h" + +namespace mooncake { +namespace tent { +namespace ub { + +struct UbParams { + // UB is opt-in until a real URMA provider and NIC topology are present. + bool enable = false; + std::vector device_filter; + + uint32_t worker_count = 6; + uint32_t poller_count = 2; + uint32_t jfc_per_context = 6; + uint32_t jetty_per_endpoint = 6; + uint32_t max_endpoints = 65536; + size_t slice_size = 64 * 1024; + uint32_t max_retries = 8; + uint32_t slice_timeout_ms = 5000; + uint32_t endpoint_cooldown_ms = 1000; + + bool enable_bandwidth_estimation = true; + bool enable_notifications = false; + + // Reads every UB setting from transports/ub/*. Numeric settings are + // deliberately read as signed JSON integers first so a negative value is + // rejected rather than silently converted to an unsigned default. + static Status FromConfig(const Config& config, UbParams& output) { + UbParams parsed; + + CHECK_STATUS(readBool(config, "transports/ub/enable", parsed.enable, + parsed.enable)); + CHECK_STATUS(readDeviceFilter(config, parsed.device_filter)); + CHECK_STATUS(readPositive(config, "transports/ub/worker_count", + parsed.worker_count, parsed.worker_count)); + CHECK_STATUS(readPositive(config, "transports/ub/poller_count", + parsed.poller_count, parsed.poller_count)); + CHECK_STATUS(readPositive(config, "transports/ub/jfc_per_context", + parsed.jfc_per_context, + parsed.jfc_per_context)); + CHECK_STATUS(readPositive(config, "transports/ub/jetty_per_endpoint", + parsed.jetty_per_endpoint, + parsed.jetty_per_endpoint)); + CHECK_STATUS(readPositive(config, "transports/ub/max_endpoints", + parsed.max_endpoints, parsed.max_endpoints)); + + uint64_t slice_size = parsed.slice_size; + CHECK_STATUS(readPositive(config, "transports/ub/slice_size", + slice_size, slice_size)); + if (slice_size > std::numeric_limits::max()) { + return Status::InvalidArgument( + "transports/ub/slice_size exceeds the URMA SGE length limit"); + } + parsed.slice_size = static_cast(slice_size); + + CHECK_STATUS(readNonNegative(config, "transports/ub/max_retries", + parsed.max_retries, parsed.max_retries)); + CHECK_STATUS(readPositive(config, "transports/ub/slice_timeout_ms", + parsed.slice_timeout_ms, + parsed.slice_timeout_ms)); + CHECK_STATUS(readPositive(config, "transports/ub/endpoint_cooldown_ms", + parsed.endpoint_cooldown_ms, + parsed.endpoint_cooldown_ms)); + CHECK_STATUS(readBool(config, + "transports/ub/enable_bandwidth_estimation", + parsed.enable_bandwidth_estimation, + parsed.enable_bandwidth_estimation)); + CHECK_STATUS(readBool(config, "transports/ub/enable_notifications", + parsed.enable_notifications, + parsed.enable_notifications)); + + output = std::move(parsed); + return Status::OK(); + } + + private: + template + static Status readPositive(const Config& config, const char* key, + T default_value, T& output) { + const json value = config.get(key, json()); + if (value.is_null()) { + output = default_value; + return Status::OK(); + } + if (!value.is_number_integer()) { + return Status::InvalidArgument(std::string(key) + + " must be a positive integer"); + } + + int64_t signed_value = 0; + try { + signed_value = value.get(); + } catch (...) { + return Status::InvalidArgument(std::string(key) + + " is outside the supported range"); + } + if (signed_value <= 0 || + static_cast(signed_value) > + static_cast(std::numeric_limits::max())) { + return Status::InvalidArgument(std::string(key) + + " must be a positive integer in " + "the supported range"); + } + output = static_cast(signed_value); + return Status::OK(); + } + + template + static Status readNonNegative(const Config& config, const char* key, + T default_value, T& output) { + const json value = config.get(key, json()); + if (value.is_null()) { + output = default_value; + return Status::OK(); + } + if (!value.is_number_integer()) { + return Status::InvalidArgument(std::string(key) + + " must be a non-negative integer"); + } + int64_t signed_value = 0; + try { + signed_value = value.get(); + } catch (...) { + return Status::InvalidArgument(std::string(key) + + " is outside the supported range"); + } + if (signed_value < 0 || + static_cast(signed_value) > + static_cast(std::numeric_limits::max())) { + return Status::InvalidArgument(std::string(key) + + " must be a non-negative integer " + "in the supported range"); + } + output = static_cast(signed_value); + return Status::OK(); + } + + static Status readBool(const Config& config, const char* key, + bool default_value, bool& output) { + const json value = config.get(key, json()); + if (value.is_null()) { + output = default_value; + return Status::OK(); + } + if (!value.is_boolean()) { + return Status::InvalidArgument(std::string(key) + + " must be a boolean"); + } + output = value.get(); + return Status::OK(); + } + + static Status readDeviceFilter(const Config& config, + std::vector& output) { + static constexpr const char* kKey = "transports/ub/device_filter"; + const json value = config.get(kKey, json()); + if (value.is_null()) { + output.clear(); + return Status::OK(); + } + + std::vector parsed; + if (value.is_string()) { + parsed.push_back(value.get()); + } else if (value.is_array()) { + for (const auto& entry : value) { + if (!entry.is_string()) { + return Status::InvalidArgument( + "transports/ub/device_filter entries must be strings"); + } + parsed.push_back(entry.get()); + } + } else { + return Status::InvalidArgument( + "transports/ub/device_filter must be a string or string " + "array"); + } + + for (const auto& entry : parsed) { + if (entry.empty()) { + return Status::InvalidArgument( + "transports/ub/device_filter entries must not be empty"); + } + } + output = std::move(parsed); + return Status::OK(); + } +}; + +} // namespace ub +} // namespace tent +} // namespace mooncake + +#endif // TENT_TRANSPORT_UB_PARAMS_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/topology_attrs.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/topology_attrs.h new file mode 100644 index 0000000000..9b4a25c2f1 --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/topology_attrs.h @@ -0,0 +1,38 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#ifndef TENT_TRANSPORT_UB_TOPOLOGY_ATTRS_H_ +#define TENT_TRANSPORT_UB_TOPOLOGY_ATTRS_H_ + +#include +#include +#include + +#include "tent/transport/ub/urma_adapter.h" + +namespace mooncake::tent::ub { + +inline constexpr std::string_view kTopologyNativeNameAttr = "ub.native_name"; +inline constexpr std::string_view kTopologyDeviceIndexAttr = "ub.device_index"; +inline constexpr std::string_view kTopologyEidIndexAttr = "ub.eid_index"; +inline constexpr std::string_view kTopologyEidAttr = "ub.eid"; +inline constexpr std::string_view kTopologyDiscoveryActiveAttr = + "ub.discovery_active"; + +inline void encodeTopologyDeviceAttributes( + const DeviceInfo& device, int device_index, + std::unordered_map& attributes) { + attributes[std::string(kTopologyNativeNameAttr)] = + device.native_device_name; + attributes[std::string(kTopologyDeviceIndexAttr)] = + std::to_string(device_index); + attributes[std::string(kTopologyEidIndexAttr)] = + std::to_string(device.eid_index); + attributes[std::string(kTopologyEidAttr)] = device.eid; + attributes[std::string(kTopologyDiscoveryActiveAttr)] = + device.active ? "true" : "false"; +} + +} // namespace mooncake::tent::ub + +#endif // TENT_TRANSPORT_UB_TOPOLOGY_ATTRS_H_ diff --git a/mooncake-transfer-engine/tent/include/tent/transport/ub/urma_adapter.h b/mooncake-transfer-engine/tent/include/tent/transport/ub/urma_adapter.h new file mode 100644 index 0000000000..3fb5166f2f --- /dev/null +++ b/mooncake-transfer-engine/tent/include/tent/transport/ub/urma_adapter.h @@ -0,0 +1,284 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef TENT_TRANSPORT_UB_URMA_ADAPTER_H_ +#define TENT_TRANSPORT_UB_URMA_ADAPTER_H_ + +#include +#include +#include +#include +#include + +#include "tent/common/status.h" + +namespace mooncake { +namespace tent { +namespace ub { + +// Adapter-neutral subset of the capabilities needed by the TENT UB data +// plane. Values of zero mean that the provider did not report a limit. +struct DeviceCapabilities { + uint32_t max_jfc = 0; + uint32_t max_jfc_depth = 0; + uint32_t max_jfr_depth = 0; + uint32_t max_jetty = 0; + uint32_t max_jetty_depth = 0; + uint32_t max_send_sge = 0; + uint32_t max_remote_sge = 0; + uint64_t max_message_size = 0; + uint32_t max_read_size = 0; + uint32_t max_write_size = 0; + uint32_t feature_flags = 0; + uint16_t transport_modes = 0; +}; + +// One DeviceInfo represents one native URMA device/EID pair. A native device +// with multiple EIDs therefore produces multiple entries. topology_name is a +// stable TENT-facing identity and native_device_name is passed only to URMA. +struct DeviceInfo { + std::string topology_name; + std::string native_device_name; + std::string native_device_path; + uint32_t eid_index = 0; + std::string eid; + bool active = false; + DeviceCapabilities capabilities; +}; + +// The segment descriptor is intentionally an explicit envelope. Raw +// urma_seg_t bytes are ABI-specific; import must reject a different API +// version, structure size, schema, or malformed hex instead of copying an +// arbitrary byte sequence into a native structure. +struct SegmentDescriptor { + static constexpr uint32_t kSchemaVersion = 1; + + uint32_t schema_version = kSchemaVersion; + uint32_t urma_api_version = 0; + uint32_t urma_abi_size = 0; + std::string hex; +}; + +// Native handles never escape through this interface. Concrete adapters own +// their raw handles and make explicit release operations retryable. All typed +// handles are reference counted so slices can retain segments until their +// completion has been consumed. +class OpaqueHandle { + public: + OpaqueHandle() = default; + virtual ~OpaqueHandle() = default; + + OpaqueHandle(const OpaqueHandle&) = delete; + OpaqueHandle& operator=(const OpaqueHandle&) = delete; + + [[nodiscard]] virtual bool valid() const noexcept = 0; +}; + +class Context : public OpaqueHandle { + public: + [[nodiscard]] virtual const DeviceInfo& deviceInfo() const noexcept = 0; + [[nodiscard]] virtual int asyncFd() const noexcept = 0; +}; + +class Jfc : public OpaqueHandle { + public: + // Returns -1 when completion events were not requested. + [[nodiscard]] virtual int eventFd() const noexcept = 0; +}; + +class LocalSegment : public OpaqueHandle { + public: + [[nodiscard]] virtual uint64_t address() const noexcept = 0; + [[nodiscard]] virtual uint64_t length() const noexcept = 0; + [[nodiscard]] virtual const SegmentDescriptor& descriptor() + const noexcept = 0; +}; + +class RemoteSegment : public OpaqueHandle { + public: + [[nodiscard]] virtual uint64_t address() const noexcept = 0; + [[nodiscard]] virtual uint64_t length() const noexcept = 0; + [[nodiscard]] virtual const SegmentDescriptor& descriptor() + const noexcept = 0; +}; + +class Jetty : public OpaqueHandle { + public: + [[nodiscard]] virtual uint32_t id() const noexcept = 0; + [[nodiscard]] virtual uint32_t uasid() const noexcept = 0; +}; + +using ContextPtr = std::shared_ptr; +using JfcPtr = std::shared_ptr; +using LocalSegmentPtr = std::shared_ptr; +using RemoteSegmentPtr = std::shared_ptr; +using JettyPtr = std::shared_ptr; + +enum SegmentAccess : uint32_t { + SEGMENT_ACCESS_READ = 1U << 0, + SEGMENT_ACCESS_WRITE = 1U << 1, + SEGMENT_ACCESS_ATOMIC = 1U << 2, + // Local CPU/URMA access remains unrestricted, but the provider must reject + // every access initiated by a remote endpoint. + SEGMENT_ACCESS_LOCAL_ONLY = 1U << 3, +}; + +struct SegmentOptions { + uint32_t access = SEGMENT_ACCESS_READ | SEGMENT_ACCESS_WRITE; + uint32_t token = 0xACFE; + bool cacheable = false; +}; + +struct JfcOptions { + uint32_t depth = 4096; + uint32_t receiver_depth = 2048; + uint32_t token = 0xACFE; + bool enable_completion_events = false; +}; + +struct JettyOptions { + uint32_t depth = 256; + uint8_t priority = 15; + uint8_t max_sge = 1; + uint8_t rnr_retry = 7; + uint8_t error_timeout = 17; +}; + +struct RemoteJettyInfo { + std::string eid; + uint32_t id = 0; + uint32_t uasid = 0; + uint32_t token = 0xACFE; +}; + +enum class Operation : uint8_t { + READ, + WRITE, +}; + +struct WorkRequest { + Operation operation = Operation::READ; + uint64_t local_address = 0; + uint64_t remote_address = 0; + size_t length = 0; + + // Zero is reserved for native completion records that do not correspond + // to a posted work request (for example a synthetic flush-done record). + uint64_t token = 0; + + LocalSegmentPtr local_segment; + RemoteSegmentPtr remote_segment; +}; + +enum class CompletionCategory : uint8_t { + SUCCESS, + LOCAL_DEVICE_ERROR, + REMOTE_PATH_ERROR, + ENDPOINT_ERROR, + MEMORY_ERROR, + TIMEOUT, + UNKNOWN_ERROR, +}; + +struct Completion { + CompletionCategory category = CompletionCategory::UNKNOWN_ERROR; + int native_status = 0; + uint64_t token = 0; + uint32_t completed_bytes = 0; + // Native Jetty ID that produced the completion. This is also populated + // for entity-level flush markers whose token is deliberately zero. + uint32_t local_jetty_id = 0; +}; + +// Pure, injectable URMA boundary. It deliberately does not know about TENT +// Request, SubBatch, UbTask, UbSlice, scheduling, retry, or rail health. +class UrmaAdapter { + public: + virtual ~UrmaAdapter() = default; + + [[nodiscard]] virtual bool available() const noexcept = 0; + [[nodiscard]] virtual uint32_t nativeApiVersion() const noexcept = 0; + [[nodiscard]] virtual size_t nativeSegmentDescriptorSize() + const noexcept = 0; + + virtual Status initialize() = 0; + virtual Status shutdown() = 0; + + virtual Status discoverDevices(std::vector& devices) = 0; + + virtual Status openContext(const DeviceInfo& device, + ContextPtr& context) = 0; + // close/delete/unregister/unimport reset the caller's shared_ptr only after + // the provider has released the native resource. If another owner still + // retains the handle, or if the provider returns busy/error, the operation + // fails and leaves the shared_ptr intact for a later retry. Calling any of + // these methods again with a null handle is successful. + virtual Status closeContext(ContextPtr& context) = 0; + + virtual Status createJfc(const ContextPtr& context, + const JfcOptions& options, JfcPtr& jfc) = 0; + virtual Status deleteJfc(JfcPtr& jfc) = 0; + + virtual Status registerLocalSegment(const ContextPtr& context, + uint64_t address, size_t length, + const SegmentOptions& options, + LocalSegmentPtr& segment) = 0; + virtual Status unregisterLocalSegment(LocalSegmentPtr& segment) = 0; + + virtual Status importRemoteSegment(const ContextPtr& context, + const SegmentDescriptor& descriptor, + const SegmentOptions& options, + RemoteSegmentPtr& segment) = 0; + virtual Status unimportRemoteSegment(RemoteSegmentPtr& segment) = 0; + + virtual Status createJetty(const ContextPtr& context, const JfcPtr& jfc, + const JettyOptions& options, + JettyPtr& jetty) = 0; + virtual Status deleteJetty(JettyPtr& jetty) = 0; + virtual Status bindJetty(const JettyPtr& jetty, + const RemoteJettyInfo& remote) = 0; + virtual Status unbindJetty(const JettyPtr& jetty) = 0; + virtual Status resetJetty(const JettyPtr& jetty) = 0; + + // Synchronously fences one Jetty. A successful return guarantees that no + // previously posted WR on this Jetty can perform any further DMA. The + // implementation transitions the native Jetty to ERROR, consumes the + // provider's flush-done marker, and returns both already-processed and + // unhandled WR completions. Callers must dispatch every non-zero token + // through their normal completion path before resetting or deleting the + // Jetty. On failure no drain guarantee is made and native resources must + // remain alive. + virtual Status quiesceJetty(const JettyPtr& jetty, uint32_t timeout_ms, + std::vector& completions) = 0; + + // On a native post error, posted_count is the number of leading requests + // accepted before the provider's bad WR. Requests counted as posted must + // still be completed through poll(). + virtual Status post(const JettyPtr& jetty, + const std::vector& requests, + size_t& posted_count) = 0; + virtual Status poll(const JfcPtr& jfc, size_t max_completions, + std::vector& completions) = 0; +}; + +// Returns the raw-liburma implementation when TENT_HAS_REAL_URMA is enabled; +// otherwise returns an injectable-compatible stub whose operational methods +// report Status::NotImplemented with an explicit unavailable message. +std::shared_ptr createDefaultUrmaAdapter(); + +} // namespace ub +} // namespace tent +} // namespace mooncake + +#endif // TENT_TRANSPORT_UB_URMA_ADAPTER_H_ diff --git a/mooncake-transfer-engine/tent/src/CMakeLists.txt b/mooncake-transfer-engine/tent/src/CMakeLists.txt index 73766ce120..6fe7ff3559 100644 --- a/mooncake-transfer-engine/tent/src/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/CMakeLists.txt @@ -28,7 +28,8 @@ if(USE_HIP) list(APPEND CMAKE_PREFIX_PATH "/opt/rocm/lib/cmake") find_package(HIP REQUIRED) message(STATUS "ROCm/HIP: Enabled") - target_compile_definitions(tent_interface INTERFACE USE_HIP __HIP_PLATFORM_AMD__) + target_compile_definitions(tent_interface INTERFACE USE_HIP + __HIP_PLATFORM_AMD__) target_include_directories(tent_interface INTERFACE ${HIP_INCLUDE_DIRS}) target_link_libraries(tent_interface INTERFACE hip::host) else() @@ -36,8 +37,8 @@ else() endif() # GDS -find_library(CUFILE_LIB cufile PATHS /usr/local/cuda/lib64) -find_path(CUFILE_INCLUDE cufile.h PATHS /usr/local/cuda/include) +find_library(CUFILE_LIB cufile HINTS ${CUDAToolkit_LIBRARY_DIR}) +find_path(CUFILE_INCLUDE cufile.h HINTS ${CUDAToolkit_INCLUDE_DIRS}) if(USE_CUDA AND CUDAToolkit_FOUND AND CUFILE_LIB @@ -80,9 +81,9 @@ endif() if(USE_SUNRISE) target_compile_definitions(tent_interface INTERFACE USE_SUNRISE) target_include_directories(tent_interface INTERFACE ${MC_TANGRT_ROOT}/include) - target_link_directories(tent_interface INTERFACE - ${MC_TANGRT_ROOT}/lib/linux-x86_64 - ${MC_TANGRT_ROOT}/lib) + target_link_directories( + tent_interface INTERFACE ${MC_TANGRT_ROOT}/lib/linux-x86_64 + ${MC_TANGRT_ROOT}/lib) target_link_libraries(tent_interface INTERFACE tangrt_shared ptml_shared dl) endif() @@ -131,6 +132,7 @@ foreach( platform_rocm platform_ascend platform_sunrise + platform_tpu tent_xport_gds tent_xport_uring tent_xport_bufio @@ -139,13 +141,23 @@ foreach( tent_xport_rdma tent_xport_shm tent_xport_tcp + tent_xport_ub tent_xport_ascend_direct tent_xport_sunrise_link + tent_xport_tpu tent_metrics) if(TARGET ${tgt}) target_link_libraries(tent_link_group INTERFACE ${tgt}) endif() endforeach() +if(TARGET mooncake_common) + # Use the archive path so CMake does not de-duplicate mooncake_common against + # an earlier transitive occurrence. TENT RPC objects reference the Transfer + # Engine client I/O pool accessor and need the archive inside this rescan + # group when all components are linked statically. + target_link_libraries(tent_link_group + INTERFACE "$") +endif() target_link_libraries(tent_link_group INTERFACE "-Wl,--end-group") install( TARGETS tent_shared diff --git a/mooncake-transfer-engine/tent/src/common/config.cpp b/mooncake-transfer-engine/tent/src/common/config.cpp index 1489cf6d29..7ddbc89ca9 100644 --- a/mooncake-transfer-engine/tent/src/common/config.cpp +++ b/mooncake-transfer-engine/tent/src/common/config.cpp @@ -55,6 +55,24 @@ static inline void setConfig(Config& config, const std::string& env_key, if (val) config.setFromString(config_key, std::string(val)); } +// Like setConfig, but parses the env value as a comma-separated list and +// stores it as a string array. Empty/whitespace-only items are dropped so a +// trailing comma or spaces around names are tolerated (e.g. "mlx5_0, mlx5_1"). +static inline void setArrayConfig(Config& config, const std::string& env_key, + const std::string& config_key) { + const char* val = std::getenv(env_key.c_str()); + if (!val) return; + std::vector items; + std::stringstream ss(val); + std::string item; + while (std::getline(ss, item, ',')) { + item.erase(0, item.find_first_not_of(" \t")); + item.erase(item.find_last_not_of(" \t") + 1); + if (!item.empty()) items.push_back(item); + } + if (!items.empty()) config.set(config_key, items); +} + Status ConfigHelper::loadFromEnv(Config& config) { const char* conf_str = std::getenv("MC_TENT_CONF"); Status status = Status::OK(); @@ -107,6 +125,7 @@ Status ConfigHelper::loadFromEnv(Config& config) { setConfig(config, "MC_PKEY_INDEX", "transports/rdma/endpoint/pkey_index"); setConfig(config, "MC_MTU", "transports/rdma/endpoint/path_mtu"); setConfig(config, "MC_IB_TC", "transports/rdma/endpoint/traffic_class"); + setConfig(config, "MC_IB_SL", "transports/rdma/endpoint/service_level"); setConfig(config, "MC_IB_PCI_RELAXED_ORDERING", "transports/rdma/pci_relaxed_ordering"); setConfig(config, "MC_WORKERS_PER_CTX", @@ -116,6 +135,16 @@ Status ConfigHelper::loadFromEnv(Config& config) { "transports/rdma/workers/max_retry_count"); setConfig(config, "MC_DISABLE_GPU_DIRECT_RDMA", "transports/rdma/disable_gpu_direct_rdma"); + setConfig(config, "MC_LOG_RDMA_SLICE_AFFINITY", + "transports/rdma/log_slice_affinity"); + // Restrict which RDMA NICs the engine discovers/uses (comma-separated + // device names). MC_TE_FILTERS is an allow-list — same name and semantics + // as the legacy Transfer Engine's device whitelist, so a single env works + // across both engines. MC_TE_FILTERS_EXCLUDE is a deny-list (new; the + // legacy engine has no deny-list). Unset = discover all (default). + // Consumed by filterInfiniBandDevices() in the platform probes. + setArrayConfig(config, "MC_TE_FILTERS", "topology/rdma_whitelist"); + setArrayConfig(config, "MC_TE_FILTERS_EXCLUDE", "topology/rdma_blacklist"); return status; } @@ -139,7 +168,14 @@ bool ConfigHelper::parseBool(const std::string& str, bool default_value) { int ConfigHelper::parseInt(const std::string& str, int default_value) { try { - return std::stoi(str); + size_t parsed = 0; + int value = std::stoi(str, &parsed); + if (parsed != str.size()) { + LOG(WARNING) << "Invalid integer value '" << str + << "', using default: " << default_value; + return default_value; + } + return value; } catch (const std::exception& e) { LOG(WARNING) << "Failed to parse integer '" << str << "': " << e.what() << ", using default: " << default_value; @@ -150,7 +186,13 @@ int ConfigHelper::parseInt(const std::string& str, int default_value) { uint16_t ConfigHelper::parsePort(const std::string& str, uint16_t default_value) { try { - int port = std::stoi(str); + size_t parsed = 0; + int port = std::stoi(str, &parsed); + if (parsed != str.size()) { + LOG(WARNING) << "Invalid port value '" << str + << "', using default: " << default_value; + return default_value; + } if (port > 0 && port <= 65535) { return static_cast(port); } else { diff --git a/mooncake-transfer-engine/tent/src/common/qos_metrics.cpp b/mooncake-transfer-engine/tent/src/common/qos_metrics.cpp new file mode 100644 index 0000000000..961c826efb --- /dev/null +++ b/mooncake-transfer-engine/tent/src/common/qos_metrics.cpp @@ -0,0 +1,379 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/common/qos_metrics.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake { +namespace tent { +namespace { + +std::vector split(const std::string& value, char delimiter) { + std::vector parts; + std::stringstream stream(value); + std::string part; + while (std::getline(stream, part, delimiter)) parts.push_back(part); + return parts; +} + +template +bool parseNumber(const std::string& value, T* result) { + if (value.empty()) return false; + std::istringstream stream(value); + stream >> std::noskipws >> *result; + return stream.eof() && !stream.fail(); +} + +double jainIndex(const std::vector& values) { + if (values.empty()) return 0.0; + double sum = 0.0; + double squared_sum = 0.0; + for (double value : values) { + sum += value; + squared_sum += value * value; + } + if (squared_sum == 0.0) return 0.0; + return sum * sum / (values.size() * squared_sum); +} + +nlohmann::json optionalJson(const std::optional& value) { + return value ? nlohmann::json(*value) : nlohmann::json(nullptr); +} + +} // namespace + +bool parseQosClasses(const std::string& spec, + std::vector* classes, std::string* error) { + classes->clear(); + if (spec.empty()) { + *error = "qos_classes must not be empty"; + return false; + } + + std::set names; + for (const auto& entry : split(spec, ',')) { + const auto fields = split(entry, ':'); + if (fields.size() != 4 && fields.size() != 5) { + *error = + "each qos class must be " + "name:threads:slo_us:weight[:isolated_gbps]"; + return false; + } + + QosClassConfig config; + config.name = fields[0]; + if (config.name.empty() || !names.insert(config.name).second) { + *error = "qos class names must be non-empty and unique"; + return false; + } + if (!parseNumber(fields[1], &config.threads) || config.threads <= 0) { + *error = "qos class threads must be a positive integer"; + return false; + } + if (fields[2].empty() || fields[2].front() == '-' || + !parseNumber(fields[2], &config.slo_us)) { + *error = "qos class slo_us must be a non-negative integer"; + return false; + } + if (!parseNumber(fields[3], &config.weight) || config.weight <= 0.0 || + !std::isfinite(config.weight)) { + *error = "qos class weight must be finite and positive"; + return false; + } + if (fields.size() == 5) { + double isolated_gbps = 0.0; + if (!parseNumber(fields[4], &isolated_gbps) || + isolated_gbps <= 0.0 || !std::isfinite(isolated_gbps)) { + *error = "isolated_gbps must be finite and positive"; + return false; + } + config.isolated_throughput_gbps = isolated_gbps; + } + classes->push_back(std::move(config)); + } + return true; +} + +bool parseQosClassesJson(const std::string& spec, + std::vector* classes, + std::string* error) { + classes->clear(); + try { + const auto root = nlohmann::json::parse(spec); + if (!root.is_array()) { + *error = "qos_classes_json must be an array"; + return false; + } + std::set names; + for (size_t i = 0; i < root.size(); ++i) { + const auto& node = root[i]; + const std::string path = + "qos_classes_json[" + std::to_string(i) + "]"; + if (!node.is_object()) { + *error = path + " must be an object"; + return false; + } + QosClassConfig config; + if (!node.contains("name") || !node["name"].is_string()) { + *error = path + ".name must be a string"; + return false; + } + config.name = node["name"].get(); + if (config.name.empty() || !names.insert(config.name).second) { + *error = "qos class names must be non-empty and unique"; + return false; + } + if (!node.contains("threads") || + !node["threads"].is_number_integer()) { + *error = path + ".threads must be an integer"; + return false; + } + config.threads = node["threads"].get(); + if (config.threads <= 0) { + *error = path + ".threads must be positive"; + return false; + } + if (!node.contains("slo_us") || + !node["slo_us"].is_number_unsigned()) { + *error = path + ".slo_us must be an unsigned integer"; + return false; + } + config.slo_us = node["slo_us"].get(); + if (!node.contains("weight") || !node["weight"].is_number()) { + *error = path + ".weight must be numeric"; + return false; + } + config.weight = node["weight"].get(); + if (config.weight <= 0.0 || !std::isfinite(config.weight)) { + *error = path + ".weight must be finite and positive"; + return false; + } + if (node.contains("isolated_gbps") && + !node["isolated_gbps"].is_null()) { + if (!node["isolated_gbps"].is_number()) { + *error = path + ".isolated_gbps must be numeric or null"; + return false; + } + const double isolated_gbps = + node["isolated_gbps"].get(); + if (isolated_gbps <= 0.0 || !std::isfinite(isolated_gbps)) { + *error = + path + ".isolated_gbps must be finite and positive"; + return false; + } + config.isolated_throughput_gbps = isolated_gbps; + } + classes->push_back(std::move(config)); + } + return true; + } catch (const std::exception& e) { + *error = std::string("failed to parse qos_classes_json: ") + e.what(); + return false; + } +} + +bool validateQosClasses(const std::vector& classes, + int num_threads, std::string* error) { + int configured_threads = 0; + for (const auto& config : classes) configured_threads += config.threads; + if (configured_threads != num_threads) { + std::ostringstream stream; + stream << "qos_classes configures " << configured_threads + << " threads, but tebench runs " << num_threads; + *error = stream.str(); + return false; + } + return true; +} + +size_t qosClassForThread(const std::vector& classes, + int thread_id) { + int boundary = 0; + for (size_t i = 0; i < classes.size(); ++i) { + boundary += classes[i].threads; + if (thread_id < boundary) return i; + } + return classes.size(); +} + +QosMetricsReport calculateQosMetrics(size_t block_size, size_t batch_size, + int num_threads, + const std::vector& classes, + const std::vector& samples, + double link_capacity_gbps) { + QosMetricsReport report; + report.block_size = block_size; + report.batch_size = batch_size; + report.num_threads = num_threads; + + std::vector normalized_throughput; + std::optional max_leakage; + for (size_t i = 0; i < classes.size(); ++i) { + const auto& config = classes[i]; + const auto& sample = samples[i]; + QosClassMetrics metrics; + metrics.name = config.name; + metrics.threads = config.threads; + metrics.slo_us = config.slo_us; + metrics.weight = config.weight; + metrics.isolated_throughput_gbps = config.isolated_throughput_gbps; + metrics.operations = sample.operations; + metrics.transferred_bytes = sample.transferred_bytes.value_or( + static_cast(block_size) * batch_size * + metrics.operations); + metrics.p99_us = sample.p99_us; + + const double duration_s = sample.total_duration_us / 1e6; + if (duration_s > 0.0) + metrics.throughput_gbps = + static_cast(metrics.transferred_bytes) / 1e9 / + duration_s; + + double attainment = 1.0; + if (config.slo_us != 0) { + attainment = sample.slo_attainment.value_or(0.0); + metrics.slo_attainment = attainment; + } + metrics.goodput_gbps = metrics.throughput_gbps * attainment; + metrics.weighted_goodput_gbps = metrics.goodput_gbps * config.weight; + normalized_throughput.push_back(metrics.throughput_gbps / + config.weight); + + if (config.isolated_throughput_gbps) { + metrics.isolation_leakage = + std::max(0.0, 1.0 - metrics.throughput_gbps / + *config.isolated_throughput_gbps); + max_leakage = + max_leakage ? std::max(*max_leakage, *metrics.isolation_leakage) + : metrics.isolation_leakage; + } + + report.aggregate_throughput_gbps += metrics.throughput_gbps; + report.weighted_goodput_gbps += metrics.weighted_goodput_gbps; + report.classes.push_back(std::move(metrics)); + } + + report.jain_fairness = jainIndex(normalized_throughput); + report.max_isolation_leakage = max_leakage; + if (link_capacity_gbps > 0.0) { + report.link_capacity_gbps = link_capacity_gbps; + report.total_utilization = + report.aggregate_throughput_gbps / link_capacity_gbps; + } + return report; +} + +void printQosMetrics(const QosMetricsReport& report) { + std::cout << " [qos-summary] throughput=" << std::fixed + << std::setprecision(6) << report.aggregate_throughput_gbps + << " GB/s weighted_goodput=" << report.weighted_goodput_gbps + << " GB/s jain_fairness=" << report.jain_fairness + << " max_isolation_leakage="; + if (report.max_isolation_leakage) { + std::cout << *report.max_isolation_leakage; + } else { + std::cout << "N/A"; + } + std::cout << " total_utilization="; + if (report.total_utilization) { + std::cout << *report.total_utilization; + } else { + std::cout << "N/A"; + } + std::cout << std::endl; + + for (const auto& metrics : report.classes) { + std::cout << " [qos-class] name=" << metrics.name + << " threads=" << metrics.threads + << " operations=" << metrics.operations + << " transferred_bytes=" << metrics.transferred_bytes + << " throughput=" << metrics.throughput_gbps + << " GB/s p99_us=" << std::setprecision(1) << metrics.p99_us + << " slo_attainment="; + std::cout << std::setprecision(6); + if (metrics.slo_attainment) { + std::cout << std::setprecision(6) << *metrics.slo_attainment; + } else { + std::cout << "N/A"; + } + std::cout << " isolation_leakage="; + if (metrics.isolation_leakage) { + std::cout << *metrics.isolation_leakage; + } else { + std::cout << "N/A"; + } + std::cout << std::endl; + } +} + +bool appendQosMetricsJsonl(const std::string& path, + const QosMetricsReport& report, std::string* error) { + nlohmann::json root = { + {"schema_version", 1}, + {"block_size", report.block_size}, + {"batch_size", report.batch_size}, + {"num_threads", report.num_threads}, + {"aggregate_throughput_gbps", report.aggregate_throughput_gbps}, + {"weighted_goodput_gbps", report.weighted_goodput_gbps}, + {"jain_fairness", report.jain_fairness}, + {"max_isolation_leakage", optionalJson(report.max_isolation_leakage)}, + {"link_capacity_gbps", optionalJson(report.link_capacity_gbps)}, + {"total_utilization", optionalJson(report.total_utilization)}, + {"classes", nlohmann::json::array()}, + }; + for (const auto& metrics : report.classes) { + root["classes"].push_back({ + {"name", metrics.name}, + {"threads", metrics.threads}, + {"slo_us", metrics.slo_us}, + {"weight", metrics.weight}, + {"operations", metrics.operations}, + {"transferred_bytes", metrics.transferred_bytes}, + {"throughput_gbps", metrics.throughput_gbps}, + {"p99_us", metrics.p99_us}, + {"slo_attainment", optionalJson(metrics.slo_attainment)}, + {"goodput_gbps", metrics.goodput_gbps}, + {"weighted_goodput_gbps", metrics.weighted_goodput_gbps}, + {"isolated_throughput_gbps", + optionalJson(metrics.isolated_throughput_gbps)}, + {"isolation_leakage", optionalJson(metrics.isolation_leakage)}, + }); + } + + std::ofstream output(path, std::ios::app); + if (!output) { + *error = "failed to open QoS JSONL output: " + path; + return false; + } + output << root.dump() << '\n'; + if (!output) { + *error = "failed to write QoS JSONL output: " + path; + return false; + } + return true; +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/metastore/http.cpp b/mooncake-transfer-engine/tent/src/metastore/http.cpp index d5f1bfc132..071d3603d1 100644 --- a/mooncake-transfer-engine/tent/src/metastore/http.cpp +++ b/mooncake-transfer-engine/tent/src/metastore/http.cpp @@ -15,10 +15,21 @@ #include "tent/metastore/http.h" #include +#include namespace mooncake { namespace tent { +static std::once_flag g_curl_global_init_flag; +static void ensureCurlGlobalInit() { + std::call_once(g_curl_global_init_flag, []() { + CURLcode rc = curl_global_init(CURL_GLOBAL_ALL); + if (rc != CURLE_OK) { + LOG(ERROR) << "curl_global_init failed: " << curl_easy_strerror(rc); + } + }); +} + HttpMetaStore::HttpMetaStore() {} HttpMetaStore::~HttpMetaStore() { disconnect(); } @@ -28,42 +39,51 @@ Status HttpMetaStore::connect(const std::string &endpoint) { return Status::MetadataError( "HTTP connection already established" LOC_MARK); } - curl_global_init(CURL_GLOBAL_ALL); - client_ = curl_easy_init(); - if (!client_) { - return Status::InternalError( - "HTTP cannot allocate curl objects" LOC_MARK); - } + ensureCurlGlobalInit(); endpoint_ = endpoint; connected_ = true; return Status::OK(); } Status HttpMetaStore::disconnect() { - if (connected_) { - curl_easy_cleanup(client_); - curl_global_cleanup(); - connected_ = false; - } + connected_ = false; return Status::OK(); } +namespace { +struct ScopedCurl { + CURL *h{nullptr}; + ScopedCurl() : h(curl_easy_init()) {} + ~ScopedCurl() { + if (h) curl_easy_cleanup(h); + } + ScopedCurl(const ScopedCurl &) = delete; + ScopedCurl &operator=(const ScopedCurl &) = delete; + operator CURL *() const { return h; } + explicit operator bool() const { return h != nullptr; } +}; +} // namespace + Status HttpMetaStore::get(const std::string &key, std::string &value) { if (!connected_) { return Status::MetadataError("HTTP connection not available" LOC_MARK); } - curl_easy_reset(client_); - curl_easy_setopt(client_, CURLOPT_TIMEOUT_MS, 3000); // 3s timeout + ScopedCurl client; + if (!client) { + return Status::InternalError( + "HTTP cannot allocate curl handle" LOC_MARK); + } + curl_easy_setopt(client.h, CURLOPT_TIMEOUT_MS, 3000); // 3s timeout - std::string url = encodeUrl(key); - curl_easy_setopt(client_, CURLOPT_URL, url.c_str()); - curl_easy_setopt(client_, CURLOPT_WRITEFUNCTION, writeCallback); + std::string url = encodeUrl(client.h, key); + curl_easy_setopt(client.h, CURLOPT_URL, url.c_str()); + curl_easy_setopt(client.h, CURLOPT_WRITEFUNCTION, writeCallback); // get response body std::string readBuffer; - curl_easy_setopt(client_, CURLOPT_WRITEDATA, &readBuffer); - CURLcode res = curl_easy_perform(client_); + curl_easy_setopt(client.h, CURLOPT_WRITEDATA, &readBuffer); + CURLcode res = curl_easy_perform(client.h); if (res != CURLE_OK) { return Status::MetadataError( std::string("HTTP failed to post request: ") + @@ -72,7 +92,7 @@ Status HttpMetaStore::get(const std::string &key, std::string &value) { // Get the HTTP response code long responseCode; - curl_easy_getinfo(client_, CURLINFO_RESPONSE_CODE, &responseCode); + curl_easy_getinfo(client.h, CURLINFO_RESPONSE_CODE, &responseCode); if (responseCode == 404) { return Status::InvalidEntry(key); } else if (responseCode != 200) { @@ -81,7 +101,7 @@ Status HttpMetaStore::get(const std::string &key, std::string &value) { std::string("HTTP received unexpected response: ") + message + LOC_MARK); } - value = std::string(readBuffer); + value = std::move(readBuffer); return Status::OK(); } @@ -90,25 +110,29 @@ Status HttpMetaStore::set(const std::string &key, const std::string &value) { return Status::MetadataError("HTTP connection not available" LOC_MARK); } - curl_easy_reset(client_); - curl_easy_setopt(client_, CURLOPT_TIMEOUT_MS, 3000); // 3s timeout + ScopedCurl client; + if (!client) { + return Status::InternalError( + "HTTP cannot allocate curl handle" LOC_MARK); + } + curl_easy_setopt(client.h, CURLOPT_TIMEOUT_MS, 3000); // 3s timeout - std::string url = encodeUrl(key); - curl_easy_setopt(client_, CURLOPT_URL, url.c_str()); - curl_easy_setopt(client_, CURLOPT_WRITEFUNCTION, writeCallback); - curl_easy_setopt(client_, CURLOPT_POSTFIELDS, value.c_str()); - curl_easy_setopt(client_, CURLOPT_POSTFIELDSIZE, value.size()); - curl_easy_setopt(client_, CURLOPT_CUSTOMREQUEST, "PUT"); + std::string url = encodeUrl(client.h, key); + curl_easy_setopt(client.h, CURLOPT_URL, url.c_str()); + curl_easy_setopt(client.h, CURLOPT_WRITEFUNCTION, writeCallback); + curl_easy_setopt(client.h, CURLOPT_POSTFIELDS, value.c_str()); + curl_easy_setopt(client.h, CURLOPT_POSTFIELDSIZE, value.size()); + curl_easy_setopt(client.h, CURLOPT_CUSTOMREQUEST, "PUT"); // get response body std::string readBuffer; - curl_easy_setopt(client_, CURLOPT_WRITEDATA, &readBuffer); + curl_easy_setopt(client.h, CURLOPT_WRITEDATA, &readBuffer); // set content-type to application/json struct curl_slist *headers = NULL; headers = curl_slist_append(headers, "Content-Type: application/json"); - curl_easy_setopt(client_, CURLOPT_HTTPHEADER, headers); - CURLcode res = curl_easy_perform(client_); + curl_easy_setopt(client.h, CURLOPT_HTTPHEADER, headers); + CURLcode res = curl_easy_perform(client.h); curl_slist_free_all(headers); // free headers if (res != CURLE_OK) { return Status::MetadataError( @@ -117,7 +141,7 @@ Status HttpMetaStore::set(const std::string &key, const std::string &value) { } long responseCode; - curl_easy_getinfo(client_, CURLINFO_RESPONSE_CODE, &responseCode); + curl_easy_getinfo(client.h, CURLINFO_RESPONSE_CODE, &responseCode); if (responseCode != 200) { std::string message = std::to_string(responseCode) + ": " + readBuffer; return Status::MetadataError( @@ -133,18 +157,22 @@ Status HttpMetaStore::remove(const std::string &key) { return Status::MetadataError("HTTP connection not available" LOC_MARK); } - curl_easy_reset(client_); - curl_easy_setopt(client_, CURLOPT_TIMEOUT_MS, 3000); // 3s timeout + ScopedCurl client; + if (!client) { + return Status::InternalError( + "HTTP cannot allocate curl handle" LOC_MARK); + } + curl_easy_setopt(client.h, CURLOPT_TIMEOUT_MS, 3000); // 3s timeout - std::string url = encodeUrl(key); - curl_easy_setopt(client_, CURLOPT_URL, url.c_str()); - curl_easy_setopt(client_, CURLOPT_WRITEFUNCTION, writeCallback); - curl_easy_setopt(client_, CURLOPT_CUSTOMREQUEST, "DELETE"); + std::string url = encodeUrl(client.h, key); + curl_easy_setopt(client.h, CURLOPT_URL, url.c_str()); + curl_easy_setopt(client.h, CURLOPT_WRITEFUNCTION, writeCallback); + curl_easy_setopt(client.h, CURLOPT_CUSTOMREQUEST, "DELETE"); // get response body std::string readBuffer; - curl_easy_setopt(client_, CURLOPT_WRITEDATA, &readBuffer); - CURLcode res = curl_easy_perform(client_); + curl_easy_setopt(client.h, CURLOPT_WRITEDATA, &readBuffer); + CURLcode res = curl_easy_perform(client.h); if (res != CURLE_OK) { return Status::MetadataError( std::string("HTTP failed to post request: ") + @@ -152,7 +180,7 @@ Status HttpMetaStore::remove(const std::string &key) { } long responseCode; - curl_easy_getinfo(client_, CURLINFO_RESPONSE_CODE, &responseCode); + curl_easy_getinfo(client.h, CURLINFO_RESPONSE_CODE, &responseCode); if (responseCode != 200) { std::string message = std::to_string(responseCode) + ": " + readBuffer; return Status::MetadataError( diff --git a/mooncake-transfer-engine/tent/src/metrics/CMakeLists.txt b/mooncake-transfer-engine/tent/src/metrics/CMakeLists.txt index ca073167fd..5b8720e05d 100644 --- a/mooncake-transfer-engine/tent/src/metrics/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/metrics/CMakeLists.txt @@ -1,20 +1,22 @@ -# Option to enable/disable TENT metrics at compile time -option(TENT_METRICS_ENABLED "Enable TENT metrics collection" OFF) - +# TENT_METRICS_ENABLED option is declared in tent/CMakeLists.txt. file(GLOB TENT_METRICS_SOURCES "*.cpp") add_library(tent_metrics STATIC ${TENT_METRICS_SOURCES}) -# Link yalantinglibs::yalantinglibs for metrics and HTTP server (ylt headers + transitive deps). -# ODR safety: yalantinglibs bundles ASIO headers but does NOT compile ASIO inline because -# ASIO_SEPARATE_COMPILATION is set globally. All ASIO symbols live exclusively in asio_shared.so, -# so there is no risk of duplicate symbols between yalantinglibs and the rest of TE. -# (See mooncake-common/src/CMakeLists.txt: "Build asio as a shared library to avoid ODR violations") -target_link_libraries(tent_metrics PUBLIC tent_common tent_interface yalantinglibs::yalantinglibs glog pthread) +# Link yalantinglibs::yalantinglibs for metrics and HTTP server (ylt headers + +# transitive deps). ODR safety: yalantinglibs bundles ASIO headers but does NOT +# compile ASIO inline because ASIO_SEPARATE_COMPILATION is set globally. All +# ASIO symbols live exclusively in asio_shared.so, so there is no risk of +# duplicate symbols between yalantinglibs and the rest of TE. (See +# mooncake-common/src/CMakeLists.txt: "Build asio as a shared library to avoid +# ODR violations") +target_link_libraries( + tent_metrics PUBLIC tent_common tent_interface yalantinglibs::yalantinglibs + glog pthread) # Pass compile definition based on option if(TENT_METRICS_ENABLED) - target_compile_definitions(tent_metrics PUBLIC TENT_METRICS_ENABLED=1) - message(STATUS "TENT metrics: ENABLED") + target_compile_definitions(tent_metrics PUBLIC TENT_METRICS_ENABLED=1) + message(STATUS "TENT metrics: ENABLED") else() - target_compile_definitions(tent_metrics PUBLIC TENT_METRICS_ENABLED=0) - message(STATUS "TENT metrics: DISABLED (zero overhead)") + target_compile_definitions(tent_metrics PUBLIC TENT_METRICS_ENABLED=0) + message(STATUS "TENT metrics: DISABLED (zero overhead)") endif() diff --git a/mooncake-transfer-engine/tent/src/metrics/config_loader.cpp b/mooncake-transfer-engine/tent/src/metrics/config_loader.cpp index e948a7c53d..acbe09f945 100644 --- a/mooncake-transfer-engine/tent/src/metrics/config_loader.cpp +++ b/mooncake-transfer-engine/tent/src/metrics/config_loader.cpp @@ -48,34 +48,6 @@ void MetricsConfigLoader::applyEnvironmentOverrides(MetricsConfig& config) { config.http_server_threads = static_cast(threads); } } - - if (const char* env_val = - std::getenv(config_keys::ENV_METRICS_ENABLE_PROMETHEUS)) { - config.enable_prometheus = - ConfigHelper::parseBool(env_val, config.enable_prometheus); - } - - if (const char* env_val = - std::getenv(config_keys::ENV_METRICS_ENABLE_JSON)) { - config.enable_json = - ConfigHelper::parseBool(env_val, config.enable_json); - } - - if (const char* env_val = - std::getenv(config_keys::ENV_METRICS_LATENCY_BUCKETS)) { - auto buckets = ConfigHelper::parseDoubleArray(env_val); - if (!buckets.empty()) { - config.latency_buckets = buckets; - } - } - - if (const char* env_val = - std::getenv(config_keys::ENV_METRICS_SIZE_BUCKETS)) { - auto buckets = ConfigHelper::parseDoubleArray(env_val); - if (!buckets.empty()) { - config.size_buckets = buckets; - } - } } MetricsConfig MetricsConfigLoader::loadFromConfig(const Config& config) { @@ -95,24 +67,6 @@ MetricsConfig MetricsConfigLoader::loadFromConfig(const Config& config) { metrics_config.report_interval_seconds = config.get(config_keys::METRICS_REPORT_INTERVAL, metrics_config.report_interval_seconds); - metrics_config.enable_prometheus = - config.get(config_keys::METRICS_ENABLE_PROMETHEUS, - metrics_config.enable_prometheus); - metrics_config.enable_json = config.get(config_keys::METRICS_ENABLE_JSON, - metrics_config.enable_json); - - // Load bucket configurations - auto latency_buckets_array = - config.getArray(config_keys::METRICS_LATENCY_BUCKETS); - if (!latency_buckets_array.empty()) { - metrics_config.latency_buckets = latency_buckets_array; - } - - auto size_buckets_array = - config.getArray(config_keys::METRICS_SIZE_BUCKETS); - if (!size_buckets_array.empty()) { - metrics_config.size_buckets = size_buckets_array; - } LOG(INFO) << "Loaded metrics config from Config object: enabled=" << metrics_config.enabled @@ -156,23 +110,6 @@ MetricsConfig MetricsConfigLoader::loadWithDefaults(const Config* config) { metrics_config.report_interval_seconds = config->get(config_keys::METRICS_REPORT_INTERVAL, metrics_config.report_interval_seconds); - metrics_config.enable_prometheus = - config->get(config_keys::METRICS_ENABLE_PROMETHEUS, - metrics_config.enable_prometheus); - metrics_config.enable_json = config->get( - config_keys::METRICS_ENABLE_JSON, metrics_config.enable_json); - - auto latency_buckets_array = - config->getArray(config_keys::METRICS_LATENCY_BUCKETS); - if (!latency_buckets_array.empty()) { - metrics_config.latency_buckets = latency_buckets_array; - } - - auto size_buckets_array = - config->getArray(config_keys::METRICS_SIZE_BUCKETS); - if (!size_buckets_array.empty()) { - metrics_config.size_buckets = size_buckets_array; - } } return metrics_config; @@ -196,36 +133,6 @@ bool MetricsConfigLoader::validateConfig(const MetricsConfig& config, return false; } - // Validate at least one output format is enabled - if (!config.enable_prometheus && !config.enable_json) { - if (error_msg) { - *error_msg = - "At least one output format (Prometheus or JSON) must be " - "enabled"; - } - return false; - } - - // Validate buckets are sorted and positive - for (size_t i = 1; i < config.latency_buckets.size(); ++i) { - if (config.latency_buckets[i] <= config.latency_buckets[i - 1]) { - if (error_msg) { - *error_msg = - "Latency buckets must be sorted in ascending order"; - } - return false; - } - } - - for (size_t i = 1; i < config.size_buckets.size(); ++i) { - if (config.size_buckets[i] <= config.size_buckets[i - 1]) { - if (error_msg) { - *error_msg = "Size buckets must be sorted in ascending order"; - } - return false; - } - } - return true; } diff --git a/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp b/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp index f678eb9f63..4c279314df 100644 --- a/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp +++ b/mooncake-transfer-engine/tent/src/metrics/tent_metrics.cpp @@ -16,8 +16,9 @@ #include #include -#include #include +#include +#include namespace mooncake::tent { @@ -30,7 +31,26 @@ TentMetrics::~TentMetrics() { shutdown(); } #if TENT_METRICS_ENABLED +namespace { +const char* operationName(Request::OpCode operation) { + return operation == Request::READ ? "read" : "write"; +} +} // namespace + Status TentMetrics::initialize(const MetricsConfig& config) { + // Validate configuration before touching initialized_. An invalid config + // (e.g. port 0, zero HTTP threads) would otherwise cause confusing + // failures inside initHttpServer(); fail fast with a clear error instead. + // Validating before the compare_exchange avoids a window where + // initialized_ is set to true and then rolled back on failure. + std::string error_msg; + if (!MetricsConfigLoader::validateConfig(config, &error_msg)) { + LOG(ERROR) << "Invalid TENT metrics config: " << error_msg + << "; metrics disabled"; + return Status::InvalidArgument( + "Invalid TENT metrics config: " + error_msg + LOC_MARK); + } + // Use compare_exchange to prevent race condition during initialization bool expected = false; if (!initialized_.compare_exchange_strong(expected, true)) { @@ -42,40 +62,20 @@ Status TentMetrics::initialize(const MetricsConfig& config) { // Set runtime enabled state from config runtime_enabled_.store(config_.enabled, std::memory_order_relaxed); - // Configure histogram buckets if provided (recreate histograms) - // Note: config latency_buckets are in seconds, convert to microseconds for - // histogram - if (!config_.latency_buckets.empty()) { - // Convert seconds to microseconds for histogram buckets - std::vector latency_buckets_us; - latency_buckets_us.reserve(config_.latency_buckets.size()); - for (double bucket_sec : config_.latency_buckets) { - latency_buckets_us.push_back(bucket_sec * - 1000000.0); // seconds -> microseconds - } - read_latency_ = ylt::metric::histogram_t( - "tent_read_latency_us", "Read latency distribution in microseconds", - latency_buckets_us); - write_latency_ = ylt::metric::histogram_t( - "tent_write_latency_us", - "Write latency distribution in microseconds", latency_buckets_us); - } - - // Configure size histogram buckets if provided - if (!config_.size_buckets.empty()) { - read_size_ = ylt::metric::histogram_t( - "tent_read_size_bytes", "Read request size distribution in bytes", - config_.size_buckets); - write_size_ = ylt::metric::histogram_t( - "tent_write_size_bytes", "Write request size distribution in bytes", - config_.size_buckets); - } - // Register all metrics to vectors for unified serialization registerMetrics(); - // Initialize and start HTTP server - initHttpServer(); + // Initialize and start HTTP server on the configured port. If the port is + // busy (e.g. another rank was given the same port), degrade to log-only + // metrics rather than falsely reporting a listening endpoint. + Status http_status = initHttpServer(); + const bool http_ok = http_status.ok(); + if (!http_ok) { + LOG(WARNING) << "TENT metrics HTTP endpoint unavailable on " + << config_.http_host << ":" << config_.http_port << " (" + << http_status.ToString() + << "); continuing with log-only metrics"; + } // Start periodic metric reporting thread if interval > 0 if (config_.report_interval_seconds > 0) { @@ -94,20 +94,54 @@ Status TentMetrics::initialize(const MetricsConfig& config) { }); } - LOG(INFO) - << "TENT metrics initialized successfully, HTTP server listening on " - << config_.http_host << ":" << config_.http_port - << ", runtime_enabled=" << (runtime_enabled_.load() ? "true" : "false"); + if (http_ok) { + LOG(INFO) << "TENT metrics initialized successfully, HTTP server " + "listening on " + << config_.http_host << ":" + << bound_http_port_.load(std::memory_order_relaxed) + << ", runtime_enabled=" + << (runtime_enabled_.load() ? "true" : "false"); + } else { + LOG(INFO) << "TENT metrics initialized in log-only mode (HTTP endpoint " + "disabled), runtime_enabled=" + << (runtime_enabled_.load() ? "true" : "false"); + } return Status::OK(); } -void TentMetrics::initHttpServer() { +Status TentMetrics::initHttpServer() { using namespace coro_http; - // Create HTTP server with configurable threads + // Create HTTP server with configurable threads on the configured port. + // Port assignment is intentionally deterministic: co-located ranks should + // be given distinct ports explicitly (e.g. base_port + local_rank), not + // auto-scanned, so a rank's metrics port stays predictable. http_server_ = std::make_unique( config_.http_server_threads, config_.http_port); + registerHttpHandlers(); + + // Start the HTTP server asynchronously. async_start() returns a future that + // already holds a result ONLY when startup failed (e.g. the port is already + // in use); on success the future stays pending while the server keeps + // running. Same idiom as mooncake-store's rpc_service.cpp / + // real_client.cpp. + auto ec = http_server_->async_start(); + if (ec.hasResult()) { + http_server_.reset(); + return Status::RpcServiceError( + "Failed to start TENT metrics HTTP server" LOC_MARK); + } + + // Record the bound port (read by httpPort() from other threads, so it must + // be the atomic, not config_). + bound_http_port_.store(config_.http_port, std::memory_order_relaxed); + return Status::OK(); +} + +void TentMetrics::registerHttpHandlers() { + using namespace coro_http; + // Register /metrics endpoint for Prometheus http_server_->set_http_handler( "/metrics", [this](coro_http_request& req, coro_http_response& resp) { @@ -140,9 +174,6 @@ void TentMetrics::initHttpServer() { resp.add_header("Content-Type", "text/plain"); resp.set_status_and_content(status_type::ok, "OK"); }); - - // Start the HTTP server asynchronously - http_server_->async_start(); } void TentMetrics::shutdown() { @@ -164,94 +195,177 @@ void TentMetrics::shutdown() { // Clear metric vectors counters_.clear(); histograms_.clear(); - histogram_boundaries_.clear(); + + // Reset bound port so httpPort() returns 0 after shutdown, not a stale + // port from a previous initialization. Without this, a re-initialize + // that fails to bind would cause httpPort() to report the old port. + bound_http_port_.store(0, std::memory_order_relaxed); initialized_ = false; LOG(INFO) << "TENT metrics shutdown complete"; } void TentMetrics::registerMetrics() { - // Pre-allocate vectors to avoid reallocation - counters_.reserve(7); - histograms_.reserve(4); - histogram_boundaries_.reserve(4); - - // Register all counters - add new counters here + // Register all counters as base metric_t* pointers so that counters with + // different label arities (N=1 per-transport, N=2 failover from→to and + // transport-attempt operation labels) share one vector for Prometheus + // serialize(). counters_ = { - &read_bytes_total_, &write_bytes_total_, &read_requests_total_, - &write_requests_total_, &read_failures_total_, &write_failures_total_, + &read_bytes_total_, + &write_bytes_total_, + &read_requests_total_, + &write_requests_total_, + &read_failures_total_, + &write_failures_total_, &failover_total_, + &transport_attempts_total_, + &transport_attempt_failures_total_, + &deadline_infeasible_total_, }; - // Register all histograms - add new histograms here - // Note: histogram_boundaries_ must match the order of histograms_ + // Register the N=1 per-transport histograms with their bucket boundaries + // and parallel sum counters. Each entry keeps the histogram, its + // compile-time boundaries, and its sum counter in sync so both the + // Prometheus and JSON serializers cannot mislabel buckets or drop _sum. + // The N=2 transport_attempt_latency_ histogram is serialized separately + // (it can't share this N=1-typed vector); see getPrometheusMetrics / + // getJsonMetrics. histograms_ = { - &read_latency_, - &write_latency_, - &read_size_, - &write_size_, - }; - histogram_boundaries_ = { - kLatencyBuckets, - kLatencyBuckets, - kSizeBuckets, - kSizeBuckets, + {&read_latency_, &kLatencyBuckets, &read_latency_sum_}, + {&write_latency_, &kLatencyBuckets, &write_latency_sum_}, + {&read_size_, &kSizeBuckets, &read_size_sum_}, + {&write_size_, &kSizeBuckets, &write_size_sum_}, + {&deadline_mlu_, &kMluPerMilleBuckets, &deadline_mlu_sum_}, + {&stage_queue_wait_, &kStageBuckets, &stage_queue_wait_sum_}, + {&stage_dispatch_, &kStageBuckets, &stage_dispatch_sum_}, + {&stage_transport_, &kStageBuckets, &stage_transport_sum_}, }; } -void TentMetrics::recordReadCompleted(size_t bytes, double latency_seconds) { - // Fast path: check runtime switch first +void TentMetrics::recordReadCompleted(TransportType tp, size_t bytes, + double latency_seconds) { if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) return; - read_bytes_total_.inc(static_cast(bytes)); - read_requests_total_.inc(); - read_size_.observe(static_cast(bytes)); + auto label = std::array{transportTypeName(tp)}; + read_bytes_total_.inc(label, static_cast(bytes)); + read_requests_total_.inc(label); + auto bytes_val = static_cast(bytes); + read_size_.observe(label, bytes_val); + read_size_sum_.inc(label, bytes_val); if (latency_seconds > 0.0) { - // Convert seconds to microseconds for histogram (int64_t internally) int64_t latency_us = static_cast(latency_seconds * 1000000.0); - read_latency_.observe(latency_us); + read_latency_.observe(label, latency_us); + read_latency_sum_.inc(label, latency_us); } } -void TentMetrics::recordWriteCompleted(size_t bytes, double latency_seconds) { - // Fast path: check runtime switch first +void TentMetrics::recordWriteCompleted(TransportType tp, size_t bytes, + double latency_seconds) { if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) return; - write_bytes_total_.inc(static_cast(bytes)); - write_requests_total_.inc(); - write_size_.observe(static_cast(bytes)); + auto label = std::array{transportTypeName(tp)}; + write_bytes_total_.inc(label, static_cast(bytes)); + write_requests_total_.inc(label); + auto bytes_val = static_cast(bytes); + write_size_.observe(label, bytes_val); + write_size_sum_.inc(label, bytes_val); if (latency_seconds > 0.0) { - // Convert seconds to microseconds for histogram (int64_t internally) int64_t latency_us = static_cast(latency_seconds * 1000000.0); - write_latency_.observe(latency_us); + write_latency_.observe(label, latency_us); + write_latency_sum_.inc(label, latency_us); } } -void TentMetrics::recordReadFailed(size_t bytes) { - // Fast path: check runtime switch first +void TentMetrics::recordDeadlineMLU(TransportType tp, double mlu) { + if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) + return; + if (mlu < 0.0) return; + auto label = std::array{transportTypeName(tp)}; + auto mlu_permille = static_cast(mlu * 1000.0); + deadline_mlu_.observe(label, mlu_permille); + deadline_mlu_sum_.inc(label, mlu_permille); +} + +void TentMetrics::recordDeadlineInfeasible(TransportType tp) { + if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) + return; + deadline_infeasible_total_.inc( + std::array{transportTypeName(tp)}); +} + +void TentMetrics::recordStageLatency(Stage stage, TransportType tp, + double latency_us) { if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) return; + if (latency_us < 0.0) return; + auto label = std::array{transportTypeName(tp)}; + int64_t val = static_cast(latency_us); + switch (stage) { + case Stage::QueueWait: + stage_queue_wait_.observe(label, val); + stage_queue_wait_sum_.inc(label, val); + break; + case Stage::Dispatch: + stage_dispatch_.observe(label, val); + stage_dispatch_sum_.inc(label, val); + break; + case Stage::Transport: + stage_transport_.observe(label, val); + stage_transport_sum_.inc(label, val); + break; + } +} - read_failures_total_.inc(); - read_requests_total_.inc(); // Count failed requests too +void TentMetrics::recordReadFailed(TransportType tp) { + if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) + return; + auto label = std::array{transportTypeName(tp)}; + read_failures_total_.inc(label); + read_requests_total_.inc(label); } -void TentMetrics::recordWriteFailed(size_t bytes) { - // Fast path: check runtime switch first +void TentMetrics::recordWriteFailed(TransportType tp) { if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) return; + auto label = std::array{transportTypeName(tp)}; + write_failures_total_.inc(label); + write_requests_total_.inc(label); +} - write_failures_total_.inc(); - write_requests_total_.inc(); // Count failed requests too +void TentMetrics::recordTransportFailover(TransportType from, + TransportType to) { + if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) + return; + failover_total_.inc(std::array{transportTypeName(from), + transportTypeName(to)}); } -void TentMetrics::recordTransportFailover() { +void TentMetrics::recordTransportAttemptStarted(TransportType tp, + Request::OpCode operation) { if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) return; + transport_attempts_total_.inc(std::array{ + transportTypeName(tp), operationName(operation)}); +} - failover_total_.inc(); +void TentMetrics::recordTransportAttemptFinished(TransportType tp, + Request::OpCode operation, + TransferStatusEnum status, + double latency_us) { + if (!initialized_ || !runtime_enabled_.load(std::memory_order_relaxed)) + return; + auto label = std::array{transportTypeName(tp), + operationName(operation)}; + if (status == FAILED) { + transport_attempt_failures_total_.inc(label); + } + if (latency_us >= 0.0) { + auto latency_val = static_cast(latency_us); + transport_attempt_latency_.observe(label, latency_val); + transport_attempt_latency_sum_.inc(label, latency_val); + } } std::string TentMetrics::getPrometheusMetrics() { @@ -259,18 +373,35 @@ std::string TentMetrics::getPrometheusMetrics() { try { std::string result; - // Pre-allocate buffer to avoid reallocation during serialization result.reserve(kPrometheusBufferSize); - // Serialize all counters + // Counters: ylt's counter_t::serialize() is reliable — no evidence of + // silent drops in practice. Kept as-is. for (auto* counter : counters_) { - counter->serialize(result); + std::string tmp; + counter->serialize(tmp); + result += tmp; } - // Serialize all histograms - for (auto* histogram : histograms_) { - histogram->serialize(result); + // Histograms: do NOT use ylt's basic_dynamic_histogram::serialize(). + // It silently drops the entire metric (including the # HELP / # TYPE + // header it already wrote) whenever every label combo has sum_==0 — + // via `if (value == 0) continue; ... if (value_str.empty()) + // str.clear();`. That condition is reachable in production: e.g. + // stage_queue_wait_us with sub-microsecond latencies that truncate to + // 0 under int64_t observation. The JSON endpoint's custom serializer + // walks get_bucket_counts() directly and is unaffected, which is why + // /metrics/json reported count=4846 while /metrics omitted the metric + // entirely. Using the same bucket-walk here closes that drift. + for (const auto& entry : histograms_) { + serializeHistogramPrometheus(entry.h, *entry.boundaries, *entry.sum, + result); } + // N=2 transport-attempt latency histogram (labels: transport, + // operation) goes through the same helper, instantiated for N=2. + serializeHistogramPrometheus(&transport_attempt_latency_, + kLatencyBuckets, + transport_attempt_latency_sum_, result); return result; } catch (const std::exception& e) { @@ -279,43 +410,227 @@ std::string TentMetrics::getPrometheusMetrics() { } } -std::string TentMetrics::getJsonMetrics() { - if (!initialized_) return "{}"; +namespace { +// Sum values across all label combos of a dynamic counter. Works with both +// raw pointers (counter members) and shared_ptr (histogram bucket counters). +template +int64_t sumCounterValues(CounterPtr counter) { + int64_t total = 0; + for (auto& e : counter->copy()) { + total += e->value.load(std::memory_order_relaxed); + } + return total; +} - try { - nlohmann::json root; +template +void serializeHistogramToJson( + nlohmann::json& root, + ylt::metric::basic_dynamic_histogram* hist, + const std::vector& boundaries, + ylt::metric::basic_dynamic_counter* sum) { + auto bucket_counts = hist->get_bucket_counts(); + int64_t total_count = 0; + int64_t total_sum = 0; + nlohmann::json buckets_obj; + for (size_t i = 0; i < bucket_counts.size(); ++i) { + int64_t bucket_total = sumCounterValues(bucket_counts[i]); + total_count += bucket_total; + if (i < boundaries.size()) { + buckets_obj[std::to_string(static_cast(boundaries[i]))] = + bucket_total; + } + } + // ylt keeps the histogram's sum_ private, so read the parallel sum counter + // maintained alongside each observe() call (matches the Prometheus path). + for (auto& e : sum->copy()) { + total_sum += e->value.load(); + } + nlohmann::json hist_obj; + hist_obj["count"] = total_count; + hist_obj["sum"] = total_sum; + hist_obj["buckets"] = buckets_obj; + root[hist->str_name()] = hist_obj; +} +} // namespace + +template +void TentMetrics::serializeHistogramPrometheus( + ylt::metric::basic_dynamic_histogram* hist, + const std::vector& boundaries, + ylt::metric::basic_dynamic_counter& sum, + std::string& out) const { + // Walk the same data the JSON path uses (get_bucket_counts() + copy()), + // so the two endpoints cannot drift. Unlike ylt's serialize() this never + // silently drops a histogram that has observed >=1 sample: ylt clears its + // output string (taking the # HELP / # TYPE header with it) whenever + // every label combo has sum_==0, which is reachable in production when + // sub-microsecond latencies truncate to 0 under int64_t observation. + // + // Templated over label arity N: the per-transport histograms are N=1 and + // the transport-attempt latency histogram is N=2. For N=1 the emitted + // text is byte-identical to the original single-label implementation. + auto bucket_counts = hist->get_bucket_counts(); + if (bucket_counts.empty()) return; + + const auto& label_names = hist->labels_name(); + + // A unique key per label tuple, used to dedup combos and look up the sum. + auto combo_key = [](const std::array& lv) { + std::string k; + for (uint8_t i = 0; i < N; ++i) { + k.append(lv[i]); + k.push_back('\x1f'); // separator that cannot appear in a label + } + return k; + }; + // Render `name0="v0",name1="v1"` for a label tuple (no surrounding braces). + auto append_labels = [&](std::string& dst, + const std::array& lv) { + for (uint8_t i = 0; i < N; ++i) { + if (i) dst.append(","); + dst.append(label_names[i]).append("=\"").append(lv[i]).append("\""); + } + }; - // Serialize all counters - for (auto* counter : counters_) { - root[counter->str_name()] = counter->value(); + // Build the union of label combos across ALL buckets. ylt's observe() + // increments exactly one bucket per observation (the bucket containing + // the value), so no single bucket sees every combo: e.g. queue_wait + // (sub-us -> bucket[0]) and transport_us (>=100us -> higher buckets) live + // in disjoint buckets. ylt itself uses sum_->copy() for this, but sum_ is + // private. Unioning across buckets gives the same set without needing sum_. + std::vector*> label_combos; + std::unordered_set seen; + for (auto& bc : bucket_counts) { + for (auto& e : bc->copy()) { + if (seen.insert(combo_key(e->label)).second) { + label_combos.push_back(&e->label); + } + } + } + if (label_combos.empty()) return; + + // Pre-compute per-combo total counts so we can (a) skip totally-empty + // combos and (b) decide whether to emit the # HELP / # TYPE header at + // all. This mirrors ylt's "emit head only if value_map non-empty" but + // uses bucket counts instead of sum, so a combo with sum==0 but real + // observations is still emitted. + std::vector*, int64_t>> + active_combos; + for (auto* labels_value : label_combos) { + int64_t total_count = 0; + for (auto& bc : bucket_counts) { + total_count += bc->value(*labels_value); } + if (total_count > 0) { + active_combos.emplace_back(labels_value, total_count); + } + } + if (active_combos.empty()) return; + + // Read back the per-combo sum from the parallel counter. ylt's histogram + // keeps sum_ private with no accessor, so we maintain a separate counter + // incremented alongside each observe() call. copy() returns a vector of + // {label, value} pairs; build a lookup map for O(1) access per combo. + std::unordered_map sum_by_combo; + for (auto& e : sum.copy()) { + sum_by_combo[combo_key(e->label)] = e->value.load(); + } - // Serialize all histograms - for (size_t h = 0; h < histograms_.size(); ++h) { - auto* histogram = histograms_[h]; - const auto& boundaries = histogram_boundaries_[h]; + const std::string& name = hist->str_name(); + const std::string help_str{hist->help()}; + + // Emit the header once per metric (matches ylt's serialize_head()). + out.append("# HELP ").append(name).append(" ").append(help_str).append( + "\n"); + out.append("# TYPE ").append(name).append(" histogram\n"); + + for (auto& [labels_value, total_count] : active_combos) { + int64_t cumulative = 0; + for (size_t i = 0; i < bucket_counts.size(); ++i) { + cumulative += bucket_counts[i]->value(*labels_value); + out.append(name).append("_bucket{"); + append_labels(out, *labels_value); + out.append(","); + if (i < boundaries.size()) { + out.append("le=\"") + .append(std::to_string(boundaries[i])) + .append("\"} "); + } else { + out.append("le=\"+Inf\"} "); + } + out.append(std::to_string(cumulative)).append("\n"); + } - auto bucket_counts = histogram->get_bucket_counts(); + // _sum: read from the parallel counter maintained alongside each + // observe() call. Falls back to 0 if the combo is not yet in the + // counter (should not happen for active combos, but defensive). + int64_t total_sum = 0; + auto it = sum_by_combo.find(combo_key(*labels_value)); + if (it != sum_by_combo.end()) { + total_sum = it->second; + } + out.append(name).append("_sum{"); + append_labels(out, *labels_value); + out.append("} ").append(std::to_string(total_sum)).append("\n"); - // Calculate total count - int64_t total_count = 0; - for (auto& bucket : bucket_counts) { - total_count += bucket->value(); - } + out.append(name).append("_count{"); + append_labels(out, *labels_value); + out.append("} ").append(std::to_string(total_count)).append("\n"); + } +} - nlohmann::json hist_obj; - hist_obj["count"] = total_count; +std::string TentMetrics::getJsonMetrics() { + if (!initialized_) return "{}"; - nlohmann::json buckets_obj; - for (size_t i = 0; - i < boundaries.size() && i < bucket_counts.size(); ++i) { - buckets_obj[std::to_string(static_cast( - boundaries[i]))] = bucket_counts[i]->value(); - } - hist_obj["buckets"] = buckets_obj; + try { + nlohmann::json root; - root[histogram->str_name()] = hist_obj; - } + // Counters: aggregate (sum) across all transport label values so the + // JSON endpoint stays a simple flat {name: total} view. Per-transport + // breakdown is available via the Prometheus endpoint. + root[read_bytes_total_.str_name()] = + sumCounterValues(&read_bytes_total_); + root[write_bytes_total_.str_name()] = + sumCounterValues(&write_bytes_total_); + root[read_requests_total_.str_name()] = + sumCounterValues(&read_requests_total_); + root[write_requests_total_.str_name()] = + sumCounterValues(&write_requests_total_); + root[read_failures_total_.str_name()] = + sumCounterValues(&read_failures_total_); + root[write_failures_total_.str_name()] = + sumCounterValues(&write_failures_total_); + root[failover_total_.str_name()] = sumCounterValues(&failover_total_); + root[transport_attempts_total_.str_name()] = + sumCounterValues(&transport_attempts_total_); + root[transport_attempt_failures_total_.str_name()] = + sumCounterValues(&transport_attempt_failures_total_); + root[deadline_infeasible_total_.str_name()] = + sumCounterValues(&deadline_infeasible_total_); + + // Histograms: sum bucket counts across all transport labels. The + // templated helper also reads back the parallel sum counter so the + // JSON endpoint emits "sum" alongside "count" (and stays in sync with + // the Prometheus endpoint). + serializeHistogramToJson(root, &read_latency_, kLatencyBuckets, + &read_latency_sum_); + serializeHistogramToJson(root, &write_latency_, kLatencyBuckets, + &write_latency_sum_); + serializeHistogramToJson(root, &read_size_, kSizeBuckets, + &read_size_sum_); + serializeHistogramToJson(root, &write_size_, kSizeBuckets, + &write_size_sum_); + serializeHistogramToJson(root, &deadline_mlu_, kMluPerMilleBuckets, + &deadline_mlu_sum_); + serializeHistogramToJson(root, &stage_queue_wait_, kStageBuckets, + &stage_queue_wait_sum_); + serializeHistogramToJson(root, &stage_dispatch_, kStageBuckets, + &stage_dispatch_sum_); + serializeHistogramToJson(root, &stage_transport_, kStageBuckets, + &stage_transport_sum_); + serializeHistogramToJson(root, &transport_attempt_latency_, + kLatencyBuckets, + &transport_attempt_latency_sum_); return root.dump(2); // Pretty print with 2-space indent } catch (const std::exception& e) { @@ -330,13 +645,16 @@ std::string TentMetrics::getSummaryString() { std::ostringstream oss; oss << std::fixed << std::setprecision(2); - double read_bytes = read_bytes_total_.value(); - double write_bytes = write_bytes_total_.value(); - double read_reqs = read_requests_total_.value(); - double write_reqs = write_requests_total_.value(); - double read_fails = read_failures_total_.value(); - double write_fails = write_failures_total_.value(); - double failovers = failover_total_.value(); + // Aggregate across all transport labels — summary is intentionally a + // single total line, not per-transport. Per-transport breakdown is via + // Prometheus. + double read_bytes = sumCounterValues(&read_bytes_total_); + double write_bytes = sumCounterValues(&write_bytes_total_); + double read_reqs = sumCounterValues(&read_requests_total_); + double write_reqs = sumCounterValues(&write_requests_total_); + double read_fails = sumCounterValues(&read_failures_total_); + double write_fails = sumCounterValues(&write_failures_total_); + double failovers = sumCounterValues(&failover_total_); // Format bytes in human-readable form auto formatBytes = [](double bytes) -> std::string { @@ -379,11 +697,18 @@ Status TentMetrics::initialize(const MetricsConfig& config) { void TentMetrics::shutdown() { initialized_ = false; } -void TentMetrics::recordReadCompleted(size_t, double) {} -void TentMetrics::recordWriteCompleted(size_t, double) {} -void TentMetrics::recordReadFailed(size_t) {} -void TentMetrics::recordWriteFailed(size_t) {} -void TentMetrics::recordTransportFailover() {} +void TentMetrics::recordReadCompleted(TransportType, size_t, double) {} +void TentMetrics::recordWriteCompleted(TransportType, size_t, double) {} +void TentMetrics::recordReadFailed(TransportType) {} +void TentMetrics::recordWriteFailed(TransportType) {} +void TentMetrics::recordTransportFailover(TransportType, TransportType) {} +void TentMetrics::recordTransportAttemptStarted(TransportType, + Request::OpCode) {} +void TentMetrics::recordTransportAttemptFinished(TransportType, Request::OpCode, + TransferStatusEnum, double) {} +void TentMetrics::recordDeadlineMLU(TransportType, double) {} +void TentMetrics::recordDeadlineInfeasible(TransportType) {} +void TentMetrics::recordStageLatency(Stage, TransportType, double) {} std::string TentMetrics::getPrometheusMetrics() { return "# TENT metrics disabled at compile time\n"; diff --git a/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt b/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt index 2f48b68696..11b7bbd9ac 100644 --- a/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/platform/CMakeLists.txt @@ -2,6 +2,7 @@ add_subdirectory(cuda) add_subdirectory(rocm) add_subdirectory(ascend) add_subdirectory(sunrise) +add_subdirectory(tpu) file(GLOB PLATFORM_SOURCES "*.cpp") add_library(tent_platform_all STATIC ${PLATFORM_SOURCES}) @@ -18,3 +19,6 @@ endif() if(TARGET platform_sunrise) target_link_libraries(tent_platform_all PUBLIC platform_sunrise) endif() +if(TARGET platform_tpu) + target_link_libraries(tent_platform_all PUBLIC platform_tpu) +endif() diff --git a/mooncake-transfer-engine/tent/src/platform/ascend/ascend_probe.cpp b/mooncake-transfer-engine/tent/src/platform/ascend/ascend_probe.cpp index e954bbcd5d..416690ba5c 100644 --- a/mooncake-transfer-engine/tent/src/platform/ascend/ascend_probe.cpp +++ b/mooncake-transfer-engine/tent/src/platform/ascend/ascend_probe.cpp @@ -67,7 +67,7 @@ static std::vector listInfiniBandDevices() { std::ifstream(path) >> numa_node; devices.push_back( - Topology::NicEntry{.name = std::move(device_name), + Topology::NicEntry{.name = device_name, .pci_bus_id = std::move(pci_bus_id), .type = Topology::NIC_RDMA, .numa_node = numa_node}); diff --git a/mooncake-transfer-engine/tent/src/platform/cpu_probe.cpp b/mooncake-transfer-engine/tent/src/platform/cpu_probe.cpp index 5e16c12997..3b97bef163 100644 --- a/mooncake-transfer-engine/tent/src/platform/cpu_probe.cpp +++ b/mooncake-transfer-engine/tent/src/platform/cpu_probe.cpp @@ -173,7 +173,7 @@ static std::vector listInfiniBandDevices() { std::ifstream(path) >> numa_node; devices.push_back( - Topology::NicEntry{.name = std::move(device_name), + Topology::NicEntry{.name = device_name, .pci_bus_id = std::move(pci_bus_id), .type = Topology::NIC_RDMA, .numa_node = numa_node}); diff --git a/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp b/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp index 09345e982b..6a52601c9a 100644 --- a/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp +++ b/mooncake-transfer-engine/tent/src/platform/cuda/cuda_probe.cpp @@ -69,7 +69,7 @@ static std::vector listInfiniBandDevices() { std::ifstream(path) >> numa_node; devices.push_back( - Topology::NicEntry{.name = std::move(device_name), + Topology::NicEntry{.name = device_name, .pci_bus_id = std::move(pci_bus_id), .type = Topology::NIC_RDMA, .numa_node = numa_node}); @@ -262,10 +262,45 @@ Status CudaPlatform::probe(std::vector& nic_list, return Status::OK(); } +namespace { +bool cudaDevicePresent() { + static const bool present = [] { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || + device_count == 0) { + LOG(WARNING) << "No CUDA device detected; treating buffers as " + "host memory"; + return false; + } + return true; + }(); + return present; +} + +bool cudaAbiMatches() { + static const bool matches = [] { + int runtime_version = 0; + // Major version only: struct layout changes across CUDA majors. + if (cudaRuntimeGetVersion(&runtime_version) == cudaSuccess && + runtime_version / 1000 != CUDART_VERSION / 1000) { + LOG(ERROR) << "CUDA ABI mismatch: built against CUDART " + << CUDART_VERSION << ", loaded libcudart is " + << runtime_version + << "; skipping pointer probe, rebuild against a " + "matching CUDA toolkit."; + return false; + } + return true; + }(); + return matches; +} +} // namespace + MemoryType CudaPlatform::getMemoryType(void* addr) { - cudaPointerAttributes attributes; - cudaError_t result; - result = cudaPointerGetAttributes(&attributes, addr); + if (!cudaDevicePresent()) return MTYPE_CPU; + if (!cudaAbiMatches()) return MTYPE_UNKNOWN; + cudaPointerAttributes attributes{}; + cudaError_t result = cudaPointerGetAttributes(&attributes, addr); if (result != cudaSuccess) { LOG(WARNING) << "cudaPointerGetAttributes: " << cudaGetErrorString(result); @@ -296,21 +331,24 @@ const std::vector CudaPlatform::getLocation(void* start, const static size_t kPageSize = 4096; std::vector entries; - cudaPointerAttributes attributes; - cudaError_t result; - - result = cudaPointerGetAttributes(&attributes, start); - if (result != cudaSuccess) { - LOG(WARNING) << "cudaPointerGetAttributes: " - << cudaGetErrorString(result); + if (cudaDevicePresent() && !cudaAbiMatches()) { entries.push_back({(uint64_t)start, len, kWildcardLocation}); return entries; } - - if (attributes.type == cudaMemoryTypeDevice) { - entries.push_back( - {(uint64_t)start, len, genCudaNodeName(attributes.device)}); - return entries; + if (cudaDevicePresent()) { + cudaPointerAttributes attributes{}; + cudaError_t result = cudaPointerGetAttributes(&attributes, start); + if (result != cudaSuccess) { + LOG(WARNING) << "cudaPointerGetAttributes: " + << cudaGetErrorString(result); + entries.push_back({(uint64_t)start, len, kWildcardLocation}); + return entries; + } + if (attributes.type == cudaMemoryTypeDevice) { + entries.push_back( + {(uint64_t)start, len, genCudaNodeName(attributes.device)}); + return entries; + } } // start and end address may not be page aligned. diff --git a/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp index f1a2918b66..2e5bc762db 100644 --- a/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp +++ b/mooncake-transfer-engine/tent/src/platform/rocm/rocm_probe.cpp @@ -65,7 +65,7 @@ static std::vector listInfiniBandDevices() { std::ifstream(path) >> numa_node; devices.push_back( - Topology::NicEntry{.name = std::move(device_name), + Topology::NicEntry{.name = device_name, .pci_bus_id = std::move(pci_bus_id), .type = Topology::NIC_RDMA, .numa_node = numa_node}); diff --git a/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_probe.cpp b/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_probe.cpp index db9e4cce61..41ff87dab5 100644 --- a/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_probe.cpp +++ b/mooncake-transfer-engine/tent/src/platform/sunrise/sunrise_probe.cpp @@ -109,7 +109,7 @@ static std::vector listInfiniBandDevices() { std::ifstream(path) >> numa_node; devices.push_back( - Topology::NicEntry{.name = std::move(device_name), + Topology::NicEntry{.name = device_name, .pci_bus_id = std::move(pci_bus_id), .type = Topology::NIC_RDMA, .numa_node = numa_node}); diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/CMakeLists.txt b/mooncake-transfer-engine/tent/src/platform/tpu/CMakeLists.txt new file mode 100644 index 0000000000..c59d55d480 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/CMakeLists.txt @@ -0,0 +1,7 @@ +if(USE_TPU) + file(GLOB TENT_PLATFORM_TPU_SOURCES "*.cpp") + add_library(platform_tpu STATIC ${TENT_PLATFORM_TPU_SOURCES}) + # No PJRT / TPU SDK link dependency: the adapter is resolved at runtime via + # dlopen() in TpuPjrtShim, so we only need libdl. + target_link_libraries(platform_tpu PUBLIC tent_common ${CMAKE_DL_LIBS}) +endif() diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/README.md b/mooncake-transfer-engine/tent/src/platform/tpu/README.md new file mode 100644 index 0000000000..1cd1f068ff --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/README.md @@ -0,0 +1,85 @@ +# TPU (PJRT) platform for TENT + +This directory implements Google TPU support for the TENT transfer engine. +It is compiled only when TENT is built with `-DUSE_TPU=ON` (OFF by default). + +## Why staging + +TPU HBM is not addressable by the NIC, so a transfer that touches TPU memory +cannot be issued directly by RDMA/TCP. Instead every such transfer is **staged +through host DRAM** and the two hops are chained by the existing +`ProxyManager` pipeline: + +``` + TPU HBM <-> host DRAM (PJRT device copy, this platform) + host DRAM <-> remote host DRAM (RDMA / TCP, existing transports) + remote host DRAM <-> remote TPU HBM (PJRT device copy on the peer) +``` + +No new networked transport is required — TENT reuses `ProxyManager`'s chunked +double-buffering, staging-buffer lifecycle, and async status tracking. + +> New to the codebase and wondering why this needs a `TpuTransport` at all +> instead of a one-line branch in `ProxyManager`? See +> [`STAGING_ARCHITECTURE.md`](STAGING_ARCHITECTURE.md). + +## Components + +| File | Role | +|------|------| +| `tpu_platform.cpp` (`TpuPlatform`) | Derives `CpuPlatform`. Host DRAM + NIC topology are inherited (RDMA NICs if present, otherwise just host NUMA nodes for TCP); only the TPU-device-aware paths are overridden (`copy`, `getMemoryType`, `getLocation`, and one `MemEntry` per device in `probe`). | +| `tpu_pjrt_shim.cpp` (`TpuPjrtShim`) | Isolates all PJRT dependency. Resolves the device-copy adapter at runtime via `dlopen`. | +| `../../transport/tpu/tpu_transport.cpp` (`TpuTransport`) | The local HBM↔host staging executor. Advertises only `gpu_to_dram` / `dram_to_gpu`; runs the copy via `Platform::copy` for `LOCAL_SEGMENT_ID` requests. | + +The staging policy that ties these together lives in +`TransferEngineImpl::findStagingPolicy` (case "TPU"). + +## The device-copy adapter (runtime dependency) + +`TpuPjrtShim` does **not** link the PJRT runtime. Instead it loads an adapter +shared library at runtime that exports the C ABI declared in +[`tpu_pjrt_abi.h`](../../../include/tent/platform/tpu_pjrt_abi.h): + +```c +int mc_tpu_pjrt_init(void); +int mc_tpu_pjrt_is_device_ptr(const void *addr); +int mc_tpu_pjrt_device_index(const void *addr); +int mc_tpu_pjrt_copy_d2h(void *host_dst, const void *device_src, size_t len); +int mc_tpu_pjrt_copy_h2d(void *device_dst, const void *host_src, size_t len); +int mc_tpu_pjrt_device_count(void); +int mc_tpu_pjrt_device_numa(int index); +``` + +- **Discovery:** the library path defaults to `libmooncake_tpu_pjrt.so` and can + be overridden with the `MC_TPU_PJRT_LIB` environment variable. +- **Graceful absence:** if the adapter cannot be loaded or does not satisfy the + ABI, `TpuPjrtShim::available()` returns false and all TPU operations return a + non-OK `Status`. A `USE_TPU` build therefore links and runs without the + runtime present (useful for CI and unit tests). +- **Pointer tokens:** the `const void *` "device pointer" is the stable token the + serving-engine integration registers with TENT for a TPU buffer (via a + `tpu:N` location). The adapter owns the mapping from that token to the + underlying PJRT buffer. The token is **not** the buffer's data and must never + be dereferenced — on real PJRT it is an internal handle that is + host-readable but reads as unrelated bytes. +- **Interior pointers:** `ProxyManager` stages a transfer in `chunk_size` + (4 MiB) pieces and passes `token + chunk_offset` for every chunk after the + first. Classification and copy entrypoints must therefore resolve an address + to the registered buffer whose range contains it. An adapter that only matches + base addresses makes TENT classify HBM as host memory, and — because the token + *is* readable — the staging copy degrades into a `memcpy` of unrelated bytes + that reports success. `TpuTransport` defends against this by requiring exactly + one TPU-device side per staging hop and failing loudly otherwise. + +A mock adapter and unit tests that exercise this ABI on any Linux host live in +[`../../../tests/tpu/`](../../../tests/tpu/). The mock hands out poisoned tokens +backed by shadow storage, so a copy that bypasses the adapter yields `0xDD` +rather than accidentally producing the right answer. + +## Not yet included (follow-ups) + +- Serving-framework (JAX / PyTorch-XLA) integration that registers TPU buffers. +- DMA-mapped (pinned) staging buffers for true async device DMA, and the + interaction between device DMA-mapping and RDMA memory registration on the + same host buffer. +- Benchmarks on real TPU hardware. diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/STAGING_ARCHITECTURE.md b/mooncake-transfer-engine/tent/src/platform/tpu/STAGING_ARCHITECTURE.md new file mode 100644 index 0000000000..119909b9de --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/STAGING_ARCHITECTURE.md @@ -0,0 +1,171 @@ +# Why TPU support needs a `TpuTransport` (and not just an `if (pjrt)` branch) + +This note explains, for someone new to the TENT codebase, **why adding TPU +staging touches a transport class, a platform class, and the routing/capability +system** rather than being a one-line change inside `ProxyManager`. Everything +here is about *existing* TENT mechanics; TPU is just the motivating example. + +All paths below are under `mooncake-transfer-engine/tent/`. + +--- + +## 1. The three layers you need to know + +TENT separates "how do I move bytes" into three cooperating layers: + +| Layer | What it is | Example | +|-------|-----------|---------| +| **Platform** | The *local* memory primitives for one accelerator family: allocate, free, `copy`, classify a pointer's memory type, discover topology. One process picks exactly one Platform at build time. | `CpuPlatform`, `CudaPlatform`, `TpuPlatform` | +| **Transport** | A *data channel* that moves a transfer request's bytes. Each advertises **capabilities** (can it do dram→dram? gpu→dram? gpu→gpu?). | `RdmaTransport`, `ShmTransport`, `NvlinkTransport`, `TpuTransport` | +| **Staging (ProxyManager)** | When no single transport can do a hop directly, it splits the transfer into stages through host DRAM and chains transports. | `ProxyManager` | + +Key headers: `include/tent/runtime/platform.h`, `include/tent/runtime/transport.h` +(the `Capabilities` struct is at the top of the latter). + +--- + +## 2. The naive expectation + +> "TPU HBM can't be reached by the NIC, so stage it through host DRAM. TENT +> already stages CUDA that way, so just find where `ProxyManager` does the +> device copy and add an `if (tpu) pjrt_copy() else cudaMemcpy()`." + +That mental model is *almost* right about the data flow but wrong about the +mechanism. There is no `cudaMemcpy` inside `ProxyManager` to branch on. + +--- + +## 3. What actually happens when you submit a transfer + +Follow one transfer from the top: + +1. **Submit** → `TransferEngineImpl::submitTransfer` → `prepareSubmit` + (`src/runtime/transfer_engine_impl.cpp`). + +2. **Should this be staged?** For each request, `prepareSubmit` calls + [`findStagingPolicy`](../../runtime/transfer_engine_impl.cpp) (defined at + `transfer_engine_impl.cpp:1278`). It returns a 3-element plan + `[server, local_stage_location, remote_stage_location]`; empty entries mean + "no staging on that side". Then: + ```cpp + owner.staging = !owner.staging_params.empty() && staging_proxy_; // ~line 1401 + ``` + +3. **If staged**, the task is handed to the proxy instead of a transport: + ```cpp + staging_proxy_->submit(&task, batch, owner.staging_params); // ~line 1477 + ``` + +4. **`ProxyManager` chops the transfer into chunks** and, per chunk, issues + *sub-transfers* for each stage — it does **not** copy bytes itself. See + `ProxyManager::transferEventLoop` (`src/runtime/proxy_manager.cpp:251`) and + `submitLocalStage` (`:71`). The local stage is submitted as an ordinary + transfer whose target is `LOCAL_SEGMENT_ID`: + ```cpp + // local_stage.source = device pointer, target_offset = host staging buffer + impl_->submitStagingTransfer(batch, {local_stage}); // -> submitTransfer(...) + ``` + +5. **That sub-transfer is routed like any other**, back through + `getTransportType`, and the *selected transport's* `submitTransferTasks` + runs. For the simple copy transports, that method is where the actual byte + movement lives — and it calls **`Platform::copy`**, not any device API + directly. Example, `ShmTransport::startTransfer` + (`src/transport/shm/shm_transport.cpp:119`): + ```cpp + status = Platform::getLoader().copy(dst, src, length); // :122 / :126 + ``` + +So the real device-copy seam is **`Platform::copy`** +(`CudaPlatform::copy` does `cudaMemcpyAsync`; `TpuPlatform::copy` will call the +PJRT adapter). That part *is* a clean override — good. + +**But there is a gate before you ever reach step 5.** + +--- + +## 4. The crux: routing is gated by transport *capabilities* + +In step 5 the sub-transfer must be *routed to some transport*. Routing asks each +candidate transport: "can you do a copy from `local_memory_type` to +`remote_memory_type`?" via a capability check: + +- Selector mode (**the default**): `TransportSelector::isTransportAvailable` + (`src/runtime/transport_selector.cpp:311`). +- Legacy mode: `checkAvailability` (`src/runtime/transfer_engine_impl.cpp:889`). + +Both map a `(local, remote)` memory-type pair to one capability bit: + +```cpp +// GPU/device local side, CPU peer -> needs caps.gpu_to_dram +// CPU local side, GPU/device peer -> needs caps.dram_to_gpu +// GPU to GPU -> needs caps.gpu_to_gpu +// CPU to CPU -> needs caps.dram_to_dram +``` + +For the TPU **local stage**, the sub-transfer is `TPU device -> host DRAM`, so +routing looks for a transport advertising `gpu_to_dram`. + +Now look at who advertises what: + +| Transport | Advertises | Set where | +|-----------|-----------|-----------| +| `RdmaTransport` | `dram_to_dram`; adds `gpu_*` **only if `nvidia_peermem` is loaded** (CUDA GPUDirect) | `rdma_transport.cpp` `install` | +| `ShmTransport` | `dram_to_dram` **only** | `shm_transport.cpp` `install` | +| `NvlinkTransport` | `dram_to_gpu` / `gpu_to_dram` / `gpu_to_gpu` (CUDA) | `nvlink_transport.cpp` `install` | + +On a TPU host there is **no transport that advertises `gpu_to_dram`**: +- `ShmTransport` = `dram_to_dram` only. +- `RdmaTransport` won't set `gpu_to_dram` (no `nvidia_peermem`), and it *couldn't* + DMA out of HBM anyway. + +**Therefore, without a new transport, the TPU local-stage sub-transfer has no +route, and it fails — even though `TpuPlatform::copy` exists and could do the +copy.** `findStagingPolicy` would also have nothing to hang the policy on. + +This is the answer to "why not just one branch": the copy *mechanism* is a clean +`Platform::copy` override, but the *routing/capability system* has no way to +select it unless some `Transport` object sits in `transport_list_[...]` and +answers "yes, I can do `gpu_to_dram`." + +For CUDA, that answering transport is `NvlinkTransport` — a full P2P transport +that TENT reuses. TPU has no equivalent to piggyback on, so we add a **thin one +whose only job is to advertise the device↔host capability and run +`Platform::copy` for `LOCAL_SEGMENT_ID` requests**: `TpuTransport` +(`src/transport/tpu/tpu_transport.cpp`). It never touches the network; the +host↔host hop is still RDMA/TCP. + +--- + +## 5. What PR #1 therefore has to touch + +| Concern | Change | File | +|---------|--------|------| +| A memory type for TPU HBM | `MTYPE_TPU` + `"tpu"` parsing | `platform.h`, `getTypeEnum` in `transfer_engine_impl.cpp` | +| Local device copy primitive | `TpuPlatform::copy` → PJRT adapter (the clean seam) | `src/platform/tpu/tpu_platform.cpp` | +| A transport that *advertises* `gpu_to_dram`/`dram_to_gpu` so routing can pick the copy | thin `TpuTransport` | `src/transport/tpu/tpu_transport.cpp` | +| Teach the capability checks that TPU is a device type | add `MTYPE_TPU` to `isGpuType` **and** the selector's `is_gpu` lambda | `transfer_engine_impl.cpp:880`, `transport_selector.cpp:337` | +| Decide to stage TPU transfers | `findStagingPolicy` TPU case | `transfer_engine_impl.cpp:1278` | + +Note the **two** capability helpers (line 880 and 337): there are two routing +modes (selector = default, legacy), each with its own device-type predicate. +Both must learn about `MTYPE_TPU`, or TPU works in one mode and silently fails in +the other. This is exactly the kind of thing that isn't visible until you trace +both paths. + +--- + +## 6. One-paragraph summary + +TENT's `ProxyManager` already implements the staged-DRAM pipeline, and +`Platform::copy` is a clean place to plug in a PJRT device copy. But a staged +sub-transfer is still *routed*, and routing only selects a transport that +*advertises* the matching capability. No existing transport advertises +`device↔host` on a TPU host, so the copy would never be reached. `TpuTransport` +exists solely to answer that capability query and dispatch to `Platform::copy`; +it is the unavoidable "interface tax" of participating in TENT's transport +system, and it accounts for most of the non-test, non-doc line count in this +change. + +See also: [`README.md`](README.md) in this directory for the component list and +the PJRT adapter ABI. diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/tpu_pjrt_shim.cpp b/mooncake-transfer-engine/tent/src/platform/tpu/tpu_pjrt_shim.cpp new file mode 100644 index 0000000000..164dc13bae --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/tpu_pjrt_shim.cpp @@ -0,0 +1,144 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/platform/tpu_pjrt_shim.h" + +#include +#include + +#include + +namespace mooncake { +namespace tent { + +namespace { +constexpr const char kDefaultAdapterLib[] = "libmooncake_tpu_pjrt.so"; + +const char *adapterLibraryPath() { + const char *env = std::getenv("MC_TPU_PJRT_LIB"); + if (env && env[0] != '\0') return env; + return kDefaultAdapterLib; +} +} // namespace + +TpuPjrtShim &TpuPjrtShim::instance() { + static TpuPjrtShim g_instance; + return g_instance; +} + +TpuPjrtShim::TpuPjrtShim() { load(); } + +TpuPjrtShim::~TpuPjrtShim() { + if (handle_) dlclose(handle_); +} + +void TpuPjrtShim::load() { + const char *lib = adapterLibraryPath(); + handle_ = dlopen(lib, RTLD_NOW | RTLD_LOCAL); + if (!handle_) { + LOG(WARNING) << "TpuPjrtShim: unable to load TPU PJRT adapter '" << lib + << "': " << dlerror() + << ". TPU transfers will be unavailable. Set " + "MC_TPU_PJRT_LIB to override the adapter path."; + available_ = false; + return; + } + + // Resolve every entrypoint; treat any missing symbol as a fatal load error + // so we never partially bind an incompatible adapter. + auto resolve = [&](const char *name) -> void * { + void *sym = dlsym(handle_, name); + if (!sym) + LOG(WARNING) << "TpuPjrtShim: adapter '" << lib + << "' is missing symbol '" << name << "'"; + return sym; + }; + + fn_init_ = reinterpret_cast(resolve("mc_tpu_pjrt_init")); + fn_is_device_ptr_ = reinterpret_cast( + resolve("mc_tpu_pjrt_is_device_ptr")); + fn_device_index_ = reinterpret_cast( + resolve("mc_tpu_pjrt_device_index")); + fn_copy_d2h_ = reinterpret_cast( + resolve("mc_tpu_pjrt_copy_d2h")); + fn_copy_h2d_ = reinterpret_cast( + resolve("mc_tpu_pjrt_copy_h2d")); + fn_device_count_ = + reinterpret_cast(resolve("mc_tpu_pjrt_device_count")); + fn_device_numa_ = + reinterpret_cast(resolve("mc_tpu_pjrt_device_numa")); + + if (!fn_init_ || !fn_is_device_ptr_ || !fn_device_index_ || !fn_copy_d2h_ || + !fn_copy_h2d_ || !fn_device_count_ || !fn_device_numa_) { + LOG(ERROR) << "TpuPjrtShim: adapter '" << lib + << "' does not satisfy the required ABI; disabling TPU."; + dlclose(handle_); + handle_ = nullptr; + available_ = false; + return; + } + + if (fn_init_() != 0) { + LOG(ERROR) << "TpuPjrtShim: mc_tpu_pjrt_init() failed; disabling TPU."; + dlclose(handle_); + handle_ = nullptr; + available_ = false; + return; + } + + available_ = true; + LOG(INFO) << "TpuPjrtShim: TPU PJRT adapter '" << lib << "' loaded (" + << fn_device_count_() << " device(s))."; +} + +bool TpuPjrtShim::isDevicePtr(const void *addr) const { + if (!available_ || !addr) return false; + return fn_is_device_ptr_(addr) != 0; +} + +int TpuPjrtShim::deviceIndex(const void *addr) const { + if (!available_ || !addr) return -1; + return fn_device_index_(addr); +} + +Status TpuPjrtShim::copyD2H(void *host_dst, const void *device_src, + size_t length) const { + if (!available_) + return Status::NotImplemented("TPU PJRT adapter unavailable" LOC_MARK); + if (fn_copy_d2h_(host_dst, device_src, length) != 0) + return Status::InternalError("TPU device->host copy failed" LOC_MARK); + return Status::OK(); +} + +Status TpuPjrtShim::copyH2D(void *device_dst, const void *host_src, + size_t length) const { + if (!available_) + return Status::NotImplemented("TPU PJRT adapter unavailable" LOC_MARK); + if (fn_copy_h2d_(device_dst, host_src, length) != 0) + return Status::InternalError("TPU host->device copy failed" LOC_MARK); + return Status::OK(); +} + +int TpuPjrtShim::deviceCount() const { + if (!available_) return 0; + return fn_device_count_(); +} + +int TpuPjrtShim::deviceNumaNode(int index) const { + if (!available_) return -1; + return fn_device_numa_(index); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/platform/tpu/tpu_platform.cpp b/mooncake-transfer-engine/tent/src/platform/tpu/tpu_platform.cpp new file mode 100644 index 0000000000..625ea8edd1 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/platform/tpu/tpu_platform.cpp @@ -0,0 +1,103 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/platform/tpu.h" + +#include + +#include + +#include "tent/common/status.h" +#include "tent/platform/tpu_pjrt_shim.h" + +namespace mooncake { +namespace tent { + +Status TpuPlatform::probe(std::vector& nic_list, + std::vector& mem_list) { + // Host DRAM + RDMA topology is identical to a CPU host. + CHECK_STATUS(CpuPlatform::probe(nic_list, mem_list)); + + // Register one memory node per visible TPU device so findNearMem("tpu:N") + // resolves to the nearest host NUMA node for staging. Device I/O never + // touches the NIC directly, so these entries only carry NIC affinity used + // to place the host staging buffers. + auto& shim = TpuPjrtShim::instance(); + int device_count = shim.deviceCount(); + for (int i = 0; i < device_count; ++i) { + Topology::MemEntry entry; + entry.name = "tpu:" + std::to_string(i); + entry.numa_node = shim.deviceNumaNode(i); + entry.type = Topology::MEM_UNKNOWN; + int nic_id = 0; + for (const auto& nic : nic_list) { + if (entry.numa_node >= 0 && nic.numa_node == entry.numa_node) + entry.device_list[0].push_back(nic_id); + else + entry.device_list[2].push_back(nic_id); + nic_id++; + } + mem_list.push_back(std::move(entry)); + } + return Status::OK(); +} + +Status TpuPlatform::allocate(void** pptr, size_t size, MemoryOptions& options) { + LocationParser location(options.location); + if (location.type() == "tpu") { + // TPU HBM buffers are owned by the serving framework and registered + // with TENT; TENT never allocates them itself. + return Status::NotImplemented( + "TpuPlatform does not allocate TPU device memory" LOC_MARK); + } + // Host DRAM staging buffers use the inherited NUMA-aware allocator. + return CpuPlatform::allocate(pptr, size, options); +} + +Status TpuPlatform::copy(void* dst, void* src, size_t length) { + auto& shim = TpuPjrtShim::instance(); + bool src_is_device = shim.isDevicePtr(src); + bool dst_is_device = shim.isDevicePtr(dst); + if (src_is_device && dst_is_device) { + return Status::NotImplemented( + "TpuPlatform: device-to-device copy is not supported; transfers " + "are staged through host DRAM" LOC_MARK); + } + if (src_is_device) return shim.copyD2H(dst, src, length); + if (dst_is_device) return shim.copyH2D(dst, src, length); + // Neither side is TPU memory: plain host copy. + return CpuPlatform::copy(dst, src, length); +} + +MemoryType TpuPlatform::getMemoryType(void* addr) { + if (TpuPjrtShim::instance().isDevicePtr(addr)) return MTYPE_TPU; + return CpuPlatform::getMemoryType(addr); +} + +const std::vector TpuPlatform::getLocation(void* start, + size_t len, + bool skip_prefault) { + auto& shim = TpuPjrtShim::instance(); + if (shim.isDevicePtr(start)) { + int index = shim.deviceIndex(start); + std::string location = + index >= 0 ? "tpu:" + std::to_string(index) : kWildcardLocation; + return {RangeLocation{reinterpret_cast(start), len, + std::move(location)}}; + } + return CpuPlatform::getLocation(start, len, skip_prefault); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/python/CMakeLists.txt b/mooncake-transfer-engine/tent/src/python/CMakeLists.txt index 5ef9b10155..5d8ea378b0 100644 --- a/mooncake-transfer-engine/tent/src/python/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/python/CMakeLists.txt @@ -23,8 +23,7 @@ set_target_properties(tentpy PROPERTIES OUTPUT_NAME "tent" INSTALL_RPATH find_package( Python3 - COMPONENTS Interpreter Development + COMPONENTS Interpreter Development.Module REQUIRED) -target_link_libraries(tentpy PRIVATE ${Python3_LIBRARIES}) target_include_directories(tentpy PRIVATE ${Python3_INCLUDE_DIRS}) target_link_libraries(tentpy PUBLIC tent_link_group) diff --git a/mooncake-transfer-engine/tent/src/python/pybind.cpp b/mooncake-transfer-engine/tent/src/python/pybind.cpp index bc1d3a23b7..6c81c90aff 100644 --- a/mooncake-transfer-engine/tent/src/python/pybind.cpp +++ b/mooncake-transfer-engine/tent/src/python/pybind.cpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include #include @@ -299,6 +300,18 @@ PYBIND11_MODULE(tent, m) { .value("TCP", TransportType::TCP) .value("AscendDirect", TransportType::AscendDirect) .value("SUNRISE_LINK", TransportType::SUNRISE_LINK) + .value("TPU", TransportType::TPU) + .value("UB", TransportType::UB) + .export_values(); + + py::enum_(m, "IntentType") + .value("INTENT_UNSPEC", IntentType::INTENT_UNSPEC) + .value("FOREGROUND_GET", IntentType::FOREGROUND_GET) + .value("BACKGROUND_PREFETCH", IntentType::BACKGROUND_PREFETCH) + .value("MIGRATION", IntentType::MIGRATION) + .value("CHECKPOINT", IntentType::CHECKPOINT) + .value("WEIGHT_LOADING", IntentType::WEIGHT_LOADING) + .value("STAGING_INTERNAL", IntentType::STAGING_INTERNAL) .export_values(); py::enum_(m, "SegmentInfoType") @@ -321,7 +334,9 @@ PYBIND11_MODULE(tent, m) { .def(py::init([](Request::OpCode opcode, uint64_t source, uint64_t target_id, uint64_t target_offset, size_t length, int priority, - TransportType transport_hint) { + TransportType transport_hint, + std::optional policy_name, + uint64_t deadline_ns, IntentType intent_type) { Request r; r.opcode = opcode; r.source = U64ToPtr(source); @@ -330,12 +345,17 @@ PYBIND11_MODULE(tent, m) { r.length = length; r.priority = priority; r.transport_hint = transport_hint; + r.policy_name = std::move(policy_name); + r.deadline_ns = deadline_ns; + r.intent_type = intent_type; return r; }), py::arg("opcode"), py::arg("source"), py::arg("target_id"), py::arg("target_offset"), py::arg("length"), py::arg("priority") = PRIO_HIGH, - py::arg("transport_hint") = TransportType::UNSPEC) + py::arg("transport_hint") = TransportType::UNSPEC, + py::arg("policy_name") = std::nullopt, py::arg("deadline_ns") = 0, + py::arg("intent_type") = IntentType::INTENT_UNSPEC) .def_property( "opcode", [](const Request& r) { return r.opcode; }, [](Request& r, Request::OpCode op) { r.opcode = op; }) @@ -346,7 +366,10 @@ PYBIND11_MODULE(tent, m) { .def_readwrite("target_offset", &Request::target_offset) .def_readwrite("length", &Request::length) .def_readwrite("priority", &Request::priority) - .def_readwrite("transport_hint", &Request::transport_hint); + .def_readwrite("transport_hint", &Request::transport_hint) + .def_readwrite("policy_name", &Request::policy_name) + .def_readwrite("deadline_ns", &Request::deadline_ns) + .def_readwrite("intent_type", &Request::intent_type); py::class_(m, "TransferStatus") .def(py::init<>()) @@ -682,6 +705,15 @@ PYBIND11_MODULE(tent, m) { py::arg("batch_id"), py::arg("request_list"), py::arg("name"), py::arg("message")) + .def( + "cancel_transfer", + [](TransferEngine& self, uint64_t batch_id, size_t task_id) { + py::gil_scoped_release release; + auto s = self.cancelTransfer((BatchID)batch_id, task_id); + ThrowStatus(s, "cancel_transfer"); + }, + py::arg("batch_id"), py::arg("task_id")) + // --------------------------------------------------------------------- // notification send/receive // --------------------------------------------------------------------- diff --git a/mooncake-transfer-engine/tent/src/rpc/rpc.cpp b/mooncake-transfer-engine/tent/src/rpc/rpc.cpp index 9091543cef..e1721188ae 100644 --- a/mooncake-transfer-engine/tent/src/rpc/rpc.cpp +++ b/mooncake-transfer-engine/tent/src/rpc/rpc.cpp @@ -16,6 +16,7 @@ #include #include +#include "transfer_engine_rpc_client_io_context.h" #include "tent/common/utils/ip.h" #include "tent/common/utils/random.h" @@ -179,7 +180,8 @@ Lazy> CoroRpcAgent::callCoroutine( ClientLease lease{pool->acquire(), pool, false}; if (!lease.client) { - lease.client = std::make_unique(); + lease.client = std::make_unique( + GetTransferEngineRpcClientIoContextPool().get_executor()); auto conn_result = co_await lease.client->connect(server_addr); if (conn_result.val() != 0) { lease.broken = true; diff --git a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp index d759b878a8..6c552dea07 100644 --- a/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/admission_queue.cpp @@ -14,8 +14,11 @@ #include "tent/runtime/admission_queue.h" +#include +#include #include #include +#include namespace mooncake { namespace tent { @@ -23,8 +26,18 @@ namespace { using PublicTaskKey = std::pair; +// Sort key for EDF: owners without a deadline (0) sort after all deadlined +// owners, so they never jump ahead of a real deadline. +inline uint64_t deadlineKey(uint64_t deadline_ns) { + return deadline_ns == 0 ? std::numeric_limits::max() + : deadline_ns; +} + bool isSupportedTerminalStatus(TransferStatusEnum status) { return status == TransferStatusEnum::COMPLETED || + status == TransferStatusEnum::INVALID || + status == TransferStatusEnum::CANCELED || + status == TransferStatusEnum::TIMEOUT || status == TransferStatusEnum::FAILED; } @@ -164,6 +177,7 @@ Status LocalTransferAdmissionQueue::tryAdmit( owner.batch_token = submit.batch_token; owner.request = owner_input.request; owner.kind = owner_input.kind; + owner.degradation_eligible = owner_input.degradation_eligible; owners_.emplace(owner_id, owner); public_to_owner_[{submit.batch_token, owner_input.owner_task_id}] = @@ -171,7 +185,27 @@ Status LocalTransferAdmissionQueue::tryAdmit( for (const auto derived_task_id : owner_input.derived_task_ids) { public_to_owner_[{submit.batch_token, derived_task_id}] = owner_id; } - fifo_.push_back(owner_id); + // RFC #2519 step 2: keep fifo_ ordered on admission so pickForDispatch + // never has to re-sort. Default (deadline_aware == false) appends in + // strict FIFO. When deadline-aware, insert at the earliest-deadline- + // first position; upper_bound places a new owner *after* existing + // owners with the same deadline, preserving FIFO order among ties. + if (limits_.deadline_aware) { + const uint64_t key = deadlineKey(owner.request.deadline_ns); + auto pos = std::upper_bound( + fifo_.begin(), fifo_.end(), key, + [this](uint64_t k, QueueOwnerId id) { + auto it = owners_.find(id); + uint64_t d = + (it == owners_.end()) + ? std::numeric_limits::max() + : deadlineKey(it->second.request.deadline_ns); + return k < d; + }); + fifo_.insert(pos, owner_id); + } else { + fifo_.push_back(owner_id); + } admitted_owner_ids.push_back(owner_id); } @@ -182,11 +216,93 @@ Status LocalTransferAdmissionQueue::tryAdmit( return Status::OK(); } +void LocalTransferAdmissionQueue::setDegradationPolicy( + BandwidthProvider bandwidth_provider, DegradationHooks hooks, + NowProvider now_provider) { + bandwidth_provider_ = std::move(bandwidth_provider); + degradation_hooks_ = std::move(hooks); + now_provider_ = std::move(now_provider); +} + std::vector LocalTransferAdmissionQueue::pickForDispatch( - size_t max_owners, size_t max_bytes) { + size_t max_owners, size_t max_bytes, + std::vector* dropped_owner_ids) { + if (dropped_owner_ids) dropped_owner_ids->clear(); std::vector picked; if (max_owners == 0 || max_bytes == 0) return picked; + // RFC #2519 step 2 (opt-in): earliest-deadline-first dispatch. When + // deadline_aware, fifo_ is kept EDF-ordered at admission time (see + // tryAdmit's ordered insert), so there is nothing to sort here — we just + // consume from the front. This keeps the hot dispatch path O(picked) + // instead of re-sorting the whole queue on every call. Default + // (deadline_aware == false) is plain FIFO. + // + // RFC #2519 step 3 (opt-in): drop is active only when a positive threshold, + // deadline awareness, and a bandwidth provider are all present. + const bool drop_enabled = limits_.deadline_aware && + limits_.mlu_local_threshold > 0.0 && + static_cast(bandwidth_provider_); + const bool promotion_enabled = + limits_.deadline_aware && limits_.promotion_slack_ns > 0; + const bool need_now = drop_enabled || promotion_enabled; + const double bw_bps = drop_enabled ? bandwidth_provider_() : 0.0; + const uint64_t now_ns = + need_now + ? (now_provider_ + ? now_provider_() + : static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now() + .time_since_epoch()) + .count())) + : 0; + + // Deadline proximity promotion: partition fifo_ so owners with critical + // slack (deadline approaching within promotion_slack_ns) appear before + // owners with comfortable slack or no deadline. stable_partition preserves + // relative EDF order within each group. + if (promotion_enabled) { + std::stable_partition(fifo_.begin(), fifo_.end(), [&](QueueOwnerId id) { + auto it = owners_.find(id); + if (it == owners_.end() || it->second.state != QueueState::Queued) { + return false; + } + const uint64_t dl = it->second.request.deadline_ns; + if (dl == 0 || dl <= now_ns) return false; + return (dl - now_ns) < limits_.promotion_slack_ns; + }); + } + + // Predicted MLU = predicted_transfer_time / remaining_window. Returns true + // if the owner is predicted to miss its deadline hard enough to drop. + auto shouldDrop = [&](const QueueOwner& owner) -> bool { + if (!drop_enabled || !owner.degradation_eligible || bw_bps <= 0.0) + return false; + const uint64_t deadline_ns = owner.request.deadline_ns; + if (deadline_ns == 0) return false; // no deadline + if (deadline_ns <= now_ns) return true; // already past + const double window_s = (deadline_ns - now_ns) / 1e9; + const double predicted_time_s = owner.request.length / bw_bps; + const double mlu = predicted_time_s / window_s; + return mlu >= limits_.mlu_local_threshold; + }; + + auto dropOwner = [&](QueueOwnerId owner_id, QueueOwner& owner) { + owner.state = QueueState::Terminal; + owner.terminal_status = TransferStatusEnum::CANCELED; + --outstanding_owners_; + outstanding_bytes_ -= owner.request.length; + if (owner.kind == QueueOwnerKind::User) { + --outstanding_user_owners_; + outstanding_user_bytes_ -= owner.request.length; + } + if (dropped_owner_ids) dropped_owner_ids->push_back(owner_id); + if (degradation_hooks_.on_local_decode_suggested) { + degradation_hooks_.on_local_decode_suggested(owner.request); + } + }; + size_t used_owners = 0; size_t used_bytes = 0; while (!fifo_.empty() && used_owners < max_owners) { @@ -201,6 +317,16 @@ std::vector LocalTransferAdmissionQueue::pickForDispatch( continue; } + // Step 3: an owner predicted to miss its deadline is dropped (not + // dispatched) and does not consume the dispatch budget. Because the + // queue is EDF-ordered, later owners have looser deadlines, so we keep + // scanning rather than stopping. + if (shouldDrop(owner_it->second)) { + fifo_.pop_front(); + dropOwner(owner_id, owner_it->second); + continue; + } + const auto& owner = owner_it->second; const size_t remaining_bytes = max_bytes - used_bytes; if (owner.request.length > remaining_bytes) break; @@ -232,9 +358,39 @@ Status LocalTransferAdmissionQueue::complete( return Status::InvalidEntry("queue owner is not dispatching" LOC_MARK); } - owner.state = terminal_status == TransferStatusEnum::COMPLETED - ? QueueState::Completed - : QueueState::Failed; + owner.state = QueueState::Terminal; + owner.terminal_status = terminal_status; + --outstanding_owners_; + outstanding_bytes_ -= owner.request.length; + if (owner.kind == QueueOwnerKind::User) { + --outstanding_user_owners_; + outstanding_user_bytes_ -= owner.request.length; + } + return Status::OK(); +} + +Status LocalTransferAdmissionQueue::cancel(QueueOwnerId owner_id) { + if (owner_id == 0) { + return Status::InvalidArgument("invalid queue owner id" LOC_MARK); + } + auto owner_it = owners_.find(owner_id); + if (owner_it == owners_.end()) { + return Status::InvalidEntry("queue owner not found" LOC_MARK); + } + auto& owner = owner_it->second; + if (owner.state == QueueState::Terminal) { + return owner.terminal_status == TransferStatusEnum::CANCELED + ? Status::OK() + : Status::InvalidEntry( + "queue owner is already terminal" LOC_MARK); + } + if (owner.state != QueueState::Queued) { + return Status::InvalidEntry( + "queue owner is already dispatching" LOC_MARK); + } + + owner.state = QueueState::Terminal; + owner.terminal_status = TransferStatusEnum::CANCELED; --outstanding_owners_; outstanding_bytes_ -= owner.request.length; if (owner.kind == QueueOwnerKind::User) { @@ -271,9 +427,7 @@ Status LocalTransferAdmissionQueue::retireBatch(uint64_t batch_token) { return Status::InternalError( "queue owner batch token mismatch" LOC_MARK); } - const bool terminal = owner.state == QueueState::Completed || - owner.state == QueueState::Failed; - if (!terminal) { + if (owner.state != QueueState::Terminal) { return Status::InvalidEntry( "batch has non-terminal queue owners" LOC_MARK); } @@ -314,11 +468,8 @@ Status LocalTransferAdmissionQueue::getPublicStatus( case QueueState::Dispatching: status = TransferStatusEnum::PENDING; break; - case QueueState::Completed: - status = TransferStatusEnum::COMPLETED; - break; - case QueueState::Failed: - status = TransferStatusEnum::FAILED; + case QueueState::Terminal: + status = owner_it->second.terminal_status; break; } return Status::OK(); diff --git a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp index d82fbb5410..ca5facb0e6 100644 --- a/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/control_plane.cpp @@ -249,7 +249,7 @@ void ControlService::onSendData(const std::string_view& request, return; } XferDataDesc* desc = (XferDataDesc*)request.data(); - auto local_desc = manager_->getLocal().get(); + auto local_desc = manager_->getLocal(); auto peer_mem_addr = le64toh(desc->peer_mem_addr); auto length = le64toh(desc->length); @@ -273,7 +273,7 @@ void ControlService::onRecvData(const std::string_view& request, return; } XferDataDesc* desc = (XferDataDesc*)request.data(); - auto local_desc = manager_->getLocal().get(); + auto local_desc = manager_->getLocal(); auto peer_mem_addr = le64toh(desc->peer_mem_addr); auto length = le64toh(desc->length); diff --git a/mooncake-transfer-engine/tent/src/runtime/platform.cpp b/mooncake-transfer-engine/tent/src/runtime/platform.cpp index 72a9d096d9..611737f532 100644 --- a/mooncake-transfer-engine/tent/src/runtime/platform.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/platform.cpp @@ -22,6 +22,8 @@ #include "tent/platform/sunrise.h" #elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) #include "tent/platform/ascend.h" +#elif defined(USE_TPU) +#include "tent/platform/tpu.h" #else #include "tent/platform/cpu.h" #endif @@ -41,6 +43,8 @@ Platform& Platform::getLoader(std::shared_ptr conf) { g_instance = std::make_shared(conf); #elif defined(USE_ASCEND) || defined(USE_ASCEND_DIRECT) g_instance = std::make_shared(conf); +#elif defined(USE_TPU) + g_instance = std::make_shared(conf); #else g_instance = std::make_shared(conf); #endif diff --git a/mooncake-transfer-engine/tent/src/runtime/progress_worker.cpp b/mooncake-transfer-engine/tent/src/runtime/progress_worker.cpp index bb52b7afd3..49a6e470e5 100644 --- a/mooncake-transfer-engine/tent/src/runtime/progress_worker.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/progress_worker.cpp @@ -14,13 +14,17 @@ #include "tent/runtime/progress_worker.h" +#include + #include "tent/common/status.h" #include "tent/runtime/transfer_engine_impl.h" namespace mooncake { namespace tent { -ProgressWorker::ProgressWorker(TransferEngineImpl* impl) : impl_(impl) {} +ProgressWorker::ProgressWorker(TransferEngineImpl* impl, + std::chrono::microseconds fallback_interval) + : impl_(impl), fallback_interval_(fallback_interval) {} ProgressWorker::~ProgressWorker() { stop(); } @@ -37,6 +41,7 @@ void ProgressWorker::stop() { // user thread's freeBatch path. order_.clear(); queued_.clear(); + queue_ready_ = false; } cv_.notify_all(); if (thread_.joinable()) thread_.join(); @@ -53,27 +58,63 @@ void ProgressWorker::notifyBatchMaybeReady(BatchID batch_id) { cv_.notify_one(); } +void ProgressWorker::notifyRuntimeQueueReady() { + if (!running_.load(std::memory_order_acquire)) return; + { + std::lock_guard lk(mu_); + queue_ready_ = true; + } + cv_.notify_one(); +} + void ProgressWorker::runner() { while (true) { BatchID batch_id = 0; + bool queue_ready = false; + bool fallback_due = false; + bool queue_active = impl_->hasActiveRuntimeQueue(); { std::unique_lock lk(mu_); - cv_.wait(lk, [&] { - return !running_.load(std::memory_order_acquire) || - !order_.empty(); - }); + if (queue_active && fallback_interval_.count() > 0) { + const bool woke = cv_.wait_for(lk, fallback_interval_, [&] { + return !running_.load(std::memory_order_acquire) || + queue_ready_ || !order_.empty(); + }); + fallback_due = !woke; + } else { + cv_.wait(lk, [&] { + return !running_.load(std::memory_order_acquire) || + queue_ready_ || !order_.empty(); + }); + } if (!running_.load(std::memory_order_acquire)) return; - batch_id = order_.front(); - order_.pop_front(); - queued_.erase(batch_id); + queue_ready = queue_ready_; + queue_ready_ = false; + if (!order_.empty()) { + batch_id = order_.front(); + order_.pop_front(); + queued_.erase(batch_id); + } + } + if (queue_ready) { + (void)impl_->progressRuntimeQueue(); + (void)impl_->lazyFreeBatch(); } // progressBatch acquires the engine's progress_mutex_ and silently // returns InvalidArgument if the batch was freed before we got here. // PENDING means "kick again later"; the next notify wakes us up. - // Terminal states leave the batch alone — freeBatch on the user - // thread is responsible for reclamation. - TransferStatus s; - (void)impl_->progressBatch(batch_id, s); + // Terminal states are observed here; lazyFreeBatch reclaims a batch + // only if freeBatch has already marked it free_requested. + if (batch_id) { + TransferStatus s; + (void)impl_->progressBatch(batch_id, s); + (void)impl_->progressRuntimeQueue(); + (void)impl_->lazyFreeBatch(); + } + if (fallback_due && !queue_ready && !batch_id && queue_active) { + (void)impl_->progressRuntimeQueue(); + (void)impl_->lazyFreeBatch(); + } } } diff --git a/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp b/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp index f5aa8019b8..53a7953757 100644 --- a/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp @@ -58,7 +58,13 @@ BatchID ProxyManager::submitCrossStage(const Request& request, inter_stage.target_id = request.target_id; inter_stage.target_offset = remote_stage_buffer; auto batch = impl_->allocateBatch(1); - impl_->submitTransfer(batch, {inter_stage}); + auto status = impl_->submitStagingTransfer(batch, {inter_stage}); + if (!status.ok()) { + LOG(WARNING) << "failed to submit cross-stage transfer: " + << status.ToString(); + if (batch) impl_->freeBatch(batch); + return 0; + } return batch; } @@ -72,7 +78,13 @@ BatchID ProxyManager::submitLocalStage(const Request& request, local_stage.target_id = LOCAL_SEGMENT_ID; local_stage.target_offset = local_stage_buffer; auto batch = impl_->allocateBatch(1); - impl_->submitTransfer(batch, {local_stage}); + auto status = impl_->submitStagingTransfer(batch, {local_stage}); + if (!status.ok()) { + LOG(WARNING) << "failed to submit local-stage transfer: " + << status.ToString(); + if (batch) impl_->freeBatch(batch); + return 0; + } return batch; } @@ -81,6 +93,7 @@ Status ProxyManager::waitLocalStage(const Request& request, uint64_t chunk_length, uint64_t offset) { auto batch = submitLocalStage(request, local_stage_buffer, chunk_length, offset); + if (!batch) return Status::TooManyRequests("submit local stage failed"); return impl_->waitTransferCompletion(batch); } @@ -119,13 +132,15 @@ Status ProxyManager::waitCrossStage(const Request& request, uint64_t chunk_length) { auto batch = submitCrossStage(request, local_stage_buffer, remote_stage_buffer, chunk_length); + if (!batch) return Status::TooManyRequests("submit cross stage failed"); return impl_->waitTransferCompletion(batch); } -Status ProxyManager::submit(TaskInfo* task, +Status ProxyManager::submit(TaskInfo* task, BatchID batch, const std::vector& params) { StagingTask staging_task; staging_task.native = task; + staging_task.batch = batch; staging_task.params = params; task->staging_status = PENDING; static std::atomic next_queue_index(0); @@ -228,6 +243,7 @@ void ProxyManager::runner(size_t id) { auto staging_status = status.ok() ? COMPLETED : FAILED; __atomic_store(&task.native->staging_status, &staging_status, __ATOMIC_RELEASE); + impl_->notifyBatchMaybeReady(task.batch); } cache.reset(); } @@ -317,6 +333,11 @@ Status ProxyManager::transferEventLoop(StagingTask& task, local_locked.insert(chunk.local_buf); chunk.batch = submitLocalStage(request, chunk.local_buf, chunk.length, chunk.offset); + if (!chunk.batch) { + chunk.state = StageState::FAILED; + event_queue.push(id); + break; + } chunk.prev_state = chunk.state; chunk.state = StageState::INFLIGHT; event_queue.push(id); @@ -356,6 +377,11 @@ Status ProxyManager::transferEventLoop(StagingTask& task, } chunk.batch = submitCrossStage(request, chunk.local_buf, chunk.remote_buf, chunk.length); + if (!chunk.batch) { + chunk.state = StageState::FAILED; + event_queue.push(id); + break; + } chunk.prev_state = chunk.state; chunk.state = StageState::INFLIGHT; event_queue.push(id); @@ -373,6 +399,11 @@ Status ProxyManager::transferEventLoop(StagingTask& task, } else if (request.opcode == Request::READ && local_staging) { chunk.batch = submitLocalStage(request, chunk.local_buf, chunk.length, chunk.offset); + if (!chunk.batch) { + chunk.state = StageState::FAILED; + event_queue.push(id); + break; + } chunk.prev_state = chunk.state; chunk.state = StageState::INFLIGHT; event_queue.push(id); @@ -589,4 +620,4 @@ Status ProxyManager::unpinStageBuffer(uint64_t addr) { } } // namespace tent -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/runtime/qos_contract.cpp b/mooncake-transfer-engine/tent/src/runtime/qos_contract.cpp new file mode 100644 index 0000000000..cc67aabb87 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/runtime/qos_contract.cpp @@ -0,0 +1,433 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/runtime/qos_contract.h" + +#include +#include +#include + +namespace mooncake { +namespace tent { +namespace { + +bool isValidPriority(int priority) { + return priority == PRIO_HIGH || priority == PRIO_MEDIUM || + priority == PRIO_LOW; +} + +bool isKnownDegradedAction(const std::string& action) { + static const std::vector kActions = { + "delay", "fallback_transport", "local_recompute", "compress", "reject"}; + return std::find(kActions.begin(), kActions.end(), action) != + kActions.end(); +} + +std::string trim(const std::string& input) { + size_t begin = 0; + while (begin < input.size() && + std::isspace(static_cast(input[begin]))) { + ++begin; + } + size_t end = input.size(); + while (end > begin && + std::isspace(static_cast(input[end - 1]))) { + --end; + } + return input.substr(begin, end - begin); +} + +bool containsKey(const std::vector& allowed, + const std::string& key) { + return std::find(allowed.begin(), allowed.end(), key) != allowed.end(); +} + +Status validateKnownKeys(const json& node, + const std::vector& allowed, + const std::string& path) { + for (auto it = node.begin(); it != node.end(); ++it) { + if (!containsKey(allowed, it.key())) { + return Status::InvalidArgument( + path + " contains unknown key: " + it.key()); + } + } + return Status::OK(); +} + +} // namespace + +std::string QosContractResolver::normalizeKey(const std::string& value) { + std::string out = trim(value); + std::transform(out.begin(), out.end(), out.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return out; +} + +std::string QosContractResolver::intentTypeName(IntentType intent) { + switch (intent) { + case IntentType::INTENT_UNSPEC: + return "unspec"; + case IntentType::FOREGROUND_GET: + return "foreground_get"; + case IntentType::BACKGROUND_PREFETCH: + return "background_prefetch"; + case IntentType::MIGRATION: + return "migration"; + case IntentType::CHECKPOINT: + return "checkpoint"; + case IntentType::WEIGHT_LOADING: + return "weight_loading"; + case IntentType::STAGING_INTERNAL: + return "staging_internal"; + } + return "unspec"; +} + +Status QosContractResolver::parsePriority(const json& node, int* out, + const std::string& path) { + if (!out) return Status::InvalidArgument("priority output is null"); + if (node.is_number_integer()) { + int value = node.get(); + if (!isValidPriority(value)) { + return Status::InvalidArgument(path + " priority out of range"); + } + *out = value; + return Status::OK(); + } + if (!node.is_string()) { + return Status::InvalidArgument(path + + " priority must be string or int"); + } + auto value = normalizeKey(node.get()); + if (value == "high" || value == "0") { + *out = PRIO_HIGH; + } else if (value == "medium" || value == "1") { + *out = PRIO_MEDIUM; + } else if (value == "low" || value == "2") { + *out = PRIO_LOW; + } else { + return Status::InvalidArgument(path + " priority is unknown: " + value); + } + return Status::OK(); +} + +Status QosContractResolver::parseBytes(const json& node, uint64_t* out, + const std::string& path) { + if (!out) return Status::InvalidArgument("bytes output is null"); + if (node.is_number_unsigned()) { + *out = node.get(); + return Status::OK(); + } + if (node.is_number_integer()) { + auto value = node.get(); + if (value < 0) return Status::InvalidArgument(path + " must be >= 0"); + *out = static_cast(value); + return Status::OK(); + } + if (!node.is_string()) { + return Status::InvalidArgument(path + " must be bytes string or int"); + } + + std::string text = normalizeKey(node.get()); + if (text.empty()) return Status::InvalidArgument(path + " is empty"); + + size_t pos = 0; + while (pos < text.size() && + std::isdigit(static_cast(text[pos]))) { + ++pos; + } + if (pos == 0) return Status::InvalidArgument(path + " missing number"); + + uint64_t base = 0; + try { + base = std::stoull(text.substr(0, pos)); + } catch (const std::exception&) { + return Status::InvalidArgument(path + " invalid number"); + } + + std::string unit = trim(text.substr(pos)); + uint64_t mul = 1; + if (unit.empty() || unit == "b") { + mul = 1; + } else if (unit == "kib" || unit == "kb") { + mul = 1024ull; + } else if (unit == "mib" || unit == "mb") { + mul = 1024ull * 1024ull; + } else if (unit == "gib" || unit == "gb") { + mul = 1024ull * 1024ull * 1024ull; + } else if (unit == "tib" || unit == "tb") { + mul = 1024ull * 1024ull * 1024ull * 1024ull; + } else { + return Status::InvalidArgument(path + " unsupported unit: " + unit); + } + if (base > std::numeric_limits::max() / mul) { + return Status::InvalidArgument(path + " overflows uint64"); + } + *out = base * mul; + return Status::OK(); +} + +Status QosContractResolver::parseUint64(const json& node, uint64_t* out, + const std::string& path) { + if (!out) return Status::InvalidArgument("uint64 output is null"); + if (node.is_number_unsigned()) { + *out = node.get(); + return Status::OK(); + } + if (node.is_number_integer()) { + auto value = node.get(); + if (value < 0) return Status::InvalidArgument(path + " must be >= 0"); + *out = static_cast(value); + return Status::OK(); + } + return Status::InvalidArgument(path + " must be unsigned integer"); +} + +Status QosContractResolver::parsePolicyFields(const json& node, + QosPolicyFields* out, + const std::string& path) { + if (!out) return Status::InvalidArgument("policy fields output is null"); + if (!node.is_object()) { + return Status::InvalidArgument(path + " must be an object"); + } + static const std::vector kAllowedPolicyKeys = { + "priority", "max_inflight_bytes", "max_inflight_requests", + "allowed_degraded_actions"}; + CHECK_STATUS(validateKnownKeys(node, kAllowedPolicyKeys, path)); + + if (node.contains("priority")) { + int priority = PRIO_LOW; + CHECK_STATUS( + parsePriority(node["priority"], &priority, path + ".priority")); + out->priority = priority; + } + + auto parse_bytes_field = [&](const char* key, + std::optional* dst) -> Status { + if (!node.contains(key)) return Status::OK(); + uint64_t bytes = 0; + CHECK_STATUS(parseBytes(node[key], &bytes, path + "." + key)); + *dst = bytes; + return Status::OK(); + }; + auto parse_uint64_field = [&](const char* key, + std::optional* dst) -> Status { + if (!node.contains(key)) return Status::OK(); + uint64_t value = 0; + CHECK_STATUS(parseUint64(node[key], &value, path + "." + key)); + *dst = value; + return Status::OK(); + }; + CHECK_STATUS( + parse_bytes_field("max_inflight_bytes", &out->max_inflight_bytes)); + CHECK_STATUS(parse_uint64_field("max_inflight_requests", + &out->max_inflight_requests)); + + if (node.contains("allowed_degraded_actions")) { + if (!node["allowed_degraded_actions"].is_array()) { + return Status::InvalidArgument( + path + ".allowed_degraded_actions must be array"); + } + std::vector actions; + for (const auto& action_node : node["allowed_degraded_actions"]) { + if (!action_node.is_string()) { + return Status::InvalidArgument( + path + ".allowed_degraded_actions entry must be string"); + } + auto action = normalizeKey(action_node.get()); + if (!isKnownDegradedAction(action)) { + return Status::InvalidArgument( + path + " unknown degraded action: " + action); + } + actions.push_back(action); + } + out->allowed_degraded_actions = std::move(actions); + } + + return Status::OK(); +} + +void QosContractResolver::mergeFields(QosPolicyFields* dst, + const QosPolicyFields& src) { + if (src.priority) dst->priority = src.priority; + if (src.max_inflight_bytes) + dst->max_inflight_bytes = src.max_inflight_bytes; + if (src.max_inflight_requests) + dst->max_inflight_requests = src.max_inflight_requests; + if (src.allowed_degraded_actions) { + dst->allowed_degraded_actions = src.allowed_degraded_actions; + } +} + +Status QosContractResolver::loadFromConfig(const Config& config) { + enabled_ = false; + global_defaults_ = QosPolicyFields{}; + tenants_.clear(); + + try { + json qos = config.get("qos", json{}); + if (qos.is_null() || qos.empty()) return Status::OK(); + if (!qos.is_object()) { + return Status::InvalidArgument("qos must be an object"); + } + static const std::vector kAllowedQosKeys = { + "version", "defaults", "tenants"}; + CHECK_STATUS(validateKnownKeys(qos, kAllowedQosKeys, "qos")); + + if (qos.contains("version") && !qos["version"].is_number_integer()) { + return Status::InvalidArgument("qos.version must be an integer"); + } + int version = qos.value("version", 1); + if (version != 1) { + return Status::InvalidArgument("unsupported qos.version: " + + std::to_string(version)); + } + if (qos.contains("defaults")) { + CHECK_STATUS(parsePolicyFields(qos["defaults"], &global_defaults_, + "qos.defaults")); + } + + if (qos.contains("tenants")) { + if (!qos["tenants"].is_array()) { + return Status::InvalidArgument("qos.tenants must be an array"); + } + static const std::vector kAllowedTenantKeys = { + "name", "defaults", "intents"}; + for (size_t i = 0; i < qos["tenants"].size(); ++i) { + const auto& tenant_node = qos["tenants"][i]; + const std::string path = + "qos.tenants[" + std::to_string(i) + "]"; + if (!tenant_node.is_object()) { + return Status::InvalidArgument(path + " must be an object"); + } + CHECK_STATUS( + validateKnownKeys(tenant_node, kAllowedTenantKeys, path)); + if (!tenant_node.contains("name") || + !tenant_node["name"].is_string()) { + return Status::InvalidArgument(path + ".name is required"); + } + auto tenant_name = + normalizeKey(tenant_node["name"].get()); + if (tenant_name.empty()) { + return Status::InvalidArgument(path + ".name is empty"); + } + if (tenants_.contains(tenant_name)) { + return Status::InvalidArgument("duplicate qos tenant: " + + tenant_name); + } + + TenantContract tenant; + if (tenant_node.contains("defaults")) { + CHECK_STATUS(parsePolicyFields(tenant_node["defaults"], + &tenant.defaults, + path + ".defaults")); + } + if (tenant_node.contains("intents")) { + if (!tenant_node["intents"].is_object()) { + return Status::InvalidArgument( + path + ".intents must be an object"); + } + for (auto it = tenant_node["intents"].begin(); + it != tenant_node["intents"].end(); ++it) { + QosPolicyFields fields; + const auto intent = normalizeKey(it.key()); + CHECK_STATUS( + parsePolicyFields(it.value(), &fields, + path + ".intents." + it.key())); + tenant.intents[intent] = std::move(fields); + } + } + tenants_[tenant_name] = std::move(tenant); + } + } + + enabled_ = qos.contains("defaults") || qos.contains("tenants"); + return Status::OK(); + } catch (const std::exception& e) { + enabled_ = false; + global_defaults_ = QosPolicyFields{}; + tenants_.clear(); + return Status::InvalidArgument( + std::string("failed to parse qos config: ") + e.what()); + } +} + +Status QosContractResolver::resolve(const QosRequestContext& context, + EffectiveQosPolicy* out) const { + if (!out) return Status::InvalidArgument("effective qos output is null"); + *out = EffectiveQosPolicy{}; + out->tenant_id = + normalizeKey(context.tenant_id.empty() ? "default" : context.tenant_id); + out->intent = + normalizeKey(context.intent.empty() ? "unspec" : context.intent); + out->requested_priority = context.requested_priority; + out->effective_priority = isValidPriority(context.requested_priority) + ? context.requested_priority + : PRIO_LOW; + + if (!enabled_) return Status::OK(); + + out->enabled = true; + QosPolicyFields fields = global_defaults_; + out->matched_contract = "global_default"; + + auto tenant_it = tenants_.find(out->tenant_id); + if (tenant_it != tenants_.end()) { + mergeFields(&fields, tenant_it->second.defaults); + out->matched_contract = out->tenant_id + ".default"; + auto intent_it = tenant_it->second.intents.find(out->intent); + if (intent_it != tenant_it->second.intents.end()) { + mergeFields(&fields, intent_it->second); + out->matched = true; + out->matched_contract = out->tenant_id + "." + out->intent; + } + } + + static_cast(*out) = fields; + if (out->priority) out->effective_priority = *out->priority; + return Status::OK(); +} + +void QosContractResolver::fieldsToJson(json* out, + const QosPolicyFields& fields) { + if (fields.priority) (*out)["priority"] = *fields.priority; + if (fields.max_inflight_bytes) { + (*out)["max_inflight_bytes"] = *fields.max_inflight_bytes; + } + if (fields.max_inflight_requests) { + (*out)["max_inflight_requests"] = *fields.max_inflight_requests; + } + if (fields.allowed_degraded_actions) { + (*out)["allowed_degraded_actions"] = *fields.allowed_degraded_actions; + } +} + +std::string QosContractResolver::explainJson(const EffectiveQosPolicy& policy, + int indent) const { + json out; + out["enabled"] = policy.enabled; + out["matched"] = policy.matched; + out["tenant_id"] = policy.tenant_id; + out["intent"] = policy.intent; + out["matched_contract"] = policy.matched_contract; + out["requested_priority"] = policy.requested_priority; + out["effective_priority"] = policy.effective_priority; + out["diagnostic_scope"] = "resolution_only"; + fieldsToJson(&out, policy); + return out.dump(indent); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp b/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp new file mode 100644 index 0000000000..5ae6c14a2e --- /dev/null +++ b/mooncake-transfer-engine/tent/src/runtime/receiver_credit.cpp @@ -0,0 +1,179 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit.h" + +namespace mooncake::tent { + +size_t CreditKeyHash::operator()(const CreditKey& k) const noexcept { + size_t h = std::hash{}(k.receiver_session.high); + auto mix = [&h](uint64_t v) { + h ^= std::hash{}(v) + 0x9e3779b97f4a7c15ULL + (h << 6) + + (h >> 2); + }; + mix(k.receiver_session.low); + mix(k.sender_peer); + mix(k.qos_class); + return h; +} + +Status SenderCreditLedger::resourceIndex(CreditResource r, size_t& i) { + auto raw = static_cast(r); + if (raw < 1 || raw > kCreditResourceCount) + return Status::InvalidArgument("unknown credit resource" LOC_MARK); + i = raw - 1; + return Status::OK(); +} + +Status SenderCreditLedger::normalize( + const CreditCharge& c, std::array& out) { + out.fill(0); + if (c.resources.empty()) + return Status::InvalidArgument("empty credit charge" LOC_MARK); + for (auto [r, amount] : c.resources) { + size_t i = 0; + CHECK_STATUS(resourceIndex(r, i)); + if (!amount || out[i]) + return Status::InvalidArgument( + "zero or duplicate credit charge" LOC_MARK); + out[i] = amount; + } + return Status::OK(); +} + +Status SenderCreditLedger::activate(const CreditKey& k, uint64_t epoch) { + if (!epoch) return Status::InvalidArgument("zero credit epoch" LOC_MARK); + std::lock_guard lock(mutex_); + auto existing = entries_.find(k); + if (existing != entries_.end()) { + if (epoch < existing->second.epoch) + return Status::InvalidEntry("stale credit activation" LOC_MARK); + if (epoch == existing->second.epoch) return Status::OK(); + Entry replacement; + replacement.epoch = epoch; + existing->second = replacement; + return Status::OK(); + } else if (entries_.size() >= max_entries_) { + return Status::TooManyRequests("credit ledger entry limit" LOC_MARK); + } + Entry e; + e.epoch = epoch; + entries_.emplace(k, e); + return Status::OK(); +} + +Status SenderCreditLedger::deactivate(const CreditKey& k, uint64_t epoch) { + if (!epoch) return Status::InvalidArgument("zero credit epoch" LOC_MARK); + std::lock_guard lock(mutex_); + auto existing = entries_.find(k); + if (existing == entries_.end()) return Status::OK(); // idempotent cleanup + if (existing->second.epoch != epoch) + return Status::InvalidEntry("credit cleanup epoch mismatch" LOC_MARK); + entries_.erase(existing); + return Status::OK(); +} + +Status SenderCreditLedger::applyUpdate(const CreditKey& k, + const ReceiverCreditUpdateV1& u, + CreditUpdateDisposition& disposition) { + if (u.schema_version != 1 || !u.epoch || !u.sequence) + return Status::InvalidArgument("invalid credit update header" LOC_MARK); + if (!(u.receiver_session_id == k.receiver_session) || + u.qos_class != k.qos_class || u.grants.size() > kCreditResourceCount) + return Status::InvalidArgument("credit update identity/size" LOC_MARK); + std::array proposed{}; + std::array present{}; + for (auto a : u.grants) { + size_t i = 0; + CHECK_STATUS(resourceIndex(a.resource, i)); + if (present[i]) + return Status::InvalidArgument("duplicate grant resource" LOC_MARK); + present[i] = true; + proposed[i] = a.grant_total; + } + std::lock_guard lock(mutex_); + auto it = entries_.find(k); + if (it == entries_.end() || it->second.epoch != u.epoch) + return Status::InvalidEntry("inactive or stale credit epoch" LOC_MARK); + auto& e = it->second; + if (e.has_update && u.sequence <= e.last_sequence) { + disposition = CreditUpdateDisposition::DuplicateOrOld; + return Status::OK(); + } + // `grants` is a partial cumulative update: each resource present in this + // message replaces that resource's cumulative grant total, while omitted + // resources retain their previous totals. This lets the receiver refresh + // only the resources whose available capacity changed. + for (size_t i = 0; i < kCreditResourceCount; ++i) + if (present[i] && + (proposed[i] < e.grants[i] || proposed[i] < e.consumed[i])) + return Status::InvalidArgument( + "decreasing or under-consumed grant" LOC_MARK); + bool gap = e.has_update && u.sequence > e.last_sequence + 1; + for (size_t i = 0; i < kCreditResourceCount; ++i) + if (present[i]) e.grants[i] = proposed[i]; + e.last_sequence = u.sequence; + e.has_update = true; + disposition = gap ? CreditUpdateDisposition::SequenceGap + : CreditUpdateDisposition::Applied; + return Status::OK(); +} + +Status SenderCreditLedger::tryReserve(const CreditKey& k, + const CreditCharge& c) { + std::array n; + CHECK_STATUS(normalize(c, n)); + std::lock_guard lock(mutex_); + auto it = entries_.find(k); + if (it == entries_.end() || !it->second.has_update) + return Status::InvalidEntry("credit unavailable" LOC_MARK); + auto& e = it->second; + for (size_t i = 0; i < kCreditResourceCount; ++i) + if (e.consumed[i] > e.grants[i] || n[i] > e.grants[i] - e.consumed[i]) + return Status::TooManyRequests("insufficient credit" LOC_MARK); + for (size_t i = 0; i < kCreditResourceCount; ++i) e.consumed[i] += n[i]; + return Status::OK(); +} + +Status SenderCreditLedger::rollbackReservation(const CreditKey& k, + const CreditCharge& c) { + std::array n; + CHECK_STATUS(normalize(c, n)); + std::lock_guard lock(mutex_); + auto it = entries_.find(k); + if (it == entries_.end()) + return Status::InvalidEntry("credit session inactive" LOC_MARK); + for (size_t i = 0; i < kCreditResourceCount; ++i) + if (n[i] > it->second.consumed[i]) + return Status::InvalidArgument( + "credit rollback underflow" LOC_MARK); + for (size_t i = 0; i < kCreditResourceCount; ++i) + it->second.consumed[i] -= n[i]; + return Status::OK(); +} + +Status SenderCreditLedger::available(const CreditKey& k, CreditResource r, + uint64_t& v) const { + size_t i = 0; + CHECK_STATUS(resourceIndex(r, i)); + std::lock_guard lock(mutex_); + auto it = entries_.find(k); + if (it == entries_.end() || !it->second.has_update) + return Status::InvalidEntry("credit unavailable" LOC_MARK); + v = it->second.grants[i] - it->second.consumed[i]; + return Status::OK(); +} + +Status SenderCreditLedger::consumed(const CreditKey& k, CreditResource r, + uint64_t& v) const { + size_t i = 0; + CHECK_STATUS(resourceIndex(r, i)); + std::lock_guard lock(mutex_); + auto it = entries_.find(k); + if (it == entries_.end()) + return Status::InvalidEntry("credit session inactive" LOC_MARK); + v = it->second.consumed[i]; + return Status::OK(); +} + +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp b/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp index 98f6331afb..6181171900 100644 --- a/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/segment_manager.cpp @@ -25,8 +25,18 @@ namespace mooncake { namespace tent { +namespace { +uint64_t nextManagerId() { + static std::atomic counter{0}; + return counter.fetch_add(1, std::memory_order_relaxed); +} +} // namespace + SegmentManager::SegmentManager(std::unique_ptr agent) - : next_id_(1), version_(0), registry_(std::move(agent)) { + : next_id_(1), + version_(0), + manager_id_(nextManagerId()), + registry_(std::move(agent)) { local_desc_ = std::make_shared(); subscribers_lock_ = std::make_shared(); subscribers_ = std::make_shared>(); @@ -60,7 +70,39 @@ Status SegmentManager::closeRemote(SegmentID handle) { return Status::OK(); } +Status SegmentManager::updateLocal( + const std::function &mutator) { + std::lock_guard g(local_update_mu_); + // No concurrent writer exists (serialized by local_update_mu_), so + // reading local_desc_ without local_desc_lock_ is safe here. + auto next = std::make_shared(*local_desc_); + CHECK_STATUS(mutator(*next)); + { + std::unique_lock guard(local_desc_lock_); + local_desc_ = std::move(next); + } + // Release pairs with the acquire load in getLocal(): a reader that + // observes the new version must also observe the pointer swap above. + // With a relaxed bump, on weakly-ordered architectures the version + // store may become visible before the swap, letting a reader tag the + // OLD snapshot with the NEW version — which its thread-local cache + // would then serve until the next publication. + local_desc_version_.fetch_add(1, std::memory_order_release); + { + std::lock_guard jg(local_json_cache_mu_); + local_json_cache_.reset(); + } + return Status::OK(); +} + Status SegmentManager::getRemoteCached(SegmentDesc *&desc, SegmentID handle) { + SegmentDescRef ref; + CHECK_STATUS(getRemoteCached(ref, handle)); + desc = ref.get(); + return Status::OK(); +} + +Status SegmentManager::getRemoteCached(SegmentDescRef &desc, SegmentID handle) { auto &cache = tl_remote_cache_.get(); auto current_ts = getCurrentTimeInNano(); auto current_version = version_.load(std::memory_order_relaxed); @@ -78,7 +120,7 @@ Status SegmentManager::getRemoteCached(SegmentDesc *&desc, SegmentID handle) { cache.id_to_desc_map[handle] = desc_ref; std::string peer_rpc_addr = desc_ref->rpc_server_addr; - std::string local_rpc_addr = local_desc_->rpc_server_addr; + std::string local_rpc_addr = getLocal()->rpc_server_addr; if (!peer_rpc_addr.empty() && !local_rpc_addr.empty()) { // Send a subscription request to enable proactive cache // invalidation. This is a best-effort mechanism to reduce stale @@ -92,7 +134,7 @@ Status SegmentManager::getRemoteCached(SegmentDesc *&desc, SegmentID handle) { << "'."; } } - desc = cache.id_to_desc_map[handle].get(); + desc = cache.id_to_desc_map[handle]; assert(desc); return Status::OK(); } @@ -164,7 +206,7 @@ Status SegmentManager::makeFileRemote(SegmentDescRef &desc, desc = std::make_shared(); desc->name = segment_name; desc->type = SegmentType::File; - desc->machine_id = local_desc_->machine_id; + desc->machine_id = getLocal()->machine_id; FileSegmentDesc detail; FileBufferDesc buffer; buffer.path = path; @@ -176,26 +218,42 @@ Status SegmentManager::makeFileRemote(SegmentDescRef &desc, } std::shared_ptr SegmentManager::getLocalDumpedJson() { + // Capture the snapshot together with its publication version so a slow + // dump of an older snapshot can never overwrite the cache entry computed + // from a newer one. Acquire keeps the tag conservative: the snapshot + // read below is then guaranteed to be at least as new as the tag. + auto version = local_desc_version_.load(std::memory_order_acquire); + auto snapshot = getLocal(); { std::lock_guard g(local_json_cache_mu_); - if (local_json_cache_) return local_json_cache_; + if (local_json_cache_ && local_json_cache_version_ == version) + return local_json_cache_; } - json j = *local_desc_; + json j = *snapshot; auto computed = std::make_shared(j.dump()); std::lock_guard g(local_json_cache_mu_); - if (!local_json_cache_) { + // Store only if no publication happened since we sampled `version`; + // otherwise this dump is already outdated and must not evict a cache + // entry computed from a newer snapshot. + if (local_desc_version_.load(std::memory_order_relaxed) == version && + !local_json_cache_) { local_json_cache_ = computed; + local_json_cache_version_ = version; } - return local_json_cache_; + return computed; } Status SegmentManager::synchronizeLocal() { { - std::lock_guard g(local_json_cache_mu_); - local_json_cache_.reset(); + // Serialize {snapshot, put} pairs: without this, a put carrying an + // older snapshot could complete after (and overwrite) one carrying a + // newer snapshot in the registry, hiding a completed registration + // from peers until the next synchronizeLocal call. + std::lock_guard g(local_sync_mu_); + auto snapshot = getLocal(); + CHECK_STATUS(registry_->putSegmentDesc(snapshot)); } - CHECK_STATUS(registry_->putSegmentDesc(local_desc_)); std::vector subscribers_snapshot; { @@ -214,7 +272,7 @@ Status SegmentManager::synchronizeLocal() { // Remove subscribers that have failed (e.g., peer might shutdown) // to avoid repeated RPC failures. ControlClient::notifySegmentUpdatedAsync( - subscriber, local_desc_->name, + subscriber, getLocal()->name, /* on_failure */ [subscribers = subscribers_, lock = subscribers_lock_, subscriber] { RWSpinlock::WriteGuard guard(*lock); @@ -225,7 +283,7 @@ Status SegmentManager::synchronizeLocal() { } Status SegmentManager::deleteLocal() { - return registry_->deleteSegmentDesc(local_desc_->name); + return registry_->deleteSegmentDesc(getLocal()->name); } } // namespace tent diff --git a/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp b/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp index 2f84486d91..0faf27ba25 100644 --- a/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/segment_tracker.cpp @@ -25,39 +25,47 @@ namespace mooncake { namespace tent { -Status SegmentTracker::query(uint64_t base, size_t length, - std::vector& result) { - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - assert(length); +namespace { +void sortBuffers(std::vector& buffers) { + std::sort(buffers.begin(), buffers.end(), + [](const BufferDesc& lhs, const BufferDesc& rhs) -> bool { + if (lhs.addr < rhs.addr) return true; + if (lhs.addr > rhs.addr) return false; + return lhs.length > rhs.length; // prefer large interval + }); +} + +bool containsBuffer(const SegmentDesc& desc, uint64_t base, size_t length) { + auto& detail = std::get(desc.detail); for (auto& buf : detail.buffers) { - if (buf.addr > base) continue; - if (buf.addr + buf.length <= base) break; - result.push_back(&buf); - if (buf.addr + buf.length >= base + length) { - return Status::OK(); - } else { - auto new_base = buf.addr + buf.length; - auto new_length = base + length - new_base; - return query(new_base, new_length, result); - } + if (buf.addr == base && buf.length == length) return true; } - return Status::InvalidArgument("Some buffers are not registered"); + return false; } +} // namespace Status SegmentTracker::add(uint64_t base, size_t length, std::function callback) { - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - mutex_.lock(); - for (auto& buf : detail.buffers) { - if (buf.addr == base && buf.length == length) { - buf.ref_count++; - mutex_.unlock(); + // Read-only pre-scan on the current snapshot: the common miss path pays + // no clone/publication. A hit is re-verified under the writer mutex; a + // racing insert of the same range degenerates to today's benign + // duplicate-registration behavior. + if (containsBuffer(*manager_.getLocal(), base, length)) { + bool found = false; + CHECK_STATUS(manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); + for (auto& buf : detail.buffers) { + if (buf.addr == base && buf.length == length) { + buf.ref_count++; + found = true; + break; + } + } return Status::OK(); - } + })); + if (found) return Status::OK(); } - mutex_.unlock(); BufferDesc new_desc; new_desc.addr = base; new_desc.length = length; @@ -71,86 +79,129 @@ Status SegmentTracker::add(uint64_t base, size_t length, new_desc.ref_count = 1; auto status = callback(new_desc); if (!status.ok()) return status; - mutex_.lock(); - detail.buffers.push_back(new_desc); - std::sort(detail.buffers.begin(), detail.buffers.end(), - [](const BufferDesc& lhs, BufferDesc& rhs) -> bool { - if (lhs.addr < rhs.addr) return true; - if (lhs.addr > rhs.addr) return false; - return lhs.length > rhs.length; // prefer large interval - }); - mutex_.unlock(); - return Status::OK(); + return manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); + detail.buffers.push_back(std::move(new_desc)); + sortBuffers(detail.buffers); + return Status::OK(); + }); } Status SegmentTracker::addInBatch( std::vector& desc_list, std::function&)> callback) { std::vector new_desc_list; - for (auto& desc : desc_list) { - bool found = false; - { - std::lock_guard lock(mutex_); - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - for (auto& buf : detail.buffers) { - if (buf.addr == desc.addr && buf.length == desc.length) { - buf.ref_count++; - found = true; - break; - } + // Read-only pre-scan (see add()): skip the ref-count publication when no + // entry duplicates an already-registered range. + bool any_dup = false; + { + auto snapshot = manager_.getLocal(); + for (auto& entry : desc_list) { + if (containsBuffer(*snapshot, entry.addr, entry.length)) { + any_dup = true; + break; } } - if (!found) new_desc_list.push_back(std::move(desc)); + } + // Ranges whose ref_count we bumped; used to roll back if the callback + // fails. The bump must happen under the writer mutex *before* the + // callback: it is what pins the duplicate entry (ref_count >= 2) so a + // concurrent unregister cannot erase it out from under this + // registration. + std::vector> bumped; + if (any_dup) { + CHECK_STATUS(manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); + for (auto& entry : desc_list) { + bool found = false; + for (auto& buf : detail.buffers) { + if (buf.addr == entry.addr && buf.length == entry.length) { + buf.ref_count++; + found = true; + break; + } + } + if (found) { + bumped.emplace_back(entry.addr, entry.length); + } else { + new_desc_list.push_back(std::move(entry)); + } + } + return Status::OK(); + })); + } else { + new_desc_list = std::move(desc_list); } auto status = callback(new_desc_list); - if (!status.ok()) return status; - { - std::lock_guard lock(mutex_); - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); + if (!status.ok()) { + // Roll back the duplicate ref-counts so a failed registration does + // not leave buffers pinned forever. + if (!bumped.empty()) { + manager_.updateLocal([&](SegmentDesc& desc) -> Status { + auto& detail = std::get(desc.detail); + for (auto& range : bumped) { + for (auto it = detail.buffers.begin(); + it != detail.buffers.end(); ++it) { + if (it->addr == range.first && + it->length == range.second) { + it->ref_count--; + // The original owner unregistered while we held + // the extra reference; drop the entry so it is + // no longer advertised. + if (it->ref_count == 0) detail.buffers.erase(it); + break; + } + } + } + return Status::OK(); + }); + } + return status; + } + return manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); for (auto& new_desc : new_desc_list) { detail.buffers.push_back(new_desc); } - std::sort(detail.buffers.begin(), detail.buffers.end(), - [](const BufferDesc& lhs, BufferDesc& rhs) -> bool { - if (lhs.addr < rhs.addr) return true; - if (lhs.addr > rhs.addr) return false; - return lhs.length > rhs.length; // prefer large interval - }); - } - return Status::OK(); + sortBuffers(detail.buffers); + return Status::OK(); + }); } Status SegmentTracker::remove(uint64_t base, size_t length, std::function callback) { - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - mutex_.lock(); - for (auto it = detail.buffers.begin(); it != detail.buffers.end(); ++it) { - if (it->addr == base && (!length || it->length == length)) { - it->ref_count--; - Status status = Status::OK(); - if (it->ref_count == 0) { - BufferDesc clone = *it; - detail.buffers.erase(it); - mutex_.unlock(); - status = callback(clone); - } else { - mutex_.unlock(); + bool removed = false; + BufferDesc removed_desc; + CHECK_STATUS(manager_.updateLocal([&](SegmentDesc& desc) -> Status { + assert(desc.type == SegmentType::Memory); + auto& detail = std::get(desc.detail); + for (auto it = detail.buffers.begin(); it != detail.buffers.end(); + ++it) { + if (it->addr == base && (!length || it->length == length)) { + it->ref_count--; + if (it->ref_count == 0) { + removed_desc = *it; + detail.buffers.erase(it); + removed = true; + } + break; } - return status; } - } - mutex_.unlock(); + return Status::OK(); + })); + if (removed) return callback(removed_desc); return Status::OK(); } -Status SegmentTracker::forEach(std::function callback) { - std::lock_guard lock(mutex_); - assert(local_desc_->type == SegmentType::Memory); - auto& detail = std::get(local_desc_->detail); - for (auto& buf : detail.buffers) { +Status SegmentTracker::forEach( + std::function callback) { + auto snapshot = manager_.getLocal(); + assert(snapshot->type == SegmentType::Memory); + for (const auto& buf : + std::get(snapshot->detail).buffers) { auto status = callback(buf); if (!status.ok()) return status; } diff --git a/mooncake-transfer-engine/tent/src/runtime/topology.cpp b/mooncake-transfer-engine/tent/src/runtime/topology.cpp index 30a263d9e3..8892ebdaea 100644 --- a/mooncake-transfer-engine/tent/src/runtime/topology.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/topology.cpp @@ -14,19 +14,24 @@ #include "tent/runtime/topology.h" +#include #include #include #include #include #include +#include +#include #include #include -#include #include "tent/common/status.h" #include "tent/runtime/platform.h" #include "tent/common/utils/random.h" #include "tent/thirdparty/nlohmann/json.h" +#ifdef USE_UB +#include "tent/transport/ub/topology_attrs.h" +#endif namespace mooncake { namespace tent { @@ -50,6 +55,9 @@ std::string Topology::toString() const { nj["pci_bus_id"] = nic.pci_bus_id; nj["type"] = nic.type; nj["numa_node"] = nic.numa_node; + if (!nic.device_attrs.empty()) { + nj["device_attrs"] = nic.device_attrs; + } j["nics"].push_back(nj); } @@ -99,10 +107,81 @@ void Topology::print() const { } Status Topology::discover(const std::vector& platforms) { + return discover(platforms, false); +} + +Status Topology::discover(const std::vector& platforms, + bool discover_ub) { clear(); for (auto& entry : platforms) { CHECK_STATUS(entry->probe(nic_list_, mem_list_)); } +#ifdef USE_UB + // UB discovery is intentionally adapter-backed instead of inferring UB + // devices from verbs/sysfs names. One topology NIC is emitted per EID and + // carries both the globally serialized identity and the native URMA name. + auto adapter = discover_ub ? ub::createDefaultUrmaAdapter() : nullptr; + if (adapter && adapter->available()) { + auto status = adapter->initialize(); + if (status.ok()) { + std::vector devices; + status = adapter->discoverDevices(devices); + if (status.ok()) { + std::unordered_map native_device_indices; + for (const auto& device : devices) { + auto [native_it, inserted] = native_device_indices.emplace( + device.native_device_name, + static_cast(native_device_indices.size())); + (void)inserted; + + int numa_node = -1; + std::string pci_bus_id; + if (!device.native_device_path.empty()) { + std::error_code error; + const auto device_path = std::filesystem::canonical( + std::filesystem::path(device.native_device_path) / + "device", + error); + if (!error) { + pci_bus_id = device_path.filename().string(); + std::ifstream(device_path / "numa_node") >> + numa_node; + } + } + + const NicID nic_id = static_cast(nic_list_.size()); + NicEntry nic{.name = device.topology_name, + .pci_bus_id = std::move(pci_bus_id), + .type = NIC_UB, + .numa_node = numa_node}; + ub::encodeTopologyDeviceAttributes( + device, native_it->second, nic.device_attrs); + nic_list_.push_back(std::move(nic)); + + for (auto& memory : mem_list_) { + const size_t rank = + numa_node >= 0 && memory.numa_node == numa_node + ? 0 + : DevicePriorityRanks - 1; + memory.device_list[rank].push_back(nic_id); + } + } + } else { + LOG(WARNING) << "Unable to discover optional UB devices: " + << status.ToString(); + } + auto shutdown_status = adapter->shutdown(); + if (!shutdown_status.ok()) { + LOG(WARNING) << "Unable to release UB discovery runtime: " + << shutdown_status.ToString(); + } + } else { + LOG(WARNING) << "Unable to initialize optional UB discovery: " + << status.ToString(); + } + } +#endif + (void)discover_ub; return Status::OK(); } @@ -118,6 +197,12 @@ Status Topology::parse(const std::string& json_content) { nic.type = static_cast(item.value("type", NIC_UNKNOWN)); nic.numa_node = item.value("numa_node", -1); + if (item.contains("device_attrs")) { + nic.device_attrs = + item.at("device_attrs") + .get< + std::unordered_map>(); + } nic_list_.push_back(nic); } } diff --git a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp index b3456e86db..b1a4f3bbbd 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp @@ -56,6 +56,9 @@ struct Batch { std::array sub_batch; std::vector task_list; size_t max_size; + size_t runtime_refs{0}; + bool free_requested{false}; + uint64_t queue_token{0}; struct SubmitHook { size_t start_task_id{0}; @@ -262,14 +265,16 @@ std::string getMachineID() { Status TransferEngineImpl::setupLocalSegment() { auto& manager = metadata_->segmentManager(); - auto segment = manager.getLocal(); - segment->name = local_segment_name_; - segment->type = SegmentType::Memory; - segment->machine_id = getMachineID(); - segment->rpc_server_addr = buildIpAddrWithPort(hostname_, port_, ipv6_); - auto& detail = std::get(segment->detail); - detail.topology = *(topology_.get()); - local_segment_tracker_ = std::make_unique(segment); + CHECK_STATUS(manager.updateLocal([&](SegmentDesc& segment) -> Status { + segment.name = local_segment_name_; + segment.type = SegmentType::Memory; + segment.machine_id = getMachineID(); + segment.rpc_server_addr = buildIpAddrWithPort(hostname_, port_, ipv6_); + auto& detail = std::get(segment.detail); + detail.topology = *(topology_.get()); + return Status::OK(); + })); + local_segment_tracker_ = std::make_unique(manager); return manager.synchronizeLocal(); } @@ -286,6 +291,37 @@ Status TransferEngineImpl::construct() { enable_auto_failover_on_poll_ = conf_->get("enable_auto_failover_on_poll", true); enable_progress_worker_ = conf_->get("enable_progress_worker", false); + runtime_queue_config_.enabled = conf_->get("enable_runtime_queue", false); + if (runtime_queue_config_.enabled) enable_progress_worker_ = true; + runtime_queue_config_.limits.max_outstanding_owners = + conf_->get("runtime_queue/max_outstanding_owners", 1024UL); + runtime_queue_config_.limits.max_outstanding_bytes = + conf_->get("runtime_queue/max_outstanding_bytes", 1UL << 30); + runtime_queue_config_.limits.staging_owner_reserve = + conf_->get("runtime_queue/staging_owner_reserve", 0UL); + runtime_queue_config_.limits.staging_byte_reserve = + conf_->get("runtime_queue/staging_byte_reserve", 0UL); + runtime_queue_config_.limits.deadline_aware = + conf_->get("runtime_queue/deadline_aware", false); + runtime_queue_config_.limits.mlu_local_threshold = + conf_->get("runtime_queue/mlu_local_threshold", 0.0); + runtime_queue_config_.limits.promotion_slack_ns = + conf_->get("runtime_queue/promotion_slack_ns", 0UL); + runtime_queue_config_.max_dispatch_owners = + conf_->get("runtime_queue/max_dispatch_owners", 64UL); + runtime_queue_config_.max_dispatch_bytes = + conf_->get("runtime_queue/max_dispatch_bytes", 64UL << 20); + runtime_queue_config_.progress_fallback_interval = + std::chrono::microseconds( + conf_->get("runtime_queue/progress_fallback_interval_us", 50000UL)); + if (runtime_queue_config_.enabled && + (runtime_queue_config_.max_dispatch_owners == 0 || + runtime_queue_config_.max_dispatch_bytes == 0)) { + return Status::InvalidArgument( + "runtime queue dispatch window must be non-zero" LOC_MARK); + } + runtime_queue_ = std::make_unique( + runtime_queue_config_.limits); if (!hostname_.empty()) CHECK_STATUS(checkLocalIpAddress(hostname_, ipv6_)); else @@ -338,8 +374,35 @@ Status TransferEngineImpl::construct() { staging_proxy_ = std::make_unique(this); + if (runtime_queue_config_.limits.deadline_aware && + runtime_queue_config_.limits.mlu_local_threshold > 0.0) { + auto rdma_xport = + transport_list_[static_cast(TransportType::RDMA)]; + if (rdma_xport) { + std::weak_ptr weak_rdma = rdma_xport; + runtime_queue_->setDegradationPolicy( + [weak_rdma]() -> double { + if (auto rdma = weak_rdma.lock()) { + return rdma->getEstimatedBandwidth(); + } + return -1.0; + }, + DegradationHooks{}, nullptr); + LOG(INFO) << "Admission queue degradation: live RDMA bw" + << ", theta_local=" + << runtime_queue_config_.limits.mlu_local_threshold; + } else { + LOG(WARNING) << "Admission queue degradation requested but RDMA " + "transport is " + "unavailable"; + } + } + if (enable_progress_worker_) { - progress_worker_ = std::make_unique(this); + progress_worker_ = std::make_unique( + this, runtime_queue_config_.enabled + ? runtime_queue_config_.progress_fallback_interval + : std::chrono::microseconds(0)); progress_worker_->start(); } @@ -386,10 +449,11 @@ Status TransferEngineImpl::deconstruct() { // Metrics cleanup is handled automatically by TentMetrics destructor // Stop the progress worker first so it cannot race with batch teardown - // below (it dereferences BatchID into Batch* via progressBatch). + // below (it dereferences BatchID into Batch* via progressBatch). Keep the + // object alive until transports are destroyed: completion paths may still + // issue a final no-op wake while their workers are joining. if (progress_worker_) { progress_worker_->stop(); - progress_worker_.reset(); } // Destroy staging_proxy_ first: its destructor calls back into @@ -398,10 +462,13 @@ Status TransferEngineImpl::deconstruct() { staging_proxy_.reset(); if (local_segment_tracker_) { - local_segment_tracker_->forEach([&](BufferDesc& desc) -> Status { + local_segment_tracker_->forEach([&](const BufferDesc& desc) -> Status { + // Snapshot entries are immutable; transports may scrub fields of + // their deregistration argument, so hand them a copy. + BufferDesc copy = desc; for (size_t type = 0; type < kSupportedTransportTypes; ++type) { if (transport_list_[type]) - transport_list_[type]->removeMemoryBuffer(desc); + transport_list_[type]->removeMemoryBuffer(copy); } return Status::OK(); }); @@ -415,17 +482,11 @@ Status TransferEngineImpl::deconstruct() { // does not access transport-internal state (workers, connections). // Callers must ensure no transfers are in-flight before calling // deconstruct(). - batch_set_.forEach([&](BatchSet& entry) { - for (auto& batch : entry.active) { - for (size_t type = 0; type < kSupportedTransportTypes; ++type) { - auto& transport = transport_list_[type]; - auto& sub_batch = batch->sub_batch[type]; - if (!transport || !sub_batch) continue; - transport->freeSubBatch(sub_batch); - } - Slab::Get().deallocate(batch); - } - for (auto& batch : entry.freelist) { + { + std::lock_guard lk(progress_mutex_); + std::unordered_set released_batches; + auto release_batch = [&](Batch* batch) { + if (!released_batches.insert(batch).second) return; for (size_t type = 0; type < kSupportedTransportTypes; ++type) { auto& transport = transport_list_[type]; auto& sub_batch = batch->sub_batch[type]; @@ -433,13 +494,17 @@ Status TransferEngineImpl::deconstruct() { transport->freeSubBatch(sub_batch); } Slab::Get().deallocate(batch); - } - entry.active.clear(); - entry.freelist.clear(); - }); + }; + for (auto& batch : batch_set_.active) release_batch(batch); + for (auto& batch : batch_set_.freelist) release_batch(batch); + batch_set_.active.clear(); + batch_set_.freelist.clear(); + alive_batches_.clear(); + } // Now safe to destroy transports (workers join here) for (auto& transport : transport_list_) transport.reset(); + progress_worker_.reset(); local_segment_tracker_.reset(); if (metadata_) { metadata_->segmentManager().deleteLocal(); @@ -484,9 +549,10 @@ Status TransferEngineImpl::closeSegment(SegmentID handle) { } Status TransferEngineImpl::getSegmentInfo(SegmentID handle, SegmentInfo& info) { - SegmentDesc* desc = nullptr; + // Owning reference: keeps the snapshot alive while we read through it. + SegmentDescRef desc; if (handle == LOCAL_SEGMENT_ID) { - desc = metadata_->segmentManager().getLocal().get(); + desc = metadata_->segmentManager().getLocal(); } else { CHECK_STATUS(metadata_->segmentManager().getRemoteCached(desc, handle)); } @@ -615,6 +681,7 @@ std::vector TransferEngineImpl::getSupportedTransports( if (transport_list_[SHM]) result.push_back(SHM); if (transport_list_[TCP]) result.push_back(TCP); if (transport_list_[GDS]) result.push_back(GDS); + if (transport_list_[TPU]) result.push_back(TPU); return result; } @@ -728,12 +795,10 @@ BatchID TransferEngineImpl::allocateBatch(size_t batch_size) { Batch* batch = Slab::Get().allocate(); if (!batch) return (BatchID)0; batch->max_size = batch_size; - batch_set_.get().active.insert(batch); BatchID batch_id = (BatchID)batch; - { - std::lock_guard lk(progress_mutex_); - alive_batches_.insert(batch_id); - } + std::lock_guard lk(progress_mutex_); + batch_set_.active.insert(batch); + alive_batches_.insert(batch_id); return batch_id; } @@ -741,38 +806,120 @@ Status TransferEngineImpl::freeBatch(BatchID batch_id) { if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK); Batch* batch = (Batch*)(batch_id); std::lock_guard lk(progress_mutex_); - batch_set_.get().freelist.push_back(batch); + if (!alive_batches_.count(batch_id)) + return Status::InvalidArgument("Batch is not alive" LOC_MARK); + if (runtime_queue_config_.enabled && batch->queue_token != 0) { + auto retire_status = retireQueueForBatch(batch); + if (!retire_status.ok() && !retire_status.IsInvalidEntry()) { + return retire_status; + } + } + if (batch->free_requested) { + CHECK_STATUS(lazyFreeBatch()); + return Status::OK(); + } + batch->free_requested = true; + batch_set_.freelist.push_back(batch); lazyFreeBatch(); return Status::OK(); } Status TransferEngineImpl::lazyFreeBatch() { - // Caller must hold progress_mutex_. - auto& batch_set = batch_set_.get(); - for (auto it = batch_set.freelist.begin(); - it != batch_set.freelist.end();) { + std::lock_guard lk(progress_mutex_); + for (auto it = batch_set_.freelist.begin(); + it != batch_set_.freelist.end();) { auto& batch = *it; + if (batch->runtime_refs > 0) { + it++; + continue; + } TransferStatus overall_status; CHECK_STATUS(getTransferStatus((BatchID)batch, overall_status)); if (overall_status.s == PENDING) { it++; continue; } + if (runtime_queue_config_.enabled && batch->queue_token != 0) { + CHECK_STATUS(retireQueueForBatch(batch)); + } for (size_t type = 0; type < kSupportedTransportTypes; ++type) { auto& transport = transport_list_[type]; auto& sub_batch = batch->sub_batch[type]; if (transport && sub_batch) transport->freeSubBatch(sub_batch); } - batch_set.active.erase(batch); + batch_set_.active.erase(batch); alive_batches_.erase((BatchID)batch); Slab::Get().deallocate(batch); - it = batch_set.freelist.erase(it); + it = batch_set_.freelist.erase(it); } return Status::OK(); } +Status TransferEngineImpl::retainBatch(BatchID batch_id, Batch*& batch) { + if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK); + std::lock_guard lk(progress_mutex_); + if (!alive_batches_.count(batch_id)) { + return Status::InvalidArgument("Batch is not alive" LOC_MARK); + } + batch = (Batch*)batch_id; + if (batch->free_requested) { + return Status::InvalidArgument("Batch is being freed" LOC_MARK); + } + ++batch->runtime_refs; + return Status::OK(); +} + +Status TransferEngineImpl::releaseBatch(Batch* batch) { + if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK); + std::lock_guard lk(progress_mutex_); + if (batch->runtime_refs == 0) { + return Status::InternalError("Batch runtime ref underflow" LOC_MARK); + } + --batch->runtime_refs; + if (batch->runtime_refs == 0 && batch->free_requested) { + CHECK_STATUS(lazyFreeBatch()); + } + return Status::OK(); +} + +class TransferEngineImpl::BatchRef { + public: + BatchRef(TransferEngineImpl& engine, Batch* batch) + : engine_(engine), batch_(batch) {} + + ~BatchRef() { + if (!batch_) return; + auto status = engine_.releaseBatch(batch_); + if (!status.ok()) { + LOG(WARNING) << "failed to release batch ref: " + << status.ToString(); + } + } + + BatchRef(const BatchRef&) = delete; + BatchRef& operator=(const BatchRef&) = delete; + + Batch* get() const { return batch_; } + + Status release() { + if (!batch_) return Status::OK(); + auto status = engine_.releaseBatch(batch_); + batch_ = nullptr; + return status; + } + + private: + TransferEngineImpl& engine_; + Batch* batch_{nullptr}; +}; + static bool isGpuType(MemoryType t) { - return t == MTYPE_CUDA || t == MTYPE_ROCM; + // TPU HBM behaves like a GPU that lacks NIC access: it is a device-side + // memory that can only reach the network by staging through host DRAM. + // Treating it as a "gpu type" makes the capability checks route its + // device<->host hop to TpuTransport (gpu_to_dram / dram_to_gpu) while + // leaving gpu_to_gpu unsatisfiable, which forces host-DRAM staging. + return t == MTYPE_CUDA || t == MTYPE_ROCM || t == MTYPE_TPU; } static bool checkAvailability(const std::shared_ptr& xport, @@ -800,6 +947,7 @@ static MemoryType getTypeEnum(const std::string& type) { if (type == "cuda") return MTYPE_CUDA; if (type == "npu") return MTYPE_CUDA; if (type == "rocm") return MTYPE_ROCM; + if (type == "tpu") return MTYPE_TPU; return MTYPE_UNKNOWN; } @@ -815,7 +963,7 @@ Status TransferEngineImpl::validateTransportHint(const Request& req, if (!transport_list_[req.transport_hint]) { return Status::InvalidArgument( "transport_hint=" + - TransportSelector::transportTypeName(req.transport_hint) + + std::string(transportTypeName(req.transport_hint)) + " is not enabled in config (request[" + std::to_string(request_index) + "])" LOC_MARK); } @@ -824,9 +972,10 @@ Status TransferEngineImpl::validateTransportHint(const Request& req, SelectionResult TransferEngineImpl::getTransportType(const Request& request, int transport_index) { - SegmentDesc* desc; + // Owning reference: keeps the snapshot alive while we read through it. + SegmentDescRef desc; if (request.target_id == LOCAL_SEGMENT_ID) { - desc = metadata_->segmentManager().getLocal().get(); + desc = metadata_->segmentManager().getLocal(); } else { auto status = metadata_->segmentManager().getRemoteCached( desc, request.target_id); @@ -859,7 +1008,11 @@ SelectionResult TransferEngineImpl::getTransportType(const Request& request, auto remote_mtype = getTypeEnum(LocationParser(entry->location).type()); for (auto type : entry->transports) { - if ((type == NVLINK || type == SHM) && !same_machine) + // NVLINK/SHM are same-machine only; TPU is a + // local-stage-only executor and must never carry a remote + // hop. + if ((type == NVLINK || type == SHM || type == TPU) && + !same_machine) continue; if (checkAvailability(transport_list_[type], local_mtype, remote_mtype)) { @@ -887,6 +1040,7 @@ SelectionResult TransferEngineImpl::getTransportType(const Request& request, ctx.priority_level = request.priority; // Use request priority for selection ctx.policy_name = request.policy_name; // Optional: bind to specific policy + ctx.intent_type = request.intent_type; // Business intent policy filter if (desc->type == SegmentType::File) { // File segment: use selector with empty buffer_transports @@ -915,32 +1069,6 @@ SelectionResult TransferEngineImpl::getTransportType(const Request& request, hint); } -static const char* transportTypeName(TransportType type) { - switch (type) { - case UNSPEC: - return "UNSPEC"; - case RDMA: - return "RDMA"; - case MNNVL: - return "MNNVL"; - case SHM: - return "SHM"; - case NVLINK: - return "NVLINK"; - case GDS: - return "GDS"; - case IOURING: - return "IOURING"; - case TCP: - return "TCP"; - case AscendDirect: - return "AscendDirect"; - case SUNRISE_LINK: - return "SUNRISE_LINK"; - } - return "UNKNOWN"; -} - std::string printRequest(const Request& request) { std::stringstream ss; ss << "opcode " << request.opcode << " source " << request.source @@ -967,6 +1095,27 @@ struct MergeResult { std::map task_lookup; }; +struct TransferEngineImpl::PreparedSubmit { + struct Task { + size_t merged_task_index{0}; + size_t task_id{0}; + }; + + struct Owner { + size_t owner_task_id{0}; + bool has_owner_task_id{false}; + std::vector derived_task_ids; + Request request{}; + SelectionResult route{}; + bool staging{false}; + std::vector staging_params; + }; + + std::chrono::steady_clock::time_point submit_time{}; + std::vector tasks; + std::vector owners; +}; + namespace { bool tryAddUint64(uint64_t lhs, uint64_t rhs, uint64_t& out) { @@ -1086,7 +1235,8 @@ std::vector resolveRequestBoundaries( // Group requests by target_id so withCachedSegment fires at most once per // peer. std::vector boundaries(requests.size()); - auto* local_desc = metadata->segmentManager().getLocal().get(); + // Owning reference: keeps the snapshot alive while we read through it. + auto local_desc = metadata->segmentManager().getLocal(); if (local_desc) { for (size_t i = 0; i < requests.size(); ++i) { @@ -1146,8 +1296,10 @@ void TransferEngineImpl::findStagingPolicy(const Request& request, SegmentDesc* desc = nullptr; BufferDesc* entry = nullptr; + // Owning reference: `entry` is used after the lambda returns. + SegmentDescRef pin; auto status = metadata_->segmentManager().withCachedSegment( - request.target_id, [&](SegmentDesc* segment) { + request.target_id, pin, [&](SegmentDesc* segment) { desc = segment; entry = desc->findBuffer(request.target_offset, request.length); if (!entry) @@ -1187,7 +1339,7 @@ void TransferEngineImpl::findStagingPolicy(const Request& request, } // case 2: pure mnnvl if (transport_list_[MNNVL] && transport_list_[NVLINK]) { - auto& xport = transport_list_[RDMA]; + auto& xport = transport_list_[MNNVL]; auto& caps = xport->capabilities(); if (local_mtype == MTYPE_CPU && remote_mtype == MTYPE_CPU && !caps.dram_to_dram) { @@ -1202,6 +1354,34 @@ void TransferEngineImpl::findStagingPolicy(const Request& request, remote, Topology::MEM_CUDA)); } } + // case 3: TPU. HBM is not NIC-addressable, so any hop touching TPU memory + // is staged through host DRAM: TpuTransport performs the local HBM<->host + // copy (via the PJRT adapter) and the host<->host hop is carried by + // whatever host-DRAM network transport is present. TPU deployments (e.g. + // cloud TPU VMs) are typically TCP/multi-NIC rather than RDMA, so we gate + // on either; the cross stage itself is routed by capability (dram_to_dram), + // so TCP is selected when RDMA is absent. We also require TpuTransport (the + // local HBM<->host executor), mirroring how the CUDA cases gate on NVLINK. + // An empty stage location means "no staging needed on that side". + if (transport_list_[TPU] && + (transport_list_[RDMA] || transport_list_[TCP])) { + if (local_mtype == MTYPE_TPU && remote_mtype == MTYPE_TPU) { + policy.clear(); + policy.push_back(server_addr); + policy.push_back(topology_->findNearMem(local)); + policy.push_back(desc->getMemory().topology.findNearMem(remote)); + } else if (local_mtype == MTYPE_TPU && remote_mtype == MTYPE_CPU) { + policy.clear(); + policy.push_back(server_addr); + policy.push_back(topology_->findNearMem(local)); + policy.push_back(""); // remote already host DRAM + } else if (local_mtype == MTYPE_CPU && remote_mtype == MTYPE_TPU) { + policy.clear(); + policy.push_back(server_addr); + policy.push_back(""); // local already host DRAM + policy.push_back(desc->getMemory().topology.findNearMem(remote)); + } + } } SelectionResult TransferEngineImpl::resolveTransport(const Request& req, @@ -1215,39 +1395,83 @@ SelectionResult TransferEngineImpl::resolveTransport(const Request& req, return result; } -Status TransferEngineImpl::submitTransfer( - BatchID batch_id, const std::vector& request_list) { - if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK); - Batch* batch = (Batch*)(batch_id); - +Status TransferEngineImpl::prepareSubmit( + Batch* batch, const std::vector& request_list, + PreparedSubmit& prepared) { + if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK); for (size_t i = 0; i < request_list.size(); ++i) { auto st = validateTransportHint(request_list[i], i); if (!st.ok()) return st; } - std::vector classified_request_list[kSupportedTransportTypes]; - std::vector task_id_list[kSupportedTransportTypes]; - std::unordered_map merged_task_id_map; - - size_t start_task_id = batch->task_list.size(); - batch->task_list.insert(batch->task_list.end(), request_list.size(), - TaskInfo{}); - - // Record start time for metrics tracking - auto submit_time = std::chrono::steady_clock::now(); - + prepared = PreparedSubmit{}; + const size_t start_task_id = batch->task_list.size(); + prepared.submit_time = std::chrono::steady_clock::now(); auto merge_boundaries = merge_requests_ ? resolveRequestBoundaries(metadata_.get(), request_list) : std::vector{}; auto merged = mergeRequests(request_list, merge_boundaries, merge_requests_); + + prepared.owners.reserve(merged.request_list.size()); + for (const auto& request : merged.request_list) { + PreparedSubmit::Owner owner; + owner.request = request; + owner.route = resolveTransport(owner.request, 0); + if (owner.route.transport == TCP) { + findStagingPolicy(owner.request, owner.staging_params); + owner.staging = !owner.staging_params.empty() && staging_proxy_; + } + prepared.owners.push_back(std::move(owner)); + } + + prepared.tasks.reserve(merged.task_lookup.size()); + for (const auto& kv : merged.task_lookup) { + const size_t public_task_index = kv.first; + const size_t merged_task_index = kv.second; + const size_t task_id = start_task_id + public_task_index; + auto& owner = prepared.owners[merged_task_index]; + if (!owner.has_owner_task_id) { + owner.owner_task_id = task_id; + owner.has_owner_task_id = true; + } else { + owner.derived_task_ids.push_back(task_id); + } + prepared.tasks.push_back({merged_task_index, task_id}); + } + return Status::OK(); +} + +uint64_t TransferEngineImpl::nextBatchToken() { return next_batch_token_++; } + +void TransferEngineImpl::attachProgressNotifier( + Batch* batch, Transport::SubBatchRef sub_batch) { + if (!batch || !sub_batch) return; + sub_batch->progress_batch_id = (BatchID)batch; + sub_batch->notify_progress = [this](BatchID batch_id) { + notifyBatchMaybeReady(batch_id); + }; +} + +Status TransferEngineImpl::commitPreparedSubmit( + Batch* batch, const PreparedSubmit& prepared) { + if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK); + + std::vector classified_request_list[kSupportedTransportTypes]; + std::vector task_id_list[kSupportedTransportTypes]; + std::unordered_map merged_task_id_map; + + batch->task_list.insert(batch->task_list.end(), prepared.tasks.size(), + TaskInfo{}); + std::unordered_map next_sub_task_id; - for (auto& kv : merged.task_lookup) { - size_t task_id = start_task_id + kv.first; - size_t merged_task_id = kv.second; + for (const auto& task_plan : prepared.tasks) { + size_t task_id = task_plan.task_id; + size_t merged_task_id = task_plan.merged_task_index; auto& task = batch->task_list[task_id]; - auto& merged_request = merged.request_list[merged_task_id]; + const auto& owner = prepared.owners[merged_task_id]; + auto& merged_request = owner.request; if (merged_task_id_map.count(merged_task_id)) { task = merged_task_id_map[merged_task_id]; task.derived = true; @@ -1261,10 +1485,11 @@ Status TransferEngineImpl::submitTransfer( task.request = merged_request; task.staging = false; task.start_time = - submit_time; // Record start time for latency tracking - auto select_result = resolveTransport(merged_request, 0); - task.type = select_result.transport; - task.device_mask = select_result.device_mask; + prepared.submit_time; // Record start time for latency tracking + task.dispatch_time = prepared.submit_time; // No queue wait on direct + task.type = owner.route.transport; + task.device_mask = owner.route.device_mask; + if (owner.route.qp_pool) task.qp_pool = *owner.route.qp_pool; if (task.type == UNSPEC) { LOG(WARNING) << "Unable to find registered buffer for request: " << printRequest(merged_request); @@ -1272,14 +1497,22 @@ Status TransferEngineImpl::submitTransfer( continue; } - if (task.type == TCP) { - std::vector staging_params; - findStagingPolicy(merged_request, staging_params); - if (!staging_params.empty() && staging_proxy_) { - task.staging = true; - staging_proxy_->submit(&task, staging_params); - continue; + if (owner.staging) { + task.staging = true; + // Staging is an orchestration step, not a concrete transport + // attempt. ProxyManager chunks the transfer and issues the real + // Transport::submitTransferTasks() calls, which are counted where + // they recurse through the non-staging path below. Only stamp the + // logical request's first post time for stage decomposition here. + auto status = staging_proxy_->submit(&task, (BatchID)batch, + owner.staging_params); + if (!status.ok()) { + task.staging = false; + task.type = UNSPEC; + } else { + task.post_time = std::chrono::steady_clock::now(); } + continue; } if (!batch->sub_batch[task.type]) { @@ -1292,6 +1525,7 @@ Status TransferEngineImpl::submitTransfer( merged_task_id_map[merged_task_id] = task; continue; } + attachProgressNotifier(batch, batch->sub_batch[task.type]); } if (!next_sub_task_id.count(task.type)) @@ -1317,21 +1551,361 @@ Status TransferEngineImpl::submitTransfer( // this batch should have the same policy) sub_batch->device_mask = batch->task_list[task_id_list[type][0]].device_mask; + sub_batch->qp_pool = + batch->task_list[task_id_list[type][0]].qp_pool; } + auto attempt_start = std::chrono::steady_clock::now(); + for (auto& task_id : task_id_list[type]) { + startTransportAttempt(batch->task_list[task_id], + static_cast(type), + attempt_start); + } auto status = transport->submitTransferTasks( sub_batch, classified_request_list[type]); if (!status.ok()) { - // LOG(WARNING) << "Failed to submit SubBatch " << type << ":" - // << status.ToString(); - for (auto& task_id : task_id_list[type]) + auto attempt_end = std::chrono::steady_clock::now(); + for (auto& task_id : task_id_list[type]) { + finishTransportAttempt(batch->task_list[task_id], FAILED, + attempt_end); batch->task_list[task_id].type = UNSPEC; + } } } return Status::OK(); } +Status TransferEngineImpl::enqueuePreparedSubmit(Batch* batch, + const PreparedSubmit& prepared, + QueueOwnerKind owner_kind) { + std::lock_guard lk(progress_mutex_); + if (prepared.tasks.empty()) return Status::OK(); + if (prepared.tasks.size() > batch->max_size - batch->task_list.size()) { + return Status::TooManyRequests( + "batch public task capacity exceeded" LOC_MARK); + } + + const uint64_t batch_token = + batch->queue_token != 0 ? batch->queue_token : nextBatchToken(); + QueueSubmit submit; + submit.batch_token = batch_token; + submit.batch_slots_left = batch->max_size - batch->task_list.size(); + submit.owners.reserve(prepared.owners.size()); + for (const auto& owner : prepared.owners) { + if (owner.request.length > runtime_queue_config_.max_dispatch_bytes) { + return Status::TooManyRequests( + "request exceeds runtime queue dispatch byte window" LOC_MARK); + } + QueueOwnerInput input; + input.owner_task_id = owner.owner_task_id; + input.derived_task_ids = owner.derived_task_ids; + input.request = owner.request; + input.kind = owner_kind; + input.degradation_eligible = + owner.route.transport == RDMA && !owner.staging; + submit.owners.push_back(std::move(input)); + } + + std::vector admitted_owner_ids; + CHECK_STATUS(runtime_queue_->tryAdmit(submit, admitted_owner_ids)); + batch->queue_token = batch_token; + + batch->task_list.insert(batch->task_list.end(), prepared.tasks.size(), + TaskInfo{}); + for (const auto& task_plan : prepared.tasks) { + auto& task = batch->task_list[task_plan.task_id]; + const auto& owner = prepared.owners[task_plan.merged_task_index]; + task.failover_count = 0; + task.xport_priority = 0; + task.status = PENDING; + task.request = owner.request; + task.staging = false; + task.start_time = prepared.submit_time; + task.type = UNSPEC; + task.sub_task_id = -1; + task.device_mask = owner.route.device_mask; + if (owner.route.qp_pool) task.qp_pool = *owner.route.qp_pool; + task.derived = task_plan.task_id != owner.owner_task_id; + } + + for (size_t i = 0; i < admitted_owner_ids.size(); ++i) { + QueuedOwnerState queued; + queued.batch = batch; + queued.owner_task_id = prepared.owners[i].owner_task_id; + queued.byte_charge = prepared.owners[i].request.length; + queued.public_task_ids.push_back(prepared.owners[i].owner_task_id); + queued.public_task_ids.insert( + queued.public_task_ids.end(), + prepared.owners[i].derived_task_ids.begin(), + prepared.owners[i].derived_task_ids.end()); + queued_owners_.emplace(admitted_owner_ids[i], queued); + } + return Status::OK(); +} + +Status TransferEngineImpl::finishQueuedOwner( + QueueOwnerId owner_id, TransferStatusEnum terminal_status) { + auto queued_it = queued_owners_.find(owner_id); + if (queued_it == queued_owners_.end()) { + return Status::InvalidEntry("queued owner not found" LOC_MARK); + } + auto& queued = queued_it->second; + if (queued.in_dispatch_window) { + if (dispatch_inflight_owners_ == 0 || + dispatch_inflight_bytes_ < queued.byte_charge) { + return Status::InternalError( + "runtime dispatch window accounting underflow" LOC_MARK); + } + } + CHECK_STATUS(runtime_queue_->complete(owner_id, terminal_status)); + if (queued.in_dispatch_window) { + --dispatch_inflight_owners_; + dispatch_inflight_bytes_ -= queued.byte_charge; + queued.in_dispatch_window = false; + } + for (const auto task_id : queued.public_task_ids) { + queued.batch->task_list[task_id].status = terminal_status; + } + queued_owners_.erase(queued_it); + return Status::OK(); +} + +Status TransferEngineImpl::cancelQueuedOwner(QueueOwnerId owner_id) { + auto queued_it = queued_owners_.find(owner_id); + if (queued_it == queued_owners_.end()) { + return Status::InvalidEntry("queued owner not found" LOC_MARK); + } + if (queued_it->second.in_dispatch_window) { + return Status::InvalidEntry( + "queued owner is already dispatching" LOC_MARK); + } + CHECK_STATUS(runtime_queue_->cancel(owner_id)); + for (const auto task_id : queued_it->second.public_task_ids) { + auto& task = queued_it->second.batch->task_list[task_id]; + task.cancel_requested = true; + task.status = CANCELED; + } + queued_owners_.erase(queued_it); + return Status::OK(); +} + +Status TransferEngineImpl::retireQueueForBatch(Batch* batch) { + if (!batch || batch->queue_token == 0) return Status::OK(); + auto status = runtime_queue_->retireBatch(batch->queue_token); + if (!status.ok()) return status; + batch->queue_token = 0; + return Status::OK(); +} + +Status TransferEngineImpl::markQueuedOwnerSubmitted(QueueOwnerId owner_id) { + auto queued_it = queued_owners_.find(owner_id); + if (queued_it == queued_owners_.end()) { + return Status::InternalError("queued owner metadata missing" LOC_MARK); + } + auto& queued = queued_it->second; + if (!queued.in_dispatch_window) { + queued.in_dispatch_window = true; + ++dispatch_inflight_owners_; + dispatch_inflight_bytes_ += queued.byte_charge; + } + return Status::OK(); +} + +Status TransferEngineImpl::dispatchQueuedOwner(QueueOwnerId owner_id) { + auto queued_it = queued_owners_.find(owner_id); + if (queued_it == queued_owners_.end()) { + return Status::InternalError("queued owner metadata missing" LOC_MARK); + } + const auto queued = queued_it->second; + auto* batch = queued.batch; + auto& task = batch->task_list[queued.owner_task_id]; + task.dispatch_time = std::chrono::steady_clock::now(); + auto route = resolveTransport(task.request, 0); + task.type = route.transport; + task.device_mask = route.device_mask; + if (route.qp_pool) task.qp_pool = *route.qp_pool; + if (task.type == UNSPEC) { + return finishQueuedOwner(owner_id, FAILED); + } + + if (task.type == TCP) { + std::vector staging_params; + findStagingPolicy(task.request, staging_params); + if (!staging_params.empty() && staging_proxy_) { + task.staging = true; + // Orchestration only; the real transport submissions issued by + // ProxyManager are counted where they recurse through the + // non-staging path below. + auto status = + staging_proxy_->submit(&task, (BatchID)batch, staging_params); + if (!status.ok()) return finishQueuedOwner(owner_id, FAILED); + task.post_time = std::chrono::steady_clock::now(); + return markQueuedOwnerSubmitted(owner_id); + } + } + + if (!batch->sub_batch[task.type]) { + auto& transport = transport_list_[task.type]; + if (!transport) return finishQueuedOwner(owner_id, FAILED); + auto status = transport->allocateSubBatch(batch->sub_batch[task.type], + batch->max_size); + if (!status.ok()) return finishQueuedOwner(owner_id, FAILED); + attachProgressNotifier(batch, batch->sub_batch[task.type]); + } + + auto& transport = transport_list_[task.type]; + if (!transport) return finishQueuedOwner(owner_id, FAILED); + auto& sub_batch = batch->sub_batch[task.type]; + if (task.type == RDMA) { + sub_batch->device_mask = task.device_mask; + sub_batch->qp_pool = task.qp_pool; + } + task.sub_task_id = sub_batch->size(); + startTransportAttempt(task, task.type, std::chrono::steady_clock::now()); + auto status = transport->submitTransferTasks(sub_batch, {task.request}); + if (!status.ok()) { + finishTransportAttempt(task, FAILED, std::chrono::steady_clock::now()); + task.type = UNSPEC; + return finishQueuedOwner(owner_id, FAILED); + } + return markQueuedOwnerSubmitted(owner_id); +} + +Status TransferEngineImpl::refillDispatchWindow() { + std::lock_guard lk(progress_mutex_); + if (!runtime_queue_config_.enabled) return Status::OK(); + if (dispatch_inflight_owners_ >= + runtime_queue_config_.max_dispatch_owners || + dispatch_inflight_bytes_ >= runtime_queue_config_.max_dispatch_bytes) { + return Status::OK(); + } + + const size_t owner_budget = + runtime_queue_config_.max_dispatch_owners - dispatch_inflight_owners_; + const size_t byte_budget = + runtime_queue_config_.max_dispatch_bytes - dispatch_inflight_bytes_; + auto picked = runtime_queue_->pickForDispatch(owner_budget, byte_budget); + for (const auto owner_id : picked) { + CHECK_STATUS(dispatchQueuedOwner(owner_id)); + } + return Status::OK(); +} + +Status TransferEngineImpl::progressRuntimeQueue() { + std::lock_guard lk(progress_mutex_); + if (!runtime_queue_config_.enabled) return Status::OK(); + + CHECK_STATUS(refillDispatchWindow()); + + std::vector owner_ids; + owner_ids.reserve(queued_owners_.size()); + for (const auto& entry : queued_owners_) { + if (entry.second.in_dispatch_window) owner_ids.push_back(entry.first); + } + + bool released_window = false; + for (const auto owner_id : owner_ids) { + auto queued_it = queued_owners_.find(owner_id); + if (queued_it == queued_owners_.end()) continue; + + auto& queued = queued_it->second; + if (!queued.in_dispatch_window) continue; + auto* batch = queued.batch; + if (!batch || !alive_batches_.count((BatchID)batch)) continue; + if (queued.owner_task_id >= batch->task_list.size()) { + return Status::InternalError( + "queued owner task id out of range" LOC_MARK); + } + + auto& task = batch->task_list[queued.owner_task_id]; + auto prev_status = task.status; + TransferStatus task_status; + CHECK_STATUS(pollTaskStatus(batch, queued.owner_task_id, task_status)); + updateTaskStatusAfterPoll(batch, queued.owner_task_id, task_status, + true); + recordTaskCompletionMetrics(task, prev_status, task_status.s); + + if (task_status.s == PENDING) continue; + + CHECK_STATUS(finishQueuedOwner(owner_id, task_status.s)); + if (task_status.s == COMPLETED) + CHECK_STATUS(maybeFireSubmitHooks(batch)); + released_window = true; + } + + if (released_window) CHECK_STATUS(refillDispatchWindow()); + return Status::OK(); +} + +bool TransferEngineImpl::hasActiveRuntimeQueue() { + std::lock_guard lk(progress_mutex_); + return runtime_queue_config_.enabled && !queued_owners_.empty(); +} + +bool TransferEngineImpl::shouldQueueSubmit(const PreparedSubmit& prepared, + QueueOwnerKind owner_kind) const { + if (!runtime_queue_config_.enabled) return false; + if (owner_kind == QueueOwnerKind::StagingInternal) return true; + return std::none_of( + prepared.owners.begin(), prepared.owners.end(), + [](const PreparedSubmit::Owner& owner) { return owner.staging; }); +} + +Status TransferEngineImpl::submitTransfer( + BatchID batch_id, const std::vector& request_list, + const Notification* notifi, QueueOwnerKind owner_kind) { + Batch* batch = nullptr; + CHECK_STATUS(retainBatch(batch_id, batch)); + BatchRef batch_ref(*this, batch); + const size_t start_task_id = batch_ref.get()->task_list.size(); + PreparedSubmit prepared; + CHECK_STATUS(prepareSubmit(batch_ref.get(), request_list, prepared)); + + if (shouldQueueSubmit(prepared, owner_kind)) { + CHECK_STATUS( + enqueuePreparedSubmit(batch_ref.get(), prepared, owner_kind)); + auto dispatch_status = refillDispatchWindow(); + if (!dispatch_status.ok()) { + LOG(WARNING) << "runtime queue dispatch failed after admission: " + << dispatch_status.ToString(); + } + notifyRuntimeQueueReady(); + } else { + CHECK_STATUS(commitPreparedSubmit(batch_ref.get(), prepared)); + } + + if (notifi) { + addSubmitHook(batch_ref.get(), start_task_id, request_list, *notifi); + } + return batch_ref.release(); +} + +Status TransferEngineImpl::submitTransfer( + BatchID batch_id, const std::vector& request_list) { + return submitTransfer(batch_id, request_list, nullptr, + QueueOwnerKind::User); +} + +Status TransferEngineImpl::submitStagingTransfer( + BatchID batch_id, const std::vector& request_list) { + return submitTransfer(batch_id, request_list, nullptr, + QueueOwnerKind::StagingInternal); +} + +void TransferEngineImpl::addSubmitHook(Batch* batch, size_t start_task_id, + const std::vector& request_list, + const Notification& notifi) { + Batch::SubmitHook hook; + hook.start_task_id = start_task_id; + hook.end_task_id = start_task_id + request_list.size(); + hook.notifi = notifi; + hook.fired = false; + for (const auto& request : request_list) + hook.targets.insert(request.target_id); + batch->submit_hooks.emplace_back(std::move(hook)); +} + Status TransferEngineImpl::maybeFireSubmitHooks(Batch* batch, bool check) { for (auto& hook : batch->submit_hooks) { if (hook.fired) continue; @@ -1367,19 +1941,76 @@ Status TransferEngineImpl::maybeFireSubmitHooks(Batch* batch, bool check) { Status TransferEngineImpl::submitTransfer( BatchID batch_id, const std::vector& request_list, const Notification& notifi) { + return submitTransfer(batch_id, request_list, ¬ifi, + QueueOwnerKind::User); +} + +Status TransferEngineImpl::cancelTransfer(BatchID batch_id, size_t task_id) { if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK); - Batch* batch = (Batch*)(batch_id); - const size_t start_task_id = batch->task_list.size(); - CHECK_STATUS(submitTransfer(batch_id, request_list)); - const size_t end_task_id = start_task_id + request_list.size(); - Batch::SubmitHook hook; - hook.start_task_id = start_task_id; - hook.end_task_id = end_task_id; - hook.notifi = notifi; - hook.fired = false; - for (const auto& request : request_list) - hook.targets.insert(request.target_id); - batch->submit_hooks.emplace_back(std::move(hook)); + std::lock_guard lk(progress_mutex_); + if (!alive_batches_.count(batch_id)) { + return Status::InvalidArgument("Batch is not alive" LOC_MARK); + } + auto* batch = reinterpret_cast(batch_id); + if (task_id >= batch->task_list.size()) { + return Status::InvalidArgument("Invalid task ID" LOC_MARK); + } + + size_t owner_task_id = task_id; + if (runtime_queue_config_.enabled && batch->queue_token != 0) { + QueueOwnerId owner_id = 0; + auto resolve_status = + runtime_queue_->resolveOwner(batch->queue_token, task_id, owner_id); + if (resolve_status.ok()) { + auto queued_it = queued_owners_.find(owner_id); + if (queued_it == queued_owners_.end()) { + TransferStatusEnum public_status = PENDING; + CHECK_STATUS(runtime_queue_->getPublicStatus( + batch->queue_token, task_id, public_status)); + return public_status != PENDING + ? Status::OK() + : Status::InvalidEntry( + "queued owner metadata missing" LOC_MARK); + } + owner_task_id = queued_it->second.owner_task_id; + if (!queued_it->second.in_dispatch_window) { + CHECK_STATUS(cancelQueuedOwner(owner_id)); + CHECK_STATUS(refillDispatchWindow()); + notifyRuntimeQueueReady(); + return Status::OK(); + } + } + } + + auto& owner = batch->task_list[owner_task_id]; + if (owner.status != PENDING) return Status::OK(); + if (owner.staging) { + return Status::NotImplemented( + "staging transfer cancellation is not implemented" LOC_MARK); + } + if (owner.type == UNSPEC) { + owner.cancel_requested = true; + owner.status = CANCELED; + return Status::OK(); + } + auto& transport = transport_list_[owner.type]; + auto& sub_batch = batch->sub_batch[owner.type]; + if (!transport || !sub_batch) { + return Status::InvalidArgument("Transport not available" LOC_MARK); + } + if (!transport->supportsCancellation()) { + return Status::NotImplemented( + "selected transport does not support cancellation" LOC_MARK); + } + + CHECK_STATUS(transport->cancelTransferTask(sub_batch, owner.sub_task_id)); + // Merged public tasks share one physical transport task. Mark every alias + // so polling any of them cannot trigger failover after cancellation. + for (auto& task : batch->task_list) { + if (task.type == owner.type && task.sub_task_id == owner.sub_task_id) { + task.cancel_requested = true; + } + } return Status::OK(); } @@ -1411,16 +2042,23 @@ Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) { LOG(INFO) << "Transport failover: " << transportTypeName(prev_type) << " -> " << transportTypeName(type) << " (attempt " << task.failover_count << "/" << max_failover_attempts_ << ")"; - TENT_RECORD_TRANSPORT_FAILOVER(); + TENT_RECORD_TRANSPORT_FAILOVER(prev_type, type); auto& transport = transport_list_[type]; - if (!batch->sub_batch[type]) + if (!batch->sub_batch[type]) { CHECK_STATUS(transport->allocateSubBatch(batch->sub_batch[type], batch->max_size)); + attachProgressNotifier(batch, batch->sub_batch[type]); + } auto& sub_batch = batch->sub_batch[type]; task.sub_task_id = sub_batch->size(); task.type = type; - return transport->submitTransferTasks(sub_batch, {task.request}); + startTransportAttempt(task, type, std::chrono::steady_clock::now()); + auto status = transport->submitTransferTasks(sub_batch, {task.request}); + if (!status.ok()) { + finishTransportAttempt(task, FAILED, std::chrono::steady_clock::now()); + } + return status; } Status TransferEngineImpl::pollTaskStatus(Batch* batch, size_t task_id, @@ -1450,9 +2088,14 @@ void TransferEngineImpl::updateTaskStatusAfterPoll(Batch* batch, size_t task_id, bool allow_failover) { auto& task = batch->task_list[task_id]; task.status = task_status.s; - if (!allow_failover || task_status.s != FAILED || task.type == UNSPEC) + if (!allow_failover || task.cancel_requested || task_status.s != FAILED || + task.type == UNSPEC) return; + // The current physical transport attempt has failed even if the logical + // request will recover through failover. Close it before task.type is + // overwritten by resubmitTransferTask(). + finishTransportAttempt(task, FAILED, std::chrono::steady_clock::now()); if (resubmitTransferTask(batch, task_id).ok()) { task_status.s = PENDING; task.status = PENDING; @@ -1507,14 +2150,52 @@ Status TransferEngineImpl::getTransferStatus(BatchID batch_id, size_t task_id, Batch* batch = (Batch*)(batch_id); if (task_id >= batch->task_list.size()) return Status::InvalidArgument("Invalid task ID" LOC_MARK); - auto& task = batch->task_list[task_id]; + const size_t public_task_id = task_id; + size_t poll_task_id = task_id; + CHECK_STATUS(refillDispatchWindow()); + if (runtime_queue_config_.enabled && batch->queue_token != 0) { + QueueOwnerId owner_id = 0; + auto resolve_status = runtime_queue_->resolveOwner( + batch->queue_token, public_task_id, owner_id); + if (resolve_status.ok()) { + TransferStatusEnum public_status = PENDING; + CHECK_STATUS(runtime_queue_->getPublicStatus( + batch->queue_token, public_task_id, public_status)); + auto queued_it = queued_owners_.find(owner_id); + if (public_status != PENDING || + (queued_it != queued_owners_.end() && + !queued_it->second.in_dispatch_window)) { + task_status.s = public_status; + task_status.transferred_bytes = + public_status == COMPLETED + ? batch->task_list[public_task_id].request.length + : 0; + return Status::OK(); + } + if (batch->task_list[public_task_id].derived && + queued_it != queued_owners_.end()) { + poll_task_id = queued_it->second.owner_task_id; + } + } + } + auto& task = batch->task_list[poll_task_id]; auto prev_status = task.status; - CHECK_STATUS(pollTaskStatus(batch, task_id, task_status)); - updateTaskStatusAfterPoll(batch, task_id, task_status, + CHECK_STATUS(pollTaskStatus(batch, poll_task_id, task_status)); + updateTaskStatusAfterPoll(batch, poll_task_id, task_status, enable_auto_failover_on_poll_); + if (runtime_queue_config_.enabled && batch->queue_token != 0 && + task_status.s != PENDING) { + QueueOwnerId owner_id = 0; + auto resolve_status = runtime_queue_->resolveOwner( + batch->queue_token, public_task_id, owner_id); + if (resolve_status.ok()) { + CHECK_STATUS(finishQueuedOwner(owner_id, task_status.s)); + CHECK_STATUS(refillDispatchWindow()); + } + } // Record metrics when task transitions to terminal state - recordTaskCompletionMetrics(batch->task_list[task_id], prev_status, + recordTaskCompletionMetrics(batch->task_list[poll_task_id], prev_status, task_status.s); if (task_status.s == COMPLETED) CHECK_STATUS(maybeFireSubmitHooks(batch)); @@ -1544,6 +2225,7 @@ Status TransferEngineImpl::getBatchStatus(BatchID batch_id, std::lock_guard lk(progress_mutex_); if (!alive_batches_.count(batch_id)) return Status::InvalidArgument("Batch is not alive" LOC_MARK); + CHECK_STATUS(refillDispatchWindow()); Batch* batch = (Batch*)(batch_id); overall_status.s = PENDING; overall_status.transferred_bytes = 0; @@ -1562,6 +2244,34 @@ Status TransferEngineImpl::getBatchStatus(BatchID batch_id, auto& task = batch->task_list[task_id]; if (task.derived) continue; // This task is performed by other tasks total_tasks++; + if (runtime_queue_config_.enabled && batch->queue_token != 0) { + QueueOwnerId owner_id = 0; + auto resolve_status = runtime_queue_->resolveOwner( + batch->queue_token, task_id, owner_id); + if (resolve_status.ok()) { + TransferStatusEnum public_status = PENDING; + CHECK_STATUS(runtime_queue_->getPublicStatus( + batch->queue_token, task_id, public_status)); + auto queued_it = queued_owners_.find(owner_id); + if (public_status == PENDING) { + if (queued_it != queued_owners_.end() && + !queued_it->second.in_dispatch_window) { + continue; + } + } + if (public_status == COMPLETED) { + success_tasks++; + overall_status.transferred_bytes += task.request.length; + continue; + } + if (public_status != PENDING) { + failed_tasks++; + if (isWorse(public_status, worst_failure)) + worst_failure = public_status; + continue; + } + } + } TransferStatus task_status; if (task.status != PENDING) { if (task.status == COMPLETED) { @@ -1577,6 +2287,16 @@ Status TransferEngineImpl::getBatchStatus(BatchID batch_id, auto prev_status = task.status; CHECK_STATUS(pollTaskStatus(batch, task_id, task_status)); updateTaskStatusAfterPoll(batch, task_id, task_status, allow_failover); + if (runtime_queue_config_.enabled && batch->queue_token != 0 && + task_status.s != PENDING) { + QueueOwnerId owner_id = 0; + auto resolve_status = runtime_queue_->resolveOwner( + batch->queue_token, task_id, owner_id); + if (resolve_status.ok()) { + CHECK_STATUS(finishQueuedOwner(owner_id, task_status.s)); + CHECK_STATUS(refillDispatchWindow()); + } + } if (task_status.s == COMPLETED) { success_tasks++; @@ -1615,10 +2335,25 @@ Status TransferEngineImpl::progressBatch(BatchID batch_id, return getBatchStatus(batch_id, overall_status, true); } +Status TransferEngineImpl::getNicLoadStats( + std::vector& stats) const { + stats.clear(); + for (const auto& transport : transport_list_) { + if (transport) { + CHECK_STATUS(transport->getNicLoadStats(stats)); + } + } + return Status::OK(); +} + void TransferEngineImpl::notifyBatchMaybeReady(BatchID batch_id) { if (progress_worker_) progress_worker_->notifyBatchMaybeReady(batch_id); } +void TransferEngineImpl::notifyRuntimeQueueReady() { + if (progress_worker_) progress_worker_->notifyRuntimeQueueReady(); +} + Status TransferEngineImpl::waitTransferCompletion(BatchID batch_id) { TransferStatus xfer_status; while (true) { @@ -1666,33 +2401,124 @@ Status TransferEngineImpl::unlockStageBuffer(uint64_t addr) { void TransferEngineImpl::recordTaskCompletionMetrics( TaskInfo& task, TransferStatusEnum prev_status, TransferStatusEnum new_status) { +#if TENT_METRICS_ENABLED if (prev_status == PENDING && new_status != PENDING && !task.derived) { + auto end_time = std::chrono::steady_clock::now(); + finishTransportAttempt(task, new_status, end_time); auto start_time = task.start_time; if (start_time.time_since_epoch().count() > 0) { - auto end_time = std::chrono::steady_clock::now(); double latency_seconds = std::chrono::duration(end_time - start_time).count(); if (new_status == COMPLETED) { if (task.request.opcode == Request::READ) { TentMetrics::instance().recordReadCompleted( - task.request.length, latency_seconds); + task.type, task.request.length, latency_seconds); } else { TentMetrics::instance().recordWriteCompleted( - task.request.length, latency_seconds); + task.type, task.request.length, latency_seconds); + } + // Causal chain stage decomposition. These stage metrics stay + // attributed to the final (task.type) transport and measure the + // full request span for backward compatibility; per-attempt and + // initial-transport breakdowns live in the additive + // tent_transport_attempt_* metrics instead. + if (task.dispatch_time.time_since_epoch().count() > 0) { + double queue_wait_us = + std::chrono::duration( + task.dispatch_time - start_time) + .count(); + TENT_RECORD_STAGE_LATENCY(TentMetrics::Stage::QueueWait, + task.type, queue_wait_us); + if (task.post_time.time_since_epoch().count() > 0) { + double dispatch_us = + std::chrono::duration( + task.post_time - task.dispatch_time) + .count(); + double transport_us = + std::chrono::duration( + end_time - task.post_time) + .count(); + TENT_RECORD_STAGE_LATENCY(TentMetrics::Stage::Dispatch, + task.type, dispatch_us); + TENT_RECORD_STAGE_LATENCY(TentMetrics::Stage::Transport, + task.type, transport_us); + } } } else if (new_status == FAILED) { if (task.request.opcode == Request::READ) { - TentMetrics::instance().recordReadFailed( - task.request.length); + TentMetrics::instance().recordReadFailed(task.type); + } else { + TentMetrics::instance().recordWriteFailed(task.type); + } + } + // Observability only (RFC #2519): deadline feasibility. The + // infeasible-at-submit case (deadline already in the past when + // the transfer was submitted) is independent of whether the + // transfer ultimately completed or failed, so it is recorded for + // both outcomes. The feasible MLU ratio requires the actual + // transfer latency, so it is only recorded on COMPLETED. + if (task.request.deadline_ns != 0) { + uint64_t start_ns = static_cast( + std::chrono::duration_cast( + start_time.time_since_epoch()) + .count()); + if (task.request.deadline_ns > start_ns) { + if (new_status == COMPLETED) { + double window_seconds = + (task.request.deadline_ns - start_ns) / 1e9; + TentMetrics::instance().recordDeadlineMLU( + task.type, latency_seconds / window_seconds); + } } else { - TentMetrics::instance().recordWriteFailed( - task.request.length); + // Deadline already in the past at submit: infeasible. + // Recorded into a dedicated counter so it does not + // pollute the MLU histogram with a sentinel value. + TentMetrics::instance().recordDeadlineInfeasible(task.type); } } // Reset start_time to prevent duplicate recording task.start_time = std::chrono::steady_clock::time_point{}; } } +#endif // TENT_METRICS_ENABLED +} + +void TransferEngineImpl::startTransportAttempt( + TaskInfo& task, TransportType type, + std::chrono::steady_clock::time_point post_time) { + if (task.derived) return; + if (task.post_time.time_since_epoch().count() == 0) { + task.post_time = post_time; + } + task.attempt_post_time = post_time; + // Capture the transport now so the attempt is attributed correctly even if + // task.type is overwritten by failover before finishTransportAttempt(). + task.attempt_type = type; + task.attempt_active = true; +#if TENT_METRICS_ENABLED + TentMetrics::instance().recordTransportAttemptStarted(type, + task.request.opcode); +#else + (void)type; +#endif +} + +void TransferEngineImpl::finishTransportAttempt( + TaskInfo& task, TransferStatusEnum status, + std::chrono::steady_clock::time_point end_time) { + if (!task.attempt_active) return; + task.attempt_active = false; +#if TENT_METRICS_ENABLED + auto post_time = task.attempt_post_time; + if (post_time.time_since_epoch().count() == 0) return; + double latency_us = + std::chrono::duration(end_time - post_time).count(); + TentMetrics::instance().recordTransportAttemptFinished( + task.attempt_type, task.request.opcode, status, latency_us); +#else + (void)status; + (void)end_time; +#endif } } // namespace tent diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp index 7ecd51276e..e4ea25f520 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transport_loader.cpp @@ -41,6 +41,10 @@ #include "tent/transport/sunrise_link/sunrise_link_transport.h" #endif +#ifdef USE_TPU +#include "tent/transport/tpu/tpu_transport.h" +#endif + namespace mooncake { namespace tent { @@ -48,7 +52,8 @@ Status TransferEngineImpl::loadTransports() { if (conf_->get("transports/tcp/enable", true)) transport_list_[TCP] = std::make_shared(); - // TODO affect the end-to-end performance because it is not numa aware + // SHM is opt-in: default false because the current path is not NUMA-aware + // (see tent/config/transfer-engine.json for an example that enables it). if (conf_->get("transports/shm/enable", false)) transport_list_[SHM] = std::make_shared(); @@ -94,6 +99,11 @@ Status TransferEngineImpl::loadTransports() { } #endif +#ifdef USE_TPU + if (conf_->get("transports/tpu/enable", true)) + transport_list_[TPU] = std::make_shared(); +#endif + return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp b/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp index e26f1e78a0..5e380c9a07 100644 --- a/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp +++ b/mooncake-transfer-engine/tent/src/runtime/transport_selector.cpp @@ -17,55 +17,58 @@ #include "tent/runtime/platform.h" #include "tent/thirdparty/nlohmann/json.h" -#include -#include - #include #include +#include +#include namespace mooncake { namespace tent { -// Transport type name mapping -static const std::unordered_map kTransportNameMap = - { - {"unspec", UNSPEC}, {"rdma", RDMA}, - {"mnnvl", MNNVL}, {"shm", SHM}, - {"nvlink", NVLINK}, {"gds", GDS}, - {"io_uring", IOURING}, {"tcp", TCP}, - {"ascend", AscendDirect}, {"sunrise_link", SUNRISE_LINK}, -}; - -static const std::unordered_map - kTransportTypeNames = { - {UNSPEC, "unspec"}, {RDMA, "rdma"}, - {MNNVL, "mnnvl"}, {SHM, "shm"}, - {NVLINK, "nvlink"}, {GDS, "gds"}, - {IOURING, "io_uring"}, {TCP, "tcp"}, - {AscendDirect, "ascend"}, {SUNRISE_LINK, "sunrise_link"}, -}; - // Memory type name mapping for pattern matching static const std::string kMemoryTypeCpu = "cpu"; static const std::string kMemoryTypeCuda = "cuda"; static const std::string kMemoryTypeNpu = "npu"; static const std::string kMemoryTypeWildcard = "*"; -std::string TransportSelector::transportTypeName(TransportType type) { - auto it = kTransportTypeNames.find(type); - if (it != kTransportTypeNames.end()) { - return it->second; +static const std::unordered_map kIntentTypeNameMap = { + {"intent_unspec", IntentType::INTENT_UNSPEC}, + {"unspec", IntentType::INTENT_UNSPEC}, + {"foreground_get", IntentType::FOREGROUND_GET}, + {"background_prefetch", IntentType::BACKGROUND_PREFETCH}, + {"migration", IntentType::MIGRATION}, + {"checkpoint", IntentType::CHECKPOINT}, + {"weight_loading", IntentType::WEIGHT_LOADING}, + {"staging_internal", IntentType::STAGING_INTERNAL}, +}; + +static std::optional parseIntentType(const json& value) { + if (value.is_string()) { + auto name = value.get(); + std::transform(name.begin(), name.end(), name.begin(), + [](unsigned char c) { return std::tolower(c); }); + auto it = kIntentTypeNameMap.find(name); + if (it != kIntentTypeNameMap.end()) return it->second; + return std::nullopt; } - return "unknown"; -} -TransportType TransportSelector::parseTransportType(const std::string& str) { - auto it = kTransportNameMap.find(str); - if (it != kTransportNameMap.end()) { - return it->second; + if (value.is_number_unsigned()) { + const auto raw = value.get(); + if (raw <= static_cast(IntentType::STAGING_INTERNAL)) { + return static_cast(raw); + } + return std::nullopt; } - LOG(WARNING) << "Unknown transport type: " << str; - return UNSPEC; + + if (value.is_number_integer()) { + const auto raw = value.get(); + if (raw >= static_cast(IntentType::INTENT_UNSPEC) && + raw <= static_cast(IntentType::STAGING_INTERNAL)) { + return static_cast(raw); + } + } + + return std::nullopt; } std::vector TransportSelector::getDefaultPolicies() { @@ -74,14 +77,18 @@ std::vector TransportSelector::getDefaultPolicies() { { "file_storage", SegmentType::File, - std::nullopt, // same_machine doesn't matter for file - std::nullopt, // local_memory_pattern - std::nullopt, // remote_memory_pattern - std::nullopt, // min_size - std::nullopt, // max_size - std::nullopt, // priority - {}, // devices (empty = all devices) - {GDS, IOURING} // File segment priority (original: GDS -> IOURING) + std::nullopt, // same_machine doesn't matter for file + std::nullopt, // local_memory_pattern + std::nullopt, // remote_memory_pattern + std::nullopt, // min_size + std::nullopt, // max_size + std::nullopt, // priority + {}, // devices (empty = all devices) + {GDS, IOURING}, // File priority (original: GDS -> IOURING) + std::nullopt, // service_level + std::nullopt, // traffic_class + std::nullopt, // qp_pool + std::nullopt // intent_type }, { "memory_default", @@ -89,11 +96,15 @@ std::vector TransportSelector::getDefaultPolicies() { std::nullopt, // any machine std::nullopt, // any local memory std::nullopt, // any remote memory - std::nullopt, // any size - std::nullopt, // min_priority + std::nullopt, // min_size + std::nullopt, // max_size + std::nullopt, // priority {}, // devices (empty = all devices) - {} // Empty priority = use buffer_transports order (original - // behavior) + {}, // Empty = use buffer_transports order + std::nullopt, // service_level + std::nullopt, // traffic_class + std::nullopt, // qp_pool + std::nullopt // intent_type }, }; } @@ -184,6 +195,19 @@ void TransportSelector::loadPolicies() { policy.priority = std::nullopt; } + // Parse the optional business-intent filter. An invalid value skips the + // entire policy instead of turning it into a catch-all rule, which + // would silently broaden its authorization scope. + if (policy_json.contains("intent_type")) { + auto intent = parseIntentType(policy_json["intent_type"]); + if (!intent.has_value()) { + LOG(WARNING) + << "Skip policy " << policy.name << ": invalid intent_type"; + continue; + } + policy.intent_type = *intent; + } + // Parse devices (optional) if (policy_json.contains("devices")) { for (const auto& device_name : policy_json["devices"]) { @@ -198,14 +222,53 @@ void TransportSelector::loadPolicies() { if (policy_json.contains("transports")) { for (const auto& transport_str : policy_json["transports"]) { if (!transport_str.is_string()) continue; - TransportType type = - parseTransportType(transport_str.get()); + const auto name = transport_str.get(); + TransportType type = parseTransportType(name); + if (type == UNSPEC && name != "unspec") { + LOG(WARNING) << "Unknown transport type: " << name; + } if (type != UNSPEC) { policy.transports.push_back(type); } } } + // Parse link-layer QoS attributes (RFC #2519 / #2568, step 1: stored + // only, not yet applied to QPs). Out-of-range values are ignored so a + // bad config never breaks selection. + if (policy_json.contains("service_level")) { + int sl = policy_json.value("service_level", -1); + if (sl >= 0 && sl <= 15) { + policy.service_level = sl; + } else { + LOG(WARNING) << "Ignore service_level in policy " << policy.name + << ", value " << sl << " out of range (0-15)"; + } + } + if (policy_json.contains("traffic_class")) { + int tc = policy_json.value("traffic_class", -1); + if (tc >= 0 && tc <= 255) { + policy.traffic_class = tc; + } else { + LOG(WARNING) << "Ignore traffic_class in policy " << policy.name + << ", value " << tc << " out of range (0-255)"; + } + } + // Reserved for step 2 (per-class QP pools); parsed for forward schema + // compatibility, no effect yet. + if (policy_json.contains("qp_pool")) { + auto& qp = policy_json["qp_pool"]; + if (!qp.is_string()) { + LOG(WARNING) << "Ignore qp_pool in policy " << policy.name + << ", expected a string"; + } else { + auto value = qp.get(); + // Treat an empty string the same as unset (use default pool) + // so a blank config value doesn't look like an explicit pool. + if (!value.empty()) policy.qp_pool = std::move(value); + } + } + policies_.push_back(std::move(policy)); LOG(INFO) << "Loaded transport policy: " << policy.name << " (segment_type=" << segment_type_str @@ -236,6 +299,9 @@ bool TransportSelector::matchesMemoryPattern(const std::string& pattern, case MTYPE_ROCM: type_str = "rocm"; break; + case MTYPE_TPU: + type_str = "tpu"; + break; default: type_str = "unknown"; break; @@ -300,6 +366,13 @@ bool TransportSelector::matchesPolicy(const SelectionPolicy& policy, } } + // Policies without an intent filter retain the historical catch-all + // behavior. Intent-specific policies require an exact match. + if (policy.intent_type.has_value() && + context.intent_type != policy.intent_type.value()) { + return false; + } + return true; } @@ -319,15 +392,21 @@ bool TransportSelector::isTransportAvailable( } // Special constraints - if ((type == NVLINK || type == SHM) && !context.same_machine) { - return false; // NVLINK and SHM only work on same machine + if ((type == NVLINK || type == SHM || type == TPU) && + !context.same_machine) { + // NVLINK/SHM only work on same machine; TPU is a local-stage-only + // executor (HBM<->host), so it must never be picked for a remote hop. + return false; } const auto& caps = transport->capabilities(); - // Helper to check if memory type is GPU/NPU + // Helper to check if memory type is a device (GPU/NPU/TPU). TPU is included + // so its device<->host staging hop routes to TpuTransport (gpu_to_dram / + // dram_to_gpu); it never satisfies gpu_to_gpu, so cross-node TPU traffic is + // always staged through host DRAM. auto is_gpu = [](MemoryType t) { - return t == MTYPE_CUDA || t == MTYPE_ROCM; + return t == MTYPE_CUDA || t == MTYPE_ROCM || t == MTYPE_TPU; }; // For file segments, check file-specific capabilities (original logic) @@ -384,6 +463,13 @@ SelectionResult TransportSelector::select( return result; // UNSPEC, all devices } + // Carry the matched policy's link-layer QoS out to the caller (RFC #2519 / + // #2568, step 1). These are plumbed but not yet applied at QP setup; that + // is the per-class QP pool follow-up (step 2). + result.service_level = matching_policy->service_level; + result.traffic_class = matching_policy->traffic_class; + result.qp_pool = matching_policy->qp_pool; + // Convert device names to mask result.device_mask = ~0ULL; // Default: all devices if (!matching_policy->devices.empty() && topology_) { diff --git a/mooncake-transfer-engine/tent/src/transfer_engine.cpp b/mooncake-transfer-engine/tent/src/transfer_engine.cpp index 5628a45c93..3b8a6387d9 100644 --- a/mooncake-transfer-engine/tent/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/tent/src/transfer_engine.cpp @@ -140,6 +140,10 @@ Status TransferEngine::submitTransfer(BatchID batch_id, return impl_->submitTransfer(batch_id, request_list, notifi); } +Status TransferEngine::cancelTransfer(BatchID batch_id, size_t task_id) { + return impl_->cancelTransfer(batch_id, task_id); +} + Status TransferEngine::sendNotification(SegmentID target_id, const Notification& notifi) { return impl_->sendNotification(target_id, notifi); @@ -174,5 +178,9 @@ Status TransferEngine::progressBatch(BatchID batch_id, return impl_->progressBatch(batch_id, overall_status); } +Status TransferEngine::getNicLoadStats(std::vector& stats) const { + return impl_->getNicLoadStats(stats); +} + } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp b/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp index 150df8b8e6..9a8fe1d2b9 100644 --- a/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp +++ b/mooncake-transfer-engine/tent/src/transfer_engine_c.cpp @@ -308,7 +308,7 @@ void tent_free_notifs(tent_notifi_info* info) { int tent_task_status(tent_engine_t engine, tent_batch_id_t batch_id, size_t task_id, tent_status_t* xfer_status) { CHECK_POINTER(engine); - CHECK_POINTER(batch_id); + if (!batch_id) return -1; CHECK_POINTER(xfer_status); mooncake::tent::TransferStatus internal_status; auto status = @@ -322,10 +322,22 @@ int tent_task_status(tent_engine_t engine, tent_batch_id_t batch_id, return 0; } +int tent_cancel_task(tent_engine_t engine, tent_batch_id_t batch_id, + size_t task_id) { + CHECK_POINTER(engine); + if (!batch_id) return -1; + auto status = CAST(engine)->cancelTransfer(batch_id, task_id); + if (!status.ok()) { + LOG(ERROR) << "tent_cancel_task: " << status.ToString(); + return -1; + } + return 0; +} + int tent_overall_status(tent_engine_t engine, tent_batch_id_t batch_id, tent_status_t* xfer_status) { CHECK_POINTER(engine); - CHECK_POINTER(batch_id); + if (!batch_id) return -1; CHECK_POINTER(xfer_status); mooncake::tent::TransferStatus internal_status; auto status = CAST(engine)->getTransferStatus(batch_id, internal_status); @@ -463,7 +475,7 @@ int tent_register_memory_batch_ex(tent_engine_t engine, void** addrs, int tent_task_status_list(tent_engine_t engine, tent_batch_id_t batch_id, tent_status_t* statuses, size_t* count) { CHECK_POINTER(engine); - CHECK_POINTER(batch_id); + if (!batch_id) return -1; CHECK_POINTER(statuses); CHECK_POINTER(count); std::vector status_list; @@ -480,3 +492,25 @@ int tent_task_status_list(tent_engine_t engine, tent_batch_id_t batch_id, *count = status_list.size(); return 0; } + +int tent_get_nic_load_stats(tent_engine_t engine, tent_nic_load_stat_t* stats, + size_t* count) { + CHECK_POINTER(engine); + CHECK_POINTER(stats); + CHECK_POINTER(count); + std::vector native_stats; + auto status = CAST(engine)->getNicLoadStats(native_stats); + if (!status.ok()) { + LOG(ERROR) << "tent_get_nic_load_stats: " << status.ToString(); + return -1; + } + size_t to_copy = std::min(native_stats.size(), *count); + for (size_t i = 0; i < to_copy; ++i) { + snprintf(stats[i].device_name, sizeof(stats[i].device_name), "%s", + native_stats[i].device_name.c_str()); + stats[i].inflight_bytes = native_stats[i].inflight_bytes; + stats[i].ewma_bandwidth_bps = native_stats[i].ewma_bandwidth_bps; + } + *count = native_stats.size(); + return 0; +} diff --git a/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt index e73f7cb063..825f3544b9 100644 --- a/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/src/transport/CMakeLists.txt @@ -8,6 +8,8 @@ add_subdirectory(io_uring) add_subdirectory(bufio) add_subdirectory(ascend) add_subdirectory(sunrise_link) +add_subdirectory(tpu) +add_subdirectory(ub) add_library(tent_transport_all INTERFACE) foreach( @@ -20,8 +22,10 @@ foreach( tent_xport_rdma tent_xport_shm tent_xport_tcp + tent_xport_ub tent_xport_ascend_direct - tent_xport_sunrise_link) + tent_xport_sunrise_link + tent_xport_tpu) if(TARGET ${tgt}) target_link_libraries(tent_transport_all INTERFACE ${tgt}) endif() diff --git a/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp b/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp index dd18c4eee4..dd93a23254 100644 --- a/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/ascend/ascend_direct_transport.cpp @@ -160,9 +160,12 @@ Status AscendDirectTransport::initHixl(const std::shared_ptr &conf) { auto hixl_name = host_ip + ":" + std::to_string(port); local_hixl_name_ = hixl_name; - auto segment = metadata_->segmentManager().getLocal(); - auto &detail = std::get(segment->detail); - detail.device_attrs["hixl_name"] = hixl_name; + CHECK_STATUS(metadata_->segmentManager().updateLocal( + [&](SegmentDesc &segment) -> Status { + auto &detail = std::get(segment.detail); + detail.device_attrs["hixl_name"] = hixl_name; + return Status::OK(); + })); hixl_ = std::make_unique(); if (!hixl_) return Status::InternalError("Create hixl failed."); @@ -380,9 +383,9 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id, return Status::InvalidArgument("Invalid task id" LOC_MARK); } auto &task = hixl_batch->task_list[task_id]; - status = TransferStatus{task.status_word, task.transferred_bytes}; if (task.status_word == TransferStatusEnum::PENDING) { if (task.req_handle == nullptr) { + status = TransferStatus{task.status_word, task.transferred_bytes}; return Status::OK(); } std::lock_guard lock(req_mutex_); @@ -396,6 +399,7 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id, if (--req_map_[task.req_handle].second == 0) { req_map_.erase(task.req_handle); } + status = TransferStatus{task.status_word, task.transferred_bytes}; return Status::OK(); } uint64_t current_ts = getCurrentTimeInNano(); @@ -406,6 +410,7 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id, req_map_[task.req_handle] = std::make_pair(task.status_word, task.batch_size - 1); } + status = TransferStatus{task.status_word, task.transferred_bytes}; return Status::OK(); } hixl::TransferStatus xfer_status; @@ -415,6 +420,7 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id, xfer_status = hixl::TransferStatus::FAILED; } if (xfer_status == hixl::TransferStatus::WAITING) { + status = TransferStatus{task.status_word, task.transferred_bytes}; return Status::OK(); } if (xfer_status == hixl::TransferStatus::COMPLETED) { @@ -427,9 +433,14 @@ Status AscendDirectTransport::getTransferStatus(SubBatchRef batch, int task_id, disconnect(task.remote_hixl, 10); task.status_word = TransferStatusEnum::FAILED; } - req_map_[task.req_handle] = - std::make_pair(task.status_word, task.batch_size - 1); + if (task.batch_size > 1) { + req_map_[task.req_handle] = + std::make_pair(task.status_word, task.batch_size - 1); + } } + // Read status AFTER the poll so a just-observed completion/failure is + // reported on this call rather than one poll cycle late. + status = TransferStatus{task.status_word, task.transferred_bytes}; return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/transport/gds/gds_transport.cpp b/mooncake-transfer-engine/tent/src/transport/gds/gds_transport.cpp index c2c90b8f99..6980d98c15 100644 --- a/mooncake-transfer-engine/tent/src/transport/gds/gds_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/gds/gds_transport.cpp @@ -279,7 +279,9 @@ Status GdsTransport::submitTransferTasks( GdsFileContext* context = findFileContext(request.target_id); if (!context || !context->ready()) return Status::InvalidArgument("Invalid remote segment" LOC_MARK); - IOParamRange range{gds_batch->io_params.size(), 0}; + size_t task_id = gds_batch->io_param_ranges.size(); + IOParamRange range; + range.base = gds_batch->io_params.size(); for (size_t offset = 0; offset < request.length; offset += kMaxSliceSize) { size_t length = std::min(kMaxSliceSize, request.length - offset); @@ -287,7 +289,8 @@ Status GdsTransport::submitTransferTasks( params.mode = CUFILE_BATCH; params.opcode = (request.opcode == Request::READ) ? CUFILE_READ : CUFILE_WRITE; - params.cookie = (void*)0; + params.cookie = + reinterpret_cast(static_cast(task_id)); params.u.batch.devPtr_base = request.source; params.u.batch.devPtr_offset = offset; params.u.batch.file_offset = request.target_offset + offset; @@ -315,26 +318,38 @@ Status GdsTransport::getTransferStatus(SubBatchRef batch, int task_id, unsigned num_tasks = gds_batch->io_param_ranges.size(); if (task_id < 0 || task_id >= (int)num_tasks) return Status::InvalidArgument("Invalid task ID"); - auto range = gds_batch->io_param_ranges[task_id]; + unsigned num_events = static_cast(gds_batch->io_params.size()); auto result = - cuFileBatchIOGetStatus(gds_batch->batch_handle->handle, 0, &num_tasks, + cuFileBatchIOGetStatus(gds_batch->batch_handle->handle, 0, &num_events, gds_batch->io_events.data(), nullptr); if (result.err != CU_FILE_SUCCESS) return Status::InternalError( std::string("Failed to get GDS batch status: Code ") + std::to_string(result.err) + LOC_MARK); - status.s = PENDING; - size_t complete_count = 0; - for (size_t index = range.base; index < range.base + range.count; ++index) { + + for (size_t index = 0; index < num_events; ++index) { auto& event = gds_batch->io_events[index]; + auto event_task_id = reinterpret_cast(event.cookie); + if (event_task_id >= gds_batch->io_param_ranges.size()) { + LOG(ERROR) << "Invalid GDS batch IO cookie: " << event_task_id; + continue; + } + + auto& range = gds_batch->io_param_ranges[event_task_id]; auto s = parseTransferStatus(event.status); - if (s == COMPLETED) - complete_count++; - else if (s != PENDING) - status.s = s; - status.transferred_bytes += event.ret; + if (s == COMPLETED) { + range.complete_count++; + range.transferred_bytes += event.ret; + } else if (s != PENDING) { + range.status = s; + } + } + + auto& range = gds_batch->io_param_ranges[task_id]; + if (range.complete_count == range.count) { + range.status = COMPLETED; } - if (complete_count == range.count) status.s = COMPLETED; + status = TransferStatus{range.status, range.transferred_bytes}; return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/transport/io_uring/io_uring_transport.cpp b/mooncake-transfer-engine/tent/src/transport/io_uring/io_uring_transport.cpp index bf7364b2d5..35e8c0927c 100644 --- a/mooncake-transfer-engine/tent/src/transport/io_uring/io_uring_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/io_uring/io_uring_transport.cpp @@ -247,26 +247,28 @@ Status IOUringTransport::getTransferStatus(SubBatchRef batch, int task_id, return Status::InternalError( std::string("io_uring_peek_cqe failed: ") + strerror(-err)); } - auto task = (IOUringTask*)cqe->user_data; - if (task) { + auto cqe_task = (IOUringTask*)cqe->user_data; + if (cqe_task) { if (cqe->res < 0) { LOG(INFO) << "Received an event with error code " << cqe->res; - task->status_word = TransferStatusEnum::FAILED; + cqe_task->status_word = TransferStatusEnum::FAILED; } else { - if (task->buffer) { - if (task->request.opcode == Request::READ) - Platform::getLoader().copy(task->request.source, - task->buffer, - task->request.length); - - free(task->buffer); - task->buffer = nullptr; + if (cqe_task->buffer) { + if (cqe_task->request.opcode == Request::READ) + Platform::getLoader().copy(cqe_task->request.source, + cqe_task->buffer, + cqe_task->request.length); + + free(cqe_task->buffer); + cqe_task->buffer = nullptr; } - task->status_word = TransferStatusEnum::COMPLETED; - task->transferred_bytes = task->request.length; + cqe_task->status_word = TransferStatusEnum::COMPLETED; + cqe_task->transferred_bytes = cqe_task->request.length; } } io_uring_cqe_seen(&io_uring_batch->ring, cqe); + batch->notifyProgress(); + status = TransferStatus{task.status_word, task.transferred_bytes}; } return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp index 414a792c0b..2e869018ea 100644 --- a/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/mnnvl/mnnvl_transport.cpp @@ -145,8 +145,8 @@ Status MnnvlTransport::allocateSubBatch(SubBatchRef &batch, size_t max_size) { batch = mnnvl_batch; mnnvl_batch->task_list.reserve(max_size); mnnvl_batch->max_size = max_size; - CHECK_STATUS(platform_->getStreamFromPool(mnnvl_batch->sync_stream)); - CHECK_STATUS(platform_->getStreamFromPool(mnnvl_batch->async_stream)); + // Streams are created lazily in submitTransferTasks where the correct + // GPU device can be inferred from request source pointers. return Status::OK(); } @@ -167,13 +167,49 @@ Status MnnvlTransport::submitTransferTasks( auto mnnvl_batch = dynamic_cast(batch); if (!mnnvl_batch) return Status::InvalidArgument("Invalid MNNVL sub-batch" LOC_MARK); + if (request_list.size() + mnnvl_batch->task_list.size() > mnnvl_batch->max_size) return Status::TooManyRequests("Exceed batch capacity" LOC_MARK); + + // Get local segment for buffer lookup + auto &segment_manager = metadata_->segmentManager(); + // Owning reference: keeps the snapshot alive while we read through it. + SegmentDescRef local_segment = segment_manager.getLocal(); + if (!local_segment) + return Status::InternalError("Local segment not found" LOC_MARK); + + // Determine device for this batch and validate all requests + int batch_device_id = -1; std::vector new_tasks; + for (auto &request : request_list) { + // Find the buffer this source pointer belongs to + BufferDesc *buf = local_segment->findBuffer( + reinterpret_cast(request.source), request.length); + if (!buf) { + return Status::InvalidArgument( + "Unregistered buffer: source pointer not in any registered " + "buffer" LOC_MARK); + } + + // Parse device ID from buffer location (e.g., "cuda:0" -> 0, "cpu" -> + // -1) + int device_id = LocationParser(buf->location).index(); + + // Capture the first GPU device encountered for stream creation. + // Mixed-GPU batches use the first GPU's stream and rely on CUDA P2P + // for cross-device access (same behavior as pre-#2569 code). + // A future refactor could group requests by device and dispatch to + // per-device streams, but that requires SubBatch structure changes. + if (batch_device_id < 0 && device_id >= 0) { + batch_device_id = device_id; + } + + // Create and populate task mnnvl_batch->task_list.push_back(MnnvlTask{}); auto &task = mnnvl_batch->task_list[mnnvl_batch->task_list.size() - 1]; + uint64_t target_addr = request.target_offset; if (request.target_id != LOCAL_SEGMENT_ID) { auto status = relocateSharedMemoryAddress( @@ -184,11 +220,27 @@ Status MnnvlTransport::submitTransferTasks( return status; } } + task.target_addr = target_addr; task.request = request; task.status_word = TransferStatusEnum::PENDING; new_tasks.push_back(&task); } + + // Get or create streams for this batch's device + if (!mnnvl_batch->async_stream.get()) { + int stream_device = batch_device_id; + if (stream_device < 0) { + // CPU-only batch: use current CUDA device + cudaGetDevice(&stream_device); + } + CHECK_STATUS(platform_->getStreamFromPool(mnnvl_batch->sync_stream, + stream_device)); + CHECK_STATUS(platform_->getStreamFromPool(mnnvl_batch->async_stream, + stream_device)); + mnnvl_batch->stream_device_id = stream_device; + } + startTransfer(new_tasks, mnnvl_batch); return Status::OK(); } @@ -261,11 +313,39 @@ void MnnvlTransport::startTransfer(std::vector &tasks, return; } + // Save and set device to match the stream's device to ensure event + // creation and recording happen on the correct device (fix for #2722). + int saved_device = -1; + if (batch->stream_device_id >= 0) { + auto err = cudaGetDevice(&saved_device); + if (err != cudaSuccess) { + LOG(ERROR) << "MnnvlTransport: cudaGetDevice failed: " + << cudaGetErrorString(err); + for (auto *task : tasks) + task->status_word = TransferStatusEnum::FAILED; + return; + } + if (saved_device != batch->stream_device_id) { + err = cudaSetDevice(batch->stream_device_id); + if (err != cudaSuccess) { + LOG(ERROR) << "MnnvlTransport: cudaSetDevice failed: " + << cudaGetErrorString(err); + for (auto *task : tasks) + task->status_word = TransferStatusEnum::FAILED; + return; + } + } + } + cudaEvent_t event; auto event_err = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); if (event_err != cudaSuccess) { LOG(ERROR) << "MnnvlTransport: cudaEventCreateWithFlags failed: " << cudaGetErrorString(event_err); + // Restore device before returning + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } for (auto *task : tasks) task->status_word = TransferStatusEnum::FAILED; return; } @@ -274,9 +354,19 @@ void MnnvlTransport::startTransfer(std::vector &tasks, LOG(ERROR) << "MnnvlTransport: cudaEventRecord failed: " << cudaGetErrorString(record_err); cudaEventDestroy(event); + // Restore device before returning + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } for (auto *task : tasks) task->status_word = TransferStatusEnum::FAILED; return; } + + // Restore original device + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } + batch->completion_events.push_back(event); for (auto *task : tasks) task->completion_event = event; } @@ -288,7 +378,6 @@ Status MnnvlTransport::getTransferStatus(SubBatchRef batch, int task_id, return Status::InvalidArgument("Invalid task id" LOC_MARK); } auto &task = mnnvl_batch->task_list[task_id]; - status = TransferStatus{task.status_word, task.transferred_bytes}; if (task.status_word == TransferStatusEnum::PENDING) { auto err = cudaEventQuery(task.completion_event); if (err == cudaSuccess) { @@ -298,6 +387,9 @@ Status MnnvlTransport::getTransferStatus(SubBatchRef batch, int task_id, task.status_word = TransferStatusEnum::FAILED; } } + // Read status AFTER the poll so a just-observed completion/failure is + // reported on this call rather than one poll cycle late. + status = TransferStatus{task.status_word, task.transferred_bytes}; return Status::OK(); } @@ -501,9 +593,11 @@ Status MnnvlTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, RWSpinlock::WriteGuard guard(relocate_lock_); BufferDesc *buffer; + // Owning reference: `buffer` is used after the lambda returns. + SegmentDescRef pin; auto &segment_manager = metadata_->segmentManager(); - CHECK_STATUS( - segment_manager.withCachedSegment(target_id, [&](SegmentDesc *segment) { + CHECK_STATUS(segment_manager.withCachedSegment( + target_id, pin, [&](SegmentDesc *segment) { buffer = segment->findBuffer(dest_addr, length); if (!buffer || buffer->mnnvl_handle.empty()) return Status::NeedsRefreshCache( diff --git a/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp b/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp index ee4602953a..f3de90d7f3 100644 --- a/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/nvlink/nvlink_transport.cpp @@ -106,8 +106,8 @@ Status NVLinkTransport::allocateSubBatch(SubBatchRef& batch, size_t max_size) { batch = shm_batch; shm_batch->task_list.reserve(max_size); shm_batch->max_size = max_size; - CHECK_STATUS(platform_->getStreamFromPool(shm_batch->sync_stream)); - CHECK_STATUS(platform_->getStreamFromPool(shm_batch->async_stream)); + // Streams are created lazily in submitTransferTasks where the correct + // GPU device can be inferred from request source pointers. return Status::OK(); } @@ -128,23 +128,75 @@ Status NVLinkTransport::submitTransferTasks( auto shm_batch = dynamic_cast(batch); if (!shm_batch) return Status::InvalidArgument("Invalid NVLink sub-batch" LOC_MARK); + if (request_list.size() + shm_batch->task_list.size() > shm_batch->max_size) return Status::TooManyRequests("Exceed batch capacity" LOC_MARK); + + // Get local segment for buffer lookup + auto& segment_manager = metadata_->segmentManager(); + // Owning reference: keeps the snapshot alive while we read through it. + SegmentDescRef local_segment = segment_manager.getLocal(); + if (!local_segment) + return Status::InternalError("Local segment not found" LOC_MARK); + + // Determine device for this batch and validate all requests + int batch_device_id = -1; std::vector new_tasks; + for (auto& request : request_list) { + // Find the buffer this source pointer belongs to + BufferDesc* buf = local_segment->findBuffer( + reinterpret_cast(request.source), request.length); + if (!buf) { + return Status::InvalidArgument( + "Unregistered buffer: source pointer not in any registered " + "buffer" LOC_MARK); + } + + // Parse device ID from buffer location (e.g., "cuda:0" -> 0, "cpu" -> + // -1) + int device_id = LocationParser(buf->location).index(); + + // Capture the first GPU device encountered for stream creation. + // Mixed-GPU batches use the first GPU's stream and rely on CUDA P2P + // for cross-device access (same behavior as pre-#2569 code). + // A future refactor could group requests by device and dispatch to + // per-device streams, but that requires SubBatch structure changes. + if (batch_device_id < 0 && device_id >= 0) { + batch_device_id = device_id; + } + + // Create and populate task shm_batch->task_list.push_back(NVLinkTask{}); auto& task = shm_batch->task_list[shm_batch->task_list.size() - 1]; + uint64_t target_addr = request.target_offset; if (request.target_id != LOCAL_SEGMENT_ID) { auto status = relocateSharedMemoryAddress( target_addr, request.length, request.target_id); if (!status.ok()) return status; } + task.target_addr = target_addr; task.request = request; task.status_word = TransferStatusEnum::PENDING; new_tasks.push_back(&task); } + + // Get or create streams for this batch's device + if (!shm_batch->async_stream.get()) { + int stream_device = batch_device_id; + if (stream_device < 0) { + // CPU-only batch: use current CUDA device + cudaGetDevice(&stream_device); + } + CHECK_STATUS(platform_->getStreamFromPool(shm_batch->sync_stream, + stream_device)); + CHECK_STATUS(platform_->getStreamFromPool(shm_batch->async_stream, + stream_device)); + shm_batch->stream_device_id = stream_device; + } + startTransfer(new_tasks, shm_batch); return Status::OK(); } @@ -222,11 +274,39 @@ void NVLinkTransport::startTransfer(std::vector& tasks, } } + // Save and set device to match the stream's device to ensure event + // creation and recording happen on the correct device (fix for #2722). + int saved_device = -1; + if (batch->stream_device_id >= 0) { + auto err = cudaGetDevice(&saved_device); + if (err != cudaSuccess) { + LOG(ERROR) << "NVLinkTransport: cudaGetDevice failed: " + << cudaGetErrorString(err); + for (auto* task : tasks) + task->status_word = TransferStatusEnum::FAILED; + return; + } + if (saved_device != batch->stream_device_id) { + err = cudaSetDevice(batch->stream_device_id); + if (err != cudaSuccess) { + LOG(ERROR) << "NVLinkTransport: cudaSetDevice failed: " + << cudaGetErrorString(err); + for (auto* task : tasks) + task->status_word = TransferStatusEnum::FAILED; + return; + } + } + } + cudaEvent_t event; auto event_err = cudaEventCreateWithFlags(&event, cudaEventDisableTiming); if (event_err != cudaSuccess) { LOG(ERROR) << "NVLinkTransport: cudaEventCreateWithFlags failed: " << cudaGetErrorString(event_err); + // Restore device before returning + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } for (auto* task : tasks) task->status_word = TransferStatusEnum::FAILED; return; } @@ -235,9 +315,19 @@ void NVLinkTransport::startTransfer(std::vector& tasks, LOG(ERROR) << "NVLinkTransport: cudaEventRecord failed: " << cudaGetErrorString(record_err); cudaEventDestroy(event); + // Restore device before returning + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } for (auto* task : tasks) task->status_word = TransferStatusEnum::FAILED; return; } + + // Restore original device + if (saved_device >= 0 && saved_device != batch->stream_device_id) { + cudaSetDevice(saved_device); + } + batch->completion_events.push_back(event); for (auto* task : tasks) task->completion_event = event; } @@ -249,7 +339,6 @@ Status NVLinkTransport::getTransferStatus(SubBatchRef batch, int task_id, return Status::InvalidArgument("Invalid task id" LOC_MARK); } auto& task = shm_batch->task_list[task_id]; - status = TransferStatus{task.status_word, task.transferred_bytes}; if (task.status_word == TransferStatusEnum::PENDING) { auto err = cudaEventQuery(task.completion_event); if (err == cudaSuccess) { @@ -259,6 +348,9 @@ Status NVLinkTransport::getTransferStatus(SubBatchRef batch, int task_id, task.status_word = TransferStatusEnum::FAILED; } } + // Read status AFTER the poll so a just-observed completion/failure is + // reported on this call rather than one poll cycle late. + status = TransferStatus{task.status_word, task.transferred_bytes}; return Status::OK(); } @@ -376,8 +468,11 @@ Status NVLinkTransport::removeMemoryBuffer(BufferDesc& desc) { std::lock_guard lock(register_mutex_); registered_base_addrs_.erase(key); } - } else if (location.type() == "cpu" && host_register_) { - CHECK_CUDA(cudaHostUnregister((void*)desc.addr)); + } else if (location.type() == "cpu" || + location.type() == kWildcardLocation) { + if (host_register_) { + CHECK_CUDA(cudaHostUnregister((void*)desc.addr)); + } } desc.shm_path.clear(); return Status::OK(); @@ -404,9 +499,11 @@ Status NVLinkTransport::relocateSharedMemoryAddress(uint64_t& dest_addr, RWSpinlock::WriteGuard guard(relocate_lock_); BufferDesc* buffer; + // Owning reference: `buffer` is used after the lambda returns. + SegmentDescRef pin; auto& segment_manager = metadata_->segmentManager(); - CHECK_STATUS( - segment_manager.withCachedSegment(target_id, [&](SegmentDesc* segment) { + CHECK_STATUS(segment_manager.withCachedSegment( + target_id, pin, [&](SegmentDesc* segment) { buffer = segment->findBuffer(dest_addr, length); if (!buffer || buffer->shm_path.empty()) return Status::NeedsRefreshCache( diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/buffers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/buffers.cpp index 0c350890d1..72fe68d9a2 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/buffers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/buffers.cpp @@ -141,6 +141,11 @@ Status LocalBufferManager::addBufferInternal(BufferDesc& desc, context->registerMemReg((void*)desc.addr, desc.length, access); } } + // NicID-keyed like context_list_, not compacted: slice dispatch subscripts + // these with a dev_id from the topology, so gaps must stay in place. + // Devices with no context contribute a zero key that is never selected. + desc.lkey.assign(context_list_.size(), 0); + desc.rkey.assign(context_list_.size(), 0); for (size_t id = 0; id < context_list_.size(); ++id) { if (!context_list_[id]) continue; if (!mem_reg_list[id]) { @@ -149,8 +154,8 @@ Status LocalBufferManager::addBufferInternal(BufferDesc& desc, } staging.mem_reg_map[context_list_[id]] = mem_reg_list[id]; auto keys = context_list_[id]->queryMemRegKey(mem_reg_list[id]); - desc.lkey.push_back(keys.first); - desc.rkey.push_back(keys.second); + desc.lkey[id] = keys.first; + desc.rkey[id] = keys.second; } staging.options = options; RWSpinlock::WriteGuard guard(lock_); diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp index 31ea13d4eb..d96081dc6f 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/context.cpp @@ -443,7 +443,11 @@ int RdmaContext::disable() { LOG(WARNING) << "RDMA context " << name() << " has been deconstructed"; return 0; } - endpoint_store_->clear(); + if (endpoint_store_->clear()) { + LOG(ERROR) << "Failed to destroy all endpoints for context " << name() + << "; preserving CQ, PD and device resources for retry"; + return -1; + } for (auto& entry : mr_set_) { int ret = verbs_.ibv_dereg_mr(entry); @@ -633,7 +637,9 @@ std::string RdmaContext::gid() const { } RdmaCQ* RdmaContext::cq(int index) { - if (index < 0 || index >= params_->device.num_cq_list) return nullptr; + // params_ is null until construct(): an inert context has no CQs. + if (!params_ || index < 0 || index >= params_->device.num_cq_list) + return nullptr; return cq_list_.empty() ? nullptr : cq_list_[index]; } diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/cq.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/cq.cpp index 4aef039485..b80d3e7021 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/cq.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/cq.cpp @@ -55,7 +55,8 @@ int RdmaCQ::construct(RdmaContext* context, int cqe_limit, int index) { } bool RdmaCQ::reserveQuota(int num_entries) { - int prev_cqe_now = __sync_fetch_and_add(&cqe_now_, num_entries); + int prev_cqe_now = + cqe_now_.fetch_add(num_entries, std::memory_order_acq_rel); if (prev_cqe_now + num_entries > cqe_limit_) { cancelQuota(num_entries); return false; @@ -64,11 +65,11 @@ bool RdmaCQ::reserveQuota(int num_entries) { } void RdmaCQ::cancelQuota(int num_entries) { - __sync_fetch_and_sub(&cqe_now_, num_entries); + cqe_now_.fetch_sub(num_entries, std::memory_order_acq_rel); } int RdmaCQ::poll(int num_entries, ibv_wc* wc) { - if (!cqe_now_) return 0; + if (cqe_now_.load(std::memory_order_relaxed) == 0) return 0; int rc = ibv_poll_cq(cq_, num_entries, wc); if (rc < 0) { PLOG(ERROR) << "ibv_poll_cq"; @@ -76,4 +77,4 @@ int RdmaCQ::poll(int num_entries, ibv_wc* wc) { return rc; } } // namespace tent -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp index 7a02d38240..4ce75a4620 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint.cpp @@ -16,6 +16,7 @@ #include +#include #include #include #include @@ -54,29 +55,57 @@ static inline const std::string statusToString( static int setupNotifyQpConnection(ibv_qp* qp, RdmaContext* ctx, const std::string& peer_gid_str, uint16_t peer_lid, uint32_t peer_qp_num, - uint16_t pkey_index); + uint16_t pkey_index, + uint8_t service_level = 0, + uint8_t traffic_class = 0); -RdmaEndPoint::RdmaEndPoint() : status_(EP_UNINIT) {} +RdmaEndPoint::RdmaEndPoint() + : status_(EP_UNINIT), + context_(nullptr), + params_(nullptr), + wr_depth_list_(nullptr), + inflight_slices_(0), + destroy_start_time_(0) {} -RdmaEndPoint::~RdmaEndPoint() { - if (status_.load(std::memory_order_relaxed) != EP_UNINIT) deconstruct(); - if (endpoints_count_) - endpoints_count_->fetch_sub(1, std::memory_order_relaxed); -} +RdmaEndPoint::~RdmaEndPoint() { deconstruct(); } int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params, - const std::string& endpoint_name, - std::atomic* endpoints_count) { + const std::string& endpoint_name) { + EndPointStatus expected = EP_UNINIT; + if (!status_.compare_exchange_strong(expected, EP_HANDSHAKING, + std::memory_order_acq_rel)) { + LOG(ERROR) << "Endpoint can only be constructed from EP_UNINIT, got " + << statusToString(expected); + return -1; + } + context_ = context; params_ = params; endpoint_name_ = endpoint_name; - inflight_slices_ = 0; - endpoints_count_ = endpoints_count; - qp_list_.resize(params_->qp_mul_factor); - wr_depth_list_ = new WrDepthBlock[params_->qp_mul_factor]; + inflight_slices_.store(0, std::memory_order_relaxed); + + // Resolve the per-pool QP layout (see computeQpPoolSegments). Empty + // qp_pools (the default) keeps the historical single homogeneous run of + // qp_mul_factor data QPs; a non-empty config lays out one contiguous + // segment per pool. qp_pool_segments_ is read-only after construct(). + QpPoolLayout layout = + computeQpPoolSegments(params_->qp_pools, params_->qp_mul_factor); + if (!layout.valid) { + LOG(ERROR) << "Invalid QP count " << layout.total_qp + << " (qp_mul_factor=" << params_->qp_mul_factor + << ", pools=" << params_->qp_pools.size() << ")"; + return -1; + } + qp_pool_segments_ = std::move(layout.segments); + const int total_qp = layout.total_qp; + + qp_list_.resize(total_qp); + // Value-initialize the full array because cleanup may run after only a + // prefix of QPs has been created. + wr_depth_list_ = new WrDepthBlock[total_qp](); - for (int i = 0; i < params_->qp_mul_factor; ++i) { - wr_depth_list_[i].value = 0; + for (int i = 0; i < total_qp; ++i) { + wr_depth_list_[i].value.store(0, std::memory_order_relaxed); ibv_qp_init_attr attr; memset(&attr, 0, sizeof(attr)); auto cq = context_->cq(i % context_->cqCount())->cq(); @@ -164,64 +193,113 @@ int RdmaEndPoint::construct(RdmaContext* context, EndPointParams* params, } } - status_.store(EP_HANDSHAKING, std::memory_order_relaxed); return 0; } int RdmaEndPoint::deconstruct() { RWSpinlock::WriteGuard guard(lock_); + if (status_.load(std::memory_order_relaxed) == EP_DESTROYED) return 0; + // Synchronous terminal destruction includes the retirement transition so + // callers cannot tear down live QPs without first blocking submissions, + // unpublishing notification state and waking notification senders. + beginDestroyNoLock(); return deconstructUnlocked(); } int RdmaEndPoint::deconstructUnlocked() { auto current_status = status_.load(std::memory_order_relaxed); - // Idempotent: if already destroyed or never initialized, skip cleanup - if (current_status == EP_DESTROYED || current_status == EP_UNINIT) return 0; - status_.store(EP_DESTROYED, std::memory_order_relaxed); + // Idempotent, including delayed destruction through an external shared_ptr. + if (current_status == EP_DESTROYED) return 0; resetInflightSlices(); peer_qp_num_list_.clear(); - // Destroy notification QP - if (notify_qp_) { - // Unregister from transport before destroying - context_->transport_.unregisterNotifyQp(notify_qp_->qp_num); - if (context_->verbs_.ibv_destroy_qp(notify_qp_)) - PLOG(ERROR) << "Failed to destroy notification QP"; - notify_qp_ = nullptr; - notify_connected_ = false; + // A default-constructed endpoint owns no verbs resources. A partially + // constructed endpoint has context_ set and must be cleaned below even + // though it never reached EP_HANDSHAKING. + if (!context_) { + status_.store(EP_DESTROYED, std::memory_order_release); + return 0; } - // Deregister and free notification memory - for (auto& mr : notify_recv_mrs_) { - if (mr) { - if (context_->verbs_.ibv_dereg_mr(mr)) - PLOG(ERROR) << "Failed to deregister notification recv MR"; - mr = nullptr; + int result = 0; + + { + std::lock_guard notify_guard(notify_resource_mutex_); + // Destroy notification QP + if (notify_qp_) { + // Unregister from transport before destroying + context_->transport_.unregisterNotifyQp(notify_qp_->qp_num); + if (context_->verbs_.ibv_destroy_qp(notify_qp_)) { + PLOG(ERROR) << "Failed to destroy notification QP"; + result = -1; + } else { + notify_qp_ = nullptr; + } } - } - notify_recv_mrs_.clear(); - if (notify_send_mr_) { - if (context_->verbs_.ibv_dereg_mr(notify_send_mr_)) - PLOG(ERROR) << "Failed to deregister notification send MR"; - notify_send_mr_ = nullptr; - } - notify_recv_buffers_.clear(); - notify_send_buffer_.clear(); + // A live QP may still reference the notification MRs. Keep both MRs + // and backing buffers intact until QP destruction succeeds. + if (!notify_qp_) { + for (auto& mr : notify_recv_mrs_) { + if (!mr) continue; + if (context_->verbs_.ibv_dereg_mr(mr)) { + PLOG(ERROR) << "Failed to deregister notification recv MR"; + result = -1; + } else { + mr = nullptr; + } + } + if (std::all_of(notify_recv_mrs_.begin(), notify_recv_mrs_.end(), + [](ibv_mr* mr) { return mr == nullptr; })) { + notify_recv_mrs_.clear(); + notify_recv_buffers_.clear(); + } + + if (notify_send_mr_) { + if (context_->verbs_.ibv_dereg_mr(notify_send_mr_)) { + PLOG(ERROR) << "Failed to deregister notification send MR"; + result = -1; + } else { + notify_send_mr_ = nullptr; + } + } + if (!notify_send_mr_) notify_send_buffer_.clear(); + } else { + result = -1; + } + } + bool all_qps_destroyed = true; for (size_t i = 0; i < qp_list_.size(); ++i) { - if (context_->verbs_.ibv_destroy_qp(qp_list_[i])) - PLOG(ERROR) << "ibv_destroy_qp"; - cancelQuota(i, wr_depth_list_[i].value); + if (wr_depth_list_) { + const int outstanding = + wr_depth_list_[i].value.load(std::memory_order_relaxed); + if (outstanding != 0) cancelQuota(i, outstanding); + } + if (!qp_list_[i]) continue; + if (context_->verbs_.ibv_destroy_qp(qp_list_[i])) { + PLOG(ERROR) << "Failed to destroy data QP[" << i << "]"; + result = -1; + all_qps_destroyed = false; + } else { + qp_list_[i] = nullptr; + } } - qp_list_.clear(); - slice_queue_.clear(); - delete[] wr_depth_list_; - wr_depth_list_ = nullptr; - peer_server_name_.clear(); - peer_nic_name_.clear(); - // Status remains EP_DESTROYED (unidirectional lifecycle) - return 0; + if (all_qps_destroyed) { + qp_list_.clear(); + slice_queue_.clear(); + delete[] wr_depth_list_; + wr_depth_list_ = nullptr; + } + + if (result == 0 && !notify_qp_ && notify_recv_mrs_.empty() && + !notify_send_mr_ && qp_list_.empty()) { + peer_server_name_.clear(); + peer_nic_name_.clear(); + status_.store(EP_DESTROYED, std::memory_order_release); + return 0; + } + return -1; } void RdmaEndPoint::beginDestroy() { @@ -237,18 +315,40 @@ void RdmaEndPoint::beginDestroyNoLock() { destroy_start_time_ = getCurrentTimeInNano(); status_.store(EP_DESTROYING, std::memory_order_release); + // Stop publishing the endpoint before QPs start flushing. A notification + // completion that already locked the weak_ptr may finish safely, while no + // later completion can acquire a retiring endpoint. + if (notify_qp_) { + context_->transport_.unregisterNotifyQp(notify_qp_->qp_num); + } + { + std::lock_guard notify_guard(notify_send_mutex_); + notify_connected_ = false; + notify_send_cv_.notify_all(); + } + + // Only EP_READY can own submitted WRs. QPs in EP_UNINIT/EP_HANDSHAKING may + // still be RESET/INIT, where a transition to ERR is invalid on providers. + if (current_status != EP_READY) return; + // Transition QPs to ERR state so hardware flushes inflight WRs to CQ ibv_qp_attr attr; memset(&attr, 0, sizeof(attr)); attr.qp_state = IBV_QPS_ERR; for (size_t i = 0; i < qp_list_.size(); ++i) { + if (!qp_list_[i]) continue; int ret = context_->verbs_.ibv_modify_qp(qp_list_[i], &attr, IBV_QP_STATE); if (ret) { PLOG(ERROR) << "Failed to modify QP to ERR in beginDestroy"; } } + if (notify_qp_ && + context_->verbs_.ibv_modify_qp(notify_qp_, &attr, IBV_QP_STATE)) { + PLOG(ERROR) << "Failed to modify notification QP to ERR in " + "beginDestroy"; + } } bool RdmaEndPoint::finishDestroy() { @@ -258,57 +358,39 @@ bool RdmaEndPoint::finishDestroy() { // Gate 1: already done if (current_status == EP_DESTROYED) return true; - // Gate 2: non-two-phase path. Endpoint reached waiting_list_ without - // going through beginDestroy(). This handles edge cases and serves as - // a safety net. Endpoints that never reached construct() own no RDMA - // resources; drop them directly. if (current_status != EP_DESTROYING) { - if (qp_list_.empty()) { - status_.store(EP_DESTROYED, std::memory_order_relaxed); - return true; - } - LOG(WARNING) << "finishDestroy called in unexpected state: " - << statusToString(current_status) - << ", forcing destruction to avoid waiting_list_ leak"; - // Fall through to the unified destroy path - } else { - // Gate 3: two-phase path. Wait for inflight WRs to drain via CQ - // polling. If ibv_modify_qp-to-ERR failed in beginDestroy, WRs may - // never be flushed; enforce a timeout to avoid leaking forever. - bool has_outstanding = false; - for (size_t i = 0; i < qp_list_.size(); ++i) { - if (wr_depth_list_[i].value != 0) { - has_outstanding = true; - break; - } + LOG(ERROR) << "finishDestroy requires EP_DESTROYING, got " + << statusToString(current_status); + return false; + } + + // Wait for inflight WRs to drain via CQ polling. If the transition to ERR + // failed, enforce a timeout so cleanup can still make progress. + bool has_outstanding = false; + for (size_t i = 0; wr_depth_list_ && i < qp_list_.size(); ++i) { + if (wr_depth_list_[i].value.load(std::memory_order_relaxed) != 0) { + has_outstanding = true; + break; } - if (has_outstanding) { - double elapsed = - (getCurrentTimeInNano() - destroy_start_time_) / 1e9; - if (elapsed < kFinishDestroyTimeoutSec) { - return false; // Still waiting for WRs to drain - } - LOG(WARNING) << "finishDestroy timed out after " << elapsed - << "s with outstanding WRs, forcing destruction"; + } + if (has_outstanding) { + double elapsed = (getCurrentTimeInNano() - destroy_start_time_) / 1e9; + if (elapsed < kFinishDestroyTimeoutSec) { + return false; } + LOG(WARNING) << "finishDestroy timed out after " << elapsed + << "s with outstanding WRs, forcing destruction"; } - // Unified destroy: tear down QPs and bound retries to avoid - // log flooding when ibv_destroy_qp fails permanently. int ret = deconstructUnlocked(); if (ret) { - finish_destroy_retries_++; - LOG(ERROR) << "Failed to finish destroying endpoint (attempt " - << finish_destroy_retries_ << "/" << kFinishDestroyMaxRetries - << "): " << ret; - if (finish_destroy_retries_ < kFinishDestroyMaxRetries) { - return false; // Retry later + destroy_error_count_++; + if (destroy_error_count_ <= kMaxDestroyErrorLogs) { + LOG(ERROR) << "Failed to finish destroying endpoint (attempt " + << destroy_error_count_ << "): " << ret; } - LOG(ERROR) << "Giving up after " << finish_destroy_retries_ - << " retries (possible resource leak)"; + return false; } - - status_.store(EP_DESTROYED, std::memory_order_relaxed); return true; } @@ -388,7 +470,26 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name, } auto bootstrap_status = ControlClient::bootstrap(rpc_server_addr, local_desc, peer_desc); - if (!bootstrap_status.ok()) return bootstrap_status; + if (!bootstrap_status.ok()) { + // With simultaneous open, the peer's bootstrap request may finish + // our passive setup while this outbound RPC is still in flight. + // A timeout therefore does not necessarily mean that connection + // establishment failed. Reuse only the exact endpoint generation + // that issued this RPC; EP_READY alone is not sufficient. + RWSpinlock::WriteGuard guard(lock_); + const bool same_peer = peer_server_name_ == peer_server_name && + peer_nic_name_ == peer_nic_name; + const bool same_local_qps = qpNum() == local_desc.qp_num; + if (status_.load(std::memory_order_relaxed) == EP_READY && + same_peer && same_local_qps) { + LOG(WARNING) + << "Bootstrap RPC failed after simultaneous-open passive " + "setup completed; reusing the established endpoint " + << endpoint_name_ << ": " << bootstrap_status.ToString(); + return mooncake::tent::Status::OK(); + } + return bootstrap_status; + } qp_num = peer_desc.qp_num; peer_gid = peer_desc.local_gid; peer_lid = peer_desc.local_lid; @@ -417,6 +518,7 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name, peer_nic_name_ = peer_nic_name; int rc = setupAllQPs(peer_gid, peer_lid, qp_num); if (rc) { + beginDestroyNoLock(); return mooncake::tent::Status::InternalError( "Failed to configure RDMA endpoint" LOC_MARK); } @@ -424,9 +526,10 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name, // Setup notification QP connection if peer supports it if (peer_desc.notify_qp_num != 0 && notify_qp_) { - rc = setupNotifyQpConnection(notify_qp_, context_, peer_gid, - peer_lid, peer_desc.notify_qp_num, - params_->pkey_index); + rc = setupNotifyQpConnection( + notify_qp_, context_, peer_gid, peer_lid, + peer_desc.notify_qp_num, params_->pkey_index, + params_->service_level, params_->traffic_class); if (rc) { LOG(WARNING) << "Failed to setup notification QP, notification disabled"; @@ -434,7 +537,8 @@ Status RdmaEndPoint::connect(const std::string& peer_server_name, } else { notify_connected_ = true; repostAllNotifyRecvs(); - context_->transport_.registerNotifyQp(notify_qp_->qp_num, this); + context_->transport_.registerNotifyQp(notify_qp_->qp_num, + shared_from_this()); } } } @@ -472,15 +576,19 @@ Status RdmaEndPoint::accept(const BootstrapDesc& peer_desc, local_desc.notify_qp_num = notifyQpNum(); return mooncake::tent::Status::OK(); } - // Endpoint already connected to a different peer - reject the request - // instead of resetting. Endpoints have unidirectional lifecycle and - // are never reset or reused. The caller should create a new endpoint. - LOG(ERROR) - << "Endpoint already established with " << peer_nic_name_ << " of " - << peer_server_name_ - << ", cannot accept new connection (unidirectional lifecycle)"; + // The bootstrap does not match the established connection: the peer + // discarded its endpoint (eviction or failure) and came back with new + // QPs, so the local QPs now point at QPs that no longer exist. + // Endpoints have unidirectional lifecycle and are never reset, so + // retire this one. The caller drops retiring endpoints from the store, + // and the peer's next bootstrap gets a newly created endpoint. + LOG(WARNING) << "Endpoint already established with " << peer_nic_name_ + << " of " << peer_server_name_ + << ", retiring it for the new connection (unidirectional " + "lifecycle)"; + beginDestroyNoLock(); return mooncake::tent::Status::InternalError( - "Endpoint already connected to different peer" LOC_MARK); + "Endpoint retired for reconnection from peer" LOC_MARK); } if (status_.load(std::memory_order_relaxed) != EP_HANDSHAKING) { LOG(ERROR) << "Endpoint not in handshaking state: " @@ -511,6 +619,7 @@ Status RdmaEndPoint::accept(const BootstrapDesc& peer_desc, int rc = setupAllQPs(peer_desc.local_gid, peer_desc.local_lid, peer_desc.qp_num); if (rc) { + beginDestroyNoLock(); return mooncake::tent::Status::InternalError( "Failed to configure RDMA endpoint" LOC_MARK); } @@ -520,13 +629,15 @@ Status RdmaEndPoint::accept(const BootstrapDesc& peer_desc, if (peer_desc.notify_qp_num != 0 && notify_qp_) { rc = setupNotifyQpConnection( notify_qp_, context_, peer_desc.local_gid, peer_desc.local_lid, - peer_desc.notify_qp_num, params_->pkey_index); + peer_desc.notify_qp_num, params_->pkey_index, + params_->service_level, params_->traffic_class); if (rc) { notify_connected_ = false; } else { notify_connected_ = true; repostAllNotifyRecvs(); - context_->transport_.registerNotifyQp(notify_qp_->qp_num, this); + context_->transport_.registerNotifyQp(notify_qp_->qp_num, + shared_from_this()); } } @@ -547,23 +658,29 @@ int RdmaEndPoint::resetConnection(const std::string& reason) { return 0; if (curr_status != EP_HANDSHAKING && curr_status != EP_READY) return 0; - destroy_start_time_ = getCurrentTimeInNano(); - status_.store(EP_DESTROYING, std::memory_order_release); + beginDestroyNoLock(); LOG(INFO) << "Endpoint marked for destruction: " << reason; } // Delete from endpoint store so endpoint() won't return this endpoint. - // remove() calls beginDestroyNoLock() to avoid deadlocking when - // caller already holds lock_. + // Store removal happens after releasing lock_ so remove() can safely take + // the endpoint lock and remains valid for async-event callers too. context_->endpointStore()->remove(endpoint_ptr); return 0; } +const QpPoolSegment* RdmaEndPoint::poolForQp(int qp_index) const { + for (const auto& seg : qp_pool_segments_) { + if (qp_index >= seg.begin && qp_index < seg.begin + seg.num_qp) + return &seg; + } + return nullptr; +} + int RdmaEndPoint::setupAllQPs(const std::string& peer_gid, uint16_t peer_lid, std::vector peer_qp_num_list, std::string* reply_msg) { if (status_.load(std::memory_order_relaxed) == EP_READY) { - status_.store(EP_DESTROYING, std::memory_order_relaxed); return -1; } @@ -574,7 +691,6 @@ int RdmaEndPoint::setupAllQPs(const std::string& peer_gid, uint16_t peer_lid, << peer_nic_name_ << " of " << peer_server_name_; LOG(ERROR) << ss.str(); if (reply_msg) *reply_msg = ss.str(); - status_.store(EP_DESTROYING, std::memory_order_relaxed); return -1; } @@ -582,7 +698,6 @@ int RdmaEndPoint::setupAllQPs(const std::string& peer_gid, uint16_t peer_lid, int ret = setupOneQP(qp_index, peer_gid, peer_lid, peer_qp_num_list[qp_index], reply_msg); if (ret) { - status_.store(EP_DESTROYING, std::memory_order_relaxed); return ret; } } @@ -608,14 +723,25 @@ int RdmaEndPoint::submitSlices(std::vector& slice_list, RWSpinlock::ReadGuard guard(lock_); if (qp_list_.empty()) return 0; if (qp_index < 0) qp_index = 0; - qp_index %= qp_list_.size(); + // Route to the QP pool this transfer asked for (RFC #2568 step 3). All + // slices in a list belong to one task, hence one pool; fold the worker-lane + // candidate into that pool's QP segment. Empty/unknown pool or no pools + // configured => unchanged global spray. + static const std::string kNoPool; + const std::string& pool_name = + (!slice_list.empty() && slice_list.front()->task) + ? slice_list.front()->task->qp_pool + : kNoPool; + qp_index = selectQpInPool(qp_pool_segments_, pool_name, qp_index, + (int)qp_list_.size()); // Check endpoint status before submitting if (status_.load(std::memory_order_relaxed) != EP_READY) return 0; auto cq = context_->cq(qp_index % context_->cqCount()); - int wr_count = - std::min(cq->maxCqe() - cq->getQuota(), - std::min(params_->max_qp_wr - wr_depth_list_[qp_index].value, - (int)slice_list.size())); + int wr_count = std::min( + cq->maxCqe() - cq->getQuota(), + std::min(params_->max_qp_wr - wr_depth_list_[qp_index].value.load( + std::memory_order_relaxed), + (int)slice_list.size())); int sge_count = wr_count * kSgeEntries; if (wr_count <= 0 || !reserveQuota(qp_index, wr_count)) return 0; @@ -694,7 +820,10 @@ int RdmaEndPoint::submitRecvImmDataRequest(int qp_index, uint64_t id) { } void RdmaEndPoint::resetInflightSlices() { - for (int qp_index = 0; qp_index < (int)qp_list_.size(); ++qp_index) { + // qp_list_ is sized up front, while slice_queue_ grows after each + // successful QP creation. Use the latter during partial-construction + // cleanup to avoid indexing queues that were never created. + for (int qp_index = 0; qp_index < (int)slice_queue_.size(); ++qp_index) { auto& queue = slice_queue_[qp_index]; while (!queue.empty()) { auto current = queue.pop(); @@ -724,19 +853,21 @@ size_t RdmaEndPoint::acknowledge(RdmaSlice* slice, TransferStatusEnum status) { std::vector RdmaEndPoint::qpNum() { std::vector ret; for (int qp_index = 0; qp_index < (int)qp_list_.size(); ++qp_index) - ret.push_back(qp_list_[qp_index]->qp_num); + if (qp_list_[qp_index]) ret.push_back(qp_list_[qp_index]->qp_num); return ret; } -int RdmaEndPoint::getInflightSlices() const { return inflight_slices_; } +int RdmaEndPoint::getInflightSlices() const { + return inflight_slices_.load(std::memory_order_relaxed); +} bool RdmaEndPoint::reserveQuota(int qp_index, int num_entries) { assert(qp_index >= 0 && qp_index < (int)qp_list_.size()); auto cq = context_->cq(qp_index % context_->cqCount()); if (!cq->reserveQuota(num_entries)) return false; - auto prev_depth_list = - __sync_fetch_and_add(&wr_depth_list_[qp_index].value, num_entries); - __sync_fetch_and_add(&inflight_slices_, num_entries); + auto prev_depth_list = wr_depth_list_[qp_index].value.fetch_add( + num_entries, std::memory_order_acq_rel); + inflight_slices_.fetch_add(num_entries, std::memory_order_acq_rel); if (prev_depth_list + num_entries > params_->max_qp_wr) { cancelQuota(qp_index, num_entries); return false; @@ -746,8 +877,9 @@ bool RdmaEndPoint::reserveQuota(int qp_index, int num_entries) { void RdmaEndPoint::cancelQuota(int qp_index, int num_entries) { assert(qp_index >= 0 && qp_index < (int)qp_list_.size()); - __sync_fetch_and_sub(&wr_depth_list_[qp_index].value, num_entries); - __sync_fetch_and_sub(&inflight_slices_, num_entries); + wr_depth_list_[qp_index].value.fetch_sub(num_entries, + std::memory_order_acq_rel); + inflight_slices_.fetch_sub(num_entries, std::memory_order_acq_rel); auto cq = context_->cq(qp_index % context_->cqCount()); cq->cancelQuota(num_entries); } @@ -758,6 +890,19 @@ int RdmaEndPoint::setupOneQP(int qp_index, const std::string& peer_gid, assert(qp_index >= 0 && qp_index < (int)qp_list_.size()); auto& qp = qp_list_[qp_index]; + // Resolve link-layer QoS for this QP. When it belongs to a pool that + // overrides SL/TC, use the pool's values; otherwise fall back to the global + // endpoint SL/TC (unchanged default behavior). + const QpPoolSegment* pool = poolForQp(qp_index); + const uint8_t qp_service_level = + (pool && pool->service_level >= 0) + ? static_cast(pool->service_level) + : params_->service_level; + const uint8_t qp_traffic_class = + (pool && pool->traffic_class >= 0) + ? static_cast(pool->traffic_class) + : params_->traffic_class; + // RESET -> INIT ibv_qp_attr attr; memset(&attr, 0, sizeof(attr)); @@ -797,9 +942,9 @@ int RdmaEndPoint::setupOneQP(int qp_index, const std::string& peer_gid, attr.ah_attr.grh.sgid_index = context().gidIndex(); attr.ah_attr.grh.hop_limit = params_->hop_limit; attr.ah_attr.grh.flow_label = params_->flow_label; - attr.ah_attr.grh.traffic_class = params_->traffic_class; + attr.ah_attr.grh.traffic_class = qp_traffic_class; attr.ah_attr.dlid = peer_lid; - attr.ah_attr.sl = params_->service_level; + attr.ah_attr.sl = qp_service_level; attr.ah_attr.src_path_bits = params_->src_path_bits; attr.ah_attr.static_rate = params_->static_rate; attr.ah_attr.is_global = 1; @@ -880,7 +1025,8 @@ void RdmaEndPoint::repostAllNotifyRecvs() { static int setupNotifyQpConnection(ibv_qp* qp, RdmaContext* ctx, const std::string& peer_gid_str, uint16_t peer_lid, uint32_t peer_qp_num, - uint16_t pkey_index) { + uint16_t pkey_index, uint8_t service_level, + uint8_t traffic_class) { // Reconnect path may call this when QP is already in RTS; force a clean // state machine: RESET -> INIT -> RTR -> RTS. ibv_qp_attr qp_attr = {}; @@ -924,14 +1070,14 @@ static int setupNotifyQpConnection(ibv_qp* qp, RdmaContext* ctx, qp_attr.min_rnr_timer = 0x12; qp_attr.ah_attr.is_global = 1; qp_attr.ah_attr.dlid = peer_lid; - qp_attr.ah_attr.sl = 0; + qp_attr.ah_attr.sl = service_level; qp_attr.ah_attr.src_path_bits = 0; qp_attr.ah_attr.port_num = ctx->portNum(); memcpy(&qp_attr.ah_attr.grh.dgid, &peer_gid, 16); qp_attr.ah_attr.grh.flow_label = 0; qp_attr.ah_attr.grh.sgid_index = ctx->gidIndex(); qp_attr.ah_attr.grh.hop_limit = 255; - qp_attr.ah_attr.grh.traffic_class = 0; + qp_attr.ah_attr.grh.traffic_class = traffic_class; ret = ibv_modify_qp(qp, &qp_attr, IBV_QP_STATE | IBV_QP_PATH_MTU | IBV_QP_DEST_QPN | @@ -965,16 +1111,18 @@ static int setupNotifyQpConnection(ibv_qp* qp, RdmaContext* ctx, bool RdmaEndPoint::sendNotification(const std::string& name, const std::string& msg) { - if (!notify_qp_ || !notify_connected_) { - LOG(ERROR) << "Notification QP not connected"; - return false; - } - // Flow control: wait for pending sends to complete std::unique_lock lock(notify_send_mutex_); notify_send_cv_.wait(lock, [this] { - return notify_pending_count_ < kNotifyMaxPendingSends; + return !notify_connected_ || + notify_pending_count_ < kNotifyMaxPendingSends; }); + if (!notify_connected_) { + LOG(ERROR) << "Notification QP not connected"; + return false; + } + std::lock_guard resource_guard(notify_resource_mutex_); + if (!notify_qp_ || !notify_send_mr_) return false; // Pick the next send slot — flow control guarantees this slot's previous // DMA has completed (at most kNotifyMaxPendingSends-1 in-flight). @@ -1032,6 +1180,7 @@ bool RdmaEndPoint::sendNotification(const std::string& name, } bool RdmaEndPoint::handleNotifyRecv(size_t buffer_idx, size_t byte_len) { + std::lock_guard resource_guard(notify_resource_mutex_); if (buffer_idx >= notify_recv_buffers_.size()) { LOG(ERROR) << "Invalid recv buffer index: " << buffer_idx; return false; diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint_store.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint_store.cpp index d1b5ae7315..c5ea69921d 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/endpoint_store.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/endpoint_store.cpp @@ -89,7 +89,7 @@ int FIFOEndpointStore::remove(RdmaEndPoint* ep) { ++iter) { if (iter->second.get() == ep) { waiting_list_.insert(iter->second); - iter->second->beginDestroyNoLock(); + iter->second->beginDestroy(); auto fifo_iter = fifo_map_[iter->first]; fifo_list_.erase(fifo_iter); fifo_map_.erase(iter->first); @@ -130,16 +130,43 @@ void FIFOEndpointStore::reclaim() { size_t FIFOEndpointStore::size() { return endpoint_map_.size(); } -void FIFOEndpointStore::clear() { - RWSpinlock::WriteGuard guard(endpoint_map_lock_); - std::vector to_delete; - for (auto& entry : endpoint_map_) to_delete.push_back(entry.first); - for (auto& key : to_delete) { - endpoint_map_.erase(key); - auto fifo_iter = fifo_map_[key]; - fifo_list_.erase(fifo_iter); - fifo_map_.erase(key); +int FIFOEndpointStore::clear() { + // clear() is used by RdmaContext::disable() immediately before CQs, PD and + // the verbs context are destroyed. Merely erasing endpoint_map_ is unsafe: + // an endpoint retained by an external shared_ptr could otherwise run its + // destructor later and access an already-destroyed RdmaContext. + std::vector> endpoints; + { + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + // Detach both published and already-retiring endpoints atomically. + // After this point no lookup can acquire one of these endpoints. + endpoints.reserve(endpoint_map_.size() + waiting_list_.size()); + for (auto& entry : endpoint_map_) endpoints.push_back(entry.second); + for (auto& endpoint : waiting_list_) endpoints.push_back(endpoint); + endpoint_map_.clear(); + waiting_list_.clear(); + fifo_list_.clear(); + fifo_map_.clear(); + } + // Do not call endpoint methods while holding endpoint_map_lock_: endpoint + // destruction also unregisters its notification QP and takes endpoint + // locks. Keeping verbs teardown outside the Store lock both shortens the + // critical section and avoids introducing a Store/Endpoint lock cycle. + // + // Context shutdown has already stopped workers, so the normal two-phase + // waiting-list path cannot drain CQ flush completions. Force synchronous + // deconstruction here, before RdmaContext destroys its CQs and PD. + // deconstruct() is idempotent, so external shared_ptr owners may release + // the already-deconstructed endpoint later without touching verbs again. + std::vector> failed; + for (auto& endpoint : endpoints) { + if (endpoint->deconstruct()) failed.push_back(endpoint); } + if (failed.empty()) return 0; + + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + waiting_list_.insert(failed.begin(), failed.end()); + return -1; } std::shared_ptr SIEVEEndpointStore::get(const std::string& key) { @@ -191,13 +218,11 @@ std::shared_ptr SIEVEEndpointStore::getOrInsert( } } endpoint = std::make_shared(); - int ret = endpoint->construct(&context_, &context_.params().endpoint, key, - &endpoints_count_); + int ret = endpoint->construct(&context_, &context_.params().endpoint, key); if (ret) { LOG(ERROR) << "Failed to construct endpoint for key " << key; return nullptr; } - endpoints_count_.fetch_add(1, std::memory_order_relaxed); while (this->size() >= max_size_) evictOne(); endpoint_map_[key] = std::make_pair(endpoint, false); fifo_list_.push_front(key); @@ -213,7 +238,7 @@ int SIEVEEndpointStore::remove(RdmaEndPoint* ep) { if (iter->second.first.get() == ep) { waiting_list_len_++; waiting_list_.insert(iter->second.first); - iter->second.first->beginDestroyNoLock(); + iter->second.first->beginDestroy(); auto fifo_iter = fifo_map_[iter->first]; if (hand_.has_value() && hand_.value() == fifo_iter) { fifo_iter == fifo_list_.begin() ? hand_ = std::nullopt @@ -276,30 +301,43 @@ void SIEVEEndpointStore::reclaim() { size_t SIEVEEndpointStore::size() { return endpoint_map_.size(); } -void SIEVEEndpointStore::clear() { - RWSpinlock::WriteGuard guard(endpoint_map_lock_); - std::vector to_delete; - for (auto& entry : endpoint_map_) to_delete.push_back(entry.first); - for (auto& key : to_delete) { - endpoint_map_.erase(key); - auto fifo_iter = fifo_map_[key]; - if (hand_.has_value() && hand_.value() == fifo_iter) { - fifo_iter == fifo_list_.begin() ? hand_ = std::nullopt - : hand_ = std::prev(fifo_iter); - } - fifo_list_.erase(fifo_iter); - fifo_map_.erase(key); +int SIEVEEndpointStore::clear() { + // Terminal shutdown semantics are intentionally different from normal + // SIEVE eviction: remove()/evictOne() retire endpoints asynchronously, + // while clear() must release all verbs objects before their parent + // RdmaContext tears down CQs, PD and the device context. + std::vector> endpoints; + { + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + // Take a strong-reference snapshot and make the Store empty in one + // critical section. This prevents concurrent lookup from observing a + // partially-cleared cache or returning an endpoint being destroyed. + endpoints.reserve(endpoint_map_.size() + waiting_list_.size()); + for (auto& entry : endpoint_map_) + endpoints.push_back(entry.second.first); + for (auto& endpoint : waiting_list_) endpoints.push_back(endpoint); + endpoint_map_.clear(); + waiting_list_.clear(); + waiting_list_len_.store(0, std::memory_order_relaxed); + fifo_list_.clear(); + fifo_map_.clear(); + hand_ = std::nullopt; } - - const int max_retries = 5000; - int retries = 0; - while (endpoints_count_.load(std::memory_order_relaxed) > 0) { - if (++retries > max_retries) { - LOG(ERROR) << "Some endpoints not cleared after 5 seconds"; - break; - } - usleep(1000); + // Run verbs operations outside endpoint_map_lock_. Workers are already + // stopped, so waiting for asynchronous CQ-driven reclaim is impossible; + // synchronously deconstruct each endpoint instead. The operation is + // idempotent, allowing outstanding external shared_ptr references to die + // safely after the Store and Context have gone away. + std::vector> failed; + for (auto& endpoint : endpoints) { + if (endpoint->deconstruct()) failed.push_back(endpoint); } + if (failed.empty()) return 0; + + RWSpinlock::WriteGuard guard(endpoint_map_lock_); + waiting_list_.insert(failed.begin(), failed.end()); + waiting_list_len_.store(waiting_list_.size(), std::memory_order_relaxed); + return -1; } } // namespace tent -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/gdr_reachability.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/gdr_reachability.cpp new file mode 100644 index 0000000000..de2ab14bc6 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/rdma/gdr_reachability.cpp @@ -0,0 +1,126 @@ +// Copyright 2025 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/transport/rdma/gdr_reachability.h" + +#include + +#include "tent/common/config.h" + +namespace mooncake { +namespace tent { + +std::atomic GdrReachability::any_exclusion_{false}; + +GdrReachability &GdrReachability::instance() { + static GdrReachability inst; + return inst; +} + +void GdrReachability::configure(const Config *conf) { + if (!conf) return; + std::unique_lock lock(mutex_); + error_threshold_ = + conf->get("transports/rdma/gdr_error_threshold", error_threshold_); + error_window_ = std::chrono::seconds(conf->get( + "transports/rdma/gdr_error_window_secs", (int)error_window_.count())); + cooldown_ = std::chrono::seconds( + conf->get("transports/rdma/gdr_cooldown_secs", (int)cooldown_.count())); + if (error_threshold_ < 1) error_threshold_ = 1; +} + +std::string GdrReachability::localKey(const std::string &nic_name, + int gpu_ordinal) { + return "L|" + nic_name + "|" + std::to_string(gpu_ordinal); +} + +std::string GdrReachability::remoteKey(const std::string &machine_id, + const std::string &nic_name, + int gpu_ordinal) { + return "R|" + machine_id + "|" + nic_name + "|" + + std::to_string(gpu_ordinal); +} + +bool GdrReachability::reachable(const std::string &key) { + std::shared_lock lock(mutex_); + auto it = state_.find(key); + if (it == state_.end()) return true; + // Not paused, or the cooldown has expired: allow a (probe) request through. + // The next success clears the entry; the next failure re-pauses with a + // doubled cooldown. + return std::chrono::steady_clock::now() >= it->second.resume_time; +} + +void GdrReachability::markFailed(const std::string &key) { + auto now = std::chrono::steady_clock::now(); + std::unique_lock lock(mutex_); + auto &st = state_[key]; + if (st.error_count == 0 || now - st.last_error > error_window_) { + st.error_count = 1; + } else { + st.error_count++; + } + st.last_error = now; + if (st.cooldown.count() == 0) { + st.cooldown = cooldown_; + } else { + st.cooldown *= 2; + if (st.cooldown > kMaxCooldown) st.cooldown = kMaxCooldown; + } + if ((int)st.error_count >= error_threshold_) { + st.resume_time = now + st.cooldown; + any_exclusion_.store(true, std::memory_order_relaxed); + } +} + +void GdrReachability::markRecovered(const std::string &key) { + std::unique_lock lock(mutex_); + state_.erase(key); +} + +void GdrReachability::reportLocalFailure(const std::string &nic_name, + int gpu_ordinal) { + markFailed(localKey(nic_name, gpu_ordinal)); +} + +void GdrReachability::reportLocalSuccess(const std::string &nic_name, + int gpu_ordinal) { + markRecovered(localKey(nic_name, gpu_ordinal)); +} + +bool GdrReachability::localReachable(const std::string &nic_name, + int gpu_ordinal) { + return reachable(localKey(nic_name, gpu_ordinal)); +} + +void GdrReachability::reportRemoteFailure(const std::string &machine_id, + const std::string &nic_name, + int gpu_ordinal) { + markFailed(remoteKey(machine_id, nic_name, gpu_ordinal)); +} + +void GdrReachability::reportRemoteSuccess(const std::string &machine_id, + const std::string &nic_name, + int gpu_ordinal) { + markRecovered(remoteKey(machine_id, nic_name, gpu_ordinal)); +} + +bool GdrReachability::remoteReachable(const std::string &machine_id, + const std::string &nic_name, + int gpu_ordinal) { + return reachable(remoteKey(machine_id, nic_name, gpu_ordinal)); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp index 53a8bb4296..9948d9b8d1 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp @@ -14,6 +14,7 @@ #include "tent/transport/rdma/quota.h" #include "tent/transport/rdma/shared_quota.h" +#include "tent/transport/rdma/gdr_reachability.h" #include "tent/common/utils/random.h" #include "tent/common/utils/os.h" @@ -59,6 +60,23 @@ Status DeviceSelector::allocate(uint64_t total_length, uint32_t num_slices, auto entry = local_topology_->getMemEntry(location); if (!entry) return Status::InvalidArgument("Unknown location" LOC_MARK); + // Exclude NICs that have proven unable to GPUDirect-DMA to this GPU. Only + // engaged once something has actually been learned (permissive fabrics pay + // nothing) and only for GPU/cuda locations. + if (GdrReachability::hasAnyExclusion()) { + LocationParser lp(location); + if (lp.type() == "cuda" && lp.index() >= 0) { + auto& gdr = GdrReachability::instance(); + for (const auto& kv : devices_) { + int dev_id = kv.first; + if (dev_id < 0 || dev_id >= 64) continue; + const auto* nic = local_topology_->getNicEntry(dev_id); + if (nic && !gdr.localReachable(nic->name, lp.index())) + device_mask &= ~(1ULL << dev_id); + } + } + } + if (!smart_selection_enabled_) { // Baseline mode: consistent with original TE behavior // Use devices from the first non-empty rank only @@ -292,7 +310,9 @@ Status DeviceSelector::release(int dev_id, uint64_t length, double latency) { auto& dev = it->second; dev.releaseInflight(length); - if (!smart_selection_enabled_) { + // Cancellation of an unposted slice must release its inflight charge but + // has no latency sample from which to learn bandwidth. + if (!smart_selection_enabled_ || latency <= 0.0) { return Status::OK(); } @@ -316,6 +336,22 @@ Status DeviceSelector::release(int dev_id, uint64_t length, double latency) { return Status::OK(); } +Status DeviceSelector::getNicLoadStats(std::vector& stats) const { + stats.reserve(stats.size() + devices_.size()); + // devices_ is populated during topology load and remains stable while + // transfers update the per-device atomic counters below. + for (const auto& [dev_id, dev] : devices_) { + std::string device_name = local_topology_->getNicName(dev_id); + if (device_name.empty()) device_name = std::to_string(dev_id); + stats.push_back(NicLoadStats{ + std::move(device_name), + dev.getInflightBytes(), + dev.getEwmaBandwidth(), + }); + } + return Status::OK(); +} + void DeviceSelector::printTrafficStats() { std::cout << "=== Device Traffic Statistics ===" << std::endl; for (const auto& [dev_id, dev] : devices_) { @@ -356,5 +392,13 @@ int DeviceSelector::getDevicePriority(int dev_id) const { return static_cast(base_index); } +double DeviceSelector::getAggregateEwmaBandwidth() const { + double total = 0.0; + for (const auto& [id, dev] : devices_) { + total += dev.getEwmaBandwidth(); + } + return total > 0.0 ? total : -1.0; +} + } // namespace tent -} // namespace mooncake \ No newline at end of file +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp index 438b211dcc..21e6f904f5 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rail_monitor.cpp @@ -323,10 +323,32 @@ void RailMonitor::updateBestMapping() { for (size_t i = 0; i < local_cnt; i++) { int local_nic = local_devices[local_numa][i]; int remote_nic = -1; - if (local_numa == remote_numa) + if (local_numa == remote_numa) { remote_nic = direct_rails_[local_nic]; - else - remote_nic = remote_devices[remote_numa][i % remote_cnt]; + } else { + // Cross-NUMA: prefer a same-name remote device (e.g. + // mlx5_5 -> mlx5_5) before falling back to positional + // assignment. loadDefault() already builds direct_rails_ + // via same-name matching (Priority 1); mirroring it here + // avoids mapping a local NIC to an unrelated remote NIC on + // a different physical/overlay network, which fails QP + // modify-to-RTR with "transport retry counter exceeded" on + // multi-bond dual-NUMA RoCEv2 fabrics (issues #2758/#2467). + auto local_entry = local_->getNicEntry(local_nic); + if (local_entry) { + for (int cand : remote_devices[remote_numa]) { + auto cand_entry = remote_->getNicEntry(cand); + if (cand_entry && + cand_entry->name == local_entry->name) { + remote_nic = cand; + break; + } + } + } + if (remote_nic < 0) + remote_nic = + remote_devices[remote_numa][i % remote_cnt]; + } if (!available(local_nic, remote_nic)) { bool found = false; for (int cand : remote_devices[remote_numa]) { diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp index 4a80f357a2..6a82670d45 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/rdma_transport.cpp @@ -37,6 +37,7 @@ #include "tent/common/utils/string_builder.h" #include "tent/runtime/topology.h" #include "tent/common/utils/random.h" +#include "tent/thirdparty/nlohmann/json.h" #define SET_DEVICE(key, param) \ param = conf->get("transports/rdma/device/" #key, param) @@ -192,12 +193,49 @@ static Status convertConfToRdmaParams(std::shared_ptr conf, else params->endpoint.path_mtu = IBV_MTU_512; + // Optional per-pool QP layout (RFC #2568 step 2). Each entry defines a + // named pool with its own QP count and link-layer SL/TC; + // SelectionPolicy.qp_pool references these by name. Absent/empty => single + // default pool (unchanged). The pool SL/TC live here in the RDMA config, + // not in SelectionPolicy, to keep the link-layer QoS definition in the + // transport layer; policies only reference a pool by name. + params->endpoint.qp_pools.clear(); + auto qp_pools_json = + conf->getArray("transports/rdma/endpoint/qp_pools"); + for (const auto& pool_json : qp_pools_json) { + if (!pool_json.is_object()) { + LOG(WARNING) << "Ignore non-object entry in qp_pools"; + continue; + } + if (!pool_json.contains("name") || !pool_json["name"].is_string()) { + LOG(WARNING) << "Ignore qp_pool entry without a string 'name'"; + continue; + } + QpPoolSegment seg; + seg.name = pool_json["name"].get(); + seg.num_qp = pool_json.value("num_qp", 0); + if (seg.num_qp <= 0) { + LOG(WARNING) << "Ignore qp_pool '" << seg.name + << "' with non-positive num_qp " << seg.num_qp; + continue; + } + seg.service_level = pool_json.value("service_level", -1); + seg.traffic_class = pool_json.value("traffic_class", -1); + params->endpoint.qp_pools.push_back(std::move(seg)); + } + if (!params->endpoint.qp_pools.empty()) { + LOG(INFO) << "Configured " << params->endpoint.qp_pools.size() + << " QP pool(s) for per-class link-layer isolation"; + } + SET_WORKERS(max_retry_count, params->workers.max_retry_count); SET_WORKERS(block_size, params->workers.block_size); SET_WORKERS(grace_period_ns, params->workers.grace_period_ns); SET_WORKERS(rail_topo_path, params->workers.rail_topo_path); params->verbose = conf->get("verbose", false); + params->log_slice_affinity = + conf->get("transports/rdma/log_slice_affinity", false); return Status::OK(); } @@ -207,10 +245,16 @@ static bool isGpuDirectRdmaSupported(std::shared_ptr conf) { if (disable_gpu_direct) { return false; } + // Detect vendor GPUDirect/peer-memory drivers from /proc/modules. + // NVIDIA: nvidia_peermem. AMD: peermem is built into amdgpu (linked with + // ib_core), so the amdgpu module itself is the presence signal. std::ifstream modules("/proc/modules"); std::string line; while (std::getline(modules, line)) { - if (line.find("nvidia_peermem") != std::string::npos) { + const auto name_end = line.find(' '); + const auto name = + name_end == std::string::npos ? line : line.substr(0, name_end); + if (name == "nvidia_peermem" || name == "amdgpu") { return true; } } @@ -224,6 +268,32 @@ RdmaTransport::RdmaTransport() RdmaTransport::~RdmaTransport() { uninstall(); } +size_t RdmaTransport::initializeContexts() { + context_set_.clear(); + context_name_lookup_.clear(); + // One slot per NicID: dev_id arrives as a NicID and subscripts both this + // and BufferDesc::lkey, so a compacted layout would name the wrong RNIC. + // Skipped NICs keep an inert context, which consumers reject via status(). + context_set_.reserve(local_topology_->getNicCount()); + size_t context_count = 0; + for (size_t i = 0; i < local_topology_->getNicCount(); ++i) { + auto entry = local_topology_->getNicEntry(i); + auto context = std::make_shared(*this); + context_set_.push_back(context); + if (entry->type != Topology::NIC_RDMA) continue; + int ret = context->construct(entry->name, params_); + if (ret) { + LOG(WARNING) << "Disable RDMA device " << entry->name << " because " + << "of initialization failure"; + continue; + } + context_name_lookup_[entry->name] = i; + ++context_count; + local_buffer_manager_.addDevice(context.get()); + } + return context_count; +} + Status RdmaTransport::install(std::string& local_segment_name, std::shared_ptr metadata, std::shared_ptr local_topology, @@ -269,22 +339,7 @@ Status RdmaTransport::install(std::string& local_segment_name, } local_buffer_manager_.setTopology(local_topology); - context_set_.clear(); - for (size_t i = 0; i < local_topology_->getNicCount(); ++i) { - auto entry = local_topology_->getNicEntry(i); - if (entry->type != Topology::NIC_RDMA) continue; - auto context = std::make_shared(*this); - int ret = context->construct(entry->name, params_); - if (ret) { - LOG(WARNING) << "Disable RDMA device " << entry->name << " because " - << "of initialization failure"; - continue; - } - context_name_lookup_[entry->name] = context_set_.size(); - context_set_.push_back(context); - local_buffer_manager_.addDevice(context.get()); - } - const bool context_empty = context_set_.empty(); + const bool context_empty = initializeContexts() == 0; const bool topology_empty = local_topology_->empty(); if (context_empty || topology_empty) { const char* error_message = "No RDMA device initialized successfully"; @@ -399,9 +454,14 @@ Status RdmaTransport::submitTransferTasks( auto* task = RdmaTaskStorage::Get().allocate(); rdma_batch->task_list.push_back(task); task->request = request; + task->qp_pool = rdma_batch->qp_pool; // RFC #2568 step 3 task->num_slices = 0; task->status_word = PENDING; task->transferred_bytes = 0; + task->success_slices.store(0, std::memory_order_relaxed); + task->resolved_slices.store(0, std::memory_order_relaxed); + task->first_error = PENDING; + task->cancel_requested.store(false, std::memory_order_relaxed); task->ref(); // Batch holds a reference to the task const double merge_ratio = 0.25; @@ -453,6 +513,8 @@ Status RdmaTransport::submitTransferTasks( slice->length = length; slice->task = task; slice->retry_count = 0; + slice->last_fallback_idx = -1; + slice->quota_charged = false; slice->ep_weak_ptr.reset(); slice->word = PENDING; slice->next = nullptr; @@ -460,8 +522,10 @@ Status RdmaTransport::submitTransferTasks( slice->priority = request.priority; // Copy priority from request task->num_slices++; task->ref(); // Each slice holds a reference to the task - if (slice_idx < slice_dev_ids.size()) + if (slice_idx < slice_dev_ids.size()) { slice->source_dev_id = slice_dev_ids[slice_idx]; + slice->quota_charged = true; + } offset += length; int part_id = next_worker_idx % num_workers; auto& list = slice_lists[part_id]; @@ -497,6 +561,23 @@ Status RdmaTransport::getTransferStatus(SubBatchRef batch, int task_id, return Status::OK(); } +Status RdmaTransport::cancelTransferTask(SubBatchRef batch, int task_id) { + auto* rdma_batch = dynamic_cast(batch); + if (!rdma_batch) { + return Status::InvalidArgument("Invalid RDMA sub-batch" LOC_MARK); + } + if (task_id < 0 || task_id >= (int)rdma_batch->task_list.size()) { + return Status::InvalidArgument("Invalid task ID" LOC_MARK); + } + auto* task = rdma_batch->task_list[task_id]; + if (task->status_word != PENDING) return Status::OK(); + return workers_->cancel(task); +} + +Status RdmaTransport::getNicLoadStats(std::vector& stats) const { + return workers_->getDeviceSelector()->getNicLoadStats(stats); +} + bool RdmaTransport::warmupMemory(void* addr, size_t length) { if (length < kMrWarmupMinBytes) return false; unsigned hwc = std::thread::hardware_concurrency(); @@ -541,22 +622,23 @@ Status RdmaTransport::removeMemoryBuffer(BufferDesc& desc) { Status RdmaTransport::setupLocalSegment() { auto& manager = metadata_->segmentManager(); - auto segment = manager.getLocal(); - assert(segment); - // Store RDMA server name for dual-NIC setups; when it differs from - // local_segment_name_ the peer will use it for NIC path construction. - if (rdma_server_name_ != local_segment_name_) { - segment->rdma_server_name = rdma_server_name_; - } - auto& detail = std::get(segment->detail); - for (auto& context : context_set_) { - if (context->status() != RdmaContext::DEVICE_ENABLED) continue; - DeviceDesc device_desc; - device_desc.name = context->name(); - device_desc.lid = context->lid(); - device_desc.gid = context->gid(); - detail.devices.push_back(device_desc); - } + CHECK_STATUS(manager.updateLocal([&](SegmentDesc& segment) -> Status { + // Store RDMA server name for dual-NIC setups; when it differs from + // local_segment_name_ the peer will use it for NIC path construction. + if (rdma_server_name_ != local_segment_name_) { + segment.rdma_server_name = rdma_server_name_; + } + auto& detail = std::get(segment.detail); + for (auto& context : context_set_) { + if (context->status() != RdmaContext::DEVICE_ENABLED) continue; + DeviceDesc device_desc; + device_desc.name = context->name(); + device_desc.lid = context->lid(); + device_desc.gid = context->gid(); + detail.devices.push_back(device_desc); + } + return Status::OK(); + })); return manager.synchronizeLocal(); } @@ -590,6 +672,10 @@ int RdmaTransport::onSetupRdmaConnections(const BootstrapDesc& peer_desc, } auto status = endpoint->accept(peer_desc, local_desc); if (!status.ok()) { + if (endpoint->status() == RdmaEndPoint::EP_DESTROYING || + endpoint->status() == RdmaEndPoint::EP_DESTROYED) { + context->endpointStore()->remove(endpoint.get()); + } LOG(ERROR) << status.ToString(); local_desc.reply_msg = status.ToString(); return -1; @@ -600,13 +686,11 @@ int RdmaTransport::onSetupRdmaConnections(const BootstrapDesc& peer_desc, std::shared_ptr RdmaTransport::getEndpoint(SegmentID target_id, int device_id) { - SegmentDesc* segment_desc = nullptr; - std::string rpc_server_addr, target_seg_name, target_dev_name; + std::string rpc_server_addr, target_seg_name, target_dev_name, + target_nic_path_name; auto status = metadata_->segmentManager().withCachedSegment( target_id, [&](SegmentDesc* segment) { - segment_desc = segment; - if (segment->type != SegmentType::Memory) { return Status::NeedsRefreshCache( "Segment type is not Memory" LOC_MARK); @@ -618,6 +702,7 @@ std::shared_ptr RdmaTransport::getEndpoint(SegmentID target_id, auto topo = &std::get(segment->detail).topology; target_seg_name = segment->name; + target_nic_path_name = segment->nicPathServerName(); target_dev_name = topo->getNicName(device_id); if (target_seg_name.empty() || target_dev_name.empty()) { return Status::NeedsRefreshCache( @@ -631,13 +716,20 @@ std::shared_ptr RdmaTransport::getEndpoint(SegmentID target_id, return nullptr; } - auto context = context_set_[0].get(); - if (context->status() != RdmaContext::DEVICE_ENABLED) { + // context_set_ is NicID-indexed, so slot 0 may be inert; take the first + // enabled context instead. + RdmaContext* context = nullptr; + for (auto& ctx : context_set_) { + if (ctx->status() == RdmaContext::DEVICE_ENABLED) { + context = ctx.get(); + break; + } + } + if (!context) { return nullptr; } std::shared_ptr endpoint; - std::string peer_name = - MakeNicPath(segment_desc->nicPathServerName(), target_dev_name); + std::string peer_name = MakeNicPath(target_nic_path_name, target_dev_name); endpoint = context->endpointStore()->getOrInsert(peer_name); if (!endpoint) { LOG(ERROR) << "Cannot allocate endpoint " << peer_name; @@ -710,22 +802,29 @@ int RdmaTransport::processNotifyCompletions() { // Process each completion for (int i = 0; i < completed; ++i) { - if (wc[i].status != IBV_WC_SUCCESS) { - LOG(ERROR) << "Notification completion failed: " << wc[i].status - << ", qp_num=" << wc[i].qp_num; - continue; - } - - // Find endpoint by QP number - RdmaEndPoint* endpoint = nullptr; + // Find endpoint by QP number before interpreting errors. A flush + // completion after endpoint unpublication is expected during + // retirement and should not flood logs. + std::shared_ptr endpoint; { RWSpinlock::ReadGuard guard(notify_endpoint_map_lock_); auto it = notify_qp_to_endpoint_.find(wc[i].qp_num); if (it != notify_qp_to_endpoint_.end()) { - endpoint = it->second; + endpoint = it->second.lock(); } } + if (wc[i].status != IBV_WC_SUCCESS) { + if (wc[i].status == IBV_WC_WR_FLUSH_ERR && + (!endpoint || + endpoint->status() != RdmaEndPoint::EP_READY)) { + continue; + } + LOG(ERROR) << "Notification completion failed: " << wc[i].status + << ", qp_num=" << wc[i].qp_num; + continue; + } + if (!endpoint) { LOG(WARNING) << "Received notification from unknown QP: " << wc[i].qp_num; @@ -745,7 +844,8 @@ int RdmaTransport::processNotifyCompletions() { return total_completions; } -void RdmaTransport::registerNotifyQp(uint32_t qp_num, RdmaEndPoint* endpoint) { +void RdmaTransport::registerNotifyQp( + uint32_t qp_num, const std::shared_ptr& endpoint) { RWSpinlock::WriteGuard guard(notify_endpoint_map_lock_); notify_qp_to_endpoint_[qp_num] = endpoint; } @@ -761,5 +861,13 @@ void RdmaTransport::notifyWorkerThread() { usleep(notify_poll_interval_us_); } } + +double RdmaTransport::getEstimatedBandwidth() const { + if (!workers_) return -1.0; + auto* sel = workers_->getDeviceSelector(); + if (!sel) return -1.0; + return sel->getAggregateEwmaBandwidth(); +} + } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp index 9b9d2bc341..77317f8d92 100644 --- a/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp +++ b/mooncake-transfer-engine/tent/src/transport/rdma/workers.cpp @@ -14,11 +14,18 @@ #include "tent/transport/rdma/workers.h" +#include "tent/transport/rdma/gdr_reachability.h" + #include +#include #include +#include +#include +#include "tent/transport/rdma/bw_arbitration.h" #include "tent/transport/rdma/endpoint_store.h" +#include "tent/transport/rdma/promotion_policy.h" #include "tent/transport/rdma/shared_quota.h" #include "tent/common/utils/ip.h" #include "tent/common/utils/string_builder.h" @@ -30,6 +37,12 @@ namespace tent { thread_local int tl_wid = -1; namespace { +struct ArbitrationEntry { + RdmaSlice* slice; + double mlu; + size_t order; +}; + // Look up (or create) the RailMonitor for `machine_id` on this worker's // map. Returning a stable reference is safe because the map stores values // via unique_ptr -- rehashes move the pointer slot, not the RailMonitor. @@ -44,10 +57,35 @@ RailMonitor& getOrCreateRail( } // namespace Workers::Workers(RdmaTransport* transport) - : transport_(transport), num_workers_(0), running_(false) { + : transport_(transport), + num_workers_(0), + running_(false), + worker_context_(nullptr) { device_selector_ = std::make_unique(); device_selector_->loadTopology(transport_->local_topology_); auto& conf = transport_->conf_; + GdrReachability::instance().configure(conf.get()); + + // RailMonitor consumes JSON text, while the public configuration is a file + // path. Load it once here instead of reopening the file for every worker + // and every remote machine. Invalid/missing files fall back to automatic + // topology matching in RailMonitor::load(). + const auto& rail_topo_path = transport_->params_->workers.rail_topo_path; + if (!rail_topo_path.empty()) { + std::ifstream input(rail_topo_path); + if (!input.is_open()) { + LOG(WARNING) << "Unable to open RDMA rail topology file " + << rail_topo_path << "; using automatic rail mapping"; + } else { + std::ostringstream contents; + contents << input.rdbuf(); + rail_topo_json_ = contents.str(); + if (rail_topo_json_.empty()) { + LOG(WARNING) << "RDMA rail topology file " << rail_topo_path + << " is empty; using automatic rail mapping"; + } + } + } // ============================================================ // Core Scheduling Configuration @@ -119,6 +157,16 @@ Workers::Workers(RdmaTransport* transport) conf->get("transports/rdma/priority_promotion_timeout_us", 10000) * 1000ull; + // Opt-in deadline-aware bandwidth arbitration within a priority tier + // (RFC #2792). Default false = original FIFO order. + deadline_bw_arbitration_ = + conf->get("transports/rdma/deadline_bw_arbitration", false); + + // Opt-in per-entry promotion (issue #2528). Default false = historical + // head-only "flush the tier" behavior. + priority_promotion_per_entry_ = + conf->get("transports/rdma/priority_promotion_per_entry", false); + // ============================================================ // Global Slot Coordination (Multi-Process) // ============================================================ @@ -225,8 +273,42 @@ Status Workers::submit(RdmaSlice* slice) { return submit(slice_list); } -Status Workers::cancel(RdmaSliceList& slice_list) { - return Status::NotImplemented("cancel not implemented" LOC_MARK); +Status Workers::cancel(RdmaTask* task) { + if (!task) return Status::InvalidArgument("Invalid RDMA task" LOC_MARK); + if (task->cancel_requested.exchange(true, std::memory_order_acq_rel)) { + return Status::OK(); + } + if (!running_.load(std::memory_order_acquire) || !worker_context_ || + !num_workers_) { + return Status::OK(); + } + // Wake every worker because one task may have slices distributed across + // several queues. Cancellation remains best effort for slices already + // posted to a QP; those drain through the normal CQ path. + for (size_t id = 0; id < num_workers_; ++id) { + auto& worker = worker_context_[id]; + std::lock_guard lock(worker.mutex); + if (worker.in_suspend) worker.cv.notify_all(); + } + return Status::OK(); +} + +bool Workers::cancelUnpostedSlice(WorkerContext& worker, RdmaSlice* slice) { + if (!slice || !slice->task || + !slice->task->cancel_requested.load(std::memory_order_acquire)) + return false; + if (slice->word == PENDING) { + releaseSliceQuota(slice); + updateSliceStatus(slice, CANCELED); + } + worker.inflight_slices.fetch_sub(1); + return true; +} + +void Workers::releaseSliceQuota(RdmaSlice* slice, double latency) { + if (!slice || !slice->quota_charged || !device_selector_) return; + device_selector_->release(slice->source_dev_id, slice->length, latency); + slice->quota_charged = false; } std::shared_ptr Workers::getEndpoint(Workers::PostPath path) { @@ -237,8 +319,8 @@ std::shared_ptr Workers::getEndpoint(Workers::PostPath path) { auto target_id = path.remote_segment_id; auto device_id = path.remote_device_id; - auto status = - segment_manager.withCachedSegment(target_id, [&](SegmentDesc* segment) { + auto status = segment_manager.withCachedSegment( + target_id, hint.pin, [&](SegmentDesc* segment) { hint.segment = segment; if (segment->type != SegmentType::Memory) { return Status::NeedsRefreshCache( @@ -324,11 +406,23 @@ void Workers::asyncPostSend() { if (slice_list.num_slices == 0) continue; auto slice = slice_list.first; for (int id = 0; id < slice_list.num_slices; ++id) { + if (cancelUnpostedSlice(worker, slice)) { + slice = slice->next; + continue; + } auto status = generatePostPath(slice); if (!status.ok()) { LOG(ERROR) << "Failed to generate post path for slice " << slice << ": " << status.ToString(); - updateSliceStatus(slice, FAILED); + releaseSliceQuota(slice); + updateSliceStatus(slice, slice->task->cancel_requested.load( + std::memory_order_acquire) + ? CANCELED + : FAILED); + worker.inflight_slices.fetch_sub(1); + } else if (cancelUnpostedSlice(worker, slice)) { + slice = slice->next; + continue; } else { PostPath path{ .local_device_id = slice->source_dev_id, @@ -344,29 +438,89 @@ void Workers::asyncPostSend() { auto& path = entry.first; auto& slices = entry.second; if (slices.empty()) continue; + slices.erase(std::remove_if(slices.begin(), slices.end(), + [&](RdmaSlice* slice) { + return cancelUnpostedSlice(worker, + slice); + }), + slices.end()); + if (slices.empty()) continue; auto endpoint = getEndpoint(path); if (!endpoint) { std::vector clone; slices.swap(clone); for (auto slice : clone) { + if (cancelUnpostedSlice(worker, slice)) continue; slice->retry_count++; if (slice->retry_count >= transport_->params_->workers.max_retry_count) { LOG(WARNING) << "Slice " << slice << " failed: retry count exceeded"; disableEndpoint(slice); + releaseSliceQuota(slice); updateSliceStatus(slice, FAILED); } else { + releaseSliceQuota(slice); submit(slice); } + worker.inflight_slices.fetch_sub(1); } continue; } + // RFC #2792 (opt-in): these slices all contend for one NIC path + // (local NIC -> remote NIC). submitSlices posts as many as the QP + // budget allows and returns num_submitted; the rest wait for the next + // round. Ordering by deadline urgency here means a flow about to miss + // its deadline claims the shared NIC's QP slots ahead of looser flows. + // Default (deadline_bw_arbitration_ == false) leaves order untouched, + // so behavior is byte-identical to today's FIFO / equal split. + if (deadline_bw_arbitration_ && slices.size() > 1) { + const uint64_t now_ns = getCurrentTimeInNano(); + const double bw_bps = device_selector_ + ? device_selector_->getSchedulingParams() + .default_bandwidth_gbps * + 1e9 / 8.0 + : 0.0; + if (bw_bps > 0.0) { + thread_local std::vector scratch; + scratch.clear(); + scratch.reserve(slices.size()); + + for (size_t i = 0; i < slices.size(); ++i) { + const RdmaSlice* s = slices[i]; + ArbFlow flow{0, 0}; + if (s && s->task) { + flow = ArbFlow{s->task->request.deadline_ns, s->length}; + } + scratch.push_back(ArbitrationEntry{ + slices[i], PredictedMlu(flow, now_ns, bw_bps), i}); + } + + std::sort( + scratch.begin(), scratch.end(), + [](const ArbitrationEntry& a, const ArbitrationEntry& b) { + if (a.mlu > b.mlu) return true; + if (a.mlu < b.mlu) return false; + return a.order < b.order; + }); + for (size_t i = 0; i < scratch.size(); ++i) { + slices[i] = scratch[i].slice; + } + } + } + int num_submitted = endpoint->submitSlices(slices, tl_wid); for (int id = 0; id < num_submitted; ++id) { auto slice = slices[id]; if (slice->failed) { + releaseSliceQuota(slice); + if (slice->task->cancel_requested.load( + std::memory_order_acquire)) { + updateSliceStatus(slice, CANCELED); + worker.inflight_slices.fetch_sub(1); + continue; + } slice->retry_count++; if (slice->retry_count >= transport_->params_->workers.max_retry_count) { @@ -377,14 +531,14 @@ void Workers::asyncPostSend() { } else { submit(slice); } + worker.inflight_slices.fetch_sub(1); } else { slice->submit_ts = getCurrentTimeInNano(); + worker.inflight_slice_set.insert(slice); } } if (num_submitted) { - worker.inflight_slice_set.insert(slices.begin(), - slices.begin() + num_submitted); slices.erase(slices.begin(), slices.begin() + num_submitted); } } @@ -397,40 +551,61 @@ void Workers::promoteTimedOutRequests(WorkerContext& worker) { // Set next check time (1ms from now) worker.next_promotion_check_ns = current_ts + 1000000ull; - // Check MEDIUM -> HIGH promotion - std::vector promoted; - worker.queues[PRIO_MEDIUM].pop(promoted); - if (!promoted.empty()) { - auto* slice = promoted.front().first; - if (slice && slice->enqueue_ts > 0 && - (current_ts - slice->enqueue_ts) >= - priority_promotion_timeout_ns_) { - for (auto& slice_list : promoted) { - worker.queues[PRIO_HIGH].push(slice_list); + // Drain one level, promote the entries the policy selects to `to`, and put + // the rest back on `from` in their original order. Returns true if anything + // was promoted (used to preserve the historical "one level per tick" stop). + auto promote_level = [&](int from, int to) -> bool { + std::vector drained; + worker.queues[from].pop(drained); + if (drained.empty()) return false; + + if (!priority_promotion_per_entry_) { + auto* slice = drained.front().first; + const bool head_timed_out = slice && slice->enqueue_ts > 0 && + current_ts >= slice->enqueue_ts && + (current_ts - slice->enqueue_ts) >= + priority_promotion_timeout_ns_; + for (auto& slice_list : drained) { + worker.queues[head_timed_out ? to : from].push(slice_list); } - return; + return head_timed_out; } - for (auto& slice_list : promoted) { - worker.queues[PRIO_MEDIUM].push(slice_list); + + std::vector enqueue_ts; + enqueue_ts.reserve(drained.size()); + for (auto& slice_list : drained) { + auto* slice = slice_list.first; + enqueue_ts.push_back(slice ? slice->enqueue_ts : 0); } - } - // Check LOW -> MEDIUM promotion - worker.queues[PRIO_LOW].pop(promoted); - if (!promoted.empty()) { - auto* slice = promoted.front().first; - if (slice && slice->enqueue_ts > 0 && - (current_ts - slice->enqueue_ts) >= - priority_promotion_timeout_ns_) { - for (auto& slice_list : promoted) { - worker.queues[PRIO_MEDIUM].push(slice_list); - } - return; + PromotionDecision decision = DecidePromotionPerEntry( + enqueue_ts, current_ts, priority_promotion_timeout_ns_); + + if (!decision.promoted_any()) { + for (auto& slice_list : drained) + worker.queues[from].push(slice_list); + return false; + } + + std::vector promote(drained.size(), false); + for (size_t idx : decision.promote_indices) { + if (idx < drained.size()) promote[idx] = true; } - for (auto& slice_list : promoted) { - worker.queues[PRIO_LOW].push(slice_list); + for (size_t i = 0; i < drained.size(); ++i) { + worker.queues[promote[i] ? to : from].push(drained[i]); } - } + return true; + }; + + // Check MEDIUM -> HIGH promotion. Preserve the historical behavior of + // handling at most one level per tick when the head-only policy is active; + // with per-entry promotion, both levels are considered each tick so a + // starving LOW entry is not stalled behind an unrelated MEDIUM promotion. + bool promoted_medium = promote_level(PRIO_MEDIUM, PRIO_HIGH); + if (promoted_medium && !priority_promotion_per_entry_) return; + + // Check LOW -> MEDIUM promotion + promote_level(PRIO_LOW, PRIO_MEDIUM); } void Workers::asyncPollCq() { @@ -465,6 +640,7 @@ void Workers::asyncPollCq() { for (int index = 0; index < num_contexts; index++) { auto& context = transport_->context_set_[index]; auto cq = context->cq(tl_wid % num_cq_list); + if (!cq) continue; // inert context for a non-RDMA or failed NIC ibv_wc wc[kPollCount]; int nr_poll = cq->poll(kPollCount, wc); if (nr_poll < 0) continue; @@ -477,10 +653,7 @@ void Workers::asyncPollCq() { (slice->submit_ts - slice->enqueue_ts) / 1000.0; double inflight_lat = (poll_ts - slice->submit_ts) / 1000.0; double overall_lat_sec = (poll_ts - slice->enqueue_ts) / 1e9; - if (slice->retry_count == 0) { - device_selector_->release(slice->source_dev_id, slice->length, - overall_lat_sec); - } + releaseSliceQuota(slice, overall_lat_sec); if (slice->word != PENDING) continue; if (!ep) { updateSliceStatus(slice, FAILED); @@ -498,6 +671,27 @@ void Workers::asyncPollCq() { << ", local_nic: " << context->name() << "): " << ibv_wc_status_str(wc[i].status); } + // GPUDirect reachability learning: a protection/access error + // on a GPU buffer means the chosen NIC cannot P2P-DMA to that + // GPU (ibv_reg_mr succeeded but the PCIe path is unusable). + // Record it so selection avoids that NIC and converges onto a + // reachable rail instead of exhausting retries. The local side + // (source NIC -> source GPU) surfaces as LOC_PROT; the remote + // side (target NIC -> target GPU) as REM_ACCESS or, for a + // remote GDR-read failure, REM_OP (observed on strict fabrics). + bool local_gdr_err = (wc[i].status == IBV_WC_LOC_PROT_ERR); + bool remote_gdr_err = (wc[i].status == IBV_WC_REM_ACCESS_ERR || + wc[i].status == IBV_WC_REM_OP_ERR); + if (local_gdr_err && slice->source_gpu_ordinal >= 0 && + slice->source_nic_name) { + GdrReachability::instance().reportLocalFailure( + slice->source_nic_name, slice->source_gpu_ordinal); + } else if (remote_gdr_err && slice->target_gpu_ordinal >= 0 && + slice->target_nic_name && slice->target_machine_id) { + GdrReachability::instance().reportRemoteFailure( + *slice->target_machine_id, slice->target_nic_name, + slice->target_gpu_ordinal); + } slice->retry_count++; if (slice->retry_count >= transport_->params_->workers.max_retry_count) { @@ -508,10 +702,31 @@ void Workers::asyncPollCq() { } else { num_slices += ep->acknowledge(slice, PENDING); disableEndpoint(slice); - submit(slice); + if (slice->task->cancel_requested.load( + std::memory_order_acquire)) { + updateSliceStatus(slice, CANCELED); + } else { + submit(slice); + } } } else { num_slices += ep->acknowledge(slice, COMPLETED); + // A successful GPU transfer re-admits any learned GDR + // unreachability for the (GPU, NIC) pair(s) it used, so a + // transient exclusion (or a recovered path) heals. Skipped + // entirely until something has actually been excluded. + if (GdrReachability::hasAnyExclusion()) { + auto& gdr = GdrReachability::instance(); + if (slice->source_gpu_ordinal >= 0 && + slice->source_nic_name) + gdr.reportLocalSuccess(slice->source_nic_name, + slice->source_gpu_ordinal); + if (slice->target_gpu_ordinal >= 0 && + slice->target_nic_name && slice->target_machine_id) + gdr.reportRemoteSuccess(*slice->target_machine_id, + slice->target_nic_name, + slice->target_gpu_ordinal); + } // A successful transfer proves this rail is healthy; clear // any accumulated error count so a previously-cooled-down // rail can be used again without waiting for the full @@ -607,6 +822,14 @@ int Workers::handleContextEvents(std::shared_ptr& context) { return 0; } +void Workers::reclaimEndpoints() { + for (auto& context : transport_->context_set_) { + // Inert contexts never built an endpoint store. + auto store = context->endpointStore(); + if (store) store->reclaim(); + } +} + void Workers::monitorThread() { // Track time for periodic endpoint reclaim (1 Hz heartbeat) auto last_reclaim_time = std::chrono::steady_clock::now(); @@ -621,9 +844,7 @@ void Workers::monitorThread() { .count(); if (time_since_last_reclaim >= 1000) { // 1 second = 1000 ms - for (auto& context : transport_->context_set_) { - context->endpointStore()->reclaim(); - } + reclaimEndpoints(); last_reclaim_time = current_time; } @@ -647,7 +868,7 @@ Status Workers::getRouteHint(RouteHint& hint, SegmentID segment_id, uint64_t addr, uint64_t length) { auto& segment_manager = transport_->metadata_->segmentManager(); CHECK_STATUS(segment_manager.withCachedSegment( - segment_id, [&](SegmentDesc* segment) { + segment_id, hint.pin, [&](SegmentDesc* segment) { hint.segment = segment; hint.buffer = segment->findBuffer(addr, length); if (!hint.buffer) @@ -685,6 +906,7 @@ Status Workers::getRouteHint(RouteHint& hint, SegmentID segment_id, auto mem_id = hint.topo->getMemId(location); if (mem_id < 0) mem_id = hint.topo->getMemId(kWildcardLocation); hint.topo_entry = hint.topo->getMemEntry(mem_id); + hint.location = std::move(location); return Status::OK(); } @@ -703,6 +925,7 @@ Status Workers::selectOptimalDevice(RouteHint& source, RouteHint& target, if (slice->source_dev_id < 0) { CHECK_STATUS(device_selector_->allocate( slice->length, source.buffer->location, slice->source_dev_id)); + slice->quota_charged = true; } if (slice->source_dev_id < 0) @@ -711,7 +934,7 @@ Status Workers::selectOptimalDevice(RouteHint& source, RouteHint& target, auto& rail = getOrCreateRail(worker.rails, target.segment->machine_id); if (!rail.ready() || target.topo != rail.remote()) - rail.load(source.topo, target.topo, /*rail_topo_json=*/"", + rail.load(source.topo, target.topo, rail_topo_json_, transport_->conf_.get()); if (slice->target_dev_id < 0) { int mapped_dev_id = rail.findBestRemoteDevice( @@ -759,7 +982,22 @@ Status Workers::selectOptimalDevice(RouteHint& source, RouteHint& target, return Status::DeviceNotFound( "No device could access the slice memory region" LOC_MARK); - if (!rail.available(slice->source_dev_id, slice->target_dev_id)) { + // Proactively steer away from a NIC/GPU pair already known to be + // GPUDirect-unreachable (learned from earlier completion errors) so we + // never post to a dead rail; the fallback path re-selects a reachable one. + // Reactive learning in asyncPollCq still covers pairs not yet observed. + bool gdr_excluded = false; + if (GdrReachability::hasAnyExclusion()) { + int src_gpu = -1, dst_gpu = -1; + LocationParser s(source.location), d(target.location); + if (s.type() == "cuda") src_gpu = s.index(); + if (d.type() == "cuda") dst_gpu = d.index(); + gdr_excluded = gdrPairExcluded(source, target, slice->source_dev_id, + slice->target_dev_id, src_gpu, dst_gpu); + } + + if (gdr_excluded || + !rail.available(slice->source_dev_id, slice->target_dev_id)) { LOG(INFO) << "Optimal device pair not available: source_dev_id " << slice->source_dev_id << ", target_dev_id " << slice->target_dev_id; @@ -778,12 +1016,41 @@ int Workers::getDeviceByFlatIndex(const RouteHint& hint, size_t flat_idx) { return -1; } +bool Workers::gdrPairExcluded(const RouteHint& source, const RouteHint& target, + int sdev, int tdev, int src_gpu, int dst_gpu) { + auto& gdr = GdrReachability::instance(); + if (src_gpu >= 0) { + const auto* lnic = source.topo->getNicEntry(sdev); + if (lnic && !gdr.localReachable(lnic->name, src_gpu)) return true; + } + // Target-GPU reachability applies whether or not the peer is remote: on a + // same-host transfer the "remote" GPU is still a physical GPU some NICs + // cannot P2P to. The failure is reported under the (target machine_id, nic, + // gpu) key either way, so the check is keyed consistently. + if (dst_gpu >= 0) { + const auto* rnic = target.topo->getNicEntry(tdev); + if (rnic && !gdr.remoteReachable(target.segment->machine_id, rnic->name, + dst_gpu)) + return true; + } + return false; +} + Status Workers::selectFallbackDevice(RouteHint& source, RouteHint& target, RdmaSlice* slice) { LOG_EVERY_N(INFO, 100) << "fallback device selection for slice " << slice; bool same_machine = (source.segment->machine_id == target.segment->machine_id); + // GPUDirect reachability filtering (only when something has been learned). + bool gdr_learned = GdrReachability::hasAnyExclusion(); + int src_gpu = -1, dst_gpu = -1; + if (gdr_learned) { + LocationParser s(source.location), d(target.location); + if (s.type() == "cuda") src_gpu = s.index(); + if (d.type() == "cuda") dst_gpu = d.index(); + } + size_t src_total = 0; for (size_t srank = 0; srank < Topology::DevicePriorityRanks; ++srank) src_total += source.topo_entry->device_list[srank].size(); @@ -793,35 +1060,48 @@ Status Workers::selectFallbackDevice(RouteHint& source, RouteHint& target, dst_total += target.topo_entry->device_list[trank].size(); size_t total_combos = src_total * dst_total; - if ((size_t)slice->retry_count >= total_combos) + if (total_combos == 0) return Status::DeviceNotFound("No available path" LOC_MARK); - size_t idx = slice->retry_count; - while (idx < total_combos) { + // Rotate through source/target combinations with wraparound, resuming just + // past the pair this slice tried last (last_fallback_idx, seeded to -1 so + // the first fallback starts at flat index 0). This keeps a retry from + // immediately re-picking the same path -- a non-GDR failure would otherwise + // burn the whole retry budget on one path before RailMonitor's error + // threshold excludes it -- while still preferring higher-priority pairs: + // getDeviceByFlatIndex walks the per-GPU priority-ranked NIC list, so flat + // index 0 is (source PIX NIC, target PIX NIC), the ideal GPUDirect-capable + // pair. Rail-down and GDR-excluded pairs are skipped, so the scan converges + // onto a reachable rail instead of exhausting retries. + auto& worker = worker_context_[tl_wid]; + RailMonitor* rail_mon = + same_machine + ? nullptr + : &getOrCreateRail(worker.rails, target.segment->machine_id); + size_t start = + static_cast(slice->last_fallback_idx + 1) % total_combos; + for (size_t k = 0; k < total_combos; ++k) { + size_t idx = (start + k) % total_combos; size_t src_idx = idx / dst_total; size_t dst_idx = idx % dst_total; int sdev = getDeviceByFlatIndex(source, src_idx); int tdev = getDeviceByFlatIndex(target, dst_idx); - bool reachable = true; + bool reachable = same_machine ? (sdev == tdev) // loopback is safe + : rail_mon->available(sdev, tdev); - if (same_machine) { - reachable = (sdev == tdev); // loopback is safe - } else { - auto& worker = worker_context_[tl_wid]; - auto& rail = - getOrCreateRail(worker.rails, target.segment->machine_id); - reachable = rail.available(sdev, tdev); - } + // Skip NICs that cannot GPUDirect-DMA to the source/target GPU. + if (reachable && gdr_learned && + gdrPairExcluded(source, target, sdev, tdev, src_gpu, dst_gpu)) + reachable = false; if (reachable) { slice->source_dev_id = sdev; slice->target_dev_id = tdev; - slice->source_lkey = source.buffer->lkey[slice->source_dev_id]; - slice->target_rkey = target.buffer->rkey[slice->target_dev_id]; + // Keys are assigned by generatePostPath() once the device pair is + // settled. + slice->last_fallback_idx = static_cast(idx); return Status::OK(); } - - ++idx; } return Status::DeviceNotFound("No available path" LOC_MARK); @@ -840,13 +1120,52 @@ Status Workers::generatePostPath(RdmaSlice* slice) { CHECK_STATUS(selectOptimalDevice(source, target, slice)); else CHECK_STATUS(selectFallbackDevice(source, target, slice)); - slice->source_lkey = source.buffer->lkey[slice->source_dev_id]; - slice->target_rkey = target.buffer->rkey[slice->target_dev_id]; + // Keys are NicID-indexed. A peer running an older build publishes a + // compacted rkey vector, so a NicID from its device_list can point past the + // end; fail the slice instead of reading out of bounds. + const auto& lkeys = source.buffer->lkey; + const auto& rkeys = target.buffer->rkey; + if (slice->source_dev_id < 0 || + (size_t)slice->source_dev_id >= lkeys.size() || + slice->target_dev_id < 0 || + (size_t)slice->target_dev_id >= rkeys.size()) + return Status::DeviceNotFound( + "Selected device has no registered memory key" LOC_MARK); + slice->source_lkey = lkeys[slice->source_dev_id]; + slice->target_rkey = rkeys[slice->target_dev_id]; // Cache the RailMonitor pointer so asyncPollCq / disableEndpoint can // update rail state without a segment lookup or string-keyed map // lookup on the hot path. slice->rail_monitor = &getOrCreateRail(worker_context_[tl_wid].rails, target.segment->machine_id); + // Stash identifiers for GPUDirect reachability learning in asyncPollCq. + // The name pointers alias stable Topology::NicEntry / segment storage and + // remain valid for the slice's lifetime. + { + LocationParser s(source.location), d(target.location); + slice->source_gpu_ordinal = (s.type() == "cuda") ? s.index() : -1; + slice->target_gpu_ordinal = (d.type() == "cuda") ? d.index() : -1; + const auto* lnic = source.topo->getNicEntry(slice->source_dev_id); + const auto* rnic = target.topo->getNicEntry(slice->target_dev_id); + slice->source_nic_name = lnic ? lnic->name.c_str() : nullptr; + slice->target_nic_name = rnic ? rnic->name.c_str() : nullptr; + slice->target_machine_id = &target.segment->machine_id; + } + if (transport_->params_->log_slice_affinity) { + const auto* local_nic = source.topo->getNicEntry(slice->source_dev_id); + const auto* remote_nic = target.topo->getNicEntry(slice->target_dev_id); + VLOG(1) << "RDMA slice affinity: source_location=" << source.location + << ", target_location=" << target.location + << ", local_device_name=" + << (local_nic ? local_nic->name : "") + << ", peer_device_name=" + << (remote_nic ? remote_nic->name : "") + << ", target_id=" << slice->task->request.target_id + << ", source_addr=" << static_cast(slice->source_addr) + << ", dest_addr=" << reinterpret_cast(slice->target_addr) + << ", length=" << slice->length + << ", retry_count=" << slice->retry_count; + } return Status::OK(); } } // namespace tent diff --git a/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp b/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp index 930df078ef..b13338462c 100644 --- a/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/shm/shm_transport.cpp @@ -15,8 +15,12 @@ #include "tent/transport/shm/shm_transport.h" #include +#include +#include #include #include +#include +#include #include #include @@ -60,11 +64,11 @@ Status ShmTransport::install(std::string &local_segment_name, Status ShmTransport::uninstall() { if (installed_) { + RWSpinlock::WriteGuard guard(relocate_lock_); metadata_.reset(); for (auto &relocate_map : relocate_map_) { for (auto &entry : relocate_map.second) { munmap(entry.second.shm_addr, entry.second.length); - close(entry.second.shm_fd); } } relocate_map_.clear(); @@ -132,6 +136,7 @@ void ShmTransport::startTransfer(ShmTask *task, ShmSubBatch *batch) { } else { task->status_word = TransferStatusEnum::FAILED; } + batch->notifyProgress(); } Status ShmTransport::getTransferStatus(SubBatchRef batch, int task_id, @@ -164,7 +169,10 @@ Status ShmTransport::removeMemoryBuffer(BufferDesc &desc) { } static inline std::string randomFileName() { - std::string result = "mooncake_"; + // Include pid so concurrent processes rarely collide and leftovers are + // easier to attribute during debugging. + std::string result = + "mooncake_" + std::to_string(static_cast(getpid())) + "_"; for (int i = 0; i < 8; ++i) result += 'a' + SimpleRandom::Get().next(26); return result; } @@ -182,13 +190,20 @@ Status ShmTransport::allocateLocalMemory(void **addr, size_t size, if (location.type() != "cpu") { return Status::InvalidArgument("ShmTransport allocates DRAM only"); } - options.shm_path = randomFileName(); options.shm_offset = 0; - *addr = createSharedMemory(options.shm_path, size); - if (!(*addr)) { - return Status::InternalError("Failed to allocate shared memory"); + for (int attempt = 0; attempt < kShmCreateMaxRetries; ++attempt) { + options.shm_path = randomFileName(); + *addr = createSharedMemory(options.shm_path, size); + if (*addr) { + return Status::OK(); + } + // createSharedMemory returns nullptr on EEXIST or other failures. + // Only retry when the name collided with an existing object. + if (errno != EEXIST) { + break; + } } - return Status::OK(); + return Status::InternalError("Failed to allocate shared memory"); } Status ShmTransport::freeLocalMemory(void *addr, size_t size) { @@ -209,20 +224,29 @@ Status ShmTransport::freeLocalMemory(void *addr, size_t size) { void *ShmTransport::createSharedMemory(const std::string &path, size_t size) { int shm_fd = -1; + // O_EXCL prevents silently opening and truncating an existing object that + // another process may still be using (which would SIGBUS that peer). if (cxl_mount_path_.empty()) - shm_fd = shm_open(path.c_str(), O_CREAT | O_RDWR, 0644); + shm_fd = shm_open(path.c_str(), O_CREAT | O_EXCL | O_RDWR, 0644); else { auto full_path = joinPath(cxl_mount_path_, path); - shm_fd = open(full_path.c_str(), O_CREAT | O_RDWR, 0644); + shm_fd = open(full_path.c_str(), O_CREAT | O_EXCL | O_RDWR, 0644); } if (shm_fd == -1) { - PLOG(ERROR) << "Failed to open shared memory file"; + // Preserve errno for allocateLocalMemory retry decisions (EEXIST). + if (errno != EEXIST) { + PLOG(ERROR) << "Failed to open shared memory file"; + } return nullptr; } if (ftruncate64(shm_fd, size) == -1) { PLOG(ERROR) << "Failed to truncate shared memory file"; close(shm_fd); + if (cxl_mount_path_.empty()) + shm_unlink(path.c_str()); + else + unlink(joinPath(cxl_mount_path_, path).c_str()); return nullptr; } @@ -231,6 +255,10 @@ void *ShmTransport::createSharedMemory(const std::string &path, size_t size) { if (mapped_addr == MAP_FAILED) { PLOG(ERROR) << "Failed to map shared memory file"; close(shm_fd); + if (cxl_mount_path_.empty()) + shm_unlink(path.c_str()); + else + unlink(joinPath(cxl_mount_path_, path).c_str()); return nullptr; } @@ -240,31 +268,53 @@ void *ShmTransport::createSharedMemory(const std::string &path, size_t size) { return mapped_addr; } +bool ShmTransport::tryResolve(const RelocateMap &relocate_map, + uint64_t &dest_addr, uint64_t length) { + for (const auto &entry : relocate_map) { + if (entry.first <= dest_addr && + dest_addr + length <= entry.first + entry.second.length) { + dest_addr = dest_addr - entry.first + + reinterpret_cast(entry.second.shm_addr); + return true; + } + } + return false; +} + Status ShmTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, uint64_t length, uint64_t target_id) { - thread_local HashMap tl_relocate_map; - if (tl_relocate_map.empty()) { + { RWSpinlock::ReadGuard guard(relocate_lock_); - tl_relocate_map = relocate_map_; - } - - auto &relocate_map = tl_relocate_map[target_id]; - for (auto &entry : relocate_map) { - if (entry.first <= dest_addr && - dest_addr + length <= entry.first + entry.second.length) { - auto shm_addr = entry.second.shm_addr; - dest_addr = dest_addr - entry.first + ((uint64_t)shm_addr); + auto target = relocate_map_.find(target_id); + if (target != relocate_map_.end() && + tryResolve(target->second, dest_addr, length)) return Status::OK(); - } } RWSpinlock::WriteGuard guard(relocate_lock_); + if (!metadata_) { + return Status::InvalidArgument( + "SHM transport is not installed" LOC_MARK); + } + // Another thread may have published this mapping while the writer lock was + // pending. Recheck before opening and mapping the same shared-memory file. + auto target = relocate_map_.find(target_id); + if (target != relocate_map_.end() && + tryResolve(target->second, dest_addr, length)) + return Status::OK(); BufferDesc *buffer; + // Owning reference: `buffer` is used after the lambda returns. + SegmentDescRef pin; auto &segment_manager = metadata_->segmentManager(); - CHECK_STATUS( - segment_manager.withCachedSegment(target_id, [&](SegmentDesc *segment) { + // Do not munmap relocate_map_ entries on NeedsRefreshCache: memcpy runs + // after this lock is released, and there is no transfer-level refcount / + // quiesce that proves no reader still holds a resolved address. POSIX + // shm_unlink already keeps the object alive for existing mappings; keep + // those mappings until uninstall(). + CHECK_STATUS(segment_manager.withCachedSegment( + target_id, pin, [&](SegmentDesc *segment) { buffer = segment->findBuffer(dest_addr, length); if (!buffer || buffer->shm_path.empty()) return Status::NeedsRefreshCache( @@ -272,25 +322,54 @@ Status ShmTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, return Status::OK(); })); - if (!relocate_map.count(buffer->addr)) { - void *shm_addr = nullptr; + void *shm_addr = nullptr; + bool mapping_found = false; + if (target != relocate_map_.end()) { + auto mapping = target->second.find(buffer->addr); + if (mapping != target->second.end()) { + shm_addr = mapping->second.shm_addr; + mapping_found = true; + } + } + + if (!mapping_found) { LocationParser location(buffer->location); if (location.type() == "cuda") { return Status::NotImplemented( "CUDA supported not enabled in this package " LOC_MARK); } else { int shm_fd = -1; + // Consumer path must never create: a missing object means the peer + // has not published (or has already unlinked) the shared memory. + // Creating here would yield a 0-byte file and SIGBUS on access. if (cxl_mount_path_.empty()) shm_fd = shm_open(buffer->shm_path.c_str(), O_RDWR, 0644); else { auto full_path = joinPath(cxl_mount_path_, buffer->shm_path); - shm_fd = open(full_path.c_str(), O_CREAT | O_RDWR, 0644); + shm_fd = open(full_path.c_str(), O_RDWR, 0644); } if (shm_fd < 0) { return Status::InternalError( std::string("Failed to open shared memory file ") + buffer->shm_path + LOC_MARK); } + // mmap can succeed even when the backing object is shorter than + // buffer->length; accessing past EOF then SIGBUS. Reject early. + struct stat st; + if (fstat(shm_fd, &st) != 0) { + close(shm_fd); + return Status::InternalError( + std::string("Failed to fstat shared memory file ") + + buffer->shm_path + LOC_MARK); + } + if (st.st_size < 0 || + static_cast(st.st_size) < buffer->length) { + close(shm_fd); + return Status::InternalError( + std::string("Shared memory file shorter than registered " + "buffer length: ") + + buffer->shm_path + LOC_MARK); + } shm_addr = mmap(nullptr, buffer->length, PROT_READ | PROT_WRITE, MAP_SHARED, shm_fd, 0); if (shm_addr == MAP_FAILED) { @@ -298,20 +377,19 @@ Status ShmTransport::relocateSharedMemoryAddress(uint64_t &dest_addr, return Status::InternalError( "Failed to map shared memory " LOC_MARK); } + close(shm_fd); LOG(INFO) << "Original shared memory: " << (void *)buffer->addr << "--" << (void *)(buffer->addr + buffer->length); LOG(INFO) << "Remapped shared memory: " << (void *)shm_addr << "--" << (void *)((uintptr_t)shm_addr + buffer->length); OpenedShmEntry shm_entry; - shm_entry.shm_fd = shm_fd; shm_entry.shm_addr = shm_addr; shm_entry.length = buffer->length; - relocate_map[buffer->addr] = shm_entry; + relocate_map_[target_id][buffer->addr] = shm_entry; } } - auto shm_addr = relocate_map[buffer->addr].shm_addr; - dest_addr = dest_addr - buffer->addr + ((uint64_t)shm_addr); + dest_addr = dest_addr - buffer->addr + reinterpret_cast(shm_addr); return Status::OK(); } diff --git a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp index 4cc36f53bf..973173ca72 100644 --- a/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp +++ b/mooncake-transfer-engine/tent/src/transport/tcp/tcp_transport.cpp @@ -152,6 +152,8 @@ Status TcpTransport::submitTransferTasks( tcp_batch->task_list.emplace_back(); auto &task = tcp_batch->task_list.back(); task.request = request; + task.progress_batch_id = batch->progress_batch_id; + task.notify_progress = batch->notify_progress; task.status_word.store(TransferStatusEnum::PENDING, std::memory_order_release); } @@ -211,6 +213,7 @@ void TcpTransport::startTransfer(TcpTask *task) { task->status_word.store(TransferStatusEnum::FAILED, std::memory_order_release); } + if (task->notify_progress) task->notify_progress(task->progress_batch_id); } Status TcpTransport::doTransferWithRetry(TcpTask *task) { diff --git a/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt new file mode 100644 index 0000000000..c7fd499ca8 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/tpu/CMakeLists.txt @@ -0,0 +1,8 @@ +if(USE_TPU) + file(GLOB XPORT_SOURCES "*.cpp") + add_library(tent_xport_tpu STATIC ${XPORT_SOURCES}) + # platform_tpu provides TpuPjrtShim, which the transport consults to + # verify that a staging copy really has a TPU-device side. + target_link_libraries(tent_xport_tpu PUBLIC tent_rpc tent_common + platform_tpu) +endif() diff --git a/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp b/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp new file mode 100644 index 0000000000..380cf1e547 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/tpu/tpu_transport.cpp @@ -0,0 +1,193 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/transport/tpu/tpu_transport.h" + +#include + +#include "tent/common/status.h" +#include "tent/platform/tpu_pjrt_shim.h" +#include "tent/runtime/platform.h" +#include "tent/runtime/slab.h" + +namespace mooncake { +namespace tent { + +TpuTransport::TpuTransport() : installed_(false) {} + +TpuTransport::~TpuTransport() { uninstall(); } + +Status TpuTransport::install(std::string &local_segment_name, + std::shared_ptr metadata, + std::shared_ptr local_topology, + std::shared_ptr conf) { + if (installed_) { + return Status::InvalidArgument( + "TPU transport has been installed" LOC_MARK); + } + metadata_ = metadata; + local_segment_name_ = local_segment_name; + local_topology_ = local_topology; + conf_ = conf; + installed_ = true; + // TPU HBM is not NIC-addressable: only the device<->host staging hop is + // supported here. gpu_to_gpu stays false so the engine always stages + // cross-node traffic through host DRAM. + caps.gpu_to_dram = true; + caps.dram_to_gpu = true; + return Status::OK(); +} + +Status TpuTransport::uninstall() { + if (installed_) { + metadata_.reset(); + installed_ = false; + } + return Status::OK(); +} + +Status TpuTransport::allocateSubBatch(SubBatchRef &batch, size_t max_size) { + auto tpu_batch = Slab::Get().allocate(); + if (!tpu_batch) + return Status::InternalError( + "Unable to allocate TPU sub-batch" LOC_MARK); + batch = tpu_batch; + tpu_batch->task_list.reserve(max_size); + tpu_batch->max_size = max_size; + return Status::OK(); +} + +Status TpuTransport::freeSubBatch(SubBatchRef &batch) { + auto tpu_batch = dynamic_cast(batch); + if (!tpu_batch) + return Status::InvalidArgument("Invalid TPU sub-batch" LOC_MARK); + Slab::Get().deallocate(tpu_batch); + batch = nullptr; + return Status::OK(); +} + +Status TpuTransport::submitTransferTasks( + SubBatchRef batch, const std::vector &request_list) { + auto tpu_batch = dynamic_cast(batch); + if (!tpu_batch) + return Status::InvalidArgument("Invalid TPU sub-batch" LOC_MARK); + if (request_list.size() + tpu_batch->task_list.size() > tpu_batch->max_size) + return Status::TooManyRequests("Exceed batch capacity" LOC_MARK); + for (auto &request : request_list) { + tpu_batch->task_list.push_back(TpuTask{}); + auto &task = tpu_batch->task_list[tpu_batch->task_list.size() - 1]; + task.request = request; + task.status_word = TransferStatusEnum::PENDING; + task.transferred_bytes = 0; + startTransfer(&task, tpu_batch); + } + return Status::OK(); +} + +void TpuTransport::startTransfer(TpuTask *task, TpuSubBatch *batch) { + // TpuTransport only handles the local device<->host staging hop, so the + // target is always the local staging buffer (LOCAL_SEGMENT_ID). Anything + // else indicates a routing bug: TPU HBM cannot be a remote transfer peer. + if (task->request.target_id != LOCAL_SEGMENT_ID) { + LOG(ERROR) << "TpuTransport: unexpected non-local target " + << task->request.target_id + << "; TPU only supports local staging copies"; + task->status_word = TransferStatusEnum::FAILED; + task->transferred_bytes = 0; + batch->notifyProgress(); + return; + } + + void *staging = reinterpret_cast(task->request.target_offset); + + // Exactly one side of a staging hop is TPU HBM: the local stage copies + // HBM<->host staging buffer, and a delegated remote stage copies the peer's + // host staging buffer<->its HBM (so `staging` is the device side there). + // Verify that here instead of relying on Platform::copy to classify: if the + // adapter fails to recognise a device pointer -- e.g. it only matches base + // addresses and ProxyManager handed us `base + chunk_offset` -- then + // Platform::copy sees two host pointers and silently memcpy()s from a token + // that is not the buffer's data, corrupting the transfer. Fail loudly. + auto &shim = TpuPjrtShim::instance(); + const bool source_is_device = shim.isDevicePtr(task->request.source); + const bool staging_is_device = shim.isDevicePtr(staging); + if (source_is_device == staging_is_device) { + LOG(ERROR) + << "TpuTransport: a staging copy must have exactly one TPU " + "device side, but source=" + << task->request.source << " (device=" << source_is_device + << ") and target=" << staging << " (device=" << staging_is_device + << "). Either the PJRT adapter is unavailable, or it does not " + "resolve interior pointers (see tpu_pjrt_abi.h)."; + task->status_word = TransferStatusEnum::FAILED; + task->transferred_bytes = 0; + batch->notifyProgress(); + return; + } + + Status status; + if (task->request.opcode == Request::READ) + // host staging buffer -> device (H2D) + status = Platform::getLoader().copy(task->request.source, staging, + task->request.length); + else + // device -> host staging buffer (D2H) + status = Platform::getLoader().copy(staging, task->request.source, + task->request.length); + + if (status.ok()) { + task->transferred_bytes = task->request.length; + task->status_word = TransferStatusEnum::COMPLETED; + } else { + LOG(WARNING) << "TpuTransport: staging copy failed: " + << status.ToString(); + task->status_word = TransferStatusEnum::FAILED; + } + batch->notifyProgress(); +} + +Status TpuTransport::getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus &status) { + auto tpu_batch = dynamic_cast(batch); + if (!tpu_batch) + return Status::InvalidArgument("Invalid TPU sub-batch" LOC_MARK); + if (task_id < 0 || task_id >= (int)tpu_batch->task_list.size()) { + return Status::InvalidArgument("Invalid task id" LOC_MARK); + } + auto &task = tpu_batch->task_list[task_id]; + status = TransferStatus{task.status_word, task.transferred_bytes}; + return Status::OK(); +} + +Status TpuTransport::addMemoryBuffer(BufferDesc &desc, + const MemoryOptions &options) { + LocationParser location(desc.location); + // Tag both TPU device buffers and host staging buffers so the staging + // policy can route the local HBM<->host hop through this transport. Routing + // keys on the target (host) buffer's transports, so the host buffer must + // carry TransportType::TPU as well (mirrors NVLinkTransport). + if (location.type() != "tpu" && location.type() != "cpu" && + location.type() != kWildcardLocation) { + return Status::OK(); // Not our buffer; leave it untagged. + } + desc.transports.push_back(TransportType::TPU); + return Status::OK(); +} + +Status TpuTransport::removeMemoryBuffer(BufferDesc &desc) { + return Status::OK(); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/src/transport/ub/CMakeLists.txt b/mooncake-transfer-engine/tent/src/transport/ub/CMakeLists.txt new file mode 100644 index 0000000000..89035a84c2 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/CMakeLists.txt @@ -0,0 +1,8 @@ +if(USE_UB) + file(GLOB TENT_UB_SOURCES CONFIGURE_DEPENDS "*.cpp") + add_library(tent_xport_ub STATIC ${TENT_UB_SOURCES}) + target_link_libraries(tent_xport_ub PUBLIC tent_common Urma::urma) + if(URMA_LIBRARY) + target_compile_definitions(tent_xport_ub PRIVATE TENT_HAS_REAL_URMA=1) + endif() +endif() diff --git a/mooncake-transfer-engine/tent/src/transport/ub/buffers.cpp b/mooncake-transfer-engine/tent/src/transport/ub/buffers.cpp new file mode 100644 index 0000000000..f9b6b971c3 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/buffers.cpp @@ -0,0 +1,610 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/buffers.h" + +#include +#include +#include +#include +#include +#include + +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake::tent::ub { +namespace { + +using json = nlohmann::json; + +bool checkedEnd(uint64_t base, uint64_t length, uint64_t& end) { + if (length > std::numeric_limits::max() - base) return false; + end = base + length; + return true; +} + +uint64_t generationSeed() { + const auto wall = static_cast( + std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch()) + .count()); + const auto steady = static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); + const uint64_t seed = + (wall ^ (steady << 17) ^ (steady >> 11)) & 0x7fffffffffffffffULL; + return seed == 0 ? 1 : seed; +} + +} // namespace + +Status encodeBufferMetadata(const UbBufferMetadata& metadata, + std::string& encoded) { + if (metadata.schema_version != UbBufferMetadata::kSchemaVersion || + metadata.generation == 0 || metadata.length == 0 || + metadata.segments.empty()) { + return Status::InvalidArgument("Invalid UB buffer metadata" LOC_MARK); + } + + uint64_t metadata_end = 0; + if (!checkedEnd(metadata.base, metadata.length, metadata_end)) { + return Status::InvalidArgument( + "UB buffer metadata range overflows" LOC_MARK); + } + json segments = json::array(); + std::unordered_set topology_ids; + for (const auto& segment : metadata.segments) { + if (segment.topology_id < 0 || segment.device_name.empty() || + segment.eid.empty() || segment.descriptor.hex.empty() || + segment.descriptor.urma_abi_size == 0 || + !topology_ids.insert(segment.topology_id).second) { + return Status::InvalidArgument( + "Incomplete UB segment metadata" LOC_MARK); + } + segments.push_back( + {{"topology_id", segment.topology_id}, + {"device_name", segment.device_name}, + {"eid", segment.eid}, + {"eid_index", segment.eid_index}, + {"descriptor", + {{"schema_version", segment.descriptor.schema_version}, + {"urma_api_version", segment.descriptor.urma_api_version}, + {"urma_abi_size", segment.descriptor.urma_abi_size}, + {"hex", segment.descriptor.hex}}}}); + } + encoded = json{{"schema_version", metadata.schema_version}, + {"generation", metadata.generation}, + {"base", metadata.base}, + {"length", metadata.length}, + {"location", metadata.location}, + {"permission", static_cast(metadata.permission)}, + {"segments", std::move(segments)}} + .dump(); + return Status::OK(); +} + +Status decodeBufferMetadata(std::string_view encoded, + UbBufferMetadata& metadata) { + try { + const auto value = json::parse(encoded); + if (!value.is_object() || !value.contains("schema_version") || + !value.contains("generation") || !value.contains("base") || + !value.contains("length") || !value.contains("permission") || + !value.contains("segments")) { + return Status::MalformedJson( + "Missing required UB buffer metadata field" LOC_MARK); + } + + UbBufferMetadata parsed; + parsed.schema_version = value.at("schema_version").get(); + if (parsed.schema_version != UbBufferMetadata::kSchemaVersion) { + return Status::InvalidMetadataType( + "Unsupported UB buffer metadata version" LOC_MARK); + } + parsed.generation = value.at("generation").get(); + parsed.base = value.at("base").get(); + parsed.length = value.at("length").get(); + parsed.location = value.value("location", ""); + const int permission = value.at("permission").get(); + if (permission < static_cast(kLocalReadWrite) || + permission > static_cast(kGlobalReadWrite)) { + return Status::InvalidMetadataType( + "Invalid UB buffer permission" LOC_MARK); + } + parsed.permission = static_cast(permission); + if (parsed.generation == 0 || parsed.length == 0 || + !value.at("segments").is_array() || value.at("segments").empty()) { + return Status::InvalidMetadataType( + "Invalid UB buffer metadata values" LOC_MARK); + } + uint64_t metadata_end = 0; + if (!checkedEnd(parsed.base, parsed.length, metadata_end)) { + return Status::InvalidMetadataType( + "UB buffer metadata range overflows" LOC_MARK); + } + + std::unordered_set topology_ids; + for (const auto& item : value.at("segments")) { + UbBufferSegmentMetadata segment; + segment.topology_id = item.at("topology_id").get(); + segment.device_name = item.at("device_name").get(); + segment.eid = item.at("eid").get(); + segment.eid_index = item.at("eid_index").get(); + const auto& descriptor = item.at("descriptor"); + segment.descriptor.schema_version = + descriptor.at("schema_version").get(); + segment.descriptor.urma_api_version = + descriptor.at("urma_api_version").get(); + segment.descriptor.urma_abi_size = + descriptor.at("urma_abi_size").get(); + segment.descriptor.hex = descriptor.at("hex").get(); + if (segment.topology_id < 0 || segment.device_name.empty() || + segment.eid.empty() || segment.descriptor.hex.empty() || + segment.descriptor.schema_version != + SegmentDescriptor::kSchemaVersion || + segment.descriptor.urma_abi_size == 0 || + !topology_ids.insert(segment.topology_id).second) { + return Status::InvalidMetadataType( + "Invalid UB segment descriptor envelope" LOC_MARK); + } + parsed.segments.push_back(std::move(segment)); + } + metadata = std::move(parsed); + return Status::OK(); + } catch (const std::exception& error) { + return Status::MalformedJson(std::string("Malformed UB metadata: ") + + error.what() + LOC_MARK); + } +} + +UbBufferManager::UbBufferManager(std::shared_ptr adapter, + std::vector contexts) + : adapter_(std::move(adapter)), contexts_(std::move(contexts)) { + next_generation_.store(generationSeed(), std::memory_order_relaxed); + for (const auto& context : contexts_) { + if (context) context_by_topology_id_[context->topologyId()] = context; + } +} + +UbBufferManager::~UbBufferManager() { (void)clear(); } + +bool UbBufferManager::AddressRange::contains(uint64_t address, + size_t size) const { + uint64_t this_end = 0; + uint64_t query_end = 0; + return checkedEnd(base, length, this_end) && + checkedEnd(address, static_cast(size), query_end) && + address >= base && query_end <= this_end; +} + +size_t UbBufferManager::ImportKeyHash::operator()( + const ImportKey& key) const noexcept { + size_t seed = std::hash{}(key.local_topology_id); + auto combine = [&seed](size_t value) { + seed ^= value + 0x9e3779b97f4a7c15ULL + (seed << 6) + (seed >> 2); + }; + combine(std::hash{}(key.remote_segment_id)); + combine(std::hash{}(key.remote_topology_id)); + combine(std::hash{}(key.buffer_base)); + combine(std::hash{}(key.generation)); + return seed; +} + +uint32_t UbBufferManager::segmentAccess(Permission permission) { + switch (permission) { + case kGlobalReadOnly: + return SEGMENT_ACCESS_READ; + case kLocalReadWrite: + return SEGMENT_ACCESS_LOCAL_ONLY; + case kGlobalReadWrite: + default: + return SEGMENT_ACCESS_READ | SEGMENT_ACCESS_WRITE; + } +} + +bool UbBufferManager::permissionAllows(Permission permission, + Request::OpCode opcode) { + if (permission == kLocalReadWrite) return false; + return opcode == Request::READ || permission == kGlobalReadWrite; +} + +UbContextPtr UbBufferManager::findContext(Topology::NicID topology_id) const { + auto it = context_by_topology_id_.find(topology_id); + return it == context_by_topology_id_.end() ? nullptr : it->second; +} + +Status UbBufferManager::addBufferInternal(BufferDesc& desc, + const MemoryOptions& options) { + if (desc.length == 0 || contexts_.empty()) { + return Status::InvalidArgument( + "Cannot register an empty UB buffer or without contexts" LOC_MARK); + } + AddressRange range{desc.addr, desc.length}; + uint64_t ignored = 0; + if (!checkedEnd(range.base, range.length, ignored)) { + return Status::InvalidArgument("UB buffer address overflow" LOC_MARK); + } + + LocalRecord record; + record.options = options; + record.generation = + next_generation_.fetch_add(1, std::memory_order_relaxed); + if (record.generation == 0) { + record.generation = + next_generation_.fetch_add(1, std::memory_order_relaxed); + } + + UbBufferMetadata metadata; + metadata.generation = record.generation; + metadata.base = desc.addr; + metadata.length = desc.length; + metadata.location = desc.location; + metadata.permission = options.perm; + + SegmentOptions segment_options; + segment_options.access = segmentAccess(options.perm); + for (const auto& context : contexts_) { + if (!context || !context->active()) continue; + LocalSegmentPtr segment; + auto status = adapter_->registerLocalSegment(context->handle(), + desc.addr, desc.length, + segment_options, segment); + if (!status.ok()) { + if (segment) { + record.segments.emplace(context->topologyId(), + std::move(segment)); + } + (void)unregisterRecord(record); + retainPendingRecord(record); + return status; + } + metadata.segments.push_back(UbBufferSegmentMetadata{ + context->topologyId(), context->deviceInfo().native_device_name, + context->deviceInfo().eid, + static_cast(context->deviceInfo().eid_index), + segment->descriptor()}); + record.segments.emplace(context->topologyId(), std::move(segment)); + } + if (record.segments.empty()) { + return Status::DeviceNotFound( + "No active UB context accepted the buffer" LOC_MARK); + } + + std::string encoded; + auto status = encodeBufferMetadata(metadata, encoded); + if (!status.ok()) { + (void)unregisterRecord(record); + retainPendingRecord(record); + return status; + } + + { + std::unique_lock lock(local_mutex_); + auto next = local_buffers_.lower_bound(range); + if ((next != local_buffers_.end() && + next->first.base < range.base + range.length) || + (next != local_buffers_.begin() && + std::prev(next)->first.base + std::prev(next)->first.length > + range.base)) { + lock.unlock(); + (void)unregisterRecord(record); + retainPendingRecord(record); + return Status::InvalidArgument( + "Overlapping UB buffer registration" LOC_MARK); + } + local_buffers_.emplace(range, std::move(record)); + } + desc.transport_attrs[TransportType::UB] = std::move(encoded); + if (std::find(desc.transports.begin(), desc.transports.end(), + TransportType::UB) == desc.transports.end()) { + desc.transports.push_back(TransportType::UB); + } + return Status::OK(); +} + +Status UbBufferManager::addBuffer(BufferDesc& desc, + const MemoryOptions& options) { + return addBufferInternal(desc, options); +} + +Status UbBufferManager::addBuffers(std::vector& descs, + const MemoryOptions& options) { + std::vector added; + added.reserve(descs.size()); + for (auto& desc : descs) { + auto status = addBufferInternal(desc, options); + if (!status.ok()) { + for (auto it = added.rbegin(); it != added.rend(); ++it) { + (void)removeBuffer(**it); + } + return status; + } + added.push_back(&desc); + } + return Status::OK(); +} + +Status UbBufferManager::unregisterRecord(LocalRecord& record) { + Status first_error = Status::OK(); + for (auto it = record.segments.begin(); it != record.segments.end();) { + auto status = adapter_->unregisterLocalSegment(it->second); + if (!status.ok()) { + if (first_error.ok()) first_error = status; + ++it; + continue; + } + if (it->second) { + if (first_error.ok()) { + first_error = Status::InternalError( + "URMA adapter retained a local segment after successful " + "unregister" LOC_MARK); + } + ++it; + continue; + } + it = record.segments.erase(it); + } + return first_error; +} + +void UbBufferManager::retainPendingRecord(LocalRecord& record) { + if (record.segments.empty()) return; + std::unique_lock lock(local_mutex_); + for (auto& [_, segment] : record.segments) { + if (segment) pending_local_segments_.push_back(std::move(segment)); + } + record.segments.clear(); +} + +Status UbBufferManager::removeBuffer(BufferDesc& desc) { + Status status = Status::OK(); + bool removed = false; + { + std::unique_lock lock(local_mutex_); + auto it = local_buffers_.find(AddressRange{desc.addr, desc.length}); + if (it != local_buffers_.end()) { + status = unregisterRecord(it->second); + if (status.ok() && it->second.segments.empty()) { + local_buffers_.erase(it); + removed = true; + } + } else { + // Idempotent retry after a previously successful removal. + removed = true; + } + } + if (!status.ok()) return status; + if (!removed) { + return Status::InternalError( + "UB local record still owns segments after unregister" LOC_MARK); + } + desc.transport_attrs.erase(TransportType::UB); + desc.transports.erase(std::remove(desc.transports.begin(), + desc.transports.end(), TransportType::UB), + desc.transports.end()); + return Status::OK(); +} + +Status UbBufferManager::clear() { + Status first_error = Status::OK(); + { + std::unique_lock lock(local_mutex_); + for (auto it = local_buffers_.begin(); it != local_buffers_.end();) { + auto status = unregisterRecord(it->second); + if (!status.ok() && first_error.ok()) first_error = status; + if (status.ok() && it->second.segments.empty()) { + it = local_buffers_.erase(it); + } else { + ++it; + } + } + for (auto it = pending_local_segments_.begin(); + it != pending_local_segments_.end();) { + auto status = adapter_->unregisterLocalSegment(*it); + if (!status.ok()) { + if (first_error.ok()) first_error = status; + ++it; + } else if (*it) { + if (first_error.ok()) { + first_error = Status::InternalError( + "URMA adapter retained a pending local segment after " + "successful unregister" LOC_MARK); + } + ++it; + } else { + it = pending_local_segments_.erase(it); + } + } + } + + { + std::unique_lock lock(import_mutex_); + for (auto it = imports_.begin(); it != imports_.end();) { + auto status = adapter_->unimportRemoteSegment(it->second); + if (!status.ok()) { + if (first_error.ok()) first_error = status; + ++it; + } else if (it->second) { + if (first_error.ok()) { + first_error = Status::InternalError( + "URMA adapter retained a remote segment after " + "successful unimport" LOC_MARK); + } + ++it; + } else { + it = imports_.erase(it); + } + } + for (auto it = pending_remote_segments_.begin(); + it != pending_remote_segments_.end();) { + auto status = adapter_->unimportRemoteSegment(*it); + if (!status.ok()) { + if (first_error.ok()) first_error = status; + ++it; + } else if (*it) { + if (first_error.ok()) { + first_error = Status::InternalError( + "URMA adapter retained a pending remote segment after " + "successful unimport" LOC_MARK); + } + ++it; + } else { + it = pending_remote_segments_.erase(it); + } + } + } + return first_error; +} + +Status UbBufferManager::findLocal(uint64_t address, size_t length, + Topology::NicID local_topology_id, + LocalSegmentRef& result) const { + std::shared_lock lock(local_mutex_); + auto it = local_buffers_.upper_bound( + AddressRange{address, std::numeric_limits::max()}); + if (it == local_buffers_.begin()) { + return Status::AddressNotRegistered( + "Local UB address is not registered" LOC_MARK); + } + --it; + if (!it->first.contains(address, length)) { + return Status::AddressNotRegistered( + "Local UB range crosses a registration boundary" LOC_MARK); + } + auto segment = it->second.segments.find(local_topology_id); + if (segment == it->second.segments.end()) { + return Status::AddressNotRegistered( + "Local UB buffer is not registered on selected device" LOC_MARK); + } + result = LocalSegmentRef{findContext(local_topology_id), segment->second, + it->second.generation, it->first.base, + it->first.length}; + return Status::OK(); +} + +Status UbBufferManager::importRemote(SegmentID remote_segment_id, + Topology::NicID local_topology_id, + Topology::NicID remote_topology_id, + const BufferDesc& remote_buffer, + Request::OpCode opcode, uint64_t address, + size_t length, + ImportedSegmentRef& result) { + auto attr = remote_buffer.transport_attrs.find(TransportType::UB); + if (attr == remote_buffer.transport_attrs.end()) { + return Status::NeedsRefreshCache( + "Remote buffer has no UB metadata" LOC_MARK); + } + UbBufferMetadata metadata; + CHECK_STATUS(decodeBufferMetadata(attr->second, metadata)); + if (metadata.base != remote_buffer.addr || + metadata.length != remote_buffer.length) { + return Status::NeedsRefreshCache( + "Remote UB metadata does not match BufferDesc" LOC_MARK); + } + uint64_t remote_end = 0; + uint64_t query_end = 0; + if (!checkedEnd(metadata.base, metadata.length, remote_end) || + !checkedEnd(address, length, query_end) || address < metadata.base || + query_end > remote_end) { + return Status::InvalidArgument( + "Remote UB request is outside the registered buffer" LOC_MARK); + } + if (!permissionAllows(metadata.permission, opcode)) { + return Status::InvalidArgument( + "Remote UB buffer permission rejects the operation" LOC_MARK); + } + + auto descriptor = std::find_if( + metadata.segments.begin(), metadata.segments.end(), + [remote_topology_id](const UbBufferSegmentMetadata& segment) { + return segment.topology_id == remote_topology_id; + }); + if (descriptor == metadata.segments.end()) { + return Status::NeedsRefreshCache( + "Remote UB buffer lacks selected device descriptor" LOC_MARK); + } + auto context = findContext(local_topology_id); + if (!context || !context->active()) { + return Status::DeviceNotFound( + "Selected local UB context is unavailable" LOC_MARK); + } + + ImportKey key{local_topology_id, remote_segment_id, remote_topology_id, + metadata.base, metadata.generation}; + RemoteSegmentPtr imported; + { + // Serialize the provider import with the second cache lookup. This + // avoids creating a duplicate handle that has no cache key capable of + // retaining it when an immediate rollback fails. + std::unique_lock lock(import_mutex_); + auto current = imports_.find(key); + if (current == imports_.end()) { + SegmentOptions options; + options.access = segmentAccess(metadata.permission); + auto import_status = adapter_->importRemoteSegment( + context->handle(), descriptor->descriptor, options, imported); + if (!import_status.ok()) { + if (imported) { + (void)adapter_->unimportRemoteSegment(imported); + if (imported) { + pending_remote_segments_.push_back(std::move(imported)); + } + } + return import_status; + } + if (!imported) { + return Status::InternalError( + "URMA adapter returned no remote segment after successful " + "import" LOC_MARK); + } + current = imports_.emplace(key, imported).first; + } else { + imported = current->second; + } + + for (auto iter = imports_.begin(); iter != imports_.end();) { + const auto& candidate = iter->first; + const bool stale = + candidate.local_topology_id == local_topology_id && + candidate.remote_segment_id == remote_segment_id && + candidate.remote_topology_id == remote_topology_id && + candidate.buffer_base == metadata.base && + candidate.generation != metadata.generation; + if (!stale) { + ++iter; + continue; + } + auto status = adapter_->unimportRemoteSegment(iter->second); + if (!status.ok()) { + // The old generation may still be retained by an in-flight + // WR. Keep it cached for a later import/clear retry, but do + // not fail the current request after its new generation was + // imported successfully. + ++iter; + } else if (iter->second) { + // Treat an adapter that reports success without releasing the + // handle conservatively: retain ownership and retry later. + ++iter; + } else { + iter = imports_.erase(iter); + } + } + } + result = + ImportedSegmentRef{context, imported, metadata.generation, + metadata.base, metadata.length, remote_topology_id}; + return Status::OK(); +} + +size_t UbBufferManager::localBufferCount() const { + std::shared_lock lock(local_mutex_); + return local_buffers_.size(); +} + +size_t UbBufferManager::importedSegmentCount() const { + std::shared_lock lock(import_mutex_); + return imports_.size(); +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/context.cpp b/mooncake-transfer-engine/tent/src/transport/ub/context.cpp new file mode 100644 index 0000000000..c63475dfa9 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/context.cpp @@ -0,0 +1,219 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/context.h" + +#include +#include +#include + +namespace mooncake::tent::ub { +namespace { + +uint64_t steadyNowNs() noexcept { + return static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); +} + +} // namespace + +UbContext::UbContext(Topology::NicID topology_id, DeviceInfo device, + std::shared_ptr adapter) + : topology_id_(topology_id), + device_(std::move(device)), + adapter_(std::move(adapter)) {} + +UbContext::~UbContext() { (void)shutdown(); } + +Status UbContext::initialize(uint32_t jfc_count, const JfcOptions& options) { + std::lock_guard lock(lifecycle_mutex_); + if (state_.load(std::memory_order_relaxed) != State::kUninitialized) { + return Status::InvalidArgument( + "UB context can only be initialized once" LOC_MARK); + } + if (!adapter_ || !device_.active || jfc_count == 0) { + return Status::InvalidArgument( + "Invalid UB context device or JFC count" LOC_MARK); + } + // Each UbJfc owns one send JFC and one receive JFC/JFR. + if (device_.capabilities.max_jfc != 0 && + (static_cast(jfc_count) * 2U) > + device_.capabilities.max_jfc) { + return Status::InvalidArgument( + "Requested UB JFC count exceeds device capability" LOC_MARK); + } + + auto status = adapter_->openContext(device_, handle_); + if (!status.ok()) { + // Providers are allowed to report an error together with a partially + // created handle. Adopt it and make cleanup retryable just like a + // later JFC creation failure. + if (handle_) { + state_.store(State::kDraining, std::memory_order_release); + (void)shutdownLocked(); + } + return status; + } + + jfcs_.reserve(jfc_count); + for (uint32_t i = 0; i < jfc_count; ++i) { + JfcPtr native_jfc; + status = adapter_->createJfc(handle_, options, native_jfc); + if (!status.ok()) { + // A provider may return an error after allocating a partial native + // JFC. Retain that handle as part of the same retryable ownership + // graph instead of letting a temporary shared_ptr destroy it. + if (native_jfc) { + jfcs_.push_back(std::make_shared(jfcs_.size(), adapter_, + std::move(native_jfc))); + } + state_.store(State::kDraining, std::memory_order_release); + (void)shutdownLocked(); + return status; + } + jfcs_.push_back( + std::make_shared(i, adapter_, std::move(native_jfc))); + } + jfc_success_epochs_.assign(jfcs_.size(), 0); + failure_cleanup_complete_ = false; + state_.store(State::kActive, std::memory_order_release); + return Status::OK(); +} + +Status UbContext::shutdown() { + std::lock_guard lock(lifecycle_mutex_); + return shutdownLocked(); +} + +Status UbContext::shutdownLocked() { + auto current = state_.load(std::memory_order_relaxed); + if (current == State::kClosed) return Status::OK(); + if (current == State::kUninitialized && !handle_ && jfcs_.empty()) { + state_.store(State::kClosed, std::memory_order_release); + return Status::OK(); + } + state_.store(State::kDraining, std::memory_order_release); + + Status first_error = Status::OK(); + for (size_t index = jfcs_.size(); index > 0;) { + --index; + auto& jfc = jfcs_[index]; + if (!jfc) { + jfcs_.erase(jfcs_.begin() + static_cast(index)); + continue; + } + auto status = jfc->close(); + if (!status.ok()) { + if (first_error.ok()) first_error = status; + continue; + } + if (jfc->handle()) { + if (first_error.ok()) { + first_error = Status::InternalError( + "URMA adapter retained a JFC after successful " + "delete" LOC_MARK); + } + continue; + } + jfcs_.erase(jfcs_.begin() + static_cast(index)); + } + // A Context is the parent of every JFC. Never ask the provider to delete + // it while a failed JFC remains owned; the next shutdown call resumes at + // the first failed child. + if (!jfcs_.empty()) return first_error; + + if (handle_) { + auto status = adapter_->closeContext(handle_); + if (!status.ok()) { + if (first_error.ok()) first_error = status; + return first_error; + } + if (handle_) { + return Status::InternalError( + "URMA adapter retained a Context after successful " + "close" LOC_MARK); + } + } + state_.store(State::kClosed, std::memory_order_release); + return first_error; +} + +std::shared_ptr UbContext::jfc(size_t index) const { + if (jfcs_.empty()) return nullptr; + return jfcs_[index % jfcs_.size()]; +} + +void UbContext::addInflight(uint64_t bytes) noexcept { + inflight_bytes_.fetch_add(bytes, std::memory_order_relaxed); + outstanding_wrs_.fetch_add(1, std::memory_order_relaxed); +} + +bool UbContext::markUnavailable() noexcept { + std::lock_guard lock(lifecycle_mutex_); + const auto current = state_.load(std::memory_order_relaxed); + if (current != State::kActive && current != State::kFailed) return false; + + const bool newly_failed = current == State::kActive; + const uint64_t now_ns = steadyNowNs(); + if (newly_failed) { + state_.store(State::kFailed, std::memory_order_release); + failure_cleanup_complete_ = false; + failure_started_ns_.store(now_ns, std::memory_order_release); + } + ++failure_epoch_; + if (failure_epoch_ == 0) ++failure_epoch_; + last_failure_ns_.store(now_ns, std::memory_order_release); + return newly_failed; +} + +void UbContext::completeFailureCleanup() noexcept { + std::lock_guard lock(lifecycle_mutex_); + if (state_.load(std::memory_order_relaxed) == State::kFailed) { + failure_cleanup_complete_ = true; + } +} + +bool UbContext::recordPollSuccess(size_t jfc_index, + uint64_t cooldown_ns) noexcept { + std::lock_guard lock(lifecycle_mutex_); + if (state_.load(std::memory_order_relaxed) != State::kFailed || + !failure_cleanup_complete_ || jfc_index >= jfc_success_epochs_.size()) { + return false; + } + jfc_success_epochs_[jfc_index] = failure_epoch_; + if (outstanding_wrs_.load(std::memory_order_acquire) != 0) return false; + + const uint64_t now_ns = steadyNowNs(); + const uint64_t failed_ns = last_failure_ns_.load(std::memory_order_acquire); + if (failed_ns == 0 || now_ns < failed_ns || + now_ns - failed_ns < cooldown_ns) { + return false; + } + if (!std::all_of( + jfc_success_epochs_.begin(), jfc_success_epochs_.end(), + [this](uint64_t epoch) { return epoch == failure_epoch_; })) { + return false; + } + + state_.store(State::kActive, std::memory_order_release); + failure_cleanup_complete_ = false; + recovery_count_.fetch_add(1, std::memory_order_relaxed); + return true; +} + +void UbContext::removeInflight(uint64_t bytes) noexcept { + auto current = inflight_bytes_.load(std::memory_order_relaxed); + while (!inflight_bytes_.compare_exchange_weak( + current, current >= bytes ? current - bytes : 0, + std::memory_order_relaxed)) { + } + current = outstanding_wrs_.load(std::memory_order_relaxed); + while (current != 0 && + !outstanding_wrs_.compare_exchange_weak(current, current - 1, + std::memory_order_relaxed)) { + } +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/jfc.cpp b/mooncake-transfer-engine/tent/src/transport/ub/jfc.cpp new file mode 100644 index 0000000000..9514f0e679 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/jfc.cpp @@ -0,0 +1,29 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/transport/ub/jfc.h" + +namespace mooncake::tent::ub { + +UbJfc::~UbJfc() { (void)close(); } + +Status UbJfc::poll(size_t max_completions, + std::vector& completions) { + if (!valid()) { + return Status::InvalidArgument("UB JFC is not active" LOC_MARK); + } + auto status = adapter_->poll(handle_, max_completions, completions); + if (!status.ok()) { + poll_error_count_.fetch_add(1, std::memory_order_relaxed); + return status; + } + completion_count_.fetch_add(completions.size(), std::memory_order_relaxed); + return Status::OK(); +} + +Status UbJfc::close() { + if (!handle_) return Status::OK(); + return adapter_->deleteJfc(handle_); +} + +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp new file mode 100644 index 0000000000..0c3f4b62b8 --- /dev/null +++ b/mooncake-transfer-engine/tent/src/transport/ub/urma_adapter.cpp @@ -0,0 +1,1626 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/transport/ub/urma_adapter.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(TENT_HAS_REAL_URMA) && TENT_HAS_REAL_URMA +#include +#endif + +namespace mooncake { +namespace tent { +namespace ub { +namespace { + +constexpr char kUnavailableMessage[] = + "native TENT UB is unavailable because this build has no liburma"; + +class UnavailableUrmaAdapter final : public UrmaAdapter { + public: + bool available() const noexcept override { return false; } + uint32_t nativeApiVersion() const noexcept override { return 0; } + size_t nativeSegmentDescriptorSize() const noexcept override { return 0; } + + Status initialize() override { return unavailable(); } + Status shutdown() override { return Status::OK(); } + + Status discoverDevices(std::vector& devices) override { + devices.clear(); + return unavailable(); + } + + Status openContext(const DeviceInfo&, ContextPtr& context) override { + context.reset(); + return unavailable(); + } + + Status closeContext(ContextPtr& context) override { + context.reset(); + return Status::OK(); + } + + Status createJfc(const ContextPtr&, const JfcOptions&, + JfcPtr& jfc) override { + jfc.reset(); + return unavailable(); + } + + Status deleteJfc(JfcPtr& jfc) override { + jfc.reset(); + return Status::OK(); + } + + Status registerLocalSegment(const ContextPtr&, uint64_t, size_t, + const SegmentOptions&, + LocalSegmentPtr& segment) override { + segment.reset(); + return unavailable(); + } + + Status unregisterLocalSegment(LocalSegmentPtr& segment) override { + segment.reset(); + return Status::OK(); + } + + Status importRemoteSegment(const ContextPtr&, const SegmentDescriptor&, + const SegmentOptions&, + RemoteSegmentPtr& segment) override { + segment.reset(); + return unavailable(); + } + + Status unimportRemoteSegment(RemoteSegmentPtr& segment) override { + segment.reset(); + return Status::OK(); + } + + Status createJetty(const ContextPtr&, const JfcPtr&, const JettyOptions&, + JettyPtr& jetty) override { + jetty.reset(); + return unavailable(); + } + + Status deleteJetty(JettyPtr& jetty) override { + jetty.reset(); + return Status::OK(); + } + + Status bindJetty(const JettyPtr&, const RemoteJettyInfo&) override { + return unavailable(); + } + + Status unbindJetty(const JettyPtr& jetty) override { + return jetty ? unavailable() : Status::OK(); + } + + Status resetJetty(const JettyPtr& jetty) override { + return jetty ? unavailable() : Status::OK(); + } + + Status quiesceJetty(const JettyPtr& jetty, uint32_t, + std::vector& completions) override { + completions.clear(); + return jetty ? unavailable() : Status::OK(); + } + + Status post(const JettyPtr&, const std::vector&, + size_t& posted_count) override { + posted_count = 0; + return unavailable(); + } + + Status poll(const JfcPtr&, size_t, + std::vector& completions) override { + completions.clear(); + return unavailable(); + } + + private: + static Status unavailable() { + return Status::NotImplemented(kUnavailableMessage); + } +}; + +#if defined(TENT_HAS_REAL_URMA) && TENT_HAS_REAL_URMA + +Status nativeError(const char* operation, int native_status) { + return Status::InternalError(std::string(operation) + + " failed with URMA status " + + std::to_string(native_status)); +} + +Status nativePointerError(const char* operation) { + const int native_status = errno == 0 ? URMA_FAIL : errno; + return nativeError(operation, native_status); +} + +Status invalidHandle(const char* handle_name) { + return Status::InvalidArgument(std::string("invalid ") + handle_name + + " handle"); +} + +bool checkedRangeContains(uint64_t base, uint64_t range_length, + uint64_t address, uint64_t requested_length) { + if (requested_length == 0 || address < base) return false; + const uint64_t offset = address - base; + return offset <= range_length && requested_length <= range_length - offset; +} + +bool isAllZero(const urma_eid_t& eid) { + for (uint8_t byte : eid.raw) { + if (byte != 0) return false; + } + return true; +} + +std::string formatEid(const urma_eid_t& eid) { + static constexpr char kHex[] = "0123456789abcdef"; + std::string result; + result.reserve(39); + for (size_t i = 0; i < URMA_EID_SIZE; ++i) { + if (i != 0 && i % 2 == 0) result.push_back(':'); + result.push_back(kHex[(eid.raw[i] >> 4) & 0x0f]); + result.push_back(kHex[eid.raw[i] & 0x0f]); + } + return result; +} + +int decodeHexDigit(char value) { + if (value >= '0' && value <= '9') return value - '0'; + if (value >= 'a' && value <= 'f') return value - 'a' + 10; + if (value >= 'A' && value <= 'F') return value - 'A' + 10; + return -1; +} + +bool parseEid(std::string_view encoded, urma_eid_t& eid) { + // Canonical URMA EID form is eight groups of four hex digits. + if (encoded.size() != 39) return false; + + urma_eid_t parsed{}; + size_t cursor = 0; + for (size_t byte_index = 0; byte_index < URMA_EID_SIZE; ++byte_index) { + if (byte_index != 0 && byte_index % 2 == 0) { + if (cursor >= encoded.size() || encoded[cursor] != ':') { + return false; + } + ++cursor; + } + if (cursor + 2 > encoded.size()) return false; + const int high = decodeHexDigit(encoded[cursor]); + const int low = decodeHexDigit(encoded[cursor + 1]); + if (high < 0 || low < 0) return false; + parsed.raw[byte_index] = static_cast((high << 4) | low); + cursor += 2; + } + if (cursor != encoded.size()) return false; + eid = parsed; + return true; +} + +std::string encodeHex(const void* data, size_t length) { + static constexpr char kHex[] = "0123456789ABCDEF"; + const auto* bytes = static_cast(data); + std::string result; + result.resize(length * 2); + for (size_t i = 0; i < length; ++i) { + result[i * 2] = kHex[(bytes[i] >> 4) & 0x0f]; + result[i * 2 + 1] = kHex[bytes[i] & 0x0f]; + } + return result; +} + +bool decodeHex(std::string_view encoded, void* output, size_t output_size) { + if (output == nullptr || encoded.size() != output_size * 2) return false; + auto* bytes = static_cast(output); + for (size_t i = 0; i < output_size; ++i) { + const int high = decodeHexDigit(encoded[i * 2]); + const int low = decodeHexDigit(encoded[i * 2 + 1]); + if (high < 0 || low < 0) return false; + bytes[i] = static_cast((high << 4) | low); + } + return true; +} + +std::string boundedString(const char* value, size_t capacity) { + return std::string(value, strnlen(value, capacity)); +} + +DeviceCapabilities convertCapabilities(const urma_device_attr_t& attr) { + DeviceCapabilities result; + result.max_jfc = attr.dev_cap.max_jfc; + result.max_jfc_depth = attr.dev_cap.max_jfc_depth; + result.max_jfr_depth = attr.dev_cap.max_jfr_depth; + result.max_jetty = attr.dev_cap.max_jetty; + result.max_jetty_depth = attr.dev_cap.max_jfs_depth; + result.max_send_sge = attr.dev_cap.max_jfs_sge; + result.max_remote_sge = attr.dev_cap.max_jfs_rsge; + result.max_message_size = attr.dev_cap.max_msg_size; + result.max_read_size = attr.dev_cap.max_read_size; + result.max_write_size = attr.dev_cap.max_write_size; + result.feature_flags = attr.dev_cap.feature.value; + result.transport_modes = attr.dev_cap.trans_mode; + return result; +} + +bool deviceIsActive(const urma_device_attr_t& attr) { + if (attr.port_cnt == 0) return true; + const size_t port_count = std::min(attr.port_cnt, MAX_PORT_CNT); + for (size_t i = 0; i < port_count; ++i) { + if (attr.port_attr[i].state == URMA_PORT_ACTIVE || + attr.port_attr[i].state == URMA_PORT_ACTIVE_DEFER) { + return true; + } + } + return false; +} + +struct DeviceListDeleter { + void operator()(urma_device_t** devices) const { + if (devices != nullptr) urma_free_device_list(devices); + } +}; +using DeviceList = std::unique_ptr; + +struct EidListDeleter { + void operator()(urma_eid_info_t* eids) const { + if (eids != nullptr) urma_free_eid_list(eids); + } +}; +using EidList = std::unique_ptr; + +std::mutex g_runtime_mutex; +size_t g_runtime_reference_count = 0; +bool g_runtime_owned = false; + +// One lease corresponds to one initialized adapter. Contexts retain their +// adapter's lease, so shutdown cannot call urma_uninit before child handles +// have released all native resources. +class RuntimeLease { + public: + static Status Acquire(std::shared_ptr& output) { + auto lease = std::shared_ptr(new RuntimeLease()); + std::lock_guard lock(g_runtime_mutex); + if (g_runtime_reference_count == 0) { + urma_init_attr_t attributes{}; + const int rc = urma_init(&attributes); + if (rc != URMA_SUCCESS && rc != URMA_EEXIST) { + return nativeError("urma_init", rc); + } + // EEXIST means another component owns the process-wide runtime. + // In that case this adapter must not uninitialize it. + g_runtime_owned = rc == URMA_SUCCESS; + } + ++g_runtime_reference_count; + lease->acquired_ = true; + output = std::move(lease); + return Status::OK(); + } + + ~RuntimeLease() { (void)release(); } + + Status release() { + if (!acquired_) return Status::OK(); + std::lock_guard lock(g_runtime_mutex); + if (g_runtime_reference_count == 0) { + return Status::InternalError( + "URMA runtime reference count underflow"); + } + if (g_runtime_reference_count == 1 && g_runtime_owned) { + const int rc = urma_uninit(); + if (rc != URMA_SUCCESS) return nativeError("urma_uninit", rc); + } + --g_runtime_reference_count; + acquired_ = false; + if (g_runtime_reference_count == 0) { + g_runtime_owned = false; + } + return Status::OK(); + } + + private: + RuntimeLease() = default; + bool acquired_ = false; +}; + +class RealContext final : public Context { + public: + RealContext(std::shared_ptr runtime, DeviceInfo info, + urma_context_t* native) + : runtime_(std::move(runtime)), + info_(std::move(info)), + native_(native) {} + + ~RealContext() override { (void)close(); } + + bool valid() const noexcept override { return native_ != nullptr; } + const DeviceInfo& deviceInfo() const noexcept override { return info_; } + int asyncFd() const noexcept override { + return native_ == nullptr ? -1 : native_->async_fd; + } + + urma_context_t* native() const noexcept { return native_; } + + Status close() { + if (native_ == nullptr) return Status::OK(); + const int rc = urma_delete_context(native_); + if (rc != URMA_SUCCESS) return nativeError("urma_delete_context", rc); + native_ = nullptr; + return Status::OK(); + } + + private: + std::shared_ptr runtime_; + DeviceInfo info_; + urma_context_t* native_ = nullptr; +}; + +class RealJfc final : public Jfc { + public: + explicit RealJfc(std::shared_ptr context) + : context_(std::move(context)) {} + + ~RealJfc() override { (void)close(); } + + Status initialize(const JfcOptions& options) { + if (options.enable_completion_events) { + jfce_ = urma_create_jfce(context_->native()); + if (jfce_ == nullptr) return nativePointerError("urma_create_jfce"); + } + + urma_jfc_cfg_t send_cfg{}; + send_cfg.depth = options.depth; + send_cfg.jfce = jfce_; + send_jfc_ = urma_create_jfc(context_->native(), &send_cfg); + if (send_jfc_ == nullptr) + return nativePointerError("urma_create_jfc(send)"); + + urma_jfc_cfg_t receive_cfg{}; + receive_cfg.depth = options.receiver_depth; + receive_jfc_ = urma_create_jfc(context_->native(), &receive_cfg); + if (receive_jfc_ == nullptr) { + return nativePointerError("urma_create_jfc(receive)"); + } + + urma_jfr_cfg_t receiver_cfg{}; + receiver_cfg.depth = options.receiver_depth; + receiver_cfg.flag.bs.tag_matching = 0; + receiver_cfg.trans_mode = URMA_TM_RC; + receiver_cfg.max_sge = 1; + receiver_cfg.min_rnr_timer = URMA_TYPICAL_MIN_RNR_TIMER; + receiver_cfg.jfc = receive_jfc_; + receiver_cfg.token_value.token = options.token; + receiver_jfr_ = urma_create_jfr(context_->native(), &receiver_cfg); + if (receiver_jfr_ == nullptr) + return nativePointerError("urma_create_jfr"); + return Status::OK(); + } + + bool valid() const noexcept override { + return send_jfc_ != nullptr && receive_jfc_ != nullptr && + receiver_jfr_ != nullptr; + } + + int eventFd() const noexcept override { + return jfce_ == nullptr ? -1 : jfce_->fd; + } + + urma_jfc_t* nativeSendJfc() const noexcept { return send_jfc_; } + urma_jfc_t* nativeReceiveJfc() const noexcept { return receive_jfc_; } + urma_jfr_t* nativeReceiver() const noexcept { return receiver_jfr_; } + const std::shared_ptr& context() const noexcept { + return context_; + } + std::mutex& pollMutex() noexcept { return poll_mutex_; } + + // Callers hold pollMutex(). Markers must survive a failed quiesce call: + // another normal poller may consume FLUSH_ERR_DONE before shutdown retries + // the same Jetty. + void rememberFlushDone(uint32_t jetty_id) { + flush_done_jetty_ids_.insert(jetty_id); + } + bool hasFlushDone(uint32_t jetty_id) const { + return flush_done_jetty_ids_.count(jetty_id) != 0; + } + void clearFlushDone(uint32_t jetty_id) { + flush_done_jetty_ids_.erase(jetty_id); + } + + Status retainSegments(uint32_t jetty_id, + const std::vector& requests) { + std::lock_guard lock(inflight_mutex_); + std::unordered_set new_tokens; + new_tokens.reserve(requests.size()); + for (const WorkRequest& request : requests) { + if (!new_tokens.insert(request.token).second || + inflight_segments_.find(request.token) != + inflight_segments_.end()) { + return Status::InvalidArgument( + "work request token is already in flight on this JFC"); + } + } + for (const WorkRequest& request : requests) { + inflight_segments_.emplace( + request.token, + InflightSegments{request.local_segment, request.remote_segment, + jetty_id}); + } + return Status::OK(); + } + + void releaseSegment(uint64_t token) { + if (token == 0) return; + std::lock_guard lock(inflight_mutex_); + inflight_segments_.erase(token); + } + + // A successful Jetty flush fence proves that no WR posted through that + // Jetty can touch its segments again. Providers are allowed to omit an + // individual completion after the fence, so drop any remaining retention + // entries by Jetty instead of waiting forever for a lost token. + void releaseSegmentsForJetty(uint32_t jetty_id) { + std::lock_guard lock(inflight_mutex_); + for (auto it = inflight_segments_.begin(); + it != inflight_segments_.end();) { + if (it->second.jetty_id == jetty_id) { + it = inflight_segments_.erase(it); + } else { + ++it; + } + } + } + + Status close() { + std::lock_guard lock(poll_mutex_); + { + std::lock_guard inflight_lock(inflight_mutex_); + if (!inflight_segments_.empty()) { + return Status::TooManyRequests( + "Jfc still retains in-flight segment handles"); + } + } + if (receiver_jfr_ != nullptr) { + const int rc = urma_delete_jfr(receiver_jfr_); + if (rc != URMA_SUCCESS) return nativeError("urma_delete_jfr", rc); + receiver_jfr_ = nullptr; + } + if (receive_jfc_ != nullptr) { + const int rc = urma_delete_jfc(receive_jfc_); + if (rc != URMA_SUCCESS) { + return nativeError("urma_delete_jfc(receive)", rc); + } + receive_jfc_ = nullptr; + } + if (send_jfc_ != nullptr) { + const int rc = urma_delete_jfc(send_jfc_); + if (rc != URMA_SUCCESS) { + return nativeError("urma_delete_jfc(send)", rc); + } + send_jfc_ = nullptr; + } + if (jfce_ != nullptr) { + const int rc = urma_delete_jfce(jfce_); + if (rc != URMA_SUCCESS) return nativeError("urma_delete_jfce", rc); + jfce_ = nullptr; + } + return Status::OK(); + } + + private: + std::shared_ptr context_; + urma_jfce_t* jfce_ = nullptr; + urma_jfc_t* send_jfc_ = nullptr; + urma_jfc_t* receive_jfc_ = nullptr; + urma_jfr_t* receiver_jfr_ = nullptr; + std::mutex poll_mutex_; + std::unordered_set flush_done_jetty_ids_; + + struct InflightSegments { + LocalSegmentPtr local; + RemoteSegmentPtr remote; + uint32_t jetty_id{0}; + }; + std::mutex inflight_mutex_; + std::unordered_map inflight_segments_; +}; + +class RealLocalSegment final : public LocalSegment { + public: + RealLocalSegment(std::shared_ptr context, + urma_target_seg_t* native, uint64_t address, + uint64_t length, SegmentDescriptor descriptor) + : context_(std::move(context)), + native_(native), + address_(address), + length_(length), + descriptor_(std::move(descriptor)) {} + + ~RealLocalSegment() override { (void)close(); } + + bool valid() const noexcept override { return native_ != nullptr; } + uint64_t address() const noexcept override { return address_; } + uint64_t length() const noexcept override { return length_; } + const SegmentDescriptor& descriptor() const noexcept override { + return descriptor_; + } + + urma_target_seg_t* native() const noexcept { return native_; } + const std::shared_ptr& context() const noexcept { + return context_; + } + + Status close() { + if (native_ == nullptr) return Status::OK(); + const int rc = urma_unregister_seg(native_); + if (rc != URMA_SUCCESS) return nativeError("urma_unregister_seg", rc); + native_ = nullptr; + return Status::OK(); + } + + private: + std::shared_ptr context_; + urma_target_seg_t* native_ = nullptr; + uint64_t address_ = 0; + uint64_t length_ = 0; + SegmentDescriptor descriptor_; +}; + +class RealRemoteSegment final : public RemoteSegment { + public: + RealRemoteSegment(std::shared_ptr context, + urma_target_seg_t* native, uint64_t address, + uint64_t length, SegmentDescriptor descriptor) + : context_(std::move(context)), + native_(native), + address_(address), + length_(length), + descriptor_(std::move(descriptor)) {} + + ~RealRemoteSegment() override { (void)close(); } + + bool valid() const noexcept override { return native_ != nullptr; } + uint64_t address() const noexcept override { return address_; } + uint64_t length() const noexcept override { return length_; } + const SegmentDescriptor& descriptor() const noexcept override { + return descriptor_; + } + + urma_target_seg_t* native() const noexcept { return native_; } + const std::shared_ptr& context() const noexcept { + return context_; + } + + Status close() { + if (native_ == nullptr) return Status::OK(); + const int rc = urma_unimport_seg(native_); + if (rc != URMA_SUCCESS) return nativeError("urma_unimport_seg", rc); + native_ = nullptr; + return Status::OK(); + } + + private: + std::shared_ptr context_; + urma_target_seg_t* native_ = nullptr; + uint64_t address_ = 0; + uint64_t length_ = 0; + SegmentDescriptor descriptor_; +}; + +class RealJetty final : public Jetty { + public: + RealJetty(std::shared_ptr context, + std::shared_ptr jfc) + : context_(std::move(context)), jfc_(std::move(jfc)) {} + + ~RealJetty() override { cleanup(); } + + Status initialize(const JettyOptions& options) { + urma_jetty_cfg_t config{}; + config.flag.bs.share_jfr = 1; + config.jfs_cfg.depth = options.depth; + config.jfs_cfg.trans_mode = URMA_TM_RC; + config.jfs_cfg.priority = options.priority; + config.jfs_cfg.max_sge = options.max_sge; + config.jfs_cfg.max_rsge = options.max_sge; + config.jfs_cfg.rnr_retry = options.rnr_retry; + config.jfs_cfg.err_timeout = options.error_timeout; + config.jfs_cfg.jfc = jfc_->nativeSendJfc(); + config.shared.jfr = jfc_->nativeReceiver(); + config.shared.jfc = nullptr; + + native_ = urma_create_jetty(context_->native(), &config); + if (native_ == nullptr) return nativePointerError("urma_create_jetty"); + depth_ = options.depth; + return Status::OK(); + } + + bool valid() const noexcept override { return native_ != nullptr; } + uint32_t id() const noexcept override { + return native_ == nullptr ? 0 : native_->jetty_id.id; + } + uint32_t uasid() const noexcept override { + return native_ == nullptr ? 0 : native_->jetty_id.uasid; + } + + const std::shared_ptr& context() const noexcept { + return context_; + } + const std::shared_ptr& jfc() const noexcept { return jfc_; } + std::mutex& mutex() noexcept { return mutex_; } + urma_jetty_t* native() const noexcept { return native_; } + urma_target_jetty_t* remote() const noexcept { return remote_; } + bool postable() const noexcept { return state_ == State::BOUND; } + uint32_t depth() const noexcept { return depth_; } + + Status beginError() { + std::lock_guard lock(mutex_); + if (native_ == nullptr) return invalidHandle("Jetty"); + if (state_ == State::RESET) { + return Status::InvalidArgument("cannot quiesce a reset Jetty"); + } + if (state_ == State::ERROR) return Status::OK(); + + urma_jetty_attr_t attributes{}; + attributes.mask = JETTY_STATE; + attributes.state = URMA_JETTY_STATE_ERROR; + const int rc = urma_modify_jetty(native_, &attributes); + if (rc != URMA_SUCCESS) return nativeError("urma_modify_jetty", rc); + state_ = State::ERROR; + flush_fenced_ = false; + return Status::OK(); + } + + void markFlushFenced() noexcept { + std::lock_guard lock(mutex_); + flush_fenced_ = true; + } + + bool flushFenced() const noexcept { + std::lock_guard lock(mutex_); + return flush_fenced_; + } + + bool errorStarted() const noexcept { + std::lock_guard lock(mutex_); + return state_ == State::ERROR; + } + + Status bind(const RemoteJettyInfo& remote_info, + const urma_eid_t& remote_eid) { + std::lock_guard lock(mutex_); + if (native_ == nullptr) return invalidHandle("Jetty"); + if (state_ == State::RESET) { + return Status::InvalidArgument("cannot bind a reset Jetty"); + } + if (remote_ != nullptr) { + if (!needs_native_unbind_ || state_ != State::BOUND) { + return Status::TooManyRequests( + "Jetty has an imported peer pending cleanup"); + } + const bool same = remote_info.id == remote_id_ && + remote_info.uasid == remote_uasid_ && + std::memcmp(remote_eid.raw, remote_eid_.raw, + URMA_EID_SIZE) == 0; + return same ? Status::OK() + : Status::InvalidArgument( + "Jetty is already bound to another peer"); + } + + urma_rjetty_t descriptor{}; + descriptor.jetty_id.eid = remote_eid; + descriptor.jetty_id.uasid = remote_info.uasid; + descriptor.jetty_id.id = remote_info.id; + descriptor.trans_mode = URMA_TM_RC; + descriptor.type = URMA_JETTY; + descriptor.tp_type = URMA_CTP; + + urma_token_t token{.token = remote_info.token}; + urma_target_jetty_t* imported = + urma_import_jetty(context_->native(), &descriptor, &token); + if (imported == nullptr) { + return nativePointerError("urma_import_jetty"); + } + + const int rc = urma_bind_jetty(native_, imported); + if (rc != URMA_SUCCESS && rc != URMA_EEXIST) { + const int rollback_rc = urma_unimport_jetty(imported); + if (rollback_rc != URMA_SUCCESS) { + // The target Jetty was imported but never bound. Preserve the + // raw handle as an explicit cleanup-only phase so endpoint + // failure teardown can retry unimport without issuing an + // invalid native unbind. + remote_ = imported; + remote_eid_ = remote_eid; + remote_id_ = remote_info.id; + remote_uasid_ = remote_info.uasid; + needs_native_unbind_ = false; + return Status::InternalError( + "urma_bind_jetty failed with URMA status " + + std::to_string(rc) + + "; rollback urma_unimport_jetty failed with URMA status " + + std::to_string(rollback_rc) + + "; imported target retained for retry"); + } + return nativeError("urma_bind_jetty", rc); + } + + remote_ = imported; + remote_eid_ = remote_eid; + remote_id_ = remote_info.id; + remote_uasid_ = remote_info.uasid; + needs_native_unbind_ = true; + state_ = State::BOUND; + return Status::OK(); + } + + Status reset() { + std::lock_guard lock(mutex_); + if (native_ == nullptr) return invalidHandle("Jetty"); + if (state_ == State::RESET) return Status::OK(); + if (state_ == State::ERROR && !flush_fenced_) { + return Status::InvalidArgument( + "cannot reset an ERROR Jetty before its flush fence"); + } + + urma_jetty_attr_t attributes{}; + attributes.mask = JETTY_STATE; + attributes.state = URMA_JETTY_STATE_RESET; + const int rc = urma_modify_jetty(native_, &attributes); + if (rc != URMA_SUCCESS) return nativeError("urma_modify_jetty", rc); + state_ = State::RESET; + return Status::OK(); + } + + Status unbind() { + std::lock_guard lock(mutex_); + if (native_ == nullptr) return invalidHandle("Jetty"); + if (remote_ == nullptr) return Status::OK(); + + if (needs_native_unbind_) { + const int unbind_rc = urma_unbind_jetty(native_); + if (unbind_rc != URMA_SUCCESS) { + return nativeError("urma_unbind_jetty", unbind_rc); + } + // UMDK clears native_->remote_jetty on success. Preserve this + // phase across an unimport failure so a retry does not issue an + // invalid second unbind and can proceed directly to unimport. + needs_native_unbind_ = false; + } + const int unimport_rc = urma_unimport_jetty(remote_); + if (unimport_rc != URMA_SUCCESS) { + return nativeError("urma_unimport_jetty", unimport_rc); + } + remote_ = nullptr; + needs_native_unbind_ = false; + if (state_ != State::RESET) state_ = State::CREATED; + return Status::OK(); + } + + Status close() { + std::lock_guard lock(mutex_); + if (native_ == nullptr) return Status::OK(); + if (remote_ != nullptr || state_ != State::RESET) { + return Status::InvalidArgument( + "Jetty must be reset and unbound before deletion"); + } + const int rc = urma_delete_jetty(native_); + if (rc != URMA_SUCCESS) return nativeError("urma_delete_jetty", rc); + native_ = nullptr; + return Status::OK(); + } + + private: + enum class State : uint8_t { CREATED, BOUND, ERROR, RESET }; + + void cleanup() noexcept { + std::lock_guard lock(mutex_); + if (native_ == nullptr) return; + + if (state_ != State::RESET) { + urma_jetty_attr_t attributes{}; + attributes.mask = JETTY_STATE; + attributes.state = URMA_JETTY_STATE_RESET; + (void)urma_modify_jetty(native_, &attributes); + } + + if (remote_ != nullptr) { + if (!needs_native_unbind_ || + urma_unbind_jetty(native_) == URMA_SUCCESS) { + needs_native_unbind_ = false; + (void)urma_unimport_jetty(remote_); + remote_ = nullptr; + } + } + (void)urma_delete_jetty(native_); + native_ = nullptr; + + // If unbind failed, deleting the local Jetty has severed the binding; + // make a final best-effort attempt to release the imported peer. + if (remote_ != nullptr) { + (void)urma_unimport_jetty(remote_); + remote_ = nullptr; + } + state_ = State::RESET; + } + + std::shared_ptr context_; + std::shared_ptr jfc_; + urma_jetty_t* native_ = nullptr; + urma_target_jetty_t* remote_ = nullptr; + urma_eid_t remote_eid_{}; + uint32_t remote_id_ = 0; + uint32_t remote_uasid_ = 0; + // True only while remote_ is natively bound. A non-null remote_ with this + // flag clear is an imported-only cleanup phase retained after rollback or + // after a successful unbind followed by a failed unimport. + bool needs_native_unbind_{false}; + uint32_t depth_ = 0; + State state_ = State::CREATED; + bool flush_fenced_{false}; + mutable std::mutex mutex_; +}; + +uint32_t nativeAccess(uint32_t access) { + uint32_t result = 0; + if ((access & SEGMENT_ACCESS_LOCAL_ONLY) != 0) { + result |= URMA_ACCESS_LOCAL_ONLY; + } + if ((access & SEGMENT_ACCESS_READ) != 0) result |= URMA_ACCESS_READ; + if ((access & SEGMENT_ACCESS_WRITE) != 0) result |= URMA_ACCESS_WRITE; + if ((access & SEGMENT_ACCESS_ATOMIC) != 0) result |= URMA_ACCESS_ATOMIC; + return result; +} + +Status validateSegmentOptions(const SegmentOptions& options) { + constexpr uint32_t kAllAccess = SEGMENT_ACCESS_READ | SEGMENT_ACCESS_WRITE | + SEGMENT_ACCESS_ATOMIC | + SEGMENT_ACCESS_LOCAL_ONLY; + if (options.access == 0 || (options.access & ~kAllAccess) != 0) { + return Status::InvalidArgument("invalid segment access mask"); + } + if ((options.access & SEGMENT_ACCESS_LOCAL_ONLY) != 0 && + options.access != SEGMENT_ACCESS_LOCAL_ONLY) { + return Status::InvalidArgument( + "local-only segment access cannot include remote permissions"); + } + return Status::OK(); +} + +CompletionCategory classifyCompletion(int status) { + switch (status) { + case URMA_CR_SUCCESS: + return CompletionCategory::SUCCESS; + case URMA_CR_LOC_OPERATION_ERR: + return CompletionCategory::LOCAL_DEVICE_ERROR; + case URMA_CR_LOC_LEN_ERR: + case URMA_CR_LOC_ACCESS_ERR: + case URMA_CR_LOC_DATA_POISON: + return CompletionCategory::MEMORY_ERROR; + case URMA_CR_REM_RESP_LEN_ERR: + case URMA_CR_REM_UNSUPPORTED_REQ_ERR: + case URMA_CR_REM_OPERATION_ERR: + case URMA_CR_REM_ACCESS_ABORT_ERR: + case URMA_CR_RNR_RETRY_CNT_EXC_ERR: + case URMA_CR_REM_DATA_POISON: + return CompletionCategory::REMOTE_PATH_ERROR; + case URMA_CR_ACK_TIMEOUT_ERR: + return CompletionCategory::TIMEOUT; + case URMA_CR_UNSUPPORTED_OPCODE_ERR: + case URMA_CR_WR_FLUSH_ERR: + case URMA_CR_WR_SUSPEND_DONE: + case URMA_CR_WR_FLUSH_ERR_DONE: + case URMA_CR_WR_UNHANDLED: + return CompletionCategory::ENDPOINT_ERROR; + default: + return CompletionCategory::UNKNOWN_ERROR; + } +} + +bool isEntityMarker(const urma_cr_t& completion) { + return completion.status == URMA_CR_WR_SUSPEND_DONE || + completion.status == URMA_CR_WR_FLUSH_ERR_DONE; +} + +Completion convertCompletion(const urma_cr_t& native) { + Completion completion; + completion.category = classifyCompletion(native.status); + completion.native_status = native.status; + completion.completed_bytes = native.completion_len; + completion.local_jetty_id = native.local_id; + // UMDK v25.12 identifies entity markers by status. Their user_ctx is + // invalid; every other CR (including WR_UNHANDLED returned by flush) is a + // real WR completion and carries the original token. + completion.token = isEntityMarker(native) ? 0 : native.user_ctx; + return completion; +} + +template +Status releaseTypedHandle(std::shared_ptr& handle, const char* name) { + if (!handle) return Status::OK(); + if (handle.use_count() != 1) { + return Status::TooManyRequests(std::string(name) + + " handle is still retained"); + } + auto native = std::dynamic_pointer_cast(handle); + if (!native) return invalidHandle(name); + CHECK_STATUS(native->close()); + handle.reset(); + return Status::OK(); +} + +class RealUrmaAdapter final : public UrmaAdapter { + public: + bool available() const noexcept override { return true; } + uint32_t nativeApiVersion() const noexcept override { + return URMA_API_VERSION; + } + size_t nativeSegmentDescriptorSize() const noexcept override { + return sizeof(urma_seg_t); + } + + Status initialize() override { + std::lock_guard lock(mutex_); + if (runtime_) return Status::OK(); + return RuntimeLease::Acquire(runtime_); + } + + Status shutdown() override { + std::lock_guard lock(mutex_); + if (!runtime_) return Status::OK(); + if (runtime_.use_count() != 1) { + return Status::TooManyRequests( + "URMA runtime is still retained by native handles"); + } + CHECK_STATUS(runtime_->release()); + runtime_.reset(); + return Status::OK(); + } + + Status discoverDevices(std::vector& devices) override { + devices.clear(); + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + + int device_count = 0; + DeviceList native_devices(urma_get_device_list(&device_count)); + if (!native_devices || device_count < 0) { + return nativePointerError("urma_get_device_list"); + } + + for (int i = 0; i < device_count; ++i) { + urma_device_t* native_device = native_devices.get()[i]; + if (native_device == nullptr) continue; + + urma_device_attr_t attributes{}; + const int query_rc = urma_query_device(native_device, &attributes); + if (query_rc != URMA_SUCCESS) { + return nativeError("urma_query_device", query_rc); + } + + uint32_t eid_count = 0; + EidList eids(urma_get_eid_list(native_device, &eid_count)); + if (!eids && eid_count != 0) { + return nativePointerError("urma_get_eid_list"); + } + + const std::string native_name = + boundedString(native_device->name, URMA_MAX_NAME); + const std::string native_path = + boundedString(native_device->path, URMA_MAX_PATH); + for (uint32_t eid_position = 0; eid_position < eid_count; + ++eid_position) { + const urma_eid_info_t& eid = eids.get()[eid_position]; + DeviceInfo info; + info.native_device_name = native_name; + info.native_device_path = native_path; + info.eid_index = eid.eid_index; + info.eid = formatEid(eid.eid); + info.topology_name = "ub:" + native_name + ":eid" + + std::to_string(eid.eid_index); + info.active = deviceIsActive(attributes) && !isAllZero(eid.eid); + info.capabilities = convertCapabilities(attributes); + devices.push_back(std::move(info)); + } + } + return Status::OK(); + } + + Status openContext(const DeviceInfo& requested, + ContextPtr& output) override { + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + if (requested.native_device_name.empty()) { + return Status::InvalidArgument("URMA native device name is empty"); + } + urma_eid_t requested_eid{}; + if (!parseEid(requested.eid, requested_eid) || + isAllZero(requested_eid)) { + return Status::InvalidArgument("invalid or null URMA EID"); + } + + int device_count = 0; + DeviceList devices(urma_get_device_list(&device_count)); + if (!devices || device_count < 0) { + return nativePointerError("urma_get_device_list"); + } + + for (int i = 0; i < device_count; ++i) { + urma_device_t* native_device = devices.get()[i]; + if (native_device == nullptr || + boundedString(native_device->name, URMA_MAX_NAME) != + requested.native_device_name) { + continue; + } + + uint32_t eid_count = 0; + EidList eids(urma_get_eid_list(native_device, &eid_count)); + if (!eids && eid_count != 0) { + return nativePointerError("urma_get_eid_list"); + } + bool eid_found = false; + for (uint32_t j = 0; j < eid_count; ++j) { + if (eids.get()[j].eid_index == requested.eid_index && + std::memcmp(eids.get()[j].eid.raw, requested_eid.raw, + URMA_EID_SIZE) == 0) { + eid_found = true; + break; + } + } + if (!eid_found) { + return Status::DeviceNotFound( + "requested EID is no longer present on URMA device " + + requested.native_device_name); + } + + urma_device_attr_t attributes{}; + const int query_rc = urma_query_device(native_device, &attributes); + if (query_rc != URMA_SUCCESS) { + return nativeError("urma_query_device", query_rc); + } + urma_context_t* native_context = + urma_create_context(native_device, requested.eid_index); + if (native_context == nullptr) { + return nativePointerError("urma_create_context"); + } + + DeviceInfo current = requested; + current.native_device_path = + boundedString(native_device->path, URMA_MAX_PATH); + current.active = deviceIsActive(attributes); + current.capabilities = convertCapabilities(attributes); + if (current.topology_name.empty()) { + current.topology_name = "ub:" + current.native_device_name + + ":eid" + + std::to_string(current.eid_index); + } + output = std::make_shared( + std::move(runtime), std::move(current), native_context); + return Status::OK(); + } + return Status::DeviceNotFound("URMA device not found: " + + requested.native_device_name); + } + + Status closeContext(ContextPtr& context) override { + return releaseTypedHandle(context, "Context"); + } + + Status createJfc(const ContextPtr& context, const JfcOptions& options, + JfcPtr& output) override { + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + auto real_context = std::dynamic_pointer_cast(context); + if (!real_context || !real_context->valid()) { + return invalidHandle("Context"); + } + if (options.depth == 0 || options.receiver_depth == 0) { + return Status::InvalidArgument("JFC depths must be positive"); + } + const auto& caps = real_context->deviceInfo().capabilities; + if (caps.max_jfc != 0 && caps.max_jfc < 2) { + return Status::InvalidArgument( + "URMA device cannot provide send and receive JFCs"); + } + if (caps.max_jfc_depth != 0 && + (options.depth > caps.max_jfc_depth || + options.receiver_depth > caps.max_jfc_depth)) { + return Status::InvalidArgument( + "requested JFC depth exceeds device capability"); + } + if (caps.max_jfr_depth != 0 && + options.receiver_depth > caps.max_jfr_depth) { + return Status::InvalidArgument( + "requested JFR depth exceeds device capability"); + } + + auto jfc = std::make_shared(std::move(real_context)); + auto status = jfc->initialize(options); + if (!status.ok()) { + // initialize may already own a JFCE/JFC/JFR prefix. Return the + // wrapper alongside the error so the caller can retain it and + // drive retryable delete instead of relying on its destructor. + output = std::move(jfc); + return status; + } + output = std::move(jfc); + return Status::OK(); + } + + Status deleteJfc(JfcPtr& jfc) override { + return releaseTypedHandle(jfc, "Jfc"); + } + + Status registerLocalSegment(const ContextPtr& context, uint64_t address, + size_t length, const SegmentOptions& options, + LocalSegmentPtr& output) override { + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + auto real_context = std::dynamic_pointer_cast(context); + if (!real_context || !real_context->valid()) { + return invalidHandle("Context"); + } + CHECK_STATUS(validateSegmentOptions(options)); + if (address == 0 || length == 0 || + length > std::numeric_limits::max() - address) { + return Status::InvalidArgument("invalid local segment range"); + } + + urma_reg_seg_flag_t flags{}; + flags.bs.token_policy = URMA_TOKEN_NONE; + flags.bs.cacheable = + options.cacheable ? URMA_CACHEABLE : URMA_NON_CACHEABLE; + flags.bs.access = nativeAccess(options.access); + + urma_seg_cfg_t config{}; + config.va = address; + config.len = length; + config.token_value.token = options.token; + config.flag = flags; + urma_target_seg_t* native_segment = + urma_register_seg(real_context->native(), &config); + if (native_segment == nullptr) { + return nativePointerError("urma_register_seg"); + } + + urma_seg_t wire_descriptor{}; + wire_descriptor.ubva = native_segment->seg.ubva; + wire_descriptor.len = native_segment->seg.len; + wire_descriptor.attr = native_segment->seg.attr; + wire_descriptor.token_id = native_segment->seg.token_id; + + SegmentDescriptor descriptor; + descriptor.urma_api_version = URMA_API_VERSION; + descriptor.urma_abi_size = sizeof(urma_seg_t); + descriptor.hex = encodeHex(&wire_descriptor, sizeof(wire_descriptor)); + output = std::make_shared( + std::move(real_context), native_segment, address, length, + std::move(descriptor)); + return Status::OK(); + } + + Status unregisterLocalSegment(LocalSegmentPtr& segment) override { + return releaseTypedHandle(segment, "LocalSegment"); + } + + Status importRemoteSegment(const ContextPtr& context, + const SegmentDescriptor& descriptor, + const SegmentOptions& options, + RemoteSegmentPtr& output) override { + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + auto real_context = std::dynamic_pointer_cast(context); + if (!real_context || !real_context->valid()) { + return invalidHandle("Context"); + } + CHECK_STATUS(validateSegmentOptions(options)); + if (descriptor.schema_version != SegmentDescriptor::kSchemaVersion) { + return Status::InvalidArgument( + "unsupported URMA segment descriptor schema"); + } + if (descriptor.urma_api_version != URMA_API_VERSION) { + return Status::InvalidArgument( + "URMA segment descriptor API version mismatch"); + } + if (descriptor.urma_abi_size != sizeof(urma_seg_t)) { + return Status::InvalidArgument( + "URMA segment descriptor ABI size mismatch"); + } + + urma_seg_t native_descriptor{}; + if (!decodeHex(descriptor.hex, &native_descriptor, + sizeof(native_descriptor))) { + return Status::InvalidArgument( + "malformed URMA segment descriptor hex"); + } + if (native_descriptor.len == 0 || + native_descriptor.len > std::numeric_limits::max() - + native_descriptor.ubva.va || + native_descriptor.attr.bs.reserved != 0) { + return Status::InvalidArgument( + "invalid URMA segment descriptor contents"); + } + + urma_import_seg_flag_t flags{}; + flags.bs.cacheable = + options.cacheable ? URMA_CACHEABLE : URMA_NON_CACHEABLE; + flags.bs.access = nativeAccess(options.access); + flags.bs.mapping = URMA_SEG_NOMAP; + urma_token_t token{.token = options.token}; + urma_target_seg_t* imported = urma_import_seg( + real_context->native(), &native_descriptor, &token, 0, flags); + if (imported == nullptr) return nativePointerError("urma_import_seg"); + + // urma_ubva_t is packed; copy its fields before passing them through + // forwarding references used by make_shared. + const uint64_t remote_address = native_descriptor.ubva.va; + const uint64_t remote_length = native_descriptor.len; + output = std::make_shared(std::move(real_context), + imported, remote_address, + remote_length, descriptor); + return Status::OK(); + } + + Status unimportRemoteSegment(RemoteSegmentPtr& segment) override { + return releaseTypedHandle(segment, "RemoteSegment"); + } + + Status createJetty(const ContextPtr& context, const JfcPtr& jfc, + const JettyOptions& options, JettyPtr& output) override { + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + auto real_context = std::dynamic_pointer_cast(context); + auto real_jfc = std::dynamic_pointer_cast(jfc); + if (!real_context || !real_context->valid()) { + return invalidHandle("Context"); + } + if (!real_jfc || !real_jfc->valid()) return invalidHandle("Jfc"); + if (real_jfc->context().get() != real_context.get()) { + return Status::InvalidArgument( + "Jfc belongs to a different Context"); + } + if (options.depth == 0 || options.max_sge == 0 || + options.priority > URMA_MAX_PRIORITY || options.rnr_retry > 7 || + options.error_timeout > 31) { + return Status::InvalidArgument("invalid Jetty options"); + } + const auto& caps = real_context->deviceInfo().capabilities; + if (caps.max_jetty_depth != 0 && options.depth > caps.max_jetty_depth) { + return Status::InvalidArgument( + "requested Jetty depth exceeds device capability"); + } + if ((caps.max_send_sge != 0 && options.max_sge > caps.max_send_sge) || + (caps.max_remote_sge != 0 && + options.max_sge > caps.max_remote_sge)) { + return Status::InvalidArgument( + "requested Jetty SGE count exceeds device capability"); + } + + auto jetty = std::make_shared(std::move(real_context), + std::move(real_jfc)); + auto status = jetty->initialize(options); + if (!status.ok()) { + output = std::move(jetty); + return status; + } + output = std::move(jetty); + return Status::OK(); + } + + Status deleteJetty(JettyPtr& jetty) override { + return releaseTypedHandle(jetty, "Jetty"); + } + + Status bindJetty(const JettyPtr& jetty, + const RemoteJettyInfo& remote) override { + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + auto real_jetty = std::dynamic_pointer_cast(jetty); + if (!real_jetty || !real_jetty->valid()) return invalidHandle("Jetty"); + if (remote.id == 0) { + return Status::InvalidArgument("remote Jetty ID must be non-zero"); + } + urma_eid_t remote_eid{}; + if (!parseEid(remote.eid, remote_eid) || isAllZero(remote_eid)) { + return Status::InvalidArgument("invalid or null remote EID"); + } + return real_jetty->bind(remote, remote_eid); + } + + Status unbindJetty(const JettyPtr& jetty) override { + if (!jetty) return Status::OK(); + auto real_jetty = std::dynamic_pointer_cast(jetty); + if (!real_jetty || !real_jetty->valid()) return invalidHandle("Jetty"); + return real_jetty->unbind(); + } + + Status resetJetty(const JettyPtr& jetty) override { + if (!jetty) return Status::OK(); + auto real_jetty = std::dynamic_pointer_cast(jetty); + if (!real_jetty || !real_jetty->valid()) return invalidHandle("Jetty"); + return real_jetty->reset(); + } + + Status quiesceJetty(const JettyPtr& jetty, uint32_t timeout_ms, + std::vector& completions) override { + completions.clear(); + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + auto real_jetty = std::dynamic_pointer_cast(jetty); + if (!real_jetty || !real_jetty->valid()) return invalidHandle("Jetty"); + if (timeout_ms == 0 || real_jetty->depth() == 0) { + return Status::InvalidArgument( + "Jetty quiesce requires a non-zero timeout and depth"); + } + if (real_jetty->flushFenced()) return Status::OK(); + + auto real_jfc = real_jetty->jfc(); + if (!real_jfc || !real_jfc->valid()) return invalidHandle("Jfc"); + + // Holding the JFC poll lock before entering ERROR makes consumption + // of the provider's fake FLUSH_ERR_DONE marker atomic with respect to + // normal pollers. Holding the Jetty's own lock inside beginError also + // fences a post already crossing the native boundary. + std::unique_lock poll_lock(real_jfc->pollMutex()); + if (real_jetty->flushFenced()) return Status::OK(); + if (!real_jetty->errorStarted()) { + // Native Jetty IDs may be reused after deletion. Only discard a + // stale marker when starting a genuinely new ERROR epoch. + real_jfc->clearFlushDone(real_jetty->id()); + } + CHECK_STATUS(real_jetty->beginError()); + + const auto deadline = std::chrono::steady_clock::now() + + std::chrono::milliseconds(timeout_ms); + bool flush_done = real_jfc->hasFlushDone(real_jetty->id()); + std::array polled{}; + while (!flush_done) { + const int count = + urma_poll_jfc(real_jfc->nativeSendJfc(), + static_cast(polled.size()), polled.data()); + if (count < 0) return nativeError("urma_poll_jfc", count); + for (int i = 0; i < count; ++i) { + const auto& native = polled[static_cast(i)]; + if (native.status == URMA_CR_WR_FLUSH_ERR_DONE) { + real_jfc->rememberFlushDone(native.local_id); + } + auto completion = convertCompletion(native); + if (completion.token != 0) { + real_jfc->releaseSegment(completion.token); + completions.push_back(completion); + } + } + flush_done = flush_done || real_jfc->hasFlushDone(real_jetty->id()); + if (flush_done) break; + if (std::chrono::steady_clock::now() >= deadline) { + return Status::RdmaError( + "timed out waiting for URMA Jetty flush-done fence"); + } + std::this_thread::sleep_for(std::chrono::microseconds(20)); + } + + const size_t flush_batch = + std::min(real_jetty->depth(), polled.size()); + size_t total_flushed = 0; + while (true) { + const int count = + urma_flush_jetty(real_jetty->native(), + static_cast(flush_batch), polled.data()); + if (count < 0) return nativeError("urma_flush_jetty", count); + if (count == 0) break; + total_flushed += static_cast(count); + if (total_flushed > real_jetty->depth()) { + return Status::InternalError( + "URMA Jetty flush exceeded its queue depth"); + } + for (int i = 0; i < count; ++i) { + auto completion = + convertCompletion(polled[static_cast(i)]); + if (completion.token == 0) { + return Status::InternalError( + "URMA Jetty flush returned an entity marker"); + } + real_jfc->releaseSegment(completion.token); + completions.push_back(completion); + } + } + real_jfc->clearFlushDone(real_jetty->id()); + real_jfc->releaseSegmentsForJetty(real_jetty->id()); + real_jetty->markFlushFenced(); + return Status::OK(); + } + + Status post(const JettyPtr& jetty, const std::vector& requests, + size_t& posted_count) override { + posted_count = 0; + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + auto real_jetty = std::dynamic_pointer_cast(jetty); + if (!real_jetty || !real_jetty->valid()) return invalidHandle("Jetty"); + if (requests.empty()) return Status::OK(); + + struct NativeWorkRequest { + urma_jfs_wr_t wr{}; + urma_sge_t local_sge{}; + urma_sge_t remote_sge{}; + }; + std::vector native_requests(requests.size()); + + std::lock_guard jetty_lock(real_jetty->mutex()); + if (!real_jetty->postable() || real_jetty->remote() == nullptr) { + return Status::InvalidArgument("Jetty is not bound and postable"); + } + for (size_t i = 0; i < requests.size(); ++i) { + const WorkRequest& request = requests[i]; + if (request.token == 0) { + return Status::InvalidArgument( + "work request token zero is reserved"); + } + if (request.length == 0 || + request.length > std::numeric_limits::max()) { + return Status::InvalidArgument("invalid work request length"); + } + + auto local = std::dynamic_pointer_cast( + request.local_segment); + auto remote = std::dynamic_pointer_cast( + request.remote_segment); + if (!local || !local->valid()) return invalidHandle("LocalSegment"); + if (!remote || !remote->valid()) { + return invalidHandle("RemoteSegment"); + } + if (local->context().get() != real_jetty->context().get() || + remote->context().get() != real_jetty->context().get()) { + return Status::InvalidArgument( + "work request segments belong to another Context"); + } + if (!checkedRangeContains(local->address(), local->length(), + request.local_address, request.length) || + !checkedRangeContains(remote->address(), remote->length(), + request.remote_address, request.length)) { + return Status::InvalidArgument( + "work request lies outside a registered segment"); + } + + NativeWorkRequest& native = native_requests[i]; + native.local_sge.addr = request.local_address; + native.local_sge.len = static_cast(request.length); + native.local_sge.tseg = local->native(); + native.remote_sge.addr = request.remote_address; + native.remote_sge.len = static_cast(request.length); + native.remote_sge.tseg = remote->native(); + + native.wr.opcode = request.operation == Operation::READ + ? URMA_OPC_READ + : URMA_OPC_WRITE; + native.wr.flag.bs.complete_enable = 1; + native.wr.tjetty = real_jetty->remote(); + native.wr.user_ctx = request.token; + if (request.operation == Operation::READ) { + native.wr.rw.src.sge = &native.remote_sge; + native.wr.rw.dst.sge = &native.local_sge; + } else { + native.wr.rw.src.sge = &native.local_sge; + native.wr.rw.dst.sge = &native.remote_sge; + } + native.wr.rw.src.num_sge = 1; + native.wr.rw.dst.num_sge = 1; + native.wr.next = + i + 1 == requests.size() ? nullptr : &native_requests[i + 1].wr; + } + + // Retain both segment handles before crossing the native post + // boundary. This prevents buffer removal from unregistering memory + // while a WR is still in flight. poll() releases the references by + // completion token. + CHECK_STATUS( + real_jetty->jfc()->retainSegments(real_jetty->id(), requests)); + + urma_jfs_wr_t* bad_wr = nullptr; + const int rc = urma_post_jetty_send_wr( + real_jetty->native(), &native_requests.front().wr, &bad_wr); + if (rc == URMA_SUCCESS) { + posted_count = requests.size(); + return Status::OK(); + } + if (bad_wr != nullptr) { + for (size_t i = 0; i < native_requests.size(); ++i) { + if (bad_wr == &native_requests[i].wr) { + posted_count = i; + break; + } + } + } + for (size_t i = posted_count; i < requests.size(); ++i) { + real_jetty->jfc()->releaseSegment(requests[i].token); + } + return nativeError("urma_post_jetty_send_wr", rc); + } + + Status poll(const JfcPtr& jfc, size_t max_completions, + std::vector& completions) override { + completions.clear(); + std::shared_ptr runtime; + CHECK_STATUS(getRuntime(runtime)); + auto real_jfc = std::dynamic_pointer_cast(jfc); + if (!real_jfc || !real_jfc->valid()) return invalidHandle("Jfc"); + if (max_completions == 0 || + max_completions > static_cast(INT_MAX)) { + return Status::InvalidArgument("invalid maximum completion count"); + } + + std::vector native_completions(max_completions); + std::lock_guard lock(real_jfc->pollMutex()); + const int count = urma_poll_jfc(real_jfc->nativeSendJfc(), + static_cast(max_completions), + native_completions.data()); + if (count < 0) return nativeError("urma_poll_jfc", count); + + completions.reserve(static_cast(count)); + for (int i = 0; i < count; ++i) { + const urma_cr_t& native = native_completions[i]; + Completion completion = convertCompletion(native); + if (native.status == URMA_CR_WR_FLUSH_ERR_DONE) { + real_jfc->rememberFlushDone(native.local_id); + } + completions.push_back(completion); + real_jfc->releaseSegment(completion.token); + } + return Status::OK(); + } + + private: + Status getRuntime(std::shared_ptr& output) const { + std::lock_guard lock(mutex_); + if (!runtime_) { + return Status::InvalidArgument( + "URMA adapter is not initialized or has been shut down"); + } + output = runtime_; + return Status::OK(); + } + + mutable std::mutex mutex_; + std::shared_ptr runtime_; +}; + +#endif // defined(TENT_HAS_REAL_URMA) && TENT_HAS_REAL_URMA + +} // namespace + +std::shared_ptr createDefaultUrmaAdapter() { +#if defined(TENT_HAS_REAL_URMA) && TENT_HAS_REAL_URMA + return std::make_shared(); +#else + return std::make_shared(); +#endif +} + +} // namespace ub +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/CMakeLists.txt b/mooncake-transfer-engine/tent/tests/CMakeLists.txt index 0b499856aa..23ad86c835 100644 --- a/mooncake-transfer-engine/tent/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tent/tests/CMakeLists.txt @@ -20,6 +20,35 @@ target_include_directories(metrics_config_loader_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME metrics_config_loader_test COMMAND metrics_config_loader_test) +# TENT Metrics HTTP server test: validates that initialize() detects a busy +# metrics port and degrades to log-only metrics instead of falsely reporting a +# listening endpoint. +add_executable(tent_metrics_http_server_test metrics_http_server_test.cpp) +target_link_libraries(tent_metrics_http_server_test + PRIVATE tent_metrics tent_common gtest gtest_main glog) +if(TARGET asio_shared) + target_link_libraries(tent_metrics_http_server_test PRIVATE asio_shared) +endif() +target_include_directories(tent_metrics_http_server_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_metrics_http_server_test + COMMAND tent_metrics_http_server_test) + +# TENT Metrics recording test: drives a real TransferEngineImpl with a +# FakeTransport and asserts that counters/histograms move correctly via the JSON +# output. Covers completed/failed/infeasible-deadline recording and histogram +# bucket invariants. Compiles to a stub test when metrics are disabled at +# compile time. +add_executable(tent_metrics_recording_test metrics_recording_test.cpp) +target_link_libraries(tent_metrics_recording_test PRIVATE gtest gtest_main + tent_link_group glog) +if(TARGET asio_shared) + target_link_libraries(tent_metrics_recording_test PRIVATE asio_shared) +endif() +target_include_directories(tent_metrics_recording_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_metrics_recording_test COMMAND tent_metrics_recording_test) + add_executable(admission_queue_test admission_queue_test.cpp ../src/runtime/admission_queue.cpp) target_link_libraries(admission_queue_test PRIVATE tent_common gtest gtest_main) @@ -27,12 +56,65 @@ target_include_directories(admission_queue_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME admission_queue_test COMMAND admission_queue_test) +add_executable(receiver_credit_test receiver_credit_test.cpp + ../src/runtime/receiver_credit.cpp) +target_link_libraries(receiver_credit_test PRIVATE tent_common gtest gtest_main) +target_include_directories(receiver_credit_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME receiver_credit_test COMMAND receiver_credit_test) + +add_executable(bw_arbitration_test bw_arbitration_test.cpp) +target_link_libraries(bw_arbitration_test PRIVATE tent_common gtest gtest_main) +target_include_directories(bw_arbitration_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME bw_arbitration_test COMMAND bw_arbitration_test) + +# Reproducible hot-path microbenchmark; intentionally not registered with ctest. +# Run manually when changing deadline promotion partitioning. +add_executable(deadline_promotion_bench deadline_promotion_bench.cpp + ../src/runtime/admission_queue.cpp) +target_link_libraries(deadline_promotion_bench PRIVATE tent_common) +target_include_directories(deadline_promotion_bench + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + +add_executable(promotion_policy_test promotion_policy_test.cpp) +target_link_libraries(promotion_policy_test PRIVATE tent_common gtest + gtest_main) +target_include_directories(promotion_policy_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME promotion_policy_test COMMAND promotion_policy_test) + +add_executable(rdma_cancel_test rdma_cancel_test.cpp ../src/runtime/slab.cpp) +target_link_libraries(rdma_cancel_test PRIVATE tent_common gtest gtest_main) +target_include_directories(rdma_cancel_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME rdma_cancel_test COMMAND rdma_cancel_test) + +add_executable(thread_local_storage_test thread_local_storage_test.cpp) +target_link_libraries(thread_local_storage_test PRIVATE gtest gtest_main) +target_include_directories(thread_local_storage_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME thread_local_storage_test COMMAND thread_local_storage_test) + +add_executable(tent_rw_spinlock_test rw_spinlock_test.cpp) +target_link_libraries(tent_rw_spinlock_test PRIVATE gtest gtest_main) +target_include_directories(tent_rw_spinlock_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_rw_spinlock_test COMMAND tent_rw_spinlock_test) + add_executable(tent_ip_utils_test ip_utils_test.cpp) target_link_libraries(tent_ip_utils_test PRIVATE tent_common gtest gtest_main) target_include_directories(tent_ip_utils_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_ip_utils_test COMMAND tent_ip_utils_test) +add_executable(tent_qp_pool_layout_test qp_pool_layout_test.cpp) +target_link_libraries(tent_qp_pool_layout_test PRIVATE tent_common gtest + gtest_main) +target_include_directories(tent_qp_pool_layout_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_qp_pool_layout_test COMMAND tent_qp_pool_layout_test) + add_executable(tent_coalesce_regions_test coalesce_regions_test.cpp) target_link_libraries(tent_coalesce_regions_test PRIVATE tent_common gtest gtest_main) @@ -40,6 +122,16 @@ target_include_directories(tent_coalesce_regions_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_coalesce_regions_test COMMAND tent_coalesce_regions_test) +add_executable(segment_manager_test segment_manager_test.cpp) +target_link_libraries(segment_manager_test PRIVATE gtest gtest_main + tent_link_group) +if(TARGET asio_shared) + target_link_libraries(segment_manager_test PRIVATE asio_shared) +endif() +target_include_directories(segment_manager_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME segment_manager_test COMMAND segment_manager_test) + add_executable(request_merge_test request_merge_test.cpp) target_link_libraries(request_merge_test PRIVATE gtest gtest_main tent_link_group) @@ -67,6 +159,13 @@ target_include_directories(tent_tcp_transport_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_tcp_transport_test COMMAND tent_tcp_transport_test) +add_executable(tent_shm_transport_test shm_transport_test.cpp) +target_link_libraries(tent_shm_transport_test PRIVATE gtest gtest_main + tent_link_group) +target_include_directories(tent_shm_transport_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_shm_transport_test COMMAND tent_shm_transport_test) + add_executable(tent_failover_test failover_test.cpp) target_link_libraries(tent_failover_test PRIVATE gtest gtest_main tent_link_group) @@ -75,15 +174,31 @@ target_include_directories(tent_failover_test add_test(NAME tent_failover_test COMMAND tent_failover_test) add_executable(tent_endpoint_lifecycle_test endpoint_lifecycle_test.cpp) -target_link_libraries(tent_endpoint_lifecycle_test PRIVATE gtest gtest_main) +target_link_libraries(tent_endpoint_lifecycle_test PRIVATE gtest gtest_main + tent_link_group) +target_include_directories(tent_endpoint_lifecycle_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_endpoint_lifecycle_test COMMAND tent_endpoint_lifecycle_test) +add_executable(tent_endpoint_store_test endpoint_store_test.cpp) +target_link_libraries(tent_endpoint_store_test PRIVATE gtest gtest_main + tent_link_group) +target_include_directories(tent_endpoint_store_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_endpoint_store_test COMMAND tent_endpoint_store_test) + +add_executable(tent_rdma_transport_test rdma_transport_test.cpp) +target_link_libraries(tent_rdma_transport_test PRIVATE gtest gtest_main + tent_link_group ibverbs) +target_include_directories(tent_rdma_transport_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_rdma_transport_test COMMAND tent_rdma_transport_test) + if(USE_HIP) find_package(HIP REQUIRED) add_executable(tent_rocm_platform_test rocm_platform_test.cpp) - target_link_libraries(tent_rocm_platform_test PRIVATE gtest gtest_main - tent_link_group - hip::host) + target_link_libraries(tent_rocm_platform_test + PRIVATE gtest gtest_main tent_link_group hip::host) target_include_directories(tent_rocm_platform_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_rocm_platform_test COMMAND tent_rocm_platform_test) @@ -92,9 +207,8 @@ endif() if(USE_SUNRISE) add_executable(tent_sunrise_link_transport_test sunrise_link_transport_test.cpp) - target_link_libraries(tent_sunrise_link_transport_test PRIVATE gtest - gtest_main - tent_link_group) + target_link_libraries(tent_sunrise_link_transport_test + PRIVATE gtest gtest_main tent_link_group) target_include_directories( tent_sunrise_link_transport_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include @@ -104,31 +218,51 @@ if(USE_SUNRISE) endif() add_executable(tent_fault_proxy_test fault_proxy_test.cpp) target_link_libraries(tent_fault_proxy_test PRIVATE gtest gtest_main - tent_link_group) + tent_link_group) target_include_directories(tent_fault_proxy_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_fault_proxy_test COMMAND tent_fault_proxy_test) add_executable(tent_rail_monitor_test rail_monitor_test.cpp) target_link_libraries(tent_rail_monitor_test PRIVATE gtest gtest_main - tent_link_group) + tent_link_group) target_include_directories(tent_rail_monitor_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_rail_monitor_test COMMAND tent_rail_monitor_test) -# Transport Selector Unit Test +# QoS Contract Schema Unit Test +add_executable(tent_qos_contract_test qos_contract_test.cpp + ../src/runtime/qos_contract.cpp) +target_link_libraries(tent_qos_contract_test PRIVATE gtest gtest_main + tent_common) +target_include_directories(tent_qos_contract_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_qos_contract_test COMMAND tent_qos_contract_test) + add_executable(tent_transport_selector_test transport_selector_test.cpp) target_link_libraries(tent_transport_selector_test PRIVATE gtest gtest_main - tent_link_group) + tent_link_group) target_include_directories(tent_transport_selector_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_transport_selector_test COMMAND tent_transport_selector_test) +if(USE_UB) + add_executable( + tent_ub_teardown_test + ub_teardown_test.cpp ../src/transport/ub/buffers.cpp + ../src/transport/ub/context.cpp ../src/transport/ub/jfc.cpp) + target_link_libraries(tent_ub_teardown_test PRIVATE tent_common gtest + gtest_main) + target_include_directories(tent_ub_teardown_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + add_test(NAME tent_ub_teardown_test COMMAND tent_ub_teardown_test) +endif() + # End-to-end failover test: drives real TransferEngineImpl with # FaultProxyTransport-wrapped fakes to exercise resubmitTransferTask. add_executable(tent_engine_failover_e2e_test engine_failover_e2e_test.cpp) -target_link_libraries(tent_engine_failover_e2e_test - PRIVATE gtest gtest_main tent_link_group) +target_link_libraries(tent_engine_failover_e2e_test PRIVATE gtest gtest_main + tent_link_group) if(TARGET asio_shared) target_link_libraries(tent_engine_failover_e2e_test PRIVATE asio_shared) endif() @@ -140,8 +274,8 @@ add_test(NAME tent_engine_failover_e2e_test # Per-request transport_hint: validates submitTransfer parameter, routing, # disabled-transport rejection, out-of-range rejection, mixed-hint batches. add_executable(tent_transport_hint_test transport_hint_test.cpp) -target_link_libraries(tent_transport_hint_test - PRIVATE gtest gtest_main tent_link_group) +target_link_libraries(tent_transport_hint_test PRIVATE gtest gtest_main + tent_link_group) if(TARGET asio_shared) target_link_libraries(tent_transport_hint_test PRIVATE asio_shared) endif() @@ -149,15 +283,81 @@ target_include_directories(tent_transport_hint_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) add_test(NAME tent_transport_hint_test COMMAND tent_transport_hint_test) +add_executable(tent_intent_type_test intent_type_test.cpp) +target_link_libraries(tent_intent_type_test PRIVATE gtest gtest_main + tent_common) +target_include_directories(tent_intent_type_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_intent_type_test COMMAND tent_intent_type_test) + # ProgressWorker skeleton test: covers default-off behavior, event-driven # progress without poll-failover, and freeBatch races (issue #2116). add_executable(tent_progress_worker_test progress_worker_test.cpp) -target_link_libraries(tent_progress_worker_test - PRIVATE gtest gtest_main tent_link_group) +target_link_libraries(tent_progress_worker_test PRIVATE gtest gtest_main + tent_link_group) if(TARGET asio_shared) target_link_libraries(tent_progress_worker_test PRIVATE asio_shared) endif() target_include_directories(tent_progress_worker_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) -add_test(NAME tent_progress_worker_test - COMMAND tent_progress_worker_test) +add_test(NAME tent_progress_worker_test COMMAND tent_progress_worker_test) + +add_executable(tent_runtime_queue_dispatch_test runtime_queue_dispatch_test.cpp) +target_link_libraries(tent_runtime_queue_dispatch_test PRIVATE gtest gtest_main + tent_link_group) +if(TARGET asio_shared) + target_link_libraries(tent_runtime_queue_dispatch_test PRIVATE asio_shared) +endif() +target_include_directories(tent_runtime_queue_dispatch_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME tent_runtime_queue_dispatch_test + COMMAND tent_runtime_queue_dispatch_test) + +# Causal chain stage decomposition: validates that dispatch_time and post_time +# timestamps are populated on both queue and direct-commit paths. +add_executable(causal_chain_test causal_chain_test.cpp) +target_link_libraries(causal_chain_test PRIVATE gtest gtest_main + tent_link_group) +if(TARGET asio_shared) + target_link_libraries(causal_chain_test PRIVATE asio_shared) +endif() +target_include_directories(causal_chain_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) +add_test(NAME causal_chain_test COMMAND causal_chain_test) + +# TPU PJRT shim test: exercises the dlopen'd adapter ABI (device-pointer +# classification, D2H/H2D copy, device topology) against an in-process mock +# adapter, so it runs on any Linux host without TPU hardware or a PJRT runtime. +if(USE_TPU) + # Mock adapter satisfying the tpu_pjrt_abi.h C ABI. Built as a shared object + # that the shim (and the test) load via dlopen. + add_library(mock_tpu_pjrt SHARED tpu/mock_tpu_pjrt_adapter.cpp) + target_include_directories(mock_tpu_pjrt + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + + add_executable(tent_tpu_pjrt_shim_test tpu/tpu_pjrt_shim_test.cpp) + add_dependencies(tent_tpu_pjrt_shim_test mock_tpu_pjrt) + target_link_libraries( + tent_tpu_pjrt_shim_test PRIVATE gtest gtest_main tent_link_group + ${CMAKE_DL_LIBS}) + target_include_directories(tent_tpu_pjrt_shim_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + target_compile_definitions( + tent_tpu_pjrt_shim_test + PRIVATE MOCK_TPU_PJRT_LIB="$") + add_test(NAME tent_tpu_pjrt_shim_test COMMAND tent_tpu_pjrt_shim_test) + + # Drives TpuTransport's staging hop (interior-offset chunks, and the + # exactly-one-device-side guard) against the same mock adapter. + add_executable(tent_tpu_transport_test tpu/tpu_transport_test.cpp) + add_dependencies(tent_tpu_transport_test mock_tpu_pjrt) + target_link_libraries( + tent_tpu_transport_test PRIVATE gtest gtest_main tent_link_group + ${CMAKE_DL_LIBS}) + target_include_directories(tent_tpu_transport_test + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include) + target_compile_definitions( + tent_tpu_transport_test + PRIVATE MOCK_TPU_PJRT_LIB="$") + add_test(NAME tent_tpu_transport_test COMMAND tent_tpu_transport_test) +endif() diff --git a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp index 1150c383e6..27f9681a2a 100644 --- a/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp +++ b/mooncake-transfer-engine/tent/tests/admission_queue_test.cpp @@ -14,6 +14,7 @@ #include "tent/runtime/admission_queue.h" +#include #include #include @@ -279,6 +280,40 @@ TEST(AdmissionQueueTest, RequiresDispatchBeforeTerminalCompletion) { EXPECT_EQ(status.code(), Status::Code::kInvalidEntry); } +TEST(AdmissionQueueTest, CancelsQueuedOwnerAndReleasesAccounting) { + LocalTransferAdmissionQueue queue({2, 128, 0, 0}); + std::vector admitted_ids; + ASSERT_TRUE( + queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids) + .ok()); + ASSERT_EQ(admitted_ids.size(), 1u); + + EXPECT_TRUE(queue.cancel(admitted_ids[0]).ok()); + EXPECT_TRUE(queue.cancel(admitted_ids[0]).ok()); + EXPECT_EQ(queue.outstandingOwners(), 0u); + EXPECT_EQ(queue.outstandingBytes(), 0u); + EXPECT_TRUE(queue.pickForDispatch(1, 16).empty()); + + TransferStatusEnum status = PENDING; + ASSERT_TRUE(queue.getPublicStatus(1, 0, status).ok()); + EXPECT_EQ(status, CANCELED); + EXPECT_TRUE(queue.retireBatch(1).ok()); +} + +TEST(AdmissionQueueTest, RejectsQueueCancelAfterDispatchStarts) { + LocalTransferAdmissionQueue queue({2, 128, 0, 0}); + std::vector admitted_ids; + ASSERT_TRUE( + queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids) + .ok()); + auto picked = queue.pickForDispatch(1, 16); + ASSERT_EQ(picked.size(), 1u); + + EXPECT_TRUE(queue.cancel(picked[0]).IsInvalidEntry()); + EXPECT_EQ(queue.outstandingOwners(), 1u); + EXPECT_TRUE(queue.complete(picked[0], COMPLETED).ok()); +} + TEST(AdmissionQueueTest, RetainsTerminalStatusUntilBatchRetire) { LocalTransferAdmissionQueue queue({2, 128, 0, 0}); std::vector admitted_ids; @@ -315,6 +350,25 @@ TEST(AdmissionQueueTest, RetainsTerminalStatusUntilBatchRetire) { EXPECT_EQ(status.code(), Status::Code::kInvalidEntry); } +TEST(AdmissionQueueTest, RetainsSpecificTerminalStatus) { + LocalTransferAdmissionQueue queue({1, 128, 0, 0}); + std::vector admitted_ids; + + auto status = + queue.tryAdmit(makeSubmit(1, 1, {makeOwner(0, 16)}), admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(1, 16); + ASSERT_EQ(picked.size(), 1u); + status = queue.complete(picked[0], TransferStatusEnum::TIMEOUT); + ASSERT_EQ(status.code(), Status::Code::kOk); + + TransferStatusEnum public_status = TransferStatusEnum::PENDING; + status = queue.getPublicStatus(1, 0, public_status); + ASSERT_EQ(status.code(), Status::Code::kOk); + EXPECT_EQ(public_status, TransferStatusEnum::TIMEOUT); +} + TEST(AdmissionQueueTest, RejectsRetireWithNonTerminalOwners) { LocalTransferAdmissionQueue queue({2, 128, 0, 0}); std::vector admitted_ids; @@ -369,6 +423,500 @@ TEST(AdmissionQueueTest, AllowsBatchTokenReuseAfterRetire) { EXPECT_EQ(resolved_owner, 2u); } +// --- RFC #2519 step 2: opt-in deadline-aware (EDF) dispatch --------------- + +QueueOwnerInput makeOwnerWithDeadline(size_t public_task_id, size_t length, + uint64_t deadline_ns) { + QueueOwnerInput owner = makeOwner(public_task_id, length); + owner.request.deadline_ns = deadline_ns; + return owner; +} + +QueueOwnerInput makeDegradationEligibleOwnerWithDeadline(size_t public_task_id, + size_t length, + uint64_t deadline_ns) { + QueueOwnerInput owner = + makeOwnerWithDeadline(public_task_id, length, deadline_ns); + owner.degradation_eligible = true; + return owner; +} + +TEST(AdmissionQueueTest, DeadlineAwareDispatchesEarliestDeadlineFirst) { + QueueLimits limits{4, 4096, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + + // Admitted in FIFO order 1,2,3 but with deadlines 300,100,200. + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 300), + makeOwnerWithDeadline(1, 16, 100), + makeOwnerWithDeadline(2, 16, 200)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + ASSERT_EQ(admitted_ids.size(), 3u); // owner ids 1,2,3 + + auto picked = queue.pickForDispatch(3, 4096); + // EDF order: owner 2 (dl 100) < owner 3 (dl 200) < owner 1 (dl 300). + const std::vector expected{2, 3, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeadlineAwareKeepsUndeadlinedOwnersLast) { + QueueLimits limits{4, 4096, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector admitted_ids; + + // owner 1: no deadline (0); owner 2: deadline 100; owner 3: no deadline. + auto status = queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 0), + makeOwnerWithDeadline(1, 16, 100), + makeOwnerWithDeadline(2, 16, 0)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(3, 4096); + // Deadlined owner 2 first; undeadlined 1,3 keep FIFO order behind it. + const std::vector expected{2, 1, 3}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, DeadlineUnawareKeepsStrictFifo) { + // Default (deadline_aware == false): FIFO regardless of deadlines. + LocalTransferAdmissionQueue queue({4, 4096, 0, 0}); + std::vector admitted_ids; + + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 300), + makeOwnerWithDeadline(1, 16, 100), + makeOwnerWithDeadline(2, 16, 200)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(3, 4096); + const std::vector expected{1, 2, + 3}; // FIFO, deadlines ignored + EXPECT_EQ(picked, expected); +} + +// fifo_ is kept EDF-ordered at admission time, so owners admitted across +// *separate* tryAdmit calls (out of deadline order) must still dispatch EDF — +// this exercises the ordered-insert path, not just a single sorted batch. +TEST(AdmissionQueueTest, DeadlineAwareOrdersAcrossSeparateAdmits) { + QueueLimits limits{8, 4096, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + std::vector ids; + + // Admit one at a time, deadlines arriving out of order: 300, 100, 200, 0. + ASSERT_EQ( + queue + .tryAdmit(makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 300)}), + ids) + .code(), + Status::Code::kOk); // owner 1 + ASSERT_EQ( + queue + .tryAdmit(makeSubmit(2, 1, {makeOwnerWithDeadline(0, 16, 100)}), + ids) + .code(), + Status::Code::kOk); // owner 2 + ASSERT_EQ( + queue + .tryAdmit(makeSubmit(3, 1, {makeOwnerWithDeadline(0, 16, 200)}), + ids) + .code(), + Status::Code::kOk); // owner 3 + ASSERT_EQ( + queue.tryAdmit(makeSubmit(4, 1, {makeOwnerWithDeadline(0, 16, 0)}), ids) + .code(), + Status::Code::kOk); // owner 4 (no deadline → last) + + auto picked = queue.pickForDispatch(8, 4096); + // EDF: 100(owner2) < 200(owner3) < 300(owner1) < no-deadline(owner4). + const std::vector expected{2, 3, 1, 4}; + EXPECT_EQ(picked, expected); +} + +// --- RFC #2519 step 3: deadline-infeasible drop + degradation hook -------- + +// Helper: build a queue with deadline_aware + a θ_local, a fixed bandwidth, +// and a fixed "now" clock so MLU is deterministic. +QueueLimits step3Limits(double theta_local) { + QueueLimits limits{4, 1 << 20, 0, 0}; + limits.deadline_aware = true; + limits.mlu_local_threshold = theta_local; + return limits; +} + +TEST(AdmissionQueueTest, Step3DropsInfeasibleAndKeepsFeasible) { + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + // Fixed now = 1e9 ns; bandwidth = 1e9 B/s (so 16 B takes 16 ns). + int hook_calls = 0; + DegradationHooks hooks; + hooks.on_local_decode_suggested = [&](const Request&) { ++hook_calls; }; + queue.setDegradationPolicy([] { return 1e9; }, hooks, + [] { return uint64_t{1'000'000'000}; }); + + std::vector admitted_ids; + // owner 1: window = 10 ns → 16 B / 1e9 = 16 ns → MLU 1.6 ≥ 1.5 → DROP. + // owner 2: window = 1e6 ns → MLU ~1.6e-5 → feasible → dispatch. + auto status = queue.tryAdmit( + makeSubmit( + 1, 2, + {makeDegradationEligibleOwnerWithDeadline(0, 16, 1'000'000'010), + makeDegradationEligibleOwnerWithDeadline(1, 16, 2'000'000'000)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + ASSERT_EQ(admitted_ids.size(), 2u); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + + const std::vector exp_pick{2}; + const std::vector exp_drop{1}; + EXPECT_EQ(picked, exp_pick); + EXPECT_EQ(dropped, exp_drop); + EXPECT_EQ(hook_calls, 1); + // Dropped owner is charged out of the outstanding accounting. + EXPECT_EQ(queue.outstandingOwners(), 1u); + EXPECT_EQ(queue.outstandingBytes(), 16u); +} + +TEST(AdmissionQueueTest, Step3DropsAlreadyExpiredDeadline) { + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, + [] { return uint64_t{2'000'000'000}; }); + + std::vector admitted_ids; + // deadline 1e9 < now 2e9 → already past → dropped. + auto status = queue.tryAdmit( + makeSubmit( + 1, 1, + {makeDegradationEligibleOwnerWithDeadline(0, 16, 1'000'000'000)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + EXPECT_TRUE(picked.empty()); + ASSERT_EQ(dropped.size(), 1u); + EXPECT_EQ(dropped[0], 1u); +} + +TEST(AdmissionQueueTest, Step3DisabledWhenThresholdZero) { + // θ_local = 0 (default off): even a hopeless deadline is dispatched, and + // the dropped vector stays empty — behavior is pure step-2 EDF. + LocalTransferAdmissionQueue queue(step3Limits(0.0)); + queue.setDegradationPolicy([] { return 1e9; }, DegradationHooks{}, + [] { return uint64_t{1'000'000'000}; }); + + std::vector admitted_ids; + auto status = queue.tryAdmit( + makeSubmit( + 1, 1, + {makeDegradationEligibleOwnerWithDeadline(0, 16, 1'000'000'001)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + ASSERT_EQ(picked.size(), 1u); + EXPECT_EQ(picked[0], 1u); + EXPECT_TRUE(dropped.empty()); +} + +TEST(AdmissionQueueTest, Step3NoDropWithoutBandwidthProvider) { + // Threshold set but no bandwidth provider → cannot predict → never drops. + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + std::vector admitted_ids; + auto status = queue.tryAdmit( + makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 1'000'000'001)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + ASSERT_EQ(picked.size(), 1u); + EXPECT_TRUE(dropped.empty()); +} + +TEST(AdmissionQueueTest, Step3DynamicBandwidthProvider) { + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + std::atomic live_bw{1e9}; + int hook_calls = 0; + DegradationHooks hooks; + hooks.on_local_decode_suggested = [&](const Request&) { ++hook_calls; }; + queue.setDegradationPolicy([&] { return live_bw.load(); }, hooks, + [] { return uint64_t{1'000'000'000}; }); + + std::vector admitted_ids; + // At 1e9 B/s: time=16ns, window=10ns, MLU=1.6 >= 1.5 -> DROP. + auto status = queue.tryAdmit( + makeSubmit( + 1, 1, + {makeDegradationEligibleOwnerWithDeadline(0, 16, 1'000'000'010)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + EXPECT_TRUE(picked.empty()); + ASSERT_EQ(dropped.size(), 1u); + EXPECT_EQ(hook_calls, 1); + + // Increase bandwidth 10x -> same profile becomes feasible. + live_bw.store(1e10); + status = queue.tryAdmit( + makeSubmit( + 2, 1, + {makeDegradationEligibleOwnerWithDeadline(0, 16, 1'000'000'010)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + dropped.clear(); + picked = queue.pickForDispatch(4, 1 << 20, &dropped); + // At 1e10 B/s: time=1.6ns, window=10ns, MLU=0.16 < 1.5 -> OK. + ASSERT_EQ(picked.size(), 1u); + EXPECT_TRUE(dropped.empty()); + EXPECT_EQ(hook_calls, 1); +} + +TEST(AdmissionQueueTest, Step3SkipsNonRdmaOwner) { + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + int hook_calls = 0; + DegradationHooks hooks; + hooks.on_local_decode_suggested = [&](const Request&) { ++hook_calls; }; + queue.setDegradationPolicy([] { return 1e9; }, hooks, + [] { return uint64_t{1'000'000'000}; }); + + auto owner = makeOwnerWithDeadline(0, 16, 1'000'000'010); + owner.degradation_eligible = false; + std::vector admitted_ids; + auto status = queue.tryAdmit(makeSubmit(1, 1, {owner}), admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + ASSERT_EQ(picked.size(), 1u); + EXPECT_TRUE(dropped.empty()); + EXPECT_EQ(hook_calls, 0); +} + +TEST(AdmissionQueueTest, Step3RequiresExplicitDegradationEligibility) { + LocalTransferAdmissionQueue queue(step3Limits(1.5)); + int hook_calls = 0; + DegradationHooks hooks; + hooks.on_local_decode_suggested = [&](const Request&) { ++hook_calls; }; + queue.setDegradationPolicy([] { return 1e9; }, hooks, + [] { return uint64_t{1'000'000'000}; }); + + std::vector admitted_ids; + auto status = queue.tryAdmit( + makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 1'000'000'010)}), + admitted_ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + ASSERT_EQ(picked.size(), 1u); + EXPECT_TRUE(dropped.empty()); + EXPECT_EQ(hook_calls, 0); +} + +// --- Deadline proximity promotion (step 4) -------------------------------- + +QueueLimits promotionLimits(uint64_t slack_ns) { + QueueLimits limits{8, 1 << 20, 0, 0}; + limits.deadline_aware = true; + limits.promotion_slack_ns = slack_ns; + return limits; +} + +TEST(AdmissionQueueTest, PromotionDisabledKeepsEdfOrder) { + QueueLimits limits{4, 4096, 0, 0}; + limits.deadline_aware = true; + LocalTransferAdmissionQueue queue(limits); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 2000), + makeOwnerWithDeadline(1, 16, 1500), + makeOwnerWithDeadline(2, 16, 1800)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 4096); + const std::vector expected{2, 3, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionMovesUrgentOwnersToFront) { + LocalTransferAdmissionQueue queue(promotionLimits(500)); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 2000), + makeOwnerWithDeadline(1, 16, 1400), + makeOwnerWithDeadline(2, 16, 1300)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 1 << 20); + const std::vector expected{3, 2, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionReordersAcrossSeparateAdmits) { + LocalTransferAdmissionQueue queue(promotionLimits(2000)); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{5000}; }); + + std::vector ids; + auto s1 = queue.tryAdmit( + makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 10000)}), ids); + ASSERT_EQ(s1.code(), Status::Code::kOk); + auto s2 = queue.tryAdmit( + makeSubmit(2, 1, {makeOwnerWithDeadline(0, 16, 6500)}), ids); + ASSERT_EQ(s2.code(), Status::Code::kOk); + auto s3 = queue.tryAdmit( + makeSubmit(3, 1, {makeOwnerWithDeadline(0, 16, 6000)}), ids); + ASSERT_EQ(s3.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 1 << 20); + const std::vector expected{3, 2, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionSkipsNoDeadlineOwners) { + LocalTransferAdmissionQueue queue(promotionLimits(5000)); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 2, + {makeOwnerWithDeadline(0, 16, 0), + makeOwnerWithDeadline(1, 16, 2000)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 1 << 20); + const std::vector expected{2, 1}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionPreservesEdfWithinPromotedGroup) { + LocalTransferAdmissionQueue queue(promotionLimits(2000)); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 2500), + makeOwnerWithDeadline(1, 16, 2200), + makeOwnerWithDeadline(2, 16, 2800)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 1 << 20); + const std::vector expected{2, 1, 3}; + EXPECT_EQ(picked, expected); +} + +TEST(AdmissionQueueTest, PromotionCoexistsWithStep3Drop) { + QueueLimits limits = promotionLimits(500); + limits.mlu_local_threshold = 1.5; + LocalTransferAdmissionQueue queue(limits); + int hook_calls = 0; + DegradationHooks hooks; + hooks.on_local_decode_suggested = [&](const Request&) { ++hook_calls; }; + queue.setDegradationPolicy([] { return 1e9; }, hooks, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = queue.tryAdmit( + makeSubmit(1, 3, + {makeDegradationEligibleOwnerWithDeadline(0, 16, 1010), + makeDegradationEligibleOwnerWithDeadline(1, 16, 1400), + makeDegradationEligibleOwnerWithDeadline(2, 16, 5000)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + std::vector dropped; + auto picked = queue.pickForDispatch(4, 1 << 20, &dropped); + + const std::vector exp_pick{2, 3}; + const std::vector exp_drop{1}; + EXPECT_EQ(picked, exp_pick); + EXPECT_EQ(dropped, exp_drop); + EXPECT_EQ(hook_calls, 1); +} + +TEST(AdmissionQueueTest, PromotionWithAdvancingTime) { + QueueLimits limits = promotionLimits(500); + LocalTransferAdmissionQueue queue(limits); + + uint64_t fake_now = 1000; + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [&] { return fake_now; }); + + std::vector ids; + auto s1 = queue.tryAdmit( + makeSubmit(1, 1, {makeOwnerWithDeadline(0, 16, 1800)}), ids); + ASSERT_EQ(s1.code(), Status::Code::kOk); + auto s2 = queue.tryAdmit( + makeSubmit(2, 1, {makeOwnerWithDeadline(0, 16, 1400)}), ids); + ASSERT_EQ(s2.code(), Status::Code::kOk); + auto s3 = queue.tryAdmit( + makeSubmit(3, 1, {makeOwnerWithDeadline(0, 16, 3000)}), ids); + ASSERT_EQ(s3.code(), Status::Code::kOk); + + auto picked1 = queue.pickForDispatch(1, 1 << 20); + ASSERT_EQ(picked1.size(), 1u); + EXPECT_EQ(picked1[0], 2u); + + auto cstatus = queue.complete(2, TransferStatusEnum::COMPLETED); + ASSERT_EQ(cstatus.code(), Status::Code::kOk); + + fake_now = 1500; + auto picked2 = queue.pickForDispatch(2, 1 << 20); + const std::vector expected2{1, 3}; + EXPECT_EQ(picked2, expected2); +} + +TEST(AdmissionQueueTest, PromotionDisabledWithoutDeadlineAware) { + QueueLimits limits{4, 4096, 0, 0}; + limits.deadline_aware = false; + limits.promotion_slack_ns = 5000; + LocalTransferAdmissionQueue queue(limits); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return uint64_t{1000}; }); + + std::vector ids; + auto status = + queue.tryAdmit(makeSubmit(1, 3, + {makeOwnerWithDeadline(0, 16, 1200), + makeOwnerWithDeadline(1, 16, 5000), + makeOwnerWithDeadline(2, 16, 1100)}), + ids); + ASSERT_EQ(status.code(), Status::Code::kOk); + + auto picked = queue.pickForDispatch(4, 4096); + const std::vector expected{1, 2, 3}; + EXPECT_EQ(picked, expected); +} + } // namespace } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/bw_arbitration_test.cpp b/mooncake-transfer-engine/tent/tests/bw_arbitration_test.cpp new file mode 100644 index 0000000000..3664362ba0 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/bw_arbitration_test.cpp @@ -0,0 +1,111 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Unit tests for the deadline-aware bandwidth arbitration ordering (#2792). + +#include "tent/transport/rdma/bw_arbitration.h" + +#include + +#include +#include + +namespace mooncake { +namespace tent { +namespace { + +constexpr uint64_t kNow = 1'000'000'000; // 1s in ns +constexpr double kBw = 1e9; // 1 GB/s -> 16 B takes 16 ns + +// A flow whose window is `window_ns` from now, transferring `len` bytes. +ArbFlow flow(uint64_t window_ns, size_t len) { + return ArbFlow{kNow + window_ns, len}; +} + +TEST(BwArbitrationTest, TighterDeadlineSortsFirst) { + std::vector flows = { + flow(1'000'000, 4096), // idx0: loose (1ms window) + flow(10'000, 4096), // idx1: tight (10us window) -> most urgent + flow(100'000, 4096), // idx2: medium + }; + auto order = OrderByUrgency(flows, kNow, kBw); + ASSERT_EQ(order.size(), 3u); + EXPECT_EQ(order[0], 1u); // tightest first + EXPECT_EQ(order[1], 2u); + EXPECT_EQ(order[2], 0u); // loosest last +} + +TEST(BwArbitrationTest, NoDeadlineSortsLast) { + std::vector flows = { + ArbFlow{0, 4096}, // idx0: no deadline -> least urgent + flow(50'000, 4096), // idx1: has deadline -> first + }; + auto order = OrderByUrgency(flows, kNow, kBw); + EXPECT_EQ(order[0], 1u); + EXPECT_EQ(order[1], 0u); +} + +TEST(BwArbitrationTest, PastDeadlineSortsFirst) { + std::vector flows = { + flow(50'000, 4096), // idx0: still feasible + ArbFlow{kNow - 1, 4096}, // idx1: already past -> most urgent + }; + auto order = OrderByUrgency(flows, kNow, kBw); + EXPECT_EQ(order[0], 1u); + EXPECT_EQ(order[1], 0u); +} + +TEST(BwArbitrationTest, TiesKeepFifoOrder) { + // Identical deadlines/lengths -> stable sort preserves original order. + std::vector flows = { + flow(50'000, 4096), // idx0 + flow(50'000, 4096), // idx1 + flow(50'000, 4096), // idx2 + }; + auto order = OrderByUrgency(flows, kNow, kBw); + EXPECT_EQ(order, (std::vector{0, 1, 2})); +} + +TEST(BwArbitrationTest, AllNoDeadlineKeepsFifoOrder) { + // No flow has a deadline -> byte-identical to today's order (no reorder). + std::vector flows = {ArbFlow{0, 1}, ArbFlow{0, 2}, ArbFlow{0, 3}}; + auto order = OrderByUrgency(flows, kNow, kBw); + EXPECT_EQ(order, (std::vector{0, 1, 2})); +} + +TEST(BwArbitrationTest, ZeroBandwidthDisablesReorder) { + // bw<=0 -> prediction disabled -> original order preserved. + std::vector flows = { + flow(1'000'000, 4096), + flow(10'000, 4096), + }; + auto order = OrderByUrgency(flows, kNow, /*bw_bps=*/0.0); + EXPECT_EQ(order, (std::vector{0, 1})); +} + +TEST(BwArbitrationTest, LongerTransferIsMoreUrgentAtSameDeadline) { + // Same window, but a bigger transfer has higher predicted MLU (needs more + // of the shared bandwidth to finish in time). + std::vector flows = { + flow(50'000, 4096), // idx0: small + flow(50'000, 65536), // idx1: large -> more urgent + }; + auto order = OrderByUrgency(flows, kNow, kBw); + EXPECT_EQ(order[0], 1u); + EXPECT_EQ(order[1], 0u); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/causal_chain_test.cpp b/mooncake-transfer-engine/tent/tests/causal_chain_test.cpp new file mode 100644 index 0000000000..057a5fe77c --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/causal_chain_test.cpp @@ -0,0 +1,291 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/common/types.h" +#if TENT_METRICS_ENABLED +#include "tent/metrics/tent_metrics.h" +#endif +#include "tent/runtime/transfer_engine_impl.h" +#include "tent/runtime/transport.h" + +namespace mooncake { +namespace tent { +namespace { + +class FakeSubBatch : public Transport::SubBatch { + public: + size_t size() const override { return task_count; } + + size_t task_count = 0; + std::vector requests; + std::vector statuses; +}; + +class FakeTransport : public Transport { + public: + explicit FakeTransport(TransportType self_type) : self_type_(self_type) { + caps.dram_to_dram = true; + } + + std::atomic submit_calls{0}; + + Status install(std::string&, std::shared_ptr, + std::shared_ptr, + std::shared_ptr) override { + return Status::OK(); + } + + Status allocateSubBatch(SubBatchRef& batch, size_t) override { + batch = new FakeSubBatch(); + return Status::OK(); + } + + Status freeSubBatch(SubBatchRef& batch) override { + delete static_cast(batch); + batch = nullptr; + return Status::OK(); + } + + Status submitTransferTasks(SubBatchRef batch, + const std::vector& requests) override { + ++submit_calls; + auto* fake = static_cast(batch); + for (const auto& request : requests) { + fake->requests.push_back(request); + fake->statuses.push_back( + {TransferStatusEnum::COMPLETED, request.length}); + ++fake->task_count; + } + batch->notifyProgress(); + return Status::OK(); + } + + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override { + auto* fake = static_cast(batch); + if (task_id < 0 || task_id >= (int)fake->statuses.size()) { + return Status::InvalidArgument("bad task_id" LOC_MARK); + } + status = fake->statuses[task_id]; + return Status::OK(); + } + + Status addMemoryBuffer(BufferDesc& desc, const MemoryOptions&) override { + desc.transports.push_back(self_type_); + return Status::OK(); + } + + Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) override { + for (auto& desc : desc_list) { + auto s = addMemoryBuffer(desc, options); + if (!s.ok()) return s; + } + return Status::OK(); + } + + Status removeMemoryBuffer(BufferDesc&) override { return Status::OK(); } + + Status allocateLocalMemory(void** addr, size_t size, + MemoryOptions&) override { + *addr = std::malloc(size); + return *addr ? Status::OK() + : Status::InternalError("malloc failed" LOC_MARK); + } + + Status freeLocalMemory(void* addr, size_t) override { + std::free(addr); + return Status::OK(); + } + + bool warmupMemory(void*, size_t) override { return false; } + + const char* getName() const override { return ""; } + + private: + TransportType self_type_; +}; + +std::shared_ptr makeCausalChainConfig(size_t max_dispatch_owners, + size_t max_dispatch_bytes) { + auto cfg = std::make_shared(); + cfg->set("metadata_type", "p2p"); + cfg->set("metadata_servers", ""); + cfg->set("rpc_server_hostname", "127.0.0.1"); + cfg->set("rpc_server_port", "0"); + cfg->set("log_level", "warning"); + cfg->set("merge_requests", false); + cfg->set("enable_runtime_queue", true); + cfg->set("runtime_queue/max_outstanding_owners", 16UL); + cfg->set("runtime_queue/max_outstanding_bytes", 1UL << 20); + cfg->set("runtime_queue/max_dispatch_owners", max_dispatch_owners); + cfg->set("runtime_queue/max_dispatch_bytes", max_dispatch_bytes); + cfg->set("runtime_queue/staging_owner_reserve", 0UL); + cfg->set("runtime_queue/staging_byte_reserve", 0UL); + cfg->set("runtime_queue/progress_fallback_interval_us", 50000UL); + + cfg->set("transports/tcp/enable", false); + cfg->set("transports/shm/enable", false); + cfg->set("transports/rdma/enable", false); + cfg->set("transports/io_uring/enable", false); + cfg->set("transports/nvlink/enable", false); + cfg->set("transports/mnnvl/enable", false); + cfg->set("transports/gds/enable", false); + cfg->set("transports/ascend_direct/enable", false); + return cfg; +} + +void installFakeRdma(TransferEngineImpl& engine, + const std::shared_ptr& fake) { + std::string seg_name = engine.getSegmentName(); + ASSERT_TRUE(fake->install(seg_name, nullptr, nullptr, nullptr).ok()); + engine.swapTransportForTest(RDMA, fake); +} + +Request makeLocalWrite(uint8_t* ptr, size_t length) { + Request request; + request.opcode = Request::WRITE; + request.source = ptr; + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(ptr); + request.length = length; + request.transport_hint = RDMA; + return request; +} + +TEST(CausalChain, TimestampsPopulatedOnQueuePath) { + auto cfg = makeCausalChainConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake = std::make_shared(RDMA); + installFakeRdma(engine, fake); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xBB); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto status = + engine.submitTransfer(batch, {makeLocalWrite(buf.data(), kLen)}); + ASSERT_TRUE(status.ok()) << status.ToString(); + + TransferStatus ts{}; + for (int i = 0; i < 200; ++i) { + engine.getTransferStatus(batch, 0, ts); + if (ts.s == TransferStatusEnum::COMPLETED) break; + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + EXPECT_EQ(ts.s, TransferStatusEnum::COMPLETED); + EXPECT_GE(fake->submit_calls.load(), 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +TEST(CausalChain, MultipleTransfersAllComplete) { + auto cfg = makeCausalChainConfig(4, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake = std::make_shared(RDMA); + installFakeRdma(engine, fake); + + constexpr size_t kLen = 4096; + constexpr int kBatchSize = 4; + std::vector buf(kLen * kBatchSize, 0xCC); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(kBatchSize); + ASSERT_NE(batch, (BatchID)0); + + std::vector requests; + for (int i = 0; i < kBatchSize; ++i) { + requests.push_back(makeLocalWrite(buf.data() + i * kLen, kLen)); + } + auto status = engine.submitTransfer(batch, requests); + ASSERT_TRUE(status.ok()) << status.ToString(); + + for (int i = 0; i < kBatchSize; ++i) { + TransferStatus ts{}; + for (int j = 0; j < 200; ++j) { + engine.getTransferStatus(batch, i, ts); + if (ts.s == TransferStatusEnum::COMPLETED) break; + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + EXPECT_EQ(ts.s, TransferStatusEnum::COMPLETED) << "task " << i; + } + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +TEST(CausalChain, TimestampsPopulatedOnDirectPath) { + auto cfg = makeCausalChainConfig(1, 1UL << 20); + cfg->set("enable_runtime_queue", false); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake = std::make_shared(RDMA); + installFakeRdma(engine, fake); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xBB); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto status = + engine.submitTransfer(batch, {makeLocalWrite(buf.data(), kLen)}); + ASSERT_TRUE(status.ok()) << status.ToString(); + + TransferStatus ts{}; + for (int i = 0; i < 200; ++i) { + engine.getTransferStatus(batch, 0, ts); + if (ts.s == TransferStatusEnum::COMPLETED) break; + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + EXPECT_EQ(ts.s, TransferStatusEnum::COMPLETED); + EXPECT_GE(fake->submit_calls.load(), 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +#if TENT_METRICS_ENABLED +TEST(CausalChain, MetricsRecordStageLatencyIsCallable) { + TENT_RECORD_STAGE_LATENCY(TentMetrics::Stage::QueueWait, UNSPEC, 100.0); + TENT_RECORD_STAGE_LATENCY(TentMetrics::Stage::Dispatch, UNSPEC, 50.0); + TENT_RECORD_STAGE_LATENCY(TentMetrics::Stage::Transport, UNSPEC, 200.0); +} +#endif + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/deadline_promotion_bench.cpp b/mooncake-transfer-engine/tent/tests/deadline_promotion_bench.cpp new file mode 100644 index 0000000000..52def2ed0f --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/deadline_promotion_bench.cpp @@ -0,0 +1,105 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include +#include +#include + +#include "tent/runtime/admission_queue.h" + +namespace mooncake { +namespace tent { +namespace { + +constexpr uint64_t kNowNs = 1'000'000; + +QueueOwnerInput makeOwner(size_t task_id, bool urgent) { + QueueOwnerInput owner; + owner.owner_task_id = task_id; + owner.request.opcode = Request::READ; + owner.request.target_id = 1; + owner.request.length = 4096; + owner.request.deadline_ns = urgent ? kNowNs + 10'000 : kNowNs + 1'000'000; + return owner; +} + +double percentile(std::vector samples, double ratio) { + std::sort(samples.begin(), samples.end()); + const size_t index = + static_cast(ratio * static_cast(samples.size() - 1)); + return samples[index]; +} + +void runDepth(size_t depth) { + const size_t repeats = std::max(100, 200'000 / depth); + std::vector samples; + samples.reserve(repeats * depth); + + for (size_t repeat = 0; repeat < repeats; ++repeat) { + QueueLimits limits{depth, depth * 4096, 0, 0}; + limits.deadline_aware = true; + limits.promotion_slack_ns = 20'000; + LocalTransferAdmissionQueue queue(limits); + queue.setDegradationPolicy(nullptr, DegradationHooks{}, + [] { return kNowNs; }); + + QueueSubmit submit; + submit.batch_token = repeat + 1; + submit.batch_slots_left = depth; + submit.owners.reserve(depth); + for (size_t i = 0; i < depth; ++i) { + // Interleave urgent and comfortable requests so every dispatch + // round exercises the stable partition. + submit.owners.push_back(makeOwner(i, i % 4 == 3)); + } + std::vector admitted; + if (!queue.tryAdmit(submit, admitted).ok()) std::abort(); + + for (size_t i = 0; i < depth; ++i) { + const auto start = std::chrono::steady_clock::now(); + auto picked = queue.pickForDispatch(1, 4096); + const auto end = std::chrono::steady_clock::now(); + if (picked.size() != 1 || + !queue.complete(picked[0], COMPLETED).ok()) { + std::abort(); + } + samples.push_back( + std::chrono::duration(end - start).count()); + } + if (!queue.retireBatch(submit.batch_token).ok()) std::abort(); + } + + const double mean = + std::accumulate(samples.begin(), samples.end(), 0.0) / samples.size(); + std::cout << depth << ',' << samples.size() << ',' << std::fixed + << std::setprecision(3) << mean << ',' + << percentile(samples, 0.50) << ',' << percentile(samples, 0.95) + << '\n'; +} + +} // namespace +} // namespace tent +} // namespace mooncake + +int main() { + std::cout << "queue_depth,samples,mean_us,p50_us,p95_us\n"; + for (size_t depth : {32, 64, 128, 256, 512}) { + mooncake::tent::runDepth(depth); + } + return 0; +} diff --git a/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp b/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp index 06c8adea8a..3243203811 100644 --- a/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp +++ b/mooncake-transfer-engine/tent/tests/endpoint_lifecycle_test.cpp @@ -16,130 +16,186 @@ #include +#include "tent/common/utils/string_builder.h" +#include "tent/transport/rdma/endpoint.h" +#include "tent/transport/rdma/slice.h" + namespace mooncake { namespace tent { + +class EndpointTestAccess { + public: + // Puts a context-less endpoint into the state a completed bootstrap + // leaves behind, so accept() can be driven without an RDMA device. + static void markConnected(RdmaEndPoint& endpoint, + const std::string& peer_server_name, + const std::string& peer_nic_name, + const std::vector& peer_qp_num_list) { + endpoint.peer_server_name_ = peer_server_name; + endpoint.peer_nic_name_ = peer_nic_name; + endpoint.peer_qp_num_list_ = peer_qp_num_list; + endpoint.status_.store(RdmaEndPoint::EP_READY, + std::memory_order_relaxed); + } +}; + namespace { -// --------------------------------------------------------------------------- -// Minimal stub to validate weak_ptr lifecycle without RDMA dependencies. -// The real RdmaEndPoint inherits enable_shared_from_this; we mirror that -// pattern here so the test proves the ownership model works. -// --------------------------------------------------------------------------- +TEST(EndpointLifecycleTest, DefaultConstructedEndpointOwnsNoResources) { + RdmaEndPoint endpoint; -class FakeEndPoint : public std::enable_shared_from_this { - public: - int acknowledge_calls = 0; - int reset_calls = 0; + EXPECT_EQ(endpoint.status(), RdmaEndPoint::EP_UNINIT); + EXPECT_TRUE(endpoint.qpNum().empty()); + EXPECT_EQ(endpoint.getInflightSlices(), 0); + EXPECT_EQ(endpoint.notifyQpNum(), 0); +} - void acknowledge() { ++acknowledge_calls; } - void reset() { ++reset_calls; } -}; +TEST(EndpointLifecycleTest, DefaultConstructedEndpointCanBeDestroyed) { + RdmaEndPoint endpoint; -// --------------------------------------------------------------------------- -// weak_ptr basic lifecycle -// --------------------------------------------------------------------------- + EXPECT_EQ(endpoint.deconstruct(), 0); + EXPECT_EQ(endpoint.status(), RdmaEndPoint::EP_DESTROYED); + EXPECT_TRUE(endpoint.qpNum().empty()); +} -TEST(EndpointLifecycleTest, WeakPtrLocksWhileAlive) { - auto ep = std::make_shared(); - std::weak_ptr weak = ep; +TEST(EndpointLifecycleTest, DeconstructIsIdempotent) { + RdmaEndPoint endpoint; - auto locked = weak.lock(); - ASSERT_NE(locked, nullptr); - locked->acknowledge(); - EXPECT_EQ(ep->acknowledge_calls, 1); + ASSERT_EQ(endpoint.deconstruct(), 0); + EXPECT_EQ(endpoint.deconstruct(), 0); + EXPECT_EQ(endpoint.status(), RdmaEndPoint::EP_DESTROYED); } -TEST(EndpointLifecycleTest, WeakPtrExpiresAfterRelease) { - std::weak_ptr weak; - { - auto ep = std::make_shared(); - weak = ep; - EXPECT_FALSE(weak.expired()); - } // ep destroyed here - EXPECT_TRUE(weak.expired()); - EXPECT_EQ(weak.lock(), nullptr); +TEST(EndpointLifecycleTest, DestroyedEndpointCannotBeConstructedAgain) { + RdmaEndPoint endpoint; + + ASSERT_EQ(endpoint.deconstruct(), 0); + EXPECT_NE(endpoint.construct(nullptr, nullptr, "reused"), 0); + EXPECT_EQ(endpoint.status(), RdmaEndPoint::EP_DESTROYED); } -TEST(EndpointLifecycleTest, SharedFromThisProducesValidWeakPtr) { - auto ep = std::make_shared(); - // Simulate what submitSlices() does: shared_from_this() assigned to - // weak_ptr - std::weak_ptr weak = ep->shared_from_this(); +TEST(EndpointLifecycleTest, TwoPhaseDestroyHandlesUninitializedEndpoint) { + RdmaEndPoint endpoint; - auto locked = weak.lock(); - ASSERT_NE(locked, nullptr); - EXPECT_EQ(locked.get(), ep.get()); + endpoint.beginDestroy(); + EXPECT_EQ(endpoint.status(), RdmaEndPoint::EP_DESTROYING); + EXPECT_TRUE(endpoint.finishDestroy()); + EXPECT_EQ(endpoint.status(), RdmaEndPoint::EP_DESTROYED); } -// --------------------------------------------------------------------------- -// Simulate the slice → endpoint dereference pattern used in workers.cpp -// --------------------------------------------------------------------------- +TEST(EndpointLifecycleTest, FinishDestroyRejectsNonRetiringEndpoint) { + RdmaEndPoint endpoint; -struct FakeSlice { - std::weak_ptr ep_weak_ptr; -}; + EXPECT_FALSE(endpoint.finishDestroy()); + EXPECT_EQ(endpoint.status(), RdmaEndPoint::EP_UNINIT); +} + +TEST(EndpointLifecycleTest, FinishDestroyIsIdempotent) { + RdmaEndPoint endpoint; + + endpoint.beginDestroy(); + ASSERT_TRUE(endpoint.finishDestroy()); + EXPECT_TRUE(endpoint.finishDestroy()); + EXPECT_EQ(endpoint.status(), RdmaEndPoint::EP_DESTROYED); +} + +TEST(EndpointLifecycleTest, NotificationFailsWhenEndpointIsNotConnected) { + RdmaEndPoint endpoint; + + EXPECT_FALSE(endpoint.sendNotification("name", "message")); +} + +TEST(EndpointLifecycleTest, SharedFromThisUsesRealEndpointOwnership) { + auto endpoint = std::make_shared(); + std::weak_ptr weak = endpoint->shared_from_this(); + + auto locked = weak.lock(); + ASSERT_NE(locked, nullptr); + EXPECT_EQ(locked.get(), endpoint.get()); + + locked.reset(); + endpoint.reset(); + EXPECT_TRUE(weak.expired()); +} TEST(EndpointLifecycleTest, SliceAccessWhileEndpointAlive) { - auto ep = std::make_shared(); - FakeSlice slice; - slice.ep_weak_ptr = ep; - - // Simulate workers.cpp completion path - if (auto locked = slice.ep_weak_ptr.lock()) { - locked->acknowledge(); - locked->reset(); - } - EXPECT_EQ(ep->acknowledge_calls, 1); - EXPECT_EQ(ep->reset_calls, 1); + auto endpoint = std::make_shared(); + RdmaSlice slice; + slice.ep_weak_ptr = endpoint; + + auto locked = slice.ep_weak_ptr.lock(); + ASSERT_NE(locked, nullptr); + EXPECT_EQ(locked.get(), endpoint.get()); } TEST(EndpointLifecycleTest, SliceAccessAfterEndpointEvicted) { - FakeSlice slice; + RdmaSlice slice; { - auto ep = std::make_shared(); - slice.ep_weak_ptr = ep; - } // endpoint evicted — shared_ptr destroyed + auto endpoint = std::make_shared(); + slice.ep_weak_ptr = endpoint; + ASSERT_FALSE(slice.ep_weak_ptr.expired()); + } - // Simulate workers.cpp: lock() returns nullptr, gracefully skip - auto locked = slice.ep_weak_ptr.lock(); - EXPECT_EQ(locked, nullptr); - // No crash — the slice safely detected endpoint destruction + EXPECT_TRUE(slice.ep_weak_ptr.expired()); + EXPECT_EQ(slice.ep_weak_ptr.lock(), nullptr); } -TEST(EndpointLifecycleTest, MultipleSlicesSameEndpoint) { - auto ep = std::make_shared(); - FakeSlice slices[3]; - for (auto& s : slices) s.ep_weak_ptr = ep; +TEST(EndpointLifecycleTest, MultipleSlicesShareEndpointLifetime) { + auto endpoint = std::make_shared(); + RdmaSlice slices[3]; + for (auto& slice : slices) slice.ep_weak_ptr = endpoint; - // All slices can lock while endpoint alive - for (auto& s : slices) { - auto locked = s.ep_weak_ptr.lock(); + for (auto& slice : slices) { + auto locked = slice.ep_weak_ptr.lock(); ASSERT_NE(locked, nullptr); - locked->acknowledge(); + EXPECT_EQ(locked.get(), endpoint.get()); } - EXPECT_EQ(ep->acknowledge_calls, 3); - - // Simulate eviction: drop the owning shared_ptr - ep.reset(); - // All slices now get nullptr - for (auto& s : slices) { - EXPECT_EQ(s.ep_weak_ptr.lock(), nullptr); + endpoint.reset(); + for (auto& slice : slices) { + EXPECT_TRUE(slice.ep_weak_ptr.expired()); + EXPECT_EQ(slice.ep_weak_ptr.lock(), nullptr); } } -TEST(EndpointLifecycleTest, WeakPtrResetClearsReference) { - auto ep = std::make_shared(); - FakeSlice slice; - slice.ep_weak_ptr = ep; +TEST(EndpointLifecycleTest, SliceWeakPtrResetClearsReference) { + auto endpoint = std::make_shared(); + RdmaSlice slice; + slice.ep_weak_ptr = endpoint; - // Simulate rdma_transport.cpp slice initialization: reset() slice.ep_weak_ptr.reset(); + EXPECT_TRUE(slice.ep_weak_ptr.expired()); EXPECT_EQ(slice.ep_weak_ptr.lock(), nullptr); + EXPECT_NE(endpoint, nullptr); +} + +TEST(EndpointLifecycleTest, BootstrapWithNewPeerQpsRetiresEstablishedEndpoint) { + // The peer dropped its endpoint (store eviction or a transfer failure) + // and bootstraps again with a fresh QP set. The established endpoint now + // points at QPs that no longer exist, so it must retire instead of + // staying EP_READY and rejecting every later bootstrap from that peer. + RdmaEndPoint endpoint; + EndpointTestAccess::markConnected(endpoint, "10.0.0.1:12345", "mlx5_0", + {100, 101}); + + BootstrapDesc peer_desc, local_desc; + peer_desc.local_nic_path = MakeNicPath("10.0.0.1:12345", "mlx5_0"); + peer_desc.qp_num = {200, 201}; + + EXPECT_FALSE(endpoint.accept(peer_desc, local_desc).ok()); + EXPECT_EQ(endpoint.status(), RdmaEndPoint::EP_DESTROYING); +} - // Original endpoint still alive - EXPECT_NE(ep, nullptr); +TEST(EndpointLifecycleTest, ExternalOwnerCanReleaseAfterExplicitDeconstruct) { + auto endpoint = std::make_shared(); + std::weak_ptr weak = endpoint; + + ASSERT_EQ(endpoint->deconstruct(), 0); + EXPECT_EQ(endpoint->status(), RdmaEndPoint::EP_DESTROYED); + + endpoint.reset(); + EXPECT_TRUE(weak.expired()); } } // namespace diff --git a/mooncake-transfer-engine/tent/tests/endpoint_store_test.cpp b/mooncake-transfer-engine/tent/tests/endpoint_store_test.cpp new file mode 100644 index 0000000000..a7fc0c3ed4 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/endpoint_store_test.cpp @@ -0,0 +1,155 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include + +#include "tent/transport/rdma/context.h" +#include "tent/transport/rdma/endpoint.h" +#include "tent/transport/rdma/endpoint_store.h" +#include "tent/transport/rdma/rdma_transport.h" + +namespace mooncake { +namespace tent { + +class EndpointStoreTestAccess { + public: + static void insertWaiting(FIFOEndpointStore& store, + std::shared_ptr endpoint) { + endpoint->beginDestroy(); + RWSpinlock::WriteGuard guard(store.endpoint_map_lock_); + store.waiting_list_.insert(std::move(endpoint)); + } + + static void insertWaiting(SIEVEEndpointStore& store, + std::shared_ptr endpoint) { + endpoint->beginDestroy(); + RWSpinlock::WriteGuard guard(store.endpoint_map_lock_); + if (store.waiting_list_.insert(std::move(endpoint)).second) { + store.waiting_list_len_.fetch_add(1, std::memory_order_relaxed); + } + } + + static size_t waitingListSize(FIFOEndpointStore& store) { + RWSpinlock::ReadGuard guard(store.endpoint_map_lock_); + return store.waiting_list_.size(); + } + + static size_t waitingListSize(SIEVEEndpointStore& store) { + RWSpinlock::ReadGuard guard(store.endpoint_map_lock_); + return store.waiting_list_.size(); + } +}; + +namespace { + +enum class StoreType { FIFO, SIEVE }; + +class EndpointStoreTest : public testing::TestWithParam { + protected: + EndpointStoreTest() : context_(transport_) {} + + std::unique_ptr makeStore() { + if (GetParam() == StoreType::FIFO) { + return std::make_unique(context_, 4); + } + return std::make_unique(context_, 4); + } + + void insertWaiting(EndpointStore& store, + std::shared_ptr endpoint) { + if (GetParam() == StoreType::FIFO) { + EndpointStoreTestAccess::insertWaiting( + static_cast(store), std::move(endpoint)); + } else { + EndpointStoreTestAccess::insertWaiting( + static_cast(store), std::move(endpoint)); + } + } + + size_t waitingListSize(EndpointStore& store) { + if (GetParam() == StoreType::FIFO) { + return EndpointStoreTestAccess::waitingListSize( + static_cast(store)); + } + return EndpointStoreTestAccess::waitingListSize( + static_cast(store)); + } + + RdmaTransport transport_; + RdmaContext context_; +}; + +TEST_P(EndpointStoreTest, ReclaimDrainsQuiescentEntries) { + auto store = makeStore(); + + constexpr size_t kEndpointCount = 10; + for (size_t i = 0; i < kEndpointCount; ++i) { + insertWaiting(*store, std::make_shared()); + } + ASSERT_EQ(waitingListSize(*store), kEndpointCount); + + store->reclaim(); + EXPECT_EQ(waitingListSize(*store), 0); +} + +TEST_P(EndpointStoreTest, ReclaimIsIdempotentWhenEmpty) { + auto store = makeStore(); + + store->reclaim(); + EXPECT_EQ(waitingListSize(*store), 0); + + insertWaiting(*store, std::make_shared()); + store->reclaim(); + ASSERT_EQ(waitingListSize(*store), 0); + + store->reclaim(); + EXPECT_EQ(waitingListSize(*store), 0); +} + +TEST_P(EndpointStoreTest, ReclaimDrainsBacklogWithoutActiveMapEntries) { + auto store = makeStore(); + + constexpr size_t kEndpointCount = 1000; + for (size_t i = 0; i < kEndpointCount; ++i) { + insertWaiting(*store, std::make_shared()); + } + ASSERT_EQ(store->size(), 0); + ASSERT_EQ(waitingListSize(*store), kEndpointCount); + + store->reclaim(); + EXPECT_EQ(waitingListSize(*store), 0); +} + +TEST_P(EndpointStoreTest, ClearDeconstructsExternallyOwnedWaitingEndpoint) { + auto store = makeStore(); + auto endpoint = std::make_shared(); + std::weak_ptr weak = endpoint; + insertWaiting(*store, endpoint); + + ASSERT_EQ(store->clear(), 0); + + EXPECT_EQ(waitingListSize(*store), 0); + EXPECT_EQ(endpoint->status(), RdmaEndPoint::EP_DESTROYED); + endpoint.reset(); + EXPECT_TRUE(weak.expired()); +} + +INSTANTIATE_TEST_SUITE_P(AllStoreTypes, EndpointStoreTest, + testing::Values(StoreType::FIFO, StoreType::SIEVE)); + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/intent_type_test.cpp b/mooncake-transfer-engine/tent/tests/intent_type_test.cpp new file mode 100644 index 0000000000..ac56a2852e --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/intent_type_test.cpp @@ -0,0 +1,95 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Unit tests for IntentType enum and its integration with Request. + +#include + +#include + +#include "tent/common/types.h" + +namespace mooncake { +namespace tent { +namespace { + +TEST(IntentTypeTest, DefaultIsUnspec) { + Request r{}; + EXPECT_EQ(r.intent_type, IntentType::INTENT_UNSPEC); +} + +TEST(IntentTypeTest, AllValuesAssignable) { + Request r{}; + r.intent_type = IntentType::FOREGROUND_GET; + EXPECT_EQ(r.intent_type, IntentType::FOREGROUND_GET); + r.intent_type = IntentType::BACKGROUND_PREFETCH; + EXPECT_EQ(r.intent_type, IntentType::BACKGROUND_PREFETCH); + r.intent_type = IntentType::MIGRATION; + EXPECT_EQ(r.intent_type, IntentType::MIGRATION); + r.intent_type = IntentType::CHECKPOINT; + EXPECT_EQ(r.intent_type, IntentType::CHECKPOINT); + r.intent_type = IntentType::WEIGHT_LOADING; + EXPECT_EQ(r.intent_type, IntentType::WEIGHT_LOADING); + r.intent_type = IntentType::STAGING_INTERNAL; + EXPECT_EQ(r.intent_type, IntentType::STAGING_INTERNAL); +} + +TEST(IntentTypeTest, IntegerValues) { + EXPECT_EQ(static_cast(IntentType::INTENT_UNSPEC), 0); + EXPECT_EQ(static_cast(IntentType::FOREGROUND_GET), 1); + EXPECT_EQ(static_cast(IntentType::BACKGROUND_PREFETCH), 2); + EXPECT_EQ(static_cast(IntentType::MIGRATION), 3); + EXPECT_EQ(static_cast(IntentType::CHECKPOINT), 4); + EXPECT_EQ(static_cast(IntentType::WEIGHT_LOADING), 5); + EXPECT_EQ(static_cast(IntentType::STAGING_INTERNAL), 6); +} + +TEST(IntentTypeTest, DoesNotAffectOtherFields) { + Request r{}; + r.opcode = Request::READ; + r.priority = PRIO_LOW; + r.deadline_ns = 12345; + r.transport_hint = RDMA; + r.intent_type = IntentType::CHECKPOINT; + + EXPECT_EQ(r.opcode, Request::READ); + EXPECT_EQ(r.priority, PRIO_LOW); + EXPECT_EQ(r.deadline_ns, 12345u); + EXPECT_EQ(r.transport_hint, RDMA); + EXPECT_EQ(r.intent_type, IntentType::CHECKPOINT); +} + +TEST(IntentTypeTest, CopyPreservesIntentType) { + Request r{}; + r.intent_type = IntentType::WEIGHT_LOADING; + Request copy = r; + EXPECT_EQ(copy.intent_type, IntentType::WEIGHT_LOADING); +} + +TEST(IntentTypeTest, VectorOfRequests) { + std::vector batch(4); + batch[0].intent_type = IntentType::FOREGROUND_GET; + batch[1].intent_type = IntentType::BACKGROUND_PREFETCH; + batch[2].intent_type = IntentType::MIGRATION; + batch[3].intent_type = IntentType::INTENT_UNSPEC; + + EXPECT_EQ(batch[0].intent_type, IntentType::FOREGROUND_GET); + EXPECT_EQ(batch[1].intent_type, IntentType::BACKGROUND_PREFETCH); + EXPECT_EQ(batch[2].intent_type, IntentType::MIGRATION); + EXPECT_EQ(batch[3].intent_type, IntentType::INTENT_UNSPEC); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/metrics_config_loader_test.cpp b/mooncake-transfer-engine/tent/tests/metrics_config_loader_test.cpp index 64f0f54540..ab2bf33075 100644 --- a/mooncake-transfer-engine/tent/tests/metrics_config_loader_test.cpp +++ b/mooncake-transfer-engine/tent/tests/metrics_config_loader_test.cpp @@ -64,10 +64,6 @@ TEST(MetricsConfigLoaderTest, GetDefaultConfigReturnsExpectedDefaults) { EXPECT_EQ(config.http_port, 9100); EXPECT_EQ(config.http_server_threads, 2); EXPECT_EQ(config.report_interval_seconds, 30); - EXPECT_TRUE(config.enable_prometheus); - EXPECT_TRUE(config.enable_json); - EXPECT_TRUE(config.latency_buckets.empty()); - EXPECT_TRUE(config.size_buckets.empty()); } //------------------------------------------------------------------------------ @@ -81,11 +77,7 @@ TEST(MetricsConfigLoaderTest, LoadFromConfigWithAllValues) { "metrics/http_port": 8080, "metrics/http_host": "127.0.0.1", "metrics/http_server_threads": 4, - "metrics/report_interval_seconds": 60, - "metrics/enable_prometheus": false, - "metrics/enable_json": true, - "metrics/latency_buckets": [0.001, 0.01, 0.1, 1.0], - "metrics/size_buckets": [1024, 10240, 102400] + "metrics/report_interval_seconds": 60 })"; ASSERT_TRUE(config.load(json_content).ok()); @@ -96,13 +88,6 @@ TEST(MetricsConfigLoaderTest, LoadFromConfigWithAllValues) { EXPECT_EQ(metrics_config.http_host, "127.0.0.1"); EXPECT_EQ(metrics_config.http_server_threads, 4); EXPECT_EQ(metrics_config.report_interval_seconds, 60); - EXPECT_FALSE(metrics_config.enable_prometheus); - EXPECT_TRUE(metrics_config.enable_json); - ASSERT_EQ(metrics_config.latency_buckets.size(), 4); - EXPECT_DOUBLE_EQ(metrics_config.latency_buckets[0], 0.001); - EXPECT_DOUBLE_EQ(metrics_config.latency_buckets[3], 1.0); - ASSERT_EQ(metrics_config.size_buckets.size(), 3); - EXPECT_DOUBLE_EQ(metrics_config.size_buckets[0], 1024); } TEST(MetricsConfigLoaderTest, LoadFromConfigWithPartialValues) { @@ -145,8 +130,6 @@ TEST(MetricsConfigLoaderTest, LoadFromEnvironmentWithAllVars) { EnvVarGuard g3(config_keys::ENV_METRICS_HTTP_HOST, "192.168.1.1"); EnvVarGuard g4(config_keys::ENV_METRICS_HTTP_SERVER_THREADS, "8"); EnvVarGuard g5(config_keys::ENV_METRICS_REPORT_INTERVAL, "120"); - EnvVarGuard g6(config_keys::ENV_METRICS_ENABLE_PROMETHEUS, "true"); - EnvVarGuard g7(config_keys::ENV_METRICS_ENABLE_JSON, "false"); MetricsConfig config = MetricsConfigLoader::loadFromEnvironment(); @@ -155,8 +138,6 @@ TEST(MetricsConfigLoaderTest, LoadFromEnvironmentWithAllVars) { EXPECT_EQ(config.http_host, "192.168.1.1"); EXPECT_EQ(config.http_server_threads, 8); EXPECT_EQ(config.report_interval_seconds, 120); - EXPECT_TRUE(config.enable_prometheus); - EXPECT_FALSE(config.enable_json); } TEST(MetricsConfigLoaderTest, LoadFromEnvironmentWithPartialVars) { @@ -170,27 +151,16 @@ TEST(MetricsConfigLoaderTest, LoadFromEnvironmentWithPartialVars) { EXPECT_EQ(config.http_host, "0.0.0.0"); } -TEST(MetricsConfigLoaderTest, LoadFromEnvironmentWithLatencyBuckets) { - EnvVarGuard g1(config_keys::ENV_METRICS_LATENCY_BUCKETS, - "0.001,0.005,0.01"); +TEST(MetricsConfigLoaderTest, InvalidNumericEnvironmentValuesUseDefaults) { + EnvVarGuard g1(config_keys::ENV_METRICS_HTTP_PORT, "8080junk"); + EnvVarGuard g2(config_keys::ENV_METRICS_HTTP_SERVER_THREADS, "2x"); + EnvVarGuard g3(config_keys::ENV_METRICS_REPORT_INTERVAL, "30s"); MetricsConfig config = MetricsConfigLoader::loadFromEnvironment(); - ASSERT_EQ(config.latency_buckets.size(), 3); - EXPECT_DOUBLE_EQ(config.latency_buckets[0], 0.001); - EXPECT_DOUBLE_EQ(config.latency_buckets[1], 0.005); - EXPECT_DOUBLE_EQ(config.latency_buckets[2], 0.01); -} - -TEST(MetricsConfigLoaderTest, LoadFromEnvironmentWithSizeBuckets) { - EnvVarGuard g1(config_keys::ENV_METRICS_SIZE_BUCKETS, "1024,2048,4096"); - - MetricsConfig config = MetricsConfigLoader::loadFromEnvironment(); - - ASSERT_EQ(config.size_buckets.size(), 3); - EXPECT_DOUBLE_EQ(config.size_buckets[0], 1024); - EXPECT_DOUBLE_EQ(config.size_buckets[1], 2048); - EXPECT_DOUBLE_EQ(config.size_buckets[2], 4096); + EXPECT_EQ(config.http_port, 9100); + EXPECT_EQ(config.http_server_threads, 2); + EXPECT_EQ(config.report_interval_seconds, 30); } //------------------------------------------------------------------------------ @@ -230,6 +200,43 @@ TEST(MetricsConfigLoaderTest, LoadWithDefaultsEnvOverridesDefault) { EXPECT_EQ(config.http_port, 5555); } +// Regression guard: the `metrics/latency_buckets` and `metrics/size_buckets` +// keys were removed from the default config (bucket boundaries are now +// compile-time constants in tent_metrics.h). The loader silently ignores +// unknown keys, so a config still carrying these stale keys must (a) not +// break loading and (b) still honor the real keys present alongside them. +// If a future change re-introduces a buckets field on MetricsConfig, this +// test should be extended with an explicit assertion on the loaded value. +TEST(MetricsConfigLoaderTest, StaleBucketKeysAreSilentlyIgnored) { + Config file_config; + // Mirror the real transfer-engine.json structure: nested "metrics" object + // carrying both real keys and the removed stale bucket keys. + std::string json_content = R"({ + "metrics": { + "enabled": true, + "http_port": 7777, + "http_host": "0.0.0.0", + "http_server_threads": 2, + "report_interval_seconds": 15, + "latency_buckets": [0.001, 0.002, 0.005], + "size_buckets": [1024, 4096, 16384] + } + })"; + ASSERT_TRUE(file_config.load(json_content).ok()); + + auto mc = MetricsConfigLoader::loadWithDefaults(&file_config); + + // Real keys must be honored despite the stale keys present above. + EXPECT_TRUE(mc.enabled); + EXPECT_EQ(mc.http_port, 7777); + EXPECT_EQ(mc.http_host, "0.0.0.0"); + EXPECT_EQ(mc.http_server_threads, 2); + EXPECT_EQ(mc.report_interval_seconds, 15); + // The stale `latency_buckets` / `size_buckets` keys are silently + // ignored: MetricsConfig has no corresponding fields (see + // config_loader.h), so there is nothing to assert about their values. +} + //------------------------------------------------------------------------------ // MetricsConfigLoader::validateConfig Tests //------------------------------------------------------------------------------ @@ -262,56 +269,6 @@ TEST(MetricsConfigLoaderTest, ValidateConfigInvalidThreadsZero) { EXPECT_NE(error_msg.find("threads"), std::string::npos); } -TEST(MetricsConfigLoaderTest, ValidateConfigNoOutputFormatEnabled) { - MetricsConfig config = MetricsConfigLoader::getDefaultConfig(); - config.enable_prometheus = false; - config.enable_json = false; - std::string error_msg; - - EXPECT_FALSE(MetricsConfigLoader::validateConfig(config, &error_msg)); - EXPECT_FALSE(error_msg.empty()); - EXPECT_NE(error_msg.find("format"), std::string::npos); -} - -TEST(MetricsConfigLoaderTest, ValidateConfigUnsortedLatencyBuckets) { - MetricsConfig config = MetricsConfigLoader::getDefaultConfig(); - config.latency_buckets = {0.1, 0.05, 0.2}; // Not sorted - std::string error_msg; - - EXPECT_FALSE(MetricsConfigLoader::validateConfig(config, &error_msg)); - EXPECT_FALSE(error_msg.empty()); - EXPECT_NE(error_msg.find("Latency"), std::string::npos); -} - -TEST(MetricsConfigLoaderTest, ValidateConfigUnsortedSizeBuckets) { - MetricsConfig config = MetricsConfigLoader::getDefaultConfig(); - config.size_buckets = {1024, 512, 2048}; // Not sorted - std::string error_msg; - - EXPECT_FALSE(MetricsConfigLoader::validateConfig(config, &error_msg)); - EXPECT_FALSE(error_msg.empty()); - EXPECT_NE(error_msg.find("Size"), std::string::npos); -} - -TEST(MetricsConfigLoaderTest, ValidateConfigDuplicateBucketValues) { - MetricsConfig config = MetricsConfigLoader::getDefaultConfig(); - config.latency_buckets = {0.1, 0.1, 0.2}; // Duplicate values - std::string error_msg; - - EXPECT_FALSE(MetricsConfigLoader::validateConfig(config, &error_msg)); - EXPECT_FALSE(error_msg.empty()); -} - -TEST(MetricsConfigLoaderTest, ValidateConfigValidSortedBuckets) { - MetricsConfig config = MetricsConfigLoader::getDefaultConfig(); - config.latency_buckets = {0.001, 0.01, 0.1, 1.0}; - config.size_buckets = {1024, 10240, 102400}; - std::string error_msg; - - EXPECT_TRUE(MetricsConfigLoader::validateConfig(config, &error_msg)); - EXPECT_TRUE(error_msg.empty()); -} - TEST(MetricsConfigLoaderTest, ValidateConfigNullErrorMsg) { MetricsConfig config = MetricsConfigLoader::getDefaultConfig(); config.http_port = 0; @@ -356,6 +313,7 @@ TEST(ConfigHelperTest, ParseIntValid) { TEST(ConfigHelperTest, ParseIntInvalid) { EXPECT_EQ(ConfigHelper::parseInt("not-a-number", 99), 99); EXPECT_EQ(ConfigHelper::parseInt("", 50), 50); + EXPECT_EQ(ConfigHelper::parseInt("42x", 99), 99); } TEST(ConfigHelperTest, ParsePortValid) { @@ -367,6 +325,7 @@ TEST(ConfigHelperTest, ParsePortValid) { TEST(ConfigHelperTest, ParsePortInvalid) { EXPECT_EQ(ConfigHelper::parsePort("not-a-port", 9100), 9100); EXPECT_EQ(ConfigHelper::parsePort("", 9100), 9100); + EXPECT_EQ(ConfigHelper::parsePort("8080junk", 9100), 9100); } TEST(ConfigHelperTest, ParseDoubleArrayValid) { @@ -401,12 +360,6 @@ TEST(ConfigKeysTest, ConfigKeyConstants) { "metrics/http_server_threads"); EXPECT_STREQ(config_keys::METRICS_REPORT_INTERVAL, "metrics/report_interval_seconds"); - EXPECT_STREQ(config_keys::METRICS_ENABLE_PROMETHEUS, - "metrics/enable_prometheus"); - EXPECT_STREQ(config_keys::METRICS_ENABLE_JSON, "metrics/enable_json"); - EXPECT_STREQ(config_keys::METRICS_LATENCY_BUCKETS, - "metrics/latency_buckets"); - EXPECT_STREQ(config_keys::METRICS_SIZE_BUCKETS, "metrics/size_buckets"); } TEST(ConfigKeysTest, EnvVarConstants) { @@ -418,14 +371,6 @@ TEST(ConfigKeysTest, EnvVarConstants) { "TENT_METRICS_HTTP_SERVER_THREADS"); EXPECT_STREQ(config_keys::ENV_METRICS_REPORT_INTERVAL, "TENT_METRICS_REPORT_INTERVAL"); - EXPECT_STREQ(config_keys::ENV_METRICS_ENABLE_PROMETHEUS, - "TENT_METRICS_ENABLE_PROMETHEUS"); - EXPECT_STREQ(config_keys::ENV_METRICS_ENABLE_JSON, - "TENT_METRICS_ENABLE_JSON"); - EXPECT_STREQ(config_keys::ENV_METRICS_LATENCY_BUCKETS, - "TENT_METRICS_LATENCY_BUCKETS"); - EXPECT_STREQ(config_keys::ENV_METRICS_SIZE_BUCKETS, - "TENT_METRICS_SIZE_BUCKETS"); } } // namespace diff --git a/mooncake-transfer-engine/tent/tests/metrics_http_server_test.cpp b/mooncake-transfer-engine/tent/tests/metrics_http_server_test.cpp new file mode 100644 index 0000000000..b15024c10f --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/metrics_http_server_test.cpp @@ -0,0 +1,103 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include "tent/metrics/tent_metrics.h" + +#if TENT_METRICS_ENABLED + +#include +#include +#include +#include + +namespace mooncake { +namespace tent { +namespace { + +// RAII helper that binds and listens on an OS-assigned port, thereby keeping +// that port busy for the duration of the test. +class PortOccupier { + public: + PortOccupier() { + fd_ = ::socket(AF_INET, SOCK_STREAM, 0); + EXPECT_GE(fd_, 0); + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = htonl(INADDR_ANY); + addr.sin_port = 0; // let the kernel pick a free port + EXPECT_EQ(::bind(fd_, reinterpret_cast(&addr), sizeof(addr)), + 0); + EXPECT_EQ(::listen(fd_, 1), 0); + + sockaddr_in bound{}; + socklen_t len = sizeof(bound); + EXPECT_EQ(::getsockname(fd_, reinterpret_cast(&bound), &len), + 0); + port_ = ntohs(bound.sin_port); + } + + ~PortOccupier() { + if (fd_ >= 0) ::close(fd_); + } + + uint16_t port() const { return port_; } + + private: + int fd_ = -1; + uint16_t port_ = 0; +}; + +// When the configured metrics port is already in use (e.g. another rank was +// mistakenly given the same port), initialize() must NOT falsely report a +// listening endpoint. It degrades to log-only metrics: init still succeeds, +// but the HTTP server is not bound, so httpPort() reports 0. +TEST(TentMetricsHttpServer, DegradesToLogOnlyWhenConfiguredPortBusy) { + PortOccupier occupier; + ASSERT_GT(occupier.port(), 0); + + MetricsConfig config; + config.enabled = true; + config.http_host = "127.0.0.1"; + config.http_port = occupier.port(); + config.report_interval_seconds = 0; // no periodic reporting thread + + auto& metrics = TentMetrics::instance(); + Status status = metrics.initialize(config); + + // Metrics subsystem still comes up (counters + log summary work without a + // port), but the scrape endpoint is not listening on the busy port. + EXPECT_TRUE(status.ok()) << status.ToString(); + EXPECT_TRUE(metrics.isInitialized()); + EXPECT_EQ(metrics.httpPort(), 0); + + metrics.shutdown(); +} + +} // namespace +} // namespace tent +} // namespace mooncake + +#else // !TENT_METRICS_ENABLED + +// With metrics compiled out, initialize() is a no-op stub. Assert only that it +// reports success so the test target still builds and runs in default builds. +TEST(TentMetricsHttpServer, DisabledAtCompileTime) { + mooncake::tent::MetricsConfig config; + EXPECT_TRUE( + mooncake::tent::TentMetrics::instance().initialize(config).ok()); +} + +#endif // TENT_METRICS_ENABLED diff --git a/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp b/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp new file mode 100644 index 0000000000..a88f9735cc --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/metrics_recording_test.cpp @@ -0,0 +1,1120 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// TDD tests for the TENT metrics recording path. Drives a real +// TransferEngineImpl with a FakeTransport (pattern borrowed from +// causal_chain_test.cpp) and asserts on TentMetrics JSON/Prometheus output. +// +// All assertions are DELTA-based (snapshot before, act, snapshot after, +// compare). TentMetrics is a process-wide singleton whose counter/histogram +// values are not reset by shutdown()/initialize(), so absolute-value +// assertions would be order-dependent across tests. + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/common/types.h" +#include "tent/metrics/config_loader.h" +#include "tent/metrics/tent_metrics.h" +#include "tent/runtime/transfer_engine_impl.h" +#include "tent/runtime/transport.h" +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake { +namespace tent { +namespace { + +#if TENT_METRICS_ENABLED + +#include // required before ylt headers (coro_io uses std::signal) +#include + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +// Bind a loopback socket on port 0, read the assigned port, close, return it. +// TOCTOU race is inherent; callers re-check via TentMetrics::httpPort(). +// Adapted from mooncake-store/src/utils.cpp getFreeTcpPort(). +uint16_t getFreeTcpPort() { + int sock = ::socket(AF_INET, SOCK_STREAM, 0); + if (sock < 0) return 0; + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + addr.sin_port = htons(0); + if (::bind(sock, reinterpret_cast(&addr), sizeof(addr)) != 0) { + ::close(sock); + return 0; + } + socklen_t len = sizeof(addr); + if (::getsockname(sock, reinterpret_cast(&addr), &len) != 0) { + ::close(sock); + return 0; + } + int port = ntohs(addr.sin_port); + ::close(sock); + return static_cast(port); +} + +// HTTP GET helper. Adapted from +// mooncake-store/tests/master_metrics_test.cpp FetchUrl(). +struct HttpResponse { + int http_status; + std::string body; +}; + +HttpResponse FetchUrl(uint16_t port, const std::string& path) { + coro_http::coro_http_client client; + auto result = client.get("http://127.0.0.1:" + std::to_string(port) + path); + return {result.status, std::string(result.resp_body)}; +} + +// Snapshot of TentMetrics JSON output. Provides delta-based accessors so +// tests are robust to singleton state accumulated by prior tests. +class MetricsSnapshot { + public: + explicit MetricsSnapshot(TentMetrics& m) { + std::string body = m.getJsonMetrics(); + prometheus_ = m.getPrometheusMetrics(); + try { + json_ = nlohmann::json::parse(body); + } catch (...) { + json_ = nlohmann::json::object(); + } + } + + // Exact Prometheus series value (0 if absent). `series` includes its full + // label set, for example: + // tent_transport_attempts_total{transport="rdma",operation="write"} + double series(const std::string& series) const { + std::istringstream input(prometheus_); + std::string line; + const std::string prefix = series + " "; + while (std::getline(input, line)) { + if (line.rfind(prefix, 0) != 0) continue; + try { + return std::stod(line.substr(prefix.size())); + } catch (...) { + return 0.0; + } + } + return 0.0; + } + + // Counter value (0 if absent). + double counter(const std::string& name) const { + if (!json_.contains(name)) return 0.0; + const auto& v = json_[name]; + return v.is_number() ? v.get() : 0.0; + } + + // Histogram total sample count (0 if absent). + int64_t histogramCount(const std::string& name) const { + if (!json_.contains(name) || !json_[name].is_object()) return 0; + if (!json_[name].contains("count")) return 0; + return json_[name]["count"].get(); + } + + // Histogram bucket value by boundary key (e.g. "100"). 0 if absent. + int64_t bucket(const std::string& name, const std::string& le) const { + if (!json_.contains(name) || !json_[name].contains("buckets")) return 0; + const auto& buckets = json_[name]["buckets"]; + if (!buckets.contains(le)) return 0; + return buckets[le].get(); + } + + // All bucket keys for a histogram (sorted ascending as integers). + std::vector bucketKeys(const std::string& name) const { + std::vector keys; + if (!json_.contains(name) || !json_[name].contains("buckets")) + return keys; + for (auto it = json_[name]["buckets"].begin(); + it != json_[name]["buckets"].end(); ++it) { + try { + keys.push_back(std::stoll(it.key())); + } catch (...) { + // skip non-numeric keys (shouldn't happen for our histograms) + } + } + std::sort(keys.begin(), keys.end()); + return keys; + } + + private: + nlohmann::json json_; + std::string prometheus_; +}; + +// --------------------------------------------------------------------------- +// FakeTransport: minimal Transport that completes every submitted task. +// Pattern borrowed from causal_chain_test.cpp. +// --------------------------------------------------------------------------- + +class FakeSubBatch : public Transport::SubBatch { + public: + size_t size() const override { return task_count; } + size_t task_count = 0; + std::vector requests; + std::vector statuses; +}; + +class FakeTransport : public Transport { + public: + explicit FakeTransport(TransportType self_type, bool force_fail = false, + bool force_submit_fail = false) + : self_type_(self_type), + force_fail_(force_fail), + force_submit_fail_(force_submit_fail) { + caps.dram_to_dram = true; + } + std::atomic submit_calls{0}; + + Status install(std::string&, std::shared_ptr, + std::shared_ptr, + std::shared_ptr) override { + return Status::OK(); + } + + Status allocateSubBatch(SubBatchRef& batch, size_t) override { + batch = new FakeSubBatch(); + return Status::OK(); + } + + Status freeSubBatch(SubBatchRef& batch) override { + delete static_cast(batch); + batch = nullptr; + return Status::OK(); + } + + Status submitTransferTasks(SubBatchRef batch, + const std::vector& requests) override { + ++submit_calls; + // Synchronous submit failure: report the error without enqueuing any + // physical task, exercising the engine's !status.ok() attempt-failure + // branches (commitPreparedSubmit / resubmitTransferTask / ...). + if (force_submit_fail_) { + return Status::InternalError("forced submit failure" LOC_MARK); + } + auto* fake = static_cast(batch); + for (const auto& request : requests) { + fake->requests.push_back(request); + fake->statuses.push_back( + force_fail_ ? TransferStatus{TransferStatusEnum::FAILED, 0} + : TransferStatus{TransferStatusEnum::COMPLETED, + request.length}); + ++fake->task_count; + } + batch->notifyProgress(); + return Status::OK(); + } + + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override { + auto* fake = static_cast(batch); + if (task_id < 0 || task_id >= (int)fake->statuses.size()) { + return Status::InvalidArgument("bad task_id" LOC_MARK); + } + status = fake->statuses[task_id]; + return Status::OK(); + } + + Status addMemoryBuffer(BufferDesc& desc, const MemoryOptions&) override { + desc.transports.push_back(self_type_); + return Status::OK(); + } + + Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) override { + for (auto& desc : desc_list) { + auto s = addMemoryBuffer(desc, options); + if (!s.ok()) return s; + } + return Status::OK(); + } + + Status removeMemoryBuffer(BufferDesc&) override { return Status::OK(); } + + Status allocateLocalMemory(void** addr, size_t size, + MemoryOptions&) override { + *addr = std::malloc(size); + return *addr ? Status::OK() + : Status::InternalError("malloc failed" LOC_MARK); + } + + Status freeLocalMemory(void* addr, size_t) override { + std::free(addr); + return Status::OK(); + } + + bool warmupMemory(void*, size_t) override { return false; } + + const char* getName() const override { return ""; } + + private: + TransportType self_type_; + bool force_fail_; + bool force_submit_fail_; +}; + +std::shared_ptr makeMetricsTestConfig() { + auto cfg = std::make_shared(); + cfg->set("metadata_type", "p2p"); + cfg->set("metadata_servers", ""); + cfg->set("rpc_server_hostname", "127.0.0.1"); + cfg->set("rpc_server_port", "0"); + cfg->set("log_level", "warning"); + cfg->set("merge_requests", false); + cfg->set("enable_runtime_queue", false); + cfg->set("transports/tcp/enable", false); + cfg->set("transports/shm/enable", false); + cfg->set("transports/rdma/enable", false); + cfg->set("transports/io_uring/enable", false); + cfg->set("transports/nvlink/enable", false); + cfg->set("transports/mnnvl/enable", false); + cfg->set("transports/gds/enable", false); + cfg->set("transports/ascend_direct/enable", false); + return cfg; +} + +void installFakeRdma(TransferEngineImpl& engine, + const std::shared_ptr& fake) { + std::string seg_name = engine.getSegmentName(); + ASSERT_TRUE(fake->install(seg_name, nullptr, nullptr, nullptr).ok()); + engine.swapTransportForTest(RDMA, fake); +} + +void installFakeTcp(TransferEngineImpl& engine, + const std::shared_ptr& fake) { + std::string seg_name = engine.getSegmentName(); + ASSERT_TRUE(fake->install(seg_name, nullptr, nullptr, nullptr).ok()); + engine.swapTransportForTest(TCP, fake); +} + +Request makeLocalWrite(uint8_t* ptr, size_t length, uint64_t deadline_ns = 0) { + Request request; + request.opcode = Request::WRITE; + request.source = ptr; + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(ptr); + request.length = length; + request.transport_hint = RDMA; + request.deadline_ns = deadline_ns; + return request; +} + +// Poll a batch/task until it reaches a terminal status or timeout (1s). +TransferStatusEnum pollUntilTerminal(TransferEngineImpl& engine, BatchID batch, + int task_id = 0) { + TransferStatus ts{}; + for (int i = 0; i < 200; ++i) { + engine.getTransferStatus(batch, task_id, ts); + if (ts.s == TransferStatusEnum::COMPLETED || + ts.s == TransferStatusEnum::FAILED) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(5)); + } + return ts.s; +} + +// --------------------------------------------------------------------------- +// Test fixture: re-initialize the TentMetrics singleton per test with a free +// HTTP port and no periodic reporting thread. Counter/histogram values are +// NOT reset (singleton members), so tests use delta assertions. +// --------------------------------------------------------------------------- +class MetricsRecordingTest : public ::testing::Test { + protected: + void SetUp() override { + TentMetrics::instance().shutdown(); + MetricsConfig config; + config.enabled = true; + config.http_host = "127.0.0.1"; + config.http_port = getFreeTcpPort(); + config.report_interval_seconds = 0; + auto status = TentMetrics::instance().initialize(config); + ASSERT_TRUE(status.ok()) << status.ToString(); + TentMetrics::setEnabled(true); + } + + void TearDown() override { TentMetrics::instance().shutdown(); } +}; + +// --------------------------------------------------------------------------- +// Completed transfer records bytes, requests, latency. Drives the real +// recordTaskCompletionMetrics path via FakeTransport. +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, CompletedTransferRecordsBytesAndLatency) { + auto cfg = makeMetricsTestConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake = std::make_shared(RDMA); + installFakeRdma(engine, fake); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xBB); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto before = MetricsSnapshot(TentMetrics::instance()); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buf.data(), kLen)}).ok()); + ASSERT_EQ(pollUntilTerminal(engine, batch, 0), + TransferStatusEnum::COMPLETED); + auto after = MetricsSnapshot(TentMetrics::instance()); + + EXPECT_EQ(after.counter("tent_write_bytes_total") - + before.counter("tent_write_bytes_total"), + kLen); + EXPECT_EQ(after.counter("tent_write_requests_total") - + before.counter("tent_write_requests_total"), + 1); + EXPECT_EQ(after.histogramCount("tent_write_latency_us") - + before.histogramCount("tent_write_latency_us"), + 1); + EXPECT_EQ(after.histogramCount("tent_write_size_bytes") - + before.histogramCount("tent_write_size_bytes"), + 1); + EXPECT_EQ(after.counter("tent_transport_attempts_total") - + before.counter("tent_transport_attempts_total"), + 1); + EXPECT_EQ(after.counter("tent_transport_attempt_failures_total") - + before.counter("tent_transport_attempt_failures_total"), + 0); + EXPECT_EQ(after.histogramCount("tent_transport_attempt_latency_us") - + before.histogramCount("tent_transport_attempt_latency_us"), + 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +// --------------------------------------------------------------------------- +// A recovered RDMA->TCP failover closes the failed RDMA physical attempt and +// opens a distinct TCP attempt. Logical request metrics remain final-transport +// compatible, while attempt metrics expose RDMA reliability and per-attempt +// latency without charging the RDMA/failover interval to TCP stage latency. +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, RecoveredFailoverRecordsDistinctAttempts) { + auto cfg = makeMetricsTestConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA, /*force_fail=*/true); + auto fake_tcp = std::make_shared(TCP); + installFakeRdma(engine, fake_rdma); + installFakeTcp(engine, fake_tcp); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xBC); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto before = MetricsSnapshot(TentMetrics::instance()); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buf.data(), kLen)}).ok()); + ASSERT_EQ(pollUntilTerminal(engine, batch, 0), + TransferStatusEnum::COMPLETED); + auto after = MetricsSnapshot(TentMetrics::instance()); + + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + EXPECT_EQ(fake_tcp->submit_calls.load(), 1); + + EXPECT_EQ(after.counter("tent_transport_attempts_total") - + before.counter("tent_transport_attempts_total"), + 2); + EXPECT_EQ(after.counter("tent_transport_attempt_failures_total") - + before.counter("tent_transport_attempt_failures_total"), + 1); + EXPECT_EQ(after.histogramCount("tent_transport_attempt_latency_us") - + before.histogramCount("tent_transport_attempt_latency_us"), + 2); + + const std::string rdma_attempts = + "tent_transport_attempts_total{transport=\"rdma\",operation=\"write\"}"; + const std::string tcp_attempts = + "tent_transport_attempts_total{transport=\"tcp\",operation=\"write\"}"; + const std::string rdma_failures = + "tent_transport_attempt_failures_total{transport=\"rdma\",operation=" + "\"write\"}"; + EXPECT_EQ(after.series(rdma_attempts) - before.series(rdma_attempts), 1); + EXPECT_EQ(after.series(tcp_attempts) - before.series(tcp_attempts), 1); + EXPECT_EQ(after.series(rdma_failures) - before.series(rdma_failures), 1); + + // Logical compatibility: the recovered request completes once and does not + // become a terminal request failure. + EXPECT_EQ(after.counter("tent_write_requests_total") - + before.counter("tent_write_requests_total"), + 1); + EXPECT_EQ(after.counter("tent_write_failures_total") - + before.counter("tent_write_failures_total"), + 0); + + // Backward-compatible stage decomposition: the causal-chain stage metrics + // remain attributed to the final transport (tcp) and measure the full + // request span, exactly as before this change. The per-attempt truth + // (which transport actually failed, and each attempt's latency) lives in + // the additive tent_transport_attempt_* metrics asserted above. + const std::string tcp_dispatch_count = + "tent_stage_dispatch_us_count{transport=\"tcp\"}"; + const std::string tcp_transport_count = + "tent_stage_transport_us_count{transport=\"tcp\"}"; + // Direct submissions have exactly zero queue wait. YLT omits dynamic + // histogram label series whose sample sum is zero from Prometheus output, + // but the JSON aggregate still preserves the observation count. + EXPECT_EQ(after.histogramCount("tent_stage_queue_wait_us") - + before.histogramCount("tent_stage_queue_wait_us"), + 1); + EXPECT_EQ( + after.series(tcp_dispatch_count) - before.series(tcp_dispatch_count), + 1); + EXPECT_EQ( + after.series(tcp_transport_count) - before.series(tcp_transport_count), + 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +// --------------------------------------------------------------------------- +// A synchronous submit failure on the direct path (submitTransferTasks returns +// non-ok) is closed as a failed physical attempt. The attempt-started counter +// still increments (the engine committed to the attempt before submitting), +// the attempt-failure counter increments, and the logical request terminates +// as a failure. This exercises the commitPreparedSubmit !status.ok() branch, +// which FakeTransport could not reach before force_submit_fail existed. +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, SynchronousSubmitFailureRecordsFailedAttempt) { + auto cfg = makeMetricsTestConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared( + RDMA, /*force_fail=*/false, /*force_submit_fail=*/true); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xCD); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto before = MetricsSnapshot(TentMetrics::instance()); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buf.data(), kLen)}).ok()); + ASSERT_EQ(pollUntilTerminal(engine, batch, 0), TransferStatusEnum::FAILED); + auto after = MetricsSnapshot(TentMetrics::instance()); + + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + // One RDMA attempt was started and failed synchronously; no bytes moved. + EXPECT_EQ(after.counter("tent_transport_attempts_total") - + before.counter("tent_transport_attempts_total"), + 1); + EXPECT_EQ(after.counter("tent_transport_attempt_failures_total") - + before.counter("tent_transport_attempt_failures_total"), + 1); + EXPECT_EQ(after.histogramCount("tent_transport_attempt_latency_us") - + before.histogramCount("tent_transport_attempt_latency_us"), + 1); + const std::string rdma_attempt_failures = + "tent_transport_attempt_failures_total{transport=\"rdma\",operation=" + "\"write\"}"; + EXPECT_EQ(after.series(rdma_attempt_failures) - + before.series(rdma_attempt_failures), + 1); + + // Logical request records exactly one terminal failure, no double counting + // from recordTaskCompletionMetrics closing an already-finished attempt. + EXPECT_EQ(after.counter("tent_write_requests_total") - + before.counter("tent_write_requests_total"), + 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +// --------------------------------------------------------------------------- +// When the failover target also fails synchronously, both the original async +// failure and the resubmit's synchronous failure are recorded as distinct +// failed attempts. This exercises the resubmitTransferTask !status.ok() branch. +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, FailoverResubmitSyncFailureRecordsBothAttempts) { + auto cfg = makeMetricsTestConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + // RDMA fails asynchronously (poll returns FAILED), triggering failover to + // TCP, whose submitTransferTasks then fails synchronously. + auto fake_rdma = std::make_shared(RDMA, /*force_fail=*/true); + auto fake_tcp = std::make_shared(TCP, /*force_fail=*/false, + /*force_submit_fail=*/true); + installFakeRdma(engine, fake_rdma); + installFakeTcp(engine, fake_tcp); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xEF); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto before = MetricsSnapshot(TentMetrics::instance()); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buf.data(), kLen)}).ok()); + ASSERT_EQ(pollUntilTerminal(engine, batch, 0), TransferStatusEnum::FAILED); + auto after = MetricsSnapshot(TentMetrics::instance()); + + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + EXPECT_EQ(fake_tcp->submit_calls.load(), 1); + + // Two attempts started (RDMA then TCP), both failed. + EXPECT_EQ(after.counter("tent_transport_attempts_total") - + before.counter("tent_transport_attempts_total"), + 2); + EXPECT_EQ(after.counter("tent_transport_attempt_failures_total") - + before.counter("tent_transport_attempt_failures_total"), + 2); + + const std::string rdma_failures = + "tent_transport_attempt_failures_total{transport=\"rdma\",operation=" + "\"write\"}"; + const std::string tcp_failures = + "tent_transport_attempt_failures_total{transport=\"tcp\",operation=" + "\"write\"}"; + EXPECT_EQ(after.series(rdma_failures) - before.series(rdma_failures), 1); + EXPECT_EQ(after.series(tcp_failures) - before.series(tcp_failures), 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +// --------------------------------------------------------------------------- +// Failed transfer increments failures and requests but NOT bytes or size +// histogram. Uses the direct API so the assertion is deterministic and does +// not depend on engine failover behavior. +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, FailedTransferCountsFailureNotBytes) { + auto before = MetricsSnapshot(TentMetrics::instance()); + + TentMetrics::instance().recordReadFailed(UNSPEC); + TentMetrics::instance().recordWriteFailed(UNSPEC); + + auto after = MetricsSnapshot(TentMetrics::instance()); + + EXPECT_EQ(after.counter("tent_read_failures_total") - + before.counter("tent_read_failures_total"), + 1); + EXPECT_EQ(after.counter("tent_write_failures_total") - + before.counter("tent_write_failures_total"), + 1); + EXPECT_EQ(after.counter("tent_read_requests_total") - + before.counter("tent_read_requests_total"), + 1); + EXPECT_EQ(after.counter("tent_write_requests_total") - + before.counter("tent_write_requests_total"), + 1); + // Bytes must NOT move on failure. + EXPECT_EQ(after.counter("tent_read_bytes_total") - + before.counter("tent_read_bytes_total"), + 0); + EXPECT_EQ(after.counter("tent_write_bytes_total") - + before.counter("tent_write_bytes_total"), + 0); + // Size histograms must NOT observe failures. + EXPECT_EQ(after.histogramCount("tent_read_size_bytes") - + before.histogramCount("tent_read_size_bytes"), + 0); + EXPECT_EQ(after.histogramCount("tent_write_size_bytes") - + before.histogramCount("tent_write_size_bytes"), + 0); +} + +// --------------------------------------------------------------------------- +// A transfer whose deadline was already in the past at submit must increment +// the dedicated tent_deadline_infeasible_total counter, and must NOT pollute +// the tent_deadline_mlu_permille histogram with a sentinel value. +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, InfeasibleDeadlineRecordsSeparateCounter) { + auto cfg = makeMetricsTestConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake = std::make_shared(RDMA); + installFakeRdma(engine, fake); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xBB); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto before = MetricsSnapshot(TentMetrics::instance()); + // deadline_ns = 1 is always in the past relative to steady_clock now. + ASSERT_TRUE(engine + .submitTransfer(batch, {makeLocalWrite(buf.data(), kLen, + /*deadline_ns=*/1)}) + .ok()); + ASSERT_EQ(pollUntilTerminal(engine, batch, 0), + TransferStatusEnum::COMPLETED); + auto after = MetricsSnapshot(TentMetrics::instance()); + + // The infeasible counter must exist and increment by exactly 1. + EXPECT_EQ(after.counter("tent_deadline_infeasible_total") - + before.counter("tent_deadline_infeasible_total"), + 1) + << "expected a dedicated tent_deadline_infeasible_total counter"; + // The MLU histogram must NOT receive a sentinel sample for the infeasible + // case (the old code observed MLU=5.0 here). + EXPECT_EQ(after.histogramCount("tent_deadline_mlu_permille") - + before.histogramCount("tent_deadline_mlu_permille"), + 0) + << "infeasible deadline must not pollute the MLU histogram"; + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +// --------------------------------------------------------------------------- +// A transfer whose deadline is in the future records a genuine MLU sample +// into the histogram (feasible or missed, but real). +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, FeasibleDeadlineRecordsMLU) { + auto cfg = makeMetricsTestConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake = std::make_shared(RDMA); + installFakeRdma(engine, fake); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xBB); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto before = MetricsSnapshot(TentMetrics::instance()); + // Deadline far in the future -> window is huge -> MLU is tiny but > 0. + uint64_t future_ns = static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch() + + std::chrono::hours(1)) + .count()); + ASSERT_TRUE(engine + .submitTransfer( + batch, {makeLocalWrite(buf.data(), kLen, future_ns)}) + .ok()); + ASSERT_EQ(pollUntilTerminal(engine, batch, 0), + TransferStatusEnum::COMPLETED); + auto after = MetricsSnapshot(TentMetrics::instance()); + + EXPECT_EQ(after.histogramCount("tent_deadline_mlu_permille") - + before.histogramCount("tent_deadline_mlu_permille"), + 1); + // And the infeasible counter must NOT move for a feasible deadline. + EXPECT_EQ(after.counter("tent_deadline_infeasible_total") - + before.counter("tent_deadline_infeasible_total"), + 0); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +// --------------------------------------------------------------------------- +// A failed transfer whose deadline was already in the past at submit must +// also increment the infeasible counter. The infeasible-at-submit condition +// is independent of the transfer outcome, so it is recorded for both +// COMPLETED and FAILED (the MLU histogram, which needs actual latency, is +// only recorded on COMPLETED). +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, InfeasibleDeadlineRecordedOnFailedTransfer) { + auto cfg = makeMetricsTestConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake = std::make_shared(RDMA, /*force_fail=*/true); + installFakeRdma(engine, fake); + + constexpr size_t kLen = 4096; + std::vector buf(kLen, 0xBB); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), buf.size()).ok()); + + BatchID batch = engine.allocateBatch(4); + ASSERT_NE(batch, (BatchID)0); + + auto before = MetricsSnapshot(TentMetrics::instance()); + // deadline_ns = 1 is always in the past relative to steady_clock now. + ASSERT_TRUE(engine + .submitTransfer(batch, {makeLocalWrite(buf.data(), kLen, + /*deadline_ns=*/1)}) + .ok()); + // The failing transport reports FAILED; poll until terminal. + ASSERT_EQ(pollUntilTerminal(engine, batch, 0), TransferStatusEnum::FAILED); + auto after = MetricsSnapshot(TentMetrics::instance()); + + // Infeasible counter must increment even though the transfer failed. + EXPECT_EQ(after.counter("tent_deadline_infeasible_total") - + before.counter("tent_deadline_infeasible_total"), + 1); + // MLU histogram must not receive a sample (no completion, no latency). + EXPECT_EQ(after.histogramCount("tent_deadline_mlu_permille") - + before.histogramCount("tent_deadline_mlu_permille"), + 0); + // And the write failure was recorded. + EXPECT_EQ(after.counter("tent_write_failures_total") - + before.counter("tent_write_failures_total"), + 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), buf.size()).ok()); +} + +// --------------------------------------------------------------------------- +// Regression guard: every histogram's JSON bucket keys must match its +// compile-time boundary vector. Locks the invariant that getJsonMetrics() +// pairs each histogram with the correct bucket boundaries. +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, HistogramJsonBucketsMatchBoundaries) { + struct HistSpec { + const char* name; + std::vector expected_keys; + }; + const std::vector specs = { + {"tent_read_latency_us", + {100, 500, 1000, 5000, 10000, 50000, 100000, 500000, 1000000}}, + {"tent_write_latency_us", + {100, 500, 1000, 5000, 10000, 50000, 100000, 500000, 1000000}}, + {"tent_read_size_bytes", + {1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, + 67108864, 268435456, 1073741824}}, + {"tent_write_size_bytes", + {1024, 4096, 16384, 65536, 262144, 1048576, 4194304, 16777216, + 67108864, 268435456, 1073741824}}, + {"tent_deadline_mlu_permille", + {100, 250, 500, 750, 900, 1000, 1250, 1500, 2000, 5000}}, + {"tent_stage_queue_wait_us", + {10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000, 500000}}, + {"tent_stage_dispatch_us", + {10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000, 500000}}, + {"tent_stage_transport_us", + {10, 50, 100, 500, 1000, 5000, 10000, 50000, 100000, 500000}}, + {"tent_transport_attempt_latency_us", + {100, 500, 1000, 5000, 10000, 50000, 100000, 500000, 1000000}}, + }; + + auto snap = MetricsSnapshot(TentMetrics::instance()); + for (const auto& s : specs) { + auto keys = snap.bucketKeys(s.name); + EXPECT_EQ(keys, s.expected_keys) + << "histogram " << s.name << " bucket keys mismatch"; + } +} + +// --------------------------------------------------------------------------- +// Histogram buckets are fixed at compile time. The runtime config knob +// (latency_buckets/size_buckets) has been removed; buckets are version- +// controlled in code for reproducible observability. This test asserts the +// latency histogram exposes exactly the kLatencyBuckets keys. +// --------------------------------------------------------------------------- +TEST_F(MetricsRecordingTest, BucketsAreCompileTimeDefaults) { + // Record one sample so the histogram is non-empty. + TentMetrics::instance().recordReadCompleted(UNSPEC, 4096, 0.001); + + auto snap = MetricsSnapshot(TentMetrics::instance()); + auto keys = snap.bucketKeys("tent_read_latency_us"); + ASSERT_FALSE(keys.empty()); + + // Compile-time kLatencyBuckets: 100, 500, 1000, 5000, 10000, 50000, + // 100000, 500000, 1000000. + std::vector expected = {100, 500, 1000, 5000, 10000, + 50000, 100000, 500000, 1000000}; + EXPECT_EQ(keys, expected) + << "latency buckets must be the compile-time defaults"; +} + +// --------------------------------------------------------------------------- +// L2 HTTP integration: scrape the real /metrics, /metrics/json, /health +// endpoints via coro_http_client and assert on status + body. Validates the +// HTTP wiring (handlers, content, status codes) on top of the L1 recording +// assertions. +// --------------------------------------------------------------------------- +class MetricsHttpTest : public ::testing::Test { + protected: + void SetUp() override { + TentMetrics::instance().shutdown(); + MetricsConfig config; + config.enabled = true; + config.http_host = "127.0.0.1"; + config.http_port = getFreeTcpPort(); + config.report_interval_seconds = 0; + ASSERT_TRUE(TentMetrics::instance().initialize(config).ok()); + TentMetrics::setEnabled(true); + port_ = TentMetrics::instance().httpPort(); + // If the port bind raced (another process grabbed it), degrade + // gracefully rather than failing the build. + if (port_ == 0) { + GTEST_SKIP() << "metrics HTTP server did not bind; skipping HTTP " + "integration test"; + } + } + + void TearDown() override { TentMetrics::instance().shutdown(); } + + // Retry a GET a few times to absorb the brief window between + // async_start() returning and the server accepting connections. + HttpResponse retryGet(const std::string& path, int attempts = 20) { + HttpResponse resp{0, ""}; + for (int i = 0; i < attempts; ++i) { + resp = FetchUrl(port_, path); + if (resp.http_status == 200) return resp; + std::this_thread::sleep_for(std::chrono::milliseconds(25)); + } + return resp; + } + + uint16_t port_ = 0; +}; + +TEST_F(MetricsHttpTest, PrometheusEndpointExposesAllCounters) { + // Record a mix of operations so the counters are non-zero. + TentMetrics::instance().recordReadCompleted(UNSPEC, 1024, 0.001); + TentMetrics::instance().recordWriteCompleted(UNSPEC, 2048, 0.002); + TentMetrics::instance().recordDeadlineInfeasible(UNSPEC); + + auto resp = retryGet("/metrics"); + EXPECT_EQ(resp.http_status, 200); + EXPECT_NE(resp.body.find("tent_read_bytes_total"), std::string::npos); + EXPECT_NE(resp.body.find("tent_write_bytes_total"), std::string::npos); + EXPECT_NE(resp.body.find("tent_deadline_infeasible_total"), + std::string::npos); + EXPECT_NE(resp.body.find("tent_write_latency_us_bucket"), + std::string::npos); +} + +TEST_F(MetricsHttpTest, PrometheusLabelsDistinguishTransports) { + // Record the same metric with different transports. + TentMetrics::instance().recordReadCompleted(RDMA, 1024, 0.001); + TentMetrics::instance().recordReadCompleted(TCP, 2048, 0.002); + TentMetrics::instance().recordTransportFailover(RDMA, TCP); + + auto resp = retryGet("/metrics"); + EXPECT_EQ(resp.http_status, 200); + + // Per-transport labels must appear as separate lines. + EXPECT_NE(resp.body.find("tent_read_bytes_total{transport=\"rdma\"}"), + std::string::npos); + EXPECT_NE(resp.body.find("tent_read_bytes_total{transport=\"tcp\"}"), + std::string::npos); + + // Failover counter must carry from/to labels. + EXPECT_NE(resp.body.find( + "tent_transport_failover_total{from=\"rdma\",to=\"tcp\"}"), + std::string::npos); +} + +TEST_F(MetricsHttpTest, JsonEndpointValid) { + TentMetrics::instance().recordReadCompleted(UNSPEC, 512, 0.0005); + + auto resp = retryGet("/metrics/json"); + EXPECT_EQ(resp.http_status, 200); + // Must be parseable JSON and contain expected keys. + auto json = + nlohmann::json::parse(resp.body, nullptr, /*allow_exceptions=*/false); + ASSERT_FALSE(json.is_discarded()) + << "body was not valid JSON: " << resp.body; + EXPECT_TRUE(json.contains("tent_read_bytes_total")); + EXPECT_TRUE(json.contains("tent_deadline_infeasible_total")); + EXPECT_TRUE(json.contains("tent_read_latency_us")); +} + +// Regression guard for the silent-drop bug: ylt's +// basic_dynamic_histogram::serialize() clears its output string (taking the +// # HELP / # TYPE header with it) when every observed value truncates to 0 +// under int64_t. Sub-microsecond queue_wait latencies trigger this in +// production — JSON reports a non-zero count but /metrics omits the metric +// entirely. The custom serializer in getPrometheusMetrics() walks bucket +// counts directly so the metric always appears when count > 0. +TEST_F(MetricsHttpTest, PrometheusExposesHistogramWhenAllSamplesAreZero) { + // Reproduce the bug condition: observe 3 sub-microsecond latencies. + // recordStageLatency casts to int64_t (val=0), so sum_ stays 0 — exactly + // the case ylt's serialize() drops. + auto& m = TentMetrics::instance(); + m.recordStageLatency(TentMetrics::Stage::QueueWait, UNSPEC, 0.3); + m.recordStageLatency(TentMetrics::Stage::QueueWait, UNSPEC, 0.5); + m.recordStageLatency(TentMetrics::Stage::QueueWait, UNSPEC, 0.9); + + auto resp = retryGet("/metrics"); + EXPECT_EQ(resp.http_status, 200); + // The bug: this substring was MISSING from Prometheus output even though + // JSON reported count=3. Must appear now. + EXPECT_NE(resp.body.find("tent_stage_queue_wait_us_bucket"), + std::string::npos) + << "stage_queue_wait_us silently dropped by ylt serialize() when all " + "samples truncate to 0 under int64_t"; + // The +Inf bucket must carry the cumulative count of 3. + EXPECT_NE( + resp.body.find("tent_stage_queue_wait_us_bucket{transport=\"unspec\"," + "le=\"+Inf\"} 3"), + std::string::npos) + << "+Inf bucket must be cumulative (== total count)"; + // _count line must be present with value 3. + EXPECT_NE(resp.body.find("tent_stage_queue_wait_us_count{transport=\"" + "unspec\"} 3"), + std::string::npos) + << "_count must match total observations"; +} + +// Same bug condition as above but with transport=TCP — this is what tebench +// produces and what surfaced the original silent-drop bug in production. +// Guard against a regression that only affects a subset of transports. +TEST_F(MetricsHttpTest, PrometheusExposesZeroValuedHistogramForTcpTransport) { + auto& m = TentMetrics::instance(); + + // TentMetrics is a process-wide singleton whose histogram values are not + // reset across tests, and other tests in this binary can complete real + // TCP transfers (e.g. failover recovery) that record tcp stage latency. + // So snapshot the tcp +Inf cumulative first and assert on the delta, the + // same convention JsonAndPrometheusAgreeOnHistogramCount uses. + auto tcp_inf_count = [&]() -> int64_t { + std::string needle = + "tent_stage_queue_wait_us_bucket{transport=\"tcp\",le=\"+Inf\"} "; + auto body = retryGet("/metrics").body; + auto pos = body.find(needle); + if (pos == std::string::npos) return 0; + return std::strtoll(body.c_str() + pos + needle.size(), nullptr, 10); + }; + int64_t inf_before = tcp_inf_count(); + + m.recordStageLatency(TentMetrics::Stage::QueueWait, TCP, 0.3); + m.recordStageLatency(TentMetrics::Stage::QueueWait, TCP, 0.5); + m.recordStageLatency(TentMetrics::Stage::QueueWait, TCP, 0.9); + + auto resp = retryGet("/metrics"); + EXPECT_EQ(resp.http_status, 200); + EXPECT_NE(resp.body.find("tent_stage_queue_wait_us_bucket"), + std::string::npos) + << "stage_queue_wait_us (transport=tcp) silently dropped when all " + "samples truncate to 0 under int64_t"; + // All three samples truncate to 0 (sub-microsecond), so the +Inf bucket + // (cumulative == total count) must still advance by exactly 3 even though + // sum_ stays 0 — the silent-drop regression this test guards against. + EXPECT_EQ(tcp_inf_count() - inf_before, 3) + << "+Inf bucket for transport=tcp must be cumulative (== total count)"; +} + +// Regression guard for Prometheus/JSON drift: both endpoints must agree on +// the histogram count for a given metric. Previously they used two completely +// different serialization paths (ylt serialize() vs. custom bucket walk), +// so any ylt behavior change would silently desync them. Both now share the +// bucket-walk code path. +TEST_F(MetricsHttpTest, JsonAndPrometheusAgreeOnHistogramCount) { + // Snapshot before so the assertion is delta-based. TentMetrics is a + // process-wide singleton whose counter/histogram values are not reset + // across tests, so absolute values would be order-dependent. + auto json_before_resp = retryGet("/metrics/json"); + ASSERT_EQ(json_before_resp.http_status, 200); + auto json_before_obj = nlohmann::json::parse(json_before_resp.body, nullptr, + /*allow_exceptions=*/false); + ASSERT_FALSE(json_before_obj.is_discarded()); + int64_t json_count_before = + json_before_obj["tent_stage_transport_us"]["count"]; + + // Parse the unspec _count line from Prometheus before-state too. + auto prom_before = retryGet("/metrics"); + ASSERT_EQ(prom_before.http_status, 200); + auto count_before = [&]() -> int64_t { + std::string needle = + "tent_stage_transport_us_count{transport=\"unspec\"} "; + auto pos = prom_before.body.find(needle); + if (pos == std::string::npos) return 0; + return std::strtoll(prom_before.body.c_str() + pos + needle.size(), + nullptr, 10); + }(); + + // Record 3 observations landing in distinct buckets. + auto& m = TentMetrics::instance(); + m.recordStageLatency(TentMetrics::Stage::Transport, UNSPEC, 750.0); + m.recordStageLatency(TentMetrics::Stage::Transport, UNSPEC, 1200.0); + m.recordStageLatency(TentMetrics::Stage::Transport, UNSPEC, 50.0); + const int64_t kDelta = 3; + + auto prom = retryGet("/metrics"); + auto json_resp = retryGet("/metrics/json"); + ASSERT_EQ(prom.http_status, 200); + ASSERT_EQ(json_resp.http_status, 200); + + // JSON aggregates across all transport labels; its count must advance by + // exactly kDelta. + auto json = nlohmann::json::parse(json_resp.body, nullptr, + /*allow_exceptions=*/false); + ASSERT_FALSE(json.is_discarded()) << "invalid JSON: " << json_resp.body; + int64_t json_count_after = json["tent_stage_transport_us"]["count"]; + EXPECT_EQ(json_count_after - json_count_before, kDelta) + << "JSON count delta must equal observations (" << kDelta + << "); before=" << json_count_before << " after=" << json_count_after; + + // Prometheus: the unspec label's _count must have advanced by kDelta too. + std::string expected = + "tent_stage_transport_us_count{transport=\"unspec\"} " + + std::to_string(count_before + kDelta); + EXPECT_NE(prom.body.find(expected), std::string::npos) + << "Prometheus unspec _count (" << (count_before + kDelta) + << ") must match JSON delta. Body length=" << prom.body.size(); +} + +TEST_F(MetricsHttpTest, HealthEndpointOk) { + auto resp = retryGet("/health"); + EXPECT_EQ(resp.http_status, 200); + EXPECT_EQ(resp.body, "OK"); +} + +#else // !TENT_METRICS_ENABLED + +// When metrics are disabled at compile time, the recording path is a no-op +// and there is nothing to assert on. This test confirms the stub initializes. +TEST(MetricsRecording, DisabledAtCompileTime) { + MetricsConfig config; + EXPECT_TRUE(TentMetrics::instance().initialize(config).ok()); + TentMetrics::instance().shutdown(); +} + +#endif // TENT_METRICS_ENABLED + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/progress_worker_test.cpp b/mooncake-transfer-engine/tent/tests/progress_worker_test.cpp index b69d8c7a25..feae180dcd 100644 --- a/mooncake-transfer-engine/tent/tests/progress_worker_test.cpp +++ b/mooncake-transfer-engine/tent/tests/progress_worker_test.cpp @@ -20,7 +20,7 @@ // progressBatch / waitTransferCompletion) still sees its batch // progress through failover; // * one notify advances the engine by exactly one progress step; -// * freeBatch racing the worker is safe (no UAF, no crash); +// * freeBatch racing the worker, including cross-thread free, is safe; // * worker shuts down cleanly on engine destruction. #include @@ -451,6 +451,46 @@ TEST(ProgressWorker, FreeBatchRacesWithWorker) { EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok()); } +TEST(ProgressWorker, BatchMayBeFreedFromDifferentThread) { + auto cfg = makeMinimalP2PConfig(); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA); + std::string seg = engine.getSegmentName(); + ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok()); + engine.swapTransportForTest(RDMA, fake_rdma); + + constexpr size_t kBufLen = 4096; + std::vector buf(kBufLen, 0xC4); + ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok()); + + BatchID batch_id = 0; + std::thread allocator([&] { batch_id = engine.allocateBatch(1); }); + allocator.join(); + ASSERT_NE(batch_id, (BatchID)0); + + Request req; + req.opcode = Request::WRITE; + req.source = buf.data(); + req.target_id = LOCAL_SEGMENT_ID; + req.target_offset = reinterpret_cast(buf.data()); + req.length = kBufLen; + ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok()); + + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch_id, status).ok()); + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED); + + std::atomic free_ok{false}; + std::thread freer([&] { free_ok.store(engine.freeBatch(batch_id).ok()); }); + freer.join(); + EXPECT_TRUE(free_ok.load()); + EXPECT_TRUE(engine.getTransferStatus(batch_id, status).IsInvalidArgument()); + + EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok()); +} + // --------------------------------------------------------------------------- // 5. Engine teardown joins the worker cleanly even with pending notifies. // --------------------------------------------------------------------------- diff --git a/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp b/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp new file mode 100644 index 0000000000..ab49faef73 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/promotion_policy_test.cpp @@ -0,0 +1,99 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// Deterministic reproduction of the promotion issues reported in #2528. The +// historical head-only policy is contrasted with a per-entry policy on the same +// inputs so the unintended behavior is unambiguous and re-runnable (no RDMA +// stack, no timing noise). + +#include "tent/transport/rdma/promotion_policy.h" + +#include + +#include +#include + +namespace mooncake { +namespace tent { +namespace { + +constexpr uint64_t kTimeout = 10'000; // 10us promotion timeout +constexpr uint64_t kNow = 1'000'000; // fixed "now" + +// One entry that has clearly timed out. +uint64_t timedOut() { return kNow - kTimeout - 1; } +// One entry that was just enqueued and is nowhere near the timeout. +uint64_t fresh() { return kNow - 1; } + +// --- Issue #2528 point 1: head-only decision, whole-queue promotion -------- + +TEST(PromotionPolicyTest, HeadOnlyPromotesFreshEntriesWhenHeadTimedOut) { + // Queue head is starving; the two entries behind it were just enqueued. + std::vector q = {timedOut(), fresh(), fresh()}; + + auto d = DecidePromotionHeadOnly(q, kNow, kTimeout); + + // BUG: all three are promoted, including the two fresh (non-starving) ones. + EXPECT_EQ(d.promote_indices, (std::vector{0, 1, 2})); + + // The per-entry policy promotes only the genuinely starving head. + auto fixed = DecidePromotionPerEntry(q, kNow, kTimeout); + EXPECT_EQ(fixed.promote_indices, (std::vector{0})); +} + +TEST(PromotionPolicyTest, HeadOnlyMissesTimedOutTailWhenHeadFresh) { + // Head was just enqueued; entries behind it have been starving. + std::vector q = {fresh(), timedOut(), timedOut()}; + + auto d = DecidePromotionHeadOnly(q, kNow, kTimeout); + + // BUG: nothing is promoted even though indices 1 and 2 are starving. + EXPECT_TRUE(d.promote_indices.empty()); + + auto fixed = DecidePromotionPerEntry(q, kNow, kTimeout); + EXPECT_EQ(fixed.promote_indices, (std::vector{1, 2})); +} + +// --- Shared behavior both policies must keep ------------------------------ + +TEST(PromotionPolicyTest, EmptyQueuePromotesNothing) { + std::vector q; + EXPECT_TRUE( + DecidePromotionHeadOnly(q, kNow, kTimeout).promote_indices.empty()); + EXPECT_TRUE( + DecidePromotionPerEntry(q, kNow, kTimeout).promote_indices.empty()); +} + +TEST(PromotionPolicyTest, ZeroTimestampNeverTimesOut) { + // enqueue_ts == 0 means "no timestamp" and must never be promoted. + std::vector q = {0, 0}; + EXPECT_TRUE( + DecidePromotionHeadOnly(q, kNow, kTimeout).promote_indices.empty()); + EXPECT_TRUE( + DecidePromotionPerEntry(q, kNow, kTimeout).promote_indices.empty()); +} + +TEST(PromotionPolicyTest, AllTimedOutPromotesAllUnderBothPolicies) { + // When every entry is starving the two policies agree — this is the case + // the historical policy was designed around. + std::vector q = {timedOut(), timedOut(), timedOut()}; + EXPECT_EQ(DecidePromotionHeadOnly(q, kNow, kTimeout).promote_indices, + (std::vector{0, 1, 2})); + EXPECT_EQ(DecidePromotionPerEntry(q, kNow, kTimeout).promote_indices, + (std::vector{0, 1, 2})); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/qos_contract_test.cpp b/mooncake-transfer-engine/tent/tests/qos_contract_test.cpp new file mode 100644 index 0000000000..47fc2ddcaa --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/qos_contract_test.cpp @@ -0,0 +1,304 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include "tent/common/config.h" +#include "tent/runtime/qos_contract.h" +#include "tent/thirdparty/nlohmann/json.h" + +namespace mooncake { +namespace tent { +namespace { + +void loadConfig(Config* config, const std::string& json_text) { + ASSERT_NE(config, nullptr); + auto status = config->load(json_text); + ASSERT_TRUE(status.ok()) << status.ToString(); +} + +TEST(QosContractTest, EmptyConfigIsDisabledAndPreservesRequestPriority) { + Config config; + QosContractResolver resolver; + auto status = resolver.loadFromConfig(config); + ASSERT_TRUE(status.ok()) << status.ToString(); + EXPECT_FALSE(resolver.enabled()); + + EffectiveQosPolicy effective; + status = resolver.resolve({.tenant_id = "tenant-a", + .intent = "foreground_get", + .requested_priority = PRIO_LOW}, + &effective); + ASSERT_TRUE(status.ok()) << status.ToString(); + EXPECT_FALSE(effective.enabled); + EXPECT_EQ(effective.effective_priority, PRIO_LOW); + EXPECT_EQ(effective.matched_contract, "compatibility_default"); +} + +TEST(QosContractTest, ResolvesGlobalTenantAndIntentLayers) { + Config config; + loadConfig(&config, R"json( +{ + "qos": { + "version": 1, + "defaults": { + "priority": "medium", + "max_inflight_requests": 128, + "allowed_degraded_actions": ["fallback_transport", "reject"] + }, + "tenants": [ + { + "name": "tenant-a", + "defaults": {"max_inflight_bytes": "1GiB"}, + "intents": { + "foreground_get": { + "priority": "high", + "max_inflight_bytes": "4GiB", + "max_inflight_requests": 4096, + "allowed_degraded_actions": ["fallback_transport", "local_recompute", "reject"] + } + } + } + ] + } +} +)json"); + + QosContractResolver resolver; + auto status = resolver.loadFromConfig(config); + ASSERT_TRUE(status.ok()) << status.ToString(); + EXPECT_TRUE(resolver.enabled()); + + EffectiveQosPolicy effective; + status = resolver.resolve({.tenant_id = " Tenant-A ", + .intent = "Foreground_Get", + .requested_priority = PRIO_LOW}, + &effective); + ASSERT_TRUE(status.ok()) << status.ToString(); + EXPECT_TRUE(effective.enabled); + EXPECT_TRUE(effective.matched); + EXPECT_EQ(effective.tenant_id, "tenant-a"); + EXPECT_EQ(effective.intent, "foreground_get"); + EXPECT_EQ(effective.matched_contract, "tenant-a.foreground_get"); + ASSERT_TRUE(effective.priority.has_value()); + EXPECT_EQ(*effective.priority, PRIO_HIGH); + EXPECT_EQ(effective.effective_priority, PRIO_HIGH); + ASSERT_TRUE(effective.max_inflight_bytes.has_value()); + EXPECT_EQ(*effective.max_inflight_bytes, 4ull * 1024 * 1024 * 1024); + ASSERT_TRUE(effective.max_inflight_requests.has_value()); + EXPECT_EQ(*effective.max_inflight_requests, 4096u); + ASSERT_TRUE(effective.allowed_degraded_actions.has_value()); + EXPECT_EQ(effective.allowed_degraded_actions->size(), 3u); + + auto explain = nlohmann::json::parse(resolver.explainJson(effective)); + EXPECT_EQ(explain["tenant_id"], "tenant-a"); + EXPECT_EQ(explain["matched_contract"], "tenant-a.foreground_get"); + EXPECT_EQ(explain["effective_priority"], PRIO_HIGH); + EXPECT_EQ(explain["diagnostic_scope"], "resolution_only"); + EXPECT_FALSE(explain.contains("enforcement")); + EXPECT_FALSE(explain.contains("max_bandwidth_gbps")); + EXPECT_FALSE(explain.contains("weight")); + EXPECT_FALSE(explain.contains("deadline_profile")); +} + +TEST(QosContractTest, TenantDefaultAppliesWithoutIntentMatch) { + Config config; + loadConfig(&config, R"json( +{ + "qos": { + "defaults": {"priority": "low", "max_inflight_requests": 16}, + "tenants": [ + { + "name": "tenant-a", + "defaults": {"priority": "medium", "max_inflight_bytes": "64MiB"} + } + ] + } +} +)json"); + QosContractResolver resolver; + ASSERT_TRUE(resolver.loadFromConfig(config).ok()); + + EffectiveQosPolicy effective; + auto status = resolver.resolve({.tenant_id = "tenant-a", + .intent = "checkpoint", + .requested_priority = PRIO_HIGH}, + &effective); + ASSERT_TRUE(status.ok()) << status.ToString(); + EXPECT_FALSE(effective.matched); + EXPECT_EQ(effective.matched_contract, "tenant-a.default"); + EXPECT_EQ(effective.effective_priority, PRIO_MEDIUM); + EXPECT_EQ(*effective.max_inflight_bytes, 64ull * 1024 * 1024); + EXPECT_EQ(*effective.max_inflight_requests, 16u); +} + +TEST(QosContractTest, UnknownTenantFallsBackToGlobalDefaults) { + Config config; + loadConfig(&config, R"json( +{"qos":{"version":1,"defaults":{"priority":"medium","max_inflight_bytes":"64MiB"}}} +)json"); + QosContractResolver resolver; + ASSERT_TRUE(resolver.loadFromConfig(config).ok()); + + EffectiveQosPolicy effective; + auto status = resolver.resolve({.tenant_id = "missing", + .intent = "foreground_get", + .requested_priority = PRIO_HIGH}, + &effective); + ASSERT_TRUE(status.ok()) << status.ToString(); + EXPECT_FALSE(effective.matched); + EXPECT_EQ(effective.matched_contract, "global_default"); + EXPECT_EQ(effective.effective_priority, PRIO_MEDIUM); + ASSERT_TRUE(effective.max_inflight_bytes.has_value()); + EXPECT_EQ(*effective.max_inflight_bytes, 64ull * 1024 * 1024); +} + +TEST(QosContractTest, EmptyIdentityNormalizesToCompatibilityKeys) { + Config config; + loadConfig(&config, R"json( +{ + "qos": { + "defaults": {"priority": "low"}, + "tenants": [ + {"name": "default", "intents": {"unspec": {"priority": "medium"}}} + ] + } +} +)json"); + QosContractResolver resolver; + ASSERT_TRUE(resolver.loadFromConfig(config).ok()); + + EffectiveQosPolicy effective; + auto status = resolver.resolve( + {.tenant_id = "", .intent = "", .requested_priority = PRIO_HIGH}, + &effective); + ASSERT_TRUE(status.ok()) << status.ToString(); + EXPECT_EQ(effective.tenant_id, "default"); + EXPECT_EQ(effective.intent, "unspec"); + EXPECT_EQ(effective.matched_contract, "default.unspec"); + EXPECT_EQ(effective.effective_priority, PRIO_MEDIUM); +} + +TEST(QosContractTest, EmptyAllowedDegradedActionsOverridesInherited) { + Config config; + loadConfig(&config, R"json( +{ + "qos": { + "defaults": { + "allowed_degraded_actions": ["fallback_transport", "reject"] + }, + "tenants": [ + { + "name": "tenant-a", + "intents": { + "foreground_get": {"allowed_degraded_actions": []} + } + } + ] + } +} +)json"); + QosContractResolver resolver; + ASSERT_TRUE(resolver.loadFromConfig(config).ok()); + + EffectiveQosPolicy effective; + auto status = resolver.resolve({.tenant_id = "tenant-a", + .intent = "foreground_get", + .requested_priority = PRIO_HIGH}, + &effective); + ASSERT_TRUE(status.ok()) << status.ToString(); + ASSERT_TRUE(effective.allowed_degraded_actions.has_value()); + EXPECT_TRUE(effective.allowed_degraded_actions->empty()); +} + +TEST(QosContractTest, DeferredM3ToM5FieldsFailClosed) { + static const std::vector kDeferredFields = { + R"json("min_bandwidth_gbps":20)json", + R"json("max_bandwidth_gbps":100)json", + R"json("burst_bytes":"1GiB")json", R"json("weight":8)json", + R"json("deadline_profile":"interactive")json"}; + + for (const auto& field : kDeferredFields) { + Config config; + loadConfig(&config, "{\"qos\":{\"defaults\":{" + field + "}}}"); + QosContractResolver resolver; + auto status = resolver.loadFromConfig(config); + EXPECT_FALSE(status.ok()) << field; + EXPECT_TRUE(status.IsInvalidArgument()) << field; + EXPECT_FALSE(resolver.enabled()) << field; + } +} + +TEST(QosContractTest, DeferredResolutionFeaturesFailClosed) { + static const std::vector kConfigs = { + R"json({"qos":{"strict_mode":true,"defaults":{"priority":"high"}}})json", + R"json({"qos":{"intent_defaults":{"foreground_get":{"priority":"high"}}}})json", + R"json({"qos":{"tenants":[{"name":"tenant-a","intents":{"foreground_get":{"name":"named-policy","priority":"high"}}}]}})json"}; + + for (const auto& text : kConfigs) { + Config config; + loadConfig(&config, text); + QosContractResolver resolver; + auto status = resolver.loadFromConfig(config); + EXPECT_FALSE(status.ok()) << text; + EXPECT_TRUE(status.IsInvalidArgument()) << text; + EXPECT_FALSE(resolver.enabled()) << text; + } +} + +TEST(QosContractTest, InvalidSchemaFailsClosedAndClearsPriorState) { + Config valid; + loadConfig(&valid, R"json({"qos":{"defaults":{"priority":"medium"}}})json"); + QosContractResolver resolver; + ASSERT_TRUE(resolver.loadFromConfig(valid).ok()); + ASSERT_TRUE(resolver.enabled()); + + static const std::vector kInvalidConfigs = { + R"json({"qos":{"defaults":{"priority":"urgent"}}})json", + R"json({"qos":{"defaults":{"allowed_degraded_actions":["teleport"]}}})json", + R"json({"qos":{"defaults":{"max_inflight_requests":"64MiB"}}})json", + R"json({"qos":{"defaults":{"unexpected_qos_key":10}}})json", + R"json({"qos":{"version":"1","defaults":{"priority":"high"}}})json"}; + + for (const auto& text : kInvalidConfigs) { + Config config; + loadConfig(&config, text); + auto status = resolver.loadFromConfig(config); + EXPECT_FALSE(status.ok()) << text; + EXPECT_TRUE(status.IsInvalidArgument()) << text; + EXPECT_FALSE(resolver.enabled()) << text; + } +} + +TEST(QosContractTest, IntentTypeNamesAreStable) { + EXPECT_EQ(QosContractResolver::intentTypeName(IntentType::INTENT_UNSPEC), + "unspec"); + EXPECT_EQ(QosContractResolver::intentTypeName(IntentType::FOREGROUND_GET), + "foreground_get"); + EXPECT_EQ( + QosContractResolver::intentTypeName(IntentType::BACKGROUND_PREFETCH), + "background_prefetch"); + EXPECT_EQ(QosContractResolver::intentTypeName(IntentType::MIGRATION), + "migration"); + EXPECT_EQ(QosContractResolver::intentTypeName(IntentType::CHECKPOINT), + "checkpoint"); + EXPECT_EQ(QosContractResolver::intentTypeName(IntentType::WEIGHT_LOADING), + "weight_loading"); + EXPECT_EQ(QosContractResolver::intentTypeName(IntentType::STAGING_INTERNAL), + "staging_internal"); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/qp_pool_layout_test.cpp b/mooncake-transfer-engine/tent/tests/qp_pool_layout_test.cpp new file mode 100644 index 0000000000..64c841c980 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/qp_pool_layout_test.cpp @@ -0,0 +1,212 @@ +// Copyright 2025 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Unit tests for computeQpPoolSegments — the pure QP-pool layout resolver used +// by RdmaEndPoint::construct() (RFC #2568 step 2). Kept free of RDMA handles so +// the layout math is testable without a device. + +#include + +#include "tent/transport/rdma/params.h" + +namespace mooncake { +namespace tent { +namespace { + +// Default path: no pools configured => a single homogeneous run of +// qp_mul_factor QPs, no explicit segments (poolForQp will return nullptr and +// callers fall back to the global SL/TC — byte-for-byte the prior behavior). +TEST(QpPoolLayoutTest, EmptyPoolsKeepsFlatQpMulFactor) { + auto layout = computeQpPoolSegments({}, 6); + EXPECT_TRUE(layout.valid); + EXPECT_EQ(layout.total_qp, 6); + EXPECT_TRUE(layout.segments.empty()); +} + +// A non-positive total (e.g. qp_mul_factor <= 0 with no pools) is rejected so +// construct() can fail cleanly instead of allocating a zero-length QP array. +TEST(QpPoolLayoutTest, EmptyPoolsWithNonPositiveFactorIsInvalid) { + auto layout = computeQpPoolSegments({}, 0); + EXPECT_FALSE(layout.valid); + EXPECT_EQ(layout.total_qp, 0); +} + +// Multiple pools lay out contiguous, non-overlapping segments; total is the +// sum of per-pool num_qp; qp_mul_factor is ignored once pools are set. +TEST(QpPoolLayoutTest, MultiplePoolsLayoutContiguousSegments) { + std::vector pools; + QpPoolSegment kv; + kv.name = "kv"; + kv.num_qp = 4; + kv.service_level = 5; + kv.traffic_class = 96; + pools.push_back(kv); + QpPoolSegment ctrl; + ctrl.name = "ctrl"; + ctrl.num_qp = 2; + pools.push_back(ctrl); + + auto layout = computeQpPoolSegments(pools, /*qp_mul_factor=*/6); + ASSERT_TRUE(layout.valid); + EXPECT_EQ(layout.total_qp, 6); // 4 + 2, not qp_mul_factor + ASSERT_EQ(layout.segments.size(), 2u); + + EXPECT_EQ(layout.segments[0].name, "kv"); + EXPECT_EQ(layout.segments[0].begin, 0); + EXPECT_EQ(layout.segments[0].num_qp, 4); + EXPECT_EQ(layout.segments[0].service_level, 5); + EXPECT_EQ(layout.segments[0].traffic_class, 96); + + EXPECT_EQ(layout.segments[1].name, "ctrl"); + EXPECT_EQ(layout.segments[1].begin, 4); // starts after kv's 4 QPs + EXPECT_EQ(layout.segments[1].num_qp, 2); + // ctrl left SL/TC unset -> sentinel -1 (setupOneQP falls back to global). + EXPECT_EQ(layout.segments[1].service_level, -1); + EXPECT_EQ(layout.segments[1].traffic_class, -1); +} + +// Segments partition [0, total_qp): every QP index maps to exactly one pool, +// mirroring RdmaEndPoint::poolForQp's linear scan. +TEST(QpPoolLayoutTest, SegmentsPartitionAllQpIndices) { + std::vector pools; + QpPoolSegment a; + a.name = "a"; + a.num_qp = 3; + pools.push_back(a); + QpPoolSegment b; + b.name = "b"; + b.num_qp = 1; + pools.push_back(b); + + auto layout = computeQpPoolSegments(pools, 6); + ASSERT_TRUE(layout.valid); + ASSERT_EQ(layout.total_qp, 4); + + auto pool_of = [&](int qp_index) -> const QpPoolSegment* { + for (const auto& seg : layout.segments) { + if (qp_index >= seg.begin && qp_index < seg.begin + seg.num_qp) + return &seg; + } + return nullptr; + }; + ASSERT_NE(pool_of(0), nullptr); + EXPECT_EQ(pool_of(0)->name, "a"); + EXPECT_EQ(pool_of(2)->name, "a"); + ASSERT_NE(pool_of(3), nullptr); + EXPECT_EQ(pool_of(3)->name, "b"); + // Out of range => no pool (default single-pool fallback in poolForQp). + EXPECT_EQ(pool_of(4), nullptr); +} + +// --- selectQpInPool: the step-3 router (slice pool -> QP index) +// --------------- + +// Helper: a two-pool layout kv=[0,4), ctrl=[4,6). +static std::vector twoPools() { + auto layout = computeQpPoolSegments( + {{"kv", 4, 0, -1, -1}, {"ctrl", 2, 0, -1, -1}}, 6); + return layout.segments; +} + +// Empty pool name => pass through, folded into the whole QP range. This is the +// default (no pool selected) behavior — identical to the pre-step-3 spray. +TEST(SelectQpInPoolTest, EmptyPoolNameSpraysAcrossAllQps) { + auto segs = twoPools(); + EXPECT_EQ(selectQpInPool(segs, "", 0, 6), 0); + EXPECT_EQ(selectQpInPool(segs, "", 5, 6), 5); + EXPECT_EQ(selectQpInPool(segs, "", 7, 6), 1); // 7 % 6 +} + +// No pools configured at all => also pass through (single default pool). +TEST(SelectQpInPoolTest, NoSegmentsSpraysAcrossAllQps) { + std::vector none; + EXPECT_EQ(selectQpInPool(none, "kv", 3, 6), 3); + EXPECT_EQ(selectQpInPool(none, "", 8, 6), 2); // 8 % 6 +} + +// A named pool folds the candidate into that pool's segment only. +TEST(SelectQpInPoolTest, NamedPoolFoldsIntoItsSegment) { + auto segs = twoPools(); // kv=[0,4), ctrl=[4,6) + // kv: begin 0, num 4 -> indices 0..3 + EXPECT_EQ(selectQpInPool(segs, "kv", 0, 6), 0); + EXPECT_EQ(selectQpInPool(segs, "kv", 3, 6), 3); + EXPECT_EQ(selectQpInPool(segs, "kv", 4, 6), 0); // 4 % 4 -> begin+0 + EXPECT_EQ(selectQpInPool(segs, "kv", 6, 6), 2); // 6 % 4 -> begin+2 + // ctrl: begin 4, num 2 -> indices 4..5 + EXPECT_EQ(selectQpInPool(segs, "ctrl", 0, 6), 4); + EXPECT_EQ(selectQpInPool(segs, "ctrl", 1, 6), 5); + EXPECT_EQ(selectQpInPool(segs, "ctrl", 3, 6), 5); // 3 % 2 -> begin+1 +} + +// Unknown pool name => fall back to the whole range (don't drop the transfer). +TEST(SelectQpInPoolTest, UnknownPoolFallsBackToWholeRange) { + auto segs = twoPools(); + EXPECT_EQ(selectQpInPool(segs, "nope", 5, 6), 5); + EXPECT_EQ(selectQpInPool(segs, "nope", 9, 6), 3); // 9 % 6 +} + +// Negative candidate is clamped to 0 before folding. +TEST(SelectQpInPoolTest, NegativeCandidateClampsToZero) { + auto segs = twoPools(); + EXPECT_EQ(selectQpInPool(segs, "ctrl", -1, 6), 4); // begin+0 + EXPECT_EQ(selectQpInPool(segs, "", -1, 6), 0); +} + +// Every result stays in [0, total_qp) regardless of pool/candidate. +TEST(SelectQpInPoolTest, ResultAlwaysInRange) { + auto segs = twoPools(); + for (int c = 0; c < 20; ++c) { + for (const char* name : {"", "kv", "ctrl", "nope"}) { + int idx = selectQpInPool(segs, name, c, 6); + EXPECT_GE(idx, 0); + EXPECT_LT(idx, 6); + } + } +} + +// A pool with a non-positive num_qp would create an empty/negative QP span and +// break the router, so the whole layout is rejected (falls back to default). +TEST(QpPoolLayoutTest, PoolWithZeroQpIsInvalid) { + std::vector pools; + QpPoolSegment ok; + ok.name = "kv"; + ok.num_qp = 4; + pools.push_back(ok); + QpPoolSegment bad; + bad.name = "ctrl"; + bad.num_qp = 0; // invalid + pools.push_back(bad); + + auto layout = computeQpPoolSegments(pools, 6); + EXPECT_FALSE(layout.valid); + EXPECT_EQ(layout.total_qp, 0); + EXPECT_TRUE(layout.segments.empty()); +} + +TEST(QpPoolLayoutTest, PoolWithNegativeQpIsInvalid) { + std::vector pools; + QpPoolSegment bad; + bad.name = "kv"; + bad.num_qp = -1; // invalid + pools.push_back(bad); + + auto layout = computeQpPoolSegments(pools, 6); + EXPECT_FALSE(layout.valid); + EXPECT_EQ(layout.total_qp, 0); + EXPECT_TRUE(layout.segments.empty()); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp b/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp index b625f39b72..fc05c7a411 100644 --- a/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp +++ b/mooncake-transfer-engine/tent/tests/rail_monitor_test.cpp @@ -58,6 +58,115 @@ static std::shared_ptr makeSingleNicTopology(const std::string& nic, return topo; } +static std::shared_ptr makeTwoNicTopology(const std::string& first, + const std::string& second) { + auto json_str = R"({ + "nics": [ + {"name": ")" + + first + R"(", "type": 0, "numa_node": 0}, + {"name": ")" + + second + R"(", "type": 0, "numa_node": 0} + ], + "mems": [{ + "name": "host0", + "type": 0, + "numa_node": 0, + "device_list": {"rank0": [0, 1]} + }] + })"; + auto topo = std::make_shared(); + auto status = topo->parse(json_str); + if (!status.ok()) { + ADD_FAILURE() << "Topology::parse failed: " << status.ToString(); + } + return topo; +} + +TEST(RailMonitorConfigTest, CustomJsonOverridesAutomaticPeerMapping) { + auto local = makeTwoNicTopology("local0", "local1"); + auto remote = makeTwoNicTopology("remote0", "remote1"); + const std::string rail_json = R"({ + "all": [ + {"local": "local0", "remote": "remote1"}, + {"local": "local1", "remote": "remote0"} + ], + "direct": [ + {"local": "local0", "remote": "remote1"}, + {"local": "local1", "remote": "remote0"} + ] + })"; + + RailMonitor rail; + ASSERT_TRUE(rail.load(local.get(), remote.get(), rail_json, nullptr).ok()); + EXPECT_EQ(rail.findBestRemoteDevice(/*local_nic=*/0, /*remote_numa=*/0), 1); + EXPECT_EQ(rail.findBestRemoteDevice(/*local_nic=*/1, /*remote_numa=*/0), 0); + EXPECT_TRUE(rail.available(/*local_nic=*/0, /*remote_nic=*/1)); + EXPECT_FALSE(rail.available(/*local_nic=*/0, /*remote_nic=*/0)); +} + +// Build a 2-NIC topology (mlx5_a, mlx5_b) with per-NIC NUMA nodes, so the two +// sides can disagree on which NUMA a same-named NIC sits in — the asymmetric +// (overlay) situation from #2467. +static std::shared_ptr makeNamedNumaTopology(const std::string& n0, + int numa0, + const std::string& n1, + int numa1) { + auto json_str = + R"({ + "nics": [ + {"name": ")" + + n0 + R"(", "type": 0, "numa_node": )" + std::to_string(numa0) + R"(}, + {"name": ")" + + n1 + R"(", "type": 0, "numa_node": )" + std::to_string(numa1) + R"(} + ], + "mems": [{ + "name": "host0", + "type": 0, + "numa_node": 0, + "device_list": {"rank0": [0, 1]} + }] + })"; + auto topo = std::make_shared(); + auto status = topo->parse(json_str); + if (!status.ok()) { + ADD_FAILURE() << "Topology::parse failed: " << status.ToString(); + } + return topo; +} + +// --------------------------------------------------------------------------- +// Cross-NUMA mapping must prefer a same-name remote device over a positional +// (i % remote_cnt) pick, so a local NIC is not routed to an unrelated remote +// NIC on a different physical/overlay network (issues #2758/#2467). +// +// Setup (asymmetric NUMA, as in #2467's overlay case): +// local : mlx5_x @ NUMA 0 (idx0), mlx5_y @ NUMA 1 (idx1) +// remote: mlx5_y @ NUMA 0 (idx0), mlx5_x @ NUMA 1 (idx1) +// Local mlx5_y sits in NUMA 1; its same-name remote mlx5_y sits in NUMA 0. +// Querying local mlx5_y (idx1) for the remote NUMA-0 domain is cross-NUMA and +// must pick the same-name remote mlx5_y (remote idx0). The positional bug would +// instead pick remote_devices[NUMA0][i]. With only one device in that domain +// they coincide, so we make the discriminating assertion below. +// --------------------------------------------------------------------------- + +TEST(RailMonitorCrossNumaTest, CrossNumaPrefersSameNameDevice) { + // local NUMA-1 domain has one NIC: mlx5_y (idx1). + auto local = makeNamedNumaTopology("mlx5_x", 0, "mlx5_y", 1); + // remote NUMA-0 domain: mlx5_y (idx0); remote NUMA-1 domain: mlx5_x (idx1). + auto remote = makeNamedNumaTopology("mlx5_y", 0, "mlx5_x", 1); + RailMonitor rail; + ASSERT_TRUE(rail.load(local.get(), remote.get()).ok()); + ASSERT_TRUE(rail.ready()); + + // local mlx5_y (idx1, NUMA 1) reaching the remote NUMA-0 domain: the only + // same-name device is remote mlx5_y at idx0. Must map there. + EXPECT_EQ(rail.findBestRemoteDevice(/*local_nic=*/1, /*remote_numa=*/0), 0); + + // local mlx5_x (idx0, NUMA 0) reaching remote NUMA-1 domain: same-name + // remote mlx5_x is at idx1. Must map there, not positionally to idx0. + EXPECT_EQ(rail.findBestRemoteDevice(/*local_nic=*/0, /*remote_numa=*/1), 1); +} + // --------------------------------------------------------------------------- // markRecovered resets error_count so failures start accumulating fresh // --------------------------------------------------------------------------- diff --git a/mooncake-transfer-engine/tent/tests/rdma_cancel_test.cpp b/mooncake-transfer-engine/tent/tests/rdma_cancel_test.cpp new file mode 100644 index 0000000000..30bca1d3bd --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/rdma_cancel_test.cpp @@ -0,0 +1,77 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "tent/transport/rdma/slice.h" + +#include + +namespace mooncake { +namespace tent { +namespace { + +TEST(RdmaCancelTest, PartialCancellationWaitsForPostedSlices) { + RdmaTask task{}; + task.num_slices = 2; + task.status_word = PENDING; + task.transferred_bytes = 0; + task.success_slices.store(0); + task.resolved_slices.store(0); + task.first_error = PENDING; + // Two slice references plus the batch reference. updateSliceStatus drops + // one reference per resolved slice; the stack object is never deallocated. + task.ref_count.store(3); + + RdmaSlice canceled{}; + canceled.task = &task; + canceled.word = PENDING; + canceled.length = 4096; + RdmaSlice posted{}; + posted.task = &task; + posted.word = PENDING; + posted.length = 8192; + + updateSliceStatus(&canceled, CANCELED); + EXPECT_EQ(task.status_word, PENDING); + EXPECT_EQ(task.transferred_bytes, 0u); + + updateSliceStatus(&posted, COMPLETED); + EXPECT_EQ(task.status_word, CANCELED); + EXPECT_EQ(task.transferred_bytes, 8192u); + EXPECT_EQ(task.resolved_slices.load(), 2); +} + +TEST(RdmaCancelTest, FullyPostedTaskMayStillComplete) { + RdmaTask task{}; + task.num_slices = 1; + task.status_word = PENDING; + task.transferred_bytes = 0; + task.success_slices.store(0); + task.resolved_slices.store(0); + task.first_error = PENDING; + task.cancel_requested.store(true); + task.ref_count.store(2); + + RdmaSlice posted{}; + posted.task = &task; + posted.word = PENDING; + posted.length = 4096; + updateSliceStatus(&posted, COMPLETED); + + EXPECT_EQ(task.status_word, COMPLETED); + EXPECT_EQ(task.transferred_bytes, 4096u); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp b/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp new file mode 100644 index 0000000000..6e918a21d2 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/rdma_transport_test.cpp @@ -0,0 +1,391 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/common/types.h" +#include "tent/transfer_engine.h" +#include "tent/runtime/topology.h" +#include "tent/transport/rdma/context.h" +#include "tent/transport/rdma/ibv_loader.h" +#include "tent/transport/rdma/params.h" +#include "tent/transport/rdma/rdma_transport.h" +#include "tent/transport/rdma/workers.h" + +namespace mooncake { +namespace tent { + +// Friend accessor for driving initializeContexts() without a full install(). +class RdmaTransportTestPeer { + public: + static void bindTopology(RdmaTransport& transport, + std::shared_ptr topology) { + transport.local_topology_ = topology; + transport.local_buffer_manager_.setTopology(topology); + transport.params_ = std::make_shared(); + transport.conf_ = std::make_shared(); + } + + // Runs the monitorThread() 1 Hz reclaim tick without starting any worker + // threads. + static void reclaimEndpoints(RdmaTransport& transport) { + Workers workers(&transport); + workers.reclaimEndpoints(); + } + + static size_t initializeContexts(RdmaTransport& transport) { + return transport.initializeContexts(); + } + + static const RdmaContextSet& contextSet(const RdmaTransport& transport) { + return transport.context_set_; + } +}; + +namespace { + +bool hasRdmaDevice() { + int count = 0; + ibv_device** devices = ibv_get_device_list(&count); + const bool available = devices != nullptr && count > 0; + if (devices) ibv_free_device_list(devices); + return available; +} + +class ChildProcessGuard { + public: + ChildProcessGuard(pid_t pid, int stop_fd) : pid_(pid), stop_fd_(stop_fd) {} + + ~ChildProcessGuard() { + if (pid_ <= 0) return; + close(stop_fd_); + (void)waitpid(pid_, nullptr, 0); + } + + int finish() { + close(stop_fd_); + int status = 0; + (void)waitpid(pid_, &status, 0); + pid_ = -1; + stop_fd_ = -1; + return status; + } + + int reap() { + int status = 0; + (void)waitpid(pid_, &status, 0); + close(stop_fd_); + pid_ = -1; + stop_fd_ = -1; + return status; + } + + private: + pid_t pid_; + int stop_fd_; +}; + +std::shared_ptr makeRdmaConfig() { + auto config = std::make_shared(); + config->set("metadata_type", "p2p"); + config->set("metadata_servers", "P2PHANDSHAKE"); + config->set("transports/rdma/enable", true); + config->set("transports/rdma/num_lanes", 1); + config->set("transports/rdma/endpoint/max_qp_wr", 1); + config->set("transports/tcp/enable", false); + config->set("transports/shm/enable", false); + return config; +} + +bool waitBatchDone(TransferEngine& engine, BatchID batch) { + TransferStatus status; + for (int i = 0; i < 10000; ++i) { + auto result = engine.getTransferStatus(batch, status); + if (!result.ok() || status.s == TransferStatusEnum::FAILED) + return false; + if (status.s == TransferStatusEnum::COMPLETED) return true; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + return false; +} + +TEST(RdmaParamsTest, DefaultsKeepLaneCountsAligned) { + RdmaParams params; + + EXPECT_EQ(params.num_lanes, 6); + EXPECT_EQ(params.device.num_cq_list, params.num_lanes); + EXPECT_EQ(params.endpoint.qp_mul_factor, params.num_lanes); + EXPECT_EQ(params.workers.num_workers, params.num_lanes); + EXPECT_EQ(params.endpoint.path_mtu, IBV_MTU_4096); +} + +TEST(RdmaSubBatchTest, ReportsTaskCount) { + RdmaSubBatch batch; + batch.max_size = 8; + + EXPECT_EQ(batch.size(), 0); + batch.task_list.push_back(nullptr); + batch.task_list.push_back(nullptr); + EXPECT_EQ(batch.size(), 2); +} + +// context_set_ is subscripted by NicID, so it must keep one slot per NIC even +// when a device is skipped. It used to push_back only on the success path, +// compacting the array so a later dev_id named the wrong RNIC or ran off it. +void expectInertContextPerNic(const RdmaContextSet& contexts, size_t expected) { + ASSERT_EQ(contexts.size(), expected); + for (size_t i = 0; i < contexts.size(); ++i) { + ASSERT_NE(contexts[i], nullptr) << "slot " << i << " must be occupied"; + EXPECT_NE(contexts[i]->status(), RdmaContext::DEVICE_ENABLED) + << "slot " << i << " must not report itself usable"; + // Inert contexts must stay safe for the whole-list consumers. + EXPECT_EQ(contexts[i]->cq(0), nullptr); + EXPECT_EQ(contexts[i]->notifyCq(), nullptr); + } +} + +// Non-RDMA entries are skipped before construct() is ever called, so this runs +// without libibverbs devices. +TEST(RdmaNicIndexAlignmentTest, ContextSetKeepsOneSlotPerNonRdmaNic) { + auto topology = std::make_shared(); + ASSERT_TRUE(topology + ->parse(R"({"nics":[ + {"name":"mc-tcp-0","type":1,"numa_node":0}, + {"name":"mc-unknown-1","type":2,"numa_node":0}, + {"name":"mc-tcp-2","type":1,"numa_node":0}]})") + .ok()); + ASSERT_EQ(topology->getNicCount(), static_cast(3)); + + RdmaTransport transport; + RdmaTransportTestPeer::bindTopology(transport, topology); + + EXPECT_EQ(RdmaTransportTestPeer::initializeContexts(transport), + static_cast(0)); + expectInertContextPerNic(RdmaTransportTestPeer::contextSet(transport), + topology->getNicCount()); +} + +// monitorThread()'s 1 Hz tick walks every slot. An inert context never built +// an endpoint store, so an unguarded endpointStore()->reclaim() would crash -- +// and only on the first heartbeat, well after startup. +TEST(RdmaNicIndexAlignmentTest, ReclaimTickSkipsInertContexts) { + auto topology = std::make_shared(); + ASSERT_TRUE(topology + ->parse(R"({"nics":[ + {"name":"mc-tcp-0","type":1,"numa_node":0}, + {"name":"mc-unknown-1","type":2,"numa_node":0}]})") + .ok()); + + RdmaTransport transport; + RdmaTransportTestPeer::bindTopology(transport, topology); + ASSERT_EQ(RdmaTransportTestPeer::initializeContexts(transport), + static_cast(0)); + + // Precondition that makes the unguarded call fatal. + for (const auto& context : RdmaTransportTestPeer::contextSet(transport)) { + ASSERT_EQ(context->endpointStore(), nullptr); + } + + RdmaTransportTestPeer::reclaimEndpoints(transport); +} + +// The construct()-failure branch. IbvLoader dlcloses libibverbs when no device +// is present, so construct() cannot be driven safely in that state. +TEST(RdmaNicIndexAlignmentTest, ContextSetKeepsOneSlotWhenConstructFails) { + if (!IbvLoader::Instance().available()) + GTEST_SKIP() << "no usable libibverbs; construct() cannot be driven"; + + // Device names that resolve to no real RNIC, so construct() fails on any + // host, with a non-RDMA entry in the middle to offset the indexes. + auto topology = std::make_shared(); + ASSERT_TRUE(topology + ->parse(R"({"nics":[ + {"name":"mc-absent-rnic-0","type":0,"numa_node":0}, + {"name":"mc-tcp-1","type":1,"numa_node":0}, + {"name":"mc-absent-rnic-2","type":0,"numa_node":0}]})") + .ok()); + ASSERT_EQ(topology->getNicCount(), static_cast(3)); + + RdmaTransport transport; + RdmaTransportTestPeer::bindTopology(transport, topology); + + EXPECT_EQ(RdmaTransportTestPeer::initializeContexts(transport), + static_cast(0)); + expectInertContextPerNic(RdmaTransportTestPeer::contextSet(transport), + topology->getNicCount()); +} + +TEST(RdmaTransportIntegrationTest, WriteThenReadAcrossProcesses) { + if (!hasRdmaDevice()) GTEST_SKIP() << "no RDMA device detected"; + + constexpr size_t kDataLength = 4 * 1024 * 1024; + constexpr size_t kCancelTaskCount = 16; + constexpr size_t kCancelStride = 8 * 1024 * 1024; + constexpr size_t kBufferLength = kCancelTaskCount * kCancelStride; + int ready_pipe[2]; + int stop_pipe[2]; + ASSERT_EQ(pipe(ready_pipe), 0); + ASSERT_EQ(pipe(stop_pipe), 0); + + pid_t child = fork(); + ASSERT_GE(child, 0); + if (child == 0) { + close(ready_pipe[0]); + close(stop_pipe[1]); + + TransferEngine server(makeRdmaConfig()); + if (!server.available()) _exit(2); + std::vector buffer(kBufferLength); + if (!server.registerLocalMemory(buffer.data(), buffer.size()).ok()) + _exit(3); + + const std::string segment = server.getSegmentName(); + uint32_t length = static_cast(segment.size()); + if (write(ready_pipe[1], &length, sizeof(length)) != sizeof(length)) + _exit(4); + if (write(ready_pipe[1], segment.data(), length) != + static_cast(length)) + _exit(5); + + char stop = 0; + const ssize_t stop_result = read(stop_pipe[0], &stop, 1); + (void)stop_result; + (void)server.unregisterLocalMemory(buffer.data(), buffer.size()); + _exit(0); + } + + close(ready_pipe[1]); + close(stop_pipe[0]); + ChildProcessGuard child_guard(child, stop_pipe[1]); + + uint32_t segment_length = 0; + ssize_t received = + read(ready_pipe[0], &segment_length, sizeof(segment_length)); + if (received != static_cast(sizeof(segment_length))) { + const int status = child_guard.reap(); + GTEST_SKIP() << "RDMA server initialization failed, child status " + << status; + } + std::string server_segment(segment_length, '\0'); + ASSERT_EQ(read(ready_pipe[0], server_segment.data(), segment_length), + static_cast(segment_length)); + + TransferEngine client(makeRdmaConfig()); + ASSERT_TRUE(client.available()); + std::vector buffer(kBufferLength); + for (size_t i = 0; i < kDataLength; ++i) { + buffer[i] = static_cast((i * 31) & 0xff); + } + ASSERT_TRUE(client.registerLocalMemory(buffer.data(), buffer.size()).ok()); + + SegmentID segment = 0; + Status result; + for (int i = 0; i < 100; ++i) { + result = client.openSegment(segment, server_segment); + if (result.ok()) break; + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + ASSERT_TRUE(result.ok()) << result.ToString(); + + SegmentInfo info; + ASSERT_TRUE(client.getSegmentInfo(segment, info).ok()); + ASSERT_FALSE(info.buffers.empty()); + + Request request{}; + request.opcode = Request::WRITE; + request.source = buffer.data(); + request.target_id = segment; + request.target_offset = info.buffers[0].base; + request.length = kDataLength; + request.transport_hint = RDMA; + + BatchID batch = client.allocateBatch(1); + ASSERT_TRUE(client.submitTransfer(batch, {request}).ok()); + ASSERT_TRUE(waitBatchDone(client, batch)); + ASSERT_TRUE(client.freeBatch(batch).ok()); + + request.opcode = Request::READ; + request.source = buffer.data() + kDataLength; + batch = client.allocateBatch(1); + ASSERT_TRUE(client.submitTransfer(batch, {request}).ok()); + ASSERT_TRUE(waitBatchDone(client, batch)); + ASSERT_TRUE(client.freeBatch(batch).ok()); + EXPECT_EQ( + std::memcmp(buffer.data(), buffer.data() + kDataLength, kDataLength), + 0); + + // Keep one QP/worker and one outstanding WR so the tail task remains in + // the worker's unposted set long enough to exercise real cancellation. + std::vector cancel_requests; + cancel_requests.reserve(kCancelTaskCount); + for (size_t i = 0; i < kCancelTaskCount; ++i) { + Request cancel_request{}; + cancel_request.opcode = Request::WRITE; + cancel_request.source = buffer.data() + i * kCancelStride; + cancel_request.target_id = segment; + cancel_request.target_offset = info.buffers[0].base + i * kCancelStride; + cancel_request.length = kDataLength; + cancel_request.transport_hint = RDMA; + cancel_requests.push_back(cancel_request); + } + + batch = client.allocateBatch(kCancelTaskCount); + ASSERT_TRUE(client.submitTransfer(batch, cancel_requests).ok()); + const size_t cancel_task_id = kCancelTaskCount - 1; + ASSERT_TRUE(client.cancelTransfer(batch, cancel_task_id).ok()); + ASSERT_TRUE(client.cancelTransfer(batch, cancel_task_id).ok()); + + std::vector statuses; + for (int poll = 0; poll < 10000; ++poll) { + ASSERT_TRUE(client.getTransferStatus(batch, statuses).ok()); + if (std::all_of(statuses.begin(), statuses.end(), + [](const TransferStatus& task_status) { + return task_status.s != TransferStatusEnum::PENDING; + })) + break; + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_EQ(statuses.size(), kCancelTaskCount); + for (size_t i = 0; i < cancel_task_id; ++i) { + EXPECT_EQ(statuses[i].s, TransferStatusEnum::COMPLETED); + } + EXPECT_EQ(statuses[cancel_task_id].s, TransferStatusEnum::CANCELED); + EXPECT_LE(statuses[cancel_task_id].transferred_bytes, kDataLength); + ASSERT_TRUE(client.freeBatch(batch).ok()); + + EXPECT_TRUE(client.closeSegment(segment).ok()); + EXPECT_TRUE( + client.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); + + const int status = child_guard.finish(); + ASSERT_TRUE(WIFEXITED(status)); + EXPECT_EQ(WEXITSTATUS(status), 0); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp b/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp new file mode 100644 index 0000000000..33faa04a2c --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/receiver_credit_test.cpp @@ -0,0 +1,236 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include "tent/runtime/receiver_credit.h" + +#include +#include + +#include + +namespace mooncake::tent { +namespace { +CreditKey key() { return {{1, 2}, 3, 4}; } +ReceiverCreditUpdateV1 update(uint64_t epoch, uint64_t seq, + std::vector grants) { + ReceiverCreditUpdateV1 u; + u.receiver_session_id = key().receiver_session; + u.qos_class = key().qos_class; + u.epoch = epoch; + u.sequence = seq; + u.grants = std::move(grants); + return u; +} +CreditCharge charge(uint64_t bytes, uint64_t slots) { + return {{{CreditResource::DataBytes, bytes}, + {CreditResource::RequestSlots, slots}}}; +} +void grant(SenderCreditLedger& l, uint64_t seq, uint64_t bytes = 100, + uint64_t slots = 2) { + CreditUpdateDisposition d; + ASSERT_TRUE(l.applyUpdate(key(), + update(7, seq, + {{CreditResource::DataBytes, bytes}, + {CreditResource::RequestSlots, slots}}), + d) + .ok()); +} + +TEST(ReceiverCredit, MultiResourceReserveIsAtomic) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + ASSERT_TRUE(l.tryReserve(key(), charge(60, 1)).ok()); + EXPECT_TRUE(l.tryReserve(key(), charge(41, 1)).IsTooManyRequests()); + uint64_t v; + ASSERT_TRUE(l.available(key(), CreditResource::RequestSlots, v).ok()); + EXPECT_EQ(v, 1); // failed byte reservation did not consume a slot +} + +TEST(ReceiverCredit, DuplicateAndReorderedUpdatesCannotMintCredit) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 2); + CreditUpdateDisposition d; + ASSERT_TRUE(l.applyUpdate( + key(), update(7, 2, {{CreditResource::DataBytes, 999}}), d) + .ok()); + EXPECT_EQ(d, CreditUpdateDisposition::DuplicateOrOld); + ASSERT_TRUE(l.applyUpdate( + key(), update(7, 1, {{CreditResource::DataBytes, 999}}), d) + .ok()); + uint64_t v; + ASSERT_TRUE(l.available(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 100); +} + +TEST(ReceiverCredit, SequenceGapIsVisibleAndSafe) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + CreditUpdateDisposition d; + ASSERT_TRUE(l.applyUpdate( + key(), update(7, 4, {{CreditResource::DataBytes, 120}}), d) + .ok()); + EXPECT_EQ(d, CreditUpdateDisposition::SequenceGap); +} + +TEST(ReceiverCredit, PartialGrantUpdateRetainsOmittedResources) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1, 100, 5); + + CreditUpdateDisposition d; + ASSERT_TRUE(l.applyUpdate( + key(), update(7, 2, {{CreditResource::DataBytes, 160}}), d) + .ok()); + EXPECT_EQ(d, CreditUpdateDisposition::Applied); + + uint64_t v; + ASSERT_TRUE(l.available(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 160); + ASSERT_TRUE(l.available(key(), CreditResource::RequestSlots, v).ok()); + EXPECT_EQ(v, 5); +} + +TEST(ReceiverCredit, StaleEpochFailsAndActivationFencesOldState) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + ASSERT_TRUE(l.tryReserve(key(), charge(60, 1)).ok()); + CreditUpdateDisposition d; + EXPECT_TRUE(l.applyUpdate( + key(), update(6, 2, {{CreditResource::DataBytes, 999}}), d) + .IsInvalidEntry()); + ASSERT_TRUE(l.activate(key(), 8).ok()); + uint64_t v; + EXPECT_TRUE( + l.available(key(), CreditResource::DataBytes, v).IsInvalidEntry()); +} + +TEST(ReceiverCredit, ActivationReplayCannotMintCredit) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + ASSERT_TRUE(l.tryReserve(key(), charge(80, 1)).ok()); + ASSERT_TRUE(l.activate(key(), 7).ok()); // idempotent, not a reset + uint64_t v; + ASSERT_TRUE(l.consumed(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 80); + EXPECT_TRUE(l.activate(key(), 6).IsInvalidEntry()); + ASSERT_TRUE(l.consumed(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 80); +} + +TEST(ReceiverCredit, LedgerEntryCountIsBounded) { + SenderCreditLedger l(1); + ASSERT_TRUE(l.activate(key(), 7).ok()); + auto other = key(); + ++other.sender_peer; + EXPECT_TRUE(l.activate(other, 7).IsTooManyRequests()); + // A new epoch for an existing key does not consume another entry. + EXPECT_TRUE(l.activate(key(), 8).ok()); +} + +TEST(ReceiverCredit, DeactivationReleasesCapacityAfterExactEpochFence) { + SenderCreditLedger l(1); + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + ASSERT_TRUE(l.tryReserve(key(), charge(80, 1)).ok()); + auto other = key(); + ++other.sender_peer; + EXPECT_TRUE(l.activate(other, 1).IsTooManyRequests()); + + EXPECT_TRUE(l.deactivate(key(), 6).IsInvalidEntry()); + EXPECT_TRUE(l.activate(other, 1).IsTooManyRequests()); + ASSERT_TRUE(l.deactivate(key(), 7).ok()); + ASSERT_TRUE(l.deactivate(key(), 7).ok()); // cleanup is idempotent + EXPECT_TRUE(l.activate(other, 1).ok()); +} + +TEST(ReceiverCredit, OldCleanupCannotEraseReactivatedEpoch) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + ASSERT_TRUE(l.activate(key(), 8).ok()); + EXPECT_TRUE(l.deactivate(key(), 7).IsInvalidEntry()); + CreditUpdateDisposition disposition; + auto fresh = update(8, 1, {{CreditResource::DataBytes, 10}}); + ASSERT_TRUE(l.applyUpdate(key(), fresh, disposition).ok()); +} + +TEST(ReceiverCredit, InvalidUpdateDoesNotPartiallyMutate) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + CreditUpdateDisposition d; + EXPECT_TRUE(l.applyUpdate(key(), + update(7, 2, + {{CreditResource::DataBytes, 200}, + {CreditResource::RequestSlots, 1}}), + d) + .IsInvalidArgument()); + uint64_t v; + ASSERT_TRUE(l.available(key(), CreditResource::DataBytes, v).ok()); + EXPECT_EQ(v, 100); +} + +TEST(ReceiverCredit, DuplicateUnknownAndZeroResourcesFailClosed) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + CreditUpdateDisposition d; + EXPECT_TRUE(l.applyUpdate(key(), + update(7, 1, + {{CreditResource::DataBytes, 1}, + {CreditResource::DataBytes, 2}}), + d) + .IsInvalidArgument()); + EXPECT_TRUE(l.tryReserve(key(), {{{static_cast(99), 1}}}) + .IsInvalidArgument()); + EXPECT_TRUE(l.tryReserve(key(), {{{CreditResource::DataBytes, 0}}}) + .IsInvalidArgument()); +} + +TEST(ReceiverCredit, RollbackChecksUnderflowAtomically) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + ASSERT_TRUE(l.tryReserve(key(), charge(60, 1)).ok()); + EXPECT_TRUE( + l.rollbackReservation(key(), charge(61, 1)).IsInvalidArgument()); + uint64_t v; + ASSERT_TRUE(l.consumed(key(), CreditResource::RequestSlots, v).ok()); + EXPECT_EQ(v, 1); + ASSERT_TRUE(l.rollbackReservation(key(), charge(60, 1)).ok()); +} + +TEST(ReceiverCredit, GrantCannotDecreaseOrFallBelowConsumption) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1); + ASSERT_TRUE(l.tryReserve(key(), charge(80, 1)).ok()); + CreditUpdateDisposition d; + EXPECT_TRUE( + l.applyUpdate(key(), update(7, 2, {{CreditResource::DataBytes, 79}}), d) + .IsInvalidArgument()); +} + +TEST(ReceiverCredit, ConcurrentReservationsNeverExceedGrant) { + SenderCreditLedger l; + ASSERT_TRUE(l.activate(key(), 7).ok()); + grant(l, 1, 100, 100); + std::atomic admitted{0}; + std::vector threads; + for (int i = 0; i < 16; ++i) { + threads.emplace_back([&] { + for (int j = 0; j < 20; ++j) + if (l.tryReserve(key(), charge(1, 1)).ok()) ++admitted; + }); + } + for (auto& thread : threads) thread.join(); + EXPECT_EQ(admitted, 100); + uint64_t consumed; + ASSERT_TRUE(l.consumed(key(), CreditResource::DataBytes, consumed).ok()); + EXPECT_EQ(consumed, 100); +} +} // namespace +} // namespace mooncake::tent diff --git a/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp new file mode 100644 index 0000000000..b164b642a3 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/runtime_queue_dispatch_test.cpp @@ -0,0 +1,659 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/common/types.h" +#include "tent/runtime/segment.h" +#include "tent/runtime/transfer_engine_impl.h" +#include "tent/runtime/transport.h" + +namespace mooncake { +namespace tent { +namespace { + +class FakeSubBatch : public Transport::SubBatch { + public: + size_t size() const override { return task_count; } + + size_t task_count = 0; + std::vector requests; + std::vector statuses; + std::vector poll_counts; +}; + +class FakeTransport : public Transport { + public: + using PollStatusFactory = + std::function; + + explicit FakeTransport(TransportType self_type, + PollStatusFactory poll_status_factory = {}, + bool notify_on_submit = false) + : self_type_(self_type), + poll_status_factory_(std::move(poll_status_factory)), + notify_on_submit_(notify_on_submit) { + caps.dram_to_dram = true; + } + + std::atomic submit_calls{0}; + std::atomic status_calls{0}; + std::atomic cancel_calls{0}; + bool cancellation_supported{true}; + + Status install(std::string&, std::shared_ptr, + std::shared_ptr, + std::shared_ptr = nullptr) override { + return Status::OK(); + } + + Status allocateSubBatch(SubBatchRef& batch, size_t) override { + batch = new FakeSubBatch(); + return Status::OK(); + } + + Status freeSubBatch(SubBatchRef& batch) override { + delete static_cast(batch); + batch = nullptr; + return Status::OK(); + } + + Status submitTransferTasks(SubBatchRef batch, + const std::vector& requests) override { + ++submit_calls; + auto* fake = static_cast(batch); + for (const auto& request : requests) { + fake->requests.push_back(request); + fake->statuses.push_back( + {TransferStatusEnum::COMPLETED, request.length}); + fake->poll_counts.push_back(0); + ++fake->task_count; + } + if (notify_on_submit_) batch->notifyProgress(); + return Status::OK(); + } + + Status getTransferStatus(SubBatchRef batch, int task_id, + TransferStatus& status) override { + ++status_calls; + auto* fake = static_cast(batch); + if (task_id < 0 || task_id >= (int)fake->statuses.size()) { + return Status::InvalidArgument("bad task_id" LOC_MARK); + } + ++fake->poll_counts[task_id]; + if (poll_status_factory_) { + status = poll_status_factory_(fake->requests[task_id], + fake->poll_counts[task_id]); + } else { + status = fake->statuses[task_id]; + } + return Status::OK(); + } + + bool supportsCancellation() const override { + return cancellation_supported; + } + + Status cancelTransferTask(SubBatchRef batch, int task_id) override { + auto* fake = static_cast(batch); + if (task_id < 0 || task_id >= (int)fake->statuses.size()) { + return Status::InvalidArgument("bad task_id" LOC_MARK); + } + if (!cancellation_supported) { + return Status::NotImplemented("cancel unsupported" LOC_MARK); + } + ++cancel_calls; + fake->statuses[task_id] = {TransferStatusEnum::CANCELED, 0}; + return Status::OK(); + } + + Status addMemoryBuffer(BufferDesc& desc, const MemoryOptions&) override { + desc.transports.push_back(self_type_); + return Status::OK(); + } + + Status addMemoryBuffer(std::vector& desc_list, + const MemoryOptions& options) override { + for (auto& desc : desc_list) { + auto status = addMemoryBuffer(desc, options); + if (!status.ok()) return status; + } + return Status::OK(); + } + + Status removeMemoryBuffer(BufferDesc&) override { return Status::OK(); } + + Status allocateLocalMemory(void** addr, size_t size, + MemoryOptions&) override { + *addr = std::malloc(size); + return *addr ? Status::OK() + : Status::InternalError("malloc failed" LOC_MARK); + } + + Status freeLocalMemory(void* addr, size_t) override { + std::free(addr); + return Status::OK(); + } + + bool warmupMemory(void*, size_t) override { return false; } + + const char* getName() const override { return ""; } + + private: + TransportType self_type_; + PollStatusFactory poll_status_factory_; + bool notify_on_submit_; +}; + +std::shared_ptr makeRuntimeQueueConfig(size_t max_dispatch_owners, + size_t max_dispatch_bytes, + bool merge_requests = false) { + auto cfg = std::make_shared(); + cfg->set("metadata_type", "p2p"); + cfg->set("metadata_servers", ""); + cfg->set("rpc_server_hostname", "127.0.0.1"); + cfg->set("rpc_server_port", "0"); + cfg->set("log_level", "warning"); + cfg->set("merge_requests", merge_requests); + cfg->set("enable_runtime_queue", true); + cfg->set("runtime_queue/max_outstanding_owners", 16UL); + cfg->set("runtime_queue/max_outstanding_bytes", 1UL << 20); + cfg->set("runtime_queue/max_dispatch_owners", max_dispatch_owners); + cfg->set("runtime_queue/max_dispatch_bytes", max_dispatch_bytes); + cfg->set("runtime_queue/staging_owner_reserve", 0UL); + cfg->set("runtime_queue/staging_byte_reserve", 0UL); + cfg->set("runtime_queue/progress_fallback_interval_us", 50000UL); + + cfg->set("transports/tcp/enable", false); + cfg->set("transports/shm/enable", false); + cfg->set("transports/rdma/enable", false); + cfg->set("transports/io_uring/enable", false); + cfg->set("transports/nvlink/enable", false); + cfg->set("transports/mnnvl/enable", false); + cfg->set("transports/gds/enable", false); + cfg->set("transports/ascend_direct/enable", false); + return cfg; +} + +void installFakeRdma(TransferEngineImpl& engine, + const std::shared_ptr& fake_rdma) { + std::string seg_name = engine.getSegmentName(); + ASSERT_TRUE(fake_rdma->install(seg_name, nullptr, nullptr).ok()); + engine.swapTransportForTest(RDMA, fake_rdma); +} + +Request makeLocalWrite(uint8_t* ptr, size_t length) { + Request request; + request.opcode = Request::WRITE; + request.source = ptr; + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = reinterpret_cast(ptr); + request.length = length; + request.transport_hint = RDMA; + return request; +} + +TEST(RuntimeQueueDispatch, RejectsOverfullBatchBeforePublishingTasks) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen * 2, 0x11); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + + auto status = engine.submitTransfer( + batch, {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}); + EXPECT_TRUE(status.IsTooManyRequests()) << status.ToString(); + EXPECT_EQ(fake_rdma->submit_calls.load(), 0); + + TransferStatus task_status{}; + EXPECT_TRUE( + engine.getTransferStatus(batch, 0, task_status).IsInvalidArgument()); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RejectsOwnerLargerThanDispatchByteWindow) { + constexpr size_t kReqLen = 4096; + auto cfg = makeRuntimeQueueConfig(1, kReqLen - 1); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + std::vector buffer(kReqLen, 0x77); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + + auto status = + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}); + EXPECT_TRUE(status.IsTooManyRequests()) << status.ToString(); + EXPECT_EQ(fake_rdma->submit_calls.load(), 0); + + TransferStatus task_status{}; + EXPECT_TRUE( + engine.getTransferStatus(batch, 0, task_status).IsInvalidArgument()); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RejectsEmptyDispatchWindowConfig) { + auto cfg = makeRuntimeQueueConfig(0, 1UL << 20); + TransferEngineImpl engine(cfg); + + EXPECT_FALSE(engine.available()); + + cfg = makeRuntimeQueueConfig(1, 0); + TransferEngineImpl byte_window_engine(cfg); + + EXPECT_FALSE(byte_window_engine.available()); +} + +TEST(RuntimeQueueDispatch, DispatchesOnlyOneWindowOnSubmit) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen * 2, 0x22); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + TransferStatus first{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_rdma->submit_calls.load(), 2); + + TransferStatus second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, KeepsDispatchWindowUntilOwnerIsTerminal) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + std::atomic complete_first{false}; + auto fake_rdma = std::make_shared( + RDMA, [&complete_first](const Request& request, int) { + if (!complete_first.load()) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen * 2, 0x55); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + TransferStatus first{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::PENDING); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + TransferStatus second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::PENDING); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + complete_first.store(true); + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_rdma->submit_calls.load(), 2); + + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CancelsQueuedOwnerWithoutDispatchingIt) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + std::atomic complete_first{false}; + auto fake_rdma = std::make_shared( + RDMA, [&complete_first](const Request& request, int) { + if (!complete_first.load()) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen * 2, 0x5a); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + ASSERT_EQ(fake_rdma->submit_calls.load(), 1); + + ASSERT_TRUE(engine.cancelTransfer(batch, 1).ok()); + EXPECT_EQ(fake_rdma->cancel_calls.load(), 0); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + TransferStatus second{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 1, second).ok()); + EXPECT_EQ(second.s, TransferStatusEnum::CANCELED); + + complete_first.store(true); + TransferStatus first{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, first).ok()); + EXPECT_EQ(first.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, CancelsDispatchedRdmaTaskIdempotently) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen, 0x6b); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + + ASSERT_TRUE(engine.cancelTransfer(batch, 0).ok()); + EXPECT_EQ(fake_rdma->cancel_calls.load(), 1); + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::CANCELED); + ASSERT_TRUE(engine.cancelTransfer(batch, 0).ok()); + EXPECT_EQ(fake_rdma->cancel_calls.load(), 1); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RejectsCancellationForUnsupportedTransport) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA); + fake_rdma->cancellation_supported = false; + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen, 0x7c); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + + EXPECT_TRUE(engine.cancelTransfer(batch, 0).IsNotImplemented()); + EXPECT_EQ(fake_rdma->cancel_calls.load(), 0); + + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED); + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, ProgressWorkerRefillsWindowFromTransportNotify) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("enable_progress_worker", true); + cfg->set("runtime_queue/progress_fallback_interval_us", 0UL); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared( + RDMA, FakeTransport::PollStatusFactory{}, true); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen * 2, 0x66); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::milliseconds(1000); + while (std::chrono::steady_clock::now() < deadline && + fake_rdma->submit_calls.load() < 2) { + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + EXPECT_EQ(fake_rdma->submit_calls.load(), 2); + + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, RuntimeQueueDrainsWithoutUserPolling) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + cfg->set("runtime_queue/progress_fallback_interval_us", 1000UL); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared( + RDMA, [](const Request& request, int poll_count) { + if (poll_count == 1) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen * 2, 0x88); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::milliseconds(1000); + while (std::chrono::steady_clock::now() < deadline && + (fake_rdma->submit_calls.load() < 2 || + fake_rdma->status_calls.load() < 4)) { + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + EXPECT_EQ(fake_rdma->submit_calls.load(), 2); + EXPECT_GE(fake_rdma->status_calls.load(), 4); + + TransferStatus status{}; + ASSERT_TRUE(engine.getTransferStatus(batch, status).ok()); + EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, EarlyFreeReclaimsAfterQueuedCompletion) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared( + RDMA, [](const Request& request, int poll_count) { + if (poll_count == 1) { + return TransferStatus{TransferStatusEnum::PENDING, 0}; + } + return TransferStatus{TransferStatusEnum::COMPLETED, + request.length}; + }); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen, 0x33); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + + BatchID batch = engine.allocateBatch(1); + ASSERT_NE(batch, (BatchID)0); + ASSERT_TRUE( + engine.submitTransfer(batch, {makeLocalWrite(buffer.data(), kReqLen)}) + .ok()); + + ASSERT_TRUE(engine.freeBatch(batch).ok()); + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::milliseconds(1000); + while (std::chrono::steady_clock::now() < deadline && + fake_rdma->status_calls.load() < 2) { + std::this_thread::sleep_for(std::chrono::milliseconds(2)); + } + EXPECT_GE(fake_rdma->status_calls.load(), 2); + + TransferStatus status{}; + EXPECT_TRUE(engine.getTransferStatus(batch, 0, status).IsInvalidArgument()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +TEST(RuntimeQueueDispatch, PollingDerivedTaskCompletesMergedOwner) { + auto cfg = makeRuntimeQueueConfig(1, 1UL << 20, true); + TransferEngineImpl engine(cfg); + ASSERT_TRUE(engine.available()); + + auto fake_rdma = std::make_shared(RDMA); + installFakeRdma(engine, fake_rdma); + + constexpr size_t kReqLen = 4096; + std::vector buffer(kReqLen * 2, 0x44); + ASSERT_TRUE(engine.registerLocalMemory(buffer.data(), buffer.size()).ok()); + + BatchID batch = engine.allocateBatch(2); + ASSERT_NE(batch, (BatchID)0); + + ASSERT_TRUE( + engine + .submitTransfer(batch, + {makeLocalWrite(buffer.data(), kReqLen), + makeLocalWrite(buffer.data() + kReqLen, kReqLen)}) + .ok()); + EXPECT_EQ(fake_rdma->submit_calls.load(), 1); + + TransferStatus derived{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 1, derived).ok()); + EXPECT_EQ(derived.s, TransferStatusEnum::COMPLETED); + + TransferStatus owner{}; + ASSERT_TRUE(engine.getTransferStatus(batch, 0, owner).ok()); + EXPECT_EQ(owner.s, TransferStatusEnum::COMPLETED); + + EXPECT_TRUE(engine.freeBatch(batch).ok()); + EXPECT_TRUE( + engine.unregisterLocalMemory(buffer.data(), buffer.size()).ok()); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/rw_spinlock_test.cpp b/mooncake-transfer-engine/tent/tests/rw_spinlock_test.cpp new file mode 100644 index 0000000000..1876907397 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/rw_spinlock_test.cpp @@ -0,0 +1,88 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include + +#include "tent/common/concurrent/rw_spinlock.h" + +namespace mooncake { +namespace tent { +namespace { + +using namespace std::chrono_literals; + +TEST(RWSpinlockTest, AggressiveWriteLockProgressesAcrossTicketWraparound) { + RWSpinlock lock; + int protected_value = 0; + + constexpr int kIterations = 65536 + 3; + for (int i = 0; i < kIterations; ++i) { + lock.writeLockAggressive(); + ++protected_value; + lock.unlock(); + } + + RWSpinlock::ReadGuard guard(lock); + EXPECT_EQ(protected_value, kIterations); +} + +TEST(RWSpinlockTest, DowngradePublishesToReadersAndBlocksWriters) { + RWSpinlock lock; + int protected_value = 0; + std::atomic writer_started{false}; + std::atomic writer_entered{false}; + std::atomic reader_observed{false}; + + lock.writeLockAggressive(); + protected_value = 42; + lock.unlockAndLockShared(); + + std::thread reader([&] { + RWSpinlock::ReadGuard guard(lock); + reader_observed.store(protected_value == 42, std::memory_order_release); + }); + + reader.join(); + EXPECT_TRUE(reader_observed.load(std::memory_order_acquire)); + + std::thread writer([&] { + writer_started.store(true, std::memory_order_release); + lock.writeLockAggressive(); + writer_entered.store(true, std::memory_order_release); + protected_value = 99; + lock.unlock(); + }); + + while (!writer_started.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + + std::this_thread::sleep_for(10ms); + EXPECT_FALSE(writer_entered.load(std::memory_order_acquire)); + + lock.unlockShared(); + writer.join(); + + RWSpinlock::ReadGuard guard(lock); + EXPECT_TRUE(writer_entered.load(std::memory_order_acquire)); + EXPECT_EQ(protected_value, 99); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/segment_manager_test.cpp b/mooncake-transfer-engine/tent/tests/segment_manager_test.cpp new file mode 100644 index 0000000000..77bb31d7fc --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/segment_manager_test.cpp @@ -0,0 +1,309 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Regression tests for issue #2477: concurrent registerLocalMemory / +// unregisterLocalMemory racing with lock-free readers of the local +// SegmentDesc. The local desc is published as immutable copy-on-write +// snapshots; these tests assert the snapshot semantics (readers never +// observe torn or unsorted buffer lists) and, when built with +// -fsanitize=thread, additionally prove the absence of data races between +// SegmentTracker writers and getLocal() / getLocalDumpedJson() readers. +// A port of ConcurrentWritersVsSnapshotReaders to the pre-fix in-place +// mutation API fails its invariants on stock builds and reports data races +// under TSAN. + +#include + +#include +#include +#include +#include +#include + +#include "tent/runtime/segment_manager.h" +#include "tent/runtime/segment_registry.h" +#include "tent/runtime/segment_tracker.h" + +namespace mooncake { +namespace tent { + +namespace { + +std::unique_ptr makeManager() { + // No registry: these tests never touch remote segments or + // synchronizeLocal(). + auto manager = std::make_unique(nullptr); + EXPECT_TRUE(manager + ->updateLocal([](SegmentDesc& desc) -> Status { + desc.name = "local_test_segment"; + desc.type = SegmentType::Memory; + desc.machine_id = "test_machine"; + return Status::OK(); + }) + .ok()); + return manager; +} + +// Location derived from the buffer address; readers use it to detect +// partially-constructed or torn BufferDesc entries. +std::string locationFor(uint64_t addr) { + return "cpu:" + std::to_string(addr % 4096); +} + +BufferDesc makeBuffer(uint64_t addr, uint64_t length) { + BufferDesc desc; + desc.addr = addr; + desc.length = length; + desc.location = locationFor(addr); + desc.ref_count = 1; + return desc; +} + +const std::vector& buffersOf(const SegmentDescRef& snapshot) { + return std::get(snapshot->detail).buffers; +} + +} // namespace + +TEST(SegmentManagerTest, UpdateLocalPublishesImmutableSnapshots) { + auto manager = makeManager(); + auto before = manager->getLocal(); + ASSERT_EQ(before->name, "local_test_segment"); + ASSERT_TRUE(buffersOf(before).empty()); + + ASSERT_TRUE(manager + ->updateLocal([](SegmentDesc& desc) -> Status { + auto& detail = std::get(desc.detail); + detail.buffers.push_back(makeBuffer(0x1000, 0x1000)); + return Status::OK(); + }) + .ok()); + + // The old snapshot is untouched; the new snapshot sees the mutation. + EXPECT_TRUE(buffersOf(before).empty()); + auto after = manager->getLocal(); + ASSERT_EQ(buffersOf(after).size(), 1u); + EXPECT_EQ(buffersOf(after)[0].addr, 0x1000u); + EXPECT_NE(before.get(), after.get()); +} + +TEST(SegmentManagerTest, UpdateLocalFailureDoesNotPublish) { + auto manager = makeManager(); + auto before = manager->getLocal(); + auto status = manager->updateLocal([](SegmentDesc& desc) -> Status { + desc.name = "must_not_be_published"; + return Status::InvalidArgument("injected failure"); + }); + EXPECT_FALSE(status.ok()); + EXPECT_EQ(manager->getLocal().get(), before.get()); + EXPECT_EQ(manager->getLocal()->name, "local_test_segment"); +} + +TEST(SegmentManagerTest, JsonCacheInvalidatedOnPublication) { + auto manager = makeManager(); + auto dump1 = manager->getLocalDumpedJson(); + auto dump2 = manager->getLocalDumpedJson(); + EXPECT_EQ(dump1.get(), dump2.get()); // served from cache + + ASSERT_TRUE(manager + ->updateLocal([](SegmentDesc& desc) -> Status { + auto& detail = std::get(desc.detail); + detail.buffers.push_back(makeBuffer(0x2000, 0x1000)); + return Status::OK(); + }) + .ok()); + + auto dump3 = manager->getLocalDumpedJson(); + EXPECT_NE(*dump1, *dump3); + EXPECT_NE(dump3->find("8192"), std::string::npos); // 0x2000 serialized +} + +TEST(SegmentTrackerTest, RefCountedAddRemove) { + auto manager = makeManager(); + SegmentTracker tracker(*manager); + + auto noop = [](std::vector&) -> Status { return Status::OK(); }; + std::vector first{makeBuffer(0x10000, 0x1000)}; + ASSERT_TRUE(tracker.addInBatch(first, noop).ok()); + std::vector second{makeBuffer(0x10000, 0x1000)}; + ASSERT_TRUE(tracker.addInBatch(second, noop).ok()); + + auto snapshot = manager->getLocal(); + ASSERT_EQ(buffersOf(snapshot).size(), 1u); + EXPECT_EQ(buffersOf(snapshot)[0].ref_count, 2); + + int remove_callbacks = 0; + auto on_remove = [&](BufferDesc&) -> Status { + remove_callbacks++; + return Status::OK(); + }; + ASSERT_TRUE(tracker.remove(0x10000, 0x1000, on_remove).ok()); + EXPECT_EQ(remove_callbacks, 0); // still referenced + ASSERT_EQ(buffersOf(manager->getLocal()).size(), 1u); + + ASSERT_TRUE(tracker.remove(0x10000, 0x1000, on_remove).ok()); + EXPECT_EQ(remove_callbacks, 1); + EXPECT_TRUE(buffersOf(manager->getLocal()).empty()); +} + +TEST(SegmentTrackerTest, AddInBatchCallbackFailureRollsBackRefCounts) { + auto manager = makeManager(); + SegmentTracker tracker(*manager); + + auto ok = [](std::vector&) -> Status { return Status::OK(); }; + std::vector first{makeBuffer(0x20000, 0x1000)}; + ASSERT_TRUE(tracker.addInBatch(first, ok).ok()); + ASSERT_EQ(buffersOf(manager->getLocal())[0].ref_count, 1); + + // A duplicate registration whose transport callback fails must not leave + // the ref-count bumped, or the buffer could never be deregistered. + auto fail = [](std::vector&) -> Status { + return Status::InternalError("injected transport failure"); + }; + std::vector dup{makeBuffer(0x20000, 0x1000)}; + EXPECT_FALSE(tracker.addInBatch(dup, fail).ok()); + ASSERT_EQ(buffersOf(manager->getLocal()).size(), 1u); + EXPECT_EQ(buffersOf(manager->getLocal())[0].ref_count, 1); + + int remove_callbacks = 0; + auto on_remove = [&](BufferDesc&) -> Status { + remove_callbacks++; + return Status::OK(); + }; + ASSERT_TRUE(tracker.remove(0x20000, 0x1000, on_remove).ok()); + EXPECT_EQ(remove_callbacks, 1); + EXPECT_TRUE(buffersOf(manager->getLocal()).empty()); +} + +TEST(SegmentTrackerTest, AddProbesRealMemoryAndRefCounts) { + auto manager = makeManager(); + SegmentTracker tracker(*manager); + + constexpr size_t kSize = 1 << 20; + void* mem = malloc(kSize); + ASSERT_NE(mem, nullptr); + auto base = reinterpret_cast(mem); + + int add_callbacks = 0; + auto on_add = [&](BufferDesc& desc) -> Status { + add_callbacks++; + EXPECT_EQ(desc.addr, base); + EXPECT_FALSE(desc.location.empty()); // NUMA probe ran + return Status::OK(); + }; + ASSERT_TRUE(tracker.add(base, kSize, on_add).ok()); + EXPECT_EQ(add_callbacks, 1); + // Re-adding the same range takes the ref-count fast path: no new probe. + ASSERT_TRUE(tracker.add(base, kSize, on_add).ok()); + EXPECT_EQ(add_callbacks, 1); + ASSERT_EQ(buffersOf(manager->getLocal()).size(), 1u); + EXPECT_EQ(buffersOf(manager->getLocal())[0].ref_count, 2); + + auto noop = [](BufferDesc&) -> Status { return Status::OK(); }; + ASSERT_TRUE(tracker.remove(base, kSize, noop).ok()); + ASSERT_TRUE(tracker.remove(base, kSize, noop).ok()); + EXPECT_TRUE(buffersOf(manager->getLocal()).empty()); + free(mem); +} + +// The actual #2477 regression: register/unregister churn concurrent with +// lock-free snapshot readers. With in-place mutation this is a data race on +// MemorySegmentDesc::buffers (vector push_back/erase/sort vs. iteration): +// ThreadSanitizer reports it and the location invariant below can observe +// partially-constructed entries. With copy-on-write snapshots every reader +// observes a fully-consistent (possibly stale) buffer list. +TEST(SegmentTrackerTest, ConcurrentWritersVsSnapshotReaders) { + auto manager = makeManager(); + SegmentTracker tracker(*manager); + + constexpr int kWriters = 4; + constexpr int kReaders = 4; + constexpr int kIterations = 300; + constexpr int kBuffersPerBatch = 8; + constexpr uint64_t kLength = 0x1000; + + std::atomic done{false}; + std::atomic failures{0}; + + auto noop = [](std::vector&) -> Status { return Status::OK(); }; + + std::vector writers; + writers.reserve(kWriters); + for (int w = 0; w < kWriters; ++w) { + writers.emplace_back([&, w] { + // Writers 0 and 1 share an address range so ref-count bumps, + // duplicate-registration races and erase-vs-bump interleavings + // are exercised concurrently, not just disjoint inserts. + const uint64_t base = (w < 2 ? 1 : w + 1) * 0x100000000ULL; + for (int iter = 0; iter < kIterations; ++iter) { + std::vector batch; + batch.reserve(kBuffersPerBatch); + for (int i = 0; i < kBuffersPerBatch; ++i) { + batch.push_back( + makeBuffer(base + i * kLength * 2, kLength)); + } + if (!tracker.addInBatch(batch, noop).ok()) failures++; + for (int i = 0; i < kBuffersPerBatch; ++i) { + auto on_remove = [](BufferDesc&) -> Status { + return Status::OK(); + }; + if (!tracker + .remove(base + i * kLength * 2, kLength, on_remove) + .ok()) + failures++; + } + } + }); + } + + std::vector readers; + readers.reserve(kReaders); + for (int r = 0; r < kReaders; ++r) { + readers.emplace_back([&, r] { + uint64_t rounds = 0; + while (!done.load(std::memory_order_acquire)) { + auto snapshot = manager->getLocal(); + const auto& buffers = buffersOf(snapshot); + uint64_t prev_addr = 0; + for (const auto& buf : buffers) { + // Entries must be fully constructed and sorted; a torn + // read of an in-place mutated vector violates these. + if (buf.length != kLength || + buf.location != locationFor(buf.addr) || + buf.addr < prev_addr) { + failures++; + } + prev_addr = buf.addr; + } + // Exercise findBuffer() through the snapshot as transports + // do, and the JSON dump path peers hit via GetSegmentDesc. + snapshot->findBuffer(0x100000000ULL, kLength); + if (rounds++ % 64 == 0) { + auto dump = manager->getLocalDumpedJson(); + if (!dump || dump->empty()) failures++; + } + } + }); + } + + for (auto& t : writers) t.join(); + done.store(true, std::memory_order_release); + for (auto& t : readers) t.join(); + + EXPECT_EQ(failures.load(), 0); + EXPECT_TRUE(buffersOf(manager->getLocal()).empty()); +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/shm_transport_test.cpp b/mooncake-transfer-engine/tent/tests/shm_transport_test.cpp new file mode 100644 index 0000000000..bf8f99ba01 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/shm_transport_test.cpp @@ -0,0 +1,371 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tent/common/config.h" +#include "tent/runtime/control_plane.h" +#include "tent/transport/shm/shm_transport.h" + +namespace mooncake { +namespace tent { + +class ShmTransportTestPeer { + public: + static Status relocate(ShmTransport& transport, uint64_t& address, + uint64_t length, SegmentID target_id) { + return transport.relocateSharedMemoryAddress(address, length, + target_id); + } + + static size_t mappingCount(ShmTransport& transport, SegmentID target_id) { + RWSpinlock::ReadGuard guard(transport.relocate_lock_); + auto it = transport.relocate_map_.find(target_id); + return it == transport.relocate_map_.end() ? 0 : it->second.size(); + } + + static bool hasTarget(ShmTransport& transport, SegmentID target_id) { + RWSpinlock::ReadGuard guard(transport.relocate_lock_); + return transport.relocate_map_.find(target_id) != + transport.relocate_map_.end(); + } + + static void* createSharedMemory(ShmTransport& transport, + const std::string& path, size_t size) { + return transport.createSharedMemory(path, size); + } +}; + +namespace { + +class ScopedShmFile { + public: + explicit ScopedShmFile(size_t length) + : name_("/mooncake_tent_shm_test_" + std::to_string(getpid())), + length_(length) { + shm_unlink(name_.c_str()); + fd_ = shm_open(name_.c_str(), O_CREAT | O_EXCL | O_RDWR, 0600); + EXPECT_GE(fd_, 0); + if (fd_ >= 0) { + EXPECT_EQ(ftruncate(fd_, length_), 0); + } + } + + ~ScopedShmFile() { + if (fd_ >= 0) close(fd_); + shm_unlink(name_.c_str()); + } + + const std::string& name() const { return name_; } + + private: + std::string name_; + size_t length_; + int fd_{-1}; +}; + +Status installLocalSegmentWithShm(ControlService& metadata, + const std::string& shm_path, + uint64_t remote_addr, size_t length) { + return metadata.segmentManager().updateLocal( + [&](SegmentDesc& segment) -> Status { + segment.name = "shm_test_segment"; + segment.machine_id = "shm_test_machine"; + segment.type = SegmentType::Memory; + auto& memory = std::get(segment.detail); + memory.buffers.clear(); + BufferDesc buffer; + buffer.addr = remote_addr; + buffer.length = length; + buffer.location = "cpu:0"; + buffer.shm_path = shm_path; + memory.buffers.push_back(std::move(buffer)); + return Status::OK(); + }); +} + +TEST(ShmTransportTest, SharesAndReleasesRelocationAcrossThreads) { + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + constexpr uint64_t kRemoteAddress = 0x10000000; + constexpr size_t kThreadCount = 8; + ScopedShmFile shm_file(page_size); + + auto metadata = std::make_shared("p2p", "", nullptr); + ASSERT_TRUE(installLocalSegmentWithShm(*metadata, shm_file.name(), + kRemoteAddress, page_size) + .ok()); + + ShmTransport transport; + std::string local_segment_name = "shm_test_segment"; + ASSERT_TRUE(transport + .install(local_segment_name, metadata, nullptr, + std::make_shared()) + .ok()); + + uint64_t missing_address = kRemoteAddress + page_size; + EXPECT_TRUE(ShmTransportTestPeer::relocate(transport, missing_address, + page_size, LOCAL_SEGMENT_ID) + .IsNeedsRefreshCache()); + EXPECT_FALSE(ShmTransportTestPeer::hasTarget(transport, LOCAL_SEGMENT_ID)); + + std::atomic ready{0}; + std::atomic start{false}; + std::vector relocated(kThreadCount, kRemoteAddress); + std::vector succeeded(kThreadCount, 0); + std::vector threads; + threads.reserve(kThreadCount); + for (size_t i = 0; i < kThreadCount; ++i) { + threads.emplace_back([&, i] { + ready.fetch_add(1, std::memory_order_release); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + succeeded[i] = + ShmTransportTestPeer::relocate(transport, relocated[i], + page_size, LOCAL_SEGMENT_ID) + .ok(); + }); + } + while (ready.load(std::memory_order_acquire) != kThreadCount) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (auto& thread : threads) thread.join(); + + for (uint8_t success : succeeded) EXPECT_TRUE(success); + for (uint64_t address : relocated) EXPECT_EQ(address, relocated.front()); + EXPECT_EQ(ShmTransportTestPeer::mappingCount(transport, LOCAL_SEGMENT_ID), + 1u); + + auto* mapped = reinterpret_cast(relocated.front()); + ASSERT_TRUE(transport.uninstall().ok()); + unsigned char residency = 0; + errno = 0; + EXPECT_EQ(mincore(mapped, page_size, &residency), -1); + EXPECT_EQ(errno, ENOMEM); + + uint64_t address_after_uninstall = kRemoteAddress; + EXPECT_TRUE(ShmTransportTestPeer::relocate(transport, + address_after_uninstall, + page_size, LOCAL_SEGMENT_ID) + .IsInvalidArgument()); +} + +// B1: consumer CXL path must not create a missing file (would SIGBUS later). +TEST(ShmTransportTest, CxlConsumerDoesNotCreateMissingFile) { + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + constexpr uint64_t kRemoteAddress = 0x20000000; + + char tmpl[] = "/tmp/mooncake_shm_cxl_XXXXXX"; + ASSERT_NE(mkdtemp(tmpl), nullptr); + const std::string cxl_dir(tmpl); + const std::string shm_name = + "mooncake_cxl_missing_" + std::to_string(getpid()); + const std::string full_path = cxl_dir + "/" + shm_name; + + auto metadata = std::make_shared("p2p", "", nullptr); + ASSERT_TRUE(installLocalSegmentWithShm(*metadata, shm_name, kRemoteAddress, + page_size) + .ok()); + + auto conf = std::make_shared(); + conf->set("transports/shm/cxl_mount_path", cxl_dir); + + ShmTransport transport; + std::string local_segment_name = "shm_test_segment"; + ASSERT_TRUE( + transport.install(local_segment_name, metadata, nullptr, conf).ok()); + + uint64_t address = kRemoteAddress; + auto status = ShmTransportTestPeer::relocate(transport, address, page_size, + LOCAL_SEGMENT_ID); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.IsInternalError()); + + // The consumer must not have created the missing backing file. + struct stat st; + EXPECT_EQ(stat(full_path.c_str(), &st), -1); + EXPECT_EQ(errno, ENOENT); + + ASSERT_TRUE(transport.uninstall().ok()); + EXPECT_EQ(rmdir(cxl_dir.c_str()), 0); +} + +// Stale/truncated backing files must be rejected before mmap; otherwise +// access past EOF can SIGBUS even without O_CREAT. +TEST(ShmTransportTest, RejectsBackingFileShorterThanBuffer) { + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + constexpr uint64_t kRemoteAddress = 0x21000000; + + char tmpl[] = "/tmp/mooncake_shm_cxl_short_XXXXXX"; + ASSERT_NE(mkdtemp(tmpl), nullptr); + const std::string cxl_dir(tmpl); + const std::string shm_name = + "mooncake_cxl_short_" + std::to_string(getpid()); + const std::string full_path = cxl_dir + "/" + shm_name; + + // Create a truncated backing file (half a page) while metadata claims + // a full page. + int fd = open(full_path.c_str(), O_CREAT | O_EXCL | O_RDWR, 0600); + ASSERT_GE(fd, 0); + ASSERT_EQ(ftruncate(fd, page_size / 2), 0); + close(fd); + + auto metadata = std::make_shared("p2p", "", nullptr); + ASSERT_TRUE(installLocalSegmentWithShm(*metadata, shm_name, kRemoteAddress, + page_size) + .ok()); + + auto conf = std::make_shared(); + conf->set("transports/shm/cxl_mount_path", cxl_dir); + + ShmTransport transport; + std::string local_segment_name = "shm_test_segment"; + ASSERT_TRUE( + transport.install(local_segment_name, metadata, nullptr, conf).ok()); + + uint64_t address = kRemoteAddress; + auto status = ShmTransportTestPeer::relocate(transport, address, page_size, + LOCAL_SEGMENT_ID); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.IsInternalError()); + EXPECT_EQ(ShmTransportTestPeer::mappingCount(transport, LOCAL_SEGMENT_ID), + 0u); + + ASSERT_TRUE(transport.uninstall().ok()); + EXPECT_EQ(unlink(full_path.c_str()), 0); + EXPECT_EQ(rmdir(cxl_dir.c_str()), 0); +} + +// B2: creating with an existing name must not truncate the live object. +TEST(ShmTransportTest, CreateSharedMemoryDoesNotTruncateExisting) { + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + const std::string name = "mooncake_excl_test_" + std::to_string(getpid()); + shm_unlink(name.c_str()); + + int fd = shm_open(name.c_str(), O_CREAT | O_EXCL | O_RDWR, 0600); + ASSERT_GE(fd, 0); + ASSERT_EQ(ftruncate(fd, page_size), 0); + void* existing = + mmap(nullptr, page_size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + ASSERT_NE(existing, MAP_FAILED); + close(fd); + std::memset(existing, 0xAB, page_size); + + ShmTransport transport; + auto metadata = std::make_shared("p2p", "", nullptr); + std::string local_segment_name = "shm_test_segment"; + ASSERT_TRUE(transport + .install(local_segment_name, metadata, nullptr, + std::make_shared()) + .ok()); + + errno = 0; + void* created = + ShmTransportTestPeer::createSharedMemory(transport, name, page_size); + EXPECT_EQ(created, nullptr); + EXPECT_EQ(errno, EEXIST); + + // Existing mapping content must be intact (not truncated to zero). + auto* bytes = static_cast(existing); + EXPECT_EQ(bytes[0], 0xAB); + EXPECT_EQ(bytes[page_size - 1], 0xAB); + + // allocateLocalMemory should still succeed by picking a different name. + MemoryOptions options; + options.location = "cpu:0"; + void* allocated = nullptr; + ASSERT_TRUE( + transport.allocateLocalMemory(&allocated, page_size, options).ok()); + ASSERT_NE(allocated, nullptr); + EXPECT_NE(options.shm_path, name); + ASSERT_TRUE(transport.freeLocalMemory(allocated, page_size).ok()); + + munmap(existing, page_size); + shm_unlink(name.c_str()); + ASSERT_TRUE(transport.uninstall().ok()); +} + +// NeedsRefreshCache must not munmap cached mappings: memcpy may still be +// in flight on a previously resolved address after the relocate lock is +// released (no transfer-level refcount/quiesce yet). +TEST(ShmTransportTest, NeedsRefreshCacheKeepsExistingMappings) { + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + constexpr uint64_t kRemoteAddress = 0x30000000; + ScopedShmFile shm_file(page_size); + + auto metadata = std::make_shared("p2p", "", nullptr); + ASSERT_TRUE(installLocalSegmentWithShm(*metadata, shm_file.name(), + kRemoteAddress, page_size) + .ok()); + + ShmTransport transport; + std::string local_segment_name = "shm_test_segment"; + ASSERT_TRUE(transport + .install(local_segment_name, metadata, nullptr, + std::make_shared()) + .ok()); + + uint64_t address = kRemoteAddress; + ASSERT_TRUE(ShmTransportTestPeer::relocate(transport, address, page_size, + LOCAL_SEGMENT_ID) + .ok()); + EXPECT_EQ(ShmTransportTestPeer::mappingCount(transport, LOCAL_SEGMENT_ID), + 1u); + auto* mapped = reinterpret_cast(address); + + ASSERT_TRUE(metadata->segmentManager() + .updateLocal([&](SegmentDesc& segment) -> Status { + auto& memory = + std::get(segment.detail); + memory.buffers.clear(); + return Status::OK(); + }) + .ok()); + + uint64_t missing_address = kRemoteAddress + page_size; + auto status = ShmTransportTestPeer::relocate(transport, missing_address, + page_size, LOCAL_SEGMENT_ID); + EXPECT_TRUE(status.IsNeedsRefreshCache()); + // Mapping must remain alive for any in-flight reader. + EXPECT_EQ(ShmTransportTestPeer::mappingCount(transport, LOCAL_SEGMENT_ID), + 1u); + EXPECT_TRUE(ShmTransportTestPeer::hasTarget(transport, LOCAL_SEGMENT_ID)); + + unsigned char residency = 0; + errno = 0; + EXPECT_EQ(mincore(mapped, page_size, &residency), 0); + + ASSERT_TRUE(transport.uninstall().ok()); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/tent_metrics_example.cpp b/mooncake-transfer-engine/tent/tests/tent_metrics_example.cpp index 05d3b74704..7bad230ae2 100644 --- a/mooncake-transfer-engine/tent/tests/tent_metrics_example.cpp +++ b/mooncake-transfer-engine/tent/tests/tent_metrics_example.cpp @@ -78,9 +78,6 @@ int main(int argc, char* argv[]) { std::cout << "\nConfiguration:" << std::endl; std::cout << " - HTTP Server: " << config.http_host << ":" << config.http_port << std::endl; - std::cout << " - Prometheus: " - << (config.enable_prometheus ? "enabled" : "disabled") - << std::endl; std::cout << " - Runtime metrics: " << (config.enabled ? "enabled" : "disabled") << std::endl; @@ -112,19 +109,19 @@ int main(int argc, char* argv[]) { // Simulate successful reads with manual latency size_t read_bytes = 1024 * 1024 * (i + 1); // 1-10 MB double read_latency = 0.001 + (i * 0.0005); // 1-5ms - TENT_RECORD_READ_COMPLETED(read_bytes, read_latency); + TENT_RECORD_READ_COMPLETED(RDMA, read_bytes, read_latency); // Simulate successful writes with manual latency size_t write_bytes = 512 * 1024 * (i + 1); // 512KB - 5MB double write_latency = 0.002 + (i * 0.0003); // 2-4.7ms - TENT_RECORD_WRITE_COMPLETED(write_bytes, write_latency); + TENT_RECORD_WRITE_COMPLETED(RDMA, write_bytes, write_latency); // Simulate some failures if (i % 3 == 0) { - TENT_RECORD_READ_FAILED(1024); + TENT_RECORD_READ_FAILED(RDMA); } if (i % 4 == 0) { - TENT_RECORD_WRITE_FAILED(512); + TENT_RECORD_WRITE_FAILED(RDMA); } std::cout << " Iteration " << (i + 1) << ": Read " @@ -145,7 +142,7 @@ int main(int argc, char* argv[]) { // Example 1: Using TENT_SCOPED_READ_LATENCY macro { - TENT_SCOPED_READ_LATENCY(2 * 1024 * 1024); // 2 MB read + TENT_SCOPED_READ_LATENCY(RDMA, 2 * 1024 * 1024); // 2 MB read // Simulate work (the latency is automatically measured) std::this_thread::sleep_for(std::chrono::milliseconds(5)); std::cout << " Scoped read: 2 MB with ~5ms simulated work" @@ -154,7 +151,7 @@ int main(int argc, char* argv[]) { // Example 2: Using TENT_SCOPED_WRITE_LATENCY macro { - TENT_SCOPED_WRITE_LATENCY(1 * 1024 * 1024); // 1 MB write + TENT_SCOPED_WRITE_LATENCY(RDMA, 1 * 1024 * 1024); // 1 MB write // Simulate work std::this_thread::sleep_for(std::chrono::milliseconds(3)); std::cout << " Scoped write: 1 MB with ~3ms simulated work" @@ -164,7 +161,7 @@ int main(int argc, char* argv[]) { // Example 3: Using ScopedLatencyRecorder directly with failure handling { ScopedLatencyRecorder recorder( - ScopedLatencyRecorder::OperationType::Read, 512 * 1024); + ScopedLatencyRecorder::OperationType::Read, RDMA, 512 * 1024); std::this_thread::sleep_for(std::chrono::milliseconds(2)); // Simulate a failure condition bool operation_failed = true; // Simulated failure @@ -178,7 +175,7 @@ int main(int argc, char* argv[]) { // Example 4: Multiple scoped operations in a loop std::cout << "\n Running 5 scoped read operations..." << std::endl; for (int i = 0; i < 5; ++i) { - TENT_SCOPED_READ_LATENCY(256 * 1024 * (i + 1)); // 256KB - 1.25MB + TENT_SCOPED_READ_LATENCY(RDMA, 256 * 1024 * (i + 1)); // 256KB - 1.25MB std::this_thread::sleep_for(std::chrono::milliseconds(1 + i)); std::cout << " Scoped read " << (i + 1) << ": " << 256 * (i + 1) << " KB with ~" << (1 + i) << "ms work" << std::endl; @@ -199,7 +196,7 @@ int main(int argc, char* argv[]) { // These calls will return immediately with minimal overhead for (int i = 0; i < 1000; ++i) { TENT_RECORD_READ_COMPLETED( - 1024, 0.001); // These are no-ops when disabled + RDMA, 1024, 0.001); // These are no-ops when disabled } std::cout << " 1000 record calls completed (no-op, metrics disabled)" << std::endl; @@ -209,7 +206,7 @@ int main(int argc, char* argv[]) { TentMetrics::setEnabled(true); // Now these will be recorded - TENT_RECORD_READ_COMPLETED(1024 * 1024, 0.005); + TENT_RECORD_READ_COMPLETED(RDMA, 1024 * 1024, 0.005); std::cout << " Recorded 1 MB read after re-enabling" << std::endl; } else { std::cout << " Skipping enable/disable demo (started with --disabled)" @@ -247,8 +244,8 @@ int main(int argc, char* argv[]) { std::this_thread::sleep_for(std::chrono::seconds(1)); // Simulate ongoing traffic - TENT_RECORD_READ_COMPLETED(1024 * 1024, 0.001); - TENT_RECORD_WRITE_COMPLETED(512 * 1024, 0.002); + TENT_RECORD_READ_COMPLETED(RDMA, 1024 * 1024, 0.001); + TENT_RECORD_WRITE_COMPLETED(RDMA, 512 * 1024, 0.002); } } diff --git a/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp b/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp new file mode 100644 index 0000000000..cb6921a461 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/thread_local_storage_test.cpp @@ -0,0 +1,234 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Regression tests for #2717: ThreadLocalStorage used one `thread_local` +// slot per template instantiation (all instances aliased each other's +// per-thread state) and never destroyed its per-thread holders (the +// deregistration path was unreachable). These tests pin the per-instance +// semantics and both teardown orders; run them under ASAN/LSAN to verify +// the holder-leak fix and the owner-destroyed-first ordering. + +#include + +#include +#include +#include +#include +#include +#include + +#include "tent/common/concurrent/thread_local_storage.h" + +namespace mooncake { +namespace tent { + +namespace { +struct Cache { + int value = 0; +}; + +size_t countValues(ThreadLocalStorage& storage) { + size_t n = 0; + storage.forEach([&](Cache&) { n++; }); + return n; +} +} // namespace + +// The original aliasing bug: two instances in the same thread must own +// distinct per-thread values. +TEST(ThreadLocalStorageTest, InstancesDoNotAlias) { + ThreadLocalStorage a; + ThreadLocalStorage b; + a.get().value = 42; + EXPECT_EQ(b.get().value, 0); + EXPECT_NE(&a.get(), &b.get()); + b.get().value = 7; + EXPECT_EQ(a.get().value, 42); +} + +// Thread exits while the owner is alive: the value deregisters (forEach no +// longer sees it) and its memory is reclaimed (LSAN would flag the previous +// implementation, which leaked one holder per thread per instantiation). +TEST(ThreadLocalStorageTest, ThreadExitDeregistersAndReclaims) { + ThreadLocalStorage storage; + storage.get().value = 1; // main thread's value + EXPECT_EQ(countValues(storage), 1u); + + std::thread t([&] { + storage.get().value = 2; + EXPECT_EQ(countValues(storage), 2u); + }); + t.join(); + + EXPECT_EQ(countValues(storage), 1u); // worker's value deregistered + EXPECT_EQ(storage.get().value, 1); // main's value untouched +} + +// Owner destroyed while a using thread is still alive: the thread's later +// exit must not touch the dead owner (the jointly-owned control block keeps +// the registry memory valid; ASAN pins this ordering). +TEST(ThreadLocalStorageTest, OwnerDestroyedBeforeThreadExitIsSafe) { + std::atomic used{false}; + std::atomic release{false}; + auto storage = std::make_unique>(); + std::thread t([&] { + storage->get().value = 7; + used.store(true); + while (!release.load()) std::this_thread::yield(); + // Thread exit here runs the node destructor against a dead owner. + }); + while (!used.load()) std::this_thread::yield(); + storage.reset(); // owner gone first + release.store(true); + t.join(); +} + +// Instance ids are never reused: a new storage that may occupy the same +// address as a destroyed one must not see the old orphaned value. +TEST(ThreadLocalStorageTest, DestroyedInstanceStateIsNotResurrected) { + for (int round = 0; round < 8; ++round) { + auto storage = std::make_unique>(); + EXPECT_EQ(storage->get().value, 0) << "round " << round; + storage->get().value = 100 + round; + } +} + +// forEach synchronizes registry membership: it sees exactly the values of +// live threads that used this instance. +TEST(ThreadLocalStorageTest, ForEachVisitsExactlyLiveRegisteredValues) { + ThreadLocalStorage storage; + constexpr int kThreads = 8; + std::atomic ready{0}; + std::atomic release{false}; + std::vector threads; + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&, i] { + storage.get().value = i + 1; + ready++; + while (!release.load()) std::this_thread::yield(); + }); + } + while (ready.load() < kThreads) std::this_thread::yield(); + + int sum = 0; + size_t n = 0; + storage.forEach([&](Cache& c) { + sum += c.value; + n++; + }); + EXPECT_EQ(n, (size_t)kThreads); + EXPECT_EQ(sum, kThreads * (kThreads + 1) / 2); + + release.store(true); + for (auto& t : threads) t.join(); + EXPECT_EQ(countValues(storage), 0u); +} + +// Concurrent churn: threads exercising get() across shared storages while +// other storages are created/destroyed, with forEach mixed in. Run under +// TSAN/ASAN for the full effect; asserts basic integrity without them. +TEST(ThreadLocalStorageTest, ConcurrentChurnStress) { + constexpr int kThreads = 8; + constexpr int kIterations = 2000; + ThreadLocalStorage shared_a; + ThreadLocalStorage shared_b; + std::atomic failures{0}; + std::vector threads; + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&, i] { + for (int iter = 0; iter < kIterations; ++iter) { + shared_a.get().value = i; + shared_b.get().value = -i; + if (shared_a.get().value != i) failures++; + if (shared_b.get().value != -i) failures++; + // Thread-private storages churn creation/destruction. + ThreadLocalStorage ephemeral; + ephemeral.get().value = iter; + if (ephemeral.get().value != iter) failures++; + if (iter % 64 == 0) { + shared_a.forEach([](Cache&) {}); + } + } + }); + } + for (auto& t : threads) t.join(); + EXPECT_EQ(failures.load(), 0u); + EXPECT_EQ(countValues(shared_a), 0u); +} + +namespace { +struct Counted { + static std::atomic live; + int value = 0; + Counted() { live.fetch_add(1, std::memory_order_relaxed); } + ~Counted() { live.fetch_sub(1, std::memory_order_relaxed); } +}; +std::atomic Counted::live{0}; +} // namespace + +// Storage churn on a long-lived thread: values of destroyed storages must be +// swept by the next first-use get() rather than accumulating until thread +// exit (one orphaned T — pinning its contents — per destroyed storage). +TEST(ThreadLocalStorageTest, OrphanedValuesSweptOnInstanceChurn) { + constexpr int kRounds = 1000; + for (int i = 0; i < kRounds; ++i) { + ThreadLocalStorage storage; + storage.get().value = i; + // Previous round's orphan must have been swept by this round's + // first-use get(): at most this round's value plus one not-yet-swept + // orphan may be alive. + ASSERT_LE(Counted::live.load(), 2) << "round " << i; + } +} + +// Informational: hot-path cost of get(). The remote-desc cache consults this +// on every transfer submit, so the common case must stay a thread_local +// access plus a compare. +TEST(ThreadLocalStorageTest, HotPathMicrobench) { + ThreadLocalStorage storage; + storage.get().value = 1; + constexpr uint64_t kOps = 20'000'000; + volatile int sink = 0; + auto t0 = std::chrono::steady_clock::now(); + for (uint64_t i = 0; i < kOps; ++i) { + sink += storage.get().value; + } + auto t1 = std::chrono::steady_clock::now(); + double ns = + (double)std::chrono::duration_cast(t1 - t0) + .count() / + (double)kOps; + printf("get_hot_path_ns_per_op %.2f\n", ns); + (void)sink; +#if defined(__SANITIZE_THREAD__) || defined(__SANITIZE_ADDRESS__) + constexpr bool kSanitized = true; +#elif defined(__has_feature) +#if __has_feature(thread_sanitizer) || __has_feature(address_sanitizer) + constexpr bool kSanitized = true; +#else + constexpr bool kSanitized = false; +#endif +#else + constexpr bool kSanitized = false; +#endif + // Wall-clock assertions flake under sanitizers (TSAN alone is ~14x); + // elsewhere keep a loose ceiling that still catches syscall- or + // contention-class regressions on the hot path. + if (!kSanitized) { + EXPECT_LT(ns, 100.0); + } +} + +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp b/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp new file mode 100644 index 0000000000..0773faffd7 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/tpu/mock_tpu_pjrt_adapter.cpp @@ -0,0 +1,167 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Mock TPU/PJRT adapter for unit testing TpuPjrtShim, TpuPlatform and +// TpuTransport without TPU hardware or a PJRT runtime. It implements the C ABI +// declared in tpu_pjrt_abi.h, so the routing, classification and staging logic +// can be exercised on any Linux host (see tpu_pjrt_shim_test.cpp and +// tpu_transport_test.cpp). +// +// Two properties of real PJRT/TPU are modelled deliberately, because both are +// load-bearing for TENT's correctness and neither is obvious: +// +// 1. The device "pointer" is an opaque TOKEN, not the buffer's data. On real +// hardware PJRT_Buffer_UnsafePointer returns an internal handle that is +// host-dereferenceable but reads as unrelated bytes. So the token here is a +// separate read-only mapping poisoned with kPoison, and the buffer contents +// live in a shadow allocation reachable only through the copy entrypoints. +// Any code path that "helpfully" memcpy()s from a device token therefore +// reads poison instead of silently producing the right answer -- which is +// what a plain-host-memory mock would have done, and why an +// interior-pointer bug could pass a green test suite. +// +// 2. Addresses may be INTERIOR to a registered buffer. TENT stages transfers +// in +// chunks and passes `token + chunk_offset` for every chunk after the first, +// so classification and copies resolve ranges, and a copy that would run +// past a buffer's end is rejected rather than truncated. + +#include + +#include +#include +#include +#include + +#include "tent/platform/tpu_pjrt_abi.h" + +namespace { +// Byte filling the token mapping. Nothing should ever read it; if a test sees +// 0xDD it means something dereferenced a device token directly. +constexpr unsigned char kPoison = 0xDD; + +struct Buffer { + uintptr_t token; // opaque handle handed out to TENT + unsigned char *shadow; // where the bytes actually live + size_t size; + int device; +}; + +std::mutex g_mutex; +std::vector g_device_registry; +int g_device_count = 4; + +// Returns the buffer containing [addr, addr + len), or nullptr. len == 0 only +// checks that `addr` itself is inside a buffer. +Buffer *findLocked(const void *addr, size_t len) { + auto a = reinterpret_cast(addr); + for (auto &b : g_device_registry) { + if (a < b.token || a >= b.token + b.size) continue; + if (len > b.token + b.size - a) return nullptr; // runs past the end + return &b; + } + return nullptr; +} +} // namespace + +extern "C" { + +// --- ABI required by TpuPjrtShim ------------------------------------------- + +int mc_tpu_pjrt_init(void) { return 0; } + +int mc_tpu_pjrt_is_device_ptr(const void *addr) { + std::lock_guard lock(g_mutex); + return findLocked(addr, 0) ? 1 : 0; +} + +int mc_tpu_pjrt_device_index(const void *addr) { + std::lock_guard lock(g_mutex); + const Buffer *b = findLocked(addr, 0); + return b ? b->device : -1; +} + +int mc_tpu_pjrt_copy_d2h(void *host_dst, const void *device_src, size_t len) { + if (!host_dst || !device_src) return 1; + std::lock_guard lock(g_mutex); + const Buffer *b = findLocked(device_src, len); + if (!b) return 1; + size_t offset = reinterpret_cast(device_src) - b->token; + std::memcpy(host_dst, b->shadow + offset, len); + return 0; +} + +int mc_tpu_pjrt_copy_h2d(void *device_dst, const void *host_src, size_t len) { + if (!device_dst || !host_src) return 1; + std::lock_guard lock(g_mutex); + const Buffer *b = findLocked(device_dst, len); + if (!b) return 1; + size_t offset = reinterpret_cast(device_dst) - b->token; + std::memcpy(b->shadow + offset, host_src, len); + return 0; +} + +int mc_tpu_pjrt_device_count(void) { return g_device_count; } + +int mc_tpu_pjrt_device_numa(int index) { + // Deterministic fake affinity: even devices on node 0, odd on node 1. + if (index < 0 || index >= g_device_count) return -1; + return index % 2; +} + +// --- Test-only helpers (not part of the shim ABI) -------------------------- + +// Creates a fake device buffer of `size` bytes on device `index` and returns +// its token. The token is readable (like PJRT's unsafe pointer) but holds +// poison, never the buffer's data. +void *mock_tpu_pjrt_register_device(size_t size, int index) { + void *token = mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (token == MAP_FAILED) return nullptr; + std::memset(token, kPoison, size); + mprotect(token, size, PROT_READ); + + std::lock_guard lock(g_mutex); + g_device_registry.push_back(Buffer{reinterpret_cast(token), + new unsigned char[size](), size, index}); + return token; +} + +// Direct access to a buffer's bytes, for seeding inputs and asserting outputs. +// Accepts an interior token address and returns the matching shadow address. +void *mock_tpu_pjrt_device_data(const void *token) { + std::lock_guard lock(g_mutex); + Buffer *b = findLocked(token, 0); + if (!b) return nullptr; + return b->shadow + (reinterpret_cast(token) - b->token); +} + +unsigned char mock_tpu_pjrt_poison_byte(void) { return kPoison; } + +void mock_tpu_pjrt_reset(void) { + std::lock_guard lock(g_mutex); + for (auto &b : g_device_registry) { + munmap(reinterpret_cast(b.token), b.size); + delete[] b.shadow; + } + g_device_registry.clear(); + g_device_count = 4; +} + +void mock_tpu_pjrt_set_device_count(int count) { + std::lock_guard lock(g_mutex); + g_device_count = count; +} + +} // extern "C" diff --git a/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp b/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp new file mode 100644 index 0000000000..aa088274ca --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/tpu/tpu_pjrt_shim_test.cpp @@ -0,0 +1,188 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Unit test for TpuPjrtShim against the mock adapter +// (mock_tpu_pjrt_adapter.cpp). Runs on any Linux host: no TPU hardware or PJRT +// runtime required. +// +// MOCK_TPU_PJRT_LIB is the path to the built mock adapter, injected by CMake. + +#include "tent/platform/tpu_pjrt_shim.h" + +#include +#include + +#include +#include +#include + +#ifndef MOCK_TPU_PJRT_LIB +#error "MOCK_TPU_PJRT_LIB must be defined by the build (path to mock adapter)" +#endif + +namespace mooncake { +namespace tent { +namespace { + +using RegisterFn = void *(*)(size_t, int); +using DeviceDataFn = void *(*)(const void *); +using ResetFn = void (*)(); +using SetCountFn = void (*)(int); + +// Loads the mock adapter a second time (dlopen is reference-counted, so this is +// the same image the shim loads) to reach the test-only registration helpers. +class TpuShimTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + // Point the shim at the mock adapter before its singleton is first + // used. + ::setenv("MC_TPU_PJRT_LIB", MOCK_TPU_PJRT_LIB, /*overwrite=*/1); + } + + void SetUp() override { + mock_ = dlopen(MOCK_TPU_PJRT_LIB, RTLD_NOW | RTLD_GLOBAL); + ASSERT_NE(mock_, nullptr) << dlerror(); + register_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_register_device")); + device_data_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_device_data")); + reset_ = reinterpret_cast(dlsym(mock_, "mock_tpu_pjrt_reset")); + set_count_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_set_device_count")); + ASSERT_NE(register_, nullptr); + ASSERT_NE(device_data_, nullptr); + ASSERT_NE(reset_, nullptr); + ASSERT_NE(set_count_, nullptr); + reset_(); + } + + void TearDown() override { + if (reset_) reset_(); + if (mock_) dlclose(mock_); + } + + // Fills a fake device buffer with a repeatable pattern. + void seedDevice(void *token, size_t size, uint8_t modulus) { + auto *data = static_cast(device_data_(token)); + ASSERT_NE(data, nullptr); + for (size_t i = 0; i < size; ++i) data[i] = (uint8_t)(i % modulus); + } + + void *mock_ = nullptr; + RegisterFn register_ = nullptr; + DeviceDataFn device_data_ = nullptr; + ResetFn reset_ = nullptr; + SetCountFn set_count_ = nullptr; +}; + +TEST_F(TpuShimTest, AdapterLoadsAndReportsDevices) { + auto &shim = TpuPjrtShim::instance(); + ASSERT_TRUE(shim.available()); + EXPECT_EQ(shim.deviceCount(), 4); + EXPECT_EQ(shim.deviceNumaNode(0), 0); + EXPECT_EQ(shim.deviceNumaNode(1), 1); + EXPECT_EQ(shim.deviceNumaNode(99), -1); +} + +TEST_F(TpuShimTest, ClassifiesRegisteredDevicePointers) { + int host_value = 0; + void *token = register_(64, /*index=*/2); + ASSERT_NE(token, nullptr); + + auto &shim = TpuPjrtShim::instance(); + EXPECT_TRUE(shim.isDevicePtr(token)); + EXPECT_EQ(shim.deviceIndex(token), 2); + + // Unregistered host memory is not device memory. + EXPECT_FALSE(shim.isDevicePtr(&host_value)); + EXPECT_EQ(shim.deviceIndex(&host_value), -1); + EXPECT_FALSE(shim.isDevicePtr(nullptr)); +} + +// Regression: ProxyManager stages a transfer in chunk_size (4 MiB) pieces and +// hands the platform `token + chunk_offset` for every chunk after the first. An +// adapter that only recognises base addresses makes TENT classify TPU HBM as +// host memory, and the staging copy degrades into a memcpy from a token that on +// real PJRT is not the buffer's data -- silently corrupting every transfer +// larger than one chunk. Interior addresses must classify as device memory. +TEST_F(TpuShimTest, ClassifiesInteriorDevicePointers) { + const size_t kSize = 8192; + auto *token = static_cast(register_(kSize, /*index=*/3)); + ASSERT_NE(token, nullptr); + + auto &shim = TpuPjrtShim::instance(); + EXPECT_TRUE(shim.isDevicePtr(token + 1)); + EXPECT_TRUE(shim.isDevicePtr(token + 4096)); + EXPECT_TRUE(shim.isDevicePtr(token + kSize - 1)); + EXPECT_EQ(shim.deviceIndex(token + 4096), 3); + + // One past the end belongs to no buffer. + EXPECT_FALSE(shim.isDevicePtr(token + kSize)); + EXPECT_EQ(shim.deviceIndex(token + kSize), -1); +} + +TEST_F(TpuShimTest, CopyRoundTripMovesBytes) { + auto &shim = TpuPjrtShim::instance(); + const std::vector src = {1, 2, 3, 4, 5, 6, 7, 8}; + void *token = register_(src.size(), /*index=*/0); + ASSERT_NE(token, nullptr); + std::vector dst(src.size(), 0); + + // host -> "device" + ASSERT_TRUE(shim.copyH2D(token, src.data(), src.size()).ok()); + EXPECT_EQ(0, std::memcmp(device_data_(token), src.data(), src.size())); + + // "device" -> host + ASSERT_TRUE(shim.copyD2H(dst.data(), token, src.size()).ok()); + EXPECT_EQ(dst, src); +} + +// The chunked staging pattern end to end: copy each 4 KiB slice of a "device" +// buffer out through an interior pointer, exactly as ProxyManager would. +TEST_F(TpuShimTest, CopyFromInteriorOffsetMovesTheRightBytes) { + auto &shim = TpuPjrtShim::instance(); + const size_t kChunk = 4096, kChunks = 4, kSize = kChunk * kChunks; + auto *token = static_cast(register_(kSize, /*index=*/1)); + ASSERT_NE(token, nullptr); + seedDevice(token, kSize, 251); + + std::vector host(kSize, 0); + for (size_t c = 0; c < kChunks; ++c) { + ASSERT_TRUE( + shim.copyD2H(host.data() + c * kChunk, token + c * kChunk, kChunk) + .ok()) + << "chunk " << c; + } + EXPECT_EQ(0, std::memcmp(host.data(), device_data_(token), kSize)); + // And the bytes are the seeded pattern, not the token's poison. + EXPECT_EQ(host[0], 0); + EXPECT_EQ(host[kChunk + 1], (uint8_t)((kChunk + 1) % 251)); +} + +// A copy must never run past the end of the buffer it started in. +TEST_F(TpuShimTest, RejectsCopyRunningPastBufferEnd) { + auto &shim = TpuPjrtShim::instance(); + auto *token = static_cast(register_(1024, /*index=*/0)); + ASSERT_NE(token, nullptr); + std::vector host(2048, 0); + + EXPECT_FALSE(shim.copyD2H(host.data(), token + 512, 1024).ok()); + EXPECT_FALSE(shim.copyH2D(token + 512, host.data(), 1024).ok()); + // Unregistered addresses are not device memory and cannot be copied. + EXPECT_FALSE(shim.copyD2H(host.data(), host.data() + 1024, 16).ok()); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/tpu/tpu_transport_test.cpp b/mooncake-transfer-engine/tent/tests/tpu/tpu_transport_test.cpp new file mode 100644 index 0000000000..ee92a45fc1 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/tpu/tpu_transport_test.cpp @@ -0,0 +1,204 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Drives TpuTransport's staging hop against the mock PJRT adapter. No TPU +// hardware or PJRT runtime required. +// +// The cases here mirror the requests ProxyManager actually issues: a local +// stage whose `source` is `device_token + chunk_offset` and whose +// `target_offset` is the host staging buffer, and the delegated remote stage +// where the device side is `target_offset` instead. +// +// The mock's device tokens are poisoned (see mock_tpu_pjrt_adapter.cpp), so a +// staging copy that bypasses the adapter and memcpy()s straight from a token +// yields 0xDD rather than the buffer's data. That makes the data assertions +// below real detectors of the "device pointer misclassified as host memory" +// failure mode, instead of accidentally passing. + +#include "tent/transport/tpu/tpu_transport.h" + +#include +#include + +#include +#include +#include + +#ifndef MOCK_TPU_PJRT_LIB +#error "MOCK_TPU_PJRT_LIB must be defined by the build (path to mock adapter)" +#endif + +namespace mooncake { +namespace tent { +namespace { + +using RegisterFn = void *(*)(size_t, int); +using DeviceDataFn = void *(*)(const void *); +using PoisonFn = unsigned char (*)(); +using ResetFn = void (*)(); + +constexpr size_t kChunk = 4ul << 20; // ProxyManager's default chunk_size +constexpr size_t kBufSize = 3 * kChunk; + +class TpuTransportTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + ::setenv("MC_TPU_PJRT_LIB", MOCK_TPU_PJRT_LIB, /*overwrite=*/1); + } + + void SetUp() override { + mock_ = dlopen(MOCK_TPU_PJRT_LIB, RTLD_NOW | RTLD_GLOBAL); + ASSERT_NE(mock_, nullptr) << dlerror(); + register_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_register_device")); + device_data_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_device_data")); + poison_ = reinterpret_cast( + dlsym(mock_, "mock_tpu_pjrt_poison_byte")); + reset_ = reinterpret_cast(dlsym(mock_, "mock_tpu_pjrt_reset")); + ASSERT_NE(register_, nullptr); + ASSERT_NE(device_data_, nullptr); + ASSERT_NE(poison_, nullptr); + ASSERT_NE(reset_, nullptr); + reset_(); + + std::string segment = "local"; + ASSERT_TRUE( + transport_.install(segment, nullptr, nullptr, nullptr).ok()); + } + + void TearDown() override { + transport_.uninstall(); + if (reset_) reset_(); + if (mock_) dlclose(mock_); + } + + // Runs one request through the transport and returns its final status. + TransferStatus run(const Request &request) { + Transport::SubBatchRef batch = nullptr; + EXPECT_TRUE(transport_.allocateSubBatch(batch, 1).ok()); + EXPECT_TRUE(transport_.submitTransferTasks(batch, {request}).ok()); + TransferStatus status{}; + EXPECT_TRUE(transport_.getTransferStatus(batch, 0, status).ok()); + EXPECT_TRUE(transport_.freeSubBatch(batch).ok()); + return status; + } + + static Request makeRequest(Request::OpCode op, void *source, + uint64_t target, size_t length) { + Request r; + r.opcode = op; + r.source = source; + r.length = length; + r.target_id = LOCAL_SEGMENT_ID; + r.target_offset = target; + return r; + } + + void *mock_ = nullptr; + RegisterFn register_ = nullptr; + DeviceDataFn device_data_ = nullptr; + PoisonFn poison_ = nullptr; + ResetFn reset_ = nullptr; + TpuTransport transport_; +}; + +// The local WRITE stage of chunk #1: source is an interior device address. +// Before interior pointers were part of the adapter contract this silently +// memcpy()d from a non-data token and still reported COMPLETED. +TEST_F(TpuTransportTest, LocalStageCopiesFromInteriorDeviceOffset) { + auto *token = static_cast(register_(kBufSize, /*index=*/0)); + ASSERT_NE(token, nullptr); + auto *data = static_cast(device_data_(token)); + for (size_t i = 0; i < kBufSize; ++i) data[i] = (uint8_t)(i % 251); + std::vector staging(kChunk, 0); + + // Chunk #1: device_token + 4 MiB -> host staging buffer. + auto status = run(makeRequest(Request::WRITE, token + kChunk, + (uint64_t)staging.data(), kChunk)); + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED); + EXPECT_EQ(status.transferred_bytes, kChunk); + EXPECT_EQ(0, std::memcmp(staging.data(), data + kChunk, kChunk)); + // Not the token's poison: the adapter really did the copy. + EXPECT_NE(staging[0], poison_()); +} + +// The mirrored remote stage: the device side is `target_offset`, at an offset. +TEST_F(TpuTransportTest, RemoteStageCopiesToInteriorDeviceOffset) { + auto *token = static_cast(register_(kBufSize, /*index=*/0)); + ASSERT_NE(token, nullptr); + std::vector staging(kChunk); + for (size_t i = 0; i < kChunk; ++i) staging[i] = (uint8_t)(i % 197); + + // WRITE with a device target: host staging -> device_token + 4 MiB. + auto status = run(makeRequest(Request::WRITE, staging.data(), + (uint64_t)(token + kChunk), kChunk)); + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED); + auto *data = static_cast(device_data_(token)); + EXPECT_EQ(0, std::memcmp(data + kChunk, staging.data(), kChunk)); +} + +// A READ stage moves host staging -> device. +TEST_F(TpuTransportTest, ReadStageCopiesHostToDevice) { + auto *token = static_cast(register_(kBufSize, /*index=*/0)); + ASSERT_NE(token, nullptr); + std::vector staging(kChunk, 0xAB); + + auto status = run(makeRequest(Request::READ, token + 2 * kChunk, + (uint64_t)staging.data(), kChunk)); + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED); + auto *data = static_cast(device_data_(token)); + EXPECT_EQ(0, std::memcmp(data + 2 * kChunk, staging.data(), kChunk)); +} + +// Neither side is device memory: the adapter is missing, or it failed to +// classify an interior pointer. Either way the transport must fail rather than +// let Platform::copy memcpy from a token that is not the buffer's data. +TEST_F(TpuTransportTest, FailsWhenNeitherSideIsDeviceMemory) { + std::vector host_a(4096, 1), host_b(4096, 2); + auto status = run(makeRequest(Request::WRITE, host_a.data(), + (uint64_t)host_b.data(), host_a.size())); + EXPECT_EQ(status.s, TransferStatusEnum::FAILED); + EXPECT_EQ(status.transferred_bytes, 0u); + // The destination is untouched -- no silent partial copy. + EXPECT_EQ(host_b[0], 2); +} + +// Both sides device: HBM<->HBM is not a staging hop and must be rejected. +TEST_F(TpuTransportTest, FailsWhenBothSidesAreDeviceMemory) { + auto *a = static_cast(register_(kBufSize, /*index=*/0)); + auto *b = static_cast(register_(kBufSize, /*index=*/1)); + ASSERT_NE(a, nullptr); + ASSERT_NE(b, nullptr); + + auto status = run(makeRequest(Request::WRITE, a, (uint64_t)b, kChunk)); + EXPECT_EQ(status.s, TransferStatusEnum::FAILED); +} + +// TPU HBM can never be a remote peer; a non-local target is a routing bug. +TEST_F(TpuTransportTest, FailsOnNonLocalTarget) { + auto *token = static_cast(register_(kBufSize, /*index=*/0)); + ASSERT_NE(token, nullptr); + std::vector staging(kChunk, 0); + + auto request = + makeRequest(Request::WRITE, token, (uint64_t)staging.data(), kChunk); + request.target_id = LOCAL_SEGMENT_ID + 1; + auto status = run(request); + EXPECT_EQ(status.s, TransferStatusEnum::FAILED); +} + +} // namespace +} // namespace tent +} // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp index 8704924b1b..5182182762 100644 --- a/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transfer_engine_config_override_test.cpp @@ -283,6 +283,59 @@ TEST(TransferEngineConfigOverrideTest, EXPECT_EQ(config.get("transports/rdma/bind_address", ""), "10.0.0.2"); } +TEST(TransferEngineConfigOverrideTest, + LegacyRdmaSliceAffinityLogEnvLoadsIntoTentConfig) { + EnvVarGuard guard("MC_LOG_RDMA_SLICE_AFFINITY", "true"); + + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + EXPECT_TRUE(config.get("transports/rdma/log_slice_affinity", false)); +} + +TEST(TransferEngineConfigOverrideTest, FilterNicEnvLoadsRdmaWhitelist) { + EnvVarGuard guard("MC_TE_FILTERS", "mlx5_0,mlx5_1"); + + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + std::vector expected{"mlx5_0", "mlx5_1"}; + EXPECT_EQ(config.getArray("topology/rdma_whitelist"), + expected); +} + +TEST(TransferEngineConfigOverrideTest, FilterNicExcludeEnvLoadsRdmaBlacklist) { + EnvVarGuard guard("MC_TE_FILTERS_EXCLUDE", "mlx5_2,mlx5_3"); + + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + std::vector expected{"mlx5_2", "mlx5_3"}; + EXPECT_EQ(config.getArray("topology/rdma_blacklist"), + expected); +} + +TEST(TransferEngineConfigOverrideTest, FilterNicEnvTrimsWhitespaceAndEmpties) { + // Spaces around names and a trailing comma must be tolerated. + EnvVarGuard guard("MC_TE_FILTERS", " mlx5_0 , mlx5_1 ,"); + + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + std::vector expected{"mlx5_0", "mlx5_1"}; + EXPECT_EQ(config.getArray("topology/rdma_whitelist"), + expected); +} + +TEST(TransferEngineConfigOverrideTest, FilterNicUnsetLeavesWhitelistEmpty) { + // Not setting the env var must leave the default (discover all NICs). + Config config; + ASSERT_TRUE(ConfigHelper().loadFromEnv(config).ok()); + + EXPECT_TRUE( + config.getArray("topology/rdma_whitelist").empty()); +} + TEST(TransferEngineConfigOverrideTest, ExplicitMetadataOverridesDriveSuccessfulHttpInitialization) { #ifdef _WIN32 diff --git a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp index a479c8e43a..d16fcee52e 100644 --- a/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp +++ b/mooncake-transfer-engine/tent/tests/transport_selector_test.cpp @@ -14,6 +14,8 @@ #include +#include +#include #include #include @@ -173,32 +175,116 @@ TEST(TransportSelectorTest, DefaultPoliciesMemorySegment) { // --------------------------------------------------------------------------- TEST(TransportSelectorTest, TransportTypeNameMapping) { - EXPECT_EQ(TransportSelector::transportTypeName(UNSPEC), "unspec"); - EXPECT_EQ(TransportSelector::transportTypeName(RDMA), "rdma"); - EXPECT_EQ(TransportSelector::transportTypeName(MNNVL), "mnnvl"); - EXPECT_EQ(TransportSelector::transportTypeName(SHM), "shm"); - EXPECT_EQ(TransportSelector::transportTypeName(NVLINK), "nvlink"); - EXPECT_EQ(TransportSelector::transportTypeName(GDS), "gds"); - EXPECT_EQ(TransportSelector::transportTypeName(IOURING), "io_uring"); - EXPECT_EQ(TransportSelector::transportTypeName(TCP), "tcp"); - EXPECT_EQ(TransportSelector::transportTypeName(AscendDirect), "ascend"); - EXPECT_EQ(TransportSelector::transportTypeName(SUNRISE_LINK), - "sunrise_link"); + EXPECT_STREQ(transportTypeName(UNSPEC), "unspec"); + EXPECT_STREQ(transportTypeName(RDMA), "rdma"); + EXPECT_STREQ(transportTypeName(MNNVL), "mnnvl"); + EXPECT_STREQ(transportTypeName(SHM), "shm"); + EXPECT_STREQ(transportTypeName(NVLINK), "nvlink"); + EXPECT_STREQ(transportTypeName(GDS), "gds"); + EXPECT_STREQ(transportTypeName(IOURING), "io_uring"); + EXPECT_STREQ(transportTypeName(TCP), "tcp"); + EXPECT_STREQ(transportTypeName(AscendDirect), "ascend"); + EXPECT_STREQ(transportTypeName(SUNRISE_LINK), "sunrise_link"); + EXPECT_STREQ(transportTypeName(UB), "ub"); } TEST(TransportSelectorTest, ParseTransportType) { - EXPECT_EQ(TransportSelector::parseTransportType("unspec"), UNSPEC); - EXPECT_EQ(TransportSelector::parseTransportType("rdma"), RDMA); - EXPECT_EQ(TransportSelector::parseTransportType("mnnvl"), MNNVL); - EXPECT_EQ(TransportSelector::parseTransportType("shm"), SHM); - EXPECT_EQ(TransportSelector::parseTransportType("nvlink"), NVLINK); - EXPECT_EQ(TransportSelector::parseTransportType("gds"), GDS); - EXPECT_EQ(TransportSelector::parseTransportType("io_uring"), IOURING); - EXPECT_EQ(TransportSelector::parseTransportType("tcp"), TCP); - EXPECT_EQ(TransportSelector::parseTransportType("ascend"), AscendDirect); - EXPECT_EQ(TransportSelector::parseTransportType("sunrise_link"), - SUNRISE_LINK); - EXPECT_EQ(TransportSelector::parseTransportType("unknown"), UNSPEC); + EXPECT_EQ(parseTransportType("unspec"), UNSPEC); + EXPECT_EQ(parseTransportType("rdma"), RDMA); + EXPECT_EQ(parseTransportType("mnnvl"), MNNVL); + EXPECT_EQ(parseTransportType("shm"), SHM); + EXPECT_EQ(parseTransportType("nvlink"), NVLINK); + EXPECT_EQ(parseTransportType("gds"), GDS); + EXPECT_EQ(parseTransportType("io_uring"), IOURING); + EXPECT_EQ(parseTransportType("tcp"), TCP); + EXPECT_EQ(parseTransportType("ascend"), AscendDirect); + EXPECT_EQ(parseTransportType("sunrise_link"), SUNRISE_LINK); + EXPECT_EQ(parseTransportType("ub"), UB); + EXPECT_EQ(parseTransportType("unknown"), UNSPEC); +} + +TEST(TransportSelectorTest, UbTransportNameRoundTrips) { + const auto name = transportTypeName(UB); + EXPECT_EQ(name, "ub"); + EXPECT_EQ(parseTransportType(name), UB); +} + +TEST(TransportTypeTest, WireValuesRemainStableWithUbAppended) { + EXPECT_EQ(static_cast(UNSPEC), 0); + EXPECT_EQ(static_cast(RDMA), 1); + EXPECT_EQ(static_cast(MNNVL), 2); + EXPECT_EQ(static_cast(SHM), 3); + EXPECT_EQ(static_cast(NVLINK), 4); + EXPECT_EQ(static_cast(GDS), 5); + EXPECT_EQ(static_cast(IOURING), 6); + EXPECT_EQ(static_cast(TCP), 7); + EXPECT_EQ(static_cast(AscendDirect), 8); + EXPECT_EQ(static_cast(SUNRISE_LINK), 9); + EXPECT_EQ(static_cast(TPU), 10); + EXPECT_EQ(static_cast(UB), 11); + EXPECT_EQ(static_cast(kNumTransportTypes), 12); +} + +// Topology::NicType is serialized as an integer. These values are therefore a +// wire-compatibility contract, not merely an implementation detail. +TEST(TopologyTest, NicTypeWireValuesRemainStableWithUbAppended) { + EXPECT_EQ(static_cast(Topology::NIC_RDMA), 0); + EXPECT_EQ(static_cast(Topology::NIC_TCP), 1); + EXPECT_EQ(static_cast(Topology::NIC_UNKNOWN), 2); + EXPECT_EQ(static_cast(Topology::NIC_UB), 3); +} + +TEST(TopologyTest, LegacyJsonDefaultsDeviceAttributes) { + constexpr const char* kLegacyTopology = R"json( + { + "nics": [ + { + "name": "legacy-nic", + "pci_bus_id": "0000:01:00.0", + "type": 2, + "numa_node": -1 + } + ], + "mems": [] + } + )json"; + + Topology topology; + ASSERT_TRUE(topology.parse(kLegacyTopology).ok()); + ASSERT_EQ(topology.getNicCount(), 1u); + const auto* nic = topology.getNicEntry(0); + ASSERT_NE(nic, nullptr); + EXPECT_EQ(nic->type, Topology::NIC_UNKNOWN); + EXPECT_TRUE(nic->device_attrs.empty()); + EXPECT_EQ(topology.toString().find("device_attrs"), std::string::npos); +} + +TEST(TopologyTest, UbDeviceAttributesRoundTripThroughJson) { + Topology source; + Topology::NicEntry ub; + ub.name = "ub-device-0/eid-2"; + ub.pci_bus_id = "0000:02:00.0"; + ub.type = Topology::NIC_UB; + ub.numa_node = 1; + ub.device_attrs = {{"ub.native_name", "ub-device-0"}, + {"ub.device_index", "7"}, + {"ub.eid_index", "2"}, + {"ub.eid", "e1:02:03:04:05:06:07:08"}, + {"ub.discovery_active", "false"}, + {"vendor.future_attribute", "preserved"}}; + source.nic_list_.push_back(ub); + + Topology parsed; + ASSERT_TRUE(parsed.parse(source.toString()).ok()); + ASSERT_EQ(parsed.getNicCount(), 1u); + ASSERT_EQ(parsed.getNicCount(Topology::NIC_UB), 1u); + const auto* round_tripped = parsed.getNicEntry(0); + ASSERT_NE(round_tripped, nullptr); + EXPECT_EQ(round_tripped->name, ub.name); + EXPECT_EQ(round_tripped->pci_bus_id, ub.pci_bus_id); + EXPECT_EQ(round_tripped->type, Topology::NIC_UB); + EXPECT_EQ(round_tripped->numa_node, ub.numa_node); + EXPECT_EQ(round_tripped->device_attrs, ub.device_attrs); } // --------------------------------------------------------------------------- @@ -641,6 +727,333 @@ TEST(TransportSelectorTest, HintNotInMatchingPolicyReturnsUnspec) { EXPECT_EQ(r.transport, UNSPEC); } +// RFC #2519 / #2568 step 1: a policy's link-layer QoS (service_level / +// traffic_class / qp_pool) is parsed from JSON and carried out via +// SelectionResult. (Step 1 only plumbs the values; applying them at QP setup +// is the per-class QP pool follow-up.) +TEST(TransportSelectorTest, PolicyLinkLayerQoSIsParsedAndCarried) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "kv-critical"; + policy["segment_type"] = "memory"; + policy["transports"] = {"rdma"}; + policy["service_level"] = 3; + policy["traffic_class"] = 96; + policy["qp_pool"] = "kv"; + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[RDMA].get())->setDramToDram(true); + + std::vector buffer_transports = {RDMA}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.buffer_transports = &buffer_transports; + ctx.policy_name = "kv-critical"; + + auto r = selector.select(ctx, transports, /*index=*/0); + ASSERT_TRUE(r.service_level.has_value()); + EXPECT_EQ(r.service_level.value(), 3); + ASSERT_TRUE(r.traffic_class.has_value()); + EXPECT_EQ(r.traffic_class.value(), 96); + ASSERT_TRUE(r.qp_pool.has_value()); + EXPECT_EQ(r.qp_pool.value(), "kv"); +} + +// Out-of-range SL/TC are ignored (left as nullopt) so a bad config never +// changes selection behavior. +TEST(TransportSelectorTest, PolicyLinkLayerQoSOutOfRangeIgnored) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "bad-qos"; + policy["segment_type"] = "memory"; + policy["transports"] = {"rdma"}; + policy["service_level"] = 99; // > 15, invalid + policy["traffic_class"] = 9999; // > 255, invalid + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[RDMA].get())->setDramToDram(true); + + std::vector buffer_transports = {RDMA}; + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.buffer_transports = &buffer_transports; + ctx.policy_name = "bad-qos"; + + auto r = selector.select(ctx, transports, /*index=*/0); + EXPECT_FALSE(r.service_level.has_value()); + EXPECT_FALSE(r.traffic_class.has_value()); +} + +// An empty-string or non-string qp_pool is treated as unset (nullopt), so a +// blank / mistyped config value never looks like an explicit pool. +TEST(TransportSelectorTest, PolicyQpPoolEmptyOrNonStringIsUnset) { + auto conf = std::make_shared(); + json empty_pool; + empty_pool["name"] = "empty-pool"; + empty_pool["segment_type"] = "memory"; + empty_pool["transports"] = {"rdma"}; + empty_pool["qp_pool"] = ""; // empty -> unset + json bad_pool; + bad_pool["name"] = "bad-pool"; + bad_pool["segment_type"] = "memory"; + bad_pool["transports"] = {"rdma"}; + bad_pool["qp_pool"] = 42; // non-string -> ignored + conf->set("policy", json::array({empty_pool, bad_pool})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[RDMA].get())->setDramToDram(true); + std::vector buffer_transports = {RDMA}; + + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.buffer_transports = &buffer_transports; + + ctx.policy_name = "empty-pool"; + EXPECT_FALSE( + selector.select(ctx, transports, /*index=*/0).qp_pool.has_value()); + ctx.policy_name = "bad-pool"; + EXPECT_FALSE( + selector.select(ctx, transports, /*index=*/0).qp_pool.has_value()); +} + +// Intent-specific policies bind Request::intent_type to transport and +// link-layer QoS selection. Policies are first-match, so the specific entry is +// deliberately placed before the catch-all fallback. +TEST(TransportSelectorTest, IntentSpecificPolicyIsSelected) { + auto conf = std::make_shared(); + json foreground; + foreground["name"] = "foreground"; + foreground["segment_type"] = "memory"; + foreground["intent_type"] = "foreground_get"; + foreground["transports"] = {"rdma"}; + foreground["service_level"] = 3; + foreground["traffic_class"] = 96; + foreground["qp_pool"] = "foreground"; + json fallback; + fallback["name"] = "fallback"; + fallback["segment_type"] = "memory"; + fallback["transports"] = {"tcp"}; + conf->set("policy", json::array({foreground, fallback})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + transports[TCP] = std::make_shared(TCP); + static_cast(transports[RDMA].get())->setDramToDram(true); + static_cast(transports[TCP].get())->setDramToDram(true); + std::vector buffer_transports = {RDMA, TCP}; + + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; + ctx.priority_level = PRIO_HIGH; + ctx.buffer_transports = &buffer_transports; + ctx.intent_type = IntentType::FOREGROUND_GET; + + auto result = selector.select(ctx, transports); + EXPECT_EQ(result.transport, RDMA); + EXPECT_EQ(result.service_level, 3); + EXPECT_EQ(result.traffic_class, 96); + EXPECT_EQ(result.qp_pool, "foreground"); +} + +TEST(TransportSelectorTest, IntentMismatchFallsThroughToCatchAll) { + auto conf = std::make_shared(); + json foreground; + foreground["name"] = "foreground"; + foreground["segment_type"] = "memory"; + foreground["intent_type"] = "foreground_get"; + foreground["transports"] = {"rdma"}; + json fallback; + fallback["name"] = "fallback"; + fallback["segment_type"] = "memory"; + fallback["transports"] = {"tcp"}; + conf->set("policy", json::array({foreground, fallback})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + transports[TCP] = std::make_shared(TCP); + static_cast(transports[RDMA].get())->setDramToDram(true); + static_cast(transports[TCP].get())->setDramToDram(true); + std::vector buffer_transports = {RDMA, TCP}; + + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; + ctx.priority_level = PRIO_LOW; + ctx.buffer_transports = &buffer_transports; + ctx.intent_type = IntentType::CHECKPOINT; + + EXPECT_EQ(selector.select(ctx, transports).transport, TCP); +} + +TEST(TransportSelectorTest, PolicyWithoutIntentMatchesAnyIntent) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "legacy"; + policy["segment_type"] = "memory"; + policy["transports"] = {"rdma"}; + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + static_cast(transports[RDMA].get())->setDramToDram(true); + std::vector buffer_transports = {RDMA}; + + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; + ctx.priority_level = PRIO_LOW; + ctx.buffer_transports = &buffer_transports; + ctx.intent_type = IntentType::CHECKPOINT; + + EXPECT_EQ(selector.select(ctx, transports).transport, RDMA); +} + +TEST(TransportSelectorTest, NumericIntentValueIsAccepted) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "checkpoint"; + policy["segment_type"] = "memory"; + policy["intent_type"] = static_cast(IntentType::CHECKPOINT); + policy["transports"] = {"tcp"}; + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[TCP] = std::make_shared(TCP); + static_cast(transports[TCP].get())->setDramToDram(true); + std::vector buffer_transports = {TCP}; + + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; + ctx.priority_level = PRIO_LOW; + ctx.buffer_transports = &buffer_transports; + ctx.intent_type = IntentType::CHECKPOINT; + + EXPECT_EQ(selector.select(ctx, transports).transport, TCP); +} + +TEST(TransportSelectorTest, InvalidIntentPolicyIsSkipped) { + auto conf = std::make_shared(); + json bad_name; + bad_name["name"] = "bad-name"; + bad_name["segment_type"] = "memory"; + bad_name["intent_type"] = "not_an_intent"; + bad_name["transports"] = {"rdma"}; + json bad_number; + bad_number["name"] = "bad-number"; + bad_number["segment_type"] = "memory"; + bad_number["intent_type"] = 999; + bad_number["transports"] = {"rdma"}; + json bad_type; + bad_type["name"] = "bad-type"; + bad_type["segment_type"] = "memory"; + bad_type["intent_type"] = true; + bad_type["transports"] = {"rdma"}; + json bad_unsigned; + bad_unsigned["name"] = "bad-unsigned"; + bad_unsigned["segment_type"] = "memory"; + bad_unsigned["intent_type"] = std::numeric_limits::max(); + bad_unsigned["transports"] = {"rdma"}; + json fallback; + fallback["name"] = "fallback"; + fallback["segment_type"] = "memory"; + fallback["transports"] = {"tcp"}; + conf->set("policy", json::array({bad_name, bad_number, bad_type, + bad_unsigned, fallback})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[RDMA] = std::make_shared(RDMA); + transports[TCP] = std::make_shared(TCP); + static_cast(transports[RDMA].get())->setDramToDram(true); + static_cast(transports[TCP].get())->setDramToDram(true); + std::vector buffer_transports = {RDMA, TCP}; + + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; + ctx.priority_level = PRIO_HIGH; + ctx.buffer_transports = &buffer_transports; + ctx.intent_type = IntentType::FOREGROUND_GET; + + EXPECT_EQ(selector.select(ctx, transports).transport, TCP); +} + +TEST(TransportSelectorTest, ExplicitPolicyNameOverridesIntentFilter) { + auto conf = std::make_shared(); + json policy; + policy["name"] = "operator-override"; + policy["segment_type"] = "memory"; + policy["intent_type"] = "checkpoint"; + policy["transports"] = {"tcp"}; + conf->set("policy", json::array({policy})); + + TransportSelector selector(conf); + std::array, kSupportedTransportTypes> + transports{}; + transports[TCP] = std::make_shared(TCP); + static_cast(transports[TCP].get())->setDramToDram(true); + std::vector buffer_transports = {TCP}; + + SelectionContext ctx; + ctx.segment_type = SegmentType::Memory; + ctx.same_machine = false; + ctx.local_memory_type = MTYPE_CPU; + ctx.remote_memory_type = MTYPE_CPU; + ctx.transfer_size = 4096; + ctx.priority_level = PRIO_HIGH; + ctx.buffer_transports = &buffer_transports; + ctx.intent_type = IntentType::FOREGROUND_GET; + ctx.policy_name = "operator-override"; + + EXPECT_EQ(selector.select(ctx, transports).transport, TCP); +} + } // namespace } // namespace tent } // namespace mooncake diff --git a/mooncake-transfer-engine/tent/tests/ub_teardown_test.cpp b/mooncake-transfer-engine/tent/tests/ub_teardown_test.cpp new file mode 100644 index 0000000000..1b4b964ce7 --- /dev/null +++ b/mooncake-transfer-engine/tent/tests/ub_teardown_test.cpp @@ -0,0 +1,533 @@ +// Copyright 2026 KVCache.AI +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "tent/transport/ub/buffers.h" +#include "tent/transport/ub/context.h" +#include "tent/transport/ub/topology_attrs.h" + +namespace mooncake::tent::ub { +namespace { + +SegmentDescriptor makeDescriptor(uint64_t address, uint64_t length) { + return SegmentDescriptor{ + SegmentDescriptor::kSchemaVersion, 1, 16, + std::to_string(address) + ":" + std::to_string(length)}; +} + +DeviceInfo makeDevice(int index) { + DeviceInfo device; + device.topology_name = "ub:test:" + std::to_string(index); + device.native_device_name = "test" + std::to_string(index); + device.eid_index = static_cast(index); + device.eid = "0001:0002:0003:0004:0005:0006:0007:0008"; + device.active = true; + device.capabilities.max_jfc = 16; + return device; +} + +TEST(UbTopologyAttributes, EncoderPreservesHardwareNeutralMap) { + auto device = makeDevice(2); + device.active = false; + std::unordered_map attributes{ + {"vendor.future_attribute", "preserved"}}; + encodeTopologyDeviceAttributes(device, 7, attributes); + + EXPECT_EQ(attributes.at(std::string(kTopologyNativeNameAttr)), + device.native_device_name); + EXPECT_EQ(attributes.at(std::string(kTopologyDeviceIndexAttr)), "7"); + EXPECT_EQ(attributes.at(std::string(kTopologyEidIndexAttr)), "2"); + EXPECT_EQ(attributes.at(std::string(kTopologyEidAttr)), device.eid); + EXPECT_EQ(attributes.at(std::string(kTopologyDiscoveryActiveAttr)), + "false"); + EXPECT_EQ(attributes.at("vendor.future_attribute"), "preserved"); +} + +class FakeContext final : public Context { + public: + explicit FakeContext(DeviceInfo device) : device_(std::move(device)) {} + + bool valid() const noexcept override { return valid_; } + const DeviceInfo& deviceInfo() const noexcept override { return device_; } + int asyncFd() const noexcept override { return -1; } + void close() noexcept { valid_ = false; } + + private: + DeviceInfo device_; + bool valid_{true}; +}; + +class FakeJfc final : public Jfc { + public: + bool valid() const noexcept override { return valid_; } + int eventFd() const noexcept override { return -1; } + void close() noexcept { valid_ = false; } + + private: + bool valid_{true}; +}; + +class FakeLocalSegment final : public LocalSegment { + public: + FakeLocalSegment(uint64_t address, uint64_t length) + : address_(address), + length_(length), + descriptor_(makeDescriptor(address, length)) {} + + bool valid() const noexcept override { return valid_; } + uint64_t address() const noexcept override { return address_; } + uint64_t length() const noexcept override { return length_; } + const SegmentDescriptor& descriptor() const noexcept override { + return descriptor_; + } + void close() noexcept { valid_ = false; } + + private: + uint64_t address_; + uint64_t length_; + SegmentDescriptor descriptor_; + bool valid_{true}; +}; + +class FakeRemoteSegment final : public RemoteSegment { + public: + explicit FakeRemoteSegment(SegmentDescriptor descriptor) + : descriptor_(std::move(descriptor)) {} + + bool valid() const noexcept override { return valid_; } + uint64_t address() const noexcept override { return 0x1000; } + uint64_t length() const noexcept override { return 0x1000; } + const SegmentDescriptor& descriptor() const noexcept override { + return descriptor_; + } + void close() noexcept { valid_ = false; } + + private: + SegmentDescriptor descriptor_; + bool valid_{true}; +}; + +class FakeJetty final : public Jetty { + public: + bool valid() const noexcept override { return valid_; } + uint32_t id() const noexcept override { return 1; } + uint32_t uasid() const noexcept override { return 0; } + void close() noexcept { valid_ = false; } + + private: + bool valid_{true}; +}; + +class FailOnceAdapter final : public UrmaAdapter { + public: + bool available() const noexcept override { return true; } + uint32_t nativeApiVersion() const noexcept override { return 1; } + size_t nativeSegmentDescriptorSize() const noexcept override { return 16; } + + Status initialize() override { return Status::OK(); } + Status shutdown() override { return Status::OK(); } + Status discoverDevices(std::vector& devices) override { + devices = {makeDevice(0), makeDevice(1)}; + return Status::OK(); + } + + Status openContext(const DeviceInfo& device, ContextPtr& context) override { + context = std::make_shared(device); + ++live_contexts; + return Status::OK(); + } + Status closeContext(ContextPtr& context) override { + ++close_context_calls; + if (close_context_failures > 0) { + --close_context_failures; + return Status::InternalError("injected close Context failure"); + } + if (auto fake = std::dynamic_pointer_cast(context)) { + fake->close(); + } + context.reset(); + --live_contexts; + return Status::OK(); + } + + Status createJfc(const ContextPtr&, const JfcOptions&, + JfcPtr& jfc) override { + jfc = std::make_shared(); + ++live_jfcs; + return Status::OK(); + } + Status deleteJfc(JfcPtr& jfc) override { + ++delete_jfc_calls; + if (delete_jfc_failures > 0) { + --delete_jfc_failures; + return Status::InternalError("injected delete JFC failure"); + } + if (auto fake = std::dynamic_pointer_cast(jfc)) fake->close(); + jfc.reset(); + --live_jfcs; + return Status::OK(); + } + + Status registerLocalSegment(const ContextPtr&, uint64_t address, + size_t length, const SegmentOptions&, + LocalSegmentPtr& segment) override { + ++register_calls; + if (fail_register_call != 0 && register_calls == fail_register_call) { + if (partial_register_on_failure) { + segment = std::make_shared(address, length); + ++live_local_segments; + } + return Status::InternalError("injected register failure"); + } + segment = std::make_shared(address, length); + ++live_local_segments; + return Status::OK(); + } + Status unregisterLocalSegment(LocalSegmentPtr& segment) override { + ++unregister_calls; + if (unregister_failures > 0) { + --unregister_failures; + return Status::InternalError("injected unregister failure"); + } + if (auto fake = std::dynamic_pointer_cast(segment)) { + fake->close(); + } + segment.reset(); + --live_local_segments; + return Status::OK(); + } + + Status importRemoteSegment(const ContextPtr&, + const SegmentDescriptor& descriptor, + const SegmentOptions&, + RemoteSegmentPtr& segment) override { + if (import_failures > 0) { + --import_failures; + if (partial_import_on_failure) { + segment = std::make_shared(descriptor); + ++live_remote_segments; + } + return Status::InternalError("injected import failure"); + } + segment = std::make_shared(descriptor); + ++live_remote_segments; + return Status::OK(); + } + Status unimportRemoteSegment(RemoteSegmentPtr& segment) override { + ++unimport_calls; + if (unimport_failures > 0) { + --unimport_failures; + return Status::InternalError("injected unimport failure"); + } + if (auto fake = std::dynamic_pointer_cast(segment)) { + fake->close(); + } + segment.reset(); + --live_remote_segments; + return Status::OK(); + } + + Status createJetty(const ContextPtr&, const JfcPtr&, const JettyOptions&, + JettyPtr& jetty) override { + jetty = std::make_shared(); + return Status::OK(); + } + Status deleteJetty(JettyPtr& jetty) override { + ++delete_jetty_calls; + if (delete_jetty_failures > 0) { + --delete_jetty_failures; + return Status::InternalError("injected delete Jetty failure"); + } + if (auto fake = std::dynamic_pointer_cast(jetty)) { + fake->close(); + } + jetty.reset(); + return Status::OK(); + } + Status bindJetty(const JettyPtr&, const RemoteJettyInfo&) override { + return Status::OK(); + } + Status unbindJetty(const JettyPtr&) override { return Status::OK(); } + Status resetJetty(const JettyPtr&) override { return Status::OK(); } + Status quiesceJetty(const JettyPtr&, uint32_t, + std::vector& completions) override { + completions.clear(); + return Status::OK(); + } + Status post(const JettyPtr&, const std::vector& requests, + size_t& posted_count) override { + posted_count = requests.size(); + return Status::OK(); + } + Status poll(const JfcPtr&, size_t, + std::vector& completions) override { + completions.clear(); + return Status::OK(); + } + + int close_context_failures{0}; + int delete_jfc_failures{0}; + int unregister_failures{0}; + int unimport_failures{0}; + int import_failures{0}; + int delete_jetty_failures{0}; + int fail_register_call{0}; + bool partial_register_on_failure{false}; + bool partial_import_on_failure{false}; + int close_context_calls{0}; + int delete_jfc_calls{0}; + int register_calls{0}; + int unregister_calls{0}; + int unimport_calls{0}; + int delete_jetty_calls{0}; + int live_contexts{0}; + int live_jfcs{0}; + int live_local_segments{0}; + int live_remote_segments{0}; +}; + +UbContextPtr makeActiveContext(const std::shared_ptr& adapter, + int topology_id, uint32_t jfc_count = 1) { + auto context = std::make_shared( + topology_id, makeDevice(topology_id), adapter); + EXPECT_TRUE(context->initialize(jfc_count, JfcOptions{}).ok()); + return context; +} + +BufferDesc makeLocalBuffer(uint64_t address) { + BufferDesc desc; + desc.addr = address; + desc.length = 4096; + desc.location = "cpu:0"; + return desc; +} + +TEST(UbTeardown, RemoveBufferRetainsOwnershipAndMetadataForRetry) { + auto adapter = std::make_shared(); + auto context = makeActiveContext(adapter, 0); + UbBufferManager manager(adapter, {context}); + auto desc = makeLocalBuffer(0x10000); + ASSERT_TRUE(manager.addBuffer(desc, MemoryOptions{}).ok()); + ASSERT_EQ(manager.localBufferCount(), 1u); + + adapter->unregister_failures = 1; + EXPECT_FALSE(manager.removeBuffer(desc).ok()); + EXPECT_EQ(manager.localBufferCount(), 1u); + EXPECT_EQ(adapter->live_local_segments, 1); + EXPECT_TRUE(desc.transport_attrs.contains(TransportType::UB)); + EXPECT_NE(std::find(desc.transports.begin(), desc.transports.end(), + TransportType::UB), + desc.transports.end()); + + EXPECT_TRUE(manager.removeBuffer(desc).ok()); + EXPECT_EQ(manager.localBufferCount(), 0u); + EXPECT_EQ(adapter->live_local_segments, 0); + EXPECT_FALSE(desc.transport_attrs.contains(TransportType::UB)); + EXPECT_TRUE(context->shutdown().ok()); +} + +TEST(UbTeardown, ClearRetainsLocalAndImportedSegmentsForRetry) { + auto adapter = std::make_shared(); + auto context = makeActiveContext(adapter, 0); + UbBufferManager manager(adapter, {context}); + auto local = makeLocalBuffer(0x20000); + ASSERT_TRUE(manager.addBuffer(local, MemoryOptions{}).ok()); + + UbBufferMetadata metadata; + metadata.generation = 7; + metadata.base = 0x1000; + metadata.length = 0x1000; + metadata.permission = kGlobalReadWrite; + metadata.segments.push_back(UbBufferSegmentMetadata{ + 9, "remote", makeDevice(9).eid, 9, + makeDescriptor(metadata.base, metadata.length)}); + BufferDesc remote; + remote.addr = metadata.base; + remote.length = metadata.length; + ASSERT_TRUE(encodeBufferMetadata(metadata, + remote.transport_attrs[TransportType::UB]) + .ok()); + ImportedSegmentRef imported; + ASSERT_TRUE(manager + .importRemote(123, 0, 9, remote, Request::READ, + metadata.base, 64, imported) + .ok()); + imported.segment.reset(); + ASSERT_EQ(manager.importedSegmentCount(), 1u); + + adapter->unregister_failures = 1; + adapter->unimport_failures = 1; + EXPECT_FALSE(manager.clear().ok()); + EXPECT_EQ(manager.localBufferCount(), 1u); + EXPECT_EQ(manager.importedSegmentCount(), 1u); + EXPECT_EQ(adapter->live_local_segments, 1); + EXPECT_EQ(adapter->live_remote_segments, 1); + + EXPECT_TRUE(manager.clear().ok()); + EXPECT_EQ(manager.localBufferCount(), 0u); + EXPECT_EQ(manager.importedSegmentCount(), 0u); + EXPECT_EQ(adapter->live_local_segments, 0); + EXPECT_EQ(adapter->live_remote_segments, 0); + EXPECT_TRUE(context->shutdown().ok()); +} + +TEST(UbTeardown, NewGenerationProceedsWhileStaleImportCleanupRetries) { + auto adapter = std::make_shared(); + auto context = makeActiveContext(adapter, 0); + UbBufferManager manager(adapter, {context}); + + UbBufferMetadata metadata; + metadata.generation = 7; + metadata.base = 0x1000; + metadata.length = 0x1000; + metadata.permission = kGlobalReadWrite; + metadata.segments.push_back(UbBufferSegmentMetadata{ + 9, "remote", makeDevice(9).eid, 9, + makeDescriptor(metadata.base, metadata.length)}); + BufferDesc remote; + remote.addr = metadata.base; + remote.length = metadata.length; + ASSERT_TRUE(encodeBufferMetadata(metadata, + remote.transport_attrs[TransportType::UB]) + .ok()); + + ImportedSegmentRef old_generation; + ASSERT_TRUE(manager + .importRemote(123, 0, 9, remote, Request::READ, + metadata.base, 64, old_generation) + .ok()); + metadata.generation = 8; + ASSERT_TRUE(encodeBufferMetadata(metadata, + remote.transport_attrs[TransportType::UB]) + .ok()); + adapter->unimport_failures = 1; + ImportedSegmentRef new_generation; + EXPECT_TRUE(manager + .importRemote(123, 0, 9, remote, Request::READ, + metadata.base, 64, new_generation) + .ok()); + EXPECT_EQ(new_generation.generation, 8u); + EXPECT_EQ(manager.importedSegmentCount(), 2u); + EXPECT_EQ(adapter->live_remote_segments, 2); + + old_generation.segment.reset(); + ImportedSegmentRef reused; + EXPECT_TRUE(manager + .importRemote(123, 0, 9, remote, Request::READ, + metadata.base, 64, reused) + .ok()); + EXPECT_EQ(reused.segment, new_generation.segment); + EXPECT_EQ(manager.importedSegmentCount(), 1u); + EXPECT_EQ(adapter->live_remote_segments, 1); + + reused.segment.reset(); + new_generation.segment.reset(); + EXPECT_TRUE(manager.clear().ok()); + EXPECT_EQ(adapter->live_remote_segments, 0); + EXPECT_TRUE(context->shutdown().ok()); +} + +TEST(UbTeardown, FailedRegistrationRollbackIsRetainedUntilClear) { + auto adapter = std::make_shared(); + auto first = makeActiveContext(adapter, 0); + auto second = makeActiveContext(adapter, 1); + UbBufferManager manager(adapter, {first, second}); + auto desc = makeLocalBuffer(0x30000); + adapter->fail_register_call = 2; + adapter->partial_register_on_failure = true; + adapter->unregister_failures = 1; + + EXPECT_FALSE(manager.addBuffer(desc, MemoryOptions{}).ok()); + EXPECT_EQ(manager.localBufferCount(), 0u); + EXPECT_EQ(adapter->live_local_segments, 1); + EXPECT_TRUE(manager.clear().ok()); + EXPECT_EQ(adapter->live_local_segments, 0); + EXPECT_TRUE(first->shutdown().ok()); + EXPECT_TRUE(second->shutdown().ok()); +} + +TEST(UbTeardown, FailedPartialImportIsRetainedUntilClear) { + auto adapter = std::make_shared(); + auto context = makeActiveContext(adapter, 0); + UbBufferManager manager(adapter, {context}); + + UbBufferMetadata metadata; + metadata.generation = 11; + metadata.base = 0x1000; + metadata.length = 0x1000; + metadata.permission = kGlobalReadWrite; + metadata.segments.push_back(UbBufferSegmentMetadata{ + 9, "remote", makeDevice(9).eid, 9, + makeDescriptor(metadata.base, metadata.length)}); + BufferDesc remote; + remote.addr = metadata.base; + remote.length = metadata.length; + ASSERT_TRUE(encodeBufferMetadata(metadata, + remote.transport_attrs[TransportType::UB]) + .ok()); + + adapter->import_failures = 1; + adapter->partial_import_on_failure = true; + adapter->unimport_failures = 1; + ImportedSegmentRef imported; + EXPECT_FALSE(manager + .importRemote(456, 0, 9, remote, Request::READ, + metadata.base, 64, imported) + .ok()); + EXPECT_EQ(manager.importedSegmentCount(), 0u); + EXPECT_EQ(adapter->live_remote_segments, 1); + + EXPECT_TRUE(manager.clear().ok()); + EXPECT_EQ(adapter->live_remote_segments, 0); + EXPECT_TRUE(context->shutdown().ok()); +} + +TEST(UbTeardown, ContextShutdownRetainsFailedJfcBeforeClosingContext) { + auto adapter = std::make_shared(); + auto context = makeActiveContext(adapter, 0); + adapter->delete_jfc_failures = 1; + + EXPECT_FALSE(context->shutdown().ok()); + EXPECT_EQ(context->state(), UbContext::State::kDraining); + EXPECT_EQ(context->jfcs().size(), 1u); + EXPECT_TRUE(context->handle()); + EXPECT_EQ(adapter->close_context_calls, 0); + EXPECT_EQ(adapter->live_jfcs, 1); + + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_EQ(context->state(), UbContext::State::kClosed); + EXPECT_TRUE(context->jfcs().empty()); + EXPECT_FALSE(context->handle()); + EXPECT_EQ(adapter->close_context_calls, 1); + EXPECT_EQ(adapter->live_jfcs, 0); + EXPECT_EQ(adapter->live_contexts, 0); +} + +TEST(UbTeardown, ContextShutdownRetainsContextAfterCloseFailure) { + auto adapter = std::make_shared(); + auto context = makeActiveContext(adapter, 0); + adapter->close_context_failures = 1; + + EXPECT_FALSE(context->shutdown().ok()); + EXPECT_EQ(context->state(), UbContext::State::kDraining); + EXPECT_TRUE(context->jfcs().empty()); + EXPECT_TRUE(context->handle()); + EXPECT_EQ(adapter->live_contexts, 1); + + EXPECT_TRUE(context->shutdown().ok()); + EXPECT_EQ(context->state(), UbContext::State::kClosed); + EXPECT_FALSE(context->handle()); + EXPECT_EQ(adapter->live_contexts, 0); +} + +} // namespace +} // namespace mooncake::tent::ub diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 028855f1b6..aa00a168e1 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -1,158 +1,296 @@ set(WORKSPACE "${CMAKE_CURRENT_SOURCE_DIR}") -if (USE_HIP) +if(USE_HIP) file(GLOB TEST_SOURCES "*.cpp") hipify_files(TEST_SOURCES) - file(RELATIVE_PATH EXAMPLE_REL_PATH "${CMAKE_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}") + file(RELATIVE_PATH EXAMPLE_REL_PATH "${CMAKE_SOURCE_DIR}" + "${CMAKE_CURRENT_SOURCE_DIR}") set(WORKSPACE "${CMAKE_BINARY_DIR}/${EXAMPLE_REL_PATH}") endif() add_executable(rdma_transport_test ${WORKSPACE}/rdma_transport_test.cpp) -target_link_libraries(rdma_transport_test PUBLIC transfer_engine gflags::gflags glog::glog) +target_link_libraries(rdma_transport_test PUBLIC transfer_engine gflags::gflags + glog::glog) # add_test(NAME rdma_transport_test COMMAND rdma_transport_test) add_executable(transport_uint_test ${WORKSPACE}/transport_uint_test.cpp) -target_link_libraries(transport_uint_test PUBLIC transfer_engine gtest gtest_main ) +target_link_libraries(transport_uint_test PUBLIC transfer_engine gtest + gtest_main) add_test(NAME transport_uint_test COMMAND transport_uint_test) add_executable(endpoint_store_test ${WORKSPACE}/endpoint_store_test.cpp) -target_link_libraries(endpoint_store_test PUBLIC transfer_engine gtest gtest_main) +target_link_libraries(endpoint_store_test PUBLIC transfer_engine gtest + gtest_main) add_test(NAME endpoint_store_test COMMAND endpoint_store_test) -# Integration test for the monitorWorker reclaim tick (issue #1845). -# Self-skips when no RDMA device is present, so safe to register with ctest. -add_executable(endpoint_store_integration_test ${WORKSPACE}/endpoint_store_integration_test.cpp) -target_link_libraries(endpoint_store_integration_test PUBLIC transfer_engine gtest gtest_main) -add_test(NAME endpoint_store_integration_test COMMAND endpoint_store_integration_test) +add_executable(rdma_endpoint_state_test + ${WORKSPACE}/rdma_endpoint_state_test.cpp) +target_link_libraries(rdma_endpoint_state_test PUBLIC transfer_engine gtest + gtest_main) +add_test(NAME rdma_endpoint_state_test COMMAND rdma_endpoint_state_test) + +# Integration test for the monitorWorker reclaim tick (issue #1845). Self-skips +# when no RDMA device is present, so safe to register with ctest. +add_executable(endpoint_store_integration_test + ${WORKSPACE}/endpoint_store_integration_test.cpp) +target_link_libraries(endpoint_store_integration_test PUBLIC transfer_engine + gtest gtest_main) +add_test(NAME endpoint_store_integration_test + COMMAND endpoint_store_integration_test) set_tests_properties(endpoint_store_integration_test PROPERTIES LABELS "rdma") add_executable(rdma_transport_test2 ${WORKSPACE}/rdma_transport_test2.cpp) -target_link_libraries(rdma_transport_test2 PUBLIC transfer_engine gtest gtest_main ) +target_link_libraries(rdma_transport_test2 PUBLIC transfer_engine gtest + gtest_main) # add_test(NAME rdma_transport_test2 COMMAND rdma_transport_test2) add_executable(rdma_loopback_test ${WORKSPACE}/rdma_loopback_test.cpp) -target_link_libraries(rdma_loopback_test PUBLIC transfer_engine gtest gtest_main ) +target_link_libraries(rdma_loopback_test PUBLIC transfer_engine gtest + gtest_main) # add_test(NAME rdma_loopback_test COMMAND rdma_loopback_test) +# Regression test for #2017 (registerLocalMemory must auto-chunk buffers larger +# than the device max_mr_size; loopback WRITE past the boundary must succeed). +add_executable(rdma_large_mr_test ${WORKSPACE}/rdma_large_mr_test.cpp) +target_link_libraries(rdma_large_mr_test PUBLIC transfer_engine gtest gtest_main ) +# add_test(NAME rdma_large_mr_test COMMAND rdma_large_mr_test) # needs an RDMA dev + metadata server + # This test verifies endpoint re-establishment in RDMATransport. -add_executable(rdma_endpoint_reestablish_test ${WORKSPACE}/rdma_endpoint_reestablish_test.cpp) -target_link_libraries(rdma_endpoint_reestablish_test PUBLIC transfer_engine gtest gtest_main ) -if (UNIX AND NOT APPLE) - target_link_options(rdma_endpoint_reestablish_test PRIVATE - "-Wl,--wrap=ibv_modify_qp" - "-Wl,--wrap=ibv_query_gid" - "-Wl,--wrap=_ibv_query_gid_ex") -endif() -add_test(NAME rdma_endpoint_reestablish_test COMMAND rdma_endpoint_reestablish_test) +add_executable(rdma_endpoint_reestablish_test + ${WORKSPACE}/rdma_endpoint_reestablish_test.cpp) +target_link_libraries(rdma_endpoint_reestablish_test PUBLIC transfer_engine + gtest gtest_main) +if(UNIX AND NOT APPLE) + target_link_options( + rdma_endpoint_reestablish_test PRIVATE "-Wl,--wrap=ibv_modify_qp" + "-Wl,--wrap=ibv_query_gid" "-Wl,--wrap=_ibv_query_gid_ex") +endif() +add_test(NAME rdma_endpoint_reestablish_test + COMMAND rdma_endpoint_reestablish_test) set_tests_properties(rdma_endpoint_reestablish_test PROPERTIES LABELS "rdma") -if (USE_CXL) - add_executable(cxl_transport_test ${WORKSPACE}/cxl_transport_test.cpp) - target_link_libraries(cxl_transport_test PUBLIC transfer_engine gtest gtest_main ) - add_test(NAME cxl_transport_test COMMAND cxl_transport_test) +# Regression test for the edge-triggered async event fd: one epoll wakeup must +# drain the whole queue. Wraps ibv_get_async_event to script the event source, +# so it needs no RDMA device. -Wl,--wrap is a GNU ld / lld feature that Apple's +# linker lacks, and the test is meaningless without it, so skip the target there +# rather than build one that calls the real symbols. +if(UNIX AND NOT APPLE) + add_executable(rdma_async_event_drain_test + ${WORKSPACE}/rdma_async_event_drain_test.cpp) + target_link_libraries(rdma_async_event_drain_test PUBLIC transfer_engine + gtest gtest_main) + target_link_options( + rdma_async_event_drain_test PRIVATE "-Wl,--wrap=ibv_get_async_event" + "-Wl,--wrap=ibv_ack_async_event") + add_test(NAME rdma_async_event_drain_test COMMAND rdma_async_event_drain_test) endif() -if (USE_NVMEOF) - add_executable(nvmeof_transport_test ${WORKSPACE}/nvmeof_transport_test.cpp) - target_link_libraries(nvmeof_transport_test PUBLIC transfer_engine gtest gtest_main ) - # add_test(NAME nvmeof_transport_test COMMAND nvmeof_transport_test) +if(USE_CXL) + add_executable(cxl_transport_test ${WORKSPACE}/cxl_transport_test.cpp) + target_link_libraries(cxl_transport_test PUBLIC transfer_engine gtest + gtest_main) + add_test(NAME cxl_transport_test COMMAND cxl_transport_test) +endif() + +if(USE_NVMEOF) + add_executable(nvmeof_status_test ${WORKSPACE}/nvmeof_status_test.cpp) + target_link_libraries(nvmeof_status_test PUBLIC transfer_engine gtest + gtest_main) + add_test(NAME nvmeof_status_test COMMAND nvmeof_status_test) + + add_executable(nvmeof_transport_test ${WORKSPACE}/nvmeof_transport_test.cpp) + target_link_libraries(nvmeof_transport_test PUBLIC transfer_engine gtest + gtest_main) + # add_test(NAME nvmeof_transport_test COMMAND nvmeof_transport_test) endif() -if (USE_TCP) -add_executable(tcp_transport_test ${WORKSPACE}/tcp_transport_test.cpp) -target_link_libraries(tcp_transport_test PUBLIC transfer_engine gtest gtest_main ) -add_test(NAME tcp_transport_test COMMAND tcp_transport_test) +if(USE_TCP) + add_executable(tcp_transport_test ${WORKSPACE}/tcp_transport_test.cpp) + target_link_libraries(tcp_transport_test PUBLIC transfer_engine gtest + gtest_main) + add_test(NAME tcp_transport_test COMMAND tcp_transport_test) + + add_executable(tcp_write_visibility_test + ${WORKSPACE}/tcp_write_visibility_test.cpp) + target_link_libraries(tcp_write_visibility_test PUBLIC transfer_engine gtest + gtest_main) + add_test(NAME tcp_write_visibility_test COMMAND tcp_write_visibility_test) endif() -add_executable(tcp_address_validation_test ${WORKSPACE}/tcp_address_validation_test.cpp) +add_executable(tcp_address_validation_test + ${WORKSPACE}/tcp_address_validation_test.cpp) target_link_libraries(tcp_address_validation_test PUBLIC gtest gtest_main) add_test(NAME tcp_address_validation_test COMMAND tcp_address_validation_test) -if (USE_MNNVL) - add_executable(nvlink_transport_test ${WORKSPACE}/nvlink_transport_test.cpp) - target_link_libraries(nvlink_transport_test PUBLIC transfer_engine gtest gtest_main ) - add_test(NAME nvlink_transport_test COMMAND nvlink_transport_test) -endif() - -if (USE_UBSHMEM) - add_executable(ubshmem_transport_test ${WORKSPACE}/ubshmem_transport_test.cpp) - target_link_libraries(ubshmem_transport_test PUBLIC transfer_engine gtest gtest_main ascendcl) - add_test(NAME ubshmem_transport_test COMMAND ubshmem_transport_test) -endif() - -if (USE_EFA) - add_executable(efa_transport_test ${WORKSPACE}/efa_transport_test.cpp) - target_link_libraries(efa_transport_test PUBLIC transfer_engine gtest gtest_main) - add_test(NAME efa_transport_test COMMAND efa_transport_test) - - add_executable(efa_c_api_test ${WORKSPACE}/efa_c_api_test.cpp) - target_link_libraries(efa_c_api_test PUBLIC transfer_engine gtest gtest_main) - add_test(NAME efa_c_api_test COMMAND efa_c_api_test) - - add_executable(efa_single_nic_large_mr_test ${WORKSPACE}/efa_single_nic_large_mr_test.cpp) - target_link_libraries(efa_single_nic_large_mr_test PUBLIC transfer_engine gflags::gflags glog::glog) - - add_executable(efa_transfer_test ${WORKSPACE}/efa_transfer_test.cpp) - target_link_libraries(efa_transfer_test PUBLIC transfer_engine gflags::gflags glog::glog) - - # GPU (CUDA device memory) loopback test — reproduces the EFA SHM - # intra-node segfault on FI_HMEM_CUDA buffers (ofiwg/libfabric#12328) - # and validates EfaContext::tryLoopbackCopy. Needs CUDA headers/libs. - if (USE_CUDA) - add_executable(efa_gpu_loopback_test ${WORKSPACE}/efa_gpu_loopback_test.cpp) - # Resolve CUDA include dirs / runtime via CUDAToolkit instead of a - # hardcoded /usr/local/cuda/include, so the test builds wherever CUDA - # lives (e.g. the DLAMI pip venv layout). find_package is idempotent - # and may not have run yet in this scope (top-level only calls it under - # WITH_EP), so request it here; fall back to the legacy path if the - # module variant is unavailable. - find_package(CUDAToolkit QUIET) - if (CUDAToolkit_FOUND) - target_include_directories(efa_gpu_loopback_test - PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) - target_link_libraries(efa_gpu_loopback_test - PUBLIC transfer_engine gtest gtest_main CUDA::cudart) - else() - target_include_directories(efa_gpu_loopback_test - PRIVATE /usr/local/cuda/include) - target_link_libraries(efa_gpu_loopback_test - PUBLIC transfer_engine gtest gtest_main cudart) - endif() - add_test(NAME efa_gpu_loopback_test COMMAND efa_gpu_loopback_test) - endif() +# Hardware-free unit test for the active-connect circuit-breaker state +# (ConnectPauseTracker is header-only and clock-injectable), runs on every CI +# runner. +add_executable(connect_pause_tracker_test + ${WORKSPACE}/connect_pause_tracker_test.cpp) +target_link_libraries(connect_pause_tracker_test PUBLIC gtest gtest_main + pthread) +add_test(NAME connect_pause_tracker_test COMMAND connect_pause_tracker_test) + +# Hardware-free unit tests for DmabufExport struct and +# RdmaContext::{exportDmabuf, closeDmabufExport}. Runs on every CI runner — +# no RDMA device or GPU required (host-memory paths only). +add_executable(dmabuf_export_test ${WORKSPACE}/dmabuf_export_test.cpp) +target_link_libraries(dmabuf_export_test PUBLIC transfer_engine gtest gtest_main) +add_test(NAME dmabuf_export_test COMMAND dmabuf_export_test) + +if(USE_MNNVL) + add_executable(nvlink_transport_test ${WORKSPACE}/nvlink_transport_test.cpp) + target_link_libraries(nvlink_transport_test PUBLIC transfer_engine gtest + gtest_main) + add_test(NAME nvlink_transport_test COMMAND nvlink_transport_test) endif() -# UB transport test with URMA endpoint and mock support -if (USE_UB) - add_executable(ub_transport_test ${WORKSPACE}/ub_transport_test.cpp) - target_link_libraries(ub_transport_test PUBLIC transfer_engine gtest gtest_main glog::glog pthread) - target_include_directories(ub_transport_test PRIVATE ${urma_INCLUDE_DIR}) - # Built but not registered with ctest: this test may still have race - # conditions and other stability issues, so keep it out of CI for now. - # Run manually with ./ub_transport_test. - # add_test(NAME ub_transport_test COMMAND ub_transport_test) +if(USE_HIP) + add_executable(hip_transport_test ${WORKSPACE}/hip_transport_test.cpp) + target_link_libraries(hip_transport_test PUBLIC transfer_engine gtest + gtest_main) + add_test(NAME hip_transport_test COMMAND hip_transport_test) endif() -if (USE_SUNRISE) - add_executable(sunrise_link_transport_test ${WORKSPACE}/sunrise_link_transport_test.cpp) - target_include_directories(sunrise_link_transport_test PRIVATE ${MC_TANGRT_ROOT}/include) - target_link_libraries(sunrise_link_transport_test PUBLIC transfer_engine gtest gtest_main) - add_test(NAME sunrise_link_transport_test COMMAND sunrise_link_transport_test) +if(USE_UBSHMEM) + add_executable(ubshmem_transport_test ${WORKSPACE}/ubshmem_transport_test.cpp) + target_link_libraries(ubshmem_transport_test PUBLIC transfer_engine gtest + gtest_main ascendcl) + add_test(NAME ubshmem_transport_test COMMAND ubshmem_transport_test) +endif() - add_executable(sunrise_link_transport_runtime_test ${WORKSPACE}/sunrise_link_transport_runtime_test.cpp) - target_include_directories(sunrise_link_transport_runtime_test PRIVATE ${MC_TANGRT_ROOT}/include) - target_link_libraries(sunrise_link_transport_runtime_test PUBLIC transfer_engine gtest gtest_main) - add_test(NAME sunrise_link_transport_runtime_test COMMAND sunrise_link_transport_runtime_test) +if(USE_EFA) + add_executable(efa_transport_test ${WORKSPACE}/efa_transport_test.cpp) + target_link_libraries(efa_transport_test PUBLIC transfer_engine gtest + gtest_main) + add_test(NAME efa_transport_test COMMAND efa_transport_test) + + add_executable(efa_c_api_test ${WORKSPACE}/efa_c_api_test.cpp) + target_link_libraries(efa_c_api_test PUBLIC transfer_engine gtest gtest_main) + add_test(NAME efa_c_api_test COMMAND efa_c_api_test) + + add_executable(efa_single_nic_large_mr_test + ${WORKSPACE}/efa_single_nic_large_mr_test.cpp) + target_link_libraries(efa_single_nic_large_mr_test + PUBLIC transfer_engine gflags::gflags glog::glog) + + add_executable(efa_transfer_test ${WORKSPACE}/efa_transfer_test.cpp) + target_link_libraries(efa_transfer_test PUBLIC transfer_engine gflags::gflags + glog::glog) + + # GPU (CUDA device memory) loopback test — reproduces the EFA SHM intra-node + # segfault on FI_HMEM_CUDA buffers (ofiwg/libfabric#12328) and validates + # EfaContext::tryLoopbackCopy. Needs CUDA headers/libs. + if(USE_CUDA) + add_executable(efa_gpu_loopback_test ${WORKSPACE}/efa_gpu_loopback_test.cpp) + # Resolve CUDA include dirs / runtime via CUDAToolkit instead of a hardcoded + # /usr/local/cuda/include, so the test builds wherever CUDA lives (e.g. the + # DLAMI pip venv layout). find_package is idempotent and may not have run + # yet in this scope (top-level only calls it under WITH_EP), so request it + # here; fall back to the legacy path if the module variant is unavailable. + find_package(CUDAToolkit QUIET) + if(CUDAToolkit_FOUND) + target_include_directories(efa_gpu_loopback_test + PRIVATE ${CUDAToolkit_INCLUDE_DIRS}) + target_link_libraries( + efa_gpu_loopback_test PUBLIC transfer_engine gtest gtest_main + CUDA::cudart) + else() + target_include_directories(efa_gpu_loopback_test + PRIVATE /usr/local/cuda/include) + target_link_libraries(efa_gpu_loopback_test PUBLIC transfer_engine gtest + gtest_main cudart) + endif() + add_test(NAME efa_gpu_loopback_test COMMAND efa_gpu_loopback_test) + endif() +endif() - add_executable(sunrise_link_transport_unit_test ${WORKSPACE}/sunrise_link_transport_unit_test.cpp) - target_include_directories(sunrise_link_transport_unit_test PRIVATE ${MC_TANGRT_ROOT}/include) - target_link_libraries(sunrise_link_transport_unit_test PUBLIC transfer_engine gtest gtest_main) - add_test(NAME sunrise_link_transport_unit_test COMMAND sunrise_link_transport_unit_test) +if(USE_CXI) + add_executable(cxi_transport_test ${WORKSPACE}/cxi_transport_test.cpp) + add_executable(cxi_transfer_test ${WORKSPACE}/cxi_transfer_test.cpp) + add_executable(cxi_unit_tests ${WORKSPACE}/cxi_unit_tests.cpp) + target_link_libraries(cxi_transport_test PUBLIC transfer_engine gtest + gtest_main) + target_link_libraries(cxi_transfer_test PUBLIC transfer_engine gtest + gtest_main) + target_link_libraries(cxi_unit_tests PUBLIC transfer_engine gtest gtest_main) + add_test(NAME cxi_transport_test COMMAND cxi_transport_test) + add_test(NAME cxi_unit_tests COMMAND cxi_unit_test) +endif() + +# UB transport test with URMA endpoint and mock support +if(USE_UB) + add_executable(ub_transport_test ${WORKSPACE}/ub_transport_test.cpp) + target_link_libraries(ub_transport_test PUBLIC transfer_engine gtest + gtest_main glog::glog pthread) + target_include_directories(ub_transport_test PRIVATE ${urma_INCLUDE_DIR}) + # Built but not registered with ctest: this test may still have race + # conditions and other stability issues, so keep it out of CI for now. Run + # manually with ./ub_transport_test. add_test(NAME ub_transport_test COMMAND + # ub_transport_test) +endif() + +if(USE_SUNRISE) + add_executable(sunrise_link_transport_test + ${WORKSPACE}/sunrise_link_transport_test.cpp) + target_include_directories(sunrise_link_transport_test + PRIVATE ${MC_TANGRT_ROOT}/include) + target_link_libraries( + sunrise_link_transport_test + PUBLIC transfer_engine gtest gtest_main + ${MC_TANGRT_ROOT}/lib/libtangrt_shared.so + ${MC_TANGRT_ROOT}/lib/libptml_shared.so dl) + add_test(NAME sunrise_link_transport_test COMMAND sunrise_link_transport_test) + + add_executable(sunrise_link_transport_runtime_test + ${WORKSPACE}/sunrise_link_transport_runtime_test.cpp) + target_include_directories(sunrise_link_transport_runtime_test + PRIVATE ${MC_TANGRT_ROOT}/include) + target_link_libraries( + sunrise_link_transport_runtime_test + PUBLIC transfer_engine gtest gtest_main + ${MC_TANGRT_ROOT}/lib/libtangrt_shared.so + ${MC_TANGRT_ROOT}/lib/libptml_shared.so dl) + add_test(NAME sunrise_link_transport_runtime_test + COMMAND sunrise_link_transport_runtime_test) + + add_executable(sunrise_link_transport_unit_test + ${WORKSPACE}/sunrise_link_transport_unit_test.cpp) + target_include_directories(sunrise_link_transport_unit_test + PRIVATE ${MC_TANGRT_ROOT}/include) + target_link_libraries( + sunrise_link_transport_unit_test + PUBLIC transfer_engine gtest gtest_main + ${MC_TANGRT_ROOT}/lib/libtangrt_shared.so + ${MC_TANGRT_ROOT}/lib/libptml_shared.so dl) + add_test(NAME sunrise_link_transport_unit_test + COMMAND sunrise_link_transport_unit_test) + + add_executable(sunrise_allocator_test ${WORKSPACE}/sunrise_allocator_test.cpp) + target_include_directories(sunrise_allocator_test + PRIVATE ${MC_TANGRT_ROOT}/include) + target_link_libraries( + sunrise_allocator_test + PUBLIC transfer_engine gtest gtest_main + ${MC_TANGRT_ROOT}/lib/libtangrt_shared.so + ${MC_TANGRT_ROOT}/lib/libptml_shared.so dl) + add_test(NAME sunrise_allocator_test COMMAND sunrise_allocator_test) + + add_executable(sunrise_link_copy_test ${WORKSPACE}/sunrise_link_copy_test.cpp) + target_include_directories(sunrise_link_copy_test + PRIVATE ${MC_TANGRT_ROOT}/include) + target_link_libraries( + sunrise_link_copy_test + PUBLIC transfer_engine gtest gtest_main + ${MC_TANGRT_ROOT}/lib/libtangrt_shared.so + ${MC_TANGRT_ROOT}/lib/libptml_shared.so dl) + add_test(NAME sunrise_link_copy_test COMMAND sunrise_link_copy_test) endif() add_executable(transfer_metadata_test ${WORKSPACE}/transfer_metadata_test.cpp) -target_link_libraries(transfer_metadata_test PUBLIC transfer_engine gtest gtest_main) +target_link_libraries(transfer_metadata_test PUBLIC transfer_engine gtest + gtest_main) add_test(NAME transfer_metadata_test COMMAND transfer_metadata_test) add_executable(config_test ${WORKSPACE}/config_test.cpp) @@ -160,48 +298,104 @@ target_link_libraries(config_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME config_test COMMAND config_test) add_executable(rdma_gid_probe_test ${WORKSPACE}/rdma_gid_probe_test.cpp) -target_link_libraries(rdma_gid_probe_test PUBLIC transfer_engine gtest gtest_main) +target_link_libraries(rdma_gid_probe_test PUBLIC transfer_engine gtest + gtest_main) add_test(NAME rdma_gid_probe_test COMMAND rdma_gid_probe_test) -add_executable(rdma_context_reprobe_test ${WORKSPACE}/rdma_context_reprobe_test.cpp) -target_link_libraries(rdma_context_reprobe_test PUBLIC transfer_engine gtest gtest_main) +add_executable(multi_transport_locality_test + ${WORKSPACE}/multi_transport_locality_test.cpp) +target_link_libraries(multi_transport_locality_test PUBLIC transfer_engine + gtest gtest_main) +add_test(NAME multi_transport_locality_test + COMMAND multi_transport_locality_test) + +add_executable(rdma_context_reprobe_test + ${WORKSPACE}/rdma_context_reprobe_test.cpp) +target_link_libraries(rdma_context_reprobe_test PUBLIC transfer_engine gtest + gtest_main) add_test(NAME rdma_context_reprobe_test COMMAND rdma_context_reprobe_test) +add_executable(rdma_transport_submit_task_test + ${WORKSPACE}/rdma_transport_submit_task_test.cpp) +target_link_libraries(rdma_transport_submit_task_test PUBLIC transfer_engine + gtest gtest_main) +add_test(NAME rdma_transport_submit_task_test + COMMAND rdma_transport_submit_task_test) + add_executable(topology_test ${WORKSPACE}/topology_test.cpp) target_link_libraries(topology_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME topology_test COMMAND topology_test) add_executable(memory_location_test ${WORKSPACE}/memory_location_test.cpp) -target_link_libraries(memory_location_test PUBLIC transfer_engine gtest gtest_main) +target_link_libraries(memory_location_test PUBLIC transfer_engine gtest + gtest_main) add_test(NAME memory_location_test COMMAND memory_location_test) add_executable(common_test ${WORKSPACE}/common_test.cpp) target_link_libraries(common_test PUBLIC transfer_engine gtest gtest_main) add_test(NAME common_test COMMAND common_test) -if (USE_ASCEND_DIRECT) - # AscendDirectTransport unit test with mock implementations - # Mock implementations are included in the test file via anonymous namespace - add_executable(ascend_direct_transport_test - ${WORKSPACE}/ascend_direct_transport_test.cpp - ) - target_link_libraries(ascend_direct_transport_test PUBLIC - transfer_engine - gtest - gtest_main - metadef - ) - # Allow shared library undefined symbols (Ascend libraries have circular deps) - set_target_properties(ascend_direct_transport_test PROPERTIES - LINK_FLAGS "-Wl,--allow-shlib-undefined" - ) - - add_test(NAME ascend_direct_transport_test COMMAND ascend_direct_transport_test) +if(USE_ASCEND_DIRECT) + # AscendDirectTransport unit test with mock implementations Mock + # implementations are included in the test file via anonymous namespace + add_executable(ascend_direct_transport_test + ${WORKSPACE}/ascend_direct_transport_test.cpp) + target_link_libraries(ascend_direct_transport_test + PUBLIC transfer_engine gtest gtest_main metadef) + # Allow shared library undefined symbols (Ascend libraries have circular deps) + set_target_properties(ascend_direct_transport_test + PROPERTIES LINK_FLAGS "-Wl,--allow-shlib-undefined") + + add_test(NAME ascend_direct_transport_test + COMMAND ascend_direct_transport_test) endif() # Multi-protocol transport test (only when ENABLE_MULTI_PROTOCOL is ON) -if (ENABLE_MULTI_PROTOCOL) - add_executable(mp_transport_test ${WORKSPACE}/mp_transport_test.cpp) - target_link_libraries(mp_transport_test PUBLIC transfer_engine gtest gtest_main) - add_test(NAME mp_transport_test COMMAND mp_transport_test) +if(ENABLE_MULTI_PROTOCOL) + add_executable(mp_transport_test ${WORKSPACE}/mp_transport_test.cpp) + target_link_libraries(mp_transport_test PUBLIC transfer_engine gtest + gtest_main) + add_test(NAME mp_transport_test COMMAND mp_transport_test) +endif() + +add_executable(graceful_shutdown_test ${WORKSPACE}/graceful_shutdown_test.cpp) +target_link_libraries(graceful_shutdown_test PUBLIC transfer_engine gtest + gtest_main) +add_test(NAME graceful_shutdown_test COMMAND graceful_shutdown_test) + +add_executable(show_links_test ${WORKSPACE}/show_links_test.cpp) +target_link_libraries(show_links_test PUBLIC transfer_engine gtest gtest_main) +add_test(NAME show_links_test COMMAND show_links_test) + +# The NCCL transport library target compiles only host-side setup, while the +# CUDA example that includes the device helpers is opt-in. This probe instantiates +# one consumer kernel against both header-only NCCL and IBGDA helpers so a +# call-shape regression fails the build. Keep nvcc optional for host-only +# DeviceTransport builds that consume JIT kernels. +if(USE_NCCL_DEVICE AND NOT USE_CXI) + if(CMAKE_VERSION VERSION_LESS 3.18) + message(STATUS + "Skipping NCCL/IBGDA device API compatibility compile test: " + "CMake 3.18+ is required for CUDA C++17") + else() + include(CheckLanguage) + check_language(CUDA) + if(CMAKE_CUDA_COMPILER) + enable_language(CUDA) + add_library(device_backend_api_compatibility_test OBJECT + ${WORKSPACE}/device_backend_api_compatibility_test.cu) + target_include_directories(device_backend_api_compatibility_test PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../include) + target_link_libraries(device_backend_api_compatibility_test PRIVATE + NCCL::nccl) + target_compile_options(device_backend_api_compatibility_test PRIVATE + $<$:--expt-relaxed-constexpr>) + set_target_properties(device_backend_api_compatibility_test PROPERTIES + CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON CUDA_EXTENSIONS OFF) + else() + message(STATUS + "Skipping NCCL/IBGDA device API compatibility compile test: " + "nvcc not found") + endif() + endif() endif() diff --git a/mooncake-transfer-engine/tests/ascend_direct_transport_test.cpp b/mooncake-transfer-engine/tests/ascend_direct_transport_test.cpp index 46c96c1b0d..071b33f47b 100644 --- a/mooncake-transfer-engine/tests/ascend_direct_transport_test.cpp +++ b/mooncake-transfer-engine/tests/ascend_direct_transport_test.cpp @@ -54,6 +54,9 @@ static int g_set_device_call_count = 0; static std::set g_set_device_ids; static std::map g_memory_locations; static std::mutex g_acl_mutex; +// 0 = no size limit; otherwise aclrtMallocPhysical fails when size > threshold. +static size_t g_malloc_physical_max_success = 0; +static int g_malloc_physical_call_count = 0; namespace mock_acl { void reset() { @@ -67,6 +70,18 @@ void reset() { g_set_device_call_count = 0; g_set_device_ids.clear(); g_memory_locations.clear(); + g_malloc_physical_max_success = 0; + g_malloc_physical_call_count = 0; +} + +void set_malloc_physical_max_success(size_t max_success) { + std::lock_guard lock(g_acl_mutex); + g_malloc_physical_max_success = max_success; +} + +int malloc_physical_call_count() { + std::lock_guard lock(g_acl_mutex); + return g_malloc_physical_call_count; } void set_device_count(int count) { @@ -256,9 +271,14 @@ aclError aclrtGetPhyDevIdByLogicDevId(int32_t logic_dev_id, aclError aclrtMallocPhysical(aclrtDrvMemHandle* handle, size_t size, aclrtPhysicalMemProp* prop, uint32_t flags) { - (void)size; (void)prop; (void)flags; + std::lock_guard lock(g_acl_mutex); + ++g_malloc_physical_call_count; + if (g_malloc_physical_max_success > 0 && + size > g_malloc_physical_max_success) { + return ACL_ERROR_FAILURE; + } *handle = reinterpret_cast(malloc(1)); return *handle ? ACL_ERROR_NONE : ACL_ERROR_FAILURE; } @@ -268,7 +288,11 @@ aclError aclrtReserveMemAddress(void** va, size_t size, size_t alignment, (void)alignment; (void)hint_addr; (void)page_type; - *va = malloc(size); + // Stub VA only — do not malloc(size) (best-effort tests use multi-GB + // sizes). + (void)size; + constexpr size_t kStubVaBytes = 64; + *va = malloc(kStubVaBytes); return *va ? ACL_ERROR_NONE : ACL_ERROR_FAILURE; } @@ -1753,6 +1777,94 @@ TEST_F(AscendDirectTransportTest, RemoteTransfer_Async_TransferTimeout) { "(GetTransferStatus stays WAITING)"; } +TEST_F(AscendDirectTransportTest, + RemoteTransfer_Sync_FailureWithAutoConnect_SkipsDisconnect) { + setenv("ASCEND_AUTO_CONNECT", "1", 1); + adxl_mock::set_transfer_result(adxl::FAILED); + + auto transport = createTransport(); + ASSERT_NE(transport, nullptr); + ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize, + "cpu:0", true, true), + 0); + addRemoteSegment(transport->meta(), + {1, "remote_server", "192.168.1.100", 30000}); + + initTestData(kTransferBufSize); + auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1, + 0x10000, kTransferBufSize); + ASSERT_TRUE(result.finished); + EXPECT_TRUE(result.failed); + EXPECT_EQ(adxl_mock::get_disconnect_count(), 0); +} + +TEST_F(AscendDirectTransportTest, + RemoteTransfer_Async_SubmitFailureWithAutoConnect_SkipsDisconnect) { + setenv("ASCEND_AUTO_CONNECT", "1", 1); + adxl_mock::set_transfer_async_result(adxl::FAILED); + + auto transport = createTransport(true); + ASSERT_NE(transport, nullptr); + ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize, + "cpu:0", true, true), + 0); + addRemoteSegment(transport->meta(), + {1, "remote_server", "192.168.1.100", 30000}); + + initTestData(kTransferBufSize); + auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1, + 0x10000, kTransferBufSize); + ASSERT_TRUE(result.finished); + EXPECT_TRUE(result.failed); + EXPECT_EQ(adxl_mock::get_disconnect_count(), 0); +} + +TEST_F(AscendDirectTransportTest, + RemoteTransfer_Async_GetStatusFailureWithAutoConnect_SkipsDisconnect) { + setenv("ASCEND_AUTO_CONNECT", "1", 1); + adxl_mock::set_transfer_async_result(adxl::SUCCESS); + adxl_mock::set_get_transfer_status_result(adxl::FAILED); + + auto transport = createTransport(true); + ASSERT_NE(transport, nullptr); + ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize, + "cpu:0", true, true), + 0); + addRemoteSegment(transport->meta(), + {1, "remote_server", "192.168.1.100", 30000}); + + initTestData(kTransferBufSize); + auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1, + 0x10000, kTransferBufSize); + ASSERT_TRUE(result.finished); + EXPECT_TRUE(result.failed); + EXPECT_EQ(adxl_mock::get_disconnect_count(), 0); +} + +TEST_F(AscendDirectTransportTest, + RemoteTransfer_Async_TimeoutWithAutoConnect_StillDisconnects) { + setenv("ASCEND_AUTO_CONNECT", "1", 1); + setenv("ASCEND_TRANSFER_TIMEOUT", "100", 1); + adxl_mock::set_transfer_async_result(adxl::SUCCESS); + adxl_mock::set_get_transfer_status_result(adxl::SUCCESS); + adxl_mock::set_transfer_status_enum(adxl::TransferStatus::WAITING); + + auto transport = createTransport(true); + ASSERT_NE(transport, nullptr); + ASSERT_EQ(transport->registerLocalMemory(test_buffer_src_, kRegisterMemSize, + "cpu:0", true, true), + 0); + addRemoteSegment(transport->meta(), + {1, "remote_server", "192.168.1.100", 30000}); + + initTestData(kTransferBufSize); + auto result = runRemoteTransfer(transport.get(), test_buffer_src_, 1, + 0x10000, kTransferBufSize); + ASSERT_TRUE(result.finished); + EXPECT_TRUE(result.failed); + EXPECT_EQ(adxl_mock::get_disconnect_count(), 1); +} + // ----------------------------------------------------------------------------- // Standalone mode tests (non-dummy-real, non-fabric-mem) // ----------------------------------------------------------------------------- @@ -1914,6 +2026,85 @@ TEST_F(AscendDirectTransportTest, // Roce mode detection (HCCL_INTRA_ROCE_ENABLE / ASCEND_GLOBAL_RESOURCE_CONFIG) // ----------------------------------------------------------------------------- +TEST(FabricMemBestEffortAllocTest, PercentileLadderFindsFeasibleSize) { + mock_acl::reset(); + constexpr size_t kGiB = 1024ULL * 1024 * 1024; + constexpr size_t kTargetGiB = 40; + constexpr size_t kMaxOkGiB = 35; + constexpr size_t kExpectGiB = 32; // 80% of 40GiB after 100%/90% fail + constexpr size_t kTarget = kTargetGiB * kGiB; + mock_acl::set_malloc_physical_max_success(kMaxOkGiB * kGiB); + + globalConfig().ascend_use_fabric_mem = true; + size_t actual = 0; + void* ptr = ascend_allocate_memory_best_effort(kTarget, "ascend", &actual); + ASSERT_NE(ptr, nullptr); + EXPECT_EQ(actual, kExpectGiB * kGiB); + EXPECT_GT(mock_acl::malloc_physical_call_count(), 1); + + ascend_free_memory("ascend", ptr); + globalConfig().ascend_use_fabric_mem = false; + mock_acl::reset(); +} + +TEST(FabricMemBestEffortAllocTest, BelowFiftyPercentReturnsNull) { + mock_acl::reset(); + constexpr size_t kGiB = 1024ULL * 1024 * 1024; + constexpr size_t kTargetGiB = 40; + constexpr size_t kMaxOkGiB = 15; // below 50% of 40GiB (=20GiB) + mock_acl::set_malloc_physical_max_success(kMaxOkGiB * kGiB); + + globalConfig().ascend_use_fabric_mem = true; + size_t actual = 0; + void* ptr = ascend_allocate_memory_best_effort(kTargetGiB * kGiB, "ascend", + &actual); + EXPECT_EQ(ptr, nullptr); + EXPECT_EQ(actual, 0); + + globalConfig().ascend_use_fabric_mem = false; + mock_acl::reset(); +} + +// 2.1 GiB target: 50% = 1.05 GiB. Align-down of lower percentiles yields 1 GiB +// (< 50%), which must be rejected even if physical alloc of 1 GiB would +// succeed. +TEST(FabricMemBestEffortAllocTest, AlignDownBelowMinPercentIsRejected) { + mock_acl::reset(); + constexpr size_t kGiB = 1024ULL * 1024 * 1024; + constexpr size_t kTarget = (2 * kGiB) + (kGiB / 10); // 2.1 GiB + mock_acl::set_malloc_physical_max_success(kGiB); // 1 GiB ok, 2 GiB fails + + globalConfig().ascend_use_fabric_mem = true; + size_t actual = 0; + void* ptr = ascend_allocate_memory_best_effort(kTarget, "ascend", &actual); + EXPECT_EQ(ptr, nullptr); + EXPECT_EQ(actual, 0); + + globalConfig().ascend_use_fabric_mem = false; + mock_acl::reset(); +} + +TEST(FabricMemBestEffortAllocTest, FullTargetFirstSuccess) { + mock_acl::reset(); + constexpr size_t kGiB = 1024ULL * 1024 * 1024; + constexpr size_t kTargetGiB = 40; + constexpr size_t kTarget = kTargetGiB * kGiB; + // Physical alloc is tried up to 3 attribute tiers per size. + constexpr int kMaxPhysicalTiersPerSize = 3; + mock_acl::set_malloc_physical_max_success(0); // unlimited + + globalConfig().ascend_use_fabric_mem = true; + size_t actual = 0; + void* ptr = ascend_allocate_memory_best_effort(kTarget, "ascend", &actual); + ASSERT_NE(ptr, nullptr); + EXPECT_EQ(actual, kTarget); // 100% first + EXPECT_LE(mock_acl::malloc_physical_call_count(), kMaxPhysicalTiersPerSize); + + ascend_free_memory("ascend", ptr); + globalConfig().ascend_use_fabric_mem = false; + mock_acl::reset(); +} + TEST(RoceModeDetectionTest, GlobalResourceConfig_StringRoceDesc) { EXPECT_TRUE(HasRoceProtocolDescInGlobalResourceConfig( R"({"comm_resource_config.protocol_desc":"roce:device"})")); diff --git a/mooncake-transfer-engine/tests/common_test.cpp b/mooncake-transfer-engine/tests/common_test.cpp index b07971be36..b158b13c4b 100644 --- a/mooncake-transfer-engine/tests/common_test.cpp +++ b/mooncake-transfer-engine/tests/common_test.cpp @@ -1,16 +1,82 @@ #include +#include +#include #include #include #include +#include +#include #include "common.h" namespace { using namespace mooncake; +using namespace std::chrono_literals; const uint16_t kDefaultPort = getDefaultHandshakePort(); +//------------------------------------------------------------------------------ +// RWSpinlock +//------------------------------------------------------------------------------ + +TEST(RWSpinlockTest, AggressiveWriteLockProgressesAcrossTicketWraparound) { + RWSpinlock lock; + int protected_value = 0; + + constexpr int kIterations = 65536 + 3; + for (int i = 0; i < kIterations; ++i) { + lock.writeLockAggressive(); + ++protected_value; + lock.unlock(); + } + + RWSpinlock::ReadGuard guard(lock); + EXPECT_EQ(protected_value, kIterations); +} + +TEST(RWSpinlockTest, DowngradePublishesToReadersAndBlocksWriters) { + RWSpinlock lock; + int protected_value = 0; + std::atomic writer_started{false}; + std::atomic writer_entered{false}; + std::atomic reader_observed{false}; + + lock.writeLockAggressive(); + protected_value = 42; + lock.unlockAndLockShared(); + + std::thread reader([&] { + RWSpinlock::ReadGuard guard(lock); + reader_observed.store(protected_value == 42, std::memory_order_release); + }); + + reader.join(); + EXPECT_TRUE(reader_observed.load(std::memory_order_acquire)); + + std::thread writer([&] { + writer_started.store(true, std::memory_order_release); + lock.writeLockAggressive(); + writer_entered.store(true, std::memory_order_release); + protected_value = 99; + lock.unlock(); + }); + + while (!writer_started.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + + std::this_thread::sleep_for(10ms); + EXPECT_FALSE(writer_entered.load(std::memory_order_acquire)); + + lock.unlockShared(); + writer.join(); + + RWSpinlock::ReadGuard guard(lock); + EXPECT_TRUE(writer_entered.load(std::memory_order_acquire)); + EXPECT_EQ(protected_value, 99); +} + //------------------------------------------------------------------------------ // parseFromString //------------------------------------------------------------------------------ @@ -281,4 +347,40 @@ TEST(GetHandshakeMaxLengthTest, ReturnsSameValueOnMultipleCalls) { EXPECT_EQ(first_call, second_call); } +//------------------------------------------------------------------------------ +// bindToSocket +//------------------------------------------------------------------------------ + +// Worker pools bind every thread they spawn, so bindToSocket() races inside +// libnuma's unlocked lazy cache fill and orphans an allocation. The leak is +// what this guards, so it only fails under ASAN/LSAN; the barrier is what makes +// it reliable there. +TEST(BindToSocketTest, ConcurrentCallsDoNotLeakNumaState) { + if (numa_available() < 0) GTEST_SKIP() << "platform does not support NUMA"; + + constexpr int kThreads = 32; + std::atomic ready{0}; + std::atomic go{false}; + std::vector threads; + std::vector results(kThreads, -1); + + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([&, i] { + ready.fetch_add(1, std::memory_order_release); + while (!go.load(std::memory_order_acquire)) + std::this_thread::yield(); + // Every thread races on the same node, maximizing the window. + results[i] = bindToSocket(0); + }); + } + while (ready.load(std::memory_order_acquire) < kThreads) + std::this_thread::yield(); + go.store(true, std::memory_order_release); + for (auto &thread : threads) thread.join(); + + // Serializing the cache fill must not change what callers observe. + for (int i = 0; i < kThreads; ++i) + EXPECT_EQ(results[i], 0) << "thread " << i << " failed to bind"; +} + } // namespace diff --git a/mooncake-transfer-engine/tests/config_test.cpp b/mooncake-transfer-engine/tests/config_test.cpp index 4b58ab7036..e3f75d5c6b 100644 --- a/mooncake-transfer-engine/tests/config_test.cpp +++ b/mooncake-transfer-engine/tests/config_test.cpp @@ -26,6 +26,8 @@ class PkeyIndexEnvTest : public ::testing::Test { void TearDown() override { ::unsetenv("MC_PKEY_INDEX"); ::unsetenv("MC_AUTO_GID_MAX_RETRIES"); + ::unsetenv("MC_IB_SL"); + ::unsetenv("MC_TE_METADATA_REFRESH_INTERVAL_SECONDS"); } }; @@ -105,5 +107,349 @@ TEST_F(PkeyIndexEnvTest, AutoGidRetriesRejectsOutOfRangeOverride) { EXPECT_EQ(config.auto_gid_max_retries, 5); } +TEST_F(PkeyIndexEnvTest, IbSlDefaultsToMinusOneWhenUnset) { + ::unsetenv("MC_IB_SL"); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.ib_service_level, -1); +} + +TEST_F(PkeyIndexEnvTest, IbSlValidOverrideIsApplied) { + ASSERT_EQ(::setenv("MC_IB_SL", "3", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.ib_service_level, 3); +} + +TEST_F(PkeyIndexEnvTest, IbSlMinBoundaryIsApplied) { + ASSERT_EQ(::setenv("MC_IB_SL", "0", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.ib_service_level, 0); +} + +TEST_F(PkeyIndexEnvTest, IbSlMaxBoundaryIsApplied) { + ASSERT_EQ(::setenv("MC_IB_SL", "15", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.ib_service_level, 15); +} + +TEST_F(PkeyIndexEnvTest, IbSlOutOfRangeIsIgnored) { + ASSERT_EQ(::setenv("MC_IB_SL", "16", 1), 0); + GlobalConfig config; + config.ib_service_level = 7; // sentinel preserved when env var is rejected + loadGlobalConfig(config); + EXPECT_EQ(config.ib_service_level, 7); +} + +TEST_F(PkeyIndexEnvTest, IbSlNegativeIsIgnored) { + ASSERT_EQ(::setenv("MC_IB_SL", "-1", 1), 0); + GlobalConfig config; + config.ib_service_level = 5; + loadGlobalConfig(config); + EXPECT_EQ(config.ib_service_level, 5); +} + +TEST_F(PkeyIndexEnvTest, IbSlNonNumericKeepsDefault) { + ASSERT_EQ(::setenv("MC_IB_SL", "abc", 1), 0); + GlobalConfig config; + config.ib_service_level = 9; + loadGlobalConfig(config); + EXPECT_EQ(config.ib_service_level, 9); +} + +TEST_F(PkeyIndexEnvTest, TeMetadataRefreshIntervalDefaultsToZeroWhenUnset) { + ::unsetenv("MC_TE_METADATA_REFRESH_INTERVAL_SECONDS"); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.te_metadata_refresh_interval_seconds, 0); +} + +TEST_F(PkeyIndexEnvTest, TeMetadataRefreshIntervalAcceptsValidOverride) { + ASSERT_EQ(::setenv("MC_TE_METADATA_REFRESH_INTERVAL_SECONDS", "5", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.te_metadata_refresh_interval_seconds, 5); +} + +TEST_F(PkeyIndexEnvTest, TeMetadataRefreshIntervalAcceptsZeroAsDisabled) { + ASSERT_EQ(::setenv("MC_TE_METADATA_REFRESH_INTERVAL_SECONDS", "0", 1), 0); + GlobalConfig config; + config.te_metadata_refresh_interval_seconds = 123; + loadGlobalConfig(config); + EXPECT_EQ(config.te_metadata_refresh_interval_seconds, 0); +} + +TEST_F(PkeyIndexEnvTest, TeMetadataRefreshIntervalRejectsNegativeOverride) { + ASSERT_EQ(::setenv("MC_TE_METADATA_REFRESH_INTERVAL_SECONDS", "-1", 1), 0); + GlobalConfig config; + config.te_metadata_refresh_interval_seconds = 123; + loadGlobalConfig(config); + EXPECT_EQ(config.te_metadata_refresh_interval_seconds, 123); +} + +TEST_F(PkeyIndexEnvTest, TeMetadataRefreshIntervalRejectsNonNumericOverride) { + ASSERT_EQ(::setenv("MC_TE_METADATA_REFRESH_INTERVAL_SECONDS", "abc", 1), 0); + GlobalConfig config; + config.te_metadata_refresh_interval_seconds = 456; + loadGlobalConfig(config); + EXPECT_EQ(config.te_metadata_refresh_interval_seconds, 456); +} + +// MC_CONN_PAUSE_TTL_MS arms the active-connect circuit-breaker: after an +// endpoint to a peer is torn down, active reconnection to that peer's address +// is paused for this many ms so the CQ poller isn't blocked re-handshaking a +// gone peer. 0 disables (and is the default); the range is capped at 600000ms. +// As with the other knobs, a typo / out-of-range value must preserve the +// default rather than silently change behavior. +class ConnPauseTtlEnvTest : public ::testing::Test { + protected: + void TearDown() override { ::unsetenv("MC_CONN_PAUSE_TTL_MS"); } +}; + +TEST_F(ConnPauseTtlEnvTest, DefaultIsZeroWhenUnset) { + ::unsetenv("MC_CONN_PAUSE_TTL_MS"); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.conn_pause_ttl_ms, 0); +} + +TEST_F(ConnPauseTtlEnvTest, ValidOverrideIsApplied) { + ASSERT_EQ(::setenv("MC_CONN_PAUSE_TTL_MS", "5000", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.conn_pause_ttl_ms, 5000); +} + +TEST_F(ConnPauseTtlEnvTest, ZeroIsAcceptedAndDisables) { + ASSERT_EQ(::setenv("MC_CONN_PAUSE_TTL_MS", "0", 1), 0); + GlobalConfig config; + config.conn_pause_ttl_ms = 99; // sentinel must be overwritten by 0 + loadGlobalConfig(config); + EXPECT_EQ(config.conn_pause_ttl_ms, 0); +} + +TEST_F(ConnPauseTtlEnvTest, MaxBoundaryIsApplied) { + ASSERT_EQ(::setenv("MC_CONN_PAUSE_TTL_MS", "600000", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.conn_pause_ttl_ms, 600000); +} + +TEST_F(ConnPauseTtlEnvTest, OutOfRangeIsIgnored) { + ASSERT_EQ(::setenv("MC_CONN_PAUSE_TTL_MS", "600001", 1), 0); + GlobalConfig config; + config.conn_pause_ttl_ms = 7; // sentinel preserved when rejected + loadGlobalConfig(config); + EXPECT_EQ(config.conn_pause_ttl_ms, 7); +} + +TEST_F(ConnPauseTtlEnvTest, NegativeIsIgnored) { + ASSERT_EQ(::setenv("MC_CONN_PAUSE_TTL_MS", "-1", 1), 0); + GlobalConfig config; + config.conn_pause_ttl_ms = 11; + loadGlobalConfig(config); + EXPECT_EQ(config.conn_pause_ttl_ms, 11); +} + +TEST_F(ConnPauseTtlEnvTest, NonNumericKeepsDefault) { + ASSERT_EQ(::setenv("MC_CONN_PAUSE_TTL_MS", "abc", 1), 0); + GlobalConfig config; + config.conn_pause_ttl_ms = 13; // a typo must NOT silently change behavior + loadGlobalConfig(config); + EXPECT_EQ(config.conn_pause_ttl_ms, 13); +} + +TEST_F(ConnPauseTtlEnvTest, NumericSuffixKeepsDefault) { + ASSERT_EQ(::setenv("MC_CONN_PAUSE_TTL_MS", "5000s", 1), 0); + GlobalConfig config; + config.conn_pause_ttl_ms = 15; + loadGlobalConfig(config); + EXPECT_EQ(config.conn_pause_ttl_ms, 15); +} + +TEST_F(ConnPauseTtlEnvTest, EmptyStringKeepsDefault) { + ASSERT_EQ(::setenv("MC_CONN_PAUSE_TTL_MS", "", 1), 0); + GlobalConfig config; + config.conn_pause_ttl_ms = 17; + loadGlobalConfig(config); + EXPECT_EQ(config.conn_pause_ttl_ms, 17); +} + +// MC_MAX_CONCURRENT_REG_MR caps how many buffers registerLocalMemoryBatch() +// registers at once; 0 (the default) means unbounded. 0 is therefore also what +// a silent atol() fallback would produce on a typo, which would read as "the +// knob was honored and asked for no cap" -- the opposite of what the operator +// wanted. So a typo must be rejected loudly and leave the field untouched. +class MaxConcurrentRegMrEnvTest : public ::testing::Test { + protected: + void TearDown() override { ::unsetenv("MC_MAX_CONCURRENT_REG_MR"); } +}; + +TEST_F(MaxConcurrentRegMrEnvTest, UnboundedWhenUnset) { + ::unsetenv("MC_MAX_CONCURRENT_REG_MR"); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.max_concurrent_reg_mr, 0u); +} + +TEST_F(MaxConcurrentRegMrEnvTest, ValidOverrideIsApplied) { + ASSERT_EQ(::setenv("MC_MAX_CONCURRENT_REG_MR", "8", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.max_concurrent_reg_mr, 8u); +} + +TEST_F(MaxConcurrentRegMrEnvTest, ExplicitZeroSelectsUnbounded) { + ASSERT_EQ(::setenv("MC_MAX_CONCURRENT_REG_MR", "0", 1), 0); + GlobalConfig config; + config.max_concurrent_reg_mr = 99; // sentinel must be overwritten by 0 + loadGlobalConfig(config); + EXPECT_EQ(config.max_concurrent_reg_mr, 0u); +} + +TEST_F(MaxConcurrentRegMrEnvTest, OneIsAcceptedAndSerializes) { + ASSERT_EQ(::setenv("MC_MAX_CONCURRENT_REG_MR", "1", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.max_concurrent_reg_mr, 1u); +} + +TEST_F(MaxConcurrentRegMrEnvTest, NegativeIsIgnored) { + ASSERT_EQ(::setenv("MC_MAX_CONCURRENT_REG_MR", "-1", 1), 0); + GlobalConfig config; + config.max_concurrent_reg_mr = 11; + loadGlobalConfig(config); + EXPECT_EQ(config.max_concurrent_reg_mr, 11u); +} + +TEST_F(MaxConcurrentRegMrEnvTest, NonNumericKeepsDefault) { + ASSERT_EQ(::setenv("MC_MAX_CONCURRENT_REG_MR", "abc", 1), 0); + GlobalConfig config; + config.max_concurrent_reg_mr = 13; + loadGlobalConfig(config); + EXPECT_EQ(config.max_concurrent_reg_mr, 13u); +} + +TEST_F(MaxConcurrentRegMrEnvTest, NumericSuffixKeepsDefault) { + ASSERT_EQ(::setenv("MC_MAX_CONCURRENT_REG_MR", "8x", 1), 0); + GlobalConfig config; + config.max_concurrent_reg_mr = 15; + loadGlobalConfig(config); + EXPECT_EQ(config.max_concurrent_reg_mr, 15u); +} + +TEST_F(MaxConcurrentRegMrEnvTest, EmptyStringKeepsDefault) { + ASSERT_EQ(::setenv("MC_MAX_CONCURRENT_REG_MR", "", 1), 0); + GlobalConfig config; + config.max_concurrent_reg_mr = 17; + loadGlobalConfig(config); + EXPECT_EQ(config.max_concurrent_reg_mr, 17u); +} + +class EfaNicSelectionEnvTest : public ::testing::Test { + protected: + void TearDown() override { ::unsetenv("MC_EFA_NIC_SELECTION"); } +}; + +TEST_F(EfaNicSelectionEnvTest, DefaultsToAll) { + ::unsetenv("MC_EFA_NIC_SELECTION"); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.efa_nic_selection, EfaNicSelection::ALL); +} + +TEST_F(EfaNicSelectionEnvTest, LocalIsApplied) { + ASSERT_EQ(::setenv("MC_EFA_NIC_SELECTION", "local", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.efa_nic_selection, EfaNicSelection::LOCAL); +} + +TEST_F(EfaNicSelectionEnvTest, AllIsAcceptedExplicitly) { + ASSERT_EQ(::setenv("MC_EFA_NIC_SELECTION", "all", 1), 0); + GlobalConfig config; + config.efa_nic_selection = EfaNicSelection::LOCAL; // must be overwritten + loadGlobalConfig(config); + EXPECT_EQ(config.efa_nic_selection, EfaNicSelection::ALL); +} + +TEST_F(EfaNicSelectionEnvTest, CaseIsIgnored) { + ASSERT_EQ(::setenv("MC_EFA_NIC_SELECTION", "LOCAL", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.efa_nic_selection, EfaNicSelection::LOCAL); +} + +TEST_F(EfaNicSelectionEnvTest, UnknownValueKeepsDefault) { + // A typo must not silently pick a policy: registering buffers on the wrong + // NIC set is a correctness-adjacent surprise, not a perf knob. + ASSERT_EQ(::setenv("MC_EFA_NIC_SELECTION", "topology", 1), 0); + GlobalConfig config; + config.efa_nic_selection = EfaNicSelection::LOCAL; + loadGlobalConfig(config); + EXPECT_EQ(config.efa_nic_selection, EfaNicSelection::LOCAL); +} + +TEST_F(EfaNicSelectionEnvTest, EmptyStringKeepsDefault) { + ASSERT_EQ(::setenv("MC_EFA_NIC_SELECTION", "", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_EQ(config.efa_nic_selection, EfaNicSelection::ALL); +} + +// max_wr_from_env distinguishes "the operator asked for this depth" from "this +// is the compiled-in default". The EFA transport needs that distinction: with +// no override it adopts the provider's per-device transmit queue depth, which +// no fixed default can match (2048 on p6-b300, 4096 on p5). A rejected value +// must NOT set the flag, or a typo would be treated as a deliberate override. +class MaxWrEnvTest : public ::testing::Test { + protected: + void TearDown() override { ::unsetenv("MC_MAX_WR"); } +}; + +TEST_F(MaxWrEnvTest, NotFromEnvWhenUnset) { + ::unsetenv("MC_MAX_WR"); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_FALSE(config.max_wr_from_env); + EXPECT_EQ(config.max_wr, 256u); // default preserved for RDMA +} + +TEST_F(MaxWrEnvTest, ValidOverrideSetsFlag) { + ASSERT_EQ(::setenv("MC_MAX_WR", "2048", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_TRUE(config.max_wr_from_env); + EXPECT_EQ(config.max_wr, 2048u); +} + +TEST_F(MaxWrEnvTest, RejectedValueDoesNotSetFlag) { + // 0 is out of range. The value is ignored, so the EFA transport must + // still treat this as "no override" and track the provider's depth. + ASSERT_EQ(::setenv("MC_MAX_WR", "0", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_FALSE(config.max_wr_from_env); + EXPECT_EQ(config.max_wr, 256u); +} + +TEST_F(MaxWrEnvTest, NonNumericDoesNotSetFlag) { + ASSERT_EQ(::setenv("MC_MAX_WR", "abc", 1), 0); + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_FALSE(config.max_wr_from_env); + EXPECT_EQ(config.max_wr, 256u); +} + +TEST_F(MaxWrEnvTest, OutOfRangeDoesNotSetFlag) { + ASSERT_EQ(::setenv("MC_MAX_WR", "70000", 1), 0); // > UINT16_MAX + GlobalConfig config; + loadGlobalConfig(config); + EXPECT_FALSE(config.max_wr_from_env); + EXPECT_EQ(config.max_wr, 256u); +} + } // namespace } // namespace mooncake diff --git a/mooncake-transfer-engine/tests/connect_pause_tracker_test.cpp b/mooncake-transfer-engine/tests/connect_pause_tracker_test.cpp new file mode 100644 index 0000000000..55ca54b7ba --- /dev/null +++ b/mooncake-transfer-engine/tests/connect_pause_tracker_test.cpp @@ -0,0 +1,143 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Hardware-free unit tests for ConnectPauseTracker (the active-connect +// circuit-breaker state). The tracker takes an injected clock, so the +// TTL/expiry/prune logic is exercised deterministically with no RDMA device +// and no real sleeps. The header is inline-only, so the test links just gtest. + +#include "transport/rdma_transport/connect_pause_tracker.h" + +#include + +#include +#include +#include +#include +#include + +using mooncake::ConnectPauseTracker; + +namespace { + +// Manually-advanced clock so expiry transitions are deterministic. +struct FakeClock { + std::atomic now{0}; + uint64_t operator()() const { return now.load(std::memory_order_relaxed); } +}; + +ConnectPauseTracker makeTracker(std::shared_ptr clk) { + return ConnectPauseTracker([clk] { return (*clk)(); }); +} + +TEST(ConnectPauseTracker, UnknownPeerNotPaused) { + auto clk = std::make_shared(); + auto t = makeTracker(clk); + EXPECT_FALSE(t.isPaused("10.0.0.1:1234")); + EXPECT_EQ(t.size(), 0u); +} + +TEST(ConnectPauseTracker, PausedUntilExpiry) { + auto clk = std::make_shared(); + auto t = makeTracker(clk); + const std::string peer = "10.0.0.1:1234"; + clk->now = 500; + t.pauseFor(peer, 1000); // paused until ts == 1500 + EXPECT_TRUE(t.isPaused(peer)); // now == 500 + clk->now = 1499; + EXPECT_TRUE(t.isPaused(peer)); // just before expiry + clk->now = 1500; // at expiry (>= until) + EXPECT_FALSE(t.isPaused(peer)); // expired + EXPECT_EQ(t.size(), 0u); // and lazily deleted on the failing check +} + +TEST(ConnectPauseTracker, RepeatedChecksDoNotExtendHardDeadline) { + auto clk = std::make_shared(); + auto t = makeTracker(clk); + const std::string peer = "10.0.0.1:1234"; + t.pauseFor(peer, 1000); + + // Simulate continuous traffic checking the pause. Lookups are not actual + // connection failures, so they must not refresh the deadline. + for (uint64_t now = 100; now < 1000; now += 100) { + clk->now = now; + EXPECT_TRUE(t.isPaused(peer)); + } + + clk->now = 1000; + EXPECT_FALSE(t.isPaused(peer)); +} + +TEST(ConnectPauseTracker, RefreshExtendsWindow) { + auto clk = std::make_shared(); + auto t = makeTracker(clk); + const std::string peer = "p"; + t.pauseFor(peer, 100); + clk->now = 50; + t.pauseFor(peer, 150); // re-arm while still paused -> extend to ts == 200 + clk->now = 150; + EXPECT_TRUE(t.isPaused(peer)); // within the extended window + clk->now = 200; + EXPECT_FALSE(t.isPaused(peer)); +} + +TEST(ConnectPauseTracker, PruneDropsOnlyExpired) { + auto clk = std::make_shared(); + auto t = makeTracker(clk); + t.pauseFor("a", 100); + t.pauseFor("b", 300); + EXPECT_EQ(t.size(), 2u); + clk->now = 200; // "a" expired, "b" still paused + t.prune(); + EXPECT_EQ(t.size(), 1u); + EXPECT_FALSE(t.isPaused("a")); + EXPECT_TRUE(t.isPaused("b")); +} + +TEST(ConnectPauseTracker, PerPeerIndependent) { + auto clk = std::make_shared(); + auto t = makeTracker(clk); + t.pauseFor("a", 100); + t.pauseFor("b", 1000); + clk->now = 150; + EXPECT_FALSE(t.isPaused("a")); // a's window lapsed + EXPECT_TRUE(t.isPaused("b")); // b's has not +} + +// Hammer all entry points concurrently; primarily a ThreadSanitizer target +// (run under -fsanitize=thread). The clock is pinned far in the future so +// entries don't expire mid-run, isolating the data-race check. +TEST(ConnectPauseTracker, ConcurrentAccessIsRaceFree) { + auto clk = std::make_shared(); + auto t = makeTracker(clk); + constexpr int kIters = 5000; + constexpr uint64_t kFarFuture = 1ull << 40; + std::vector threads; + for (int i = 0; i < 4; ++i) + threads.emplace_back([&t, i] { + std::string s = "peer" + std::to_string(i % 3); + for (int k = 0; k < kIters; ++k) t.pauseFor(s, kFarFuture); + }); + for (int i = 0; i < 4; ++i) + threads.emplace_back([&t] { + for (int k = 0; k < kIters; ++k) (void)t.isPaused("peer0"); + }); + threads.emplace_back([&t] { + for (int k = 0; k < kIters; ++k) t.prune(); + }); + for (auto& th : threads) th.join(); + SUCCEED(); // TSan asserts the absence of data races +} + +} // namespace diff --git a/mooncake-transfer-engine/tests/cxi_transfer_test.cpp b/mooncake-transfer-engine/tests/cxi_transfer_test.cpp new file mode 100644 index 0000000000..763058f462 --- /dev/null +++ b/mooncake-transfer-engine/tests/cxi_transfer_test.cpp @@ -0,0 +1,605 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// CXI multi-NIC transfer test +// +// Target node: allocate buffers, register on all NICs, wait. +// Initiator node: allocate receive buffer, register, pull from target. +// +// Usage: +// # Target (holds KV cache): +// ./cxi_transfer_test --mode target --server :12345 \ +// --num_bufs 4 --buf_size_gb 1 +// +// # Initiator (pulls data): +// ./cxi_transfer_test --mode initiator --server :12346 \ +// --target :12345 --num_bufs 4 --buf_size_gb 1 \ +// --transfer_mb 368 + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "transfer_engine.h" +#include "transport/cxi_transport/cxi_transport.h" + +#if defined(USE_CUDA) || defined(USE_HIP) +#include +#endif + +using namespace mooncake; + +DEFINE_string(mode, "target", "Running mode: target or initiator"); +DEFINE_string(server, "", "Local server name, e.g. 172.31.6.162:12345"); +DEFINE_string(target, "", "Target server name (initiator mode only)"); +DEFINE_string(metadata, "P2PHANDSHAKE", "Metadata server"); +DEFINE_int32(num_bufs, 200, "Number of buffers to allocate and register"); +DEFINE_double(buf_size_gb, 2.0, "Size of each buffer in GB"); +DEFINE_double(transfer_mb, 368.0, + "Transfer size per iteration in MB (initiator)"); +DEFINE_int32(iterations, 50, "Number of benchmark iterations"); +DEFINE_int32(warmup, 5, "Number of warmup iterations"); +DEFINE_int32(batch_size, 1, "Batch size for each transfer submission"); +DEFINE_int32(threads, 1, "Number of initiator worker threads"); +DEFINE_uint64(block_size, 65536, "Block size for transfer requests (64KB)"); +DEFINE_bool(use_device, false, "Use CUDA device memory for transfer"); + +static std::atomic g_running(true); + +static void signalHandler(int) { + g_running.store(false, std::memory_order_relaxed); +} + +static void setupSignalHandler() { + struct sigaction sa; + sa.sa_handler = signalHandler; + sigemptyset(&sa.sa_mask); + sa.sa_flags = 0; + sigaction(SIGINT, &sa, nullptr); + sigaction(SIGTERM, &sa, nullptr); +} + +static void* allocateHugepage(size_t size, int node_id) { + void* buf = mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0); + if (buf == MAP_FAILED) { + LOG(WARNING) << "Hugepage mmap failed (" << strerror(errno) + << "), falling back to regular pages"; + buf = mmap(nullptr, size, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (buf == MAP_FAILED) return nullptr; + } + unsigned long nodemask = 1UL << node_id; + + mbind(buf, size, MPOL_BIND, &nodemask, + /* maxnode */ sizeof(nodemask) * 8, 0); + return buf; +} + +#if defined(USE_CUDA) || defined(USE_HIP) +static constexpr bool can_use_device = true; + +static void* allocateDevice(size_t size, int device) { + cudaSetDevice(device); + void* buf = nullptr; + cudaError_t err = cudaMalloc(&buf, size); + if (err != cudaSuccess) { + return nullptr; + } + return buf; +} + +static int setDeviceMem(void* buf, size_t size, int device, int val) { + cudaSetDevice(device); + cudaError_t err = cudaMemset(buf, val, size); + if (err != cudaSuccess) { + return -1; + } + return 0; +} + +static bool checkDeviceMem(void* buf, int device, uint8_t expected) { + cudaSetDevice(device); + uint8_t actual; + size_t offset = 42; + cudaError_t err = + cudaMemcpy(&actual, buf + offset, 1, cudaMemcpyDeviceToHost); + if (err != cudaSuccess) return false; + return expected == actual; +} + +static int freeDeviceMem(void* buf, int device) { + cudaSetDevice(device); + cudaError_t err = cudaFree(buf); + if (err != cudaSuccess) return -1; + return 0; +} + +#else +static constexpr bool can_use_device = false; +#endif + +static int runTarget(TransferEngine* engine) { + size_t buf_bytes = + static_cast(FLAGS_buf_size_gb * 1024 * 1024 * 1024); + int num_bufs = FLAGS_num_bufs; + + LOG(INFO) << "=== Target Node ==="; + LOG(INFO) << "Registering " << num_bufs << " x " << FLAGS_buf_size_gb + << " GB = " << num_bufs * FLAGS_buf_size_gb << " GB"; + + // Allocate buffers + std::vector bufs; + bufs.reserve(num_bufs); + LOG(INFO) << "Allocating " << num_bufs << " buffers..."; + for (int i = 0; i < num_bufs; ++i) { + void* buf; + if (!FLAGS_use_device) { + buf = allocateHugepage(buf_bytes, 0); + if (!buf) { + LOG(ERROR) << "Allocation failed at buffer " << i; + for (auto* p : bufs) munmap(p, buf_bytes); + return 1; + } + buf = memset(buf, (uint8_t)42, buf_bytes); + if (!buf) { + LOG(ERROR) << "Memset failed at buffer " << i; + for (auto* p : bufs) munmap(p, buf_bytes); + return 1; + } + } else { +#if defined(USE_CUDA) || defined(USE_HIP) + int buf_device = i % 4; + buf = allocateDevice(buf_bytes, buf_device); + if (!buf) { + LOG(ERROR) << "allocation failed at buffer " << i; + for (int j = 0; j < i; j++) { + freeDeviceMem(bufs[j], j % 4); + } + return 1; + } + if (setDeviceMem(buf, buf_bytes, buf_device, 42) != 0) { + LOG(ERROR) << "failed to set memory for buffer " << i; + for (int j = 0; j < i; j++) { + freeDeviceMem(bufs[j], j % 4); + } + return 1; + } +#endif + } + + bufs.push_back(buf); + if ((i + 1) % 50 == 0 || i == num_bufs - 1) + LOG(INFO) << " allocated " << (i + 1) << "/" << num_bufs; + } + + // Register all buffers + LOG(INFO) << "Registering " << num_bufs << " buffers on all NICs..."; + auto t0 = std::chrono::steady_clock::now(); + for (int i = 0; i < num_bufs; ++i) { + int ret = engine->registerLocalMemory(bufs[i], buf_bytes, "*", true); + if (ret != 0) { + LOG(ERROR) << "registerLocalMemory failed at buffer " << i + << ": ret=" << ret; + for (int j = 0; j < i; ++j) engine->unregisterLocalMemory(bufs[j]); + for (auto* p : bufs) munmap(p, buf_bytes); + return 1; + } + if ((i + 1) % 50 == 0 || i == num_bufs - 1) { + auto now = std::chrono::steady_clock::now(); + double elapsed = std::chrono::duration(now - t0).count(); + LOG(INFO) << " registered " << (i + 1) << "/" << num_bufs << " (" + << elapsed << "s)"; + } + } + auto t1 = std::chrono::steady_clock::now(); + LOG(INFO) << "Registration complete: " + << std::chrono::duration(t1 - t0).count() << "s"; + + LOG(INFO) << "Target ready. First buffer at " << bufs[0] + << ". Waiting for initiator (Ctrl+C to stop)..."; + + if (FLAGS_use_device) cudaSetDevice(0); + while (g_running) sleep(1); + + LOG(INFO) << "Shutting down target..."; + for (auto* p : bufs) { + engine->unregisterLocalMemory(p); + if (!FLAGS_use_device) { + munmap(p, buf_bytes); + } else { +#if defined(USE_CUDA) || defined(USE_HIP) + freeDeviceMem(p, buf_bytes); +#endif + } + } + return 0; +} + +struct LatencyStats { + double avg_ms; + double p50_ms; + double p99_ms; + double throughput_gbs; +}; + +static LatencyStats computeStats(std::vector& latencies_ms, + size_t transfer_bytes) { + std::sort(latencies_ms.begin(), latencies_ms.end()); + double sum = std::accumulate(latencies_ms.begin(), latencies_ms.end(), 0.0); + size_t n = latencies_ms.size(); + LatencyStats stats; + stats.avg_ms = sum / n; + stats.p50_ms = latencies_ms[n / 2]; + stats.p99_ms = latencies_ms[std::min(n - 1, (size_t)(n * 0.99))]; + stats.throughput_gbs = (transfer_bytes / 1e9) / (stats.p50_ms / 1000.0); + return stats; +} + +static int runInitiator(TransferEngine* engine) { + size_t transfer_bytes = + static_cast(FLAGS_transfer_mb * 1024 * 1024); + + LOG(INFO) << "=== Initiator Node ==="; + LOG(INFO) << "Target: " << FLAGS_target; + LOG(INFO) << "Transfer size: " << FLAGS_transfer_mb << " MB"; + LOG(INFO) << "Threads: " << FLAGS_threads; + + if (FLAGS_target.empty()) { + LOG(ERROR) << "--target required in initiator mode"; + return 1; + } + + // Allocate local receive buffer (one per thread) + size_t recv_bytes = transfer_bytes; + std::vector recv_bufs(FLAGS_threads); + for (int t = 0; t < FLAGS_threads; ++t) { + if (!FLAGS_use_device) { + recv_bufs[t] = allocateHugepage(recv_bytes, 0); + if (!recv_bufs[t]) { + LOG(ERROR) << "Failed to allocate receive buffer for thread " + << t; + return 1; + } + recv_bufs[t] = memset(recv_bufs[t], 66, recv_bytes); + if (!recv_bufs[t]) { + LOG(ERROR) << "Failed to memset receive buffer for thread " + << t; + return 1; + } + } else { +#if defined(USE_CUDA) || defined(USE_HIP) + int tid_device = t % 4; + recv_bufs[t] = allocateDevice(recv_bytes, tid_device); + if (!recv_bufs[t]) { + LOG(ERROR) << "Failed to allocate receive buffer for thread " + << t; + return 1; + } + if (setDeviceMem(recv_bufs[t], recv_bytes, tid_device, 66) != 0) { + LOG(ERROR) << "Failed to set data on receive buffer for thread " + << t; + } +#endif + } + + int ret = + engine->registerLocalMemory(recv_bufs[t], recv_bytes, "*", true); + if (ret != 0) { + LOG(ERROR) << "Failed to register receive buffer: " << ret; + return 1; + } + } + LOG(INFO) << "Allocated and registered " << FLAGS_threads + << " receive buffers of " << recv_bytes / 1e6 << " MB each"; + + // Open remote segment + auto segment_id = engine->openSegment(FLAGS_target); + if (segment_id < 0) { + LOG(ERROR) << "openSegment failed for " << FLAGS_target; + return 1; + } + + // Get remote buffer info + auto segment_desc = engine->getMetadata()->getSegmentDescByID(segment_id); + if (!segment_desc || segment_desc->buffers.empty()) { + LOG(ERROR) << "No remote buffers found"; + return 1; + } + size_t num_remote_bufs = segment_desc->buffers.size(); + LOG(INFO) << "Remote has " << num_remote_bufs << " buffers"; + LOG(INFO) << "First buffer: addr=0x" << std::hex + << segment_desc->buffers[0].addr << std::dec + << " size=" << segment_desc->buffers[0].length; + + for (auto buffer : segment_desc->buffers) { + for (int i = 0; i <= 4; i++) { + LOG(INFO) << "retry " << i << " selected device: " + << segment_desc->topology.selectDevice(buffer.name, i); + } + for (int i = 0; i < buffer.rkey.size(); i++) { + LOG(INFO) << "remote device " << i << " has rkey " << buffer.rkey[i] + << "\n"; + } + } + + // Connection warmup: small transfer to establish endpoints + LOG(INFO) << "Warming up connection..."; + { + size_t warmup_size = std::min(transfer_bytes, (size_t)(64 * 1024)); + auto batch_id = engine->allocateBatchID(1); + TransferRequest req; + req.opcode = TransferRequest::READ; + req.source = (uint8_t*)recv_bufs[0]; + req.target_id = segment_id; + req.target_offset = segment_desc->buffers[0].addr; + req.length = warmup_size; + auto s = engine->submitTransfer(batch_id, {req}); + if (!s.ok()) { + LOG(ERROR) << "Warmup transfer failed: " << s.ToString(); + return 1; + } + while (true) { + TransferStatus status; + engine->getTransferStatus(batch_id, 0, status); + if (status.s == TransferStatusEnum::COMPLETED) break; + if (status.s == TransferStatusEnum::FAILED) { + LOG(ERROR) << "Warmup transfer FAILED"; + return 1; + } + } + if (recv_bufs[0]) engine->freeBatchID(batch_id); + } + LOG(INFO) << "Connection ready."; + + // Refresh segment desc after connection warmup (endpoints are now up) + engine->syncSegmentCache(FLAGS_target); + segment_desc = engine->getMetadata()->getSegmentDescByID(segment_id); + + // Worker function: each thread runs its own transfer loop + struct ThreadResult { + std::vector latencies; + int errors = 0; + }; + + auto workerFn = [&](int tid, int warmup_iters, int bench_iters, + ThreadResult* result) { + void* my_recv = recv_bufs[tid]; + size_t buf_idx = + (tid + 1) % num_remote_bufs; // one thread <-> one buffer + for (int w = 0; w < warmup_iters; ++w) { + uint64_t raddr = segment_desc->buffers[buf_idx].addr; + size_t rlen = segment_desc->buffers[buf_idx].length; + size_t xfer = std::min(transfer_bytes, rlen); + + auto bid = engine->allocateBatchID(1); + TransferRequest req; + req.opcode = TransferRequest::READ; + req.source = (uint8_t*)my_recv; + req.target_id = segment_id; + req.target_offset = raddr; + req.length = xfer; + + engine->submitTransfer(bid, {req}); + while (true) { + TransferStatus st; + engine->getTransferStatus(bid, 0, st); + if (st.s == TransferStatusEnum::COMPLETED || + st.s == TransferStatusEnum::FAILED) + break; + } + engine->freeBatchID(bid); + } + + for (int i = 0; i < bench_iters; ++i) { + uint64_t raddr = segment_desc->buffers[buf_idx].addr; + size_t rlen = segment_desc->buffers[buf_idx].length; + size_t xfer = std::min(transfer_bytes, rlen); + + auto t0 = std::chrono::steady_clock::now(); + auto bid = engine->allocateBatchID(1); + TransferRequest req; + req.opcode = TransferRequest::READ; + req.source = (uint8_t*)my_recv; + req.target_id = segment_id; + req.target_offset = raddr; + req.length = xfer; + auto s = engine->submitTransfer(bid, {req}); + if (!s.ok()) { + result->errors++; + engine->freeBatchID(bid); + continue; + } + bool ok = false; + while (true) { + TransferStatus st; + engine->getTransferStatus(bid, 0, st); + if (st.s == TransferStatusEnum::COMPLETED) { + ok = true; + break; + } + if (st.s == TransferStatusEnum::FAILED) { + LOG(ERROR) + << "thread " << tid << " has failed the transfer\n"; + result->errors++; + break; + } + } + engine->freeBatchID(bid); + if (ok) { + auto t1 = std::chrono::steady_clock::now(); + result->latencies.push_back( + std::chrono::duration(t1 - t0).count()); + } + } + }; + + // Run warmup + benchmark with threads + int num_threads = FLAGS_threads; + LOG(INFO) << "Running with " << num_threads << " threads, " << FLAGS_warmup + << " warmup + " << FLAGS_iterations + << " bench iterations per thread..."; + + std::vector results(num_threads); + std::vector threads; + + auto wall_t0 = std::chrono::steady_clock::now(); + for (int t = 0; t < num_threads; ++t) { + threads.emplace_back(workerFn, t, FLAGS_warmup, FLAGS_iterations, + &results[t]); + } + for (auto& th : threads) th.join(); + auto wall_t1 = std::chrono::steady_clock::now(); + double wall_ms = + std::chrono::duration(wall_t1 - wall_t0).count(); + + // Aggregate results + std::vector all_latencies; + int total_errors = 0; + for (auto& r : results) { + all_latencies.insert(all_latencies.end(), r.latencies.begin(), + r.latencies.end()); + total_errors += r.errors; + } + + if (all_latencies.empty()) { + LOG(ERROR) << "All transfers failed"; + return 1; + } + + auto stats = computeStats(all_latencies, transfer_bytes); + size_t total_xfers = all_latencies.size(); + double total_bytes = (double)total_xfers * transfer_bytes; + double agg_throughput = total_bytes / 1e9 / (wall_ms / 1000.0); + + LOG(INFO) << "=== Results (" << num_threads << " threads) ==="; + LOG(INFO) << "Transfer: " << FLAGS_transfer_mb << " MB x " << total_xfers + << " transfers"; + LOG(INFO) << "Wall time: " << wall_ms << " ms"; + LOG(INFO) << "Per-transfer p50: " << stats.p50_ms << " ms" + << " p99: " << stats.p99_ms << " ms"; + LOG(INFO) << "Per-transfer throughput: " << stats.throughput_gbs << " GB/s"; + LOG(INFO) << "Aggregate throughput: " << agg_throughput << " GB/s"; + LOG(INFO) << "Errors: " << total_errors; + + // Per-thread stats + for (int t = 0; t < num_threads; ++t) { + if (results[t].latencies.empty()) continue; + auto ts = computeStats(results[t].latencies, transfer_bytes); + LOG(INFO) << " Thread " << t << ": p50=" << ts.p50_ms + << "ms tput=" << ts.throughput_gbs << " GB/s" + << " iters=" << results[t].latencies.size() + << " errors=" << results[t].errors; + } + + // Cleanup + for (int t = 0; t < FLAGS_threads; ++t) { + if (!FLAGS_use_device) { + if (((uint8_t*)recv_bufs[t])[66] != 42) { + LOG(ERROR) << "transfer has corrupted data, expected to find " + "42 but found " + << ((int)((uint8_t*)recv_bufs[t])[66]); + } else { + LOG(INFO) << "transfer was successful"; + } + engine->unregisterLocalMemory(recv_bufs[t]); + munmap(recv_bufs[t], recv_bytes); + } else { +#if defined(USE_CUDA) || defined(USE_HIP) + bool result = checkDeviceMem(recv_bufs[t], t % 4, 42); + if (!result) { + LOG(ERROR) + << "transfer has corrupted data, expected to find 42"; + } else { + LOG(INFO) << "transfer was successful\n"; + } + engine->unregisterLocalMemory(recv_bufs[t]); + freeDeviceMem(recv_bufs[t], t % 4); +#endif + } + } + return total_errors > 0 ? 1 : 0; +} + +int main(int argc, char** argv) { + google::InitGoogleLogging(argv[0]); + gflags::ParseCommandLineFlags(&argc, &argv, true); + FLAGS_logtostderr = 1; + + setupSignalHandler(); + + if (FLAGS_use_device && !can_use_device) { + LOG(ERROR) + << "Cannot use device memory, recompile with cmake -DUSE_CUDA!"; + return 1; + } + + if (FLAGS_use_device) { + LOG(INFO) << "Using device memory"; + } + + if (FLAGS_server.empty()) { + LOG(ERROR) << "--server is required"; + return 1; + } + + // Parse host:port + auto colon = FLAGS_server.rfind(':'); + std::string host = FLAGS_server.substr(0, colon); + uint64_t port = 12345; + if (colon != std::string::npos) + port = std::stoull(FLAGS_server.substr(colon + 1)); + + auto engine = std::make_unique(false); + int ret = engine->init(FLAGS_metadata, FLAGS_server, host, port); + if (ret != 0) { + LOG(ERROR) << "Engine init failed: " << ret; + return 1; + } + + // Discover topology and install CXI transport (all NICs) + engine->getLocalTopology()->discover({}); + auto* xport = engine->installTransport("cxi", nullptr); + if (!xport) { + LOG(ERROR) << "installTransport(cxi) failed"; + return 1; + } + + std::string actual_server = engine->getLocalIpAndPort(); + LOG(INFO) << "Actual server name (use this for --target): " + << actual_server; + + if (FLAGS_mode == "target") { + ret = runTarget(engine.get()); + } else if (FLAGS_mode == "initiator") { + ret = runInitiator(engine.get()); + } else { + LOG(ERROR) << "Unknown mode: " << FLAGS_mode; + ret = 1; + } + + return ret; +} diff --git a/mooncake-transfer-engine/tests/cxi_transport_test.cpp b/mooncake-transfer-engine/tests/cxi_transport_test.cpp new file mode 100644 index 0000000000..a8cb4cc0d4 --- /dev/null +++ b/mooncake-transfer-engine/tests/cxi_transport_test.cpp @@ -0,0 +1,640 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +#include +#include + +#include "transfer_engine.h" +#include "transport/cxi_transport/cxi_transport.h" +#include "transport/transport.h" +#include "cuda_alike.h" + +#if defined(USE_CUDA) || defined(USE_HIP) +#define USE_GPU +#endif + +using namespace mooncake; + +namespace mooncake { + +static void *allocateMemoryPool(size_t size, int socket_id) { + return numa_alloc_onnode(size, socket_id); +} + +static void freeMemoryPool(void *addr, size_t size) { numa_free(addr, size); } + +#ifdef USE_GPU +static void *allocateMemoryPoolDevice(size_t size, int device) { + void *devPtr = nullptr; + cudaError_t err = cudaMalloc(&devPtr, size); + cudaDeviceSynchronize(); + if (err != cudaSuccess) { + LOG(ERROR) << "failed to alloc " << size << " bytes on cuda device " + << device; + return nullptr; + } + return devPtr; +} + +static void freeMemoryPoolDevice(void *addr) { + cudaError_t err = cudaFree(addr); + if (err != cudaSuccess) { + LOG(ERROR) << "failed to free cuda memory @ " << addr; + } +} +#endif + +// --------------------------------------------------------------------------- +// CXI Transport Test Fixture +// +// This test uses the P2PHANDSHAKE metadata backend and performs loopback +// transfers (local_server_name == segment_id), similar to the TCP transport +// tests. It requires CXI hardware to be present (fi_info -p cxi must succeed). +// +// Environment variables: +// MC_METADATA_SERVER - metadata backend (default: P2PHANDSHAKE) +// MC_LOCAL_SERVER_NAME - local server name (default: 127.0.0.1:12345) +// --------------------------------------------------------------------------- +class CXITransportTest : public ::testing::Test { + protected: + void SetUp() override { + google::InitGoogleLogging("CXITransportTest"); + FLAGS_logtostderr = 1; + + const char *env = std::getenv("MC_METADATA_SERVER"); + metadata_server_ = env ? env : "P2PHANDSHAKE"; + LOG(INFO) << "metadata_server: " << metadata_server_; + + env = std::getenv("MC_LOCAL_SERVER_NAME"); + local_server_name_ = env ? env : "127.0.0.1:12345"; + LOG(INFO) << "local_server_name: " << local_server_name_; + } + + void TearDown() override { google::ShutdownGoogleLogging(); } + + // Helper: create engine, install CXI transport, register memory + struct EngineSetup { + std::unique_ptr engine; + Transport *xport; + void *addr; + size_t buffer_size; + SegmentID segment_id; + }; + +#ifdef USE_GPU + + EngineSetup createEngineDevice(size_t buffer_size = 1ull << 30) { + EngineSetup s; + s.buffer_size = buffer_size; + + s.engine = std::make_unique(false); + // Manually discover topology to populate CXI device list + // (same pattern as the Python binding in transfer_engine_py.cpp) + s.engine->getLocalTopology()->discover({}); + auto hp = parseHostNameWithPort(local_server_name_); + int rc = s.engine->init(metadata_server_, local_server_name_, + hp.first.c_str(), hp.second); + EXPECT_EQ(rc, 0) << "engine->init failed"; + + s.xport = s.engine->installTransport("cxi", nullptr); + EXPECT_NE(s.xport, nullptr) << "installTransport(\"cxi\") failed"; + + s.addr = allocateMemoryPoolDevice(buffer_size, + 0); // allocate mempool on cuda:0 + EXPECT_NE(s.addr, nullptr) << "allocateMemoryPool failed"; + + rc = s.engine->registerLocalMemory(s.addr, buffer_size, + GPU_PREFIX + "0"); + EXPECT_EQ(rc, 0) << "registerLocalMemory failed"; + + // Use actual RPC address (P2PHANDSHAKE picks a random port) + auto actual_addr = s.engine->getLocalIpAndPort(); + s.segment_id = s.engine->openSegment(actual_addr); + return s; + } + + void destroyEngineDevice(EngineSetup &s) { + if (s.engine && s.addr) { + s.engine->unregisterLocalMemory(s.addr); + } + if (s.addr) { + freeMemoryPoolDevice(s.addr); + s.addr = nullptr; + } + } + +#endif + + EngineSetup createEngine(size_t buffer_size = 1ull << 30) { + EngineSetup s; + s.buffer_size = buffer_size; + + s.engine = std::make_unique(false); + // Manually discover topology to populate CXI device list + // (same pattern as the Python binding in transfer_engine_py.cpp) + s.engine->getLocalTopology()->discover({}); + auto hp = parseHostNameWithPort(local_server_name_); + int rc = s.engine->init(metadata_server_, local_server_name_, + hp.first.c_str(), hp.second); + EXPECT_EQ(rc, 0) << "engine->init failed"; + + s.xport = s.engine->installTransport("cxi", nullptr); + EXPECT_NE(s.xport, nullptr) << "installTransport(\"cxi\") failed"; + + s.addr = allocateMemoryPool(buffer_size, 0); + EXPECT_NE(s.addr, nullptr) << "allocateMemoryPool failed"; + + rc = s.engine->registerLocalMemory(s.addr, buffer_size, "cpu:0"); + EXPECT_EQ(rc, 0) << "registerLocalMemory failed"; + + // Use actual RPC address (P2PHANDSHAKE picks a random port) + auto actual_addr = s.engine->getLocalIpAndPort(); + s.segment_id = s.engine->openSegment(actual_addr); + return s; + } + + void destroyEngine(EngineSetup &s) { + if (s.engine && s.addr) { + s.engine->unregisterLocalMemory(s.addr); + } + if (s.addr) { + freeMemoryPool(s.addr, s.buffer_size); + s.addr = nullptr; + } + } + + // Helper: submit a single transfer and poll until completion + bool submitAndWait(TransferEngine *engine, SegmentID segment_id, + void *source, uint64_t target_offset, size_t length, + TransferRequest::OpCode opcode) { + auto batch_id = engine->allocateBatchID(1); + + TransferRequest entry; + entry.opcode = opcode; + entry.length = length; + entry.source = (uint8_t *)source; + entry.target_id = segment_id; + entry.target_offset = target_offset; + + Status s = engine->submitTransfer(batch_id, {entry}); + if (!s.ok()) { + LOG(ERROR) << "submitTransfer failed: " << s.ToString(); + engine->freeBatchID(batch_id); + return false; + } + + // Poll for completion with timeout + const int kMaxPollIterations = 1000000; + TransferStatus status; + for (int i = 0; i < kMaxPollIterations; ++i) { + s = engine->getTransferStatus(batch_id, 0, status); + if (!s.ok()) { + LOG(ERROR) << "getTransferStatus failed: " << s.ToString(); + engine->freeBatchID(batch_id); + return false; + } + if (status.s == TransferStatusEnum::COMPLETED) { + engine->freeBatchID(batch_id); + return true; + } + if (status.s == TransferStatusEnum::FAILED) { + LOG(ERROR) << "Transfer FAILED"; + engine->freeBatchID(batch_id); + return false; + } + } + LOG(ERROR) << "Transfer timed out"; + engine->freeBatchID(batch_id); + return false; + } + + std::string metadata_server_; + std::string local_server_name_; +}; + +// Test 1: Verify CXI transport can be installed +TEST_F(CXITransportTest, InstallTransport) { + auto engine = std::make_unique(false); + engine->getLocalTopology()->discover({}); + auto hp = parseHostNameWithPort(local_server_name_); + int rc = engine->init(metadata_server_, local_server_name_, + hp.first.c_str(), hp.second); + ASSERT_EQ(rc, 0); + + Transport *xport = engine->installTransport("cxi", nullptr); + ASSERT_NE(xport, nullptr) + << "CXI transport should be installable on CXI hardware"; +} + +// Test 2: Basic loopback write +TEST_F(CXITransportTest, LoopbackWrite) { + auto setup = createEngine(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 4096; + + // Fill source buffer with known data + memset(setup.addr, 0xAB, kDataLength); + + bool ok = submitAndWait(setup.engine.get(), setup.segment_id, setup.addr, + remote_base, kDataLength, TransferRequest::WRITE); + EXPECT_TRUE(ok) << "Loopback write should succeed"; + + destroyEngine(setup); +} + +#ifdef USE_GPU +TEST_F(CXITransportTest, LoopbackWriteDevice) { + cudaSetDevice(0); + auto setup = createEngineDevice(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 4096; + + // Fill source buffer with known data + cudaMemset(setup.addr, 0xAB, kDataLength); + cudaDeviceSynchronize(); + + bool ok = submitAndWait(setup.engine.get(), setup.segment_id, setup.addr, + remote_base, kDataLength, TransferRequest::WRITE); + EXPECT_TRUE(ok) << "Loopback write should succeed"; + + destroyEngineDevice(setup); +} +#endif + +// Test 3: Write then read, verify data integrity +TEST_F(CXITransportTest, WriteAndRead) { + auto setup = createEngine(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 4096000; + uint8_t *buf = (uint8_t *)setup.addr; + + // Fill first half with random data + for (size_t i = 0; i < kDataLength; ++i) buf[i] = 'a' + lrand48() % 26; + + // Write local -> remote (loopback) + bool ok = submitAndWait(setup.engine.get(), setup.segment_id, buf, + remote_base, kDataLength, TransferRequest::WRITE); + ASSERT_TRUE(ok) << "Write should succeed"; + + // Read remote -> local (into second half of buffer) + ok = submitAndWait(setup.engine.get(), setup.segment_id, buf + kDataLength, + remote_base, kDataLength, TransferRequest::READ); + ASSERT_TRUE(ok) << "Read should succeed"; + + // Verify data integrity + EXPECT_EQ(0, memcmp(buf, buf + kDataLength, kDataLength)) + << "Read-back data should match written data"; + + destroyEngine(setup); +} + +#ifdef USE_GPU +TEST_F(CXITransportTest, WriteAndReadDevice) { + cudaSetDevice(0); + auto setup = createEngineDevice(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 4096000; + uint8_t *buf = (uint8_t *)setup.addr; + + uint8_t random_data = 'a' + lrand48() % 26; + cudaMemset(buf, random_data, kDataLength); + cudaDeviceSynchronize(); + + // Write local -> remote (loopback) + bool ok = submitAndWait(setup.engine.get(), setup.segment_id, buf, + remote_base, kDataLength, TransferRequest::WRITE); + ASSERT_TRUE(ok) << "Write should succeed"; + + // Read remote -> local (into second half of buffer) + ok = submitAndWait(setup.engine.get(), setup.segment_id, buf + kDataLength, + remote_base, kDataLength, TransferRequest::READ); + ASSERT_TRUE(ok) << "Read should succeed"; + // copy data back to host + uint8_t *dataHost = (uint8_t *)allocateMemoryPool(2 * kDataLength, 0); + cudaMemcpy(dataHost, buf, 2 * kDataLength, cudaMemcpyDeviceToHost); + cudaDeviceSynchronize(); + // Verify data integrity + EXPECT_EQ(0, memcmp(dataHost, dataHost + kDataLength, kDataLength)) + << "Read-back data should match written data"; + + destroyEngineDevice(setup); +} +#endif + +// Test 4: Multiple sequential writes in a batch +TEST_F(CXITransportTest, MultiWrite) { + auto setup = createEngine(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 65536; + const int kBatchSize = 16; + + auto batch_id = setup.engine->allocateBatchID(kBatchSize); + + std::vector requests; + for (int i = 0; i < kBatchSize; ++i) { + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = (uint8_t *)setup.addr + i * kDataLength; + entry.target_id = setup.segment_id; + entry.target_offset = remote_base + i * kDataLength; + requests.push_back(entry); + } + + Status s = setup.engine->submitTransfer(batch_id, requests); + ASSERT_TRUE(s.ok()) << "submitTransfer failed: " << s.ToString(); + + // Poll all tasks until completion + for (int task_id = 0; task_id < kBatchSize; ++task_id) { + TransferStatus status; + const int kMaxPollIterations = 1000000; + for (int i = 0; i < kMaxPollIterations; ++i) { + s = setup.engine->getTransferStatus(batch_id, task_id, status); + ASSERT_TRUE(s.ok()); + if (status.s == TransferStatusEnum::COMPLETED) break; + ASSERT_NE(status.s, TransferStatusEnum::FAILED) + << "Task " << task_id << " failed"; + } + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED) + << "Task " << task_id << " did not complete"; + } + + s = setup.engine->freeBatchID(batch_id); + ASSERT_TRUE(s.ok()); + + destroyEngine(setup); +} + +// Test 5: Stress test - multiple batches to verify no CQ overflow +TEST_F(CXITransportTest, StressMultipleBatches) { + auto setup = createEngine(); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kDataLength = 65536; + const int kBatchSize = 8; + const int kNumBatches = 20; + + for (int batch = 0; batch < kNumBatches; ++batch) { + auto batch_id = setup.engine->allocateBatchID(kBatchSize); + + std::vector requests; + for (int i = 0; i < kBatchSize; ++i) { + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = + (uint8_t *)setup.addr + (i + batch * kBatchSize) * kDataLength; + entry.target_id = setup.segment_id; + entry.target_offset = + remote_base + (i + batch * kBatchSize) * kDataLength; + requests.push_back(entry); + } + + Status s = setup.engine->submitTransfer(batch_id, requests); + ASSERT_TRUE(s.ok()) + << "Batch " << batch << " submitTransfer failed: " << s.ToString(); + + // Wait for all tasks in batch + for (int task_id = 0; task_id < kBatchSize; ++task_id) { + TransferStatus status; + const int kMaxPollIterations = 1000000; + for (int i = 0; i < kMaxPollIterations; ++i) { + s = setup.engine->getTransferStatus(batch_id, task_id, status); + ASSERT_TRUE(s.ok()); + if (status.s == TransferStatusEnum::COMPLETED) break; + ASSERT_NE(status.s, TransferStatusEnum::FAILED) + << "Batch " << batch << " task " << task_id << " failed"; + } + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED) + << "Batch " << batch << " task " << task_id + << " did not complete"; + } + + s = setup.engine->freeBatchID(batch_id); + ASSERT_TRUE(s.ok()); + } + + destroyEngine(setup); +} + +// Test 6: warmupSegment on loopback peer +// +// Exercises CxiTransport::warmupSegment() which is the C++ entry point +// behind the warmup_efa_segment() Python binding / warmupSegment() C API. +// Loopback is enough to cover the handshake + fi_av_insert path AND the +// idempotent short-circuit on the second call. +TEST_F(CXITransportTest, WarmupSegmentLoopback) { + auto setup = createEngine(); + + auto *cxi = dynamic_cast(setup.xport); + ASSERT_NE(cxi, nullptr) << "installTransport did not return a CxiTransport"; + + // First call: should connect every (local NIC x peer NIC) pair. + int rc = cxi->warmupSegment(setup.engine->getLocalIpAndPort()); + EXPECT_EQ(rc, 0) << "warmupSegment should succeed on loopback"; + + // Second call: should short-circuit (all endpoints already connected). + rc = cxi->warmupSegment(setup.engine->getLocalIpAndPort()); + EXPECT_EQ(rc, 0) << "warmupSegment should be idempotent"; + + // Empty / self-name: short-circuit path returning 0 without touching AV. + rc = cxi->warmupSegment(""); + EXPECT_EQ(rc, 0) << "warmupSegment(\"\") should be a no-op"; + + destroyEngine(setup); +} + +// Test 7: warmupSegment on a non-existent segment name should fail cleanly +// (no crash, no hang) rather than blocking for the poll timeout. +TEST_F(CXITransportTest, WarmupSegmentNotFound) { + auto setup = createEngine(); + + auto *cxi = dynamic_cast(setup.xport); + ASSERT_NE(cxi, nullptr); + + int rc = cxi->warmupSegment("127.0.0.1:1"); // not openSegment'd + EXPECT_NE(rc, 0) << "warmupSegment should fail for unknown segment"; + + destroyEngine(setup); +} + +// Test 8: registerLocalMemoryBatch / unregisterLocalMemoryBatch round-trip. +// Covers the batched MR path which the single-buffer tests above never hit. +TEST_F(CXITransportTest, RegisterMemoryBatch) { + auto engine = std::make_unique(false); + engine->getLocalTopology()->discover({}); + auto hp = parseHostNameWithPort(local_server_name_); + int rc = engine->init(metadata_server_, local_server_name_, + hp.first.c_str(), hp.second); + ASSERT_EQ(rc, 0); + + Transport *xport = engine->installTransport("cxi", nullptr); + ASSERT_NE(xport, nullptr); + + const size_t kBufSize = 4ull << 20; // 4 MB each + const int kNumBufs = 4; + std::vector addrs; + std::vector entries; + for (int i = 0; i < kNumBufs; ++i) { + void *a = allocateMemoryPool(kBufSize, 0); + ASSERT_NE(a, nullptr); + addrs.push_back(a); + entries.push_back({a, kBufSize}); + } + + rc = engine->registerLocalMemoryBatch(entries, "cpu:0"); + EXPECT_EQ(rc, 0) << "registerLocalMemoryBatch should succeed"; + + rc = engine->unregisterLocalMemoryBatch(addrs); + EXPECT_EQ(rc, 0) << "unregisterLocalMemoryBatch should succeed"; + + for (void *a : addrs) freeMemoryPool(a, kBufSize); +} + +// Test 9: Larger transfer (64 MB total split into 1 MB slices) to exercise +// the WR / CQ pacing logic in CxiContext::submitSlicesOnPeer beyond what the +// 16 x 64 KB MultiWrite test reaches. +TEST_F(CXITransportTest, LargeTransfer) { + const size_t kBufSize = 128ull << 20; // 128 MB + auto setup = createEngine(kBufSize); + + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + + const size_t kSliceLen = 1ull << 20; // 1 MB per slice + const int kNumSlices = 64; // 64 MB total + ASSERT_LE(static_cast(kNumSlices) * kSliceLen, kBufSize / 2); + + // Fill first half with known data + uint8_t *buf = (uint8_t *)setup.addr; + for (size_t i = 0; i < static_cast(kNumSlices) * kSliceLen; ++i) + buf[i] = (uint8_t)(i & 0xFF); + + auto batch_id = setup.engine->allocateBatchID(kNumSlices); + std::vector requests; + requests.reserve(kNumSlices); + for (int i = 0; i < kNumSlices; ++i) { + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kSliceLen; + entry.source = buf + i * kSliceLen; + entry.target_id = setup.segment_id; + entry.target_offset = remote_base + (kBufSize / 2) + i * kSliceLen; + requests.push_back(entry); + } + + Status s = setup.engine->submitTransfer(batch_id, requests); + ASSERT_TRUE(s.ok()) << "submitTransfer failed: " << s.ToString(); + + for (int task_id = 0; task_id < kNumSlices; ++task_id) { + TransferStatus status; + const int kMaxPollIterations = 2000000; + int i = 0; + for (; i < kMaxPollIterations; ++i) { + s = setup.engine->getTransferStatus(batch_id, task_id, status); + ASSERT_TRUE(s.ok()); + if (status.s == TransferStatusEnum::COMPLETED) break; + ASSERT_NE(status.s, TransferStatusEnum::FAILED); + } + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED) + << "task " << task_id << " did not complete"; + } + + s = setup.engine->freeBatchID(batch_id); + ASSERT_TRUE(s.ok()); + + // Verify byte-level integrity of the last slice (spot check). + EXPECT_EQ(0, memcmp(buf + (kNumSlices - 1) * kSliceLen, + buf + (kBufSize / 2) + (kNumSlices - 1) * kSliceLen, + kSliceLen)); + + destroyEngine(setup); +} + +// Test 10: Repeated open/close of the same remote segment must not leak AV +// slots or break loopback transfers — this is the setPeerNicPath-detach path +// that target restarts depend on under the SRD shared-endpoint model. +TEST_F(CXITransportTest, RepeatedOpenSegment) { + auto setup = createEngine(); + + auto actual_addr = setup.engine->getLocalIpAndPort(); + + // First write via setup.segment_id (from createEngine()). + auto segment_desc = + setup.engine->getMetadata()->getSegmentDescByID(setup.segment_id); + ASSERT_NE(segment_desc, nullptr); + uint64_t remote_base = (uint64_t)segment_desc->buffers[0].addr; + memset(setup.addr, 0xCD, 4096); + EXPECT_TRUE(submitAndWait(setup.engine.get(), setup.segment_id, setup.addr, + remote_base, 4096, TransferRequest::WRITE)); + + // Re-open same segment several times; each should return a working handle + // and subsequent writes should still succeed. + for (int i = 0; i < 5; ++i) { + SegmentID sid = setup.engine->openSegment(actual_addr); + ASSERT_NE(sid, (SegmentID)-1); + auto desc = setup.engine->getMetadata()->getSegmentDescByID(sid); + ASSERT_NE(desc, nullptr); + uint64_t base = (uint64_t)desc->buffers[0].addr; + ASSERT_TRUE(submitAndWait(setup.engine.get(), sid, setup.addr, base, + 4096, TransferRequest::WRITE)) + << "write #" << i << " after re-open failed"; + } + + destroyEngine(setup); +} + +} // namespace mooncake + +int main(int argc, char **argv) { + gflags::ParseCommandLineFlags(&argc, &argv, false); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-transfer-engine/tests/cxi_unit_tests.cpp b/mooncake-transfer-engine/tests/cxi_unit_tests.cpp new file mode 100644 index 0000000000..d5e954c1f2 --- /dev/null +++ b/mooncake-transfer-engine/tests/cxi_unit_tests.cpp @@ -0,0 +1,124 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "transfer_engine.h" +#include "transport/transport.h" +#include "transport/cxi_transport/cxi_transport.h" +#include "common.h" + +using namespace mooncake; + +namespace mooncake { + +class CxiTransportUnitTest : public ::testing::Test { + protected: + void SetUp() override { + metadata_server_ = "P2PHANDSHAKE"; + LOG(INFO) << "metadata_server: " << metadata_server_; + + local_server_name_ = "127.0.0.1:12345"; + LOG(INFO) << "local_server_name: " << local_server_name_; + } + + void TearDown() override {} + + std::string metadata_server_; + std::string local_server_name_; +}; + +TEST_F(CxiTransportUnitTest, TestInit) { + std::unique_ptr engine; + std::string device_filter = "cxi0"; + engine = std::make_unique( + false, std::vector{device_filter}); + EXPECT_NE(engine, nullptr) << "failed to allocate TransferEngine"; + + engine->getLocalTopology()->discover({}); + auto hp = parseHostNameWithPort(local_server_name_); + int rc = engine->init(metadata_server_, local_server_name_, + hp.first.c_str(), hp.second); + EXPECT_EQ(rc, 0) << "engine->init failed"; + + Transport* xport = engine->installTransport("cxi", nullptr); + EXPECT_NE(xport, nullptr) << "installTransport(\"cxi\") failed"; + CxiTransport* cxi_xport = dynamic_cast(xport); + EXPECT_NE(cxi_xport, nullptr) << "initialization for cxi failed"; + + engine.reset(); + + LOG(INFO) << "default init ok"; + + device_filter = "invalid_device"; // dummy device name + engine = std::make_unique( + false, std::vector{device_filter}); + EXPECT_NE(engine, nullptr) << "failed to allocate TransferEngine"; + engine->getLocalTopology()->discover({device_filter}); + hp = parseHostNameWithPort(local_server_name_); + rc = engine->init(metadata_server_, local_server_name_, hp.first.c_str(), + hp.second); + EXPECT_EQ(rc, 0) << "engine->init failed"; + xport = engine->installTransport("cxi", nullptr); + EXPECT_EQ(xport, nullptr) + << "installTransport(\"cxi\") should fail for invalid_device"; +} + +TEST_F(CxiTransportUnitTest, TestAllocate) { + std::unique_ptr engine; + engine = std::make_unique(false); + EXPECT_NE(engine, nullptr) << "failed to allocate TransferEngine"; + + engine->getLocalTopology()->discover({}); + auto hp = parseHostNameWithPort(local_server_name_); + int rc = engine->init(metadata_server_, local_server_name_, + hp.first.c_str(), hp.second); + EXPECT_EQ(rc, 0) << "engine->init failed"; + + Transport* xport = engine->installTransport("cxi", nullptr); + EXPECT_NE(xport, nullptr) << "installTransport(\"cxi\") failed"; + CxiTransport* cxi_xport = dynamic_cast(xport); + EXPECT_NE(cxi_xport, nullptr) << "initialization for cxi failed"; + + size_t buffer_size = 128 * 1024; + auto buffer = std::vector(buffer_size); + size_t dummy_buffer_size = 128; + auto dummy_buffer = std::vector(dummy_buffer_size); + + // check invalid device + int retcode = engine->registerLocalMemory( + (void*)dummy_buffer.data(), dummy_buffer_size, "invalid_device:0"); + EXPECT_EQ(retcode, ERR_DEVICE_NOT_FOUND) + << "invalid device memory should not be registered!"; + + retcode = + engine->registerLocalMemory((void*)buffer.data(), buffer_size, "cpu:0"); + EXPECT_EQ(retcode, 0) << "unable to register memory"; + + retcode = engine->registerLocalMemory( + (void*)(buffer.data() + (buffer_size >> 1)), buffer_size >> 1, "cpu:1"); + EXPECT_EQ(retcode, ERR_ADDRESS_OVERLAPPED) + << "should not be able to register overlapping MRs!"; + + engine->unregisterLocalMemory(buffer.data()); +} + +}; // namespace mooncake + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, false); + ::testing::InitGoogleTest(&argc, argv); + google::InitGoogleLogging("CxiTransportUnitTest"); + FLAGS_logtostderr = 1; + int rc = RUN_ALL_TESTS(); + google::ShutdownGoogleLogging(); + return rc; +} diff --git a/mooncake-transfer-engine/tests/device_backend_api_compatibility_test.cu b/mooncake-transfer-engine/tests/device_backend_api_compatibility_test.cu new file mode 100644 index 0000000000..7485fa5748 --- /dev/null +++ b/mooncake-transfer-engine/tests/device_backend_api_compatibility_test.cu @@ -0,0 +1,80 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "transport/device/comm_device.cuh" +#include "transport/device/nccl_device.cuh" + +#include + +namespace mooncake { +namespace device { +namespace compatibility_test { + +struct NcclOps { + using Context = NcclDeviceContext; + + __device__ __forceinline__ static void put(const Context& ctx, int channel, + int peer, int qps_per_rank, + const void* source, + void* destination, + uint32_t bytes, int lane) { + mc_nccl_put(ctx, channel, peer, qps_per_rank, source, destination, + bytes, lane); + } +}; + +struct IbgdaOps { + using Context = CommCtx; + + __device__ __forceinline__ static void put(const Context& ctx, int channel, + int peer, int qps_per_rank, + const void* source, + void* destination, + uint32_t bytes, int lane) { + mc_rdma_put(ctx, channel, peer, qps_per_rank, source, destination, + bytes, lane); + } +}; + +template +__device__ __forceinline__ void exchange(const typename Ops::Context& ctx, + int channel, int peer, + int qps_per_rank, const void* source, + void* destination, uint32_t bytes, + int lane) { + Ops::put(ctx, channel, peer, qps_per_rank, source, destination, bytes, + lane); +} + +// The NCCL transport library target does not instantiate these header-only +// device helpers; its CUDA example is opt-in. These compile probes deliberately +// share the same templated operation and are never executed, so the test needs +// neither GPUs nor NICs. +__global__ void instantiateNccl(NcclOps::Context ctx, int channel, int peer, + int qps_per_rank, const void* source, + void* destination, uint32_t bytes, int lane) { + exchange(ctx, channel, peer, qps_per_rank, source, destination, + bytes, lane); +} + +__global__ void instantiateIbgda(IbgdaOps::Context ctx, int channel, int peer, + int qps_per_rank, const void* source, + void* destination, uint32_t bytes, int lane) { + exchange(ctx, channel, peer, qps_per_rank, source, destination, + bytes, lane); +} + +} // namespace compatibility_test +} // namespace device +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/dmabuf_export_test.cpp b/mooncake-transfer-engine/tests/dmabuf_export_test.cpp new file mode 100644 index 0000000000..2e477a044d --- /dev/null +++ b/mooncake-transfer-engine/tests/dmabuf_export_test.cpp @@ -0,0 +1,154 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Hardware-free unit tests for DmabufExport and +// RdmaContext::{exportDmabuf, closeDmabufExport}. +// +// No RDMA device or GPU is required. The tests cover: +// - DmabufExport struct defaults +// - closeDmabufExport: closes a real fd, is idempotent, no-ops when fd == -1 +// - exportDmabuf: host memory always yields kHostReg / fd == -1 on every +// build (CUDA, HIP, or non-GPU), so those cases run in CI without hardware. + +#include "transport/rdma_transport/rdma_context.h" + +#include +#include +#include +#include + +#include +#include +#include + +using mooncake::DmabufExport; +using mooncake::RdmaContext; + +namespace { + +// Open a real fd via pipe() (POSIX, no _GNU_SOURCE required). +// Returns via reference so ASSERT_EQ can abort the test on failure, avoiding +// undefined behaviour from an uninitialized pipefd if pipe() fails. +static void make_test_fd(int &out_fd) { + int pipefd[2]; + ASSERT_EQ(pipe(pipefd), 0) << "pipe() failed: " << strerror(errno); + close(pipefd[1]); // write end not needed + out_fd = pipefd[0]; +} + +static bool fd_is_closed(int fd) { + return fcntl(fd, F_GETFD) == -1 && errno == EBADF; +} + +// ── DmabufExport struct defaults ───────────────────────────────────────────── + +TEST(DmabufExport, DefaultIsHostRegWithNoFd) { + DmabufExport exp; + EXPECT_EQ(exp.method, DmabufExport::Method::kHostReg); + EXPECT_EQ(exp.fd, -1); + EXPECT_EQ(exp.offset, 0u); +} + +// ── closeDmabufExport ──────────────────────────────────────────────────────── + +TEST(CloseDmabufExport, NoOpWhenFdIsNegative) { + DmabufExport exp; // fd == -1 by default + RdmaContext::closeDmabufExport(exp); + EXPECT_EQ(exp.fd, -1); // still -1, no crash +} + +TEST(CloseDmabufExport, ClosesLiveFdAndClearsIt) { + int fd = -1; + make_test_fd(fd); + ASSERT_GE(fd, 0); + + DmabufExport exp; + exp.method = DmabufExport::Method::kDmabufReg; + exp.fd = fd; + + RdmaContext::closeDmabufExport(exp); + + EXPECT_EQ(exp.fd, -1); + EXPECT_TRUE(fd_is_closed(fd)) << "fd " << fd << " should be closed"; +} + +TEST(CloseDmabufExport, Idempotent) { + int fd = -1; + make_test_fd(fd); + ASSERT_GE(fd, 0); + + DmabufExport exp; + exp.method = DmabufExport::Method::kDmabufReg; + exp.fd = fd; + + RdmaContext::closeDmabufExport(exp); // first close + RdmaContext::closeDmabufExport( + exp); // second call: fd == -1, must not crash + EXPECT_EQ(exp.fd, -1); +} + +// ── exportDmabuf on host memory ────────────────────────────────────────────── +// +// malloc'd memory is host memory on every supported build: +// - Non-GPU build: the #else branch returns kHostReg immediately. +// - CUDA build: cuPointerGetAttribute fails for host addrs → kHostReg. +// - HIP build: hipPointerGetAttributes fails / returns hipMemoryTypeHost. +// +// So these tests exercise the "not GPU memory" fast-path on all CI runners. + +TEST(ExportDmabuf, HostMemoryYieldsHostReg) { + std::vector buf(4096); + DmabufExport exp; + int ret = RdmaContext::exportDmabuf(buf.data(), exp); + EXPECT_EQ(ret, 0); + EXPECT_EQ(exp.method, DmabufExport::Method::kHostReg); + EXPECT_EQ(exp.fd, -1); + // Closing a kHostReg export is always safe. + RdmaContext::closeDmabufExport(exp); +} + +TEST(ExportDmabuf, LargeHostBufferYieldsHostReg) { + constexpr size_t kSize = 8ULL * 1024 * 1024; // 8 MiB + std::vector buf(kSize); + DmabufExport exp; + int ret = RdmaContext::exportDmabuf(buf.data(), exp); + EXPECT_EQ(ret, 0); + EXPECT_EQ(exp.method, DmabufExport::Method::kHostReg); + EXPECT_EQ(exp.fd, -1); +} + +TEST(ExportDmabuf, MmapAnonymousYieldsHostReg) { + void *p = mmap(nullptr, 4096, PROT_READ | PROT_WRITE, + MAP_ANONYMOUS | MAP_PRIVATE, -1, 0); + ASSERT_NE(p, MAP_FAILED); + + DmabufExport exp; + int ret = RdmaContext::exportDmabuf(p, exp); + EXPECT_EQ(ret, 0); + EXPECT_EQ(exp.method, DmabufExport::Method::kHostReg); + EXPECT_EQ(exp.fd, -1); + + munmap(p, 4096); +} + +TEST(ExportDmabuf, StackAddressYieldsHostReg) { + char stack_buf[128]; + DmabufExport exp; + int ret = RdmaContext::exportDmabuf(stack_buf, exp); + EXPECT_EQ(ret, 0); + EXPECT_EQ(exp.method, DmabufExport::Method::kHostReg); + EXPECT_EQ(exp.fd, -1); +} + +} // namespace diff --git a/mooncake-transfer-engine/tests/efa_transport_test.cpp b/mooncake-transfer-engine/tests/efa_transport_test.cpp index 704a338409..e88372028a 100644 --- a/mooncake-transfer-engine/tests/efa_transport_test.cpp +++ b/mooncake-transfer-engine/tests/efa_transport_test.cpp @@ -19,6 +19,8 @@ #include #include +#include +#include #include "transfer_engine.h" #include "transport/efa_transport/efa_transport.h" @@ -405,6 +407,47 @@ TEST_F(EFATransportTest, RegisterMemoryBatch) { for (void *a : addrs) freeMemoryPool(a, kBufSize); } +// MC_EFA_NIC_SELECTION=local narrows a device buffer to the NICs the topology +// reports as closest to its GPU. Pure map inversion, so no hardware needed. +TEST(EFALocalNicMapTest, InvertsPreferredHca) { + // Two GPUs, two NICs each, shaped like a p5 rail group. + TopologyMatrix matrix; + matrix["cuda:0"] = TopologyEntry{.name = "cuda:0", + .preferred_hca = {"efa0", "efa1"}, + .avail_hca = {"efa2", "efa3"}}; + matrix["cuda:1"] = TopologyEntry{.name = "cuda:1", + .preferred_hca = {"efa2", "efa3"}, + .avail_hca = {"efa0", "efa1"}}; + std::vector devices{"efa0", "efa1", "efa2", "efa3"}; + + auto map = EfaTransport::buildLocalNicMap(matrix, devices); + ASSERT_EQ(map.size(), 2u); + EXPECT_EQ(map["cuda:0"], (std::vector{0, 1})); + EXPECT_EQ(map["cuda:1"], (std::vector{2, 3})); + + // A preferred HCA absent from the opened set is skipped, not mapped to a + // bogus index -- this happens when a device's construct() failed. + auto partial = EfaTransport::buildLocalNicMap(matrix, {"efa0", "efa3"}); + EXPECT_EQ(partial["cuda:0"], (std::vector{0})); + EXPECT_EQ(partial["cuda:1"], (std::vector{1})); + + // A location with no surviving preferred NIC is omitted entirely, so the + // caller's lookup misses and falls back to all NICs rather than registering + // the buffer on none. + auto none = EfaTransport::buildLocalNicMap(matrix, {"efa2"}); + EXPECT_EQ(none.count("cuda:0"), 0u); + EXPECT_EQ(none["cuda:1"], (std::vector{0})); + + // An entry that has only avail_hca contributes nothing: avail_hca is "every + // other NIC", which is the all-NICs behavior this exists to avoid. + TopologyMatrix avail_only; + avail_only["cuda:0"] = TopologyEntry{ + .name = "cuda:0", .preferred_hca = {}, .avail_hca = {"efa0", "efa1"}}; + EXPECT_TRUE(EfaTransport::buildLocalNicMap(avail_only, devices).empty()); + + EXPECT_TRUE(EfaTransport::buildLocalNicMap({}, devices).empty()); +} + // Test 9: Larger transfer (64 MB total split into 1 MB slices) to exercise // the WR / CQ pacing logic in EfaContext::submitSlicesOnPeer beyond what the // 16 x 64 KB MultiWrite test reaches. diff --git a/mooncake-transfer-engine/tests/fault-tolerant/fault_test.py b/mooncake-transfer-engine/tests/fault-tolerant/fault_test.py index e764cc771b..25f9b922e1 100644 --- a/mooncake-transfer-engine/tests/fault-tolerant/fault_test.py +++ b/mooncake-transfer-engine/tests/fault-tolerant/fault_test.py @@ -1,8 +1,6 @@ import subprocess import time -import signal import sys -import os from typing import Optional def start_process(cmd: list[str]) -> subprocess.Popen: diff --git a/mooncake-transfer-engine/tests/fault-tolerant/transfer_engine.py b/mooncake-transfer-engine/tests/fault-tolerant/transfer_engine.py index e1482b4ea1..9eeadad060 100644 --- a/mooncake-transfer-engine/tests/fault-tolerant/transfer_engine.py +++ b/mooncake-transfer-engine/tests/fault-tolerant/transfer_engine.py @@ -1,6 +1,4 @@ -import json import logging -from dataclasses import dataclass from typing import Optional logger = logging.getLogger(__name__) diff --git a/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp b/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp new file mode 100644 index 0000000000..f0741c82bf --- /dev/null +++ b/mooncake-transfer-engine/tests/graceful_shutdown_test.cpp @@ -0,0 +1,156 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include + +#include + +#include "transfer_engine.h" + +using namespace mooncake; + +namespace { + +void waitChildWithTimeout(pid_t pid, int* status) { + for (int i = 0; i < 50; ++i) { + pid_t ret = waitpid(pid, status, WNOHANG); + ASSERT_NE(ret, -1) << "waitpid() failed"; + if (ret == pid) return; + usleep(100000); + } + kill(pid, SIGKILL); + waitpid(pid, status, 0); + FAIL() << "child did not exit before timeout"; +} + +} // namespace + +TEST(GracefulShutdownTest, SigtermTriggersCleanExit) { + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGTERM); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)) + << "Child did not exit normally (signaled: " << WIFSIGNALED(status) + << ")"; + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGTERM); +} + +TEST(GracefulShutdownTest, SigintTriggersCleanExit) { + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGINT); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)) << "Child did not exit normally"; + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGINT); +} + +TEST(GracefulShutdownTest, IdempotentEnable) { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + engine->enableGracefulShutdown(); + engine->enableGracefulShutdown(); +} + +TEST(GracefulShutdownTest, EngineDestroyedBeforeSignal) { + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + } + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGTERM); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)); + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGTERM); +} + +TEST(GracefulShutdownTest, ForkAfterInstallDoesNotHangChildSignal) { + auto engine = std::make_unique(false); + engine->enableGracefulShutdown(); + + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGTERM); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)) + << "Child did not exit normally (signaled: " << WIFSIGNALED(status) + << ")"; + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGTERM); +} + +TEST(GracefulShutdownTest, MultipleEngines) { + pid_t pid = fork(); + ASSERT_NE(pid, -1) << "fork() failed"; + + if (pid == 0) { + auto engine1 = std::make_unique(false); + auto engine2 = std::make_unique(false); + engine1->enableGracefulShutdown(); + engine2->enableGracefulShutdown(); + pause(); + _exit(99); + } + + usleep(100000); + kill(pid, SIGTERM); + + int status; + waitChildWithTimeout(pid, &status); + ASSERT_TRUE(WIFEXITED(status)); + EXPECT_EQ(WEXITSTATUS(status), 128 + SIGTERM); +} diff --git a/mooncake-transfer-engine/tests/hip_transport_test.cpp b/mooncake-transfer-engine/tests/hip_transport_test.cpp new file mode 100644 index 0000000000..0da4f88d70 --- /dev/null +++ b/mooncake-transfer-engine/tests/hip_transport_test.cpp @@ -0,0 +1,133 @@ +// Copyright 2025 Mooncake Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include + +#include +#include + +#include "cuda_alike.h" +#include "transfer_engine.h" +#include "transfer_metadata.h" // P2PHANDSHAKE +#include "transport/transport.h" + +using namespace mooncake; + +// startAsyncTransfer() switches the active device to the source GPU; if it does +// not restore it, the calling thread (which also launches the engine's compute +// kernels) is left on the wrong GPU and the next kernel fails with +// hipErrorInvalidDevice. This pins the caller to one GPU, transfers from a +// buffer on a different GPU, and asserts the active device is unchanged. + +namespace { +constexpr size_t kLen = 64 * 1024; + +void* allocOnDevice(size_t size, int device) { + EXPECT_EQ(cudaSetDevice(device), cudaSuccess); + void* ptr = nullptr; + EXPECT_EQ(cudaMalloc(&ptr, size), cudaSuccess); + return ptr; +} +} // namespace + +TEST(HipTransportTest, RestoresActiveDeviceAfterTransfer) { + int device_count = 0; + ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess); + if (device_count < 2) { + GTEST_SKIP() << "Needs >= 2 GPUs: the transfer source must live on a " + "different device than the calling thread's device."; + } + + const int kCallerDevice = 0; // device the engine thread runs on + const int kSourceDevice = 1; // KV source lives on a different GPU + + // P2PHANDSHAKE: self-contained loopback, no external metadata server. + auto engine = std::make_unique(false); + const std::string server_name = "127.0.0.1:17813"; + if (engine->init(P2PHANDSHAKE, server_name, "127.0.0.1", 17813) != 0) { + GTEST_SKIP() << "TransferEngine init failed in this environment."; + } + + Transport* transport = engine->installTransport("hip", nullptr); + if (transport == nullptr) { + GTEST_SKIP() + << "HIP transport unavailable (built without -DUSE_HIP=ON?)."; + } + + // Both buffers on kSourceDevice so the source GPU differs from the caller. + void* src = allocOnDevice(kLen, kSourceDevice); + void* dst = allocOnDevice(kLen, kSourceDevice); + ASSERT_EQ(engine->registerLocalMemory( + src, kLen, GPU_PREFIX + std::to_string(kSourceDevice)), + 0); + ASSERT_EQ(engine->registerLocalMemory( + dst, kLen, GPU_PREFIX + std::to_string(kSourceDevice)), + 0); + + auto segment_id = engine->openSegment(server_name); + ASSERT_GE(segment_id, 0); + + ASSERT_EQ(cudaSetDevice(kSourceDevice), cudaSuccess); + ASSERT_EQ(cudaMemset(src, 0xAB, kLen), cudaSuccess); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + + // Pin the caller to a different device than the source, as the engine does. + ASSERT_EQ(cudaSetDevice(kCallerDevice), cudaSuccess); + + auto batch_id = engine->allocateBatchID(1); + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kLen; + entry.source = src; + entry.target_id = segment_id; + entry.target_offset = reinterpret_cast(dst); + Status s = engine->submitTransfer(batch_id, {entry}); + ASSERT_TRUE(s.ok()); + + // Invariant: the caller's device is unchanged (== kSourceDevice before + // fix). + int active_after_submit = -1; + ASSERT_EQ(cudaGetDevice(&active_after_submit), cudaSuccess); + EXPECT_EQ(active_after_submit, kCallerDevice) + << "startAsyncTransfer() must restore the caller's active device. " + "Leaving it on the source GPU corrupts the engine thread's HIP " + "context, so its next kernel launch fails with " + "hipErrorInvalidDevice."; + + // Drain the transfer (also a basic functional check). + TransferStatus status; + do { + ASSERT_TRUE(engine->getTransferStatus(batch_id, 0, status).ok()); + } while (status.s == TransferStatusEnum::WAITING); + EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED); + + int active_after_wait = -1; + ASSERT_EQ(cudaGetDevice(&active_after_wait), cudaSuccess); + EXPECT_EQ(active_after_wait, kCallerDevice); + + engine->freeBatchID(batch_id); + engine->unregisterLocalMemory(src); + engine->unregisterLocalMemory(dst); + (void)cudaSetDevice(kSourceDevice); + (void)cudaFree(src); + (void)cudaFree(dst); +} + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, false); + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-transfer-engine/tests/multi_transport_locality_test.cpp b/mooncake-transfer-engine/tests/multi_transport_locality_test.cpp new file mode 100644 index 0000000000..750d9e60bd --- /dev/null +++ b/mooncake-transfer-engine/tests/multi_transport_locality_test.cpp @@ -0,0 +1,84 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include + +#include "multi_transport_locality.h" + +using namespace mooncake; + +TEST(MultiTransportLocalityTest, SegmentHostStripsPort) { + EXPECT_EQ(segmentHost("10.0.0.1:12345"), "10.0.0.1"); + EXPECT_EQ(segmentHost("node-a:8000"), "node-a"); + // No port: whole string is the host. + EXPECT_EQ(segmentHost("node-a"), "node-a"); + EXPECT_EQ(segmentHost(""), ""); +} + +TEST(MultiTransportLocalityTest, HipReachableForSameHostDifferentPort) { + // Two engines co-located on one host share the host, differ only in port. + EXPECT_TRUE(isHipReachableTarget("10.0.0.1:20000", "10.0.0.1:20001")); + EXPECT_TRUE(isHipReachableTarget("node-a:8000", "node-a:9000")); +} + +TEST(MultiTransportLocalityTest, HipReachableForIdenticalName) { + EXPECT_TRUE(isHipReachableTarget("node-a:8000", "node-a:8000")); +} + +TEST(MultiTransportLocalityTest, HipNotReachableForRemoteHost) { + // Cross-host target: hip (GPU IPC) is not usable, must fall back to rdma. + EXPECT_FALSE(isHipReachableTarget("10.0.0.2:20000", "10.0.0.1:20000")); + EXPECT_FALSE(isHipReachableTarget("node-b:8000", "node-a:8000")); +} + +TEST(MultiTransportLocalityTest, HandlesMissingPort) { + // Locality decision still works when one side has no explicit port. + EXPECT_TRUE(isHipReachableTarget("node-a", "node-a:8000")); + EXPECT_FALSE(isHipReachableTarget("node-b", "node-a:8000")); +} + +TEST(MultiTransportLocalityTest, SegmentHostParsesIPv6) { + // Bracketed IPv6 with and without a port. + EXPECT_EQ(segmentHost("[2001:db8::1]:8000"), "2001:db8::1"); + EXPECT_EQ(segmentHost("[2001:db8::1]"), "2001:db8::1"); + EXPECT_EQ(segmentHost("[::1]:20000"), "::1"); + // Bare IPv6 literal without a port: the whole string is the host. + EXPECT_EQ(segmentHost("2001:db8::1"), "2001:db8::1"); + EXPECT_EQ(segmentHost("::1"), "::1"); +} + +TEST(MultiTransportLocalityTest, HipReachableForIPv6) { + // Same IPv6 host, different ports (bracketed) -> intra-node hip. + EXPECT_TRUE( + isHipReachableTarget("[2001:db8::1]:20000", "[2001:db8::1]:20001")); + // Bracketed-with-port vs bare literal for the same host must still match. + EXPECT_TRUE(isHipReachableTarget("[2001:db8::1]:20000", "2001:db8::1")); + // Different IPv6 hosts -> cross-node, fall back to rdma. + EXPECT_FALSE( + isHipReachableTarget("[2001:db8::1]:20000", "[2001:db8::2]:20000")); +} + +TEST(MultiTransportLocalityTest, HostMatchIsCaseInsensitive) { + // Hostnames and IPv6 hex literals are case-insensitive. + EXPECT_TRUE(isHipReachableTarget("Node-A:8000", "node-a:9000")); + EXPECT_TRUE( + isHipReachableTarget("[2001:DB8::1]:8000", "[2001:db8::1]:9000")); +} + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-transfer-engine/tests/nvmeof_status_test.cpp b/mooncake-transfer-engine/tests/nvmeof_status_test.cpp new file mode 100644 index 0000000000..775deeb6f9 --- /dev/null +++ b/mooncake-transfer-engine/tests/nvmeof_status_test.cpp @@ -0,0 +1,127 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include + +#include "transport/nvmeof_transport/nvmeof_transport.h" + +namespace mooncake { + +class CUFileDescPoolTestPeer { + public: + static bool cachePolledEvent(std::vector& events, + const CUfileIOEvents_t& event) { + return CUFileDescPool::cachePolledEvent(events, event); + } +}; + +class NVMeoFTransportTestPeer { + public: + static std::unique_ptr createWithoutDriver() { + return std::unique_ptr( + new NVMeoFTransport(std::make_shared())); + } + + static Transport::TransferStatus aggregate( + const std::vector& statuses, + bool& is_finished) { + return NVMeoFTransport::aggregateTransferStatus(statuses, is_finished); + } +}; + +TEST(NVMeoFStatusTest, RejectsUnsupportedMultiTransportSubmission) { + auto transport = NVMeoFTransportTestPeer::createWithoutDriver(); + Transport::TransferTask task; + + auto status = transport->submitTransferTask({&task}); + + EXPECT_TRUE(status.IsNotImplemented()); + EXPECT_EQ(status.message(), + "NVMeoFTransport does not support MultiTransport batches"); + EXPECT_TRUE(task.is_finished); +} + +TEST(NVMeoFStatusTest, ReportsKnownFailureBeforeAllSlicesFinish) { + bool is_finished = true; + auto waiting_status = NVMeoFTransportTestPeer::aggregate( + {{Transport::FAILED, 0}, {Transport::WAITING, 0}}, is_finished); + EXPECT_EQ(waiting_status.s, Transport::FAILED); + EXPECT_FALSE(is_finished); + + auto pending_status = NVMeoFTransportTestPeer::aggregate( + {{Transport::FAILED, 0}, {Transport::PENDING, 0}}, is_finished); + EXPECT_EQ(pending_status.s, Transport::FAILED); + EXPECT_FALSE(is_finished); +} + +TEST(NVMeoFStatusTest, ReportsPendingOnlyWhenNoFailureIsKnown) { + bool is_finished = true; + auto status = NVMeoFTransportTestPeer::aggregate( + {{Transport::COMPLETED, 1024}, {Transport::PENDING, 0}}, is_finished); + + EXPECT_EQ(status.s, Transport::PENDING); + EXPECT_FALSE(is_finished); +} + +TEST(NVMeoFStatusTest, AggregatesCompletedBytes) { + bool is_finished = false; + auto status = NVMeoFTransportTestPeer::aggregate( + {{Transport::COMPLETED, 1024}, {Transport::COMPLETED, 2048}}, + is_finished); + + EXPECT_EQ(status.s, Transport::COMPLETED); + EXPECT_EQ(status.transferred_bytes, 3072); + EXPECT_TRUE(is_finished); +} + +TEST(NVMeoFStatusTest, UsesDeterministicTerminalFailurePrecedence) { + bool first_finished = false; + auto first = NVMeoFTransportTestPeer::aggregate({{Transport::INVALID, 0}, + {Transport::FAILED, 0}, + {Transport::TIMEOUT, 0}}, + first_finished); + + bool second_finished = false; + auto second = NVMeoFTransportTestPeer::aggregate({{Transport::TIMEOUT, 0}, + {Transport::FAILED, 0}, + {Transport::INVALID, 0}}, + second_finished); + + EXPECT_EQ(first.s, Transport::FAILED); + EXPECT_EQ(second.s, Transport::FAILED); + EXPECT_TRUE(first_finished); + EXPECT_TRUE(second_finished); +} + +TEST(NVMeoFStatusTest, CorrelatesPartialCompletionsByCookie) { + std::vector cached = { + {.cookie = reinterpret_cast(1), + .status = CUFILE_WAITING, + .ret = 0}, + {.cookie = reinterpret_cast(2), + .status = CUFILE_WAITING, + .ret = 0}}; + CUfileIOEvents_t second = {.cookie = reinterpret_cast(2), + .status = CUFILE_COMPLETE, + .ret = 4096}; + + ASSERT_TRUE(CUFileDescPoolTestPeer::cachePolledEvent(cached, second)); + EXPECT_EQ(cached[0].status, CUFILE_WAITING); + EXPECT_EQ(cached[1].status, CUFILE_COMPLETE); + EXPECT_EQ(cached[1].ret, 4096); +} + +} // namespace mooncake diff --git a/mooncake-transfer-engine/tests/rdma_async_event_drain_test.cpp b/mooncake-transfer-engine/tests/rdma_async_event_drain_test.cpp new file mode 100644 index 0000000000..fdc1df1682 --- /dev/null +++ b/mooncake-transfer-engine/tests/rdma_async_event_drain_test.cpp @@ -0,0 +1,225 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// WorkerPool::doProcessContextEvents() must empty the async event fd on every +// call: the fd is edge-triggered, so anything left queued is stranded until the +// next event arrives. These tests replace the event source with a scripted +// queue via linker wrapping, so no RDMA device is needed. + +#include + +#include +#include +#include +#include + +#include +#include + +#include "error.h" +#include "transport/rdma_transport/rdma_context.h" +#include "transport/rdma_transport/rdma_transport.h" +#include "transport/rdma_transport/worker_pool.h" + +#if defined(__has_feature) +#define MC_HAS_FEATURE(x) __has_feature(x) +#else +#define MC_HAS_FEATURE(x) 0 +#endif +#if defined(__SANITIZE_ADDRESS__) || MC_HAS_FEATURE(address_sanitizer) || \ + MC_HAS_FEATURE(leak_sanitizer) +#include +#define MC_LSAN_IGNORE_OBJECT(p) __lsan_ignore_object(p) + +// Suppress false positives from libnuma.so's process-wide static cache +// allocated in numa_node_to_cpus() which is intentionally retained until exit. +extern "C" __attribute__((weak, visibility("default"))) const char * +__lsan_default_suppressions() { + return "leak:libnuma.so\n"; +} +#else +#define MC_LSAN_IGNORE_OBJECT(p) ((void)(p)) +#endif + +using namespace mooncake; + +namespace { + +// Stands in for the kernel's async event queue behind async_fd. +struct AsyncEventScript { + std::deque pending; + // Reported once `pending` runs dry. EAGAIN mimics a drained non-blocking + // fd; anything else mimics a real read failure. + int drained_errno = EAGAIN; + // Fail this many reads with EINTR before serving the queue. + int pending_eintr = 0; + int get_calls = 0; + int ack_calls = 0; +}; + +AsyncEventScript g_script; + +} // namespace + +extern "C" int __wrap_ibv_get_async_event(struct ibv_context *context, + struct ibv_async_event *event) { + (void)context; // Never dereferenced by the code under test. + g_script.get_calls++; + if (g_script.pending_eintr > 0) { + g_script.pending_eintr--; + errno = EINTR; + return -1; + } + if (g_script.pending.empty()) { + errno = g_script.drained_errno; + return -1; + } + memset(event, 0, sizeof(*event)); + event->event_type = g_script.pending.front(); + g_script.pending.pop_front(); + return 0; +} + +extern "C" void __wrap_ibv_ack_async_event(struct ibv_async_event *event) { + (void)event; + g_script.ack_calls++; +} + +namespace mooncake { + +class WorkerPoolTestPeer { + public: + static int processContextEvents(WorkerPool &worker_pool) { + return worker_pool.doProcessContextEvents(); + } +}; + +} // namespace mooncake + +namespace { + +// Only event types whose handlers touch nothing a device-less context lacks. +// IBV_EVENT_COMM_EST is not claimed by handleContextEvent(), so the caller acks +// it; IBV_EVENT_PORT_ACTIVE is claimed and only stores a recovery timestamp, so +// the handler acks it. Together they cover both ack paths. +constexpr ibv_event_type kUnclaimedEvent = IBV_EVENT_COMM_EST; +constexpr ibv_event_type kClaimedEvent = IBV_EVENT_PORT_ACTIVE; + +class AsyncEventDrainTest : public ::testing::Test { + protected: + void SetUp() override { + // The pool's monitorWorker busy-polls epoll_wait() on an invalid fd + // for its whole lifetime because this context was never opened. + // Silence it so gtest failures stay readable. + previous_min_log_level_ = FLAGS_minloglevel; + FLAGS_minloglevel = google::GLOG_FATAL; + + // Intentional leak: ~RdmaTransport dereferences metadata_, which is + // null until install(). We only need it as RdmaContext's owner, same + // as rdma_endpoint_state_test. + transport_ = new RdmaTransport(); + MC_LSAN_IGNORE_OBJECT(transport_); + context_ = std::make_unique(*transport_, "unused"); + // Always fails, on any host, because no device is named "unused" -- + // which keeps the context from starting a second worker pool that + // would race us for the scripted queue. It still creates the endpoint + // store first, and monitorWorker's reclaim tick needs that to exist. + context_->construct(); + + g_script = AsyncEventScript{}; + worker_pool_ = std::make_unique(*context_); + } + + void TearDown() override { + worker_pool_.reset(); + context_.reset(); + FLAGS_minloglevel = previous_min_log_level_; + } + + int processContextEvents() { + return WorkerPoolTestPeer::processContextEvents(*worker_pool_); + } + + int previous_min_log_level_ = 0; + RdmaTransport *transport_ = nullptr; + std::unique_ptr context_; + std::unique_ptr worker_pool_; +}; + +// The regression: a burst arriving between two epoll wakeups must be consumed +// in full, not one event at a time. +TEST_F(AsyncEventDrainTest, DrainsEveryQueuedEvent) { + constexpr int kBurst = 5; + for (int i = 0; i < kBurst; ++i) + g_script.pending.push_back(kUnclaimedEvent); + + EXPECT_EQ(processContextEvents(), 0); + + EXPECT_TRUE(g_script.pending.empty()) + << "one epoll wakeup must drain the whole async event queue; " + << g_script.pending.size() << " event(s) were left stranded"; + // kBurst reads plus the trailing EAGAIN that ends the drain. + EXPECT_EQ(g_script.get_calls, kBurst + 1); + EXPECT_EQ(g_script.ack_calls, kBurst); +} + +// Both ack paths must drain, and neither may ack twice or not at all. +TEST_F(AsyncEventDrainTest, DrainsClaimedAndUnclaimedEventsAlike) { + g_script.pending.push_back(kUnclaimedEvent); + g_script.pending.push_back(kClaimedEvent); + g_script.pending.push_back(kUnclaimedEvent); + + EXPECT_EQ(processContextEvents(), 0); + + EXPECT_TRUE(g_script.pending.empty()); + EXPECT_EQ(g_script.get_calls, 4); + EXPECT_EQ(g_script.ack_calls, 3); +} + +// A wakeup with nothing pending is a drained fd, not a failure. +TEST_F(AsyncEventDrainTest, EmptyQueueIsNotAnError) { + EXPECT_EQ(processContextEvents(), 0); + + EXPECT_EQ(g_script.get_calls, 1); + EXPECT_EQ(g_script.ack_calls, 0); +} + +// A signal must not end the drain early, or the events behind it are stranded +// exactly as they were before the fix. +TEST_F(AsyncEventDrainTest, InterruptedReadIsRetried) { + g_script.pending.push_back(kUnclaimedEvent); + g_script.pending.push_back(kUnclaimedEvent); + g_script.pending_eintr = 1; + + EXPECT_EQ(processContextEvents(), 0); + + EXPECT_TRUE(g_script.pending.empty()); + // One EINTR, two reads, one trailing EAGAIN. + EXPECT_EQ(g_script.get_calls, 4); + EXPECT_EQ(g_script.ack_calls, 2); +} + +// A genuine read failure still aborts and propagates, so the loop cannot spin +// forever on a broken fd. +TEST_F(AsyncEventDrainTest, RealReadErrorStopsTheDrain) { + g_script.pending.push_back(kUnclaimedEvent); + g_script.drained_errno = EIO; + + EXPECT_EQ(processContextEvents(), ERR_CONTEXT); + + EXPECT_EQ(g_script.get_calls, 2); + EXPECT_EQ(g_script.ack_calls, 1); +} + +} // namespace diff --git a/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp b/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp index 319d57b242..66a9454a9e 100644 --- a/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp +++ b/mooncake-transfer-engine/tests/rdma_context_reprobe_test.cpp @@ -14,13 +14,19 @@ #include +#include #include #include #include #include +#ifdef __linux__ +#include +#endif + #include "common.h" #include "error.h" +#include "rdma_test_peers.h" #include "transfer_metadata.h" #include "transport/rdma_transport/rdma_context.h" #include "transport/rdma_transport/rdma_transport.h" @@ -39,39 +45,176 @@ using namespace mooncake; -namespace mooncake { +#ifdef __linux__ +namespace { + +struct FakeVerbsDevice { + bool enabled = false; + ibv_device device = {}; + ibv_device *device_list[2] = {&device, nullptr}; + ibv_context context = {}; + size_t alloc_pd_calls = 0; +}; + +FakeVerbsDevice fake_verbs; -class RdmaTransportTestPeer { +class FakeVerbsDeviceScope { public: - static void bindMetadata(RdmaTransport &transport, - std::shared_ptr metadata, - std::string local_server_name) { - transport.metadata_ = std::move(metadata); - transport.local_server_name_ = std::move(local_server_name); + explicit FakeVerbsDeviceScope(int num_comp_vectors) { + fake_verbs.enabled = true; + fake_verbs.context = {}; + fake_verbs.context.num_comp_vectors = num_comp_vectors; + fake_verbs.alloc_pd_calls = 0; } + + ~FakeVerbsDeviceScope() { fake_verbs.enabled = false; } + + size_t allocPdCalls() const { return fake_verbs.alloc_pd_calls; } }; -class RdmaContextTestPeer { - public: - static void seedAutoGidState(RdmaContext &context, ibv_context *verbs_ctx, - uint8_t port, uint16_t lid, const ibv_gid &gid, - int gid_index) { - context.context_ = verbs_ctx; - context.port_ = port; - context.lid_ = lid; - context.gid_ = gid; - context.gid_index_ = gid_index; - context.auto_gid_selection_enabled_ = true; +} // namespace + +#undef ibv_query_port + +// Interpose the libibverbs boundary for this test binary so construct() can +// exercise device validation without RDMA hardware. +extern "C" { + +ibv_device **ibv_get_device_list(int *num_devices) { + if (!fake_verbs.enabled) { + *num_devices = 0; + return nullptr; } + *num_devices = 1; + return fake_verbs.device_list; +} + +void ibv_free_device_list(ibv_device **) {} + +const char *ibv_get_device_name(ibv_device *device) { + if (fake_verbs.enabled && device == &fake_verbs.device) + return "nonexistent-device"; + return ""; +} + +ibv_context *ibv_open_device(ibv_device *device) { + if (fake_verbs.enabled && device == &fake_verbs.device) + return &fake_verbs.context; + return nullptr; +} + +int ibv_query_port(ibv_context *context, uint8_t, + _compat_ibv_port_attr *compat_port_attr) { + if (!fake_verbs.enabled || context != &fake_verbs.context) return EINVAL; + auto *port_attr = reinterpret_cast(compat_port_attr); + *port_attr = {}; + port_attr->state = IBV_PORT_ACTIVE; + port_attr->lid = 1; + port_attr->active_mtu = IBV_MTU_4096; + return 0; +} + +int ibv_query_device(ibv_context *context, ibv_device_attr *device_attr) { + if (!fake_verbs.enabled || context != &fake_verbs.context) return EINVAL; + *device_attr = {}; + device_attr->max_qp = std::numeric_limits::max(); + device_attr->max_cq = std::numeric_limits::max(); + device_attr->max_qp_wr = std::numeric_limits::max(); + device_attr->max_sge = std::numeric_limits::max(); + device_attr->max_cqe = std::numeric_limits::max(); + device_attr->max_mr_size = std::numeric_limits::max(); + return 0; +} + +int ibv_query_gid(ibv_context *context, uint8_t, int, ibv_gid *gid) { + if (!fake_verbs.enabled || context != &fake_verbs.context) return EINVAL; + *gid = {}; + gid->raw[15] = 1; + return 0; +} + +ibv_pd *ibv_alloc_pd(ibv_context *context) { + if (!fake_verbs.enabled || context != &fake_verbs.context) return nullptr; + ++fake_verbs.alloc_pd_calls; + return nullptr; +} + +int ibv_close_device(ibv_context *) { return 0; } - static void disableContextForTeardown(RdmaContext &context) { - context.context_ = nullptr; +} // extern "C" +#endif // __linux__ + +namespace { + +class RdmaContextConstructionTest : public ::testing::Test { + protected: + void SetUp() override { + // RdmaTransport teardown requires metadata initialized by install(). + // These tests only need the constructor reference, so match the + // existing uninstalled-transport test setup below. + transport_ = new RdmaTransport(); + MC_LSAN_IGNORE_OBJECT(transport_); + context_ = + std::make_unique(*transport_, "nonexistent-device"); } + + RdmaTransport *transport_ = nullptr; + std::unique_ptr context_; }; -} // namespace mooncake +TEST_F(RdmaContextConstructionTest, RejectsZeroCompletionQueuesBeforeSetup) { + EXPECT_EQ(context_->construct(/*num_cq_list=*/0, + /*num_comp_channels=*/1), + ERR_INVALID_ARGUMENT); + EXPECT_FALSE(RdmaContextTestPeer::hasEndpointStore(*context_)); +} -namespace { +TEST_F(RdmaContextConstructionTest, RejectsZeroCompletionChannelsBeforeSetup) { + EXPECT_EQ(context_->construct(/*num_cq_list=*/1, + /*num_comp_channels=*/0), + ERR_INVALID_ARGUMENT); + EXPECT_FALSE(RdmaContextTestPeer::hasEndpointStore(*context_)); +} + +TEST_F(RdmaContextConstructionTest, + RejectsDeviceWithoutCompletionVectorsBeforeAllocatingResources) { +#ifdef __linux__ + FakeVerbsDeviceScope fake_device(/*num_comp_vectors=*/0); + + EXPECT_EQ(context_->construct(/*num_cq_list=*/1, + /*num_comp_channels=*/1, + /*port=*/1, + /*gid_index=*/0), + ERR_CONTEXT); + EXPECT_EQ(fake_device.allocPdCalls(), 0); + RdmaContextTestPeer::disableContextForTeardown(*context_); +#else + GTEST_SKIP() << "Requires Linux libibverbs symbol interposition"; +#endif +} + +TEST(RdmaMemoryRegistrationPolicyTest, LocalOnlyBufferHasNoPublishedRkey) { + auto metadata = std::make_shared(P2PHANDSHAKE); + auto local_desc = std::make_shared(); + local_desc->name = "local-rdma-segment"; + local_desc->protocol = "rdma"; + ASSERT_EQ(metadata->addLocalSegment(LOCAL_SEGMENT_ID, "local-rdma-segment", + std::move(local_desc)), + 0); + + RdmaTransport transport; + RdmaTransportTestPeer::bindMetadata(transport, metadata, + "local-rdma-segment"); + std::array buffer{}; + ASSERT_EQ(transport.registerLocalMemory(buffer.data(), buffer.size(), + "cpu:0", false, false), + 0); + + auto desc = metadata->getSegmentDescByID(LOCAL_SEGMENT_ID); + ASSERT_NE(desc, nullptr); + ASSERT_EQ(desc->buffers.size(), 1); + EXPECT_TRUE(desc->buffers[0].rkey.empty()); +} ibv_gid makeGid(const std::array &bytes) { ibv_gid gid = {}; @@ -141,4 +284,70 @@ TEST_F(RdmaContextReprobeTest, EXPECT_EQ(after_desc.get(), before_desc.get()); } +// Regression tests for the HCA-index / context_list_ alignment invariant. +// An HCA's id is its position in getHcaList(), and disableDevice() keeps the +// surviving ids there. initializeRdmaResources() used to append a context only +// for devices that came up, compacting context_list_ so a later device_id +// named the wrong RNIC or ran past its end. + +constexpr const char *kAlignmentTopologyJson = + R"({"cpu:0": [["mlx5_0", "mlx5_1", "mlx5_2"], []]})"; + +TEST(RdmaHcaIndexAlignmentTest, ContextListKeepsOneSlotPerHcaWhenInitFails) { +#ifndef __linux__ + GTEST_SKIP() << "Requires Linux libibverbs symbol interposition"; +#else + // No FakeVerbsDeviceScope: the interposed ibv_get_device_list reports zero + // devices, so every construct() fails regardless of the host's hardware. + auto topology = std::make_shared(); + ASSERT_EQ(topology->parse(kAlignmentTopologyJson), 0); + const auto hca_list = topology->getHcaList(); + ASSERT_EQ(hca_list.size(), static_cast(3)); + + auto metadata = std::make_shared(P2PHANDSHAKE); + RdmaTransport transport; + RdmaTransportTestPeer::bindMetadata(transport, metadata, + "local-rdma-segment"); + RdmaTransportTestPeer::bindTopology(transport, topology); + + // Every device fails, so the topology ends up empty -- the slot layout + // must line up with getHcaList() anyway. + EXPECT_EQ(RdmaTransportTestPeer::initializeResources(transport), + ERR_DEVICE_NOT_FOUND); + + const auto &contexts = transport.getContextList(); + ASSERT_EQ(contexts.size(), hca_list.size()); + for (size_t i = 0; i < contexts.size(); ++i) { + ASSERT_NE(contexts[i], nullptr) << "slot " << i << " must be occupied"; + EXPECT_EQ(contexts[i]->deviceName(), hca_list[i]) + << "slot " << i << " names the wrong RNIC"; + EXPECT_FALSE(contexts[i]->active()) + << "placeholder for a failed RNIC must not report itself active"; + } +#endif // __linux__ +} + +// Characterizes the contract the fix above relies on. +TEST(RdmaHcaIndexAlignmentTest, DisabledDeviceKeepsRemainingHcaIndexesStable) { + Topology topology; + ASSERT_EQ(topology.parse(kAlignmentTopologyJson), 0); + const auto hca_list = topology.getHcaList(); + ASSERT_EQ(hca_list.size(), static_cast(3)); + ASSERT_NE(std::find(hca_list.begin(), hca_list.end(), "mlx5_1"), + hca_list.end()); + + ASSERT_EQ(topology.disableDevice("mlx5_1"), 0); + // getHcaList() must not shrink: ids are positions in it. + ASSERT_EQ(topology.getHcaList().size(), hca_list.size()); + + for (int retry_count = 0; retry_count < 16; ++retry_count) { + const int device_id = topology.selectDevice("cpu:0", retry_count); + ASSERT_GE(device_id, 0); + ASSERT_LT(static_cast(device_id), hca_list.size()) + << "device_id must index an hca_list-sized context array"; + EXPECT_NE(hca_list[device_id], "mlx5_1") + << "a disabled device must never be selected"; + } +} + } // namespace diff --git a/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp b/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp index 77f5819288..c98fc20b44 100644 --- a/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp +++ b/mooncake-transfer-engine/tests/rdma_endpoint_reestablish_test.cpp @@ -29,6 +29,8 @@ */ #include +#include +#include #include #include #include @@ -37,6 +39,7 @@ #include #include #include +#include #include #include @@ -50,10 +53,82 @@ #include "common.h" #include "transfer_engine.h" +#include "transport/rdma_transport/rdma_context.h" +#include "transport/rdma_transport/rdma_endpoint.h" +#include "transport/rdma_transport/rdma_transport.h" +#include "transport/rdma_transport/worker_pool.h" #include "transport/transport.h" using namespace mooncake; +namespace mooncake { + +class RdmaTransportTestPeer { + public: + static bool setContextActive(RdmaTransport* transport, + const std::string& device_name, bool active) { + if (!transport) return false; + for (auto& context : transport->context_list_) { + if (context && context->deviceName() == device_name) { + context->set_active(active); + return true; + } + } + return false; + } + + static bool contextActive(RdmaTransport* transport, + const std::string& device_name) { + if (!transport) return false; + for (auto& context : transport->context_list_) { + if (context && context->deviceName() == device_name) { + return context->active(); + } + } + return false; + } + + static bool injectContextEvent(RdmaTransport* transport, + const std::string& device_name, + ibv_event_type event_type); +}; + +class WorkerPoolTestPeer { + public: + static void processContextEvent(WorkerPool& worker_pool, + ibv_event_type event_type) { + worker_pool.processContextEventForTest(event_type); + } +}; + +class RdmaContextTestPeer { + public: + static bool injectContextEvent(RdmaContext* context, + ibv_event_type event_type) { + if (context && context->worker_pool_) { + WorkerPoolTestPeer::processContextEvent(*context->worker_pool_, + event_type); + return true; + } + return false; + } +}; + +bool RdmaTransportTestPeer::injectContextEvent(RdmaTransport* transport, + const std::string& device_name, + ibv_event_type event_type) { + if (!transport) return false; + for (auto& context : transport->context_list_) { + if (context && context->deviceName() == device_name) { + return RdmaContextTestPeer::injectContextEvent(context.get(), + event_type); + } + } + return false; +} + +} // namespace mooncake + namespace { constexpr size_t kRAMBufSize = 256ull << 24; @@ -82,8 +157,11 @@ struct RtrFaultInjectionState { std::mutex mu; bool synthetic_gid_swap_enabled = false; std::string synthetic_gid_device; + int synthetic_gid_a = 0; + int synthetic_gid_b = 1; bool fail_first_rtr_einval = false; std::string fail_rtr_device; + bool first_rtr_attempted = false; int injected_failures = 0; std::unordered_map> rtr_sgid_history; std::unordered_map> rtr_gid_history; @@ -105,8 +183,11 @@ void resetRtrFaultInjectionState() { std::lock_guard guard(g_rtr_fault_injection_state.mu); g_rtr_fault_injection_state.synthetic_gid_swap_enabled = false; g_rtr_fault_injection_state.synthetic_gid_device.clear(); + g_rtr_fault_injection_state.synthetic_gid_a = 0; + g_rtr_fault_injection_state.synthetic_gid_b = 1; g_rtr_fault_injection_state.fail_first_rtr_einval = false; g_rtr_fault_injection_state.fail_rtr_device.clear(); + g_rtr_fault_injection_state.first_rtr_attempted = false; g_rtr_fault_injection_state.injected_failures = 0; g_rtr_fault_injection_state.rtr_sgid_history.clear(); g_rtr_fault_injection_state.rtr_gid_history.clear(); @@ -116,8 +197,11 @@ void configureRtrFaultInjection(const std::string& device_name) { std::lock_guard guard(g_rtr_fault_injection_state.mu); g_rtr_fault_injection_state.synthetic_gid_swap_enabled = true; g_rtr_fault_injection_state.synthetic_gid_device = device_name; + g_rtr_fault_injection_state.synthetic_gid_a = 0; + g_rtr_fault_injection_state.synthetic_gid_b = 1; g_rtr_fault_injection_state.fail_first_rtr_einval = true; g_rtr_fault_injection_state.fail_rtr_device = device_name; + g_rtr_fault_injection_state.first_rtr_attempted = false; g_rtr_fault_injection_state.injected_failures = 0; g_rtr_fault_injection_state.rtr_sgid_history.clear(); g_rtr_fault_injection_state.rtr_gid_history.clear(); @@ -160,22 +244,40 @@ int maybeSwapSyntheticGidIndex(const std::string& device_name, int gid_index) { g_rtr_fault_injection_state.synthetic_gid_device != device_name) { return gid_index; } - if (gid_index == 0) return 1; - if (gid_index == 1) return 0; + if (gid_index == g_rtr_fault_injection_state.synthetic_gid_a) + return g_rtr_fault_injection_state.synthetic_gid_b; + if (gid_index == g_rtr_fault_injection_state.synthetic_gid_b) + return g_rtr_fault_injection_state.synthetic_gid_a; return gid_index; } -bool shouldInjectRtrEinval(const std::string& device_name, int sgid_index) { +int injectedRtrErrnoOrZero(const std::string& device_name, int sgid_index) { std::lock_guard guard(g_rtr_fault_injection_state.mu); if (!g_rtr_fault_injection_state.fail_first_rtr_einval || g_rtr_fault_injection_state.fail_rtr_device != device_name || - sgid_index != 0) { - return false; + sgid_index != g_rtr_fault_injection_state.synthetic_gid_a || + g_rtr_fault_injection_state.first_rtr_attempted) { + return 0; } g_rtr_fault_injection_state.fail_first_rtr_einval = false; g_rtr_fault_injection_state.synthetic_gid_swap_enabled = false; + g_rtr_fault_injection_state.first_rtr_attempted = true; ++g_rtr_fault_injection_state.injected_failures; - return true; + return EINVAL; +} + +void expectRtrEinvalRecoveredWithRetry(const std::string& device_name) { + EXPECT_EQ(getInjectedFailureCount(), 1); + auto sgid_history = getRtrSgidHistory(device_name); + auto gid_history = getRtrGidHistory(device_name); + ASSERT_EQ(sgid_history.size(), gid_history.size()); + ASSERT_GE(sgid_history.size(), 2u) + << "RTR/EINVAL was injected, but the endpoint did not retry RTR"; + EXPECT_FALSE(gid_history.front().empty()); + + EXPECT_TRUE(std::any_of( + gid_history.begin() + 1, gid_history.end(), + [&](const std::string& gid) { return gid != gid_history.front(); })); } std::string formatDeviceNames(const std::string& device_names) { @@ -204,6 +306,16 @@ std::string makeNicPriorityMatrix(const std::string& device_name) { formatted_devices + "],[]]}"; } +std::string makeNicPriorityMatrix(const std::string& preferred_devices, + const std::string& fallback_devices) { + auto formatted_preferred = formatDeviceNames(preferred_devices); + auto formatted_fallback = formatDeviceNames(fallback_devices); + return "{\"cpu:0\": [[" + formatted_preferred + "],[" + formatted_fallback + + "]], " + " \"cpu:1\": [[" + + formatted_preferred + "],[" + formatted_fallback + "]]}"; +} + void waitForTransfer(TransferEngine* engine, BatchID batch_id, const std::string& op_name) { bool completed = false; @@ -221,8 +333,46 @@ void waitForTransfer(TransferEngine* engine, BatchID batch_id, EXPECT_EQ(s, Status::OK()); } +bool waitForTransferWithTimeout(TransferEngine* engine, BatchID batch_id, + const std::string& op_name, + std::chrono::seconds timeout) { + auto deadline = std::chrono::steady_clock::now() + timeout; + TransferStatus status; + while (std::chrono::steady_clock::now() < deadline) { + Status s = engine->getTransferStatus(batch_id, 0, status); + if (s != Status::OK()) { + ADD_FAILURE() << op_name + << " getTransferStatus failed: " << s.ToString(); + return false; + } + if (status.s == TransferStatusEnum::COMPLETED) { + s = engine->freeBatchID(batch_id); + if (s != Status::OK()) { + ADD_FAILURE() + << op_name << " freeBatchID failed: " << s.ToString(); + return false; + } + return true; + } + if (status.s == TransferStatusEnum::FAILED) { + ADD_FAILURE() << op_name << " FAILED"; + engine->freeBatchID(batch_id); + return false; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + + ADD_FAILURE() << op_name << " timed out after " << timeout.count() + << " seconds"; + engine->freeBatchID(batch_id); + return false; +} + +struct RawNicPriorityMatrix {}; + struct TEContext { std::unique_ptr engine_{}; + RdmaTransport* rdma_transport_{}; uint8_t* local_addr_{}; bool segment_opened_{false}; SegmentHandle segment_handle_{}; @@ -230,17 +380,25 @@ struct TEContext { TEContext(const std::string& local_server_name, const std::string& metadata_server, const std::string& segment_id, - const std::string& device_name) { + const std::string& device_name) + : TEContext(local_server_name, metadata_server, segment_id, + makeNicPriorityMatrix(device_name), + RawNicPriorityMatrix{}) {} + + TEContext(const std::string& local_server_name, + const std::string& metadata_server, const std::string& segment_id, + const std::string& nic_priority_matrix, RawNicPriorityMatrix) { engine_ = std::make_unique(false); auto hostname_port = parseHostNameWithPort(local_server_name); engine_->init(metadata_server, local_server_name, hostname_port.first, hostname_port.second); - auto nic_priority_matrix = makeNicPriorityMatrix(device_name); void* args[2] = {const_cast(nic_priority_matrix.c_str()), nullptr}; Transport* xport = engine_->installTransport("rdma", args); LOG_ASSERT(xport); + rdma_transport_ = dynamic_cast(xport); + LOG_ASSERT(rdma_transport_); local_addr_ = static_cast(numa_alloc_onnode(kRAMBufSize, 0)); memset(local_addr_, 0, kDataLength); @@ -401,34 +559,341 @@ TEST_F(RDMAEndpointReestablishTest, EndpointReestablishReverseDevices) { TEST_F(RDMAEndpointReestablishTest, ActiveHandshakeRetriesAfterAutoGidReprobe) { configureRtrFaultInjection(initiator_device_name); runEndpointReestablishScenario(target_device_name, initiator_device_name); - - EXPECT_EQ(getInjectedFailureCount(), 1); - auto sgid_history = getRtrSgidHistory(initiator_device_name); - auto gid_history = getRtrGidHistory(initiator_device_name); - ASSERT_EQ(sgid_history.size(), gid_history.size()); - ASSERT_GE(sgid_history.size(), 2u); - EXPECT_EQ(sgid_history.front(), 0); - EXPECT_FALSE(gid_history.front().empty()); - EXPECT_TRUE(std::any_of( - gid_history.begin() + 1, gid_history.end(), - [&](const std::string& gid) { return gid != gid_history.front(); })); + expectRtrEinvalRecoveredWithRetry(initiator_device_name); } TEST_F(RDMAEndpointReestablishTest, PassiveHandshakeRetriesAfterAutoGidReprobe) { configureRtrFaultInjection(target_device_name); runEndpointReestablishScenario(target_device_name, initiator_device_name); + expectRtrEinvalRecoveredWithRetry(target_device_name); +} - EXPECT_EQ(getInjectedFailureCount(), 1); - auto sgid_history = getRtrSgidHistory(target_device_name); - auto gid_history = getRtrGidHistory(target_device_name); - ASSERT_EQ(sgid_history.size(), gid_history.size()); - ASSERT_GE(sgid_history.size(), 2u); - EXPECT_EQ(sgid_history.front(), 0); - EXPECT_FALSE(gid_history.front().empty()); - EXPECT_TRUE(std::any_of( - gid_history.begin() + 1, gid_history.end(), - [&](const std::string& gid) { return gid != gid_history.front(); })); +TEST_F(RDMAEndpointReestablishTest, SenderSingleRnicDownUsesFallbackLocalRnic) { + if (target_device_name == initiator_device_name) { + GTEST_SKIP() << "Need two distinct RDMA devices for sender failover"; + } + + const std::string both_devices = + target_device_name + "," + initiator_device_name; + const std::string target_matrix = makeNicPriorityMatrix(both_devices); + LOG(INFO) << "========== Setting up dual-RNIC Target =========="; + TEContext target_ctx(target_server_name, metadata_server, "", target_matrix, + RawNicPriorityMatrix{}); + const std::string target_segment_name = usesP2PHandshake(metadata_server) + ? target_ctx.localSegmentName() + : target_server_name; + + // Put the soon-to-be-down sender RNIC in the preferred tier and the + // surviving sender RNIC in the fallback tier. This proves TE can skip the + // inactive preferred local context and still complete the transfer through + // another local RNIC. + const std::string sender_matrix = + makeNicPriorityMatrix(initiator_device_name, target_device_name); + TEContext init_ctx(initiator_server_name, metadata_server, + target_segment_name, sender_matrix, + RawNicPriorityMatrix{}); + ASSERT_TRUE(RdmaTransportTestPeer::setContextActive( + init_ctx.rdma_transport_, initiator_device_name, false)); + ASSERT_FALSE(RdmaTransportTestPeer::contextActive(init_ctx.rdma_transport_, + initiator_device_name)); + ASSERT_TRUE(RdmaTransportTestPeer::contextActive(init_ctx.rdma_transport_, + target_device_name)); + + for (size_t i = 0; i < kDataLength; ++i) { + init_ctx.local_addr_[i] = static_cast((i * 17) % 251); + } + + auto batch_id = init_ctx.engine_->allocateBatchID(1); + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = init_ctx.local_addr_; + entry.target_id = init_ctx.segment_handle_; + entry.target_offset = init_ctx.remote_base_; + + Status s = init_ctx.engine_->submitTransfer(batch_id, {entry}); + ASSERT_EQ(s, Status::OK()); + waitForTransfer(init_ctx.engine_.get(), batch_id, + "WRITE with one sender RNIC down"); + + for (size_t i = 0; i < kDataLength; ++i) { + ASSERT_EQ(target_ctx.local_addr_[i], + static_cast((i * 17) % 251)) + << "Data mismatch at offset " << i; + } +} + +TEST_F(RDMAEndpointReestablishTest, + InjectedSenderRnicDownKeepsFallbackTransfersSmoothForTwentySeconds) { + auto devices = getAvailableRdmaDevices(); + for (const auto* device : {"mlx5_1", "mlx5_2", "mlx5_3", "mlx5_4"}) { + if (std::find(devices.begin(), devices.end(), device) == + devices.end()) { + GTEST_SKIP() << "Need " << device << " for this 4-RNIC test"; + } + } + + const std::string target_matrix = makeNicPriorityMatrix("mlx5_3,mlx5_4"); + TEContext target_ctx(target_server_name, metadata_server, "", target_matrix, + RawNicPriorityMatrix{}); + const std::string target_segment_name = usesP2PHandshake(metadata_server) + ? target_ctx.localSegmentName() + : target_server_name; + + const std::string sender_matrix = makeNicPriorityMatrix("mlx5_2", "mlx5_1"); + TEContext init_ctx(initiator_server_name, metadata_server, + target_segment_name, sender_matrix, + RawNicPriorityMatrix{}); + + std::atomic injected_down{false}; + std::atomic injected_up{false}; + std::atomic inject_down_ok{false}; + std::atomic inject_up_ok{false}; + std::thread injector([&] { + std::this_thread::sleep_for(std::chrono::seconds(5)); + inject_down_ok.store( + RdmaTransportTestPeer::injectContextEvent( + init_ctx.rdma_transport_, "mlx5_2", IBV_EVENT_PORT_ERR), + std::memory_order_release); + injected_down.store(true, std::memory_order_release); + + std::this_thread::sleep_for(std::chrono::seconds(7)); + inject_up_ok.store( + RdmaTransportTestPeer::injectContextEvent( + init_ctx.rdma_transport_, "mlx5_2", IBV_EVENT_PORT_ACTIVE), + std::memory_order_release); + injected_up.store(true, std::memory_order_release); + }); + + constexpr size_t kFaultLoopLength = 4ull << 20; + auto start = std::chrono::steady_clock::now(); + auto deadline = start + std::chrono::seconds(20); + size_t completed_transfers = 0; + uint8_t last_pattern = 0; + bool transfer_ok = true; + + while (std::chrono::steady_clock::now() < deadline) { + last_pattern = static_cast((completed_transfers * 37) % 251); + memset(init_ctx.local_addr_, last_pattern, kFaultLoopLength); + + auto batch_id = init_ctx.engine_->allocateBatchID(1); + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kFaultLoopLength; + entry.source = init_ctx.local_addr_; + entry.target_id = init_ctx.segment_handle_; + entry.target_offset = init_ctx.remote_base_; + + Status s = init_ctx.engine_->submitTransfer(batch_id, {entry}); + if (s != Status::OK()) { + ADD_FAILURE() << "submitTransfer failed during injected sender " + "RNIC down/fallback: " + << s.ToString(); + transfer_ok = false; + break; + } + if (!waitForTransferWithTimeout( + init_ctx.engine_.get(), batch_id, + "20s WRITE during injected sender RNIC down/fallback", + std::chrono::seconds(5))) { + transfer_ok = false; + break; + } + ++completed_transfers; + } + + injector.join(); + ASSERT_TRUE(injected_down.load(std::memory_order_acquire)); + ASSERT_TRUE(injected_up.load(std::memory_order_acquire)); + ASSERT_TRUE(inject_down_ok.load(std::memory_order_acquire)); + ASSERT_TRUE(inject_up_ok.load(std::memory_order_acquire)); + ASSERT_TRUE(transfer_ok); + ASSERT_GT(completed_transfers, 0u); + + for (size_t i = 0; i < kFaultLoopLength; ++i) { + ASSERT_EQ(target_ctx.local_addr_[i], last_pattern) + << "Data mismatch at offset " << i; + } + LOG(INFO) << "Completed " << completed_transfers + << " transfers during injected mlx5_2 down/fallback"; +} + +TEST_F(RDMAEndpointReestablishTest, + InjectedReceiverRnicDownKeepsFallbackTransfersSmoothForTwentySeconds) { + auto devices = getAvailableRdmaDevices(); + for (const auto* device : {"mlx5_1", "mlx5_2", "mlx5_3", "mlx5_4"}) { + if (std::find(devices.begin(), devices.end(), device) == + devices.end()) { + GTEST_SKIP() << "Need " << device << " for this 4-RNIC test"; + } + } + + const std::string target_matrix = makeNicPriorityMatrix("mlx5_3", "mlx5_4"); + TEContext target_ctx(target_server_name, metadata_server, "", target_matrix, + RawNicPriorityMatrix{}); + const std::string target_segment_name = usesP2PHandshake(metadata_server) + ? target_ctx.localSegmentName() + : target_server_name; + + const std::string sender_matrix = makeNicPriorityMatrix("mlx5_1,mlx5_2"); + TEContext init_ctx(initiator_server_name, metadata_server, + target_segment_name, sender_matrix, + RawNicPriorityMatrix{}); + + std::atomic injected_down{false}; + std::atomic injected_up{false}; + std::atomic inject_down_ok{false}; + std::atomic inject_up_ok{false}; + std::thread injector([&] { + std::this_thread::sleep_for(std::chrono::seconds(5)); + inject_down_ok.store( + RdmaTransportTestPeer::injectContextEvent( + target_ctx.rdma_transport_, "mlx5_3", IBV_EVENT_PORT_ERR), + std::memory_order_release); + injected_down.store(true, std::memory_order_release); + + std::this_thread::sleep_for(std::chrono::seconds(7)); + inject_up_ok.store( + RdmaTransportTestPeer::injectContextEvent( + target_ctx.rdma_transport_, "mlx5_3", IBV_EVENT_PORT_ACTIVE), + std::memory_order_release); + injected_up.store(true, std::memory_order_release); + }); + + constexpr size_t kFaultLoopLength = 4ull << 20; + auto start = std::chrono::steady_clock::now(); + auto deadline = start + std::chrono::seconds(20); + size_t completed_transfers = 0; + uint8_t last_pattern = 0; + bool transfer_ok = true; + + while (std::chrono::steady_clock::now() < deadline) { + last_pattern = static_cast((completed_transfers * 41) % 251); + memset(init_ctx.local_addr_, last_pattern, kFaultLoopLength); + + auto batch_id = init_ctx.engine_->allocateBatchID(1); + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kFaultLoopLength; + entry.source = init_ctx.local_addr_; + entry.target_id = init_ctx.segment_handle_; + entry.target_offset = init_ctx.remote_base_; + + Status s = init_ctx.engine_->submitTransfer(batch_id, {entry}); + if (s != Status::OK()) { + ADD_FAILURE() << "submitTransfer failed during injected receiver " + "RNIC down/fallback: " + << s.ToString(); + transfer_ok = false; + break; + } + if (!waitForTransferWithTimeout( + init_ctx.engine_.get(), batch_id, + "20s WRITE during injected receiver RNIC down/fallback", + std::chrono::seconds(5))) { + transfer_ok = false; + break; + } + ++completed_transfers; + } + + injector.join(); + ASSERT_TRUE(injected_down.load(std::memory_order_acquire)); + ASSERT_TRUE(injected_up.load(std::memory_order_acquire)); + ASSERT_TRUE(inject_down_ok.load(std::memory_order_acquire)); + ASSERT_TRUE(inject_up_ok.load(std::memory_order_acquire)); + ASSERT_TRUE(transfer_ok); + ASSERT_GT(completed_transfers, 0u); + + for (size_t i = 0; i < kFaultLoopLength; ++i) { + ASSERT_EQ(target_ctx.local_addr_[i], last_pattern) + << "Data mismatch at offset " << i; + } + LOG(INFO) << "Completed " << completed_transfers + << " transfers during injected receiver mlx5_3 down/fallback"; +} + +} // namespace + +#if defined(__has_feature) +#define MC_HAS_FEATURE(x) __has_feature(x) +#else +#define MC_HAS_FEATURE(x) 0 +#endif +#if defined(__SANITIZE_ADDRESS__) || MC_HAS_FEATURE(address_sanitizer) +#include +#define MC_LSAN_IGNORE_OBJECT(p) __lsan_ignore_object(p) +#else +#define MC_LSAN_IGNORE_OBJECT(p) ((void)(p)) +#endif + +namespace mooncake { +class RdmaEndPointTestPeer { + public: + static void addDummyQp(RdmaEndPoint &endpoint, ibv_qp *qp) { + endpoint.qp_list_.push_back(qp); + } + + static void clearDummyQp(RdmaEndPoint &endpoint) { + endpoint.qp_list_.clear(); + } + + static int doSetupConnection(RdmaEndPoint &endpoint, int qp_index, + const ibv_gid &peer_gid, uint16_t peer_lid, + uint32_t peer_qp_num, int local_gid_index, + std::string *reply_msg, + int &out_stage, int &out_sys_errno) { + RdmaEndPoint::SetupConnectionFailureInfo failure_info = {}; + int rc = endpoint.doSetupConnection(qp_index, peer_gid, peer_lid, + peer_qp_num, local_gid_index, + reply_msg, &failure_info); + out_stage = static_cast(failure_info.stage); + out_sys_errno = failure_info.sys_errno; + return rc; + } + + static int getResetStageValue() { + return static_cast(RdmaEndPoint::SetupConnectionFailureStage::kReset); + } +}; +} + +extern std::atomic g_inject_modify_qp_error; + +namespace { + +TEST(RDMAEndpointSetupErrorTest, VerifyVerbsErrorCorrectlyCaptured) { + auto* transport = new RdmaTransport(); + // Ignore memory leak of transport during destruction since it's a test + MC_LSAN_IGNORE_OBJECT(transport); + auto context = std::make_unique(*transport, "unused"); + auto endpoint = std::make_unique(*context); + + // Add a dummy non-null QP pointer to trigger doSetupConnection + ibv_qp* dummy_qp = reinterpret_cast(0xdeadbeef); + mooncake::RdmaEndPointTestPeer::addDummyQp(*endpoint, dummy_qp); + + // Inject EINVAL error for ibv_modify_qp + g_inject_modify_qp_error.store(EINVAL); + + ibv_gid dummy_gid = {}; + std::string reply_msg; + int stage = 0; + int sys_errno = 0; + int rc = mooncake::RdmaEndPointTestPeer::doSetupConnection( + *endpoint, 0, dummy_gid, 0, 0, 0, &reply_msg, stage, sys_errno); + + // Clear injection immediately + g_inject_modify_qp_error.store(0); + + // Clear dummy qp so destructor doesn't try to destroy 0xdeadbeef and crash + mooncake::RdmaEndPointTestPeer::clearDummyQp(*endpoint); + + EXPECT_EQ(rc, ERR_ENDPOINT); + EXPECT_EQ(stage, mooncake::RdmaEndPointTestPeer::getResetStageValue()); + EXPECT_EQ(sys_errno, EINVAL); + EXPECT_TRUE(reply_msg.find("EINVAL") != std::string::npos || reply_msg.find("Invalid argument") != std::string::npos || reply_msg.find("22") != std::string::npos); } } // namespace @@ -461,8 +926,14 @@ extern "C" int __wrap_ibv_query_gid(ibv_context* context, uint8_t port_num, extern "C" int __real_ibv_modify_qp(ibv_qp* qp, ibv_qp_attr* attr, int attr_mask); +std::atomic g_inject_modify_qp_error{0}; + extern "C" int __wrap_ibv_modify_qp(ibv_qp* qp, ibv_qp_attr* attr, int attr_mask) { + int error_to_inject = g_inject_modify_qp_error.load(); + if (error_to_inject != 0) { + return error_to_inject; + } if (qp != nullptr && attr != nullptr && attr->qp_state == IBV_QPS_RTR && (attr_mask & IBV_QP_AV)) { const std::string device_name = @@ -475,9 +946,9 @@ extern "C" int __wrap_ibv_modify_qp(ibv_qp* qp, ibv_qp_attr* attr, gid_string = formatGidBytes(actual_gid.raw); } recordRtrAttempt(device_name, sgid_index, gid_string); - if (shouldInjectRtrEinval(device_name, sgid_index)) { - errno = EINVAL; - return -1; + int injected_errno = injectedRtrErrnoOrZero(device_name, sgid_index); + if (injected_errno != 0) { + return injected_errno; } } return __real_ibv_modify_qp(qp, attr, attr_mask); diff --git a/mooncake-transfer-engine/tests/rdma_endpoint_state_test.cpp b/mooncake-transfer-engine/tests/rdma_endpoint_state_test.cpp new file mode 100644 index 0000000000..5a2f4fe818 --- /dev/null +++ b/mooncake-transfer-engine/tests/rdma_endpoint_state_test.cpp @@ -0,0 +1,148 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include +#include +#include + +#include "error.h" +#include "transport/rdma_transport/rdma_context.h" +#include "transport/rdma_transport/rdma_endpoint.h" +#include "transport/rdma_transport/rdma_transport.h" + +#if defined(__has_feature) +#define MC_HAS_FEATURE(x) __has_feature(x) +#else +#define MC_HAS_FEATURE(x) 0 +#endif +#if defined(__SANITIZE_ADDRESS__) || MC_HAS_FEATURE(address_sanitizer) +#include +#define MC_LSAN_IGNORE_OBJECT(p) __lsan_ignore_object(p) +#else +#define MC_LSAN_IGNORE_OBJECT(p) ((void)(p)) +#endif + +using namespace mooncake; + +namespace mooncake { + +class RdmaEndPointTestPeer { + public: + static void setStatus(RdmaEndPoint &endpoint, RdmaEndPoint::Status status) { + endpoint.status_.store(status, std::memory_order_relaxed); + } + + static void setReadyWaitStartTs(RdmaEndPoint &endpoint, uint64_t start_ts) { + endpoint.ready_wait_start_ts_.store(start_ts, + std::memory_order_relaxed); + } + + static void setPeerQpNums(RdmaEndPoint &endpoint, + std::vector peer_qp_nums) { + endpoint.peer_qp_num_list_ = std::move(peer_qp_nums); + } +}; + +} // namespace mooncake + +namespace { + +class RdmaEndPointStateTest : public ::testing::Test { + protected: + void SetUp() override { + transport_ = new RdmaTransport(); + // Intentional leak: ~RdmaTransport dereferences metadata_, which is + // null until install(). We only need it as RdmaContext's owner. + MC_LSAN_IGNORE_OBJECT(transport_); + context_ = std::make_unique(*transport_, "unused"); + endpoint_ = std::make_unique(*context_); + } + + RdmaTransport *transport_ = nullptr; + std::unique_ptr context_; + std::unique_ptr endpoint_; +}; + +TEST_F(RdmaEndPointStateTest, WaitingReadyAckIsConnectedButNotReadyToSend) { + RdmaEndPointTestPeer::setStatus(*endpoint_, + RdmaEndPoint::CONNECTED_WAIT_READY_ACK); + + EXPECT_TRUE(endpoint_->connected()); + EXPECT_FALSE(endpoint_->readyToSend()); +} + +TEST_F(RdmaEndPointStateTest, ConnectedIsReadyToSend) { + RdmaEndPointTestPeer::setStatus(*endpoint_, RdmaEndPoint::CONNECTED); + + EXPECT_TRUE(endpoint_->connected()); + EXPECT_TRUE(endpoint_->readyToSend()); +} + +TEST_F(RdmaEndPointStateTest, ReadyAckTimeoutOnlyAppliesToWaitingState) { + RdmaEndPointTestPeer::setReadyWaitStartTs(*endpoint_, 1); + + RdmaEndPointTestPeer::setStatus(*endpoint_, + RdmaEndPoint::CONNECTED_WAIT_READY_ACK); + EXPECT_TRUE(endpoint_->readyAckTimedOut()); + + RdmaEndPointTestPeer::setStatus(*endpoint_, RdmaEndPoint::CONNECTED); + EXPECT_FALSE(endpoint_->readyAckTimedOut()); +} + +TEST_F(RdmaEndPointStateTest, ReadyAckWithSamePeerQpMarksEndpointReady) { + endpoint_->setPeerNicPath("peer@nic"); + RdmaEndPointTestPeer::setPeerQpNums(*endpoint_, {11, 22}); + RdmaEndPointTestPeer::setReadyWaitStartTs(*endpoint_, 1); + RdmaEndPointTestPeer::setStatus(*endpoint_, + RdmaEndPoint::CONNECTED_WAIT_READY_ACK); + + RdmaEndPoint::HandShakeDesc peer_desc; + peer_desc.ready_ack = true; + peer_desc.ready_ack_supported = true; + peer_desc.qp_num = {11, 22}; + RdmaEndPoint::HandShakeDesc local_desc; + + EXPECT_EQ(0, endpoint_->setupConnectionsByPassive(peer_desc, local_desc)); + EXPECT_TRUE(local_desc.reply_msg.empty()); + EXPECT_TRUE(endpoint_->connected()); + EXPECT_TRUE(endpoint_->readyToSend()); + EXPECT_FALSE(endpoint_->readyAckTimedOut()); +} + +TEST_F(RdmaEndPointStateTest, StaleReadyAckWithDifferentPeerQpDoesNotReset) { + endpoint_->setPeerNicPath("peer@nic"); + RdmaEndPointTestPeer::setPeerQpNums(*endpoint_, {11, 22}); + RdmaEndPointTestPeer::setReadyWaitStartTs(*endpoint_, 1); + RdmaEndPointTestPeer::setStatus(*endpoint_, + RdmaEndPoint::CONNECTED_WAIT_READY_ACK); + + RdmaEndPoint::HandShakeDesc peer_desc; + peer_desc.ready_ack = true; + peer_desc.ready_ack_supported = true; + peer_desc.qp_num = {33, 44}; + RdmaEndPoint::HandShakeDesc local_desc; + + EXPECT_EQ(ERR_REJECT_HANDSHAKE, + endpoint_->setupConnectionsByPassive(peer_desc, local_desc)); + EXPECT_FALSE(local_desc.reply_msg.empty()); + EXPECT_TRUE(endpoint_->connected()); + EXPECT_FALSE(endpoint_->readyToSend()); + EXPECT_TRUE(endpoint_->readyAckTimedOut()); +} + +} // namespace diff --git a/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp b/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp index 740222bcd4..0cbb75b819 100644 --- a/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp +++ b/mooncake-transfer-engine/tests/rdma_gid_probe_test.cpp @@ -14,6 +14,8 @@ #include +#include +#include #include #include "transport/rdma_transport/rdma_gid_probe.h" @@ -399,4 +401,161 @@ TEST(RdmaGidProbeTest, RetryActionRequiresObservedOrReprobedChange) { AutoGidRetryAction::kRetryWithObservedChange); } +// Regression tests for #2729: a routable-fabric private-range IPv4 GID must +// outrank a link-local IPv6 GID instead of tying with it in the degraded +// tier (where the lowest-index tie-break used to pick fe80::). + +// Exactly the GID table from the #2729 report: fe80 v1/v2 at indices 0/1, +// 10.14.x-mapped v1/v2 at indices 2/3, all on the same netdev. RoCE v1 +// entries are filtered by type; index 3 (private v4, RoCE v2) must win over +// index 1 (link-local, RoCE v2). +TEST(RdmaGidProbeTest, PrefersPrivateRangeV4OverLinkLocal) { + std::vector candidates = { + makeCandidate(/*gid_index=*/0, IBV_GID_TYPE_ROCE_V1, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + makeCandidate(/*gid_index=*/2, IBV_GID_TYPE_ROCE_V1, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + makeCandidate(/*gid_index=*/3, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 3); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkPrivateV4); +} + +// A genuinely routable (non-private) v4 GID still outranks a private-range +// one: the new tier sits strictly between routable and degraded. +TEST(RdmaGidProbeTest, RoutableV4StillOutranksPrivateRangeV4) { + std::vector candidates = { + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + makeCandidate(/*gid_index=*/5, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 5); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkRoutable); +} + +// An overlay-NAMED interface (docker0/cni/...) stays demoted below +// private-range v4 even when its address is v4-mapped: the interface-name +// heuristic remains the strongest demotion signal. +TEST(RdmaGidProbeTest, OverlayInterfaceStaysDemotedBelowPrivateRangeV4) { + std::vector candidates = { + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/true, + /*is_overlay_ipv4=*/true), + makeCandidate(/*gid_index=*/4, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 4); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkPrivateV4); +} + +// Cross-tier order is preserved from before the split: a link-local GID +// with a netdev still outranks a private-range v4 GID without one, exactly +// as network-degraded outranked no-network-degraded before. +TEST(RdmaGidProbeTest, NetworkLinkLocalStillOutranksNoNetworkPrivateV4) { + std::vector candidates = { + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + makeCandidate(/*gid_index=*/3, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/false, + /*is_ipv4_mapped=*/true, + /*is_link_local_ipv6=*/false, + /*is_overlay_network=*/false, + /*is_overlay_ipv4=*/true), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 1); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkDegraded); +} + +// Completeness pin for the tier tables: every class has a distinct +// priority consistent with its enum order and a unique display name, so a +// future class insertion cannot silently create a tie or reuse a label. +TEST(RdmaGidProbeTest, ClassPriorityAndNameTablesAreComplete) { + const AutoGidCandidateClass all[] = { + AutoGidCandidateClass::kNetworkRoutable, + AutoGidCandidateClass::kNoNetworkRoutable, + AutoGidCandidateClass::kNetworkPrivateV4, + AutoGidCandidateClass::kNetworkDegraded, + AutoGidCandidateClass::kNoNetworkPrivateV4, + AutoGidCandidateClass::kNoNetworkDegraded, + AutoGidCandidateClass::kFallbackNonzero, + }; + int expected_priority = 0; + std::set names; + for (auto cls : all) { + EXPECT_EQ(autoGidCandidateClassPriority(cls), expected_priority++); + std::string name = autoGidCandidateClassToString(cls); + EXPECT_NE(name, "unknown"); + EXPECT_TRUE(names.insert(name).second) + << "duplicate class name: " << name; + } +} + +// When only link-local candidates exist, behavior is unchanged: lowest +// index wins within the tier. +TEST(RdmaGidProbeTest, LinkLocalOnlyKeepsLowestIndexTieBreak) { + std::vector candidates = { + makeCandidate(/*gid_index=*/1, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + makeCandidate(/*gid_index=*/2, IBV_GID_TYPE_ROCE_V2, + /*has_network_device=*/true, + /*is_ipv4_mapped=*/false, + /*is_link_local_ipv6=*/true), + }; + + auto selection = selectBestAutoGidCandidate(candidates); + ASSERT_TRUE(selection.has_value()); + EXPECT_EQ(selection->gid_index, 1); + EXPECT_EQ(selection->candidate_class, + AutoGidCandidateClass::kNetworkDegraded); +} + } // namespace diff --git a/mooncake-transfer-engine/tests/rdma_large_mr_test.cpp b/mooncake-transfer-engine/tests/rdma_large_mr_test.cpp new file mode 100644 index 0000000000..fdf413703b --- /dev/null +++ b/mooncake-transfer-engine/tests/rdma_large_mr_test.cpp @@ -0,0 +1,308 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Regression test for issue #2017: RdmaTransport::registerLocalMemory must NOT +// silently truncate a buffer larger than the device max_mr_size. +// +// Bug: registerMemoryRegionInternal shrank `length` to max_mr_size and +// registered a single MR of that size, but registerLocalMemory still published +// BufferDesc.length = full length with that one (truncated) rkey. A remote RDMA +// op whose target address fell past max_mr_size then failed with +// IBV_WC_REM_ACCESS_ERR (e.g. ionic CQE error 10). It is ADDRESS-driven: ops to +// low addresses succeed, ops past the boundary fail -> PD KV transfer ran for a +// while then a worker died mid-run (seen on MI355x/ionic, per-layer KV ~3.25 +// GiB > the ionic 2 GiB max_mr_size). +// +// Fix: split buffers > max_mr_size into <= max_mr_size chunks, register each as +// its own MR, and publish one BufferDesc per chunk. +// +// This test forces the condition at a small, HW-independent size by setting +// MC_MAX_MR_SIZE, then does a LOOPBACK RDMA WRITE whose target lands PAST that +// boundary. Pre-fix: the transfer FAILS (REM_ACCESS_ERR). Post-fix: it +// COMPLETES and the bytes match. Runs on any RDMA device (incl. rdma_rxe / +// loopback). + +#include +#include +#include + +#include +#include +#include +#include + +#if defined(USE_CUDA) || defined(USE_HIP) +#include "cuda_alike.h" +#include "environ.h" +#endif +#include "config.h" +#include "transfer_engine.h" +#include "transport/transport.h" + +using namespace mooncake; + +namespace mooncake { + +// Small max_mr_size so the >max_mr_size path is exercised without needing a +// multi-GB allocation. Must be set before TransferEngine init (config reads +// env). +static constexpr size_t kMaxMrSize = 64ull << 20; // 64 MiB +static constexpr size_t kBufferSize = 256ull << 20; // 256 MiB -> 4 chunks + +class RDMALargeMrTest : public ::testing::Test { + public: + void *addr = nullptr; + std::unique_ptr engine; + + protected: + void SetUp() override { + engine = std::make_unique(true); + const char *env_meta = std::getenv("MC_METADATA_SERVER"); + std::string metadata_server = env_meta ? env_meta : "P2PHANDSHAKE"; + ASSERT_EQ(engine->init(metadata_server, "test_node_large_mr"), 0); + ASSERT_EQ(globalConfig().max_mr_size, kMaxMrSize) + << "Launch this test process with MC_MAX_MR_SIZE=" << kMaxMrSize; + addr = numa_alloc_onnode(kBufferSize, 0); + ASSERT_NE(addr, nullptr); + // Register a buffer LARGER than max_mr_size. Pre-fix this silently + // truncates the MR to kMaxMrSize; post-fix it splits into 4 chunks. + int rc = engine->registerLocalMemory(addr, kBufferSize, "cpu:0"); + ASSERT_EQ(rc, 0); + } + + void TearDown() override { + if (engine && addr) engine->unregisterLocalMemory(addr); + if (addr) numa_free(addr, kBufferSize); + } +}; + +// Loopback RDMA WRITE whose TARGET lands past max_mr_size. The source stays in +// the first MR; only the destination address exercises the truncation boundary. +TEST_F(RDMALargeMrTest, WritePastMaxMrSizeBoundary) { + const size_t kDataLength = 1ull << 20; // 1 MiB + // Target offset is well past kMaxMrSize (in the 4th chunk). Pre-fix: the + // single 64 MiB MR does not cover this address -> IBV_WC_REM_ACCESS_ERR. + const size_t kTargetOffset = kBufferSize - kDataLength; // ~255 MiB + ASSERT_GT(kTargetOffset, kMaxMrSize); + + for (size_t i = 0; i < kDataLength; ++i) + *((char *)addr + i) = (char)('a' + (lrand48() % 26)); + + auto batch_id = engine->allocateBatchID(1); + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = (uint8_t *)addr; // src in chunk 0 + entry.target_id = LOCAL_SEGMENT_ID; + entry.target_offset = (uint64_t)addr + kTargetOffset; // dst past 64 MiB + + Status s = engine->submitTransfer(batch_id, {entry}); + ASSERT_TRUE(s.ok()); + + TransferStatus status; + bool completed = false; + while (!completed) { + Status gs = engine->getTransferStatus(batch_id, 0, status); + ASSERT_EQ(gs, Status::OK()); + if (status.s == TransferStatusEnum::COMPLETED) + completed = true; + else if (status.s == TransferStatusEnum::FAILED) + break; + } + ASSERT_EQ(engine->freeBatchID(batch_id), Status::OK()); + + // The regression assertion: pre-fix this is FAILED (remote access error); + // post-fix it COMPLETES and the bytes match. + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED) + << "RDMA WRITE to an address past max_mr_size failed -- buffer MR was " + "truncated and not auto-chunked (issue #2017)."; + ASSERT_EQ(0, memcmp(addr, (char *)addr + kTargetOffset, kDataLength)); +} + +// Loopback RDMA WRITE whose TARGET range STRADDLES a chunk boundary (the seam +// between chunk 0 and chunk 1 at kMaxMrSize). The single logical transfer spans +// two MRs, so it must be split across the two chunks' rkeys -- otherwise the +// same IBV_WC_REM_ACCESS_ERR class re-surfaces at chunk seams (the concern this +// test guards, complementing WritePastMaxMrSizeBoundary which stays within one +// chunk). +TEST_F(RDMALargeMrTest, WriteStraddlesChunkBoundary) { + const size_t kDataLength = 1ull << 20; // 1 MiB + // Center the target on the first chunk boundary: half in chunk 0, half in + // chunk 1. + const size_t kTargetOffset = kMaxMrSize - kDataLength / 2 + 1; + ASSERT_LT(kTargetOffset, kMaxMrSize); // starts in chunk 0 + ASSERT_GT(kTargetOffset + kDataLength, kMaxMrSize); // ends in chunk 1 + ASSERT_NE(kTargetOffset % globalConfig().slice_size, 0u); + + // Distinct source bytes, and poison the destination so a dropped/partial + // write past the seam is caught by the memcmp below. + for (size_t i = 0; i < kDataLength; ++i) + *((char *)addr + i) = (char)('A' + (lrand48() % 26)); + memset((char *)addr + kTargetOffset, 0, kDataLength); + + auto batch_id = engine->allocateBatchID(1); + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = (uint8_t *)addr; // src in chunk 0 + entry.target_id = LOCAL_SEGMENT_ID; + entry.target_offset = (uint64_t)addr + kTargetOffset; // straddles the seam + + Status s = engine->submitTransfer(batch_id, {entry}); + ASSERT_TRUE(s.ok()); + + TransferStatus status; + bool completed = false; + while (!completed) { + Status gs = engine->getTransferStatus(batch_id, 0, status); + ASSERT_EQ(gs, Status::OK()); + if (status.s == TransferStatusEnum::COMPLETED) + completed = true; + else if (status.s == TransferStatusEnum::FAILED) + break; + } + ASSERT_EQ(engine->freeBatchID(batch_id), Status::OK()); + + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED) + << "RDMA WRITE straddling a chunk boundary failed -- a cross-chunk " + "transfer was not split across the two chunks' MRs (issue #2017)."; + ASSERT_EQ(0, memcmp(addr, (char *)addr + kTargetOffset, kDataLength)); +} + +// Verify RDMA WRITE when local source buffer range straddles an MR boundary. +TEST_F(RDMALargeMrTest, WriteWithSourceStraddlingChunkBoundary) { + const size_t kDataLength = 1ull << 20; + const size_t kSourceOffset = kMaxMrSize - kDataLength / 2 + 1; + const size_t kTargetOffset = 2 * kMaxMrSize; + ASSERT_NE(kSourceOffset % globalConfig().slice_size, 0u); + + for (size_t i = 0; i < kDataLength; ++i) + *((char *)addr + kSourceOffset + i) = (char)('A' + (lrand48() % 26)); + memset((char *)addr + kTargetOffset, 0, kDataLength); + + auto batch_id = engine->allocateBatchID(1); + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = (uint8_t *)addr + kSourceOffset; + entry.target_id = LOCAL_SEGMENT_ID; + entry.target_offset = (uint64_t)addr + kTargetOffset; + + Status s = engine->submitTransfer(batch_id, {entry}); + ASSERT_TRUE(s.ok()); + + TransferStatus status; + while (true) { + ASSERT_EQ(engine->getTransferStatus(batch_id, 0, status), Status::OK()); + if (status.s != TransferStatusEnum::WAITING) break; + } + ASSERT_EQ(engine->freeBatchID(batch_id), Status::OK()); + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED); + ASSERT_EQ(0, memcmp((char *)addr + kSourceOffset, + (char *)addr + kTargetOffset, kDataLength)); +} + +#if defined(USE_CUDA) || defined(USE_HIP) +class RDMAGpuDmabufChunkTest : public ::testing::Test { + protected: + std::unique_ptr engine; + void *gpu_addr = nullptr; + + void SetUp() override { + int device_count = 0; + ASSERT_EQ(cudaGetDeviceCount(&device_count), cudaSuccess); + if (device_count == 0) GTEST_SKIP() << "No GPU available"; + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + +#if defined(USE_CUDA) + ASSERT_FALSE(Environ::Get().GetWithNvidiaPeermem()) + << "Launch with WITH_NVIDIA_PEERMEM=0 so this test exercises " + "ibv_reg_dmabuf_mr instead of the nvidia-peermem fallback"; +#endif + + engine = std::make_unique(true); + const char *env_meta = std::getenv("MC_METADATA_SERVER"); + std::string metadata_server = env_meta ? env_meta : "P2PHANDSHAKE"; + ASSERT_EQ(engine->init(metadata_server, "test_node_gpu_dmabuf_chunk"), + 0); + ASSERT_EQ(globalConfig().max_mr_size, kMaxMrSize) + << "Launch this test process with MC_MAX_MR_SIZE=" << kMaxMrSize; + + ASSERT_EQ(cudaMalloc(&gpu_addr, kBufferSize), cudaSuccess); + ASSERT_EQ(engine->registerLocalMemory(gpu_addr, kBufferSize, + GPU_PREFIX + "0"), + 0); + } + + void TearDown() override { + if (engine && gpu_addr) engine->unregisterLocalMemory(gpu_addr); + if (gpu_addr) cudaFree(gpu_addr); + } +}; + +// The allocation is split into four MRs. Before the fix, every chunk is +// registered with the dma-buf offset of chunk 0. A loopback WRITE targeting +// chunk 3 therefore maps the wrong GPU pages, fails registration/transfer, or +// completes without updating the requested destination bytes. +TEST_F(RDMAGpuDmabufChunkTest, LaterChunkUsesItsOwnDmabufOffset) { + const size_t kDataLength = 1ull << 20; + const size_t kTargetOffset = kBufferSize - kDataLength; + std::vector source(kDataLength); + std::vector result(kDataLength, 0); + for (size_t i = 0; i < source.size(); ++i) + source[i] = (char)('a' + (lrand48() % 26)); + + ASSERT_EQ(cudaMemcpy(gpu_addr, source.data(), kDataLength, + cudaMemcpyHostToDevice), + cudaSuccess); + ASSERT_EQ(cudaMemset((char *)gpu_addr + kTargetOffset, 0, kDataLength), + cudaSuccess); + + auto batch_id = engine->allocateBatchID(1); + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kDataLength; + entry.source = gpu_addr; + entry.target_id = LOCAL_SEGMENT_ID; + entry.target_offset = (uint64_t)gpu_addr + kTargetOffset; + ASSERT_TRUE(engine->submitTransfer(batch_id, {entry}).ok()); + + TransferStatus status; + while (true) { + ASSERT_EQ(engine->getTransferStatus(batch_id, 0, status), Status::OK()); + if (status.s != TransferStatusEnum::WAITING) break; + } + ASSERT_EQ(engine->freeBatchID(batch_id), Status::OK()); + ASSERT_EQ(status.s, TransferStatusEnum::COMPLETED); + + ASSERT_EQ(cudaMemcpy(result.data(), (char *)gpu_addr + kTargetOffset, + kDataLength, cudaMemcpyDeviceToHost), + cudaSuccess); + EXPECT_EQ(result, source); +} +#endif + +} // namespace mooncake + +int main(int argc, char **argv) { + gflags::ParseCommandLineFlags(&argc, &argv, false); + // Initialize logging once for the whole binary — calling + // InitGoogleLogging() per-test in SetUp() aborts on the 2nd TEST_F. + google::InitGoogleLogging(argv[0]); + FLAGS_logtostderr = 1; + ::testing::InitGoogleTest(&argc, argv); + int rc = RUN_ALL_TESTS(); + google::ShutdownGoogleLogging(); + return rc; +} diff --git a/mooncake-transfer-engine/tests/rdma_loopback_test.cpp b/mooncake-transfer-engine/tests/rdma_loopback_test.cpp index 791d3af269..466eaa085f 100644 --- a/mooncake-transfer-engine/tests/rdma_loopback_test.cpp +++ b/mooncake-transfer-engine/tests/rdma_loopback_test.cpp @@ -27,11 +27,11 @@ using namespace mooncake; -namespace mooncake { - DEFINE_string(metadata_server, "127.0.0.1:2379", "central metadata server for transfer engine"); +namespace mooncake { + class RDMALoopbackTest : public ::testing::Test { public: void *addr = nullptr; @@ -96,4 +96,4 @@ int main(int argc, char **argv) { gflags::ParseCommandLineFlags(&argc, &argv, false); ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); -} \ No newline at end of file +} diff --git a/mooncake-transfer-engine/tests/rdma_test_peers.h b/mooncake-transfer-engine/tests/rdma_test_peers.h new file mode 100644 index 0000000000..8fa2840ce7 --- /dev/null +++ b/mooncake-transfer-engine/tests/rdma_test_peers.h @@ -0,0 +1,83 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Shared test-only friend accessors for exercising RdmaTransport/RdmaContext +// without a real RDMA device. Include this instead of redeclaring +// RdmaTransportTestPeer/RdmaContextTestPeer in individual test files. + +#ifndef MOONCAKE_TESTS_RDMA_TEST_PEERS_H +#define MOONCAKE_TESTS_RDMA_TEST_PEERS_H + +#include +#include + +#include "transfer_metadata.h" +#include "transport/rdma_transport/rdma_context.h" +#include "transport/rdma_transport/rdma_transport.h" + +namespace mooncake { + +class RdmaTransportTestPeer { + public: + static void bindMetadata(RdmaTransport &transport, + std::shared_ptr metadata, + std::string local_server_name) { + transport.metadata_ = std::move(metadata); + transport.local_server_name_ = std::move(local_server_name); + } + + // Registers a (possibly bare, unconstructed) RdmaContext as one of the + // transport's local devices, bypassing initializeRdmaResources(). + static void addContext(RdmaTransport &transport, + std::shared_ptr context) { + transport.context_list_.push_back(std::move(context)); + } + + static void bindTopology(RdmaTransport &transport, + std::shared_ptr topology) { + transport.local_topology_ = std::move(topology); + } + + // Drives the device-initialization loop directly so tests can assert on + // the resulting context_list_ layout without a full install(). + static int initializeResources(RdmaTransport &transport) { + return transport.initializeRdmaResources(); + } +}; + +class RdmaContextTestPeer { + public: + static bool hasEndpointStore(const RdmaContext &context) { + return context.endpoint_store_ != nullptr; + } + + static void seedAutoGidState(RdmaContext &context, ibv_context *verbs_ctx, + uint8_t port, uint16_t lid, const ibv_gid &gid, + int gid_index) { + context.context_ = verbs_ctx; + context.port_ = port; + context.lid_ = lid; + context.gid_ = gid; + context.gid_index_ = gid_index; + context.auto_gid_selection_enabled_ = true; + } + + static void disableContextForTeardown(RdmaContext &context) { + context.context_ = nullptr; + } +}; + +} // namespace mooncake + +#endif // MOONCAKE_TESTS_RDMA_TEST_PEERS_H diff --git a/mooncake-transfer-engine/tests/rdma_transport_submit_task_test.cpp b/mooncake-transfer-engine/tests/rdma_transport_submit_task_test.cpp new file mode 100644 index 0000000000..fc2d2fe8d4 --- /dev/null +++ b/mooncake-transfer-engine/tests/rdma_transport_submit_task_test.cpp @@ -0,0 +1,168 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Regression tests for RdmaTransport::submitTransferTask()'s +// "memory region not registered" (!found_device) error path. A slice that +// already succeeded device selection and was queued into the function-local +// slices_to_post accumulator remains owned by TransferTask::slice_list. It +// must not be deallocated twice, and it must reach a terminal FAILED state if +// a later slice causes submitTransferTask() to return an error. + +#include + +#include +#include + +#include "config.h" +#include "rdma_test_peers.h" +#include "transfer_metadata.h" +#include "transport/rdma_transport/rdma_context.h" +#include "transport/rdma_transport/rdma_transport.h" + +using namespace mooncake; + +namespace { + +using SegmentDesc = TransferMetadata::SegmentDesc; +using BufferDesc = TransferMetadata::BufferDesc; + +class SubmitTransferTaskTest : public ::testing::Test { + protected: + static constexpr uint64_t kBufferAddr = 0x10000; + + std::shared_ptr metadata_; + std::unique_ptr transport_; + std::shared_ptr context_; + uint64_t block_size_ = 0; + + void SetUp() override { + block_size_ = globalConfig().slice_size; + + metadata_ = std::make_shared(P2PHANDSHAKE); + transport_ = std::make_unique(); + RdmaTransportTestPeer::bindMetadata(*transport_, metadata_, + "unit-test-server:1234"); + + // construct() is never called: no real device is opened. active() + // defaults to true, which is all submitTransferTask() checks. + context_ = std::make_shared(*transport_, "mlx5_unit_test"); + RdmaTransportTestPeer::addContext(*transport_, context_); + + auto desc = std::make_shared(); + desc->name = "unit-test-server:1234"; + desc->protocol = "rdma"; + BufferDesc buffer; + buffer.name = "cpu:0"; + buffer.addr = kBufferAddr; + buffer.length = block_size_; + buffer.lkey = {1}; + buffer.rkey = {1}; + desc->buffers.push_back(buffer); + ASSERT_EQ( + desc->topology.parse(R"({"cpu:0": [["mlx5_unit_test"], []]})"), 0); + metadata_->addLocalSegment(LOCAL_SEGMENT_ID, desc->name, + std::move(desc)); + } + + void triggerError(Transport::TransferRequest &req, + Transport::TransferTask &task) { + req.opcode = Transport::TransferRequest::WRITE; + req.source = reinterpret_cast(kBufferAddr); + req.length = 2 * block_size_; + req.target_id = LOCAL_SEGMENT_ID; + req.target_offset = 0; + task.request = &req; + + // markFailed() needs a valid BatchDesc in event-driven builds. The + // direct task is deliberately not inserted into that BatchDesc; the ID + // is only used by Slice::check_batch_completion(). + task.batch_id = transport_->allocateBatchID(1); + auto status = transport_->submitTransferTask({&task}); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.IsAddressNotRegistered()); + ASSERT_EQ(task.slice_list.size(), 2u); + EXPECT_EQ(transport_->freeBatchID(task.batch_id), Status::OK()); + } +}; + +TEST_F(SubmitTransferTaskTest, NoDuplicateSlice) { + Transport::Slice *original = nullptr; + { + Transport::TransferRequest req; + Transport::TransferTask task; + triggerError(req, task); + original = task.slice_list[0]; + } + + Transport::TransferRequest req; + Transport::TransferTask task; + triggerError(req, task); + + EXPECT_EQ(task.slice_list[0], original) + << "the cache should legitimately reuse the first released slice"; + EXPECT_NE(task.slice_list[0], task.slice_list[1]) + << "two independent slices must never share the same Slice object"; +} + +TEST_F(SubmitTransferTaskTest, PartialSubmitFailsBatch) { + auto batch_id = transport_->allocateBatchID(1); + Transport::TransferRequest request; + request.opcode = Transport::TransferRequest::WRITE; + request.source = reinterpret_cast(kBufferAddr); + request.length = 2 * block_size_; + request.target_id = LOCAL_SEGMENT_ID; + request.target_offset = 0; + + auto submit_status = transport_->submitTransfer(batch_id, {request}); + ASSERT_FALSE(submit_status.ok()); + ASSERT_TRUE(submit_status.IsAddressNotRegistered()); + + Transport::TransferStatus transfer_status; + ASSERT_EQ(transport_->getTransferStatus(batch_id, 0, transfer_status), + Status::OK()); + EXPECT_EQ(transfer_status.s, Transport::TransferStatusEnum::FAILED); + EXPECT_EQ(transport_->freeBatchID(batch_id), Status::OK()); +} + +TEST_F(SubmitTransferTaskTest, PartialSubmitFailsAllTasks) { + auto batch_id = transport_->allocateBatchID(2); + Transport::TransferRequest failing_request; + failing_request.opcode = Transport::TransferRequest::WRITE; + failing_request.source = reinterpret_cast(kBufferAddr); + failing_request.length = 2 * block_size_; + failing_request.target_id = LOCAL_SEGMENT_ID; + failing_request.target_offset = 0; + + Transport::TransferRequest unstarted_request; + unstarted_request.opcode = Transport::TransferRequest::WRITE; + unstarted_request.source = reinterpret_cast(kBufferAddr); + unstarted_request.length = block_size_; + unstarted_request.target_id = LOCAL_SEGMENT_ID; + unstarted_request.target_offset = kBufferAddr; + + auto submit_status = transport_->submitTransfer( + batch_id, {failing_request, unstarted_request}); + ASSERT_FALSE(submit_status.ok()); + ASSERT_TRUE(submit_status.IsAddressNotRegistered()); + + std::vector transfer_status; + ASSERT_EQ(transport_->getTransferStatus(batch_id, transfer_status), + Status::OK()); + ASSERT_EQ(transfer_status.size(), 2u); + EXPECT_EQ(transfer_status[0].s, Transport::TransferStatusEnum::FAILED); + EXPECT_EQ(transfer_status[1].s, Transport::TransferStatusEnum::FAILED); + EXPECT_EQ(transport_->freeBatchID(batch_id), Status::OK()); +} + +} // namespace diff --git a/mooncake-transfer-engine/tests/rdma_transport_test.cpp b/mooncake-transfer-engine/tests/rdma_transport_test.cpp index 6c6e747bce..d9ecd35052 100644 --- a/mooncake-transfer-engine/tests/rdma_transport_test.cpp +++ b/mooncake-transfer-engine/tests/rdma_transport_test.cpp @@ -50,7 +50,7 @@ #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ - defined(USE_COREX) + defined(USE_COREX) || defined(USE_SUPA) #include @@ -100,8 +100,9 @@ std::string pickBackend() { } #if defined(USE_MLU) return "mlu"; -#elif defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ - defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) +#elif defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ + defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ + defined(USE_SUPA) return FLAGS_use_vram ? "gpu" : "cpu"; #else return "cpu"; @@ -112,8 +113,9 @@ int pickDevId(const std::string &backend) { if (FLAGS_device_id >= 0) { return FLAGS_device_id; } -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ - defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ + defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ + defined(USE_SUPA) if (backend == "gpu") { return FLAGS_gpu_id; } @@ -127,8 +129,9 @@ void validateBackend(const std::string &backend) { if (backend == "cpu") { return; } -#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ - defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) +#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ + defined(USE_MACA) || defined(USE_HYGON) || defined(USE_COREX) || \ + defined(USE_SUPA) if (backend == "gpu") { return; } @@ -177,7 +180,7 @@ bool validateTransferSizes() { } #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ - defined(USE_COREX) + defined(USE_COREX) || defined(USE_SUPA) checkCudaError(cudaSetDevice(pickDevId(backend)), "Failed to set device"); #else LOG(FATAL) << "Device memory backend is not available in this build"; @@ -189,7 +192,7 @@ void *allocateMemoryPool(size_t size, int socket_id, if (usesDeviceMemory(backend)) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ - defined(USE_COREX) + defined(USE_COREX) || defined(USE_SUPA) setBackendDevice(backend); void *d_buf = nullptr; checkCudaError(cudaMalloc(&d_buf, size), @@ -207,7 +210,7 @@ void freeMemoryPool(void *addr, size_t size, const std::string &backend) { if (usesDeviceMemory(backend)) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ - defined(USE_COREX) + defined(USE_COREX) || defined(USE_SUPA) cudaFree(addr); return; #else @@ -222,7 +225,7 @@ void copyFromHost(void *dst, const void *src, size_t size, if (usesDeviceMemory(backend)) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ - defined(USE_COREX) + defined(USE_COREX) || defined(USE_SUPA) checkCudaError(cudaMemcpy(dst, src, size, cudaMemcpyHostToDevice), "Failed to copy host data to device"); return; @@ -238,7 +241,7 @@ void copyToHost(void *dst, const void *src, size_t size, if (usesDeviceMemory(backend)) { #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \ defined(USE_MLU) || defined(USE_MACA) || defined(USE_HYGON) || \ - defined(USE_COREX) + defined(USE_COREX) || defined(USE_SUPA) checkCudaError(cudaMemcpy(dst, src, size, cudaMemcpyDeviceToHost), "Failed to copy device data to host"); return; diff --git a/mooncake-transfer-engine/tests/rdma_transport_test2.cpp b/mooncake-transfer-engine/tests/rdma_transport_test2.cpp index 1c48ac46ad..0ea85938f5 100644 --- a/mooncake-transfer-engine/tests/rdma_transport_test2.cpp +++ b/mooncake-transfer-engine/tests/rdma_transport_test2.cpp @@ -28,8 +28,6 @@ using namespace mooncake; -namespace mooncake { - DEFINE_string(local_server_name, getHostname(), "Local server name for segment discovery"); DEFINE_string(metadata_server, "127.0.0.1:2379", "etcd server host address"); @@ -47,6 +45,8 @@ DEFINE_string(nic_priority_matrix, "", DEFINE_string(segment_id, "127.0.0.2", "Segment ID to access data"); +namespace mooncake { + std::string formatDeviceNames(const std::string &device_names) { std::stringstream ss(device_names); std::string item; @@ -254,4 +254,4 @@ int main(int argc, char **argv) { gflags::ParseCommandLineFlags(&argc, &argv, false); ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); -} \ No newline at end of file +} diff --git a/mooncake-transfer-engine/tests/show_links_test.cpp b/mooncake-transfer-engine/tests/show_links_test.cpp new file mode 100644 index 0000000000..27706ddd6e --- /dev/null +++ b/mooncake-transfer-engine/tests/show_links_test.cpp @@ -0,0 +1,120 @@ +// Copyright 2024 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include + +#include +#include +#include + +#include "transfer_engine.h" +#include "transfer_engine_c.h" + +using namespace mooncake; + +TEST(ShowLinksTest, NoInitReturnsEmptyOrPlaceholder) { + auto engine = std::make_unique(false); + auto result = engine->showLinks(); + EXPECT_FALSE(result.empty()); + + auto json_result = engine->showLinks(true); + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::unique_ptr reader(builder.newCharReader()); + ASSERT_TRUE(reader->parse(json_result.data(), + json_result.data() + json_result.size(), &root, + &errors)) + << errors; + EXPECT_TRUE(root.isMember("local_nics")); +} + +TEST(ShowLinksTest, CApiSupportsJsonAndRejectsInvalidArguments) { + auto engine = std::make_unique(false); + char output[4096] = {}; + auto handle = reinterpret_cast(engine.get()); + + ASSERT_EQ(::showLinks(handle, output, sizeof(output), 1), 0); + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::unique_ptr reader(builder.newCharReader()); + ASSERT_TRUE( + reader->parse(output, output + std::strlen(output), &root, &errors)) + << errors; + EXPECT_TRUE(root.isMember("local_nics")); + + EXPECT_NE(::showLinks(nullptr, output, sizeof(output), 0), 0); + EXPECT_NE(::showLinks(handle, nullptr, sizeof(output), 0), 0); + EXPECT_NE(::showLinks(handle, output, 0, 0), 0); +} + +TEST(ShowLinksTest, AutoDiscoverShowsNics) { + auto engine = std::make_unique(true); + auto result = engine->showLinks(); + EXPECT_NE(result.find("Local NICs"), std::string::npos); +} + +TEST(ShowLinksTest, OutputContainsTopologySection) { + auto engine = std::make_unique(true); + auto result = engine->showLinks(); + // If RDMA devices exist, should show topology + // If not, gracefully show empty + EXPECT_FALSE(result.empty()); +} + +TEST(ShowLinksTest, JsonOutputIsValid) { + auto engine = std::make_unique(true); + auto result = engine->showLinks(true); + EXPECT_FALSE(result.empty()); + + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::unique_ptr reader(builder.newCharReader()); + ASSERT_TRUE(reader->parse(result.data(), result.data() + result.size(), + &root, &errors)) + << errors; + EXPECT_TRUE(root.isMember("local_nics")); +} + +TEST(ShowLinksTest, TopologyOnlyNicsAppearWithoutTransport) { + auto engine = std::make_unique(false); + auto topology = engine->getLocalTopology(); + ASSERT_NE(topology, nullptr); + ASSERT_EQ(topology->parse("{\"cpu:0\" : [[\"erdma_0\"],[\"erdma_1\"]]}"), + 0); + + auto readable = engine->showLinks(); + EXPECT_NE(readable.find("erdma_0"), std::string::npos); + EXPECT_NE(readable.find("erdma_1"), std::string::npos); + EXPECT_NE(readable.find("topology only"), std::string::npos); + + auto json_result = engine->showLinks(true); + Json::Value root; + Json::CharReaderBuilder builder; + std::string errors; + std::unique_ptr reader(builder.newCharReader()); + ASSERT_TRUE(reader->parse(json_result.data(), + json_result.data() + json_result.size(), &root, + &errors)) + << errors; + + ASSERT_TRUE(root["local_nics"].isArray()); + ASSERT_EQ(root["local_nics"].size(), 2u); + for (const auto& nic : root["local_nics"]) { + EXPECT_EQ(nic["source"].asString(), "topology"); + } +} diff --git a/mooncake-transfer-engine/tests/sunrise_allocator_test.cpp b/mooncake-transfer-engine/tests/sunrise_allocator_test.cpp new file mode 100644 index 0000000000..579044a4dd --- /dev/null +++ b/mooncake-transfer-engine/tests/sunrise_allocator_test.cpp @@ -0,0 +1,215 @@ +#include + +#include +#include +#include +#include + +#include "sunrise_allocator.h" + +using namespace mooncake::sunrise_alloc_detail; + +class SunriseAllocatorTest : public ::testing::Test { + protected: + void SetUp() override { + std::lock_guard lock(tangAllocMutex()); + saved_host_set_ = tangHostAllocatedSet(); + saved_dev_set_ = tangDeviceAllocatedSet(); + saved_ranges_ = storeMemRanges(); + tangHostAllocatedSet().clear(); + tangDeviceAllocatedSet().clear(); + storeMemRanges().clear(); + } + + void TearDown() override { + std::lock_guard lock(tangAllocMutex()); + tangHostAllocatedSet() = std::move(saved_host_set_); + tangDeviceAllocatedSet() = std::move(saved_dev_set_); + storeMemRanges() = std::move(saved_ranges_); + } + + private: + std::unordered_set saved_host_set_; + std::unordered_set saved_dev_set_; + std::vector saved_ranges_; +}; + +TEST_F(SunriseAllocatorTest, IsDeviceMemoryRangeNullptr) { + EXPECT_FALSE(sunrise_is_device_memory_range(nullptr)); +} + +TEST_F(SunriseAllocatorTest, IsDeviceMemoryRangeUntracked) { + int x; + EXPECT_FALSE(sunrise_is_device_memory_range(&x)); +} + +TEST_F(SunriseAllocatorTest, IsDeviceMemoryRangeBasePointer) { + void* dev_ptr = reinterpret_cast(0x1000); + size_t size = 4096; + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().insert(dev_ptr); + } + addStoreMemRange(dev_ptr, size); + + EXPECT_TRUE(sunrise_is_device_memory_range(dev_ptr)); + + removeStoreMemRange(dev_ptr); + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().erase(dev_ptr); + } +} + +TEST_F(SunriseAllocatorTest, IsDeviceMemoryRangeSubPointer) { + void* base = reinterpret_cast(0x2000); + size_t size = 8192; + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().insert(base); + } + addStoreMemRange(base, size); + + void* mid = reinterpret_cast(0x2000 + 100); + void* end_minus_one = reinterpret_cast(0x2000 + size - 1); + EXPECT_TRUE(sunrise_is_device_memory_range(mid)); + EXPECT_TRUE(sunrise_is_device_memory_range(end_minus_one)); + + void* past_end = reinterpret_cast(0x2000 + size); + EXPECT_FALSE(sunrise_is_device_memory_range(past_end)); + + removeStoreMemRange(base); + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().erase(base); + } +} + +TEST_F(SunriseAllocatorTest, IsHostAllocatedBasePointer) { + void* host_ptr = reinterpret_cast(0x3000); + size_t size = 4096; + { + std::lock_guard lock(tangAllocMutex()); + tangHostAllocatedSet().insert(host_ptr); + } + addStoreMemRange(host_ptr, size); + + EXPECT_TRUE(sunrise_is_host_allocated(host_ptr)); + EXPECT_FALSE(sunrise_is_device_memory_range(host_ptr)); + + removeStoreMemRange(host_ptr); + { + std::lock_guard lock(tangAllocMutex()); + tangHostAllocatedSet().erase(host_ptr); + } +} + +TEST_F(SunriseAllocatorTest, IsHostAllocatedSubPointer) { + void* base = reinterpret_cast(0x4000); + size_t size = 8192; + { + std::lock_guard lock(tangAllocMutex()); + tangHostAllocatedSet().insert(base); + } + addStoreMemRange(base, size); + + void* mid = reinterpret_cast(0x4000 + 500); + EXPECT_TRUE(sunrise_is_host_allocated(mid)); + EXPECT_FALSE(sunrise_is_device_memory_range(mid)); + + removeStoreMemRange(base); + { + std::lock_guard lock(tangAllocMutex()); + tangHostAllocatedSet().erase(base); + } +} + +TEST_F(SunriseAllocatorTest, HostAllocNotConfusedWithDevice) { + void* host_ptr = reinterpret_cast(0x5000); + size_t size = 4096; + { + std::lock_guard lock(tangAllocMutex()); + tangHostAllocatedSet().insert(host_ptr); + } + addStoreMemRange(host_ptr, size); + + EXPECT_TRUE(sunrise_is_host_allocated(host_ptr)); + EXPECT_FALSE(sunrise_is_device_memory_range(host_ptr)); + + removeStoreMemRange(host_ptr); + { + std::lock_guard lock(tangAllocMutex()); + tangHostAllocatedSet().erase(host_ptr); + } +} + +TEST_F(SunriseAllocatorTest, DeviceAllocNotConfusedWithHost) { + void* dev_ptr = reinterpret_cast(0x6000); + size_t size = 4096; + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().insert(dev_ptr); + } + addStoreMemRange(dev_ptr, size); + + EXPECT_TRUE(sunrise_is_device_memory_range(dev_ptr)); + EXPECT_FALSE(sunrise_is_host_allocated(dev_ptr)); + + removeStoreMemRange(dev_ptr); + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().erase(dev_ptr); + } +} + +TEST_F(SunriseAllocatorTest, RemoveStoreMemRangeCleansUp) { + void* base = reinterpret_cast(0x8000); + size_t size = 4096; + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().insert(base); + } + addStoreMemRange(base, size); + EXPECT_TRUE(sunrise_is_device_memory_range(base)); + + removeStoreMemRange(base); + EXPECT_FALSE(sunrise_is_device_memory_range(base)); + + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().erase(base); + } +} + +TEST_F(SunriseAllocatorTest, MultipleRanges) { + void* dev_base = reinterpret_cast(0x9000); + size_t dev_size = 4096; + void* host_base = reinterpret_cast(0xA000); + size_t host_size = 8192; + + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().insert(dev_base); + tangHostAllocatedSet().insert(host_base); + } + addStoreMemRange(dev_base, dev_size); + addStoreMemRange(host_base, host_size); + + EXPECT_TRUE(sunrise_is_device_memory_range(dev_base)); + EXPECT_FALSE(sunrise_is_device_memory_range(host_base)); + EXPECT_TRUE(sunrise_is_host_allocated(host_base)); + EXPECT_FALSE(sunrise_is_host_allocated(dev_base)); + + void* dev_mid = reinterpret_cast(0x9000 + 50); + void* host_mid = reinterpret_cast(0xA000 + 50); + EXPECT_TRUE(sunrise_is_device_memory_range(dev_mid)); + EXPECT_TRUE(sunrise_is_host_allocated(host_mid)); + + removeStoreMemRange(dev_base); + removeStoreMemRange(host_base); + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().erase(dev_base); + tangHostAllocatedSet().erase(host_base); + } +} diff --git a/mooncake-transfer-engine/tests/sunrise_link_copy_test.cpp b/mooncake-transfer-engine/tests/sunrise_link_copy_test.cpp new file mode 100644 index 0000000000..9fcc547146 --- /dev/null +++ b/mooncake-transfer-engine/tests/sunrise_link_copy_test.cpp @@ -0,0 +1,186 @@ +#include + +#include +#include +#include +#include + +#include "gpu_vendor/sunrise.h" +#include "sunrise_allocator.h" +#include "transfer_engine.h" +#include "transport/transport.h" + +using namespace mooncake; +using namespace mooncake::sunrise_alloc_detail; + +namespace { + +static constexpr size_t kTestSize = 4096; +static constexpr size_t kLargeTestSize = 4 * 1024 * 1024; + +class SunriseLinkCopyTest : public ::testing::Test { + protected: + void SetUp() override { + int gpu_count = 0; + if (tangGetDeviceCount(&gpu_count) != tangSuccess || gpu_count <= 0) { + GTEST_SKIP() << "Sunrise device is unavailable"; + } + ASSERT_EQ(tangSetDevice(0), tangSuccess); + } + + static std::vector makePattern(size_t size, char seed) { + std::vector data(size); + for (size_t i = 0; i < size; ++i) + data[i] = static_cast((seed + i) % 256); + return data; + } + + static bool verifyHostBuffer(const void* host_ptr, size_t size, char seed) { + auto expected = makePattern(size, seed); + return memcmp(host_ptr, expected.data(), size) == 0; + } +}; + +} // namespace + +TEST_F(SunriseLinkCopyTest, DeviceToHostAllocStaging) { + void* dev_buf = nullptr; + ASSERT_EQ(tangMalloc(&dev_buf, kTestSize), tangSuccess); + void* host_buf = nullptr; + ASSERT_EQ(tangHostAlloc(&host_buf, kTestSize, 0), tangSuccess); + + auto pattern = makePattern(kTestSize, 'A'); + ASSERT_EQ( + tangMemcpy(dev_buf, pattern.data(), kTestSize, tangMemcpyHostToDevice), + tangSuccess); + + std::vector staging(kTestSize); + ASSERT_EQ( + tangMemcpy(staging.data(), dev_buf, kTestSize, tangMemcpyDeviceToHost), + tangSuccess); + memcpy(host_buf, staging.data(), kTestSize); + EXPECT_TRUE(verifyHostBuffer(host_buf, kTestSize, 'A')); + + tangFree(dev_buf); + tangFreeHost(host_buf); +} + +TEST_F(SunriseLinkCopyTest, HostAllocToDeviceStaging) { + void* host_buf = nullptr; + ASSERT_EQ(tangHostAlloc(&host_buf, kTestSize, 0), tangSuccess); + void* dev_buf = nullptr; + ASSERT_EQ(tangMalloc(&dev_buf, kTestSize), tangSuccess); + + auto pattern = makePattern(kTestSize, 'B'); + memcpy(host_buf, pattern.data(), kTestSize); + + std::vector staging(kTestSize); + memcpy(staging.data(), host_buf, kTestSize); + ASSERT_EQ( + tangMemcpy(dev_buf, staging.data(), kTestSize, tangMemcpyHostToDevice), + tangSuccess); + + std::vector verify(kTestSize); + ASSERT_EQ( + tangMemcpy(verify.data(), dev_buf, kTestSize, tangMemcpyDeviceToHost), + tangSuccess); + EXPECT_EQ(memcmp(verify.data(), pattern.data(), kTestSize), 0); + + tangFree(dev_buf); + tangFreeHost(host_buf); +} + +TEST_F(SunriseLinkCopyTest, DeviceToHostAllocLargeData) { + void* dev_buf = nullptr; + ASSERT_EQ(tangMalloc(&dev_buf, kLargeTestSize), tangSuccess); + void* host_buf = nullptr; + ASSERT_EQ(tangHostAlloc(&host_buf, kLargeTestSize, 0), tangSuccess); + + auto pattern = makePattern(kLargeTestSize, 'C'); + ASSERT_EQ(tangMemcpy(dev_buf, pattern.data(), kLargeTestSize, + tangMemcpyHostToDevice), + tangSuccess); + + std::vector staging(kLargeTestSize); + ASSERT_EQ(tangMemcpy(staging.data(), dev_buf, kLargeTestSize, + tangMemcpyDeviceToHost), + tangSuccess); + memcpy(host_buf, staging.data(), kLargeTestSize); + EXPECT_TRUE(verifyHostBuffer(host_buf, kLargeTestSize, 'C')); + + tangFree(dev_buf); + tangFreeHost(host_buf); +} + +TEST_F(SunriseLinkCopyTest, DeviceToHostAllocStagingConsistency) { + for (int trial = 0; trial < 5; ++trial) { + void* dev_buf = nullptr; + ASSERT_EQ(tangMalloc(&dev_buf, kTestSize), tangSuccess); + void* host_buf = nullptr; + ASSERT_EQ(tangHostAlloc(&host_buf, kTestSize, 0), tangSuccess); + + char seed = static_cast('D' + trial); + auto pattern = makePattern(kTestSize, seed); + ASSERT_EQ(tangMemcpy(dev_buf, pattern.data(), kTestSize, + tangMemcpyHostToDevice), + tangSuccess); + + std::vector staging(kTestSize); + ASSERT_EQ(tangMemcpy(staging.data(), dev_buf, kTestSize, + tangMemcpyDeviceToHost), + tangSuccess); + memcpy(host_buf, staging.data(), kTestSize); + ASSERT_TRUE(verifyHostBuffer(host_buf, kTestSize, seed)); + + tangFree(dev_buf); + tangFreeHost(host_buf); + } +} + +TEST_F(SunriseLinkCopyTest, IsDeviceMemoryRangeOnRealAlloc) { + void* dev_buf = nullptr; + ASSERT_EQ(tangMalloc(&dev_buf, kTestSize), tangSuccess); + addStoreMemRange(dev_buf, kTestSize); + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().insert(dev_buf); + } + + EXPECT_TRUE(sunrise_is_device_memory_range(dev_buf)); + EXPECT_TRUE( + sunrise_is_device_memory_range(static_cast(dev_buf) + 100)); + EXPECT_FALSE(sunrise_is_host_allocated(dev_buf)); + + removeStoreMemRange(dev_buf); + { + std::lock_guard lock(tangAllocMutex()); + tangDeviceAllocatedSet().erase(dev_buf); + } + tangFree(dev_buf); +} + +TEST_F(SunriseLinkCopyTest, IsHostAllocatedOnRealAlloc) { + void* host_buf = nullptr; + ASSERT_EQ(tangHostAlloc(&host_buf, kTestSize, 0), tangSuccess); + addStoreMemRange(host_buf, kTestSize); + { + std::lock_guard lock(tangAllocMutex()); + tangHostAllocatedSet().insert(host_buf); + } + + EXPECT_TRUE(sunrise_is_host_allocated(host_buf)); + EXPECT_TRUE(sunrise_is_host_allocated(static_cast(host_buf) + 100)); + EXPECT_FALSE(sunrise_is_device_memory_range(host_buf)); + + removeStoreMemRange(host_buf); + { + std::lock_guard lock(tangAllocMutex()); + tangHostAllocatedSet().erase(host_buf); + } + tangFreeHost(host_buf); +} + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp b/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp new file mode 100644 index 0000000000..6bc9d23408 --- /dev/null +++ b/mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp @@ -0,0 +1,638 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Contract tests for TcpTransport completion semantics (issue #2086). +// +// Under the legacy (v1) framing, a WRITE was reported COMPLETED when the +// final chunk reached the initiator's kernel socket buffer: destination +// memory could still be mutating megabytes later (measured 166/400 +// iterations torn, worst case the entire 2.4 MB descriptor undelivered), +// and a server-side rejection was invisible to the initiator (a +// single-chunk WRITE to an unregistered address "succeeded"). The v2 +// acknowledged framing makes COMPLETED mean "applied at the destination" +// and failures mean failures; these tests pin that contract and the +// mixed-version behavior. The main visibility test fails against the +// legacy framing (which the MC_TCP_PROTO=1 escape hatch still selects). + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "transfer_engine.h" +#include "transport/transport.h" + +using namespace mooncake; + +namespace { + +constexpr size_t kBigLength = 2432 * 1024; // ~2.4 MB, as reported in #2086 +constexpr size_t kSmallLength = 16 * 1024; // 16 KB control size +constexpr size_t kRegionAlign = 4 * 1024 * 1024; +constexpr int kIterations = 400; +constexpr int kNoiseThreads = 3; + +class ScopedEnvVar { + public: + ScopedEnvVar(const char* name, const char* value) : name_(name) { + if (const char* old = std::getenv(name)) { + had_old_value_ = true; + old_value_ = old; + } + setenv(name_.c_str(), value, 1); + } + + ~ScopedEnvVar() { + if (had_old_value_) + setenv(name_.c_str(), old_value_.c_str(), 1); + else + unsetenv(name_.c_str()); + } + + ScopedEnvVar(const ScopedEnvVar&) = delete; + ScopedEnvVar& operator=(const ScopedEnvVar&) = delete; + + private: + std::string name_; + std::string old_value_; + bool had_old_value_ = false; +}; + +struct TestSessionHeader { + uint64_t size; + uint64_t addr; + uint8_t opcode; +}; + +static_assert(sizeof(TestSessionHeader) == 24, + "legacy TCP header ABI changed unexpectedly"); + +// Minimal legacy-server behavior for a flagged WRITE: v1 does not recognize +// opcode 0x81 as WRITE, so it treats the request as READ and streams payload +// bytes without consuming the initiator's body. A sequential v2 client then +// deadlocks once both socket directions fill; the concurrent status read must +// reject the non-status bytes and cancel the body promptly. +class LegacyReadServer { + public: + // keep_open=true models the real v1 server loop: after streaming the + // "READ payload" it does NOT close, it waits for the next 24-byte + // header. For flagged requests shorter than a status frame this is the + // configuration that used to hang a v2 initiator forever (fewer than 8 + // payload bytes ever arrive, no EOF, no deadline). + explicit LegacyReadServer(bool keep_open = false) : keep_open_(keep_open) { + listen_fd_ = socket(AF_INET, SOCK_STREAM, 0); + if (listen_fd_ < 0) return; + int one = 1; + if (setsockopt(listen_fd_, SOL_SOCKET, SO_REUSEADDR, &one, + sizeof(one)) != 0) + return; + + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = 0; + // The production HTTP-metadata path advertises the engine-selected + // LAN address in RPC metadata even when local_server_name is a + // loopback test name. Listen on every local interface so the stale + // descriptor can redirect either that path or P2PHANDSHAKE's + // loopback path to this fake legacy peer. + addr.sin_addr.s_addr = htonl(INADDR_ANY); + if (bind(listen_fd_, reinterpret_cast(&addr), + sizeof(addr)) != 0) + return; + if (listen(listen_fd_, 1) != 0) return; + + socklen_t len = sizeof(addr); + if (getsockname(listen_fd_, reinterpret_cast(&addr), &len) != + 0) + return; + port_ = ntohs(addr.sin_port); + ok_ = true; + thread_ = std::thread([this] { serve(); }); + } + + ~LegacyReadServer() { join(); } + + uint16_t port() const { return port_; } + bool ok() const { return ok_; } + + void join() { + // Wake a blocked accept if the client failed before connecting. This + // keeps a failed assertion from turning into a hung test process. + if (listen_fd_ >= 0) (void)shutdown(listen_fd_, SHUT_RDWR); + if (thread_.joinable()) thread_.join(); + if (listen_fd_ >= 0) { + close(listen_fd_); + listen_fd_ = -1; + } + } + + bool sawFlaggedWrite() const { return saw_flagged_write_.load(); } + + private: + static bool recvExact(int fd, void* buffer, size_t size) { + char* out = static_cast(buffer); + while (size) { + ssize_t n = recv(fd, out, size, 0); + if (n <= 0) return false; + out += n; + size -= static_cast(n); + } + return true; + } + + void serve() { + int fd = accept(listen_fd_, nullptr, nullptr); + if (fd < 0) return; + + timeval timeout{8, 0}; + (void)setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &timeout, + sizeof(timeout)); + (void)setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, + sizeof(timeout)); + + TestSessionHeader header{}; + if (!recvExact(fd, &header, sizeof(header))) { + close(fd); + return; + } + saw_flagged_write_.store(header.opcode == 0x81); + + const uint64_t total = le64toh(header.size); + std::vector payload(64 * 1024, static_cast(0xA5)); + uint64_t sent = 0; + while (sent < total) { + size_t chunk = std::min(payload.size(), total - sent); + ssize_t n = send(fd, payload.data(), chunk, MSG_NOSIGNAL); + if (n <= 0) break; + sent += static_cast(n); + } + if (keep_open_) { + // v1 loop: wait for the next header; leaves only when the + // client drops the connection (or the test tears down the + // listener, which shuts the accepted fd's peer down too). + TestSessionHeader next{}; + (void)recvExact(fd, &next, sizeof(next)); + } + close(fd); + } + + int listen_fd_ = -1; + uint16_t port_ = 0; + bool ok_ = false; + bool keep_open_ = false; + std::thread thread_; + std::atomic saw_flagged_write_{false}; +}; + +struct EngineHandle { + std::unique_ptr engine; + void* pool = nullptr; + + ~EngineHandle() { + engine.reset(); // unregisters memory before the pool goes away + free(pool); + } + Transport::SegmentID segment_id = 0; + uint64_t remote_base = 0; + bool ok = false; // ASSERT_* in a helper only aborts the helper + + void init(const std::string& metadata_server, + const std::string& server_name, size_t pool_size) { + // Exercise the pooled-connection path (default off): connection + // reuse vs. discard-on-unclean-exchange is part of the contract + // under test. + setenv("MC_TCP_ENABLE_CONNECTION_POOL", "1", 1); + engine = std::make_unique(false); + auto hp = parseHostNameWithPort(server_name); + int rc = engine->init(metadata_server, server_name, hp.first.c_str(), + hp.second); + ASSERT_EQ(rc, 0); + ASSERT_NE(engine->installTransport("tcp", nullptr), nullptr); + pool = malloc(pool_size); + ASSERT_NE(pool, nullptr); + memset(pool, 0, pool_size); + rc = engine->registerLocalMemory(pool, pool_size, "cpu:0"); + ASSERT_EQ(rc, 0); + // The descriptor is fetchable under the name it was registered + // with: in P2P-handshake mode the RPC port is auto-assigned, so + // that is the engine-reported ip:port; against a real metadata + // service (CI runs one at http://...) it is the requested + // server_name, matching how production callers open segments. + std::string segment_name = (metadata_server == P2PHANDSHAKE) + ? engine->getLocalIpAndPort() + : server_name; + segment_id = engine->openSegment(segment_name); + auto desc = engine->getMetadata()->getSegmentDescByID(segment_id); + ASSERT_NE(desc, nullptr); + remote_base = (uint64_t)desc->buffers[0].addr; + ok = true; + } +}; + +// Submit one request and poll until terminal state; returns final status. +TransferStatusEnum runOne(TransferEngine* engine, TransferRequest entry) { + auto batch_id = engine->allocateBatchID(1); + Status s = engine->submitTransfer(batch_id, {entry}); + if (!s.ok()) return TransferStatusEnum::FAILED; + TransferStatus status; + status.s = TransferStatusEnum::WAITING; + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(15); + while (status.s != TransferStatusEnum::COMPLETED && + status.s != TransferStatusEnum::FAILED) { + if (std::chrono::steady_clock::now() >= deadline) + return TransferStatusEnum::TIMEOUT; + s = engine->getTransferStatus(batch_id, 0, status); + if (!s.ok()) return TransferStatusEnum::FAILED; + std::this_thread::yield(); + } + (void)engine->freeBatchID(batch_id); + return status.s; +} + +} // namespace + +TEST(TcpWriteVisibilityTest, CompletedWriteIsVisibleToSubsequentRead) { + const char* env = std::getenv("MC_METADATA_SERVER"); + std::string metadata_server = env ? env : "P2PHANDSHAKE"; + const char* name_env = std::getenv("MC_LOCAL_SERVER_NAME"); + std::string server_name = name_env ? name_env : "127.0.0.2:17901"; + + const size_t pool_size = 64ull << 20; + EngineHandle h; + h.init(metadata_server, server_name, pool_size); + ASSERT_TRUE(h.ok) << "engine/segment setup failed"; + + // Region layout inside the registered pool (all offsets from pool base): + // [0, kBigLength) : WRITE target region + // [kRegionAlign, +kBigLength) : local staging for WRITE + // [3*kRegionAlign + t*kRegionAlign, ...): per-noise-thread scratch + char* base = (char*)h.pool; + char* write_src = base + kRegionAlign; + + std::atomic stop{false}; + std::atomic noise_failures{0}; + std::vector noise; + for (int t = 0; t < kNoiseThreads; ++t) { + noise.emplace_back([&, t] { + char* src = base + (3 + 2 * t) * kRegionAlign; + uint64_t dst_off = (4 + 2 * t) * kRegionAlign; + memset(src, 0x5A + t, kSmallLength); + while (!stop.load(std::memory_order_relaxed)) { + TransferRequest entry; + entry.opcode = TransferRequest::WRITE; + entry.length = kSmallLength; + entry.source = src; + entry.target_id = h.segment_id; + entry.target_offset = h.remote_base + dst_off; + if (runOne(h.engine.get(), entry) != + TransferStatusEnum::COMPLETED) + noise_failures++; + } + }); + } + + uint64_t torn_reads = 0; + uint64_t torn_bytes_worst = 0; + int first_bad_iter = -1; + for (int iter = 1; iter <= kIterations; ++iter) { + // Generation-stamped pattern: every byte identifies the iteration. + memset(write_src, iter & 0xFF, kBigLength); + + TransferRequest w; + w.opcode = TransferRequest::WRITE; + w.length = kBigLength; + w.source = write_src; + w.target_id = h.segment_id; + w.target_offset = h.remote_base; // region at pool offset 0 + ASSERT_EQ(runOne(h.engine.get(), w), TransferStatusEnum::COMPLETED) + << "WRITE failed at iteration " << iter; + + // The WRITE is COMPLETED. The API contract (matching RDMA WRITE + // semantics, which disaggregated-serving integrations rely on when + // they notify the consumer out-of-band) is that destination memory + // is now fully written. Verify by direct local inspection of the + // destination region — this is exactly what a decode instance does + // after the prefill side signals transfer completion. Scan backwards: + // the tail chunks are the ones still in flight when the initiator's + // final local send completes. + size_t bad = 0; + for (size_t i = kBigLength; i-- > 0;) { + if ((unsigned char)base[i] != (unsigned char)(iter & 0xFF)) { + bad = i + 1; // bytes [0, i] not yet guaranteed; count prefix + break; + } + } + if (bad) { + torn_reads++; + torn_bytes_worst = std::max(torn_bytes_worst, bad); + if (first_bad_iter < 0) first_bad_iter = iter; + // Show it is a visibility delay, not data loss: wait for the + // server-side drain to finish before the next iteration so + // generations do not overlap. + while (memcmp(base, write_src, kBigLength) != 0) + std::this_thread::yield(); + } + } + + stop = true; + for (auto& t : noise) t.join(); + + EXPECT_EQ(torn_reads, 0u) + << torn_reads << "/" << kIterations + << " reads observed destination bytes not matching the COMPLETED " + "write (worst: " + << torn_bytes_worst << " stale bytes; first at iteration " + << first_bad_iter << "; noise failures: " << noise_failures.load() + << ")"; +} + +// A server-side rejection must surface as FAILED, not silent success: under +// v1 framing a single-chunk WRITE to an unregistered address reported +// COMPLETED because the protocol had no channel for the server to say no. +TEST(TcpWriteVisibilityTest, RejectedWriteMustFail) { + const char* env = std::getenv("MC_METADATA_SERVER"); + std::string metadata_server = env ? env : "P2PHANDSHAKE"; + EngineHandle h; + h.init(metadata_server, "127.0.0.2:17902", 8ull << 20); + ASSERT_TRUE(h.ok) << "engine/segment setup failed"; + + char* src = (char*)h.pool; + memset(src, 0xAB, kSmallLength); + TransferRequest w; + w.opcode = TransferRequest::WRITE; + w.length = kSmallLength; + w.source = src; + w.target_id = h.segment_id; + // One page past the registered pool: the server rejects it in address + // validation. + w.target_offset = h.remote_base + (8ull << 20) + 4096; + EXPECT_EQ(runOne(h.engine.get(), w), TransferStatusEnum::FAILED); + + // The connection that carried the rejected request must not poison + // subsequent transfers. + w.target_offset = h.remote_base + kRegionAlign; + EXPECT_EQ(runOne(h.engine.get(), w), TransferStatusEnum::COMPLETED); +} + +// A v2 server rejects an invalid destination immediately after the header, +// while a large client body is still in flight. FAILED is also the caller's +// source-buffer lifetime boundary, so it must not be published until closing +// the socket has quiesced the outstanding async_write. +TEST(TcpWriteVisibilityTest, LargeRejectedWriteQuiescesSourceBeforeFailure) { + const char* env = std::getenv("MC_METADATA_SERVER"); + std::string metadata_server = env ? env : "P2PHANDSHAKE"; + EngineHandle h; + h.init(metadata_server, "127.0.0.2:17906", 8ull << 20); + ASSERT_TRUE(h.ok) << "engine/segment setup failed"; + + constexpr size_t kLength = 32ull << 20; // exceed the socket buffers + void* source = mmap(nullptr, kLength, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + ASSERT_NE(source, MAP_FAILED); + memset(source, 0x6D, kLength); + + TransferRequest w; + w.opcode = TransferRequest::WRITE; + w.length = kLength; + w.source = source; + w.target_id = h.segment_id; + w.target_offset = h.remote_base + (8ull << 20) + 4096; + + auto start = std::chrono::steady_clock::now(); + EXPECT_EQ(runOne(h.engine.get(), w), TransferStatusEnum::FAILED); + EXPECT_LT(std::chrono::steady_clock::now() - start, + std::chrono::seconds(3)); + + ASSERT_EQ(mprotect(source, kLength, PROT_NONE), 0); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + ASSERT_EQ(mprotect(source, kLength, PROT_READ | PROT_WRITE), 0); + ASSERT_EQ(munmap(source, kLength), 0); + + // The rejected exchange is unclean, so the next request must use a fresh + // connection rather than inheriting the server's mid-session state. + char* small_source = static_cast(h.pool); + memset(small_source, 0x42, kSmallLength); + w.length = kSmallLength; + w.source = small_source; + w.target_offset = h.remote_base + kRegionAlign; + EXPECT_EQ(runOne(h.engine.get(), w), TransferStatusEnum::COMPLETED); +} + +// v2 READ round-trip: exercises the status-frame-then-data framing in both +// directions, including content integrity and a rejected READ surfacing as +// FAILED (v1 could only signal that by dropping the connection). +TEST(TcpWriteVisibilityTest, V2ReadRoundTripAndRejectedRead) { + const char* env = std::getenv("MC_METADATA_SERVER"); + std::string metadata_server = env ? env : "P2PHANDSHAKE"; + EngineHandle h; + h.init(metadata_server, "127.0.0.2:17904", 16ull << 20); + ASSERT_TRUE(h.ok) << "engine/segment setup failed"; + + char* base = (char*)h.pool; + char* src = base + kRegionAlign; + char* dst = base + 2 * kRegionAlign; + for (size_t i = 0; i < kBigLength; ++i) src[i] = (char)(i * 131 + 7); + + TransferRequest w; + w.opcode = TransferRequest::WRITE; + w.length = kBigLength; + w.source = src; + w.target_id = h.segment_id; + w.target_offset = h.remote_base; + ASSERT_EQ(runOne(h.engine.get(), w), TransferStatusEnum::COMPLETED); + + TransferRequest r; + r.opcode = TransferRequest::READ; + r.length = kBigLength; + r.source = dst; + r.target_id = h.segment_id; + r.target_offset = h.remote_base; + ASSERT_EQ(runOne(h.engine.get(), r), TransferStatusEnum::COMPLETED); + // v2 WRITE completion means the destination was already applied, so the + // read-back must match immediately — no drain wait. + EXPECT_EQ(memcmp(dst, src, kBigLength), 0); + + // A READ of an unregistered range must fail via the error status frame. + r.target_offset = h.remote_base + (16ull << 20) + 4096; + EXPECT_EQ(runOne(h.engine.get(), r), TransferStatusEnum::FAILED); + + // And the pool must still be usable afterwards. + r.target_offset = h.remote_base; + EXPECT_EQ(runOne(h.engine.get(), r), TransferStatusEnum::COMPLETED); +} + +// Mixed-version quadrant: a legacy (v1) initiator against the v2 server — +// selected via the MC_TCP_PROTO=1 escape hatch — still transfers data +// correctly (with the old weaker completion semantics). +TEST(TcpWriteVisibilityTest, LegacyInitiatorInteropWithV2Server) { + const char* env = std::getenv("MC_METADATA_SERVER"); + std::string metadata_server = env ? env : "P2PHANDSHAKE"; + // Set the process environment before the engine starts any threads, and + // restore it only after those threads have stopped. POSIX does not require + // setenv()/unsetenv() to synchronize with concurrent getenv() calls. + ScopedEnvVar legacy_proto("MC_TCP_PROTO", "1"); + { + EngineHandle h; + h.init(metadata_server, "127.0.0.2:17903", 16ull << 20); + ASSERT_TRUE(h.ok) << "engine/segment setup failed"; + + char* base = (char*)h.pool; + char* src = base + kRegionAlign; + memset(src, 0x3C, kSmallLength); + TransferRequest w; + w.opcode = TransferRequest::WRITE; + w.length = kSmallLength; + w.source = src; + w.target_id = h.segment_id; + w.target_offset = h.remote_base; + EXPECT_EQ(runOne(h.engine.get(), w), TransferStatusEnum::COMPLETED); + + // v1 completion does not guarantee destination visibility; wait for + // the server drain before checking content. + auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(10); + while (memcmp(base, src, kSmallLength) != 0 && + std::chrono::steady_clock::now() < deadline) + std::this_thread::yield(); + EXPECT_EQ(memcmp(base, src, kSmallLength), 0); + + // Read-back over the transport under v1 framing. + char* dst = base + 2 * kRegionAlign; + TransferRequest r; + r.opcode = TransferRequest::READ; + r.length = kSmallLength; + r.source = dst; + r.target_id = h.segment_id; + r.target_offset = h.remote_base; + EXPECT_EQ(runOne(h.engine.get(), r), TransferStatusEnum::COMPLETED); + EXPECT_EQ(memcmp(dst, src, kSmallLength), 0); + } +} + +// A cached v2 descriptor can briefly outlive a server downgrade/restart. A +// legacy server interprets the flagged WRITE opcode as READ and sends payload +// while the client sends its body. The concurrent status read must break this +// full-duplex deadlock, and FAILED must not become visible until asio has +// released the caller-owned source buffer. +TEST(TcpWriteVisibilityTest, + StaleV2DescriptorAgainstLegacyServerQuiescesWriteBeforeFailure) { + const char* env = std::getenv("MC_METADATA_SERVER"); + std::string metadata_server = env ? env : "P2PHANDSHAKE"; + EngineHandle h; + h.init(metadata_server, "127.0.0.2:17905", 8ull << 20); + ASSERT_TRUE(h.ok) << "engine/segment setup failed"; + + auto desc = h.engine->getMetadata()->getSegmentDescByID(h.segment_id); + ASSERT_NE(desc, nullptr); + + LegacyReadServer legacy_server; + ASSERT_TRUE(legacy_server.ok()); + desc->tcp_data_port = legacy_server.port(); + desc->tcp_proto_version = 2; // deliberately stale capability advertisement + + constexpr size_t kLength = 32ull << 20; // exceed both socket buffers + void* source = mmap(nullptr, kLength, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + ASSERT_NE(source, MAP_FAILED); + memset(source, 0x5C, kLength); + + TransferRequest w; + w.opcode = TransferRequest::WRITE; + w.length = kLength; + w.source = source; + w.target_id = h.segment_id; + w.target_offset = h.remote_base; + + auto start = std::chrono::steady_clock::now(); + EXPECT_EQ(runOne(h.engine.get(), w), TransferStatusEnum::FAILED); + auto elapsed = std::chrono::steady_clock::now() - start; + // The fake legacy peer waits up to 8s for the old mutual-write deadlock. + // Leave generous CI headroom while still proving the concurrent-read path. + EXPECT_LT(elapsed, std::chrono::seconds(3)); + + // Terminal status is the source-buffer lifetime boundary. Protecting the + // pages immediately after FAILED would crash if an async_write still owned + // them and attempted further progress. + ASSERT_EQ(mprotect(source, kLength, PROT_NONE), 0); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + ASSERT_EQ(mprotect(source, kLength, PROT_READ | PROT_WRITE), 0); + ASSERT_EQ(munmap(source, kLength), 0); + + legacy_server.join(); + EXPECT_TRUE(legacy_server.sawFlaggedWrite()); +} + +// A stale v2 descriptor can also point at a legacy server with a request +// SHORTER than a status frame (1-7 bytes). The v1 peer treats the flagged +// opcode as READ, streams fewer than 8 "payload" bytes, and then keeps the +// connection open waiting for the next header — no EOF ever arrives, and for +// WRITE it is symmetrically stuck parsing our body bytes as a partial +// header. Without a status-frame deadline both directions wait forever; +// with it they must fail, and only after actually waiting the deadline out +// (an early failure would mean something else broke). +TEST(TcpWriteVisibilityTest, StaleV2DescriptorShortRequestFailsWithinDeadline) { + ScopedEnvVar fast_deadline("MC_TCP_STATUS_TIMEOUT_SEC", "2"); + const char* env = std::getenv("MC_METADATA_SERVER"); + std::string metadata_server = env ? env : "P2PHANDSHAKE"; + EngineHandle h; + h.init(metadata_server, "127.0.0.2:17906", 8ull << 20); + ASSERT_TRUE(h.ok) << "engine/segment setup failed"; + + auto desc = h.engine->getMetadata()->getSegmentDescByID(h.segment_id); + ASSERT_NE(desc, nullptr); + desc->tcp_proto_version = 2; // deliberately stale capability advertisement + + char buf[4] = {0x11, 0x22, 0x33, 0x44}; + for (auto opcode : {TransferRequest::WRITE, TransferRequest::READ}) { + // One server per direction: the previous connection was (correctly) + // discarded rather than re-pooled, so each request dials anew. + LegacyReadServer legacy_server(/*keep_open=*/true); + ASSERT_TRUE(legacy_server.ok()); + desc->tcp_data_port = legacy_server.port(); + + TransferRequest r; + r.opcode = opcode; + r.length = sizeof(buf); + r.source = buf; + r.target_id = h.segment_id; + r.target_offset = h.remote_base; + + auto start = std::chrono::steady_clock::now(); + EXPECT_EQ(runOne(h.engine.get(), r), TransferStatusEnum::FAILED) + << "opcode " << static_cast(opcode); + auto elapsed = std::chrono::steady_clock::now() - start; + EXPECT_GE(elapsed, std::chrono::seconds(1)) + << "failure arrived before the status deadline could have fired; " + "the wrong path failed (opcode " + << static_cast(opcode) << ")"; + EXPECT_LT(elapsed, std::chrono::seconds(6)) + << "deadline did not bound the stale-descriptor wait (opcode " + << static_cast(opcode) << ")"; + legacy_server.join(); + } +} diff --git a/mooncake-transfer-engine/tests/topology_test.cpp b/mooncake-transfer-engine/tests/topology_test.cpp index f8e38b3455..732fb2984e 100644 --- a/mooncake-transfer-engine/tests/topology_test.cpp +++ b/mooncake-transfer-engine/tests/topology_test.cpp @@ -3,7 +3,11 @@ #include #include #include +#include +#include +#include "config.h" +#include "cuda_alike.h" #include "transfer_metadata.h" #include "memory_location.h" @@ -114,6 +118,79 @@ TEST(ToplogyTest, TestSelectDeviceAny) { ASSERT_TRUE(items.empty()); } +TEST(ToplogyTest, TestSelectDeviceEmptyEntry) { + mooncake::Topology topology; + std::string json_str = "{\"gpu:0\" : [[],[]]}"; + topology.clear(); + ASSERT_EQ(topology.parse(json_str), 0); + + ASSERT_EQ(topology.selectDevice("gpu:0", 0), ERR_DEVICE_NOT_FOUND); + ASSERT_EQ(topology.selectDevice("gpu:0", 1), ERR_DEVICE_NOT_FOUND); +} + +TEST(ToplogyTest, TestDisableOnlyPreferredDeviceLeavesNoSelection) { + mooncake::Topology topology; + std::string json_str = "{\"gpu:0\" : [[\"mlx5_2\"],[]]}"; + topology.clear(); + ASSERT_EQ(topology.parse(json_str), 0); + ASSERT_EQ(topology.selectDevice("gpu:0", 0), 0); + + ASSERT_EQ(topology.disableDevice("mlx5_2"), 0); + ASSERT_EQ(topology.selectDevice("gpu:0", 0), ERR_DEVICE_NOT_FOUND); + ASSERT_EQ(topology.selectDevice("gpu:0", 1), ERR_DEVICE_NOT_FOUND); +} + +TEST(ToplogyTest, TestDisableDeviceRemovesLocalHcaAffinityCandidate) { + mooncake::Topology topology; + std::string json_str = + "{\"cpu:0\" : [[\"mlx5_1\"],[\"mlx5_2\"]]," + "\"cpu:1\" : [[\"mlx5_2\"],[\"mlx5_1\"]]}"; + topology.clear(); + ASSERT_EQ(topology.parse(json_str), 0); + + const auto &hca_list = topology.getHcaList(); + auto disabled_iter = std::find(hca_list.begin(), hca_list.end(), "mlx5_2"); + ASSERT_NE(disabled_iter, hca_list.end()); + const int disabled_index = + static_cast(std::distance(hca_list.begin(), disabled_iter)); + + ASSERT_EQ(topology.disableDevice("mlx5_2"), 0); + ASSERT_NE(topology.selectDeviceByLocalHca("cpu:0", "mlx5_2", 0), + disabled_index); +} + +// HCA peer affinity must key off GPU_PREFIX (cuda:/hip:/...), not a +// hardcoded "cuda:" string — otherwise USE_HIP builds never resolve affinity +// for discovered hip:N topology entries. +TEST(ToplogyTest, HcaPeerAffinityAppliesToGpuPrefixEntries) { + auto &cfg = mooncake::globalConfig(); + const bool old_enable = cfg.enable_hca_peer_affinity; + const auto old_map = cfg.nic_peer_affinity; + cfg.enable_hca_peer_affinity = true; + cfg.nic_peer_affinity = {{"L", {"P0"}}}; + + const std::string gpu_loc = GPU_PREFIX + "0"; + const std::string json_str = "{\"" + gpu_loc + "\" : [[\"P0\",\"P1\"],[]]}"; + + mooncake::Topology topology; + ASSERT_EQ(topology.parse(json_str), 0); + ASSERT_EQ(topology.getHcaList().size(), static_cast(2)); + ASSERT_EQ(topology.getHcaList()[0], "P0"); + + std::unordered_map hist; + for (int i = 0; i < 64; ++i) { + int id = topology.selectDeviceByLocalHca(gpu_loc, "L", 0); + hist[id]++; + } + + cfg.enable_hca_peer_affinity = old_enable; + cfg.nic_peer_affinity = old_map; + + ASSERT_EQ(hist.size(), static_cast(1)); + EXPECT_EQ(hist[0], 64) << "peer affinity should pin " << gpu_loc + << " to P0 for local HCA L"; +} + int main(int argc, char **argv) { ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); diff --git a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp index d8fd0d889c..a890489a3e 100644 --- a/mooncake-transfer-engine/tests/transfer_metadata_test.cpp +++ b/mooncake-transfer-engine/tests/transfer_metadata_test.cpp @@ -17,10 +17,18 @@ #include #include #include +#include +#include #include +#include +#include #include +#include +#include "common.h" +#include "config.h" +#include "transfer_metadata_plugin.h" #include "transport/transport.h" using namespace mooncake; @@ -107,6 +115,75 @@ TEST_F(TransferMetadataTest, LocalMemoryBufferTest) { ASSERT_EQ(re, 0); } +TEST_F(TransferMetadataTest, NcclMetadataAndHandshakePayloadRoundTrip) { + TransferMetadata server(P2PHANDSHAKE); + int sockfd = -1; + const uint16_t port = findAvailableTcpPort(sockfd); + ASSERT_NE(port, 0); + const std::string host = globalConfig().use_ipv6 ? "::1" : "127.0.0.1"; + + auto server_segment = std::make_shared(); + server_segment->name = maybeWrapIpV6(host) + ":" + std::to_string(port); + server_segment->protocol = "nccl"; + TransferMetadata::BufferDesc server_buffer; + server_buffer.name = "cuda:2"; + server_buffer.addr = 0x10000; + server_buffer.length = 4096; + server_buffer.device_id = 2; + server_segment->buffers.push_back(server_buffer); + const std::string server_name = server_segment->name; + TransferMetadata::BufferDesc server_buffer_2; + server_buffer_2.name = "cuda:2-aux"; + server_buffer_2.addr = 0x08000; + server_buffer_2.length = 8192; + server_buffer_2.device_id = 2; + server_segment->buffers.push_back(server_buffer_2); + ASSERT_EQ(server.addLocalSegment(LOCAL_SEGMENT_ID, server_name, + std::move(server_segment)), + 0); + + TransferMetadata::RpcMetaDesc rpc{}; + rpc.ip_or_host_name = host; + rpc.rpc_port = port; + rpc.sockfd = sockfd; + ASSERT_EQ(server.addRpcMetaEntry("nccl-metadata-test", rpc), 0); + ASSERT_EQ(server.startHandshakeDaemon( + [](const TransferMetadata::HandShakeDesc& peer, + TransferMetadata::HandShakeDesc& local) { + local.payload = "reply:" + peer.payload; + return 0; + }, + port, sockfd), + 0); + + TransferMetadata client(P2PHANDSHAKE); + auto client_segment = std::make_shared(); + client_segment->name = "client"; + client_segment->protocol = "nccl"; + ASSERT_EQ(client.addLocalSegment(LOCAL_SEGMENT_ID, "client", + std::move(client_segment)), + 0); + + auto peer_segment = client.getSegmentDesc(server_name); + ASSERT_NE(peer_segment, nullptr); + ASSERT_EQ(peer_segment->protocol, "nccl"); + ASSERT_EQ(peer_segment->buffers.size(), 2U); + EXPECT_EQ(peer_segment->buffers[0].name, "cuda:2"); + EXPECT_EQ(peer_segment->buffers[0].addr, 0x10000U); + EXPECT_EQ(peer_segment->buffers[0].length, 4096U); + EXPECT_EQ(peer_segment->buffers[0].device_id, 2); + EXPECT_EQ(peer_segment->buffers[1].name, "cuda:2-aux"); + EXPECT_EQ(peer_segment->buffers[1].addr, 0x08000U); + EXPECT_EQ(peer_segment->buffers[1].length, 8192U); + EXPECT_EQ(peer_segment->buffers[1].device_id, 2); + + TransferMetadata::HandShakeDesc request; + request.payload = "bootstrap"; + TransferMetadata::HandShakeDesc response; + ASSERT_EQ(client.sendHandshake(server_name, request, response), 0); + EXPECT_EQ(response.payload, "reply:bootstrap"); +} + // add, get and remove RPCMetaEntryMeta TEST_F(TransferMetadataTest, RpcMetaEntryTest) { auto hostname_port = parseHostNameWithPort(local_server_name); @@ -123,9 +200,364 @@ TEST_F(TransferMetadataTest, RpcMetaEntryTest) { ASSERT_EQ(re, 0); } +namespace { + +struct ScopedMetadataRefreshConfig { + uint64_t old_interval_seconds; + bool old_metacache; + + ScopedMetadataRefreshConfig(uint64_t interval_seconds, bool metacache) + : old_interval_seconds( + globalConfig().te_metadata_refresh_interval_seconds), + old_metacache(globalConfig().metacache) { + globalConfig().te_metadata_refresh_interval_seconds = interval_seconds; + globalConfig().metacache = metacache; + } + + ~ScopedMetadataRefreshConfig() { + globalConfig().te_metadata_refresh_interval_seconds = + old_interval_seconds; + globalConfig().metacache = old_metacache; + } +}; + +TransferMetadata::BufferDesc makeRdmaBufferDesc(uint64_t addr) { + TransferMetadata::BufferDesc buffer_desc; + buffer_desc.name = "buffer"; + buffer_desc.addr = addr; + buffer_desc.length = 1024; + buffer_desc.lkey.push_back(1); + buffer_desc.rkey.push_back(2); + return buffer_desc; +} + +std::shared_ptr makeRdmaSegmentDesc( + const std::string& name, uint64_t addr) { + auto segment_desc = std::make_shared(); + segment_desc->name = name; + segment_desc->protocol = "rdma"; + segment_desc->tcp_data_port = 0; + + TransferMetadata::DeviceDesc device_desc; + device_desc.name = "mlx5_0"; + device_desc.lid = 1; + device_desc.gid = "00000000000000000000ffff7f000001"; + segment_desc->devices.push_back(device_desc); + + segment_desc->buffers.push_back(makeRdmaBufferDesc(addr)); + return segment_desc; +} + +} // namespace + +TEST(TransferMetadataPollingTest, PollingRefreshesCachedRemoteSegmentDesc) { + constexpr uint64_t kInitialAddr = 0x1000; + constexpr uint64_t kUpdatedAddr = 0x2000; + + ScopedMetadataRefreshConfig restore(1, true); + TransferMetadata server(P2PHANDSHAKE); + TransferMetadata client(P2PHANDSHAKE); + + int sockfd = -1; + const uint16_t port = findAvailableTcpPort(sockfd); + ASSERT_GT(port, 0); + const std::string remote_segment_name = "127.0.0.1:" + std::to_string(port); + + ASSERT_EQ(server.addLocalSegment( + LOCAL_SEGMENT_ID, remote_segment_name, + makeRdmaSegmentDesc(remote_segment_name, kInitialAddr)), + 0); + TransferMetadata::RpcMetaDesc rpc_desc; + rpc_desc.ip_or_host_name = "127.0.0.1"; + rpc_desc.rpc_port = port; + rpc_desc.sockfd = sockfd; + ASSERT_EQ(server.addRpcMetaEntry(remote_segment_name, rpc_desc), 0); + + ASSERT_EQ( + client.addLocalSegment(LOCAL_SEGMENT_ID, "127.0.0.1:0", + makeRdmaSegmentDesc("127.0.0.1:0", 0x3000)), + 0); + + const auto segment_id = client.getSegmentID(remote_segment_name); + ASSERT_NE(segment_id, static_cast(-1)); + auto cached_desc = client.getSegmentDescByID(segment_id); + ASSERT_TRUE(cached_desc); + ASSERT_EQ(cached_desc->buffers[0].addr, kInitialAddr); + + ASSERT_EQ(server.removeLocalMemoryBuffer( + reinterpret_cast(kInitialAddr), false), + 0); + ASSERT_EQ( + server.addLocalMemoryBuffer(makeRdmaBufferDesc(kUpdatedAddr), false), + 0); + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(3); + while (std::chrono::steady_clock::now() < deadline) { + cached_desc = client.getSegmentDescByID(segment_id); + ASSERT_TRUE(cached_desc); + if (!cached_desc->buffers.empty() && + cached_desc->buffers[0].addr == kUpdatedAddr) { + return; + } + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + } + + FAIL() << "TE metadata refresh polling did not refresh cached descriptor"; +} + +TEST(TransferMetadataPublicationTest, PreservesLocalOnlyBufferWithoutRkey) { + constexpr uint64_t kRemoteAddr = 0x1000; + constexpr uint64_t kLocalOnlyAddr = 0x2000; + + TransferMetadata server(P2PHANDSHAKE); + TransferMetadata client(P2PHANDSHAKE); + + int sockfd = -1; + const uint16_t port = findAvailableTcpPort(sockfd); + ASSERT_GT(port, 0); + const std::string remote_segment_name = "127.0.0.1:" + std::to_string(port); + + auto server_desc = makeRdmaSegmentDesc(remote_segment_name, kRemoteAddr); + auto local_only_buffer = makeRdmaBufferDesc(kLocalOnlyAddr); + local_only_buffer.rkey.clear(); + server_desc->buffers.push_back(local_only_buffer); + ASSERT_EQ(server.addLocalSegment(LOCAL_SEGMENT_ID, remote_segment_name, + std::move(server_desc)), + 0); + + TransferMetadata::RpcMetaDesc rpc_desc; + rpc_desc.ip_or_host_name = "127.0.0.1"; + rpc_desc.rpc_port = port; + rpc_desc.sockfd = sockfd; + ASSERT_EQ(server.addRpcMetaEntry(remote_segment_name, rpc_desc), 0); + + ASSERT_EQ( + client.addLocalSegment(LOCAL_SEGMENT_ID, "127.0.0.1:0", + makeRdmaSegmentDesc("127.0.0.1:0", 0x3000)), + 0); + + const auto segment_id = client.getSegmentID(remote_segment_name); + ASSERT_NE(segment_id, static_cast(-1)); + auto remote_desc = client.getSegmentDescByID(segment_id, true); + ASSERT_NE(remote_desc, nullptr); + ASSERT_EQ(remote_desc->buffers.size(), 2); + EXPECT_EQ(remote_desc->buffers[0].addr, kRemoteAddr); + EXPECT_EQ(remote_desc->buffers[1].addr, kLocalOnlyAddr); + EXPECT_TRUE(remote_desc->buffers[1].rkey.empty()); +} + +// A peer descriptor whose key vector is longer than its device list lets the +// topology-selected device_id pass the rkey bound in selectPeerDevice() and +// still index devices[] out of bounds. Such a descriptor must be rejected at +// decode time. +TEST(TransferMetadataValidationTest, RejectsMoreKeysThanDevices) { + constexpr uint64_t kRemoteAddr = 0x1000; + constexpr size_t kKeyCount = 64; + + TransferMetadata server(P2PHANDSHAKE); + TransferMetadata client(P2PHANDSHAKE); + + int sockfd = -1; + const uint16_t port = findAvailableTcpPort(sockfd); + ASSERT_GT(port, 0); + const std::string remote_segment_name = "127.0.0.1:" + std::to_string(port); + + auto server_desc = makeRdmaSegmentDesc(remote_segment_name, kRemoteAddr); + ASSERT_EQ(server_desc->devices.size(), 1u); + auto& buffer = server_desc->buffers[0]; + while (buffer.rkey.size() < kKeyCount) { + buffer.lkey.push_back(1); + buffer.rkey.push_back(2); + } + ASSERT_EQ(server.addLocalSegment(LOCAL_SEGMENT_ID, remote_segment_name, + std::move(server_desc)), + 0); + + TransferMetadata::RpcMetaDesc rpc_desc; + rpc_desc.ip_or_host_name = "127.0.0.1"; + rpc_desc.rpc_port = port; + rpc_desc.sockfd = sockfd; + ASSERT_EQ(server.addRpcMetaEntry(remote_segment_name, rpc_desc), 0); + + ASSERT_EQ( + client.addLocalSegment(LOCAL_SEGMENT_ID, "127.0.0.1:0", + makeRdmaSegmentDesc("127.0.0.1:0", 0x3000)), + 0); + + EXPECT_EQ(client.getSegmentID(remote_segment_name), + static_cast(-1)); +} + +// The multi-NIC case a real publisher produces: one key per device. It must +// still decode. +TEST(TransferMetadataValidationTest, AcceptsOneKeyPerDevice) { + constexpr uint64_t kRemoteAddr = 0x1000; + constexpr size_t kDeviceCount = 4; + + TransferMetadata server(P2PHANDSHAKE); + TransferMetadata client(P2PHANDSHAKE); + + int sockfd = -1; + const uint16_t port = findAvailableTcpPort(sockfd); + ASSERT_GT(port, 0); + const std::string remote_segment_name = "127.0.0.1:" + std::to_string(port); + + auto server_desc = makeRdmaSegmentDesc(remote_segment_name, kRemoteAddr); + auto& buffer = server_desc->buffers[0]; + while (server_desc->devices.size() < kDeviceCount) { + TransferMetadata::DeviceDesc device_desc; + device_desc.name = + "mlx5_" + std::to_string(server_desc->devices.size()); + device_desc.lid = 1; + device_desc.gid = "00000000000000000000ffff7f000001"; + server_desc->devices.push_back(device_desc); + buffer.lkey.push_back(1); + buffer.rkey.push_back(2); + } + ASSERT_EQ(buffer.rkey.size(), kDeviceCount); + ASSERT_EQ(server.addLocalSegment(LOCAL_SEGMENT_ID, remote_segment_name, + std::move(server_desc)), + 0); + + TransferMetadata::RpcMetaDesc rpc_desc; + rpc_desc.ip_or_host_name = "127.0.0.1"; + rpc_desc.rpc_port = port; + rpc_desc.sockfd = sockfd; + ASSERT_EQ(server.addRpcMetaEntry(remote_segment_name, rpc_desc), 0); + + ASSERT_EQ( + client.addLocalSegment(LOCAL_SEGMENT_ID, "127.0.0.1:0", + makeRdmaSegmentDesc("127.0.0.1:0", 0x3000)), + 0); + + const auto segment_id = client.getSegmentID(remote_segment_name); + ASSERT_NE(segment_id, static_cast(-1)); + auto remote_desc = client.getSegmentDescByID(segment_id, true); + ASSERT_NE(remote_desc, nullptr); + EXPECT_EQ(remote_desc->devices.size(), kDeviceCount); + EXPECT_EQ(remote_desc->buffers[0].rkey.size(), kDeviceCount); +} + +TEST(HandshakeFrameTest, ValidFrameRoundTrips) { + int fds[2]; + ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0); + + ASSERT_EQ(writeString(fds[0], HandShakeRequestType::Metadata, + "{\"name\":\"segment\"}"), + 0); + auto [type, payload] = readString(fds[1]); + EXPECT_EQ(type, HandShakeRequestType::Metadata); + EXPECT_EQ(payload, "{\"name\":\"segment\"}"); + + close(fds[0]); + close(fds[1]); +} + +TEST(HandshakeFrameTest, ValidTypedFrameWithTlsLikeNativeEndianLength) { + int fds[2]; + ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0); + + // 790 is encoded as 0x16 0x03 0x00 ... on little-endian machines, which + // collides with the first two bytes of a TLS ClientHello record. It is + // still a valid native-endian handshake frame length. + const std::string payload(789, 'x'); + ASSERT_EQ(writeString(fds[0], HandShakeRequestType::Metadata, payload), 0); + + auto [type, read_payload] = readString(fds[1]); + EXPECT_EQ(type, HandShakeRequestType::Metadata); + EXPECT_EQ(read_payload, payload); + + close(fds[0]); + close(fds[1]); +} + +TEST(HandshakeFrameTest, OldProtocolFrameStillWorks) { + int fds[2]; + ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0); + + const std::string old_payload = "{\"name\":\"segment\"}"; + uint64_t old_length = old_payload.size(); + ASSERT_EQ(writeFully(fds[0], &old_length, sizeof(old_length)), + static_cast(sizeof(old_length))); + ASSERT_EQ(writeFully(fds[0], old_payload.data(), old_payload.size()), + static_cast(old_payload.size())); + + auto [type, payload] = readString(fds[1]); + EXPECT_EQ(type, HandShakeRequestType::OldProtocol); + EXPECT_EQ(payload, old_payload); + + close(fds[0]); + close(fds[1]); +} + +TEST(HandshakeFrameTest, RejectsHttpProbe) { + int fds[2]; + ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0); + + const std::string request = "GET / HTTP/1.1\r\nHost: localhost\r\n\r\n"; + ASSERT_EQ(writeFully(fds[0], request.data(), request.size()), + static_cast(request.size())); + + auto [type, payload] = readString(fds[1]); + EXPECT_EQ(type, HandShakeRequestType::Invalid); + EXPECT_TRUE(payload.empty()); + + close(fds[0]); + close(fds[1]); +} + +TEST(HandshakeFrameTest, RejectsTlsProbe) { + int fds[2]; + ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0); + + const uint8_t client_hello_prefix[] = {0x16, 0x03, 0x01, 0x05, + 0xc2, 0x01, 0x00, 0x05}; + + ASSERT_EQ( + writeFully(fds[0], client_hello_prefix, sizeof(client_hello_prefix)), + static_cast(sizeof(client_hello_prefix))); + + auto [read_type, payload] = readString(fds[1]); + EXPECT_EQ(read_type, HandShakeRequestType::Invalid); + EXPECT_TRUE(payload.empty()); + + close(fds[0]); + close(fds[1]); +} + +TEST(HandshakeFrameTest, RejectsInvalidLength) { + int fds[2]; + ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0); + + const uint64_t oversized_length = kMaxHandshakeMaxLength + 1; + ASSERT_EQ(writeFully(fds[0], &oversized_length, sizeof(oversized_length)), + static_cast(sizeof(oversized_length))); + + auto [oversized_type, oversized_payload] = readString(fds[1]); + EXPECT_EQ(oversized_type, HandShakeRequestType::Invalid); + EXPECT_TRUE(oversized_payload.empty()); + + close(fds[0]); + close(fds[1]); + + ASSERT_EQ(socketpair(AF_UNIX, SOCK_STREAM, 0, fds), 0); + + const uint64_t zero_length = 0; + ASSERT_EQ(writeFully(fds[0], &zero_length, sizeof(zero_length)), + static_cast(sizeof(zero_length))); + + auto [zero_type, zero_payload] = readString(fds[1]); + EXPECT_EQ(zero_type, HandShakeRequestType::Invalid); + EXPECT_TRUE(zero_payload.empty()); + + close(fds[0]); + close(fds[1]); +} + } // namespace mooncake int main(int argc, char** argv) { ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); -} \ No newline at end of file +} diff --git a/mooncake-transfer-engine/tests/transport_uint_test.cpp b/mooncake-transfer-engine/tests/transport_uint_test.cpp index 599617c6e6..91c3897277 100644 --- a/mooncake-transfer-engine/tests/transport_uint_test.cpp +++ b/mooncake-transfer-engine/tests/transport_uint_test.cpp @@ -17,18 +17,220 @@ #include #include +#include +#include +#include #include #include +#include #include #include +#include +#include #include "transfer_engine.h" +#include "transfer_engine_impl.h" #include "transport/transport.h" using namespace mooncake; namespace mooncake { +class TransferEngineImplTestPeer { + public: + static void replaceTransports(TransferEngineImpl& engine, + std::shared_ptr transport) { + engine.multi_transports_->transport_map_.clear(); + engine.multi_transports_->transport_map_.emplace("blocking", + std::move(transport)); + } + + static void replaceTransports( + TransferEngineImpl& engine, + const std::vector>>& + transports) { + engine.multi_transports_->transport_map_.clear(); + for (const auto& [name, transport] : transports) { + engine.multi_transports_->transport_map_.emplace(name, transport); + } + } + + static AutoDiscoverConfig autoDiscoverConfig( + const TransferEngineImpl& engine) { + return engine.auto_discover_config_; + } + + static std::string autoDiscoverTransport(const TransferEngineImpl& engine) { + return engine.autoDiscoverTransport(); + } + + static void setUseBarex(TransferEngineImpl& engine, bool use_barex) { + engine.use_barex_ = use_barex; + } +}; + +TEST(TransferEngineAutoDiscoverTest, SelectsEfaForEfaProtocol) { + TransferEngineImpl engine(false); + engine.setAutoDiscover({.enabled = true, .protocol = "efa"}); + + const auto config = TransferEngineImplTestPeer::autoDiscoverConfig(engine); + EXPECT_TRUE(config.enabled); + EXPECT_EQ(config.protocol, "efa"); + EXPECT_EQ(TransferEngineImplTestPeer::autoDiscoverTransport(engine), "efa"); +} + +TEST(TransferEngineAutoDiscoverTest, BarexOverrideTakesPrecedence) { + TransferEngineImpl engine(false); + engine.setAutoDiscover({.enabled = true, .protocol = "efa"}); + TransferEngineImplTestPeer::setUseBarex(engine, true); + + EXPECT_EQ(TransferEngineImplTestPeer::autoDiscoverTransport(engine), + "barex"); +} + +TEST(TransferEngineAutoDiscoverTest, BoolSetterPreservesDefaultSelection) { + TransferEngineImpl engine(false); + engine.setAutoDiscover({.enabled = true, .protocol = "efa"}); + engine.setAutoDiscover(true); + + const auto config = TransferEngineImplTestPeer::autoDiscoverConfig(engine); + EXPECT_TRUE(config.enabled); + EXPECT_TRUE(config.protocol.empty()); + EXPECT_EQ(TransferEngineImplTestPeer::autoDiscoverTransport(engine), + "rdma"); +} + +class BatchResultTransport : public Transport { + public: + explicit BatchResultTransport(int unregister_result = 0) + : unregister_result_(unregister_result) {} + + int unregisterBatchCalls() const { return unregister_batch_calls_; } + size_t registeredBufferCount() const { return registered_buffers_.size(); } + void setRegisterResult(int result) { register_result_ = result; } + + Status submitTransfer(BatchID, + const std::vector&) override { + return Status::OK(); + } + + Status getTransferStatus(BatchID, size_t, TransferStatus&) override { + return Status::OK(); + } + + private: + int registerLocalMemory(void*, size_t, const std::string&, bool, + bool) override { + return 0; + } + + int unregisterLocalMemory(void*, bool) override { return 0; } + + int registerLocalMemoryBatch(const std::vector& buffer_list, + const std::string&) override { + if (register_result_) { + if (!buffer_list.empty()) { + registered_buffers_.push_back(buffer_list.front().addr); + } + return register_result_; + } + for (const auto& buffer : buffer_list) { + registered_buffers_.push_back(buffer.addr); + } + return 0; + } + + int unregisterLocalMemoryBatch( + const std::vector& addr_list) override { + ++unregister_batch_calls_; + for (void* addr : addr_list) { + registered_buffers_.erase( + std::remove(registered_buffers_.begin(), + registered_buffers_.end(), addr), + registered_buffers_.end()); + } + return unregister_result_; + } + + const char* getName() const override { return "batch-result"; } + + int register_result_ = 0; + int unregister_result_; + int unregister_batch_calls_ = 0; + std::vector registered_buffers_; +}; + +class BlockingRegistrationTransport : public Transport { + public: + explicit BlockingRegistrationTransport(int first_registration_result = 0) + : first_registration_result_(first_registration_result) {} + + void waitForFirstRegistration() { + std::unique_lock lock(mutex_); + cv_.wait(lock, [this] { return first_registration_started_; }); + } + + void releaseFirstRegistration() { + { + std::lock_guard lock(mutex_); + release_first_registration_ = true; + } + cv_.notify_all(); + } + + int registrationCalls() { + std::lock_guard lock(mutex_); + return registration_calls_; + } + + Status submitTransfer(BatchID, + const std::vector&) override { + return Status::OK(); + } + + Status getTransferStatus(BatchID, size_t, TransferStatus&) override { + return Status::OK(); + } + + private: + int waitOnFirstRegistration() { + std::unique_lock lock(mutex_); + ++registration_calls_; + if (registration_calls_ == 1) { + first_registration_started_ = true; + cv_.notify_all(); + cv_.wait(lock, [this] { return release_first_registration_; }); + return first_registration_result_; + } + return 0; + } + + int registerLocalMemory(void*, size_t, const std::string&, bool, + bool) override { + return waitOnFirstRegistration(); + } + + int unregisterLocalMemory(void*, bool) override { return 0; } + + int registerLocalMemoryBatch(const std::vector&, + const std::string&) override { + return waitOnFirstRegistration(); + } + + int unregisterLocalMemoryBatch(const std::vector&) override { + return 0; + } + + const char* getName() const override { return "blocking"; } + + std::mutex mutex_; + std::condition_variable cv_; + int first_registration_result_; + int registration_calls_ = 0; + bool first_registration_started_ = false; + bool release_first_registration_ = false; +}; + class TransportTest : public ::testing::Test { protected: void SetUp() override { @@ -76,6 +278,46 @@ TEST_F(TransportTest, parseHostNameWithPortTest) { ASSERT_EQ(res.second, 12001); } +TEST_F(TransportTest, TransferTaskDestructorRunsSliceCleanup) { + int cleanup_count = 0; + { + Transport::TransferTask task; + auto* slice = new Transport::Slice(); + slice->source_addr = &cleanup_count; + slice->cleanup_callback = [](Transport::Slice* released) { + auto* count = static_cast(released->source_addr); + ++*count; + }; + task.slice_list.push_back(slice); + } + + EXPECT_EQ(cleanup_count, 1); +} + +TEST_F(TransportTest, SliceCleanupRunsOnceBeforeCacheReuse) { + Transport::ThreadLocalSliceCache cache; + int cleanup_count = 0; + + Transport::Slice* slice = cache.allocate(); + slice->source_addr = &cleanup_count; + slice->cleanup_callback = [](Transport::Slice* released) { + auto* count = static_cast(released->source_addr); + ++*count; + }; + + cache.deallocate(slice); + EXPECT_EQ(cleanup_count, 1); + + Transport::Slice* reused = cache.allocate(); + EXPECT_EQ(reused, slice); + EXPECT_EQ(reused->cleanup_callback, nullptr); + + // A backend that does not install a callback must not inherit the callback + // from the previous owner of this cached slice. + cache.deallocate(reused); + EXPECT_EQ(cleanup_count, 1); +} + TEST_F(TransportTest, WriteSuccess) { int fd = CreateTempFile(); ASSERT_NE(fd, -1) << "Failed to create temporary file"; @@ -171,9 +413,203 @@ TEST_F(TransportTest, ReadEmptyFile) { close(fd); } + +TEST_F(TransportTest, RegisterLocalMemoryBatchRejectsOverlappingBuffers) { + TransferEngine engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + + std::array buffer{}; + std::vector entries = { + {buffer.data() + 64, 128}, + {buffer.data(), 128}, + }; + + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), + ERR_ADDRESS_OVERLAPPED); +} + +TEST_F(TransportTest, RegisterLocalMemoryBatchRejectsZeroLengthBuffer) { + TransferEngine engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + + std::array buffer{}; + std::vector entries = { + {buffer.data(), 0}, + }; + + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), + ERR_INVALID_ARGUMENT); +} + +TEST_F(TransportTest, RegisterLocalMemoryBatchAllowsAdjacentBuffers) { + TransferEngine engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + + std::array buffer{}; + std::vector entries = { + {buffer.data() + 128, 128}, + {buffer.data(), 128}, + }; + + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), 0); +} + +TEST_F(TransportTest, ConcurrentRegisterLocalMemoryRejectsOverlap) { + TransferEngineImpl engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + auto transport = std::make_shared(); + TransferEngineImplTestPeer::replaceTransports(engine, transport); + + std::array buffer{}; + auto first = std::async(std::launch::async, [&] { + return engine.registerLocalMemory(buffer.data(), buffer.size(), + "cpu:0"); + }); + transport->waitForFirstRegistration(); + + int second = + engine.registerLocalMemory(buffer.data(), buffer.size(), "cpu:0"); + int registration_calls = transport->registrationCalls(); + transport->releaseFirstRegistration(); + + EXPECT_EQ(second, ERR_ADDRESS_OVERLAPPED); + EXPECT_EQ(registration_calls, 1); + EXPECT_EQ(first.get(), 0); + EXPECT_EQ(engine.unregisterLocalMemory(buffer.data()), 0); +} + +TEST_F(TransportTest, ConcurrentRegisterLocalMemoryBatchRejectsOverlap) { + TransferEngineImpl engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + auto transport = std::make_shared(); + TransferEngineImplTestPeer::replaceTransports(engine, transport); + + std::array buffer{}; + std::vector entries = {{buffer.data(), buffer.size()}}; + auto first = std::async(std::launch::async, [&] { + return engine.registerLocalMemoryBatch(entries, "cpu:0"); + }); + transport->waitForFirstRegistration(); + + int second = engine.registerLocalMemoryBatch(entries, "cpu:0"); + int registration_calls = transport->registrationCalls(); + transport->releaseFirstRegistration(); + + EXPECT_EQ(second, ERR_ADDRESS_OVERLAPPED); + EXPECT_EQ(registration_calls, 1); + EXPECT_EQ(first.get(), 0); + EXPECT_EQ(engine.unregisterLocalMemoryBatch({buffer.data()}), 0); +} + +TEST_F(TransportTest, FailedRegistrationReleasesReservedRegion) { + TransferEngineImpl engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + auto transport = + std::make_shared(ERR_MEMORY); + TransferEngineImplTestPeer::replaceTransports(engine, transport); + + std::array buffer{}; + auto first = std::async(std::launch::async, [&] { + return engine.registerLocalMemory(buffer.data(), buffer.size(), + "cpu:0"); + }); + transport->waitForFirstRegistration(); + transport->releaseFirstRegistration(); + + EXPECT_EQ(first.get(), ERR_MEMORY); + EXPECT_EQ(engine.registerLocalMemory(buffer.data(), buffer.size(), "cpu:0"), + 0); + EXPECT_EQ(engine.unregisterLocalMemory(buffer.data()), 0); +} + +TEST_F(TransportTest, UnregisterLocalMemoryBatchPropagatesTransportError) { + TransferEngine engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + ASSERT_NE(engine.installTransport("tcp", nullptr), nullptr); + + std::array buffer{}; + EXPECT_EQ(engine.unregisterLocalMemoryBatch({buffer.data()}), + ERR_ADDRESS_NOT_REGISTERED); +} + +TEST_F(TransportTest, UnregisterLocalMemoryBatchContinuesAcrossTransports) { + TransferEngineImpl engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + auto failing = std::make_shared(ERR_MEMORY); + auto succeeding = std::make_shared(); + TransferEngineImplTestPeer::replaceTransports( + engine, {{"a-failing", failing}, {"b-succeeding", succeeding}}); + + std::array buffer{}; + EXPECT_EQ(engine.unregisterLocalMemoryBatch({buffer.data()}), ERR_MEMORY); + EXPECT_EQ(failing->unregisterBatchCalls(), 1); + EXPECT_EQ(succeeding->unregisterBatchCalls(), 1); +} + +TEST_F(TransportTest, UnregisterLocalMemoryBatchContinuesAfterAddressError) { + TransferEngineImpl engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + ASSERT_NE(engine.installTransport("tcp", nullptr), nullptr); + + std::array registered{}; + std::array missing{}; + std::vector entries = { + {registered.data(), 1}, + {registered.data() + 1, 1}, + }; + ASSERT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), 0); + + auto metadata = engine.getMetadata(); + ASSERT_NE(metadata, nullptr); + auto contains_buffer = [&](void* addr) { + auto desc = metadata->getSegmentDescByID(LOCAL_SEGMENT_ID); + if (!desc) return false; + auto value = reinterpret_cast(addr); + return std::any_of( + desc->buffers.begin(), desc->buffers.end(), + [value](const auto& buffer) { return buffer.addr == value; }); + }; + ASSERT_TRUE(contains_buffer(registered.data())); + ASSERT_TRUE(contains_buffer(registered.data() + 1)); + + EXPECT_EQ(engine.unregisterLocalMemoryBatch( + {missing.data(), registered.data(), registered.data() + 1}), + ERR_ADDRESS_NOT_REGISTERED); + EXPECT_FALSE(contains_buffer(registered.data())); + EXPECT_FALSE(contains_buffer(registered.data() + 1)); +} + +TEST_F(TransportTest, RegisterLocalMemoryBatchRollsBackAttemptedTransports) { + TransferEngineImpl engine(false); + ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0); + auto succeeding = std::make_shared(); + auto failing = std::make_shared(); + failing->setRegisterResult(ERR_MEMORY); + TransferEngineImplTestPeer::replaceTransports( + engine, {{"a-succeeding", succeeding}, {"b-failing", failing}}); + + std::array buffer{}; + std::vector entries = { + {buffer.data(), 1}, + {buffer.data() + 1, 1}, + }; + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), ERR_MEMORY); + EXPECT_EQ(succeeding->registeredBufferCount(), 0); + EXPECT_EQ(failing->registeredBufferCount(), 0); + EXPECT_EQ(succeeding->unregisterBatchCalls(), 1); + EXPECT_EQ(failing->unregisterBatchCalls(), 1); + + failing->setRegisterResult(0); + EXPECT_EQ(engine.registerLocalMemoryBatch(entries, "cpu:0"), 0); + EXPECT_EQ(succeeding->registeredBufferCount(), entries.size()); + EXPECT_EQ(failing->registeredBufferCount(), entries.size()); + EXPECT_EQ( + engine.unregisterLocalMemoryBatch({buffer.data(), buffer.data() + 1}), + 0); +} } // namespace mooncake int main(int argc, char** argv) { ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); -} \ No newline at end of file +} diff --git a/mooncake-transfer-engine/ubshmem-allocator/CMakeLists.txt b/mooncake-transfer-engine/ubshmem-allocator/CMakeLists.txt index f6389f745a..dc5a880a55 100644 --- a/mooncake-transfer-engine/ubshmem-allocator/CMakeLists.txt +++ b/mooncake-transfer-engine/ubshmem-allocator/CMakeLists.txt @@ -10,6 +10,13 @@ add_fabric_allocator_build_target( build_ubshmem_allocator BUILD_SCRIPT ${CMAKE_CURRENT_SOURCE_DIR}/build.sh + OUTPUT_NAME + ubshmem_fabric_allocator.so + BUILD_DEPENDS + ${CMAKE_CURRENT_SOURCE_DIR}/ubshmem_fabric_allocator.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/../include/cuda_alike.h + ${CMAKE_CURRENT_SOURCE_DIR}/../include/gpu_vendor/ubshmem.h + ${CMAKE_CURRENT_SOURCE_DIR}/../scripts/allocator_build_common.sh COMMENT "Building ubshmem allocator to ${CMAKE_CURRENT_BINARY_DIR}" ENABLE_BUILD diff --git a/mooncake-wheel/mooncake/_fast_copy.c b/mooncake-wheel/mooncake/_fast_copy.c new file mode 100644 index 0000000000..017d5a5996 --- /dev/null +++ b/mooncake-wheel/mooncake/_fast_copy.c @@ -0,0 +1,195 @@ +#define PY_SSIZE_T_CLEAN +#include +#include +#include +#include +#include + +#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION +#include + +typedef struct { + void **src_ptrs; + size_t *src_sizes; + int count; + char *dst; + size_t offset; + size_t bytes_copied; +} ThreadWork; + +static void *copy_thread_func(void *arg) { + ThreadWork *w = (ThreadWork *)arg; + char *d = w->dst + w->offset; + size_t off = 0; + for (int i = 0; i < w->count; i++) { + if (w->src_sizes[i] > 0) { + memcpy(d + off, w->src_ptrs[i], w->src_sizes[i]); + off += w->src_sizes[i]; + } + } + w->bytes_copied = off; + return NULL; +} + +static PyObject *concat_arrays_into(PyObject *self, PyObject *args) { + PyObject *list_obj; + unsigned long long dest_ptr_val; + unsigned long long dest_size_val; + Py_ssize_t start = 0; + Py_ssize_t count = -1; + int nthreads = 1; + ThreadWork *works = NULL; + pthread_t *threads = NULL; + char *thread_started = NULL; + + if (!PyArg_ParseTuple(args, "O!KK|nni", &PyList_Type, &list_obj, + &dest_ptr_val, &dest_size_val, &start, &count, + &nthreads)) + return NULL; + + Py_ssize_t list_len = PyList_GET_SIZE(list_obj); + if (start < 0) start = 0; + if (start > list_len) start = list_len; + if (count < 0 || start + count > list_len) count = list_len - start; + if (count == 0) return PyLong_FromSize_t(0); + if (nthreads < 1) nthreads = 1; + if (nthreads > (int)count) nthreads = (int)count; + + size_t dest_size = (size_t)dest_size_val; + void **ptrs = (void **)malloc(count * sizeof(void *)); + size_t *sizes = (size_t *)malloc(count * sizeof(size_t)); + PyObject **items = (PyObject **)calloc(count, sizeof(PyObject *)); + if (!ptrs || !sizes || !items) { + free(ptrs); + free(sizes); + free(items); + return PyErr_NoMemory(); + } + + Py_ssize_t held = 0; + size_t total_input_bytes = 0; + for (Py_ssize_t i = 0; i < count; i++) { + PyObject *item = PyList_GET_ITEM(list_obj, start + i); + if (!PyArray_Check(item)) { + PyErr_Format(PyExc_TypeError, "arrays[%zd] is not an ndarray", + start + i); + goto fail; + } + Py_INCREF(item); + items[held++] = item; + PyArrayObject *arr = (PyArrayObject *)item; + if (!PyArray_IS_C_CONTIGUOUS(arr)) { + PyErr_Format(PyExc_ValueError, "arrays[%zd] is not C-contiguous", + start + i); + goto fail; + } + size_t nbytes = (size_t)PyArray_NBYTES(arr); + if (nbytes > SIZE_MAX - total_input_bytes) { + PyErr_SetString(PyExc_OverflowError, "array byte sizes overflow"); + goto fail; + } + ptrs[i] = PyArray_DATA(arr); + sizes[i] = nbytes; + total_input_bytes += nbytes; + } + if (total_input_bytes > dest_size) { + PyErr_Format(PyExc_ValueError, + "destination buffer too small: need %zu bytes, got %zu", + total_input_bytes, dest_size); + goto fail; + } + + /* Partition work across threads. */ + works = (ThreadWork *)calloc(nthreads, sizeof(ThreadWork)); + threads = (pthread_t *)malloc(nthreads * sizeof(pthread_t)); + thread_started = (char *)calloc(nthreads, sizeof(char)); + if (!works || !threads || !thread_started) { + PyErr_NoMemory(); + goto fail; + } + + int n = (int)count; + int per_t = (n + nthreads - 1) / nthreads; + size_t offset = 0; + int actual_threads = 0; + for (int t = 0; t < nthreads; t++) { + int s = t * per_t; + int c = per_t; + if (s + c > n) c = n - s; + if (c <= 0) break; + + works[t].src_ptrs = ptrs + s; + works[t].src_sizes = sizes + s; + works[t].count = c; + works[t].dst = (char *)(uintptr_t)dest_ptr_val; + works[t].offset = offset; + + size_t tb = 0; + for (int j = s; j < s + c; j++) tb += sizes[j]; + offset += tb; + actual_threads++; + } + + Py_BEGIN_ALLOW_THREADS; + if (actual_threads == 1) { + copy_thread_func(&works[0]); + } else { + for (int t = 1; t < actual_threads; t++) { + if (pthread_create(&threads[t], NULL, copy_thread_func, + &works[t]) != 0) { + /* Copy inline if worker creation fails after other workers + * start. */ + copy_thread_func(&works[t]); + } else { + thread_started[t] = 1; + } + } + copy_thread_func(&works[0]); + for (int t = 1; t < actual_threads; t++) { + if (thread_started[t]) { + pthread_join(threads[t], NULL); + } + } + } + Py_END_ALLOW_THREADS; + + size_t total = 0; + for (int t = 0; t < actual_threads; t++) total += works[t].bytes_copied; + + free(ptrs); + free(sizes); + for (Py_ssize_t i = 0; i < held; i++) Py_DECREF(items[i]); + free(items); + free(works); + free(threads); + free(thread_started); + return PyLong_FromSize_t(total); + +fail: + free(ptrs); + free(sizes); + for (Py_ssize_t i = 0; i < held; i++) Py_DECREF(items[i]); + free(items); + free(works); + free(threads); + free(thread_started); + return NULL; +} + +static PyMethodDef module_methods[] = { + {"concat_arrays_into", concat_arrays_into, METH_VARARGS, + "Scatter-copy arrays[start:start+count] into dest_ptr (GIL released)."}, + {NULL, NULL, 0, NULL}}; + +static struct PyModuleDef moduledef = { + PyModuleDef_HEAD_INIT, + "_fast_copy", + "Fast scatter-gather copy for ndarray lists.", + -1, + module_methods, +}; + +PyMODINIT_FUNC PyInit__fast_copy(void) { + import_array(); + return PyModule_Create(&moduledef); +} diff --git a/mooncake-wheel/mooncake/cli.py b/mooncake-wheel/mooncake/cli.py index e252d5a313..2fba48a38f 100644 --- a/mooncake-wheel/mooncake/cli.py +++ b/mooncake-wheel/mooncake/cli.py @@ -4,6 +4,7 @@ """ import os +import stat import sys import subprocess @@ -16,10 +17,12 @@ def main(): # Get the path to the mooncake_master binary package_dir = os.path.dirname(os.path.abspath(__file__)) bin_path = os.path.join(package_dir, "mooncake_master") - + # Make sure the binary is executable - os.chmod(bin_path, 0o755) - + if not os.access(bin_path, os.X_OK): + st = os.stat(bin_path) + os.chmod(bin_path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + # Run the binary with all arguments passed through return subprocess.call([bin_path] + sys.argv[1:]) diff --git a/mooncake-wheel/mooncake/cli_bench.py b/mooncake-wheel/mooncake/cli_bench.py index 2eb5ea986b..f4bf27d766 100644 --- a/mooncake-wheel/mooncake/cli_bench.py +++ b/mooncake-wheel/mooncake/cli_bench.py @@ -4,6 +4,7 @@ """ import os +import stat import sys import subprocess @@ -16,10 +17,12 @@ def main(): # Get the path to the transfer_engine_bench binary package_dir = os.path.dirname(os.path.abspath(__file__)) bin_path = os.path.join(package_dir, "transfer_engine_bench") - + # Make sure the binary is executable - os.chmod(bin_path, 0o755) - + if not os.access(bin_path, os.X_OK): + st = os.stat(bin_path) + os.chmod(bin_path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + # Run the binary with all arguments passed through return subprocess.call([bin_path] + sys.argv[1:]) diff --git a/mooncake-wheel/mooncake/cli_client.py b/mooncake-wheel/mooncake/cli_client.py index c32cefebe9..0cf20ec717 100644 --- a/mooncake-wheel/mooncake/cli_client.py +++ b/mooncake-wheel/mooncake/cli_client.py @@ -4,6 +4,7 @@ """ import os +import stat import sys import subprocess @@ -18,7 +19,9 @@ def main(): bin_path = os.path.join(package_dir, "mooncake_client") # Make sure the binary is executable - os.chmod(bin_path, 0o755) + if not os.access(bin_path, os.X_OK): + st = os.stat(bin_path) + os.chmod(bin_path, st.st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) # Run the binary with all arguments passed through return subprocess.call([bin_path] + sys.argv[1:]) diff --git a/mooncake-wheel/mooncake/http_metadata_server.py b/mooncake-wheel/mooncake/http_metadata_server.py index cf3cf21d43..cc98eba563 100644 --- a/mooncake-wheel/mooncake/http_metadata_server.py +++ b/mooncake-wheel/mooncake/http_metadata_server.py @@ -9,7 +9,6 @@ import argparse import asyncio import logging -import os import signal import sys import threading @@ -61,7 +60,10 @@ def _setup_routes(self): async def _handle_metadata(self, request: web.Request): """Handle metadata requests.""" - key = request.query.get('key', '') + key = request.query.get('key', '').strip() + if not key: + return web.Response(text='metadata key is required', status=400, + content_type='application/json') if request.method == 'GET': return await self._handle_get(key) diff --git a/mooncake-wheel/mooncake/mooncake_config.py b/mooncake-wheel/mooncake/mooncake_config.py index 642557cc9f..9b494ebe42 100644 --- a/mooncake-wheel/mooncake/mooncake_config.py +++ b/mooncake-wheel/mooncake/mooncake_config.py @@ -20,6 +20,7 @@ Transfer Engine C++ Level (Advanced): ------------------------------------- In addition to tcp and rdma, the C++ Transfer Engine also supports: +- efa: AWS Elastic Fabric Adapter transport (libfabric-based) - nvmeof: NVMe over Fabric for direct NVMe storage access - nvlink: NVIDIA NVLink for inter-GPU communication across nodes - nvlink_intra: NVIDIA NVLink for intra-node GPU communication @@ -28,6 +29,19 @@ - cxl: Compute Express Link for memory pooling and sharing - ascend: Huawei Ascend NPU communication (HCCL and direct transport) +Store Surface (mooncake-store): +------------------------------- +MooncakeConfig also drives the Mooncake Store, which additionally accepts: +- ub / ubshmem: Unified Bus transport and its shared-memory variant +- maca: MetaX MACA GPU transport +- sunrise_link: SunriseLink interconnect transport +- rpc_only: Store-only mode with no Transfer Engine attached + +Protocol names are matched case-sensitively by the C++ engine, so the value is +normalised to lowercase when a MooncakeConfig is constructed (e.g. "RDMA" -> +"rdma"). Because the authoritative, build-flag-dependent set lives in C++, an +unrecognised protocol is passed through with a warning rather than rejected. + For most use cases, 'tcp' or 'rdma' is recommended. The default is 'tcp'. For RDMA, you also need to specify the device_name (e.g., 'mlx5_0', 'erdma_0') or use auto-discovery. @@ -45,26 +59,68 @@ export MOONCAKE_PROTOCOL="rdma" export MOONCAKE_DEVICE="auto-discovery" """ + import json +import logging import os from dataclasses import dataclass from typing import Optional +logger = logging.getLogger(__name__) + DEFAULT_GLOBAL_SEGMENT_SIZE = 3355443200 # 3.125 GiB DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 # 1.0 GiB _SIZE_SUFFIXES = [ ("kb", 1024), - ("mb", 1024 ** 2), - ("gb", 1024 ** 3), - ("tb", 1024 ** 4), + ("mb", 1024**2), + ("gb", 1024**3), + ("tb", 1024**4), ("k", 1024), - ("m", 1024 ** 2), - ("g", 1024 ** 3), - ("t", 1024 ** 4), + ("m", 1024**2), + ("g", 1024**3), + ("t", 1024**4), ("b", 1), ] +# Protocols Mooncake is known to accept. This is a *diagnostic hint*, not an +# authoritative gate: MooncakeConfig drives both the Transfer Engine and the +# Store, and the C++ layer is the source of truth for what a given build +# actually supports (several transports are gated behind USE_* build flags). +# The C++ comparisons are exact and lowercase (multi_transport.cpp +# installTransport, mooncake-store client_service.cpp), so the protocol is +# canonicalised to lowercase below and an unrecognised value only warns. +# Union of Transfer Engine transports and Store-only modes. Keep in sync with +# the protocol list documented in the module docstring above. +_KNOWN_PROTOCOLS = frozenset( + { + # Transfer Engine transports (mooncake-transfer-engine installTransport) + "tcp", + "rdma", + "efa", + "nvmeof", + "nvlink", + "nvlink_intra", + "hip", + "barex", + "cxl", + "ascend", + "ub", + "ubshmem", + "maca", + "sunrise_link", + # Store-only mode (mooncake-store client_service.cpp: no transfer engine) + "rpc_only", + } +) + +# Required fields that must be present AND non-empty. +_REQUIRED_NON_EMPTY_FIELDS = ( + "local_hostname", + "metadata_server", + "master_server_address", +) + def _parse_segment_size(value) -> int: if isinstance(value, int): @@ -85,6 +141,31 @@ def _parse_segment_size(value) -> int: return int(value) +def _parse_bool(value) -> bool: + """Interpret a config boolean that may arrive as a real bool, number, or string. + + Config files and environment variables sometimes carry booleans as strings + ("true", "false", "yes", "no", "on", "off", "0", "1"). Using bool() directly + would treat any non-empty string (including "false") as True, so parse the + common textual forms and reject anything unrecognized instead of silently + defaulting to False, mirroring how _parse_segment_size rejects bad input. + """ + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + if value is None: + return False + s = str(value).strip().lower() + if not s: + return False + if s in ("true", "1", "yes", "on", "enable"): + return True + if s in ("false", "0", "no", "off", "disable"): + return False + raise ValueError(f"Invalid boolean value: {value!r}") + + @dataclass class MooncakeConfig: """The configuration class for Mooncake. @@ -94,16 +175,25 @@ class MooncakeConfig: metadata_server (str): The address of the metadata server. global_segment_size (int): The size of each global segment in bytes. local_buffer_size (int): The size of the local buffer in bytes. - protocol (str): The communication protocol to use. Supported values: + protocol (str): The communication protocol to use. Common values: - "tcp" (default): Standard TCP/IP protocol - "rdma": RDMA protocol (requires RDMA-capable NICs and device_name) - See module docstring for full list of supported protocols. + The value is normalised to lowercase (the C++ engine matches + protocol names case-sensitively). The Transfer Engine and Store + together accept a wider set; see the module docstring. An + unrecognised value is passed through to the engine with a warning, + not rejected. device_name (Optional[str]): The name of the RDMA device to use (e.g., "mlx5_0", "erdma_0", or "auto-discovery"). Required when protocol is "rdma", optional for other protocols. master_server_address (str): The address of the master server. enable_ssd_offload (bool): Enable SSD offload. Default is False. ssd_offload_path (str): The path to the SSD directory for offloading. + tenant_id (str): Tenant identifier. Default is "default". + enable_client_http_server (bool): Enable the client HTTP health/metrics + endpoints. Default is False. + client_http_port (int): Port for the client HTTP endpoints. + Defaults to 9300. Example of configuration file: { @@ -115,9 +205,12 @@ class MooncakeConfig: "device_name": "", "master_server_address": "localhost:8081", "enable_ssd_offload": true, - "ssd_offload_path": "/nvme/mooncake_offload" + "ssd_offload_path": "/nvme/mooncake_offload", + "tenant_id": "default", + "enable_client_http_server": false, + "client_http_port": 9300 } - + For RDMA: { "local_hostname": "node1", @@ -128,9 +221,13 @@ class MooncakeConfig: "device_name": "mlx5_0", "master_server_address": "master:8081", "enable_ssd_offload": true, - "ssd_offload_path": "/nvme/mooncake_offload" + "ssd_offload_path": "/nvme/mooncake_offload", + "tenant_id": "default", + "enable_client_http_server": false, + "client_http_port": 9300 } """ + local_hostname: str metadata_server: str global_segment_size: int @@ -140,9 +237,62 @@ class MooncakeConfig: master_server_address: str enable_ssd_offload: bool = False ssd_offload_path: str = "" + tenant_id: str = "default" + enable_client_http_server: bool = False + client_http_port: int = 9300 + + def __post_init__(self): + """Validate and normalise configuration invariants. + + This runs for every ``MooncakeConfig`` instance regardless of how it is + constructed (``from_file``, ``load_from_env`` or direct instantiation). + The protocol is canonicalised to lowercase (the C++ engine is + case-sensitive) and an unrecognised protocol is warned about but passed + through, because the authoritative set is decided in the + build-flag-dependent C++ layer. Genuine structural problems (a + non-string or empty protocol, a negative size, an empty required field) + are still reported here with an actionable message instead of surfacing + as a cryptic failure deep inside the C++ engine. + """ + if not isinstance(self.protocol, str) or not self.protocol.strip(): + raise ValueError( + f"Invalid protocol: {self.protocol!r}. Protocol must be a " + f"non-empty string, e.g. 'tcp' or 'rdma'." + ) + # Canonicalise to lowercase. The C++ engine matches protocol names + # case-sensitively against lowercase literals, so e.g. "RDMA" would be + # silently accepted here but rejected deep in the engine; normalising + # turns that hidden misconfiguration into a working configuration. + self.protocol = self.protocol.strip().lower() + # Warn (do not reject) on values we do not recognise: MooncakeConfig + # drives both Transfer Engine and Store paths, and a given build may + # support protocols beyond this list. The C++ layer is authoritative and + # will reject a genuinely unsupported protocol with its own error. + if self.protocol not in _KNOWN_PROTOCOLS: + logger.warning( + "Unrecognised protocol %r; passing it through to the Mooncake " + "engine unchanged. Known protocols are: %s.", + self.protocol, + ", ".join(sorted(_KNOWN_PROTOCOLS)), + ) + + for field_name in ("global_segment_size", "local_buffer_size"): + value = getattr(self, field_name) + if value < 0: + raise ValueError( + f"Invalid {field_name}: {value}. Size must be non-negative." + ) + + for field_name in _REQUIRED_NON_EMPTY_FIELDS: + value = getattr(self, field_name) + if not isinstance(value, str) or not value.strip(): + raise ValueError( + f"Config field {field_name!r} must be a non-empty string, " + f"got {value!r}." + ) @staticmethod - def from_file(file_path: str) -> 'MooncakeConfig': + def from_file(file_path: str) -> "MooncakeConfig": """Load the config from a JSON file.""" with open(file_path) as fin: config = json.load(fin) @@ -154,6 +304,8 @@ def from_file(file_path: str) -> 'MooncakeConfig': for field in required_fields: if field not in config: raise ValueError(f"Missing required config field: {field}") + ssd_offload_path = config.get("ssd_offload_path") + tenant_id = config.get("tenant_id") return MooncakeConfig( local_hostname=config.get("local_hostname"), metadata_server=config.get("metadata_server"), @@ -166,27 +318,40 @@ def from_file(file_path: str) -> 'MooncakeConfig': protocol=config.get("protocol", "tcp"), device_name=config.get("device_name", ""), master_server_address=config.get("master_server_address"), - enable_ssd_offload=bool(config.get("enable_ssd_offload", False)), - ssd_offload_path=str(config.get("ssd_offload_path", "")), + enable_ssd_offload=_parse_bool(config.get("enable_ssd_offload", False)), + ssd_offload_path=str(ssd_offload_path) + if ssd_offload_path is not None + else "", + tenant_id=str(tenant_id) if tenant_id is not None else "default", + enable_client_http_server=_parse_bool( + config.get("enable_client_http_server", False) + ), + client_http_port=int(config.get("client_http_port", 9300)), ) @staticmethod - def load_from_env() -> 'MooncakeConfig': + def load_from_env() -> "MooncakeConfig": """Load config from a file specified in the environment variable. export MOONCAKE_MASTER=10.13.3.232:50051 export MOONCAKE_PROTOCOL="rdma" export MOONCAKE_DEVICE="" export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE" """ - config_file_path = os.getenv('MOONCAKE_CONFIG_PATH') + config_file_path = os.getenv("MOONCAKE_CONFIG_PATH") if config_file_path is None: if not os.getenv("MOONCAKE_MASTER"): - raise ValueError("Neither the environment variable 'MOONCAKE_CONFIG_PATH' nor 'MOONCAKE_MASTER' is set.") + raise ValueError( + "Neither the environment variable 'MOONCAKE_CONFIG_PATH' nor 'MOONCAKE_MASTER' is set." + ) return MooncakeConfig( local_hostname=os.getenv("MOONCAKE_LOCAL_HOSTNAME", "localhost"), - metadata_server=os.getenv("MOONCAKE_TE_META_DATA_SERVER", "P2PHANDSHAKE"), + metadata_server=os.getenv( + "MOONCAKE_TE_META_DATA_SERVER", "P2PHANDSHAKE" + ), global_segment_size=_parse_segment_size( - os.getenv("MOONCAKE_GLOBAL_SEGMENT_SIZE", DEFAULT_GLOBAL_SEGMENT_SIZE) + os.getenv( + "MOONCAKE_GLOBAL_SEGMENT_SIZE", DEFAULT_GLOBAL_SEGMENT_SIZE + ) ), local_buffer_size=_parse_segment_size( os.getenv("MOONCAKE_LOCAL_BUFFER_SIZE", DEFAULT_LOCAL_BUFFER_SIZE) @@ -194,7 +359,14 @@ def load_from_env() -> 'MooncakeConfig': protocol=os.getenv("MOONCAKE_PROTOCOL", "tcp"), device_name=os.getenv("MOONCAKE_DEVICE", ""), master_server_address=os.getenv("MOONCAKE_MASTER"), - enable_ssd_offload=os.getenv("MOONCAKE_OFFLOAD_ENABLED", "false").lower() in ("true", "1"), + enable_ssd_offload=_parse_bool( + os.getenv("MOONCAKE_OFFLOAD_ENABLED", "false") + ), ssd_offload_path=os.getenv("MOONCAKE_OFFLOAD_FILE_STORAGE_PATH", ""), + tenant_id=os.getenv("MOONCAKE_TENANT_ID", "default"), + enable_client_http_server=_parse_bool( + os.getenv("MOONCAKE_ENABLE_CLIENT_HTTP_SERVER", "false") + ), + client_http_port=int(os.getenv("MOONCAKE_CLIENT_HTTP_PORT", 9300)), ) - return MooncakeConfig.from_file(config_file_path) \ No newline at end of file + return MooncakeConfig.from_file(config_file_path) diff --git a/mooncake-wheel/mooncake/mooncake_connector_v1.py b/mooncake-wheel/mooncake/mooncake_connector_v1.py index 18873b38fc..f3b9a51fd4 100644 --- a/mooncake-wheel/mooncake/mooncake_connector_v1.py +++ b/mooncake-wheel/mooncake/mooncake_connector_v1.py @@ -15,7 +15,6 @@ from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from queue import Queue from os import getenv from typing import TYPE_CHECKING, Any, Optional diff --git a/mooncake-wheel/mooncake/mooncake_elastic_buffer.py b/mooncake-wheel/mooncake/mooncake_elastic_buffer.py new file mode 100644 index 0000000000..df751aef59 --- /dev/null +++ b/mooncake-wheel/mooncake/mooncake_elastic_buffer.py @@ -0,0 +1,582 @@ +import os +import warnings +from typing import Any, List, Optional, Tuple, Union + +import torch +import torch.distributed as dist + +from .mooncake_ep_buffer import EventOverlap + + +def _using_musa_backend() -> bool: + return os.getenv("MOONCAKE_EP_USE_MUSA", "").upper() in { + "1", + "ON", + "TRUE", + "YES", + } + + +def _dist_barrier(group: dist.ProcessGroup) -> None: + if _using_musa_backend(): + dist.barrier(group=group, device_ids=[torch.cuda.current_device()]) + else: + group.barrier() + + +def _ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +def _align(x: int, alignment: int) -> int: + return _ceil_div(x, alignment) * alignment + + +class EPHandle: + """ + Official DeepEP elastic-compatible communication handle. + + The field names and semantics intentionally follow the official DeepEP elastic + handle contract so that model code can select Mooncake ElasticBuffer without + switching back to the legacy Buffer tuple handle. Mooncake stores the native + legacy handle as an implementation detail while the elastic kernels are being + wired to the Device API backend. + """ + + def __init__( + self, + do_expand: bool, + num_experts: int, + expert_alignment: int, + num_max_tokens_per_rank: int, + num_sms: int, + topk_idx: torch.Tensor, + num_recv_tokens_per_expert_list: List[int], + psum_num_recv_tokens_per_scaleup_rank: torch.Tensor, + psum_num_recv_tokens_per_expert: torch.Tensor, + recv_src_metadata: torch.Tensor, + dst_buffer_slot_idx: torch.Tensor, + token_metadata_at_forward: Optional[torch.Tensor], + channel_linked_list: Optional[torch.Tensor], + native_handle: Optional[Tuple[Any, ...]] = None, + ) -> None: + assert topk_idx is not None + self.do_expand = do_expand + self.num_experts = num_experts + self.expert_alignment = expert_alignment + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.num_sms = num_sms + self.topk_idx = topk_idx + self.psum_num_recv_tokens_per_scaleup_rank = psum_num_recv_tokens_per_scaleup_rank + self.psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert + self.num_recv_tokens_per_expert_list = num_recv_tokens_per_expert_list + self.recv_src_metadata = recv_src_metadata + self.dst_buffer_slot_idx = dst_buffer_slot_idx + self.token_metadata_at_forward = token_metadata_at_forward + self.channel_linked_list = channel_linked_list + self.native_handle = native_handle + + # Same convention as DeepEP: without a CPU sync this is an inferred upper + # bound; after CPU sync it tracks the actual received-token count. + self.num_recv_tokens = int(recv_src_metadata.shape[0]) + + +class ElasticBuffer: + """ + Official DeepEP elastic EP API backed by Mooncake EP transports. + + Public API source of truth: official DeepEP `ElasticBuffer`. The implementation is + deliberately separate from Mooncake's legacy `Buffer` API, while reusing the + existing Mooncake Device API transport/bootstrap path for the native data + movement backend. + """ + + # Mirrors DeepEP's fixed workspace assumptions closely enough for sizing and + # keeping one reusable buffer for all elastic EP shapes. + _NUM_MAX_RANKS = 1024 + _NUM_MAX_EXPERTS = 2048 + _NUM_MAX_CHANNELS = 8 * 160 + _NUM_BARRIER_TAGS = 16 + _NUM_MAX_INFLIGHT_AGRS = 32 + + def __init__( + self, + group: dist.ProcessGroup, + num_bytes: Optional[int] = None, + num_max_tokens_per_rank: int = 0, + hidden: int = 0, + num_topk: int = 0, + use_fp8_dispatch: bool = False, + deterministic: bool = False, + allow_hybrid_mode: bool = True, + allow_multiple_reduction: bool = True, + prefer_overlap_with_compute: bool = True, + sl_idx: int = 3, + num_allocated_qps: int = 0, + num_cpu_timeout_secs: int = 300, + num_gpu_timeout_secs: int = 100, + explicitly_destroy: bool = False, + ) -> None: + if not allow_multiple_reduction: + raise NotImplementedError( + "Mooncake ElasticBuffer currently supports only " + "allow_multiple_reduction=True" + ) + self.group = group + self.rank_idx = group.rank() + self.num_ranks = group.size() + self.allow_hybrid_mode = allow_hybrid_mode + self.allow_multiple_reduction = allow_multiple_reduction + self.prefer_overlap_with_compute = prefer_overlap_with_compute + self.deterministic = deterministic + self.sl_idx = int(os.getenv("EP_OVERRIDE_RDMA_SL", sl_idx)) + self.num_allocated_qps = num_allocated_qps + self.num_cpu_timeout_secs = num_cpu_timeout_secs + self.num_gpu_timeout_secs = num_gpu_timeout_secs + self.explicitly_destroy = explicitly_destroy + + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.hidden = hidden + self.num_topk = num_topk + self.use_fp8_dispatch = use_fp8_dispatch + + if num_bytes is None: + num_bytes = self.get_buffer_size_hint( + group, + num_max_tokens_per_rank, + hidden, + num_topk=num_topk, + use_fp8_dispatch=use_fp8_dispatch, + allow_hybrid_mode=allow_hybrid_mode, + allow_multiple_reduction=allow_multiple_reduction, + ) + self.num_bytes = num_bytes + + ( + self.num_scaleout_ranks, + self.num_scaleup_ranks, + ) = self._calculate_logical_domain_size(group, allow_hybrid_mode) + self.scaleout_rank_idx = self.rank_idx // self.num_scaleup_ranks + self.scaleup_rank_idx = self.rank_idx % self.num_scaleup_ranks + self.num_rdma_ranks, self.num_nvlink_ranks = self._calculate_physical_domain_size(group) + + self.backend = group + + # Native Mooncake transport/runtime. This keeps the legacy Buffer ABI + # untouched while giving ElasticBuffer users a dedicated native entrypoint. + from mooncake import ep + + self.runtime = ep.ElasticBuffer( + self.rank_idx, + self.num_ranks, + num_bytes, + num_max_tokens_per_rank, + hidden, + num_topk, + use_fp8_dispatch, + deterministic, + allow_hybrid_mode, + allow_multiple_reduction, + prefer_overlap_with_compute, + self.sl_idx, + num_allocated_qps, + num_cpu_timeout_secs, + num_gpu_timeout_secs, + ) + self._connect_native() + + torch.cuda.synchronize() + _dist_barrier(group) + torch.cuda.synchronize() + + def _active_ranks_mask(self) -> list: + # `mooncake.ep.get_active_ranks` is a Mooncake PG helper and performs a + # native static cast to MooncakeBackend. ElasticBuffer transport + # bootstrap can also be driven by a regular NCCL/Gloo ProcessGroup; in + # that case every rank in the supplied group is active by definition. + if "Mooncake" not in type(self.backend).__name__: + return [1] * self.num_ranks + + from mooncake.ep import get_active_ranks + + return get_active_ranks(self.backend).tolist() + + def _connect_native(self, is_update: bool = False) -> None: + from mooncake import ep + + if not bool(self.runtime.ibgda_disabled()): + raddr, rkey = self.runtime.get_mr_info() + raddr_tensor = torch.tensor([raddr], dtype=torch.int64, device="cuda") + raddrs = [torch.empty(1, dtype=torch.int64, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(raddrs, raddr_tensor, self.group) + raddrs_list = torch.cat(raddrs).tolist() + + rkey_tensor = torch.tensor([rkey], dtype=torch.int32, device="cuda") + rkeys = [torch.empty(1, dtype=torch.int32, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(rkeys, rkey_tensor, self.group) + rkeys_list = torch.cat(rkeys).tolist() + + all_to_all_size = ep.MAX_QP_COUNT // self.num_ranks + if is_update: + self.runtime.update_local_qpns() + + local_qpns = torch.tensor(self.runtime.get_local_qpns(), dtype=torch.int32, device="cuda").view( + -1, all_to_all_size + ) + remote_qpns = [torch.empty(all_to_all_size, dtype=torch.int32, device="cuda") for _ in range(self.num_ranks)] + dist.all_to_all(remote_qpns, list(torch.unbind(local_qpns)), self.group) + peer_qpns = [remote_qpns[r].tolist() for r in range(self.num_ranks)] + + local_lids = torch.tensor(self.runtime.get_local_lids(), dtype=torch.int32, device="cuda").view( + -1, all_to_all_size + ) + remote_lids = [torch.empty(all_to_all_size, dtype=torch.int32, device="cuda") for _ in range(self.num_ranks)] + dist.all_to_all(remote_lids, list(torch.unbind(local_lids)), self.group) + peer_lids = [remote_lids[r].tolist() for r in range(self.num_ranks)] + + subnet_prefix, interface_id = self.runtime.get_gid() + subnet_prefix_tensor = torch.tensor([subnet_prefix], dtype=torch.int64, device="cuda") + subnet_prefixes = [torch.empty(1, dtype=torch.int64, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(subnet_prefixes, subnet_prefix_tensor, self.group) + subnet_prefixes_list = torch.cat(subnet_prefixes).tolist() + + interface_id_tensor = torch.tensor([interface_id], dtype=torch.int64, device="cuda") + interface_ids = [torch.empty(1, dtype=torch.int64, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(interface_ids, interface_id_tensor, self.group) + interface_ids_list = torch.cat(interface_ids).tolist() + + active_ranks_mask = self._active_ranks_mask() + self.runtime.sync_ibgda_peers( + raddrs_list, + rkeys_list, + peer_qpns, + peer_lids, + subnet_prefixes_list, + interface_ids_list, + active_ranks_mask, + ) + + try: + local_handle_ints = self.runtime.get_ipc_handle() + local_handle_tensor = torch.tensor(local_handle_ints, dtype=torch.int32, device="cuda") + handles = [torch.empty(len(local_handle_ints), dtype=torch.int32, device="cuda") for _ in range(self.num_ranks)] + dist.all_gather(handles, local_handle_tensor, self.group) + remote_handles = [h.tolist() for h in handles] + + active_ranks_mask = self._active_ranks_mask() + self.runtime.sync_nvlink_ipc_handles(remote_handles, active_ranks_mask) + except Exception as exc: + if bool(self.runtime.ibgda_disabled()): + raise RuntimeError( + f"[Rank {self.rank_idx}] Failed to exchange IPC handles " + "for ElasticBuffer and RDMA is disabled; native elastic " + "mode cannot continue safely." + ) from exc + warnings.warn( + f"[Rank {self.rank_idx}] Failed to exchange IPC handles for ElasticBuffer: {exc}. " + "Continuing with RDMA-only routing.", + RuntimeWarning, + stacklevel=2, + ) + + def update_ep_member(self) -> None: + self._connect_native(True) + + def destroy(self) -> None: + # Existing Mooncake Buffer owns native resources through object lifetime. + # Keep the method to match the official ElasticBuffer API. + self.runtime = None + + @staticmethod + def _workspace_num_bytes() -> int: + num_bytes = 0 + num_bytes += ElasticBuffer._NUM_BARRIER_TAGS * ( + 8 + 2 * ElasticBuffer._NUM_MAX_RANKS * 4 + ) + num_bytes += (ElasticBuffer._NUM_MAX_RANKS + ElasticBuffer._NUM_MAX_EXPERTS) * 8 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * 8 * 2 + num_bytes += ElasticBuffer._NUM_MAX_EXPERTS * 8 * 2 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * 4 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * 4 * 2 + num_bytes += ElasticBuffer._NUM_MAX_EXPERTS * 4 * 2 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * ElasticBuffer._NUM_MAX_CHANNELS * 8 + num_bytes += ElasticBuffer._NUM_MAX_RANKS * ElasticBuffer._NUM_MAX_CHANNELS * 4 + num_bytes += 2 * 2 * 8 + num_bytes += (ElasticBuffer._NUM_MAX_INFLIGHT_AGRS + 1) * ElasticBuffer._NUM_MAX_RANKS * 4 + return _align(num_bytes, 32) + + @staticmethod + def _atomic_scratch_num_bytes() -> int: + # Mirrors the native runtime: RDMA atomics need a local response area + # separate from the remote-visible workspace. + return ElasticBuffer._workspace_num_bytes() + + @staticmethod + def get_buffer_size_hint( + group: dist.ProcessGroup, + num_max_tokens_per_rank: int, + hidden: int, + num_topk: int = 0, + use_fp8_dispatch: bool = False, + allow_hybrid_mode: bool = True, + allow_multiple_reduction: bool = True, + ) -> int: + try: + from mooncake import ep + + return int( + ep.calculate_elastic_buffer_size( + group.size(), + num_max_tokens_per_rank, + hidden, + num_topk, + use_fp8_dispatch, + allow_hybrid_mode, + allow_multiple_reduction, + ) + ) + except Exception: + pass + + num_ranks = group.size() + num_topk = max(1, num_topk) + dtype_bytes = 1 if use_fp8_dispatch else 2 + scale_bytes = _ceil_div(hidden, 128) * 4 if use_fp8_dispatch else 0 + token_bytes = _align(hidden * dtype_bytes, 32) + _align(scale_bytes, 32) + metadata_bytes = _align(num_topk * (4 + 4) + (1 + num_topk) * 4, 32) + per_slot_bytes = token_bytes + metadata_bytes + + # Direct elastic send/recv buffers plus room for combine reduce buffers. + dispatch_bytes = num_ranks * num_max_tokens_per_rank * num_topk * per_slot_bytes * 2 + combine_factor = 3 if allow_multiple_reduction else 4 + combine_bytes = dispatch_bytes * combine_factor + hybrid_factor = 2 if allow_hybrid_mode and num_ranks > 1 else 1 + return int( + ElasticBuffer._workspace_num_bytes() + + ElasticBuffer._atomic_scratch_num_bytes() + + hybrid_factor * (dispatch_bytes + combine_bytes) + ) + + @staticmethod + def get_engram_storage_size_hint( + num_entries: int, + hidden: int, + num_max_tokens_per_rank: int, + dtype: torch.dtype = torch.bfloat16, + ) -> int: + num_sf_packs = _ceil_div(hidden, 128) if dtype.itemsize <= 1 else 0 + num_bytes_per_entry = _align(hidden * dtype.itemsize + num_sf_packs * 4, 32) + return num_bytes_per_entry * (num_entries + num_max_tokens_per_rank) + + @staticmethod + def get_pp_buffer_size_hint(num_max_tensor_bytes: int, num_max_inflight_tensors: int) -> int: + return _align(num_max_tensor_bytes, 32) * num_max_inflight_tensors * 2 * 2 + + @staticmethod + def get_agrs_buffer_size_hint(group: dist.ProcessGroup, num_max_session_bytes: int) -> int: + return num_max_session_bytes + + @staticmethod + def _calculate_physical_domain_size(group: dist.ProcessGroup) -> Tuple[int, int]: + num_ranks = group.size() + num_local_ranks = int(os.getenv("MOONCAKE_EP_NUM_LOCAL_RANKS", "0")) + if num_local_ranks <= 0: + try: + num_local_ranks = max(1, min(num_ranks, torch.cuda.device_count())) + except Exception: + num_local_ranks = 1 + num_local_ranks = max(1, min(num_local_ranks, num_ranks)) + return _ceil_div(num_ranks, num_local_ranks), num_local_ranks + + @staticmethod + def _calculate_logical_domain_size(group: dist.ProcessGroup, allow_hybrid_mode: bool = True) -> Tuple[int, int]: + num_ranks = group.size() + num_rdma_ranks, num_nvlink_ranks = ElasticBuffer._calculate_physical_domain_size(group) + if allow_hybrid_mode and num_rdma_ranks > 1: + return num_rdma_ranks, num_nvlink_ranks + return 1, num_ranks + + def get_physical_domain_size(self) -> Tuple[int, int]: + return self.num_rdma_ranks, self.num_nvlink_ranks + + def get_logical_domain_size(self) -> Tuple[int, int]: + return self.num_scaleout_ranks, self.num_scaleup_ranks + + def barrier(self, use_comm_stream: bool = True, with_cpu_sync: bool = False) -> None: + if with_cpu_sync: + torch.cuda.synchronize() + _dist_barrier(self.group) + if with_cpu_sync: + torch.cuda.synchronize() + + @staticmethod + def capture() -> Any: + from mooncake import ep + + return ep.EventHandle() + + def get_theoretical_num_sms(self, num_experts: int, num_topk: int) -> int: + device = torch.cuda.current_device() + sm_count = torch.cuda.get_device_properties(device).multi_processor_count + if self.prefer_overlap_with_compute: + return max(1, min(24, sm_count // 4)) + return max(1, min(40, sm_count // 2, num_experts * max(1, num_topk))) + + def dispatch( + self, + x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + topk_idx: Optional[torch.Tensor] = None, + topk_weights: Optional[torch.Tensor] = None, + num_experts: Optional[int] = None, + num_max_tokens_per_rank: Optional[int] = None, + expert_alignment: Optional[int] = None, + handle: Optional[EPHandle] = None, + do_expand: bool = False, + do_cpu_sync: Optional[bool] = None, + num_sms: Optional[int] = None, + async_with_compute_stream: bool = False, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor], Optional[torch.Tensor], EPHandle, EventOverlap]: + if self.runtime is None: + raise RuntimeError("ElasticBuffer has been destroyed") + if handle is not None: + if topk_idx is not None or topk_weights is not None: + raise AssertionError("topk_idx and topk_weights must be None when cached handle is provided") + if do_cpu_sync: + raise AssertionError("Cannot do CPU sync with cached handle") + if handle.native_handle is None: + raise RuntimeError("Cached EPHandle is missing its native Mooncake handle") + topk_idx = handle.topk_idx + num_max_tokens_per_rank = num_max_tokens_per_rank or handle.num_max_tokens_per_rank + num_experts = num_experts or handle.num_experts + expert_alignment = handle.expert_alignment if expert_alignment is None else expert_alignment + num_sms = handle.num_sms if num_sms is None else num_sms + do_cpu_sync = False + else: + if topk_idx is None: + raise AssertionError("topk_idx must be provided when cached handle is not provided") + expert_alignment = 1 if expert_alignment is None else expert_alignment + do_cpu_sync = True if do_cpu_sync is None else do_cpu_sync + if do_expand: + warnings.warn( + "do_expand=True was requested. Mooncake currently returns the native packed expert layout; " + "expanded contiguous expert layout will be produced by the native elastic kernels.", + RuntimeWarning, + stacklevel=2, + ) + + x_data = x[0] if isinstance(x, tuple) else x + sf = x[1] if isinstance(x, tuple) else None + if num_experts is None: + num_experts = int(torch.max(topk_idx).item()) + 1 + if num_max_tokens_per_rank is None: + num_max_tokens_per_rank = self.num_max_tokens_per_rank or x_data.shape[0] + if num_sms is None: + num_sms = self.get_theoretical_num_sms(num_experts, topk_idx.shape[1]) + + active_ranks = torch.ones(self.num_ranks, dtype=torch.int32, device=x_data.device) + output = self.runtime.dispatch( + x_data, + sf, + topk_idx, + topk_weights, + active_ranks, + num_experts, + num_max_tokens_per_rank, + expert_alignment, + num_sms, + do_expand, + do_cpu_sync, + async_with_compute_stream, + handle.native_handle if handle is not None else None, + ) + native_handle = output.handle + + elastic_handle = EPHandle( + do_expand=native_handle.do_expand, + num_experts=native_handle.num_experts, + expert_alignment=native_handle.expert_alignment, + num_max_tokens_per_rank=native_handle.num_max_tokens_per_rank, + num_sms=native_handle.num_sms, + topk_idx=native_handle.topk_idx, + num_recv_tokens_per_expert_list=list(native_handle.num_recv_tokens_per_expert_list), + psum_num_recv_tokens_per_scaleup_rank=native_handle.psum_num_recv_tokens_per_scaleup_rank, + psum_num_recv_tokens_per_expert=native_handle.psum_num_recv_tokens_per_expert, + recv_src_metadata=native_handle.recv_src_metadata, + dst_buffer_slot_idx=native_handle.dst_buffer_slot_idx, + token_metadata_at_forward=native_handle.token_metadata_at_forward, + channel_linked_list=native_handle.channel_linked_list, + native_handle=native_handle, + ) + recv_x = (output.recv_x, output.recv_x_scales) if output.recv_x_scales is not None else output.recv_x + tensors_to_record = ( + x_data, + topk_idx, + active_ranks, + output.recv_x, + output.recv_topk_idx, + native_handle.topk_idx, + native_handle.psum_num_recv_tokens_per_scaleup_rank, + native_handle.psum_num_recv_tokens_per_expert, + native_handle.recv_src_metadata, + native_handle.dst_buffer_slot_idx, + *(() if sf is None else (sf,)), + *(() if topk_weights is None else (topk_weights,)), + *(() if output.recv_x_scales is None else (output.recv_x_scales,)), + *(() if output.recv_topk_weights is None else (output.recv_topk_weights,)), + *(() if native_handle.token_metadata_at_forward is None else (native_handle.token_metadata_at_forward,)), + *(() if native_handle.channel_linked_list is None else (native_handle.channel_linked_list,)), + ) + return ( + recv_x, + output.recv_topk_idx, + output.recv_topk_weights, + elastic_handle, + EventOverlap(output.event, tensors_to_record if async_with_compute_stream else None), + ) + + def combine( + self, + x: torch.Tensor, + handle: EPHandle, + topk_weights: Optional[torch.Tensor] = None, + num_sms: Optional[int] = None, + async_with_compute_stream: bool = False, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], EventOverlap]: + if self.runtime is None: + raise RuntimeError("ElasticBuffer has been destroyed") + if handle.native_handle is None: + raise RuntimeError("Mooncake EPHandle does not contain a native handle") + active_ranks = torch.ones(self.num_ranks, dtype=torch.int32, device=x.device) + if topk_weights is None: + topk_weights = torch.ones_like(handle.topk_idx, dtype=torch.float32, device=x.device) + output = self.runtime.combine( + x, + handle.native_handle, + topk_weights, + active_ranks, + num_sms if num_sms is not None else handle.num_sms, + async_with_compute_stream, + None, + ) + native_handle = handle.native_handle + tensors_to_record = ( + x, + topk_weights, + active_ranks, + output.combined_x, + native_handle.topk_idx, + native_handle.psum_num_recv_tokens_per_scaleup_rank, + native_handle.psum_num_recv_tokens_per_expert, + native_handle.recv_src_metadata, + native_handle.dst_buffer_slot_idx, + *(() if native_handle.token_metadata_at_forward is None else (native_handle.token_metadata_at_forward,)), + *(() if native_handle.channel_linked_list is None else (native_handle.channel_linked_list,)), + ) + return ( + output.combined_x, + output.combined_topk_weights, + EventOverlap(output.event, tensors_to_record if async_with_compute_stream else None), + ) + + +__all__ = ["ElasticBuffer", "EPHandle", "EventOverlap"] diff --git a/mooncake-wheel/mooncake/mooncake_ep_buffer.py b/mooncake-wheel/mooncake/mooncake_ep_buffer.py index b3deb7e762..8f9a720027 100644 --- a/mooncake-wheel/mooncake/mooncake_ep_buffer.py +++ b/mooncake-wheel/mooncake/mooncake_ep_buffer.py @@ -4,6 +4,23 @@ from typing import Any, Callable, List, Tuple, Optional, Union +def _env_enabled(name: str, default: bool = False) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.upper() in {"1", "ON", "TRUE", "YES"} + + +_USE_MACA = ( + _env_enabled("MOONCAKE_EP_USE_MACA") + or bool(getattr(torch.version, "maca", None)) +) +_USE_SPLIT_SEND_RECV = ( + _env_enabled("MOONCAKE_EP_USE_MUSA") + or _USE_MACA +) + + class EventOverlap: """ A wrapper class to manage CUDA events, also for better overlapping convenience. @@ -81,8 +98,75 @@ def __init__(self, group: dist.ProcessGroup, num_ep_buffer_bytes: int = 0): # P2P+IPC succeeds for all ranks). We re-evaluate after IPC sync below. self._use_fallback = bool(self.runtime.ibgda_disabled()) self._fallback_next_combine_buffer: Optional[torch.Tensor] = None + self._maca_phase_token: Optional[torch.Tensor] = None + self._maca_phase_recv_tokens: Optional[List[torch.Tensor]] = None + self._warned_active_ranks_without_mooncake_backend = False self.connect() - + + def _maca_phase_fence(self, send_event: Optional[Any] = None) -> None: + if not _USE_MACA: + return + + backend = dist.get_backend(self.group) + fence_device = torch.device("cpu" if backend == "gloo" else "cuda") + + def wait_send_done() -> None: + if send_event is not None: + send_event.synchronize() + else: + torch.cuda.synchronize() + + # Compatibility fence between SEND and RECV. The EP payload still + # uses the P2P fast path; this only keeps rank phases aligned on MACA. + wait_send_done() + if ( + self._maca_phase_token is None + or self._maca_phase_token.device != fence_device + ): + self._maca_phase_token = torch.empty( + 1, dtype=torch.int32, device=fence_device + ) + if ( + self._maca_phase_recv_tokens is None + or self._maca_phase_recv_tokens[0].device != fence_device + ): + self._maca_phase_recv_tokens = [ + torch.empty(1, dtype=torch.int32, device=fence_device) + for _ in range(self.group_size) + ] + self._maca_phase_token.fill_(1) + ops = [] + for peer in range(self.group_size): + if peer == self.rank: + continue + ops.append( + dist.P2POp( + dist.isend, self._maca_phase_token, peer, self.group + ) + ) + ops.append( + dist.P2POp( + dist.irecv, + self._maca_phase_recv_tokens[peer], + peer, + self.group, + ) + ) + if not ops: + return + for work in dist.batch_isend_irecv(ops): + work.wait() + + def _wrap_maca_recv_hook( + self, hook: Optional[Callable], send_event: Optional[Any] + ) -> Callable: + def wrapped_hook() -> None: + self._maca_phase_fence(send_event) + if hook is not None: + hook() + + return wrapped_hook + def connect(self, is_update: bool = False): from mooncake import ep @@ -157,37 +241,40 @@ def connect(self, is_update: bool = False): dist.all_gather(interface_ids_list, interface_id_t, self.group) interface_ids = torch.cat(interface_ids_list).tolist() - from mooncake.ep import get_active_ranks - active_ranks_mask = get_active_ranks(self.backend).tolist() + active_ranks_mask = self._active_ranks_list(torch.device("cuda")) self.runtime.sync_ibgda_peers( raddrs, rkeys, peer_qpns, peer_lids, subnet_prefixes, interface_ids, active_ranks_mask ) - try: - local_handle_ints = self.runtime.get_ipc_handle() - # pybind11 converts std::vector to a list of integers - local_handle_tensor = torch.tensor( - local_handle_ints, dtype=torch.int32, device="cuda" - ) - handles = [ - torch.empty(len(local_handle_ints), dtype=torch.int32, device="cuda") - for _ in range(self.group_size) - ] - dist.all_gather(handles, local_handle_tensor, self.group) - remote_handles = [h.tolist() for h in handles] - from mooncake.ep import get_active_ranks - active_ranks_mask = get_active_ranks(self.backend).tolist() - self.runtime.sync_nvlink_ipc_handles(remote_handles, - active_ranks_mask) - except Exception as e: - import warnings - - warnings.warn( - f"[Rank {self.rank}] Failed to exchange IPC handles: {e}. Falling back.", - RuntimeWarning, - stacklevel=2, - ) + if self.group_size == 1: + # No peer can import this IPC handle in single-rank EP. Skipping + # export also avoids unnecessary driver IPC calls on MACA. + self._use_fallback = False + return + else: + try: + local_handle_ints = self.runtime.get_ipc_handle() + # pybind11 converts std::vector to a list of integers + local_handle_tensor = torch.tensor( + local_handle_ints, dtype=torch.int32, device="cuda" + ) + handles = [ + torch.empty(len(local_handle_ints), dtype=torch.int32, device="cuda") + for _ in range(self.group_size) + ] + dist.all_gather(handles, local_handle_tensor, self.group) + remote_handles = [h.tolist() for h in handles] + active_ranks_mask = self._active_ranks_list(torch.device("cuda")) + self.runtime.sync_nvlink_ipc_handles(remote_handles, active_ranks_mask) + except Exception as e: + import warnings + + warnings.warn( + f"[Rank {self.rank}] Failed to exchange IPC handles: {e}. Falling back.", + RuntimeWarning, + stacklevel=2, + ) use_fast_path = False try: @@ -198,10 +285,46 @@ def connect(self, is_update: bool = False): self._use_fallback = not use_fast_path - def update_ep_member(self): self.connect(True) + def _is_mooncake_backend(self) -> bool: + try: + return dist.get_backend(self.group) == "mooncake" + except Exception: + return False + + def _active_ranks_tensor( + self, device: torch.device, dtype: torch.dtype = torch.int32 + ) -> torch.Tensor: + if not self._is_mooncake_backend(): + if not self._warned_active_ranks_without_mooncake_backend: + import warnings + + try: + backend = dist.get_backend(self.group) + except Exception: + backend = "unknown" + warnings.warn( + "Mooncake EP active_ranks is only available with the " + f"mooncake process group; got backend={backend}. " + "Treating all ranks as active.", + RuntimeWarning, + stacklevel=2, + ) + self._warned_active_ranks_without_mooncake_backend = True + return torch.ones((self.group_size,), dtype=dtype, device=device) + + try: + from mooncake.ep import get_active_ranks + + return get_active_ranks(self.backend).to(device=device, dtype=dtype) + except Exception: + return torch.ones((self.group_size,), dtype=dtype, device=device) + + def _active_ranks_list(self, device: torch.device) -> List[int]: + return self._active_ranks_tensor(device=device, dtype=torch.int32).tolist() + @staticmethod def get_ep_buffer_size_hint( num_max_dispatch_tokens_per_rank: int, @@ -224,7 +347,7 @@ def dispatch( num_max_dispatch_tokens_per_rank: int, num_experts: int, timeout_us: int, - use_fp8: bool = True, + use_fp8: Optional[bool] = None, async_finish: bool = False, return_recv_hook: bool = False, ) -> Tuple[ @@ -234,23 +357,30 @@ def dispatch( EventOverlap, Callable, ]: - # MUSA does not support cooperative grid sync, so the C++ runtime + if use_fp8 is None: + use_fp8 = not _USE_MACA + elif _USE_MACA and use_fp8: + raise NotImplementedError("FP8 dispatch is not supported on MACA") + + # MUSA/MACA do not support cooperative grid sync, so the C++ runtime # splits no-hook calls into SEND -> phase-ack -> RECV instead of using # a single cooperative kernel. async_finish still returns a stream # event, but it is not the CUDA single-kernel cooperative path. - if os.getenv("MOONCAKE_EP_USE_MUSA") and async_finish: + if _USE_SPLIT_SEND_RECV and async_finish: import warnings warnings.warn( - "MUSA async_finish uses split SEND/RECV kernels plus a stream " + "async_finish uses split SEND/RECV kernels plus a stream " "event, not CUDA cooperative single-kernel async semantics.", RuntimeWarning, stacklevel=2, ) - if self._use_fallback: - from mooncake.ep import get_active_ranks + runtime_return_recv_hook = return_recv_hook or ( + _USE_MACA and not self._use_fallback + ) + if self._use_fallback: ( packed_recv_x, packed_recv_x_scales, @@ -267,7 +397,7 @@ def dispatch( use_fp8, return_recv_hook, ) - backend_active_ranks = get_active_ranks(self.backend).to( + backend_active_ranks = self._active_ranks_tensor( device=active_ranks.device, dtype=active_ranks.dtype ) if active_ranks.numel() == backend_active_ranks.numel(): @@ -290,8 +420,13 @@ def dispatch( timeout_us, use_fp8, async_finish, - return_recv_hook, + runtime_return_recv_hook, ) + if _USE_MACA: + hook = self._wrap_maca_recv_hook(hook, event) + if not return_recv_hook: + hook() + hook = None handle = ( packed_recv_src_info, packed_recv_layout_range, @@ -330,12 +465,12 @@ def combine( return_recv_hook: bool = False, out: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, EventOverlap, Callable]: - # Same MUSA split-kernel behavior as dispatch(). - if os.getenv("MOONCAKE_EP_USE_MUSA") and async_finish: + # Same split-kernel behavior as dispatch(). + if _USE_SPLIT_SEND_RECV and async_finish: import warnings warnings.warn( - "MUSA async_finish uses split SEND/RECV kernels plus a stream " + "async_finish uses split SEND/RECV kernels plus a stream " "event, not CUDA cooperative single-kernel async semantics.", RuntimeWarning, stacklevel=2, @@ -348,9 +483,11 @@ def combine( hidden, num_experts, ) = handle - if self._use_fallback: - from mooncake.ep import get_active_ranks + runtime_return_recv_hook = return_recv_hook or ( + _USE_MACA and not self._use_fallback + ) + if self._use_fallback: combined_x, event, hook = self._fallback_combine( x, topk_idx, @@ -363,7 +500,7 @@ def combine( return_recv_hook, out, ) - backend_active_ranks = get_active_ranks(self.backend).to( + backend_active_ranks = self._active_ranks_tensor( device=active_ranks.device, dtype=active_ranks.dtype ) if active_ranks.numel() == backend_active_ranks.numel(): @@ -381,9 +518,14 @@ def combine( timeout_us, zero_copy, async_finish, - return_recv_hook, + runtime_return_recv_hook, out, ) + if _USE_MACA: + hook = self._wrap_maca_recv_hook(hook, event) + if not return_recv_hook: + hook() + hook = None tensors_to_record = ( x, topk_idx, @@ -458,8 +600,6 @@ def _fallback_dispatch( use_fp8: bool, return_recv_hook: bool, ): - from mooncake.ep import get_active_ranks - with torch.profiler.record_function("dispatch"): num_tokens, hidden = x.shape k = topk_idx.size(1) @@ -476,7 +616,7 @@ def _fallback_dispatch( ] dist.all_gather(num_tokens_list, num_tokens_tensor, group=self.group) num_tokens_per_rank = [t.item() for t in num_tokens_list] - backend_active_ranks = get_active_ranks(self.backend).tolist() + backend_active_ranks = self._active_ranks_list(x.device) for i in range(num_ranks): if backend_active_ranks[i] == 0: num_tokens_per_rank[i] = 0 @@ -682,8 +822,6 @@ def _fallback_combine( return_recv_hook: bool, out: Optional[torch.Tensor], ): - from mooncake.ep import get_active_ranks - with torch.profiler.record_function("combine"): num_tokens = topk_idx.size(0) hidden = (x if not zero_copy else self._fallback_next_combine_buffer).size( @@ -702,7 +840,7 @@ def _fallback_combine( ] dist.all_gather(num_tokens_list, num_tokens_tensor, group=self.group) num_tokens_per_rank = [t.item() for t in num_tokens_list] - backend_active_ranks = get_active_ranks(self.backend).tolist() + backend_active_ranks = self._active_ranks_list(topk_idx.device) for i in range(num_ranks): if backend_active_ranks[i] == 0: num_tokens_per_rank[i] = 0 diff --git a/mooncake-wheel/mooncake/mooncake_ssd_register.py b/mooncake-wheel/mooncake/mooncake_ssd_register.py index 3a5e6e18cf..bd49ee58e9 100644 --- a/mooncake-wheel/mooncake/mooncake_ssd_register.py +++ b/mooncake-wheel/mooncake/mooncake_ssd_register.py @@ -4,7 +4,6 @@ import argparse import json import logging -import time import re import shlex from typing import List, Dict, Any @@ -12,7 +11,6 @@ import paramiko from mooncake.store import MooncakeDistributedNoFRegister -from mooncake.mooncake_config import MooncakeConfig class MooncakeNoFRegister: diff --git a/mooncake-wheel/mooncake/mooncake_ssd_unregister.py b/mooncake-wheel/mooncake/mooncake_ssd_unregister.py index ec84b21b98..3c8a42d330 100644 --- a/mooncake-wheel/mooncake/mooncake_ssd_unregister.py +++ b/mooncake-wheel/mooncake/mooncake_ssd_unregister.py @@ -4,14 +4,12 @@ import argparse import logging import json -import re import shlex from typing import List, Dict, Any import paramiko from mooncake.store import MooncakeDistributedNoFRegister -from mooncake.mooncake_config import MooncakeConfig class MooncakeNoFUnregister: @@ -150,7 +148,6 @@ def _parse_spdk_targets(self, master_server_address: str) -> List[Dict[str, Any] # Default transport parameters trsvcid = int(target.get('port', '4420')) - trtype = target.get('trtype', 'RDMA') # Create SSD config for each namespace (or specified ns only) if specified_ns is not None: diff --git a/mooncake-wheel/mooncake/mooncake_store_service.py b/mooncake-wheel/mooncake/mooncake_store_service.py index 73f5fed2f0..a78bfb4659 100644 --- a/mooncake-wheel/mooncake/mooncake_store_service.py +++ b/mooncake-wheel/mooncake/mooncake_store_service.py @@ -5,6 +5,7 @@ import asyncio import json import logging +import signal import time from aiohttp import web @@ -20,6 +21,7 @@ async def wrapper(request): finally: elapsed_ms = (time.perf_counter() - start_time) * 1000 logging.info(f"{operation_name} operation completed in {elapsed_ms:.2f} ms") + return wrapper @@ -32,6 +34,32 @@ def _shm_name_to_path(name): return f"/dev/shm/{normalized}" +def _unblock_shutdown_signals(): + try: + signal.pthread_sigmask( + signal.SIG_UNBLOCK, {signal.SIGINT, signal.SIGTERM} + ) + except AttributeError: + pass + + +def _install_shutdown_signal_handlers(loop, shutdown_event): + def request_shutdown(signum): + logging.info("Received signal %s, shutting down", signum) + shutdown_event.set() + + for sig in (signal.SIGINT, signal.SIGTERM): + try: + loop.add_signal_handler(sig, request_shutdown, sig) + except (NotImplementedError, RuntimeError): + signal.signal( + sig, + lambda signum, _frame: loop.call_soon_threadsafe( + request_shutdown, signum + ), + ) + + class MooncakeStoreService: """ Mooncake Store Service with REST API. @@ -63,10 +91,12 @@ def __init__(self, config_path: str = None, cli_config: dict = None): self._setup_logging() # State for /api/reconfigure (Prefill/Decode mode switch) - self.current_mode = "prefill" # "prefill" or "decode" - self.mounted_segment_ids = [] # persisted segment_ids from last decode mount - self.last_mount_info = {} # last mount parameters for debugging - self._state_lock = asyncio.Lock() # serialize reconfigure/mount/unmount state changes + self.current_mode = "prefill" # "prefill" or "decode" + self.mounted_segment_ids = [] # persisted segment_ids from last decode mount + self.last_mount_info = {} # last mount parameters for debugging + self._state_lock = ( + asyncio.Lock() + ) # serialize reconfigure/mount/unmount state changes try: if config_path: @@ -88,15 +118,18 @@ def __init__(self, config_path: str = None, cli_config: dict = None): def _setup_logging(self): logging.basicConfig( level=logging.INFO, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) - async def start_store_service(self, max_wait_time: float = 60): + async def start_store_service( + self, max_wait_time: float = 60, shutdown_event=None + ): """ Start the store service with retry mechanism. Args: max_wait_time: Maximum total wait time in seconds (default: 60) + shutdown_event: Optional asyncio event used to cancel startup Returns: True if successful, False otherwise @@ -105,7 +138,15 @@ async def start_store_service(self, max_wait_time: float = 60): retry_interval = 1.0 # Fixed retry interval: 1 second start_time = time.perf_counter() + if shutdown_event is not None: + # Process any signal callback queued immediately before startup. + await asyncio.sleep(0) + while True: + if shutdown_event is not None and shutdown_event.is_set(): + logging.info("Store startup cancelled by shutdown request") + return False + elapsed = time.perf_counter() - start_time if elapsed >= max_wait_time: logging.error( @@ -122,22 +163,41 @@ async def start_store_service(self, max_wait_time: float = 60): self.store = MooncakeDistributedStore() ret = self.store.setup( - self.config.local_hostname, - self.config.metadata_server, - self.config.global_segment_size, - self.config.local_buffer_size, - self.config.protocol, - self.config.device_name, - self.config.master_server_address, - None, - self.config.enable_ssd_offload, - self.config.ssd_offload_path + { + "local_hostname": self.config.local_hostname, + "metadata_server": self.config.metadata_server, + "global_segment_size": self.config.global_segment_size, + "local_buffer_size": self.config.local_buffer_size, + "protocol": self.config.protocol, + "rdma_devices": self.config.device_name, + "master_server_addr": self.config.master_server_address, + "enable_ssd_offload": self.config.enable_ssd_offload, + "ssd_offload_path": self.config.ssd_offload_path, + "tenant_id": self.config.tenant_id, + "enable_client_http_server": ( + self.config.enable_client_http_server + ), + "client_http_port": self.config.client_http_port, + } ) + if shutdown_event is not None: + # setup() is synchronous. Give asyncio signal callbacks a + # chance to publish a shutdown requested while it ran. + await asyncio.sleep(0) + if shutdown_event.is_set(): + logging.info( + "Store startup cancelled by shutdown request" + ) + await self.stop() + return False + if ret != 0: raise RuntimeError("Store initialization failed") - logging.info(f"Store service started successfully on {self.config.local_hostname}") + logging.info( + f"Store service started successfully on {self.config.local_hostname}" + ) return True except Exception as e: @@ -146,7 +206,9 @@ async def start_store_service(self, max_wait_time: float = 60): remaining_time = max_wait_time - elapsed_after_attempt # Calculate actual sleep duration - actual_sleep_time = min(retry_interval, remaining_time) if remaining_time > 0 else 0 + actual_sleep_time = ( + min(retry_interval, remaining_time) if remaining_time > 0 else 0 + ) logging.warning( f"Store startup failed (attempt {retry_count}): {e}. " @@ -155,26 +217,55 @@ async def start_store_service(self, max_wait_time: float = 60): # Wait before retry, but don't exceed max_wait_time if actual_sleep_time > 0: - await asyncio.sleep(actual_sleep_time) - + if shutdown_event is not None: + try: + await asyncio.wait_for( + shutdown_event.wait(), + timeout=actual_sleep_time, + ) + logging.info( + "Store startup cancelled by shutdown request" + ) + return False + except asyncio.TimeoutError: + pass + else: + await asyncio.sleep(actual_sleep_time) async def start_http_service(self, port: int = 8080): app = web.Application(client_max_size=1024 * 1024 * 100) # 100MB limit - app.add_routes([ - web.post('/api/reconfigure', _timed_handler("RECONFIGURE", self.handle_reconfigure)), - web.post('/api/mount_shm', _timed_handler("MOUNT_SHM", self.handle_mount_shm)), - web.post('/api/unmount_shm', _timed_handler("UNMOUNT_SHM", self.handle_unmount_shm)), - web.post('/api/mount', _timed_handler("MOUNT", self.handle_mount)), - web.post('/api/unmount', _timed_handler("UNMOUNT", self.handle_unmount)), - web.put('/api/put', _timed_handler("PUT", self.handle_put)), - web.get('/api/get/{key}', _timed_handler("GET", self.handle_get)), - web.get('/api/exist/{key}', _timed_handler("EXIST", self.handle_exist)), - web.delete('/api/remove/{key}', _timed_handler("REMOVE", self.handle_remove)), - web.delete('/api/remove_all', _timed_handler("REMOVE_ALL", self.handle_remove_all)) - ]) + app.add_routes( + [ + web.post( + "/api/reconfigure", + _timed_handler("RECONFIGURE", self.handle_reconfigure), + ), + web.post( + "/api/mount_shm", _timed_handler("MOUNT_SHM", self.handle_mount_shm) + ), + web.post( + "/api/unmount_shm", + _timed_handler("UNMOUNT_SHM", self.handle_unmount_shm), + ), + web.post("/api/mount", _timed_handler("MOUNT", self.handle_mount)), + web.post( + "/api/unmount", _timed_handler("UNMOUNT", self.handle_unmount) + ), + web.put("/api/put", _timed_handler("PUT", self.handle_put)), + web.get("/api/get/{key}", _timed_handler("GET", self.handle_get)), + web.get("/api/exist/{key}", _timed_handler("EXIST", self.handle_exist)), + web.delete( + "/api/remove/{key}", _timed_handler("REMOVE", self.handle_remove) + ), + web.delete( + "/api/remove_all", + _timed_handler("REMOVE_ALL", self.handle_remove_all), + ), + ] + ) runner = web.AppRunner(app) await runner.setup() - site = web.TCPSite(runner, '0.0.0.0', port) + site = web.TCPSite(runner, "0.0.0.0", port) await site.start() logging.info(f"REST API started on port {port}") return True @@ -195,25 +286,57 @@ async def handle_reconfigure(self, request): if not path or size is None: return web.Response( status=400, - text=json.dumps({"error": "Missing path or size for decode mode"}), - content_type="application/json" + text=json.dumps( + {"error": "Missing path or size for decode mode"} + ), + content_type="application/json", ) async with self._state_lock: - # If already in decode mode with mounted segments, unmount them first - if self.mounted_segment_ids: - logging.info("Reconfigure decode: unmounting previous segments before remount") - ret = self.store.unmount_segment(self.mounted_segment_ids) - if ret != 0: + # Make-before-break remount: capture the currently serving + # segments so a failed remount can keep them alive instead of + # dropping all capacity. + previous_segment_ids = list(self.mounted_segment_ids) + # /api/reconfigure binds through store.mount_segment -> + # MooncakeDistributedStore.mount_segment -> RealClient::mountSegment + # -> Client::MountSegmentAndGetId -> MasterClient::MountSegment, + # the standard allocator path: each mount mints a fresh UUID + # and the duplicate check is UUID-keyed. That path never calls + # MountNoFSegment or enters ScopedNoFSegmentAccess, so the NoF + # te_endpoint dedup restriction does not apply here even in a + # USE_NOF build. A same-path make-before-break therefore cannot + # collide, and the old and new segments can coexist; keeping + # the old segments live across a failed replacement mount + # preserves decode capacity (the bug this PR fixes). Mount the + # new segment first, and only retire the previous segments + # per-id once the new mount succeeds. + result = self.store.mount_segment( + path, size, offset, protocol, location + ) + if result["ret"] != 0: + if previous_segment_ids: + # New mount failed but the previous segments are still + # healthy; keep serving from them instead of demoting. + logging.warning( + "Reconfigure decode: mount of %s failed (ret=%s); " + "keeping previous decode segments", + path, + result["ret"], + ) return web.Response( status=500, - text=json.dumps({"error": f"Unmount of previous segments failed, ret={ret}"}), - content_type="application/json" + text=json.dumps( + { + "error": ( + f"Mount failed, ret={result['ret']}; " + "keeping previous decode segments" + ), + "mode": self.current_mode, + } + ), + content_type="application/json", ) - self.mounted_segment_ids.clear() - - result = self.store.mount_segment(path, size, offset, protocol, location) - if result["ret"] != 0: + # Nothing healthy to fall back to; roll back to prefill. self.current_mode = "prefill" self.mounted_segment_ids.clear() self.last_mount_info.clear() @@ -228,24 +351,54 @@ async def handle_reconfigure(self, request): "mode": self.current_mode, } ), - content_type="application/json" + content_type="application/json", + ) + + # New mount succeeded; retire any previously serving segments. + # unmount_segment returns the first error for the whole batch, + # so a single call can't tell which ids were actually freed. + # Unmount one id at a time (as handle_unmount_shm does) and keep + # only the ids whose cleanup genuinely failed: this neither leaks + # them (dropping all ids on failure) nor retains stale ids for + # segments that were already removed (keeping all ids on failure). + failed_unmount_ids = [] + for sid in previous_segment_ids: + ret = self.store.unmount_segment([sid]) + if ret != 0: + failed_unmount_ids.append(sid) + if failed_unmount_ids: + # The new segment is already live; a failed cleanup only + # leaves the old ids around. Keep them tracked so a later + # unmount (or a switch back to prefill) can retry them. + logging.warning( + "Reconfigure decode: new mount succeeded but unmount of " + "previous segments %s failed; keeping them tracked for " + "future cleanup", + failed_unmount_ids, ) - self.mounted_segment_ids = list(result["segment_ids"]) + self.mounted_segment_ids = ( + list(result["segment_ids"]) + failed_unmount_ids + ) self.current_mode = "decode" self.last_mount_info = { - "path": path, "offset": offset, "size": size, - "protocol": protocol, "location": location + "path": path, + "offset": offset, + "size": size, + "protocol": protocol, + "location": location, } return web.Response( status=200, - text=json.dumps({ - "status": "success", - "mode": self.current_mode, - "segment_ids": self.mounted_segment_ids, - }), - content_type="application/json" + text=json.dumps( + { + "status": "success", + "mode": self.current_mode, + "segment_ids": self.mounted_segment_ids, + } + ), + content_type="application/json", ) elif mode == "prefill": @@ -255,8 +408,10 @@ async def handle_reconfigure(self, request): if ret != 0: return web.Response( status=500, - text=json.dumps({"error": f"Unmount failed, ret={ret}"}), - content_type="application/json" + text=json.dumps( + {"error": f"Unmount failed, ret={ret}"} + ), + content_type="application/json", ) self.mounted_segment_ids.clear() @@ -266,21 +421,23 @@ async def handle_reconfigure(self, request): return web.Response( status=200, text=json.dumps({"status": "success", "mode": self.current_mode}), - content_type="application/json" + content_type="application/json", ) else: return web.Response( status=400, - text=json.dumps({"error": "Invalid mode. Use 'decode' or 'prefill'"}), - content_type="application/json" + text=json.dumps( + {"error": "Invalid mode. Use 'decode' or 'prefill'"} + ), + content_type="application/json", ) except Exception as e: logging.error("RECONFIGURE error: %s", e) return web.Response( status=500, text=json.dumps({"error": str(e)}), - content_type="application/json" + content_type="application/json", ) async def handle_mount_shm(self, request): @@ -297,7 +454,7 @@ async def handle_mount_shm(self, request): return web.Response( status=400, text=json.dumps({"error": "Missing or invalid name or size"}), - content_type="application/json" + content_type="application/json", ) result = self.store.mount_segment(path, size, offset, protocol, location) @@ -305,7 +462,7 @@ async def handle_mount_shm(self, request): return web.Response( status=500, text=json.dumps({"error": f"Mount failed, ret={result['ret']}"}), - content_type="application/json" + content_type="application/json", ) return web.Response( @@ -323,7 +480,7 @@ async def handle_mount_shm(self, request): return web.Response( status=500, text=json.dumps({"error": str(e)}), - content_type="application/json" + content_type="application/json", ) async def handle_unmount_shm(self, request): @@ -374,7 +531,7 @@ async def handle_unmount_shm(self, request): return web.Response( status=500, text=json.dumps({"error": str(e)}), - content_type="application/json" + content_type="application/json", ) async def handle_mount(self, request): @@ -387,16 +544,20 @@ async def handle_mount(self, request): if type(size) is not int or size <= 0: return web.Response( status=400, - text=json.dumps({"error": "Invalid size, must be a positive integer"}), - content_type="application/json" + text=json.dumps( + {"error": "Invalid size, must be a positive integer"} + ), + content_type="application/json", ) result = self.store.allocate_and_mount_segment(size, protocol, location) if result["ret"] != 0: return web.Response( status=500, - text=json.dumps({"error": f"Allocate and mount failed, ret={result['ret']}"}), - content_type="application/json" + text=json.dumps( + {"error": f"Allocate and mount failed, ret={result['ret']}"} + ), + content_type="application/json", ) return web.Response( @@ -415,7 +576,7 @@ async def handle_mount(self, request): return web.Response( status=500, text=json.dumps({"error": str(e)}), - content_type="application/json" + content_type="application/json", ) async def handle_unmount(self, request): @@ -432,15 +593,11 @@ async def handle_unmount(self, request): ) grace_period_seconds = data.get("grace_period_seconds", 0) - ret = self.store.unmount_and_free_segment( - segment_ids, grace_period_seconds - ) + ret = self.store.unmount_and_free_segment(segment_ids, grace_period_seconds) if ret != 0: return web.Response( status=500, - text=json.dumps( - {"error": f"Unmount and free failed, ret={ret}"} - ), + text=json.dumps({"error": f"Unmount and free failed, ret={ret}"}), content_type="application/json", ) @@ -454,20 +611,20 @@ async def handle_unmount(self, request): return web.Response( status=500, text=json.dumps({"error": str(e)}), - content_type="application/json" + content_type="application/json", ) async def handle_put(self, request): try: data = await request.json() - key = data.get('key') - raw_value = data.get('value') + key = data.get("key") + raw_value = data.get("value") if not key or raw_value is None: return web.Response( status=400, - text=json.dumps({'error': 'Missing key or value'}), - content_type='application/json' + text=json.dumps({"error": "Missing key or value"}), + content_type="application/json", ) value = raw_value.encode() @@ -475,89 +632,122 @@ async def handle_put(self, request): if ret != 0: return web.Response( status=500, - text=json.dumps({'error': 'PUT operation failed'}), - content_type='application/json' + text=json.dumps({"error": "PUT operation failed"}), + content_type="application/json", ) return web.Response( status=200, - text=json.dumps({'status': 'success'}), - content_type='application/json' + text=json.dumps({"status": "success"}), + content_type="application/json", ) except Exception as e: logging.error("PUT error: %s", e) return web.Response( status=500, - text=json.dumps({'error': str(e)}), - content_type='application/json' + text=json.dumps({"error": str(e)}), + content_type="application/json", ) async def handle_get(self, request): try: - key = request.match_info['key'] - value = self.store.get(key) + key = request.match_info["key"] + exists = self.store.is_exist(key) - if not value: + if exists == 0: return web.Response( status=404, - text=json.dumps({'error': 'Key not found'}), - content_type='application/json' + text=json.dumps({"error": "Key not found"}), + content_type="application/json", + ) + if exists < 0: + return web.Response( + status=500, + text=json.dumps({"error": "Exist check failed"}), + content_type="application/json", ) + value = self.store.get(key) + if value is None: + return web.Response( + status=500, + text=json.dumps({"error": "GET operation failed"}), + content_type="application/json", + ) + if value == b"": + exists = self.store.is_exist(key) + if exists == 0: + return web.Response( + status=404, + text=json.dumps({"error": "Key not found"}), + content_type="application/json", + ) + if exists < 0: + return web.Response( + status=500, + text=json.dumps({"error": "Exist check failed"}), + content_type="application/json", + ) + return web.Response( - status=200, - body=value, - content_type='application/octet-stream' + status=200, body=value, content_type="application/octet-stream" ) except Exception as e: logging.error("GET error: %s", e) return web.Response( status=500, - text=json.dumps({'error': str(e)}), - content_type='application/json' + text=json.dumps({"error": str(e)}), + content_type="application/json", ) async def handle_exist(self, request): try: - key = request.match_info['key'] + key = request.match_info["key"] exists = self.store.is_exist(key) + if exists < 0: + return web.Response( + status=500, + text=json.dumps({"error": "Exist check failed"}), + content_type="application/json", + ) + return web.Response( status=200, - text=json.dumps({'exists': bool(exists)}), - content_type='application/json' + text=json.dumps({"exists": exists > 0}), + content_type="application/json", ) except Exception as e: logging.error("EXIST error: %s", e) return web.Response( status=500, - text=json.dumps({'error': str(e)}), - content_type='application/json' + text=json.dumps({"error": str(e)}), + content_type="application/json", ) async def handle_remove(self, request): try: - key = request.match_info['key'] + key = request.match_info["key"] ret = self.store.remove(key) if ret != 0: return web.Response( status=500, - text=json.dumps({'error': 'Remove operation failed'}), - content_type='application/json' + text=json.dumps({"error": "Remove operation failed"}), + content_type="application/json", ) return web.Response( status=200, - text=json.dumps({'status': 'success'}), - content_type='application/json' + text=json.dumps({"status": "success"}), + content_type="application/json", ) except Exception as e: logging.error("REMOVE error: %s", e) return web.Response( status=500, - text=json.dumps({'error': str(e)}), - content_type='application/json' + text=json.dumps({"error": str(e)}), + content_type="application/json", ) async def handle_remove_all(self, request): @@ -567,78 +757,108 @@ async def handle_remove_all(self, request): if ret < 0: return web.Response( status=500, - text=json.dumps({'error': 'RemoveAll operation failed'}), - content_type='application/json' + text=json.dumps({"error": "RemoveAll operation failed"}), + content_type="application/json", ) return web.Response( status=200, - text=json.dumps({'status': 'success removed ' + str(ret) + ' keys'}), - content_type='application/json' + text=json.dumps({"status": "success removed " + str(ret) + " keys"}), + content_type="application/json", ) except Exception as e: logging.error("REMOVE_ALL error: %s", e) return web.Response( status=500, - text=json.dumps({'error': str(e)}), - content_type='application/json' + text=json.dumps({"error": str(e)}), + content_type="application/json", ) async def stop(self): if self.store: - self.store.close() - logging.info("Mooncake service stopped") + ret = self.store.close() + self.store = None + if ret != 0: + logging.warning("Mooncake service close returned %s", ret) + else: + logging.info("Mooncake service stopped") + def parse_arguments(): - parser = argparse.ArgumentParser(description='Mooncake Store Service with REST API') - parser.add_argument('--config', type=str, - help='Path to Mooncake config file', - required=False) - parser.add_argument('-D', '--define', action='append', - help='Override configuration with key=value pairs (e.g., -Dlocal_hostname=example.com)', - default=[]) - parser.add_argument('--port', type=int, - help='HTTP API port (default: 8080)', - default=8080, - required=False) - parser.add_argument('--max-wait-time', type=float, - help='Maximum total wait time in seconds (default: 60)', - default=60, - required=False) + parser = argparse.ArgumentParser(description="Mooncake Store Service with REST API") + parser.add_argument( + "--config", type=str, help="Path to Mooncake config file", required=False + ) + parser.add_argument( + "-D", + "--define", + action="append", + help="Override configuration with key=value pairs (e.g., -Dlocal_hostname=example.com)", + default=[], + ) + parser.add_argument( + "--port", + type=int, + help="HTTP API port (default: 8080)", + default=8080, + required=False, + ) + parser.add_argument( + "--max-wait-time", + type=float, + help="Maximum total wait time in seconds (default: 60)", + default=60, + required=False, + ) return parser.parse_args() + async def main(): args = parse_arguments() # Parse -D key=value pairs into a dictionary cli_config = {} for item in args.define: - if '=' in item: - key, value = item.split('=', 1) + if "=" in item: + key, value = item.split("=", 1) cli_config[key] = value else: logging.warning(f"Ignoring invalid CLI config: {item}") service = MooncakeStoreService(args.config, cli_config) + shutdown_event = asyncio.Event() + loop = asyncio.get_running_loop() + + _install_shutdown_signal_handlers(loop, shutdown_event) + _unblock_shutdown_signals() try: - if not await service.start_store_service(max_wait_time=args.max_wait_time): + if not await service.start_store_service( + max_wait_time=args.max_wait_time, + shutdown_event=shutdown_event, + ): + if shutdown_event.is_set(): + return raise RuntimeError("Failed to start store service") + _unblock_shutdown_signals() + await asyncio.sleep(0) + if shutdown_event.is_set(): + return + if not await service.start_http_service(args.port): raise RuntimeError("Failed to start HTTP service") - logging.info("Mooncake Store Service is running. Press Ctrl+C to stop.") - while True: - await asyncio.sleep(1) + logging.info("Mooncake Store Service is running") + await shutdown_event.wait() except KeyboardInterrupt: - logging.info("Received shutdown signal") - await service.stop() + logging.info("Received keyboard interrupt, shutting down") except Exception as e: logging.error("Service error: %s", e) - await service.stop() raise + finally: + await service.stop() def sync_main(): diff --git a/mooncake-wheel/mooncake/structured_object_store.py b/mooncake-wheel/mooncake/structured_object_store.py index a50871c273..9b6a7fbda9 100644 --- a/mooncake-wheel/mooncake/structured_object_store.py +++ b/mooncake-wheel/mooncake/structured_object_store.py @@ -1,26 +1,47 @@ from __future__ import annotations import ctypes +import io import json import uuid from concurrent.futures import Future, ThreadPoolExecutor, as_completed from contextlib import contextmanager from dataclasses import dataclass, field -from typing import Any, Iterator, Literal, Mapping, Optional, Protocol, Sequence +from typing import Any, Callable, Iterator, Literal, Mapping, Optional, Protocol, Sequence import numpy as np -DEFAULT_BUNDLE_CHUNK_BYTES = 512 * 1024**2 -AUTO_PARALLEL_MIN_BYTES = 4 * 1024**3 +try: + import mooncake.store as _mooncake_store +except Exception: # pragma: no cover - depends on built extension + _mooncake_store = None # type: ignore[assignment] + +import msgpack as _msgpack + +# -- C fast-path: concat_arrays_into(list[ndarray], dest_ptr) ---------------- +try: + from mooncake._fast_copy import concat_arrays_into as _concat_arrays_into +except Exception: # pragma: no cover + _concat_arrays_into = None + +DEFAULT_BUNDLE_CHUNK_BYTES = 64 * 1024**2 +AUTO_PARALLEL_MIN_BYTES = DEFAULT_BUNDLE_CHUNK_BYTES AUTO_PARALLEL_MIN_CHUNKS = 8 MISSING_OBJECT_ERROR = ( -704 ) # Mooncake remove returns -704 for an already-missing object. STRUCTURED_FIELD_SPECS_KEY = "__mooncake_structured_fields__" +_ENCODING_FALLBACK_ERRORS = ( + TypeError, + ValueError, + OverflowError, + RuntimeError, + RecursionError, +) class BundleStore(Protocol): - def put(self, key: str, value: Any) -> int: ... + def put(self, key: str, value: Any, config: Any = None) -> int: ... def get(self, key: str) -> bytes: ... @@ -33,6 +54,20 @@ class BundleTransferPolicy: max_inflight_put: int = 1 put_mode: Literal["auto", "batch", "parallel"] = "auto" + copy_mode: Literal["auto", "zero_copy", "copy"] = "auto" + + +@dataclass(frozen=True) +class FieldSchema: + """Schema hint for DataProto fields. + + ``metadata["section"]`` may pin a field to ``batch``, + ``non_tensor_batch``, or ``meta_info``. + """ + + codec: str + nullable: bool = True + metadata: dict[str, Any] = field(default_factory=dict) @dataclass @@ -69,6 +104,135 @@ class StructuredObjectResult: objects: dict[str, Any] +@dataclass(frozen=True) +class StructuredFieldLocation: + """Location of a DataProto field inside stage-level structured objects.""" + + stage: str + member: str + section: Literal["batch", "non_tensor_batch"] + + +@dataclass +class MooncakeDataProtoRef: + """Lightweight DataProto handle backed by structured object refs.""" + + batch_size: int + stage_refs: dict[str, RemoteBundleRef] + field_index: dict[str, StructuredFieldLocation] + meta_info: dict[str, Any] + namespace: str = "default" + partition: str = "default" + global_indexes: list[int] | None = None + encoded_non_tensor: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class _TensorPayload: + tensor: Any + + +@dataclass(frozen=True) +class _TensorObjectBufferPayload: + ptr: int + size: int + owner: Any + batch_size: int | None = None + + def __len__(self) -> int: + if self.batch_size is None: + raise TypeError( + "tensor object buffer requires batch_size for DataProto use" + ) + return self.batch_size + + +@dataclass(frozen=True) +class _MultiBufferPayload: + buffers: tuple[memoryview, ...] + owners: tuple[Any, ...] = () + dtype: str | None = None + shape: tuple[int, ...] | None = None + + @property + def nbytes(self) -> int: + return _buffer_group_nbytes(self.buffers) + + +@dataclass(frozen=True) +class _DirectCopyPayload: + """Deferred-copy ndarray list copied directly into pool memory at PUT time.""" + arrays: list[np.ndarray] + total_bytes: int + dtype: str | None = None + shape: tuple[int, ...] | None = None + + @property + def nbytes(self) -> int: + return self.total_bytes + + @staticmethod + def from_flat_arrays( + flat_arrays: list[np.ndarray], dtype: np.dtype, total_elems: int, + ) -> "_DirectCopyPayload": + return _DirectCopyPayload( + arrays=flat_arrays, + total_bytes=sum(a.nbytes for a in flat_arrays), + dtype=np.dtype(dtype).str, + shape=(total_elems,), + ) + + +class _RawDestinationBuffer: + ptr: int + size: int + owner: Any + pre_registered: bool = False + + +class _PoolLeaseOwner: + def __init__(self, lease: Any) -> None: + self.lease = lease + self.released = False + + def release(self) -> None: + if not self.released: + self.lease.release() + self.released = True + + def __del__(self) -> None: + self.release() + + +class _PoolBackedNdarray(np.ndarray): + def __new__( + cls, owner: _PoolLeaseOwner, dtype: np.dtype[Any], shape: tuple[int, ...] + ) -> "_PoolBackedNdarray": + nbytes = ( + int(np.prod(shape, dtype=np.int64)) * dtype.itemsize + if shape + else dtype.itemsize + ) + array = ( + np.ctypeslib.as_array( + (ctypes.c_uint8 * nbytes).from_address(owner.lease.ptr) + ) + .view(dtype) + .reshape(shape) + .view(cls) + ) + array._mooncake_pool_owner = owner + return array + + def __array_finalize__(self, obj: Any) -> None: + # WARNING: slicing/viewing propagates the same _PoolLeaseOwner to the + # derived array. If the original array is GC'd first its __del__ + # releases the pool lease, leaving the slice pointing at freed memory. + # Callers must ensure the original array outlives any derived views. + if obj is not None: + self._mooncake_pool_owner = getattr(obj, "_mooncake_pool_owner", None) + + @dataclass(frozen=True) class StructuredMemberSlice: """Slice description for one structured member.""" @@ -115,103 +279,2797 @@ def slice_member( member_slices=tuple(member_slices.items()), ) - def member_slice(self, name: str) -> StructuredMemberSlice | None: - """Return the configured slice for a member, if any.""" - return dict(self.member_slices).get(name) + def member_slice(self, name: str) -> StructuredMemberSlice | None: + """Return the configured slice for a member, if any.""" + return dict(self.member_slices).get(name) + + +@dataclass(frozen=True) +class _NdarrayReadPlan: + dtype: np.dtype[Any] + full_shape: tuple[int, ...] + output_shape: tuple[int, ...] + byte_offset: int + byte_length: int + step: int + cover_row_count: int + + +@dataclass(frozen=True) +class _DataProtoRowSelection: + count: int + member_slice: StructuredMemberSlice | None = None + indices: tuple[int, ...] | None = None + + +DATAPROTO_REF_HANDLE_TYPE = "mooncake_dataproto_ref" +DATAPROTO_REF_HANDLE_VERSION = 1 +DataProtoRefLike = MooncakeDataProtoRef | Mapping[str, Any] + + +def is_dataproto_ref_handle(value: Any) -> bool: + return isinstance(value, Mapping) and value.get("type") == DATAPROTO_REF_HANDLE_TYPE + + +def export_dataproto_ref(ref: MooncakeDataProtoRef) -> dict[str, Any]: + """Export a DataProto ref as a JSON-safe transport handle.""" + if not isinstance(ref, MooncakeDataProtoRef): + raise TypeError(f"expected MooncakeDataProtoRef, got {type(ref).__name__}") + handle = { + "type": DATAPROTO_REF_HANDLE_TYPE, + "version": DATAPROTO_REF_HANDLE_VERSION, + "kind": "bundle_stages", + "batch_size": int(ref.batch_size), + "namespace": ref.namespace, + "partition": ref.partition, + "global_indexes": _json_safe_value(ref.global_indexes), + "stage_refs": { + stage: {"manifest_key": stage_ref.manifest_key} + for stage, stage_ref in ref.stage_refs.items() + }, + "field_index": { + name: { + "stage": location.stage, + "member": location.member, + "section": location.section, + } + for name, location in ref.field_index.items() + }, + "meta_info": _json_safe_value(ref.meta_info), + "encoded_non_tensor": _json_safe_value(ref.encoded_non_tensor), + } + json.dumps(handle, ensure_ascii=False) + return handle + + +def import_dataproto_ref(handle: Mapping[str, Any]) -> MooncakeDataProtoRef: + """Import a JSON-safe DataProto transport handle into a lazy ref.""" + if not is_dataproto_ref_handle(handle): + raise ValueError("not a Mooncake DataProto ref handle") + if int(handle.get("version", -1)) != DATAPROTO_REF_HANDLE_VERSION: + raise ValueError( + f"unsupported DataProto ref handle version: {handle.get('version')!r}" + ) + if handle.get("kind") != "bundle_stages": + raise ValueError( + f"unsupported DataProto ref handle kind: {handle.get('kind')!r}" + ) + stage_refs: dict[str, RemoteBundleRef] = {} + for stage, stage_ref in _require_mapping( + handle.get("stage_refs"), "stage_refs" + ).items(): + manifest_key = _require_mapping(stage_ref, f"stage_refs[{stage!r}]").get( + "manifest_key" + ) + if not isinstance(stage, str) or not isinstance(manifest_key, str): + raise ValueError( + "DataProto ref handle stage refs must contain string manifest_key values" + ) + stage_refs[stage] = RemoteBundleRef(manifest_key=manifest_key, manifest={}) + field_index: dict[str, StructuredFieldLocation] = {} + for name, raw_location in _require_mapping( + handle.get("field_index"), "field_index" + ).items(): + location = _require_mapping(raw_location, f"field_index[{name!r}]") + section = location.get("section") + if section not in {"batch", "non_tensor_batch"}: + raise ValueError(f"invalid DataProto field section: {section!r}") + stage = location.get("stage") + member = location.get("member") + if ( + not isinstance(name, str) + or not isinstance(stage, str) + or not isinstance(member, str) + ): + raise ValueError("DataProto ref handle field locations must be strings") + if stage not in stage_refs: + raise ValueError( + f"DataProto ref handle field {name!r} references unknown stage {stage!r}" + ) + field_index[name] = StructuredFieldLocation( + stage=stage, member=member, section=section + ) + return MooncakeDataProtoRef( + batch_size=int(handle["batch_size"]), + stage_refs=stage_refs, + field_index=field_index, + meta_info=dict(_require_mapping(handle.get("meta_info", {}), "meta_info")), + namespace=str(handle.get("namespace", "default")), + partition=str(handle.get("partition", "default")), + global_indexes=_import_global_indexes(handle.get("global_indexes")), + encoded_non_tensor=dict( + _require_mapping(handle.get("encoded_non_tensor", {}), "encoded_non_tensor") + ), + ) + + +export_ref = export_dataproto_ref +import_ref = import_dataproto_ref + + +def _resolve_dataproto_ref(ref: DataProtoRefLike) -> MooncakeDataProtoRef: + if isinstance(ref, MooncakeDataProtoRef): + return ref + if is_dataproto_ref_handle(ref): + return import_dataproto_ref(ref) + raise TypeError( + f"expected MooncakeDataProtoRef or DataProto handle, got {type(ref).__name__}" + ) + + +def _require_mapping(value: Any, name: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"DataProto ref handle {name} must be a mapping") + return value + + +def _import_global_indexes(value: Any) -> list[int] | None: + if value is None: + return None + if not isinstance(value, list): + raise ValueError("DataProto ref handle global_indexes must be a list or null") + indexes: list[int] = [] + for index in value: + if isinstance(index, bool) or not isinstance(index, int): + raise ValueError( + "DataProto ref handle global_indexes must contain integers" + ) + indexes.append(index) + return indexes + + +def _json_safe_value(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, np.generic): + return value.item() + if isinstance(value, np.ndarray): + return _json_safe_value(value.tolist()) + if isinstance(value, Mapping): + return {str(key): _json_safe_value(val) for key, val in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe_value(item) for item in value] + raise TypeError( + f"DataProto ref handles support only JSON-safe values, got {type(value).__name__}" + ) + + +class MooncakeBundleTransfer: + """Transfer structured objects through Mooncake, with a low-level bundle fallback.""" + + def __init__( + self, + store: BundleStore, + key_prefix: str = "bundle", + default_chunk_bytes: int = DEFAULT_BUNDLE_CHUNK_BYTES, + buffer_pool: Any = None, + ) -> None: + """Initialize a bundle transfer helper with a configurable default chunk size.""" + self.store = store + self.key_prefix = _normalize_key_prefix(key_prefix) + self.default_chunk_bytes = _validate_chunk_bytes(default_chunk_bytes) + self._transport = _MooncakePayloadTransport(store, buffer_pool=buffer_pool) + self._bundle_store = _BundleManifestStore( + store=store, + transport=self._transport, + key_prefix=self.key_prefix, + default_chunk_bytes=self.default_chunk_bytes, + ) + self._structured_store = _StructuredObjectLayer(self._bundle_store) + + def put_bundle( + self, + meta: bytes | bytearray | memoryview, + buffers: Mapping[str, Any], + partition: str = "default", + chunk_bytes: Optional[int] = None, + policy: Optional[BundleTransferPolicy] = None, + max_inflight_put: Optional[int] = None, + config: Any = None, + ) -> RemoteBundleRef: + """Store raw metadata bytes plus named buffers as a low-level bundle.""" + return self._bundle_store.put_bundle( + meta=meta, + buffers=buffers, + partition=partition, + chunk_bytes=chunk_bytes, + policy=policy, + max_inflight_put=max_inflight_put, + config=config, + ) + + def remove_bundle(self, ref: RemoteBundleRef | Mapping[str, Any]) -> None: + """Remove all Mooncake objects that belong to a stored bundle.""" + self._bundle_store.remove_bundle(ref) + + def put_structured_object( + self, + payload: StructuredObjectPayload, + partition: str = "default", + chunk_bytes: Optional[int] = None, + policy: Optional[BundleTransferPolicy] = None, + max_inflight_put: Optional[int] = None, + config: Any = None, + ) -> RemoteBundleRef: + """Store a structured object described by JSON metadata plus named members.""" + return self._structured_store.put_structured_object( + payload=payload, + partition=partition, + chunk_bytes=chunk_bytes, + policy=policy, + max_inflight_put=max_inflight_put, + config=config, + ) + + def put_object( + self, + obj: Any, + partition: str = "default", + chunk_bytes: Optional[int] = None, + policy: Optional[BundleTransferPolicy] = None, + max_inflight_put: Optional[int] = None, + config: Any = None, + ) -> RemoteBundleRef: + """Store a mapping object or a single tensor/array value.""" + if isinstance(obj, Mapping): + payload = StructuredObjectPayload(metadata={}, buffers=obj) + else: + payload = StructuredObjectPayload( + metadata={"__mooncake_wrapped_object__": True}, + buffers={"value": obj}, + ) + return self.put_structured_object( + payload, + partition=partition, + chunk_bytes=chunk_bytes, + policy=policy, + max_inflight_put=max_inflight_put, + config=config, + ) + + def get_object(self, ref: RemoteBundleRef | Mapping[str, Any]) -> Any: + """Materialize an object stored by put_object().""" + result = self.materialize(self.read_spec(ref)) + if result.metadata.get("__mooncake_wrapped_object__"): + return result.objects["value"] + return result.objects + + def read_spec( + self, ref: RemoteBundleRef | Mapping[str, Any] + ) -> StructuredObjectReadSpec: + """Create a lazy read spec for a structured object reference.""" + return self._structured_store.read_spec(ref) + + def materialize(self, spec: StructuredObjectReadSpec) -> StructuredObjectResult: + """Materialize a structured object read spec.""" + return self._structured_store.materialize(spec) + + def materialize_into( + self, + spec: StructuredObjectReadSpec, + destinations: Optional[Mapping[str, Any]], + ) -> StructuredObjectResult: + """Materialize a structured object read spec into caller-provided destinations when possible.""" + return self._structured_store.materialize_into(spec, destinations) + + @staticmethod + def release_result(result: Any) -> None: + """Release pool-backed buffers in a GET result. + + After get_dataproto / get_dict, ndarray payloads may be backed by + the BufferPool. Call this to release those leases deterministically + instead of waiting for GC ``__del__``. + + Works for both flat dicts and nested envelope dicts + (dataproto: {batch: {...}, non_tensor_batch: {...}, meta_info: {...}}). + """ + released_owners: set[int] = set() + visited_containers: set[int] = set() + + def release_owner(owner: Any) -> None: + owner_id = id(owner) + if owner_id in released_owners: + return + released_owners.add(owner_id) + owner.release() + + def visit(value: Any) -> None: + owner = getattr(value, "_mooncake_pool_owner", None) + if owner is not None: + release_owner(owner) + return + + items_to_visit = None + if isinstance(value, Mapping): + items_to_visit = value.values() + elif isinstance(value, (list, tuple)): + items_to_visit = value + elif isinstance(value, np.ndarray) and value.dtype == object: + items_to_visit = value.flat + + if items_to_visit is not None: + container_id = id(value) + if container_id in visited_containers: + return + visited_containers.add(container_id) + for item in items_to_visit: + visit(item) + + visit(result) + + def put( + self, + data: Any, + *, + type: Literal["dataproto", "dict"] = "dataproto", + namespace: str = "default", + partition: str = "default", + stage: str = "default", + chunk_bytes: Optional[int] = None, + policy: Optional[BundleTransferPolicy] = None, + config: Any = None, + field_schemas: Optional[Mapping[str, FieldSchema]] = None, + ) -> MooncakeDataProtoRef: + """Store a DataProto-like object or flat dict as a structured object.""" + if type == "dataproto": + stage_data = data + elif type == "dict": + stage_data = _flat_dict_to_envelope(data, field_schemas) + else: + raise ValueError(f"unsupported Mooncake payload type: {type!r}") + return self._put_dataproto_stage( + None, + stage_data, + namespace=namespace, + partition=partition, + stage=stage, + chunk_bytes=chunk_bytes, + policy=policy, + overwrite=False, + config=config, + field_schemas=field_schemas, + ) + + def get( + self, + ref: DataProtoRefLike, + *, + type: Literal["dataproto", "dict"] = "dataproto", + fields: Optional[Sequence[str]] = None, + batch_fields: Optional[Sequence[str]] = None, + non_tensor_fields: Optional[Sequence[str]] = None, + meta_info_keys: Optional[Sequence[str]] = None, + data_cls: Optional[Any] = None, + destinations: Optional[Mapping[str, Any]] = None, + rows: slice | StructuredMemberSlice | Sequence[int] | None = None, + ) -> Any: + """Materialize a DataProto-like object or flat dict.""" + if type not in {"dataproto", "dict"}: + raise ValueError(f"unsupported Mooncake payload type: {type!r}") + result = self.get_dataproto( + ref, + fields=fields, + batch_fields=batch_fields, + non_tensor_fields=non_tensor_fields, + meta_info_keys=meta_info_keys, + data_cls=dict if type == "dict" else data_cls, + destinations=destinations, + rows=rows, + ) + return _envelope_to_flat_dict(result) if type == "dict" else result + + + def put_dataproto( + self, + data: Any, + *, + namespace: str = "default", + partition: str = "default", + stage: str = "default", + chunk_bytes: Optional[int] = None, + policy: Optional[BundleTransferPolicy] = None, + field_schemas: Optional[Mapping[str, FieldSchema]] = None, + config: Any = None, + ) -> MooncakeDataProtoRef: + """Store a DataProto-like object as a stage-level structured object.""" + return self.put( + data, + type="dataproto", + namespace=namespace, + partition=partition, + stage=stage, + chunk_bytes=chunk_bytes, + policy=policy, + field_schemas=field_schemas, + config=config, + ) + + def append_dataproto_fields( + self, + ref: DataProtoRefLike, + data: Any, + *, + stage: str, + overwrite: bool = False, + chunk_bytes: Optional[int] = None, + policy: Optional[BundleTransferPolicy] = None, + field_schemas: Optional[Mapping[str, FieldSchema]] = None, + config: Any = None, + ) -> MooncakeDataProtoRef: + """Append DataProto fields, optionally using schema hints for new fields.""" + ref = _resolve_dataproto_ref(ref) + return self._put_dataproto_stage( + ref, + data, + namespace=ref.namespace, + partition=ref.partition, + stage=stage, + chunk_bytes=chunk_bytes, + policy=policy, + overwrite=overwrite, + field_schemas=field_schemas, + config=config, + ) + + def dataproto_manifest_view(self, ref: DataProtoRefLike) -> dict[str, Any]: + """Return a DataProto field view derived from stage structured-object manifests.""" + return _dataproto_manifest_view(self, _resolve_dataproto_ref(ref)) + + def get_dataproto( + self, + ref: DataProtoRefLike, + *, + fields: Optional[Sequence[str]] = None, + batch_fields: Optional[Sequence[str]] = None, + non_tensor_fields: Optional[Sequence[str]] = None, + meta_info_keys: Optional[Sequence[str]] = None, + data_cls: Optional[Any] = None, + destinations: Optional[Mapping[str, Any]] = None, + rows: slice | StructuredMemberSlice | Sequence[int] | None = None, + ) -> Any: + """Materialize selected DataProto fields from structured object refs.""" + ref = _resolve_dataproto_ref(ref) + row_selection = _coerce_dataproto_row_selection(rows, ref.batch_size) + row_slice = None if row_selection is None else row_selection.member_slice + row_indices = None if row_selection is None else row_selection.indices + output_rows = ref.batch_size if row_selection is None else row_selection.count + batch_names, non_tensor_names = _resolve_dataproto_field_selection( + ref, fields, batch_fields, non_tensor_fields + ) + batch: dict[str, Any] = {} + non_tensor_batch: dict[str, Any] = {} + destination_map = destinations or {} + requested = [ + *[("batch", name) for name in batch_names], + *[("non_tensor_batch", name) for name in non_tensor_names], + ] + by_stage: dict[str, list[tuple[str, StructuredFieldLocation]]] = {} + encoded_requests: list[tuple[str, StructuredFieldLocation]] = [] + for _section, name in requested: + location = ref.field_index[name] + if ( + location.section == "non_tensor_batch" + and name in ref.encoded_non_tensor + ): + encoded_requests.append((name, location)) + continue + by_stage.setdefault(location.stage, []).append((name, location)) + for stage, entries in by_stage.items(): + stage_ref = ref.stage_refs[stage] + members = [location.member for _name, location in entries] + stage_destinations = { + location.member: destination_map[name] + for name, location in entries + if name in destination_map + } + if row_indices is None: + spec = self.read_spec(stage_ref).select_members(members) + if row_slice is not None: + for member in members: + spec = spec.slice_member( + member, + axis=row_slice.axis, + start=row_slice.start, + end=row_slice.end, + step=row_slice.step, + ) + result = self.materialize_into(spec, stage_destinations) + for name, location in entries: + value = result.objects[location.member] + if location.section == "batch": + batch[name] = value + else: + non_tensor_batch[name] = value + continue + for name, location in entries: + value = self._read_dataproto_member_indices( + stage_ref, + location.member, + row_indices, + destination_map.get(name), + ) + if location.section == "batch": + batch[name] = value + else: + non_tensor_batch[name] = value + for name, location in encoded_requests: + stage_ref = ref.stage_refs[location.stage] + encoded = ref.encoded_non_tensor[name] + if row_selection is None: + members = list(encoded["payload_members"].values()) + result = self.materialize( + self.read_spec(stage_ref).select_members(members) + ) + payload = { + payload_name: result.objects[member] + for payload_name, member in encoded["payload_members"].items() + } + values = _decode_structured_non_tensor_encoded( + encoded, payload, ref.batch_size, encoded.get("metadata") + ) + elif row_indices is not None: + payload, metadata = self._read_structured_non_tensor_payload_indices( + stage_ref, + encoded, + row_indices, + ) + values = _decode_structured_non_tensor_encoded( + encoded, payload, output_rows, metadata + ) + else: + payload, metadata = self._read_structured_non_tensor_payload_slice( + stage_ref, + encoded, + row_slice, + ref.batch_size, + ) + values = _decode_structured_non_tensor_encoded( + encoded, payload, output_rows, metadata + ) + non_tensor_batch[name] = _object_array_from_decoded_values(values) + meta_info = _select_mapping(ref.meta_info, meta_info_keys) + return _build_dataproto_like_result( + batch, non_tensor_batch, meta_info, data_cls + ) + + def _read_dataproto_member_indices( + self, + stage_ref: RemoteBundleRef, + member: str, + indices: Sequence[int], + destination: Any, + ) -> Any: + manifest = self._bundle_store.resolve_manifest(stage_ref) + metadata = _decode_structured_metadata( + self._bundle_store.read_payload(manifest["meta"]) + ) + field_spec = _structured_field_specs(metadata).get(member, {"encoding": "bytes"}) + payload_spec = manifest["buffers"][member] + encoding = field_spec.get("encoding", "bytes") + if encoding == "ndarray": + return self._read_ndarray_member_indices( + member, payload_spec, field_spec, indices, destination + ) + if encoding == "torch_tensor": + return self._read_torch_tensor_member_indices( + member, payload_spec, field_spec, indices, destination + ) + raise ValueError(f"structured member {member} does not support indexed rows") + + def _read_ndarray_member_indices( + self, + name: str, + payload_spec: Mapping[str, Any], + field_spec: Mapping[str, Any], + indices: Sequence[int], + destination: Any, + ) -> np.ndarray: + dtype = field_spec.get("dtype") + shape = field_spec.get("shape") + if not isinstance(dtype, str) or not isinstance(shape, list): + raise ValueError( + f"structured ndarray field {name} is missing dtype or shape" + ) + dtype_obj = np.dtype(dtype) + full_shape = tuple(int(dim) for dim in shape) + if not full_shape: + raise ValueError("structured ndarray indexed read requires at least one dimension") + output_shape = (len(indices), *full_shape[1:]) + if isinstance(destination, _RawDestinationBuffer): + nbytes = int(np.prod(output_shape, dtype=np.int64)) * dtype_obj.itemsize + if destination.size < nbytes: + raise ValueError( + f"raw destination has {destination.size} bytes, expected at least {nbytes}" + ) + target = np.ctypeslib.as_array( + (ctypes.c_uint8 * nbytes).from_address(destination.ptr) + ).view(dtype_obj).reshape(output_shape) + if not indices: + return target + row_width = dtype_obj.itemsize * int(np.prod(full_shape[1:], dtype=np.int64)) + ranges = [ + (row * row_width, out * row_width, row_width) + for out, row in enumerate(indices) + ] + if destination.pre_registered: + self._bundle_store.read_payload_ranges_into_raw_destination( + payload_spec, destination.ptr, ranges + ) + else: + data = bytearray(nbytes) + self._bundle_store.read_payload_ranges_into_bytearray( + payload_spec, data, ranges + ) + ctypes.memmove(destination.ptr, bytes(data), nbytes) + _ = destination.owner + return target + target = _resolve_ndarray_destination(name, destination, dtype_obj, output_shape) + if not indices: + return target + row_width = dtype_obj.itemsize * int(np.prod(full_shape[1:], dtype=np.int64)) + self._bundle_store.read_payload_ranges_into_array( + payload_spec, + target.view(np.uint8).reshape(-1), + [ + (row * row_width, out * row_width, row_width) + for out, row in enumerate(indices) + ], + ) + return target + + def _read_torch_tensor_member_indices( + self, + name: str, + payload_spec: Mapping[str, Any], + field_spec: Mapping[str, Any], + indices: Sequence[int], + destination: Any, + ) -> Any: + metadata_bytes = int(payload_spec.get("metadata_bytes", -1)) + shape = field_spec.get("shape") + element_size = int(field_spec.get("element_size", 0)) + if metadata_bytes < 0 or not isinstance(shape, list) or element_size <= 0: + raise ValueError( + f"structured tensor field {name} is missing slice metadata" + ) + if not shape: + raise ValueError("structured tensor indexed read requires at least one dimension") + row_width = element_size * int(np.prod(shape[1:], dtype=np.int64)) + data_length = len(indices) * row_width + metadata = self._bundle_store.read_payload_range(payload_spec, 0, metadata_bytes) + sliced_metadata = _slice_tensor_metadata( + metadata, (len(indices), *shape[1:]), data_length + ) + if destination is not None: + if not isinstance(destination, (_TensorObjectBufferPayload, _RawDestinationBuffer)): + raise ValueError( + f"structured tensor member {name} only supports tensor_object_buffer or raw_destination destinations" + ) + expected_bytes = metadata_bytes + data_length + if destination.size < expected_bytes: + raise ValueError( + f"tensor destination has {destination.size} bytes, expected at least {expected_bytes}" + ) + ctypes.memmove(destination.ptr, sliced_metadata, metadata_bytes) + if data_length: + ranges = [ + ( + metadata_bytes + row * row_width, + metadata_bytes + out * row_width, + row_width, + ) + for out, row in enumerate(indices) + ] + if isinstance(destination, _RawDestinationBuffer) and destination.pre_registered: + self._bundle_store.read_payload_ranges_into_raw_destination( + payload_spec, destination.ptr, ranges + ) + else: + data = bytearray(data_length) + self._bundle_store.read_payload_ranges_into_bytearray( + payload_spec, + data, + [ + (source_offset, destination_offset - metadata_bytes, size) + for source_offset, destination_offset, size in ranges + ], + ) + ctypes.memmove( + destination.ptr + metadata_bytes, bytes(data), data_length + ) + _ = destination.owner + return destination + data = bytearray(data_length) + if data_length: + self._bundle_store.read_payload_ranges_into_bytearray( + payload_spec, + data, + [(metadata_bytes + row * row_width, out * row_width, row_width) for out, row in enumerate(indices)], + ) + return _deserialize_tensor_payload(sliced_metadata + data) + + def _read_structured_non_tensor_payload_indices( + self, + stage_ref: RemoteBundleRef, + encoded: Mapping[str, Any], + indices: Sequence[int], + ) -> tuple[dict[str, Any], Mapping[str, Any]]: + codec = encoded["codec"] + payload_members = encoded["payload_members"] + manifest = self._bundle_store.resolve_manifest(stage_ref) + stage_metadata = _decode_structured_metadata( + self._bundle_store.read_payload(manifest["meta"]) + ) + field_specs = _structured_field_specs(stage_metadata) + metadata = dict(encoded.get("metadata") or {}) + + def member(payload_name: str) -> str: + return payload_members[payload_name] + + def member_dtype(payload_name: str) -> np.dtype[Any]: + spec = field_specs.get(member(payload_name), {}) + dtype = spec.get("dtype") + if not isinstance(dtype, str): + raise ValueError( + f"structured non-tensor payload {member(payload_name)} is missing dtype" + ) + return np.dtype(dtype) + + def read_member_indices(payload_name: str, member_indices: Sequence[int]) -> Any: + return self._read_dataproto_member_indices( + stage_ref, member(payload_name), member_indices, None + ) + + def read_data_ranges( + payload_name: str, + ranges: Sequence[tuple[int, int, int]], + dtype: np.dtype[Any] | None = None, + ) -> Any: + total_bytes = sum(byte_length for _src, _dst, byte_length in ranges) + payload_spec = manifest["buffers"][member(payload_name)] + if dtype is None: + data = bytearray(total_bytes) + self._bundle_store.read_payload_ranges_into_bytearray( + payload_spec, data, ranges + ) + return bytes(data) + array = np.empty(total_bytes // dtype.itemsize, dtype=dtype) + self._bundle_store.read_payload_ranges_into_array( + payload_spec, array.view(np.uint8).reshape(-1), ranges + ) + return array + + if _is_recursive_encoded_non_tensor(encoded): + metadata = _copy_recursive_metadata_for_leaf_updates(metadata) + recursive_payload: dict[str, Any] = {} + for node in metadata.get("nodes", []): + for key in ("missing_payload", "row_mask_payload", "lengths_payload"): + payload_name = node.get(key) + if payload_name is not None: + recursive_payload[payload_name] = read_member_indices( + payload_name, indices + ) + for leaf in metadata.get("leaves", []): + leaf_payload_members = leaf["payload_members"] + flat_payload_members = { + name: payload_members[global_name] + for name, global_name in leaf_payload_members.items() + if name != "missing" + } + leaf_payload, _leaf_metadata = self._read_structured_non_tensor_payload_indices( + stage_ref, + { + "codec": leaf["codec"], + "metadata": leaf.get("metadata") or {}, + "payload_members": flat_payload_members, + }, + indices, + ) + leaf["metadata"] = _leaf_metadata + for name, value in leaf_payload.items(): + recursive_payload[leaf_payload_members[name]] = value + recursive_payload[leaf_payload_members["missing"]] = read_member_indices( + leaf_payload_members["missing"], indices + ) + return recursive_payload, metadata + + if codec == "ndarray": + return { + "data": read_member_indices("data", indices), + "nulls": read_member_indices("nulls", indices), + }, metadata + + if codec in {"ragged_tensor", "typed_ragged"}: + offsets = read_member_indices( + "offsets", [index for row in indices for index in (row, row + 1)] + ) + begins = offsets[0::2] + ends = offsets[1::2] + item_counts = [int(end) - int(begin) for begin, end in zip(begins, ends)] + gathered_offsets = np.empty(len(indices) + 1, dtype=offsets.dtype) + gathered_offsets[0] = 0 + for index, count in enumerate(item_counts): + gathered_offsets[index + 1] = int(gathered_offsets[index]) + count + dtype = member_dtype("data") + ranges = [] + destination_item = 0 + for begin, count in zip(begins, item_counts): + if count: + ranges.append( + ( + int(begin) * dtype.itemsize, + destination_item * dtype.itemsize, + count * dtype.itemsize, + ) + ) + destination_item += count + return { + "data": read_data_ranges("data", ranges, dtype), + "offsets": gathered_offsets, + "shapes": read_member_indices("shapes", indices), + "ndims": read_member_indices("ndims", indices), + "nulls": read_member_indices("nulls", indices), + }, metadata + + if codec == "ragged_tensor_dict": + return self._read_ragged_tensor_dict_payload( + payload_members, + metadata, + read_null_mask=lambda: read_member_indices("null_mask", indices), + read_key_payload=lambda encoded: self._read_structured_non_tensor_payload_indices( + stage_ref, encoded, indices + ), + ) + + if codec in {"media_bytes", "bytes_ragged", "utf8_ragged", "msgpack_ragged", "json_ragged"}: + offsets = read_member_indices( + "offsets", [index for row in indices for index in (row, row + 1)] + ) + begins = offsets[0::2] + ends = offsets[1::2] + byte_counts = [int(end) - int(begin) for begin, end in zip(begins, ends)] + gathered_offsets = np.empty(len(indices) + 1, dtype=offsets.dtype) + gathered_offsets[0] = 0 + ranges = [] + destination_offset = 0 + for index, (begin, count) in enumerate(zip(begins, byte_counts)): + gathered_offsets[index + 1] = destination_offset + count + if count: + ranges.append((int(begin), destination_offset, count)) + destination_offset += count + if "media_encodings" in metadata: + metadata["media_encodings"] = [ + metadata["media_encodings"][index] for index in indices + ] + return { + "data": read_data_ranges("data", ranges), + "offsets": gathered_offsets, + "nulls": read_member_indices("nulls", indices), + }, metadata + + if codec == "media_list_ragged": + row_offsets = read_member_indices( + "row_offsets", [index for row in indices for index in (row, row + 1)] + ) + item_begins = row_offsets[0::2] + item_ends = row_offsets[1::2] + item_counts = [ + int(end) - int(begin) for begin, end in zip(item_begins, item_ends) + ] + gathered_row_offsets = np.empty(len(indices) + 1, dtype=row_offsets.dtype) + gathered_row_offsets[0] = 0 + for index, count in enumerate(item_counts): + gathered_row_offsets[index + 1] = int(gathered_row_offsets[index]) + count + boundary_indices = [ + boundary + for begin, end in zip(item_begins, item_ends) + for boundary in range(int(begin), int(end) + 1) + ] + byte_offsets = read_member_indices("byte_offsets", boundary_indices) + ranges = [] + gathered_byte_offsets = np.empty( + int(gathered_row_offsets[-1]) + 1, dtype=byte_offsets.dtype + ) + gathered_byte_offsets[0] = 0 + destination_offset = 0 + source_index = 0 + destination_boundary = 1 + media_encodings = [] + for begin, count in zip(item_begins, item_counts): + for item in range(count): + byte_begin = int(byte_offsets[source_index + item]) + byte_end = int(byte_offsets[source_index + item + 1]) + byte_count = byte_end - byte_begin + if byte_count: + ranges.append((byte_begin, destination_offset, byte_count)) + destination_offset += byte_count + gathered_byte_offsets[destination_boundary] = destination_offset + destination_boundary += 1 + if "media_encodings" in metadata: + media_encodings.extend( + metadata["media_encodings"][int(begin) : int(begin) + count] + ) + source_index += count + 1 + if "media_encodings" in metadata: + metadata["media_encodings"] = media_encodings + return { + "data": read_data_ranges("data", ranges), + "row_offsets": gathered_row_offsets, + "byte_offsets": gathered_byte_offsets, + "nulls": read_member_indices("nulls", indices), + }, metadata + + raise ValueError(f"unknown structured non-tensor codec: {codec}") + + def _read_structured_non_tensor_payload_slice( + self, + stage_ref: RemoteBundleRef, + encoded: Mapping[str, Any], + row_slice: StructuredMemberSlice, + total_rows: int, + ) -> tuple[dict[str, Any], Mapping[str, Any]]: + start, end, step = _normalized_member_slice(row_slice, total_rows) + if step != 1: + raise ValueError( + "DataProto structured non-tensor slicing currently requires step=1" + ) + codec = encoded["codec"] + payload_members = encoded["payload_members"] + manifest = self._bundle_store.resolve_manifest(stage_ref) + metadata = dict(encoded.get("metadata") or {}) + + def read_member(payload_name: str, row_start: int, row_end: int) -> Any: + member = payload_members[payload_name] + spec = ( + self.read_spec(stage_ref) + .select_members([member]) + .slice_member(member, axis=0, start=row_start, end=row_end) + ) + return self.materialize(spec).objects[member] + + def read_bytes(payload_name: str, byte_start: int, byte_end: int) -> bytes: + member = payload_members[payload_name] + payload_spec = manifest["buffers"][member] + return self._bundle_store.read_payload_range( + payload_spec, byte_start, byte_end - byte_start + ) + + if _is_recursive_encoded_non_tensor(encoded): + metadata = _copy_recursive_metadata_for_leaf_updates(metadata) + recursive_payload: dict[str, Any] = {} + for node in metadata.get("nodes", []): + for key in ("missing_payload", "row_mask_payload", "lengths_payload"): + payload_name = node.get(key) + if payload_name is not None: + recursive_payload[payload_name] = read_member(payload_name, start, end) + for leaf in metadata.get("leaves", []): + leaf_payload_members = leaf["payload_members"] + flat_payload_members = { + name: payload_members[global_name] + for name, global_name in leaf_payload_members.items() + if name != "missing" + } + leaf_payload, _leaf_metadata = self._read_structured_non_tensor_payload_slice( + stage_ref, + { + "codec": leaf["codec"], + "metadata": leaf.get("metadata") or {}, + "payload_members": flat_payload_members, + }, + row_slice, + total_rows, + ) + leaf["metadata"] = _leaf_metadata + for name, value in leaf_payload.items(): + recursive_payload[leaf_payload_members[name]] = value + recursive_payload[leaf_payload_members["missing"]] = read_member( + leaf_payload_members["missing"], start, end + ) + return recursive_payload, metadata + + if codec == "ndarray": + return { + "data": read_member("data", start, end), + "nulls": read_member("nulls", start, end), + }, metadata + + if codec in {"ragged_tensor", "typed_ragged"}: + offsets = read_member("offsets", start, end + 1) + base = int(offsets[0]) + limit = int(offsets[-1]) + offsets = offsets - base + return { + "data": read_member("data", base, limit), + "offsets": offsets, + "shapes": read_member("shapes", start, end), + "ndims": read_member("ndims", start, end), + "nulls": read_member("nulls", start, end), + }, metadata + + if codec == "ragged_tensor_dict": + return self._read_ragged_tensor_dict_payload( + payload_members, + metadata, + read_null_mask=lambda: read_member("null_mask", start, end), + read_key_payload=lambda encoded: self._read_structured_non_tensor_payload_slice( + stage_ref, encoded, row_slice, total_rows + ), + ) + + if codec in {"media_bytes", "bytes_ragged", "utf8_ragged", "msgpack_ragged", "json_ragged"}: + offsets = read_member("offsets", start, end + 1) + base = int(offsets[0]) + limit = int(offsets[-1]) + offsets = offsets - base + if "media_encodings" in metadata: + metadata["media_encodings"] = metadata["media_encodings"][start:end] + return { + "data": read_bytes("data", base, limit), + "offsets": offsets, + "nulls": read_member("nulls", start, end), + }, metadata + + if codec == "media_list_ragged": + row_offsets = read_member("row_offsets", start, end + 1) + item_start = int(row_offsets[0]) + item_end = int(row_offsets[-1]) + row_offsets = row_offsets - item_start + byte_offsets = read_member("byte_offsets", item_start, item_end + 1) + byte_start = int(byte_offsets[0]) + byte_end = int(byte_offsets[-1]) + byte_offsets = byte_offsets - byte_start + if "media_encodings" in metadata: + metadata["media_encodings"] = metadata["media_encodings"][ + item_start:item_end + ] + return { + "data": read_bytes("data", byte_start, byte_end), + "row_offsets": row_offsets, + "byte_offsets": byte_offsets, + "nulls": read_member("nulls", start, end), + }, metadata + + raise ValueError(f"unknown structured non-tensor codec: {codec}") + + def _read_ragged_tensor_dict_payload( + self, + payload_members: Mapping[str, str], + metadata: dict[str, Any], + *, + read_null_mask: Callable[[], Any], + read_key_payload: Callable[ + [Mapping[str, Any]], tuple[dict[str, Any], Mapping[str, Any]] + ], + ) -> tuple[dict[str, Any], Mapping[str, Any]]: + key_codecs = dict(metadata.get("key_codecs") or {}) + metadata["key_codecs"] = key_codecs + dict_payload: dict[str, Any] = {"null_mask": read_null_mask()} + for key in _normalize_ragged_tensor_dict_keys(metadata.get("keys", [])): + key_members = _ragged_tensor_dict_payload_items( + payload_members, key, kind="manifest" + ) + key_payload, key_metadata = read_key_payload( + { + "codec": "ragged_tensor", + "metadata": key_codecs.get(key) or {}, + "payload_members": key_members, + } + ) + key_codecs[key] = key_metadata + for name, value in key_payload.items(): + dict_payload[f"{key}.{name}"] = value + return dict_payload, metadata + + def cleanup_dataproto(self, ref: DataProtoRefLike) -> None: + """Remove all structured object stages referenced by a DataProto handle.""" + ref = _resolve_dataproto_ref(ref) + seen: set[str] = set() + for stage_ref in ref.stage_refs.values(): + if stage_ref.manifest_key in seen: + continue + seen.add(stage_ref.manifest_key) + self.remove_bundle(stage_ref) + + def _append_dataproto_stage_manifest( + self, + old_stage_ref: RemoteBundleRef, + payload: StructuredObjectPayload, + *, + partition: str, + chunk_bytes: Optional[int], + policy: Optional[BundleTransferPolicy], + config: Any = None, + ) -> RemoteBundleRef: + new_stage_ref = self.put_structured_object( + payload, + partition=partition, + chunk_bytes=chunk_bytes, + policy=policy, + config=config, + ) + try: + old_manifest = self._bundle_store.resolve_manifest(old_stage_ref) + new_manifest = self._bundle_store.resolve_manifest(new_stage_ref) + old_metadata = _decode_structured_metadata( + self._bundle_store.read_payload(old_manifest["meta"]) + ) + new_metadata = _decode_structured_metadata( + self._bundle_store.read_payload(new_manifest["meta"]) + ) + merged_metadata = _merge_structured_stage_metadata( + old_metadata, new_metadata + ) + merged_buffers = dict(old_manifest["buffers"]) + collisions = sorted(set(merged_buffers) & set(new_manifest["buffers"])) + if collisions: + raise ValueError(f"structured members already exist: {collisions}") + merged_buffers.update(new_manifest["buffers"]) + merged_ref = self._bundle_store.put_bundle_manifest( + _encode_structured_metadata(merged_metadata), + merged_buffers, + partition=partition, + chunk_bytes=chunk_bytes, + policy=policy, + cleanup_keys=[ + self._bundle_store.manifest_key(old_stage_ref), + *self._bundle_store.payload_keys(old_manifest["meta"]), + ], + config=config, + ) + except Exception: + self.remove_bundle(new_stage_ref) + raise + obsolete_keys = [ + self._bundle_store.manifest_key(new_stage_ref), + *self._bundle_store.payload_keys(new_manifest["meta"]), + ] + self._bundle_store.remove_keys(obsolete_keys, strict=False) + return merged_ref + + def _put_dataproto_stage( + self, + ref: MooncakeDataProtoRef | None, + data: Any, + *, + namespace: str, + partition: str, + stage: str, + chunk_bytes: Optional[int], + policy: Optional[BundleTransferPolicy], + overwrite: bool, + field_schemas: Optional[Mapping[str, FieldSchema]] = None, + config: Any = None, + ) -> MooncakeDataProtoRef: + batch, non_tensor_batch, meta_info = _split_dataproto_like(data) + _validate_dataproto_schema_sections( + batch, non_tensor_batch, meta_info, field_schemas + ) + batch_size = _dataproto_batch_size(batch, non_tensor_batch) + if ref is not None and batch_size != ref.batch_size: + raise ValueError( + f"DataProto append batch size {batch_size} does not match ref batch size {ref.batch_size}" + ) + cross_section_fields = sorted(set(batch) & set(non_tensor_batch)) + if cross_section_fields: + raise ValueError( + f"DataProto batch and non_tensor_batch fields overlap: {cross_section_fields}" + ) + buffers: dict[str, Any] = {} + field_updates: dict[str, StructuredFieldLocation] = {} + for name, value in batch.items(): + member = f"batch.{name}" + buffers[member] = value + field_updates[name] = StructuredFieldLocation(stage, member, "batch") + encoded_updates: dict[str, Any] = {} + for name, value in non_tensor_batch.items(): + schema = _schema_for_section( + field_schemas, name, "non_tensor_batch" + ) + encoded: _EncodedStructuredLeaf | None = None + if schema is not None: + try: + encode_value = _coerce_schema_non_tensor_value(name, value) + if schema.codec == "auto": + _validate_schema_nullable( + f"non_tensor_batch.{name}", encode_value, schema + ) + else: + encoded = _encode_with_schema( + f"non_tensor_batch.{name}", encode_value, schema + ) + except (TypeError, ValueError, RuntimeError, AttributeError) as exc: + raise type(exc)( + f"failed to encode non_tensor_batch field {name!r} " + f"with FieldSchema codec {schema.codec!r}: {exc}" + ) from exc + if encoded is None and _should_encode_non_tensor_field(value): + encoded = _encode_structured_non_tensor_field( + f"non_tensor_batch.{name}", value + ) + if encoded is not None: + payload_members: dict[str, str] = {} + for payload_name, payload_value in encoded.payload.items(): + member = f"non_tensor_batch.{name}.{payload_name}" + buffers[member] = payload_value + payload_members[payload_name] = member + field_updates[name] = StructuredFieldLocation( + stage, f"non_tensor_batch.{name}", "non_tensor_batch" + ) + encoded_updates[name] = { + "codec": encoded.codec, + "rows": encoded.rows, + "metadata": encoded.metadata, + "payload_members": payload_members, + } + continue + member = f"non_tensor_batch.{name}" + buffers[member] = value + field_updates[name] = StructuredFieldLocation( + stage, member, "non_tensor_batch" + ) + if not buffers: + if ref is not None: + merged_meta_info = dict(ref.meta_info) + merged_meta_info.update(meta_info) + return MooncakeDataProtoRef( + batch_size=ref.batch_size, + stage_refs=dict(ref.stage_refs), + field_index=dict(ref.field_index), + meta_info=merged_meta_info, + namespace=ref.namespace, + partition=ref.partition, + global_indexes=ref.global_indexes, + encoded_non_tensor=dict(ref.encoded_non_tensor), + ) + return MooncakeDataProtoRef( + batch_size=batch_size, + stage_refs={}, + field_index={}, + meta_info=dict(meta_info), + namespace=namespace, + partition=partition, + encoded_non_tensor={}, + ) + duplicates = ( + sorted(set(ref.field_index) & set(field_updates)) if ref is not None else [] + ) + if duplicates and not overwrite: + raise ValueError(f"DataProto fields already exist: {duplicates}") + if ref is not None and overwrite: + dangling = [ + name + for name, location in ref.field_index.items() + if location.stage == stage and name not in field_updates + ] + if dangling: + raise ValueError( + f"DataProto overwrite for stage {stage!r} must include existing fields: {sorted(dangling)}" + ) + payload = StructuredObjectPayload( + metadata={ + "layout": "dataproto_stage", + "dataproto": { + "version": 1, + "namespace": namespace, + "partition": partition, + "stage": stage, + "batch_size": batch_size, + }, + }, + buffers=buffers, + ) + existing_stage_ref = ref.stage_refs.get(stage) if ref is not None else None + if existing_stage_ref is not None and not overwrite: + stage_ref = self._append_dataproto_stage_manifest( + existing_stage_ref, + payload, + partition=partition, + chunk_bytes=chunk_bytes, + policy=policy, + config=config, + ) + else: + stage_ref = self.put_structured_object( + payload, + partition=partition, + chunk_bytes=chunk_bytes, + policy=policy, + config=config, + ) + if ( + existing_stage_ref is not None + and overwrite + and existing_stage_ref.manifest_key != stage_ref.manifest_key + ): + self.remove_bundle(existing_stage_ref) + if ref is None: + stage_refs = {stage: stage_ref} + field_index = dict(field_updates) + merged_meta_info = dict(meta_info) + encoded_non_tensor = dict(encoded_updates) + else: + stage_refs = dict(ref.stage_refs) + stage_refs[stage] = stage_ref + field_index = dict(ref.field_index) + field_index.update(field_updates) + merged_meta_info = dict(ref.meta_info) + merged_meta_info.update(meta_info) + encoded_non_tensor = dict(ref.encoded_non_tensor) + for name in field_updates: + encoded_non_tensor.pop(name, None) + encoded_non_tensor.update(encoded_updates) + return MooncakeDataProtoRef( + batch_size=batch_size, + stage_refs=stage_refs, + field_index=field_index, + meta_info=merged_meta_info, + namespace=namespace, + partition=partition, + global_indexes=None if ref is None else ref.global_indexes, + encoded_non_tensor=encoded_non_tensor, + ) + + +def _copy_transfer_policy(policy: BundleTransferPolicy) -> BundleTransferPolicy: + if policy.copy_mode == "copy": + return policy + return BundleTransferPolicy( + max_inflight_put=policy.max_inflight_put, + put_mode=policy.put_mode, + copy_mode="copy", + ) + + +def _dataproto_manifest_view( + transfer: MooncakeBundleTransfer, ref: MooncakeDataProtoRef +) -> dict[str, Any]: + stage_manifests = { + stage: transfer._bundle_store.resolve_manifest(stage_ref) + for stage, stage_ref in ref.stage_refs.items() + } + stage_metadata = { + stage: _decode_structured_metadata( + transfer._bundle_store.read_payload(stage_manifest["meta"]) + ) + for stage, stage_manifest in stage_manifests.items() + } + fields: dict[str, dict[str, Any]] = {} + for name, location in ref.field_index.items(): + metadata = stage_metadata[location.stage] + field_specs = _structured_field_specs(metadata) + member_spec = dict(field_specs.get(location.member, {"encoding": "bytes"})) + encoded = ref.encoded_non_tensor.get(name) + if encoded is not None: + member_spec = { + "encoding": "structured_non_tensor", + "codec": encoded["codec"], + "metadata": encoded.get("metadata"), + "payload_members": dict(encoded["payload_members"]), + "payload_specs": { + payload_name: dict(field_specs.get(member, {"encoding": "bytes"})) + for payload_name, member in encoded["payload_members"].items() + }, + } + fields[name] = { + "section": location.section, + "stage": location.stage, + "member": location.member, + "spec": member_spec, + } + return { + "namespace": ref.namespace, + "partition": ref.partition, + "batch_size": ref.batch_size, + "batch_fields": { + name: info for name, info in fields.items() if info["section"] == "batch" + }, + "non_tensor_fields": { + name: info + for name, info in fields.items() + if info["section"] == "non_tensor_batch" + }, + "meta_info_keys": list(ref.meta_info), + "stages": { + stage: { + "manifest_key": stage_ref.manifest_key, + "dataproto": stage_metadata[stage].get("dataproto", {}), + } + for stage, stage_ref in ref.stage_refs.items() + }, + } + + +def _flat_dict_to_envelope( + data: Mapping[str, Any], + field_schemas: Optional[Mapping[str, FieldSchema]] = None, +) -> dict[str, Any]: + if not isinstance(data, Mapping): + raise TypeError("flat dict payload must be a mapping") + if not field_schemas: + field_schemas = {} + schema_row_count = _flat_dict_schema_row_count(data, field_schemas) + row_count = schema_row_count + if row_count == 0: + schema_meta_fields = { + name + for name, schema in field_schemas.items() + if name in data and _schema_section(name, schema) == "meta_info" + } + row_count = _flat_dict_auto_row_count(data, exclude=schema_meta_fields) + batch: dict[str, Any] = {} + non_tensor_batch: dict[str, Any] = {} + meta_info: dict[str, Any] = {} + for key, value in data.items(): + schema = field_schemas.get(key) + section = None if schema is None else _schema_section(key, schema) + if section is None: + if _is_row_aligned_dense_field(value, row_count): + batch[key] = value + elif ( + schema_row_count == 0 + and _is_non_string_sequence(value) + and len(value) == row_count + ): + non_tensor_batch[key] = _coerce_flat_dict_non_tensor_field( + key, value, row_count, schema + ) + else: + meta_info[key] = value + continue + if section == "batch": + batch[key] = value + elif section == "non_tensor_batch": + non_tensor_batch[key] = _coerce_flat_dict_non_tensor_field( + key, value, row_count, schema + ) + else: + meta_info[key] = value + return { + "batch": batch, + "non_tensor_batch": non_tensor_batch, + "meta_info": meta_info, + } + + +def _flat_dict_schema_row_count( + data: Mapping[str, Any], field_schemas: Mapping[str, FieldSchema] +) -> int: + sizes = { + _field_len(name, data[name]) + for name, schema in field_schemas.items() + if name in data and _schema_section(name, schema) in {"batch", "non_tensor_batch"} + } + if not sizes: + return 0 + if len(sizes) != 1: + raise ValueError(f"flat dict fields have inconsistent batch sizes: {sorted(sizes)}") + return sizes.pop() + + +def _field_len(name: str, value: Any) -> int: + try: + return len(value) + except TypeError as error: + raise ValueError(f"flat dict row-aligned field {name!r} must be sized") from error + + +def _flat_dict_auto_row_count( + data: Mapping[str, Any], exclude: set[str] | frozenset[str] = frozenset() +) -> int: + dense_sizes = { + len(value) + for key, value in data.items() + if key not in exclude + and ( + (_torch is not None and isinstance(value, _torch.Tensor) and value.ndim > 0) + or (isinstance(value, np.ndarray) and value.dtype != object and value.ndim > 0) + ) + } + if len(dense_sizes) > 1: + raise ValueError( + f"flat dict dense fields have ambiguous batch sizes: {sorted(dense_sizes)}" + ) + if dense_sizes: + return dense_sizes.pop() + + sizes = { + len(value) + for key, value in data.items() + if key not in exclude + and ( + (isinstance(value, np.ndarray) and value.ndim > 0) + or _is_non_string_sequence(value) + ) + } + if not sizes: + return 0 + if len(sizes) != 1: + raise ValueError( + f"flat dict fields have ambiguous batch sizes: {sorted(sizes)}; " + "pass FieldSchema metadata['section'] for list-valued metadata fields" + ) + return sizes.pop() + + +def _coerce_flat_dict_non_tensor_field( + name: str, value: Any, row_count: int, schema: Optional[FieldSchema] = None +) -> Any: + if isinstance(value, np.ndarray): + if len(value) != row_count: + raise ValueError( + f"flat dict non_tensor_batch field {name!r} has batch size {len(value)}, expected {row_count}" + ) + schema_dtype = _schema_ndarray_dtype(schema) if schema is not None else None + if ( + schema is not None + and schema.codec == "ndarray" + and schema_dtype is not None + and (value.dtype != object or all(item is not None for item in value)) + ): + return np.asarray(value, dtype=schema_dtype) + return value + if _is_non_string_sequence(value): + if len(value) != row_count: + raise ValueError( + f"flat dict non_tensor_batch field {name!r} has batch size {len(value)}, expected {row_count}" + ) + array = np.empty(row_count, dtype=object) + array[:] = list(value) + return array + raise TypeError( + f"flat dict non_tensor_batch field {name!r} must be an ndarray or non-string sequence" + ) + + +def _is_non_string_sequence(value: Any) -> bool: + return isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)) + + +def _is_row_aligned_dense_field(value: Any, row_count: int) -> bool: + if row_count == 0: + return False + if _torch is not None and isinstance(value, _torch.Tensor): + return value.shape[:1] == (row_count,) + return ( + isinstance(value, np.ndarray) + and value.dtype != object + and value.shape[:1] == (row_count,) + ) + + +def _envelope_to_flat_dict(data: Mapping[str, Any]) -> dict[str, Any]: + meta_info = _mapping_to_dict(data.get("meta_info")) + batch = _mapping_to_dict(data.get("batch")) + non_tensor_batch = _mapping_to_dict(data.get("non_tensor_batch")) + overlap = ( + (set(meta_info) & set(batch)) + | (set(meta_info) & set(non_tensor_batch)) + | (set(batch) & set(non_tensor_batch)) + ) + if overlap: + raise ValueError( + f"Duplicate keys found across DataProto sections: {sorted(overlap)}" + ) + + result = dict(meta_info) + result.update(batch) + for name, value in non_tensor_batch.items(): + result[name] = ( + list(value) + if isinstance(value, np.ndarray) and value.dtype == object + else value + ) + return result + +def _build_dataproto_like_result( + batch: dict[str, Any], + non_tensor_batch: dict[str, Any], + meta_info: dict[str, Any], + data_cls: Optional[Any], +) -> Any: + payload = { + "batch": batch, + "non_tensor_batch": non_tensor_batch, + "meta_info": meta_info, + } + if data_cls is None or data_cls is dict: + return payload + from_dict = getattr(data_cls, "from_dict", None) + try: + if callable(from_dict): + return from_dict(batch, non_tensor_batch, meta_info=meta_info) + return data_cls( + batch=batch, non_tensor_batch=non_tensor_batch, meta_info=meta_info + ) + except TypeError as error: + raise TypeError( + f"{getattr(data_cls, '__name__', data_cls)!r} cannot be constructed from " + "batch, non_tensor_batch, and meta_info" + ) from error + + +def _split_dataproto_like( + data: Any, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: + if isinstance(data, Mapping): + if _is_dataproto_envelope_mapping(data): + return ( + _mapping_to_dict(data.get("batch")), + _mapping_to_dict(data.get("non_tensor_batch")), + _mapping_to_dict(data.get("meta_info")), + ) + return dict(data), {}, {} + batch = _mapping_to_dict(getattr(data, "batch", None)) + non_tensor_batch = _mapping_to_dict(getattr(data, "non_tensor_batch", None)) + meta_info = _mapping_to_dict(getattr(data, "meta_info", None)) + return batch, non_tensor_batch, meta_info + + +def _is_dataproto_envelope_mapping(data: Mapping[str, Any]) -> bool: + envelope_keys = {"batch", "non_tensor_batch", "meta_info"} + if not data or not set(data).issubset(envelope_keys): + return False + return all(_is_mapping_like_or_none(value) for value in data.values()) + + +def _is_mapping_like_or_none(value: Any) -> bool: + return ( + value is None + or isinstance(value, Mapping) + or callable(getattr(value, "items", None)) + ) + + +def _mapping_to_dict(value: Any) -> dict[str, Any]: + if value is None: + return {} + if isinstance(value, Mapping): + return dict(value) + items = getattr(value, "items", None) + if callable(items): + return dict(items()) + raise TypeError(f"expected mapping-like value, got {type(value).__name__}") + + +def _dataproto_batch_size( + batch: Mapping[str, Any], non_tensor_batch: Mapping[str, Any] +) -> int: + sizes = { + len(value) for values in (batch, non_tensor_batch) for value in values.values() + } + if not sizes: + return 0 + if len(sizes) != 1: + raise ValueError( + f"DataProto fields have inconsistent batch sizes: {sorted(sizes)}" + ) + return sizes.pop() + + +def _resolve_dataproto_field_selection( + ref: MooncakeDataProtoRef, + fields: Optional[Sequence[str]], + batch_fields: Optional[Sequence[str]], + non_tensor_fields: Optional[Sequence[str]], +) -> tuple[list[str], list[str]]: + if fields is not None and ( + batch_fields is not None or non_tensor_fields is not None + ): + raise ValueError( + "fields cannot be combined with batch_fields or non_tensor_fields" + ) + if fields is not None: + requested = list(fields) + _validate_dataproto_fields_exist(ref, requested) + return ( + [name for name in requested if ref.field_index[name].section == "batch"], + [ + name + for name in requested + if ref.field_index[name].section == "non_tensor_batch" + ], + ) + if batch_fields is None and non_tensor_fields is None: + batch_names = [ + name + for name, location in ref.field_index.items() + if location.section == "batch" + ] + non_tensor_names = [ + name + for name, location in ref.field_index.items() + if location.section == "non_tensor_batch" + ] + else: + batch_names = [] if batch_fields is None else list(batch_fields) + non_tensor_names = ( + [] if non_tensor_fields is None else list(non_tensor_fields) + ) + _validate_dataproto_fields_exist(ref, [*batch_names, *non_tensor_names]) + return batch_names, non_tensor_names + + +def _validate_dataproto_fields_exist( + ref: MooncakeDataProtoRef, names: Sequence[str] +) -> None: + missing = [name for name in names if name not in ref.field_index] + if missing: + raise KeyError(f"unknown DataProto fields: {missing}") + + +def _coerce_dataproto_row_selection( + rows: slice | StructuredMemberSlice | Sequence[int] | None, total_rows: int +) -> _DataProtoRowSelection | None: + if rows is None: + return None + if isinstance(rows, StructuredMemberSlice): + if rows.axis != 0: + raise ValueError("DataProto row slicing currently supports axis=0 only") + _normalized_member_slice(rows, total_rows) + return _DataProtoRowSelection( + count=_slice_length(rows, total_rows), member_slice=rows + ) + if isinstance(rows, slice): + start, end, step = rows.indices(total_rows) + if step <= 0: + raise ValueError("DataProto row slicing step must be positive") + member_slice = StructuredMemberSlice(axis=0, start=start, end=end, step=step) + return _DataProtoRowSelection( + count=_slice_length(member_slice, total_rows), member_slice=member_slice + ) + if isinstance(rows, Sequence) and not isinstance(rows, (str, bytes, bytearray)): + indices = tuple(_normalize_dataproto_row_index(index, total_rows) for index in rows) + return _DataProtoRowSelection(count=len(indices), indices=indices) + raise TypeError("DataProto rows must be a slice, StructuredMemberSlice, or row index sequence") + + +def _normalize_dataproto_row_index(index: Any, total_rows: int) -> int: + if not isinstance(index, (int, np.integer)): + raise TypeError("DataProto row indices must be integers") + normalized = int(index) + if normalized < 0: + normalized += total_rows + if normalized < 0 or normalized >= total_rows: + raise IndexError(f"DataProto row index {index} out of range for {total_rows} rows") + return normalized + + +def _slice_length(member_slice: StructuredMemberSlice, total_rows: int) -> int: + start, end, step = _normalized_member_slice(member_slice, total_rows) + if start >= end: + return 0 + return 1 + (end - 1 - start) // step + + +def _select_mapping( + value: Mapping[str, Any], keys: Optional[Sequence[str]] +) -> dict[str, Any]: + if keys is None: + return dict(value) + return {key: value[key] for key in keys if key in value} + + +_DATAPROTO_SCHEMA_SECTIONS = frozenset({"batch", "non_tensor_batch", "meta_info"}) +_RAGGED_TENSOR_PAYLOAD_NAMES = frozenset({"data", "offsets", "shapes", "ndims", "nulls"}) + + +def _schema_section(name: str, schema: FieldSchema) -> str | None: + section = schema.metadata.get("section") + if section is None: + return None + if section not in _DATAPROTO_SCHEMA_SECTIONS: + raise ValueError( + f"FieldSchema for {name!r} metadata['section'] must be one of " + f"{sorted(_DATAPROTO_SCHEMA_SECTIONS)}" + ) + return section + + +def _schema_for_section( + field_schemas: Optional[Mapping[str, FieldSchema]], + name: str, + section: str, +) -> FieldSchema | None: + if not field_schemas: + return None + schema = field_schemas.get(name) + if schema is None: + return None + declared = _schema_section(name, schema) + return schema if declared is None or declared == section else None + + +def _schema_ndarray_dtype( + schema: FieldSchema, *, required: bool = False +) -> np.dtype[Any] | None: + dtype_name = schema.metadata.get("dtype") + if dtype_name is None: + if required: + raise ValueError("FieldSchema metadata['dtype'] is required") + return None + try: + dtype = np.dtype(dtype_name) + except (TypeError, ValueError) as exc: + raise ValueError(f"invalid FieldSchema metadata['dtype']: {dtype_name!r}") from exc + if dtype.kind not in "biufc": + raise ValueError( + f"FieldSchema metadata['dtype'] must be numeric or bool, got {dtype}" + ) + return dtype + + +def _validate_dataproto_schema_sections( + batch: Mapping[str, Any], + non_tensor_batch: Mapping[str, Any], + meta_info: Mapping[str, Any], + field_schemas: Optional[Mapping[str, FieldSchema]], +) -> None: + if not field_schemas: + return + sections = { + "batch": batch, + "non_tensor_batch": non_tensor_batch, + "meta_info": meta_info, + } + actual_sections: dict[str, set[str]] = {} + for section, fields in sections.items(): + for name in fields: + actual_sections.setdefault(name, set()).add(section) + for name, schema in field_schemas.items(): + declared = _schema_section(name, schema) + if declared is None: + continue + actual = actual_sections.get(name) + if not actual or declared in actual: + continue + actual_text = ", ".join(repr(section) for section in sorted(actual)) + raise ValueError( + f"FieldSchema for {name!r} declares section {declared!r}, " + f"but data contains it in {actual_text}" + ) + +def _coerce_schema_non_tensor_value(name: str, value: Any) -> np.ndarray: + if isinstance(value, np.ndarray): + return value + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + result = np.empty(len(value), dtype=object) + result[:] = list(value) + return result + raise TypeError( + f"non_tensor_batch field {name!r} with FieldSchema must be an ndarray or non-string sequence" + ) + + +def _validate_schema_nullable(path: str, value: np.ndarray, schema: FieldSchema) -> None: + if not schema.nullable and any(item is None for item in value): + raise ValueError(f"FieldSchema for {path!r} is not nullable") + + +def _should_encode_non_tensor_field(value: Any) -> bool: + return isinstance(value, np.ndarray) and value.dtype == object + + +def _encode_structured_non_tensor_field( + path: str, value: np.ndarray +) -> _EncodedStructuredLeaf: + values = list(value) + leaves: list[_InferredLeaf] = [] + nodes: list[_InferredNode] = [] + infer_structure(path, values, leaves, nodes) + if nodes and _should_encode_recursive_structure(leaves): + return _encode_recursive_structured_non_tensor_field(path, values, leaves, nodes) + decision = _choose_leaf_codec(values) + return _encode_structured_leaf(values, decision) + + +def _encode_with_schema( + path: str, value: np.ndarray, schema: FieldSchema +) -> "_EncodedStructuredLeaf": + """Encode a non_tensor_batch field using an explicit schema (no inference).""" + values = list(value) + _validate_schema_nullable(path, values, schema) + codec = schema.codec + if codec == "auto": + return _encode_with_fallback(path, value) + if codec == "ragged_tensor_dict": + return _encode_ragged_tensor_dict_values(values, schema) + if codec == "ragged_tensor": + payload, metadata = _encode_ragged_tensor_values(values) + elif codec == "typed_ragged": + payload, metadata = _encode_typed_ragged_values( + values, dtype_hint=_schema_ndarray_dtype(schema) + ) + elif codec == "ndarray": + dtype = _schema_ndarray_dtype(schema) + if dtype is None: + return _encode_with_fallback(path, value) + decision = _CodecDecision( + True, "ndarray", "schema", "numeric scalar", {"dtype": str(dtype)} + ) + payload, metadata = _encode_numeric_scalar_values(values, decision) + elif codec in ("bytes_ragged", "media_bytes"): + payload, metadata = _encode_bytes_like_values(values) + elif codec == "media_list_ragged": + payload, metadata = _encode_media_list_values(values) + elif codec == "utf8_ragged": + payload, metadata = _encode_bytes_like_values( + [None if v is None else v.encode("utf-8") for v in values] + ) + elif codec == "msgpack_ragged": + payload, metadata = _encode_msgpack_ragged_values(path, values) + elif codec == "json_ragged": + payload, metadata = _encode_bytes_like_values( + [ + None + if v is None + else json.dumps(v, ensure_ascii=False, separators=(",", ":")).encode( + "utf-8" + ) + for v in values + ] + ) + else: + raise ValueError(f"unsupported schema codec: {codec!r}") + return _EncodedStructuredLeaf( + codec=codec, rows=len(values), payload=payload, metadata=metadata + ) + + +def _should_encode_recursive_structure(leaves: Sequence[_InferredLeaf]) -> bool: + return any(_should_encode_recursive_leaf(leaf) for leaf in leaves) + + +def _should_encode_recursive_leaf(leaf: _InferredLeaf) -> bool: + codec = leaf.decision.codec + if codec in {"ragged_tensor", "bytes_ragged", "media_bytes", "media_list_ragged"}: + return True + if codec == "typed_ragged": + return any(isinstance(value, np.ndarray) for value in _non_null(leaf.values)) + return False + + +def _recursive_leaf_decision(values: list[Any], decision: _CodecDecision) -> _CodecDecision: + if decision.accepted: + return decision + if all(value is None or isinstance(value, _Missing) for value in values): + return _CodecDecision( + True, + "json_ragged", + "all rows are null or missing", + "json", + ) + return decision + + +def _encode_recursive_structured_non_tensor_field( + path: str, + values: list[Any], + leaves: Sequence[_InferredLeaf], + nodes: Sequence[_InferredNode], +) -> _EncodedStructuredLeaf: + payload: dict[str, Any] = {} + node_specs: list[dict[str, Any]] = [] + for node_id, node in enumerate(nodes): + spec: dict[str, Any] = { + "id": node_id, + "path": node.path, + "node_type": node.node_type, + "children": list(node.children), + } + missing_payload_name = f"node.{node_id}.missing" + payload[missing_payload_name] = np.asarray( + [_lookup_structured_path(value, path, node.path) is MISSING for value in values], + dtype=np.bool_, + ) + spec["missing_payload"] = missing_payload_name + if node.row_mask is not None: + payload_name = f"node.{node_id}.row_mask" + payload[payload_name] = np.asarray(node.row_mask, dtype=np.bool_) + spec["row_mask_payload"] = payload_name + if node.lengths is not None: + payload_name = f"node.{node_id}.lengths" + payload[payload_name] = np.asarray(node.lengths, dtype=np.int64) + spec["lengths_payload"] = payload_name + node_specs.append(spec) + + leaf_specs: list[dict[str, Any]] = [] + for leaf_id, leaf in enumerate(leaves): + missing = np.asarray( + [isinstance(value, _Missing) for value in leaf.values], dtype=np.bool_ + ) + codec_values = [None if is_missing else value for is_missing, value in zip(missing, leaf.values)] + decision = _recursive_leaf_decision(leaf.values, leaf.decision) + encoded = _encode_structured_leaf(codec_values, decision) + leaf_payload_members: dict[str, str] = {} + for payload_name, payload_value in encoded.payload.items(): + recursive_payload_name = f"leaf.{leaf_id}.{payload_name}" + payload[recursive_payload_name] = payload_value + leaf_payload_members[payload_name] = recursive_payload_name + missing_payload_name = f"leaf.{leaf_id}.missing" + payload[missing_payload_name] = missing + leaf_payload_members["missing"] = missing_payload_name + leaf_specs.append( + { + "id": leaf_id, + "path": leaf.path, + "codec": encoded.codec, + "rows": encoded.rows, + "metadata": encoded.metadata, + "payload_members": leaf_payload_members, + } + ) + + return _EncodedStructuredLeaf( + codec="structured_recursive", + rows=len(values), + payload=payload, + metadata={ + "schema_source": "inferred_from_runtime_values", + "structure_version": 1, + "root_path": path, + "nodes": node_specs, + "leaves": leaf_specs, + }, + ) + + +def _is_recursive_encoded_non_tensor(encoded: Mapping[str, Any]) -> bool: + return encoded.get("codec") == "structured_recursive" + + +def _structured_path_tokens(path: str) -> list[Any]: + tokens: list[Any] = [] + index = 0 + current: list[str] = [] + while index < len(path): + char = path[index] + if char == "\\" and index + 1 < len(path): + current.append(path[index + 1]) + index += 2 + continue + if char == ".": + if current: + tokens.append("".join(current)) + current = [] + index += 1 + continue + if char == "[": + if current: + tokens.append("".join(current)) + current = [] + end = path.index("]", index) + tokens.append(int(path[index + 1 : end])) + index = end + 1 + continue + current.append(char) + index += 1 + if current: + tokens.append("".join(current)) + return tokens + + +def _lookup_structured_path(value: Any, root_path: str, target_path: str) -> Any: + suffix = target_path[len(root_path) :] + if suffix.startswith("."): + suffix = suffix[1:] + if not suffix: + return value + current = value + for token in _structured_path_tokens(suffix): + if isinstance(current, _Missing): + return MISSING + if isinstance(token, str): + if not isinstance(current, dict) or token not in current: + return MISSING + current = current[token] + else: + if not isinstance(current, (list, tuple)) or token >= len(current): + return MISSING + current = current[token] + return current +def _encode_structured_leaf( + values: list[Any], decision: _CodecDecision +) -> _EncodedStructuredLeaf: + if not decision.accepted: + raise ValueError(f"unsupported structured non-tensor field: {decision.reason}") + codec = decision.codec + if codec == "ragged_tensor": + payload, metadata = _encode_ragged_tensor_values(values) + elif codec == "typed_ragged": + payload, metadata = _encode_typed_ragged_values(values) + elif codec == "media_list_ragged": + payload, metadata = _encode_media_list_values(values) + elif codec == "ndarray": + payload, metadata = _encode_numeric_scalar_values(values, decision) + elif codec == "bytes_ragged": + payload, metadata = _encode_bytes_like_values(values) + elif codec == "media_bytes": + payload, metadata = _encode_bytes_like_values(values) + elif codec == "utf8_ragged": + payload, metadata = _encode_bytes_like_values( + [None if value is None else value.encode("utf-8") for value in values] + ) + elif codec == "msgpack_ragged": + payload, metadata = _encode_bytes_like_values( + [ + None + if value is None + else _msgpack.packb(value, use_bin_type=True, strict_types=True) + for value in values + ] + ) + elif codec == "json_ragged": + payload, metadata = _encode_bytes_like_values( + [ + None + if value is None + else json.dumps(value, ensure_ascii=False, separators=(",", ":")).encode( + "utf-8" + ) + for value in values + ] + ) + else: + raise ValueError(f"unsupported structured non-tensor codec: {codec}") + metadata.update( + { + "schema_source": "inferred_from_runtime_values" + if decision.accepted + else "fallback", + "decision_reason": decision.reason, + "normalized_type": decision.normalized_type, + **decision.metadata, + } + ) + return _EncodedStructuredLeaf( + codec=codec, rows=len(values), payload=payload, metadata=metadata + ) + + + +def _decode_structured_non_tensor_encoded( + encoded: Mapping[str, Any], + payload: dict[str, Any], + rows: int, + metadata: Optional[Mapping[str, Any]] = None, +) -> list[Any]: + codec = encoded.get("codec") + if codec == "ragged_tensor_dict": + return _decode_ragged_tensor_dict_values( + payload, rows, metadata or encoded.get("metadata", {}) + ) + if encoded.get("codec") == "structured_recursive": + if metadata is not None: + encoded = {**encoded, "metadata": metadata} + return _decode_structured_recursive_field(encoded, payload, rows) + return _decode_structured_leaf(encoded["codec"], payload, rows, metadata) + + +def _copy_recursive_metadata_for_leaf_updates( + metadata: Mapping[str, Any] +) -> dict[str, Any]: + copied = dict(metadata) + copied["leaves"] = [dict(leaf) for leaf in metadata.get("leaves", [])] + return copied + + +def _decode_structured_recursive_field( + encoded: Mapping[str, Any], payload: dict[str, Any], rows: int +) -> list[Any]: + metadata = encoded.get("metadata") or {} + leaf_values: dict[str, list[Any]] = {} + for leaf in metadata.get("leaves", []): + leaf_payload_members = leaf["payload_members"] + leaf_payload = { + name: payload[global_name] + for name, global_name in leaf_payload_members.items() + if name != "missing" + } + values = _decode_structured_leaf( + leaf["codec"], leaf_payload, rows, leaf.get("metadata") + ) + missing = payload[leaf_payload_members["missing"]] + leaf_values[leaf["path"]] = [ + MISSING if bool(missing[row]) else values[row] for row in range(rows) + ] + return _reconstruct_structured_rows(metadata, payload, leaf_values, rows) + + +def _child_path(node: Mapping[str, Any], child: Any) -> str: + if node["node_type"] == "dict": + return f"{node['path']}.{_escape_key(str(child))}" + return f"{node['path']}[{int(child)}]" + + +def _path_depth(path: str) -> int: + return path.count(".") + path.count("[") + + +def _reconstruct_structured_rows( + metadata: Mapping[str, Any], + payload: Mapping[str, Any], + leaf_values: Mapping[str, list[Any]], + rows: int, +) -> list[Any]: + values_by_path: dict[str, list[Any]] = dict(leaf_values) + nodes = sorted( + metadata.get("nodes", []), key=lambda node: _path_depth(node["path"]), reverse=True + ) + for node in nodes: + missing_payload = node.get("missing_payload") + missing = payload[missing_payload] if missing_payload is not None else None + row_mask_payload = node.get("row_mask_payload") + row_mask = payload[row_mask_payload] if row_mask_payload is not None else None + lengths_payload = node.get("lengths_payload") + lengths = payload[lengths_payload] if lengths_payload is not None else None + node_values = [] + for row in range(rows): + if missing is not None and bool(missing[row]): + node_values.append(MISSING) + continue + if row_mask is not None and not bool(row_mask[row]): + node_values.append(None) + continue + if node["node_type"] == "dict": + item = {} + for child in node["children"]: + child_values = values_by_path[_child_path(node, child)] + child_value = child_values[row] + if not isinstance(child_value, _Missing): + item[child] = child_value + node_values.append(item) + continue + length = int(lengths[row]) if lengths is not None else len(node["children"]) + item = [None] * length + for child in node["children"]: + index = int(child) + if index >= length: + continue + child_values = values_by_path[_child_path(node, child)] + child_value = child_values[row] + if not isinstance(child_value, _Missing): + item[index] = child_value + node_values.append(item) + values_by_path[node["path"]] = node_values + return values_by_path[metadata["root_path"]] + + +def _decode_structured_leaf( + codec: str, + payload: dict[str, Any], + rows: int, + metadata: Optional[Mapping[str, Any]] = None, +) -> list[Any]: + if codec == "ragged_tensor": + return _decode_ragged_tensor_values(payload, rows, metadata) + if codec == "typed_ragged": + return _decode_typed_ragged_values(payload, rows, metadata) + if codec == "media_list_ragged": + return _decode_media_list_values(payload, rows, metadata) + if codec == "ndarray": + return _decode_numeric_scalar_values(payload) + if codec == "media_bytes": + return _decode_bytes_like_values(payload, rows, metadata) + if codec == "bytes_ragged": + return _decode_bytes_like_values(payload, rows) + if codec == "utf8_ragged": + return [ + None if value is None else value.decode("utf-8") + for value in _decode_bytes_like_values(payload, rows) + ] + if codec == "msgpack_ragged": + return _decode_msgpack_ragged_values(payload, rows) + if codec == "json_ragged": + return [ + None if value is None else json.loads(value) + for value in _decode_bytes_like_values(payload, rows) + ] + raise ValueError(f"unknown structured non-tensor codec: {codec}") + + +def _encode_ragged_tensor_values( + values: list[Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + if _torch is None: + raise RuntimeError("torch is required to encode ragged tensor fields") + tensors: list[Any] = [] + dtype = None + max_ndim = 0 + for value in values: + if value is None: + tensors.append(None) + continue + tensor = value.detach() + if tensor.device.type != "cpu" or not tensor.is_contiguous(): + tensor = tensor.cpu().contiguous() + dtype = tensor.dtype if dtype is None else dtype + if tensor.dtype != dtype: + raise ValueError(f"mixed tensor dtype: {dtype} vs {tensor.dtype}") + max_ndim = max(max_ndim, tensor.dim()) + tensors.append(tensor) + offsets = _torch.zeros(len(tensors) + 1, dtype=_torch.int64) + ndims = _torch.zeros(len(tensors), dtype=_torch.int16) + shapes = _torch.zeros((len(tensors), max(max_ndim, 1)), dtype=_torch.int64) + nulls = np.asarray([tensor is None for tensor in tensors], dtype=np.bool_) + flat_parts = [] + offset = 0 + for row, tensor in enumerate(tensors): + if tensor is None: + offsets[row + 1] = offset + continue + flat = tensor.reshape(-1) + flat_parts.append(flat) + offset += flat.numel() + offsets[row + 1] = offset + ndims[row] = tensor.dim() + if tensor.dim() > 0: + shapes[row, : tensor.dim()] = _torch.tensor( + list(tensor.shape), dtype=_torch.int64 + ) + data_dtype = dtype or _torch.float32 + data = ( + _torch.cat(flat_parts) if flat_parts else _torch.empty((0,), dtype=data_dtype) + ) + payload = { + "data": data.numpy(), + "offsets": offsets.numpy(), + "shapes": shapes.numpy(), + "ndims": ndims.numpy(), + "nulls": nulls, + } + return payload, { + "dtype": str(data_dtype), + "max_ndim": int(max_ndim), + "shape_policy": "ragged", + } + + +def _decode_ragged_tensor_values( + payload: dict[str, Any], + rows: int, + metadata: Optional[Mapping[str, Any]] = None, +) -> list[Any]: + if _torch is None: + raise RuntimeError("torch is required to decode ragged tensor fields") + data = _torch.from_numpy(payload["data"]) + offsets = payload["offsets"] + shapes = payload["shapes"] + ndims = payload["ndims"] + nulls = payload["nulls"] + values = [] + for row in range(rows): + if bool(nulls[row]): + values.append(None) + continue + begin = int(offsets[row]) + end = int(offsets[row + 1]) + ndim = int(ndims[row]) + shape = tuple(int(v) for v in shapes[row, :ndim].tolist()) + values.append(data[begin:end].reshape(shape)) + return values + +def _normalize_ragged_tensor_dict_keys(keys: Any) -> list[str]: + if isinstance(keys, Mapping): + iterable = keys.keys() + else: + iterable = keys + normalized: list[str] = [] + seen: set[str] = set() + for key in iterable: + if not isinstance(key, str): + raise TypeError("ragged_tensor_dict keys must be strings") + if not key: + raise ValueError("ragged_tensor_dict keys must not be empty") + if any(separator in key for separator in (".", "/", "\\")): + raise ValueError("ragged_tensor_dict keys must not contain '.', '/', or '\\'") + if key in seen: + raise ValueError(f"ragged_tensor_dict keys contain duplicate key {key!r}") + seen.add(key) + normalized.append(key) + return normalized + + +def _ragged_tensor_dict_payload_items( + payload: Mapping[str, Any], key: str, *, kind: str +) -> dict[str, Any]: + prefix = f"{key}." + items = { + name[len(prefix) :]: value + for name, value in payload.items() + if name.startswith(prefix) + } + missing = sorted(_RAGGED_TENSOR_PAYLOAD_NAMES - set(items)) + if missing: + raise ValueError( + f"ragged_tensor_dict {kind} for key {key!r} is missing payloads: {missing}" + ) + return items + + +def _encode_ragged_tensor_dict_values( + values: list[Any], + schema: FieldSchema, +) -> "_EncodedStructuredLeaf": + """Encode list[dict[str, Tensor] | None] directly without inference. + + Each dict key's tensors are encoded as a separate ragged_tensor sub-payload, + giving per-key zero-copy RDMA transfer and independent partial-read access. + """ + rows = len(values) + null_mask = np.asarray([v is None for v in values], dtype=np.bool_) + + schema_keys = schema.metadata.get("keys") + keys = None if schema_keys is None else _normalize_ragged_tensor_dict_keys(schema_keys) + declared_keys = None if keys is None else set(keys) + inferred_keys: set[str] = set() + for row, item in enumerate(values): + if item is None: + continue + if not isinstance(item, Mapping): + raise TypeError( + "ragged_tensor_dict rows must be mappings or None; " + f"row {row} is {type(item).__name__}" + ) + explicit_null_keys = sorted(key for key, value in item.items() if value is None) + if explicit_null_keys: + raise TypeError( + "ragged_tensor_dict rows cannot contain explicit None tensor values; " + f"row {row} has explicit None for {explicit_null_keys}" + ) + row_keys = _normalize_ragged_tensor_dict_keys(item.keys()) + if declared_keys is None: + inferred_keys.update(row_keys) + continue + extra_keys = sorted(set(row_keys) - declared_keys) + if extra_keys: + raise ValueError( + "ragged_tensor_dict rows contain keys not declared in " + f"FieldSchema metadata['keys']; row {row} has keys not declared: {extra_keys}" + ) + if keys is None: + keys = _normalize_ragged_tensor_dict_keys(sorted(inferred_keys)) + if keys and _torch is None: + raise RuntimeError("torch is required to encode ragged_tensor_dict fields") + + payload: dict[str, Any] = {"null_mask": null_mask} + key_codecs: dict[str, Any] = {} + for key in keys: + try: + key_values = [None if v is None else v.get(key) for v in values] + sub_payload, sub_metadata = _encode_ragged_tensor_values(key_values) + except _ENCODING_FALLBACK_ERRORS as exc: + raise ValueError( + f"failed to encode ragged_tensor_dict key {key!r}" + ) from exc + for payload_name, payload_value in sub_payload.items(): + payload[f"{key}.{payload_name}"] = payload_value + key_codecs[key] = sub_metadata + + return _EncodedStructuredLeaf( + codec="ragged_tensor_dict", + rows=rows, + payload=payload, + metadata={"keys": keys, "key_codecs": key_codecs}, + ) + +def _decode_ragged_tensor_dict_values( + payload: dict[str, Any], + rows: int, + metadata: Mapping[str, Any], +) -> list[Any]: + """Decode ragged_tensor_dict payload back to list[dict[str, Tensor] | None].""" + null_mask = payload["null_mask"] + if isinstance(null_mask, (bytes, bytearray)): + null_mask = np.frombuffer(null_mask, dtype=np.bool_) + keys = metadata.get("keys", []) + + key_values: dict[str, list[Any]] = {} + for key in keys: + prefix = f"{key}." + sub_payload = { + name[len(prefix) :]: payload[name] + for name in payload + if name.startswith(prefix) + } + sub_metadata = metadata.get("key_codecs", {}).get(key) + key_values[key] = _decode_ragged_tensor_values(sub_payload, rows, sub_metadata) + + result: list[Any] = [] + for row in range(rows): + if bool(null_mask[row]): + result.append(None) + else: + d: dict[str, Any] = {} + for key in keys: + val = key_values[key][row] + if val is not None: + d[key] = val + result.append(d) + return result -@dataclass(frozen=True) -class _NdarrayReadPlan: - dtype: np.dtype[Any] - full_shape: tuple[int, ...] - output_shape: tuple[int, ...] - byte_offset: int - byte_length: int - step: int - cover_row_count: int +def _encode_typed_ragged_values( + values: list[Any], dtype_hint: np.dtype[Any] | None = None +) -> tuple[dict[str, Any], dict[str, Any]]: + if dtype_hint is None: + source_arrays = [np.asarray(value) for value in values if value is not None] + dtype = np.result_type(*source_arrays) if source_arrays else np.dtype(np.int64) + else: + dtype = np.dtype(dtype_hint) + if dtype.hasobject: + raise ValueError("typed_ragged codec requires non-object dtype") + arrays = [ + np.asarray([], dtype=dtype) + if value is None + else np.ascontiguousarray(np.asarray(value, dtype=dtype)) + for value in values + ] + max_ndim = max((array.ndim for array in arrays), default=0) + offsets = np.zeros(len(arrays) + 1, dtype=np.int64) + ndims = np.zeros(len(arrays), dtype=np.int16) + shapes = np.zeros((len(arrays), max(max_ndim, 1)), dtype=np.int64) + nulls = np.asarray([value is None for value in values], dtype=np.bool_) + flat_arrays = [] + offset = 0 + for row, array in enumerate(arrays): + flat = array.reshape(-1) + flat_arrays.append(flat) + offset += flat.size + offsets[row + 1] = offset + ndims[row] = array.ndim + if array.ndim > 0: + shapes[row, : array.ndim] = array.shape + if flat_arrays: + if _concat_arrays_into is not None: + data = _DirectCopyPayload.from_flat_arrays( + flat_arrays, dtype, int(offset) + ) + else: + buffers = tuple(memoryview(flat.data).cast("B") for flat in flat_arrays) + data = _MultiBufferPayload( + buffers=buffers, + owners=tuple(flat_arrays), + dtype=np.dtype(dtype).str, + shape=(int(offset),), + ) + else: + empty = np.empty(0, dtype=dtype) + data = _MultiBufferPayload( + buffers=(memoryview(empty.data).cast("B"),), + owners=(empty,), + dtype=np.dtype(dtype).str, + shape=(0,), + ) + return ( + { + "data": data, + "offsets": offsets, + "shapes": shapes, + "ndims": ndims, + "nulls": nulls, + }, + {"dtype": str(dtype), "max_ndim": int(max_ndim), "shape_policy": "ragged"}, + ) -class MooncakeBundleTransfer: - """Transfer structured objects through Mooncake, with a low-level bundle fallback.""" +def _decode_typed_ragged_values( + payload: dict[str, Any], rows: int, metadata: Optional[Mapping[str, Any]] = None +) -> list[Any]: + data = payload["data"] + offsets = payload["offsets"] + shapes = payload["shapes"] + ndims = payload["ndims"] + nulls = payload["nulls"] + values = [] + for row in range(rows): + if bool(nulls[row]): + values.append(None) + continue + begin = int(offsets[row]) + end = int(offsets[row + 1]) + ndim = int(ndims[row]) + shape = tuple(int(v) for v in shapes[row, :ndim].tolist()) + values.append(data[begin:end].reshape(shape).tolist()) + return values - def __init__( - self, - store: BundleStore, - key_prefix: str = "bundle", - default_chunk_bytes: int = DEFAULT_BUNDLE_CHUNK_BYTES, - ) -> None: - """Initialize a bundle transfer helper with a configurable default chunk size.""" - self.store = store - self.key_prefix = _normalize_key_prefix(key_prefix) - self.default_chunk_bytes = _validate_chunk_bytes(default_chunk_bytes) - self._transport = _MooncakePayloadTransport(store) - self._bundle_store = _BundleManifestStore( - store=store, - transport=self._transport, - key_prefix=self.key_prefix, - default_chunk_bytes=self.default_chunk_bytes, - ) - self._structured_store = _StructuredObjectLayer(self._bundle_store) - def put_bundle( - self, - meta: bytes | bytearray | memoryview, - buffers: Mapping[str, Any], - partition: str = "default", - chunk_bytes: Optional[int] = None, - policy: Optional[BundleTransferPolicy] = None, - max_inflight_put: Optional[int] = None, - pre_registered_buffers: Optional[Mapping[str, bool]] = None, - ) -> RemoteBundleRef: - """Store raw metadata bytes plus named buffers as a low-level bundle.""" - return self._bundle_store.put_bundle( - meta=meta, - buffers=buffers, - partition=partition, - chunk_bytes=chunk_bytes, - policy=policy, - max_inflight_put=max_inflight_put, - pre_registered_buffers=pre_registered_buffers, +def _encode_numeric_scalar_values( + values: list[Any], decision: _CodecDecision +) -> tuple[dict[str, Any], dict[str, Any]]: + dtype = np.dtype(decision.metadata["dtype"]) + nulls = np.asarray([value is None for value in values], dtype=np.bool_) + fill_value = False if dtype == np.dtype(np.bool_) else 0 + data = np.asarray( + [fill_value if value is None else value for value in values], dtype=dtype + ) + return {"data": data, "nulls": nulls}, { + "dtype": str(data.dtype), + "shape": list(data.shape), + } + + +def _decode_numeric_scalar_values(payload: dict[str, Any]) -> list[Any]: + data = payload["data"] + nulls = payload["nulls"] + if not bool(nulls.any()): + return data.tolist() + values = data.astype(object) + values[nulls] = None + return values.tolist() + + +def _value_to_media_bytes(value: Any) -> tuple[bytes, str | None, dict[str, Any]]: + if value is None: + return b"", None, {"kind": "null"} + if _is_bytes_like(value): + return bytes(value), None, {"kind": "bytes"} + if _is_pil_image(value): + image = ( + value.convert(value.mode) if getattr(value, "readonly", False) else value + ) + return ( + image.tobytes(), + "image/raw", + { + "kind": "pil_raw", + "mode": image.mode, + "size": list(image.size), + "format": getattr(value, "format", None), + }, ) + return bytes(value), None, {"kind": "bytes"} - def remove_bundle(self, ref: RemoteBundleRef | Mapping[str, Any]) -> None: - """Remove all Mooncake objects that belong to a stored bundle.""" - self._bundle_store.remove_bundle(ref) +def _encode_bytes_like_values( + values: list[Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + offsets = [0] + parts = [] + media_types = [] + encodings = [] + for value in values: + data, media_type, encoding = _value_to_media_bytes(value) + parts.append(data) + media_types.append(media_type) + encodings.append(encoding) + offsets.append(offsets[-1] + len(data)) + payload_bytes = _multi_buffer_bytes_payload(parts) + metadata: dict[str, Any] = {} + if any(encoding.get("kind") == "pil_raw" for encoding in encodings): + metadata["media_encodings"] = encodings + non_null_media_types = sorted( + {media_type for media_type in media_types if media_type} + ) + if non_null_media_types: + metadata["media_types"] = non_null_media_types + metadata["encode_source"] = "raw_pixels_from_object" + return ( + { + "data": payload_bytes, + "offsets": np.asarray(offsets, dtype=np.int64), + "nulls": np.asarray([value is None for value in values], dtype=np.bool_), + }, + metadata, + ) - def put_structured_object( - self, - payload: StructuredObjectPayload, - partition: str = "default", - chunk_bytes: Optional[int] = None, - policy: Optional[BundleTransferPolicy] = None, - max_inflight_put: Optional[int] = None, - ) -> RemoteBundleRef: - """Store a structured object described by JSON metadata plus named members.""" - return self._structured_store.put_structured_object( - payload=payload, - partition=partition, - chunk_bytes=chunk_bytes, - policy=policy, - max_inflight_put=max_inflight_put, - ) - def read_spec( - self, ref: RemoteBundleRef | Mapping[str, Any] - ) -> StructuredObjectReadSpec: - """Create a lazy read spec for a structured object reference.""" - return self._structured_store.read_spec(ref) +def _multi_buffer_bytes_payload( + parts: Sequence[bytes | memoryview], +) -> _MultiBufferPayload: + buffers = tuple(_bytes_view(part, "payload part") for part in parts if part) + return _MultiBufferPayload(buffers, tuple(parts)) - def materialize(self, spec: StructuredObjectReadSpec) -> StructuredObjectResult: - """Materialize a structured object read spec.""" - return self._structured_store.materialize(spec) - def materialize_into( - self, - spec: StructuredObjectReadSpec, - destinations: Optional[Mapping[str, Any]], - ) -> StructuredObjectResult: - """Materialize a structured object read spec into caller-provided destinations when possible.""" - return self._structured_store.materialize_into(spec, destinations) +def _decode_media_bytes(data: Any, encoding: Optional[Mapping[str, Any]]) -> Any: + if not encoding or encoding.get("kind") != "pil_raw": + return bytes(data) + try: + from PIL import Image + except ImportError: + return bytes(data) + image = Image.frombuffer( + encoding["mode"], tuple(encoding["size"]), data, "raw", encoding["mode"], 0, 1 + ) + image.format = encoding.get("format") + return image + + +def _decode_bytes_like_values( + payload: dict[str, Any], rows: int, metadata: Optional[Mapping[str, Any]] = None +) -> list[Any]: + data = payload["data"] + offsets = payload["offsets"] + nulls = payload["nulls"] + encodings = (metadata or {}).get("media_encodings", []) + values = [] + for row in range(rows): + if bool(nulls[row]): + values.append(None) + continue + item = memoryview(data)[int(offsets[row]) : int(offsets[row + 1])] + encoding = encodings[row] if row < len(encodings) else None + values.append(_decode_media_bytes(item, encoding)) + return values + + +def _encode_media_list_values( + values: list[Any], +) -> tuple[dict[str, Any], dict[str, Any]]: + row_offsets = [0] + byte_offsets = [0] + parts = [] + media_types = [] + encodings = [] + for value in values: + items = [] if value is None else list(value) + for item in items: + data, media_type, encoding = _value_to_media_bytes(item) + parts.append(data) + media_types.append(media_type) + encodings.append(encoding) + byte_offsets.append(byte_offsets[-1] + len(data)) + row_offsets.append(row_offsets[-1] + len(items)) + metadata: dict[str, Any] = { + "encode_source": "raw_pixels_from_object", + "media_encodings": encodings, + } + non_null_media_types = sorted( + {media_type for media_type in media_types if media_type} + ) + if non_null_media_types: + metadata["media_types"] = non_null_media_types + return ( + { + "data": b"".join(parts), + "row_offsets": np.asarray(row_offsets, dtype=np.int64), + "byte_offsets": np.asarray(byte_offsets, dtype=np.int64), + "nulls": np.asarray([value is None for value in values], dtype=np.bool_), + }, + metadata, + ) + + +def _decode_media_list_values( + payload: dict[str, Any], rows: int, metadata: Optional[Mapping[str, Any]] = None +) -> list[Any]: + data = payload["data"] + row_offsets = payload["row_offsets"] + byte_offsets = payload["byte_offsets"] + nulls = payload["nulls"] + encodings = (metadata or {}).get("media_encodings", []) + values = [] + for row in range(rows): + if bool(nulls[row]): + values.append(None) + continue + items = [] + for item_index in range(int(row_offsets[row]), int(row_offsets[row + 1])): + item = memoryview(data)[ + int(byte_offsets[item_index]) : int(byte_offsets[item_index + 1]) + ] + encoding = encodings[item_index] if item_index < len(encodings) else None + items.append(_decode_media_bytes(item, encoding)) + values.append(items) + return values class _StructuredObjectLayer: @@ -225,15 +3083,20 @@ def put_structured_object( chunk_bytes: Optional[int], policy: Optional[BundleTransferPolicy], max_inflight_put: Optional[int], + config: Any = None, ) -> RemoteBundleRef: metadata, buffers = _encode_structured_fields(payload.metadata, payload.buffers) + transfer_policy = self._bundle_store._policy( + policy, max_inflight_put=max_inflight_put + ) return self._bundle_store.put_bundle( meta=_encode_structured_metadata(metadata), buffers=buffers, partition=partition, chunk_bytes=chunk_bytes, - policy=policy, - max_inflight_put=max_inflight_put, + policy=transfer_policy, + max_inflight_put=None, + config=config, ) def read_spec( @@ -297,25 +3160,134 @@ def _read_structured_member( return self._read_bytes_member( name, payload_spec, member_slice, destination ) + if encoding == "torch_tensor": + return self._read_torch_tensor_member( + name, payload_spec, field_spec, member_slice, destination + ) if encoding != "ndarray": raise ValueError(f"unsupported structured field encoding: {encoding}") return self._read_ndarray_member( name, payload_spec, field_spec, member_slice, destination ) + def _read_torch_tensor_member( + self, + name: str, + payload_spec: Mapping[str, Any], + field_spec: Mapping[str, Any], + member_slice: StructuredMemberSlice | None, + destination: Any, + ) -> Any: + if member_slice is not None: + return self._read_sliced_torch_tensor_member( + name, payload_spec, field_spec, member_slice, destination + ) + if destination is not None: + if not isinstance(destination, (_TensorObjectBufferPayload, _RawDestinationBuffer)): + raise ValueError( + f"structured tensor member {name} only supports tensor_object_buffer or raw_destination destinations" + ) + materialized = self._bundle_store.read_tensor_payload_into( + payload_spec, destination + ) + return destination if materialized is None else materialized + if payload_spec.get("kind") == "tensor": + return self._bundle_store.read_tensor_payload(payload_spec) + if payload_spec.get("format") == "torch_save": + return _deserialize_torch_save_payload( + self._bundle_store.read_payload(payload_spec) + ) + return _deserialize_tensor_payload( + self._bundle_store.read_payload(payload_spec) + ) + + def _read_sliced_torch_tensor_member( + self, + name: str, + payload_spec: Mapping[str, Any], + field_spec: Mapping[str, Any], + member_slice: StructuredMemberSlice, + destination: Any, + ) -> Any: + metadata_bytes = int(payload_spec.get("metadata_bytes", -1)) + shape = field_spec.get("shape") + element_size = int(field_spec.get("element_size", 0)) + if metadata_bytes < 0 or not isinstance(shape, list) or element_size <= 0: + raise ValueError( + f"structured tensor field {name} is missing slice metadata" + ) + if member_slice.axis != 0: + raise ValueError("structured tensor slicing currently supports axis=0 only") + if not shape: + raise ValueError( + "structured tensor slicing requires at least one dimension" + ) + start, end, step = _normalized_member_slice(member_slice, int(shape[0])) + if step != 1: + raise ValueError("structured tensor slicing currently requires step=1") + row_width = element_size * int(np.prod(shape[1:], dtype=np.int64)) + data_offset = metadata_bytes + start * row_width + data_length = (end - start) * row_width + metadata = self._bundle_store.read_payload_range( + payload_spec, 0, metadata_bytes + ) + sliced_metadata = _slice_tensor_metadata( + metadata, (end - start, *shape[1:]), data_length + ) + if destination is not None: + if not isinstance(destination, (_TensorObjectBufferPayload, _RawDestinationBuffer)): + raise ValueError( + f"structured tensor member {name} only supports tensor_object_buffer or raw_destination destinations" + ) + expected_bytes = metadata_bytes + data_length + if destination.size < expected_bytes: + raise ValueError( + f"tensor destination has {destination.size} bytes, expected at least {expected_bytes}" + ) + ctypes.memmove(destination.ptr, sliced_metadata, metadata_bytes) + if isinstance(destination, _RawDestinationBuffer) and destination.pre_registered: + self._bundle_store.read_payload_range_into_raw_destination( + payload_spec, + destination.ptr, + metadata_bytes, + data_offset, + data_length, + ) + else: + data = self._bundle_store.read_payload_range( + payload_spec, data_offset, data_length + ) + ctypes.memmove(destination.ptr + metadata_bytes, data, data_length) + _ = destination.owner + return destination + data = self._bundle_store.read_payload_range( + payload_spec, data_offset, data_length + ) + return _deserialize_tensor_payload(sliced_metadata + data) + def _read_bytes_member( self, name: str, payload_spec: Mapping[str, Any], member_slice: StructuredMemberSlice | None, destination: Any, - ) -> bytes: + ) -> bytes | np.ndarray: if member_slice is not None: raise ValueError(f"structured bytes member {name} does not support slicing") if destination is not None: raise ValueError( f"structured bytes member {name} does not support materialize_into" ) + # Try pool-backed read first: pool memory is pre-registered, + # so RDMA reads go directly into pool memory without extra copies. + expected_bytes = int(payload_spec.get("bytes", 0)) + if expected_bytes > 0: + result = self._bundle_store.read_payload_range_into_pool_array( + payload_spec, np.dtype(np.uint8), (expected_bytes,), 0 + ) + if result is not None: + return result + # Pool unavailable or size unknown - fall back to plain bytes. return self._bundle_store.read_payload(payload_spec) def _read_ndarray_member( @@ -335,12 +3307,21 @@ def _read_ndarray_member( read_plan = _resolve_ndarray_read_plan( tuple(int(dim) for dim in shape), np.dtype(dtype), member_slice ) + if destination is None and read_plan.byte_length > 0 and read_plan.step == 1: + target = self._bundle_store.read_payload_range_into_pool_array( + payload_spec, + read_plan.dtype, + read_plan.output_shape, + read_plan.byte_offset, + ) + if target is not None: + return target target = _resolve_ndarray_destination( name, destination, read_plan.dtype, read_plan.output_shape ) - destination_view = target.view(np.uint8).reshape(-1) if read_plan.byte_length == 0: return target + destination_view = target.view(np.uint8).reshape(-1) if read_plan.step == 1: self._bundle_store.read_payload_range_into_destination( payload_spec, @@ -379,7 +3360,7 @@ def put_bundle( chunk_bytes: Optional[int], policy: Optional[BundleTransferPolicy], max_inflight_put: Optional[int], - pre_registered_buffers: Optional[Mapping[str, bool]] = None, + config: Any = None, ) -> RemoteBundleRef: _validate_key_segment(partition, "partition") meta_view = _bytes_view(meta, "meta") @@ -392,26 +3373,49 @@ def put_bundle( manifest_key = f"{base_key}/manifest" written_keys: list[str] = [] buffer_specs: dict[str, Any] = {} - pre_registered_map = dict(pre_registered_buffers or {}) try: meta_spec, meta_keys = self._put_payload( f"{base_key}/meta", meta_view, target_chunk_bytes, - transfer_policy, - pre_registered=False, + _copy_transfer_policy(transfer_policy), + config=config, ) written_keys.extend(meta_keys) for name, value in buffers.items(): _validate_key_segment(name, "buffer name") - payload_view = _bytes_view(value, name) - payload_spec, payload_keys = self._put_payload( - f"{base_key}/buffer/{name}", - payload_view, - target_chunk_bytes, - transfer_policy, - pre_registered=bool(pre_registered_map.get(name, False)), - ) + payload_key = f"{base_key}/buffer/{name}" + if isinstance(value, (_TensorPayload, _TensorObjectBufferPayload)): + payload_spec, payload_keys = self._put_tensor_payload( + payload_key, + value, + transfer_policy, + config=config, + ) + elif isinstance(value, _MultiBufferPayload): + payload_spec, payload_keys = self._put_multi_buffer_payload( + payload_key, + value, + target_chunk_bytes, + transfer_policy, + config=config, + ) + elif isinstance(value, _DirectCopyPayload): + payload_spec, payload_keys = self._put_direct_copy_payload( + payload_key, + value, + target_chunk_bytes, + transfer_policy, + config=config, + ) + else: + payload_spec, payload_keys = self._put_payload( + payload_key, + _bytes_view(value, name), + target_chunk_bytes, + transfer_policy, + config=config, + ) buffer_specs[name] = payload_spec written_keys.extend(payload_keys) manifest = { @@ -423,7 +3427,69 @@ def put_bundle( } manifest_blob = _encode_manifest(manifest) _check_status( - self._store.put(manifest_key, manifest_blob), "put", manifest_key + _put_with_optional_config(self._store, manifest_key, manifest_blob, config), + "put", + manifest_key, + ) + written_keys.append(manifest_key) + except Exception: + _cleanup_keys(self._store, written_keys, strict=False) + raise + return RemoteBundleRef(manifest_key=manifest_key, manifest=manifest) + + def put_bundle_manifest( + self, + meta: bytes | bytearray | memoryview, + buffers: Mapping[str, Any], + *, + partition: str, + chunk_bytes: Optional[int], + policy: Optional[BundleTransferPolicy], + cleanup_keys: Optional[Sequence[str]] = None, + config: Any = None, + ) -> RemoteBundleRef: + _validate_key_segment(partition, "partition") + meta_view = _bytes_view(meta, "meta") + target_chunk_bytes = _validate_chunk_bytes( + self._default_chunk_bytes if chunk_bytes is None else chunk_bytes + ) + transfer_policy = self._policy(policy) + object_id = f"{partition}/{uuid.uuid4().hex}" + base_key = f"{self._key_prefix}/{object_id}" + manifest_key = f"{base_key}/manifest" + written_keys: list[str] = [] + try: + meta_spec, meta_keys = self._put_payload( + f"{base_key}/meta", + meta_view, + target_chunk_bytes, + _copy_transfer_policy(transfer_policy), + config=config, + ) + written_keys.extend(meta_keys) + manifest = { + "version": 1, + "layout": "bundle", + "object_id": object_id, + "meta": meta_spec, + "buffers": dict(buffers), + "buffer_object_ids": sorted( + { + str(payload_spec["key"]) + .removeprefix(f"{self._key_prefix}/") + .split("/buffer/", 1)[0] + for payload_spec in buffers.values() + } + ), + } + if cleanup_keys: + manifest["cleanup_keys"] = list(dict.fromkeys(cleanup_keys)) + self._validate_manifest(manifest) + manifest_blob = _encode_manifest(manifest) + _check_status( + _put_with_optional_config(self._store, manifest_key, manifest_blob, config), + "put", + manifest_key, ) written_keys.append(manifest_key) except Exception: @@ -434,14 +3500,26 @@ def put_bundle( def remove_bundle(self, ref: RemoteBundleRef | Mapping[str, Any]) -> None: manifest = self.resolve_manifest(ref) keys = self._payload_keys(manifest) + keys.extend(manifest.get("cleanup_keys", [])) keys.append(self._manifest_key(ref, manifest)) _cleanup_keys(self._store, keys, strict=True) + def manifest_key(self, ref: RemoteBundleRef | Mapping[str, Any]) -> str: + return self._manifest_key(ref, self.resolve_manifest(ref)) + + def payload_keys(self, payload_spec: Mapping[str, Any]) -> list[str]: + return [chunk["key"] for chunk in payload_spec["chunks"]] + + def remove_keys(self, keys: Sequence[str], *, strict: bool) -> None: + _cleanup_keys(self._store, keys, strict=strict) + def resolve_manifest( self, ref: RemoteBundleRef | Mapping[str, Any] ) -> dict[str, Any]: if isinstance(ref, RemoteBundleRef): manifest = ref.manifest + if not manifest: + manifest = _decode_manifest(self._store.get(ref.manifest_key)) else: manifest = ref.get("manifest") if manifest is None: @@ -455,19 +3533,151 @@ def resolve_manifest( def read_payload(self, payload_spec: Mapping[str, Any]) -> bytes: return self._transport.read_payload(payload_spec) - def read_payload_range_into_destination( + def read_payload_range( + self, payload_spec: Mapping[str, Any], byte_offset: int, byte_length: int + ) -> bytes: + return self._transport.read_payload_range( + payload_spec, byte_offset, byte_length + ) + + def read_payload_range_into_pool_array( + self, + payload_spec: Mapping[str, Any], + dtype: np.dtype[Any], + shape: tuple[int, ...], + byte_offset: int, + ) -> np.ndarray | None: + return self._transport.read_payload_range_into_pool_array( + payload_spec, dtype, shape, byte_offset + ) + + def read_tensor_payload(self, payload_spec: Mapping[str, Any]) -> Any: + return self._transport.read_tensor_payload(payload_spec) + + def read_tensor_payload_into( + self, payload_spec: Mapping[str, Any], destination: _TensorObjectBufferPayload + ) -> Any: + return self._transport.read_tensor_payload_into(payload_spec, destination) + + def read_payload_range_into_destination( + self, + payload_spec: Mapping[str, Any], + destination: np.ndarray, + byte_offset: int, + destination_pre_registered: bool = False, + ) -> None: + self._transport.read_payload_range_into_destination( + payload_spec, + destination, + byte_offset, + destination_pre_registered=destination_pre_registered, + ) + + def read_payload_range_into_raw_destination( + self, + payload_spec: Mapping[str, Any], + destination_ptr: int, + destination_offset: int, + byte_offset: int, + byte_length: int, + ) -> None: + if not self._transport.read_payload_range_into_raw_destination( + payload_spec, destination_ptr, destination_offset, byte_offset, byte_length + ): + data = self.read_payload_range(payload_spec, byte_offset, byte_length) + ctypes.memmove(destination_ptr + destination_offset, data, byte_length) + + def read_payload_ranges_into_array( + self, + payload_spec: Mapping[str, Any], + destination: np.ndarray, + ranges: Sequence[tuple[int, int, int]], + ) -> None: + if not ranges: + return + self._transport.read_payload_ranges_into_array(payload_spec, destination, ranges) + + def read_payload_ranges_into_bytearray( + self, + payload_spec: Mapping[str, Any], + destination: bytearray, + ranges: Sequence[tuple[int, int, int]], + ) -> None: + if not ranges: + return + self._transport.read_payload_ranges_into_bytearray( + payload_spec, destination, ranges + ) + + def read_payload_ranges_into_raw_destination( self, payload_spec: Mapping[str, Any], - destination: np.ndarray, - byte_offset: int, - destination_pre_registered: bool = False, + destination_ptr: int, + ranges: Sequence[tuple[int, int, int]], ) -> None: - self._transport.read_payload_range_into_destination( - payload_spec, - destination, - byte_offset, - destination_pre_registered=destination_pre_registered, + if not ranges: + return + if self._transport.read_payload_ranges_into_raw_destination( + payload_spec, destination_ptr, ranges + ): + return + for byte_offset, destination_offset, byte_length in ranges: + data = self.read_payload_range(payload_spec, byte_offset, byte_length) + ctypes.memmove(destination_ptr + destination_offset, data, byte_length) + + def _put_tensor_payload( + self, + key: str, + value: _TensorPayload | _TensorObjectBufferPayload, + transfer_policy: BundleTransferPolicy, + config: Any = None, + ) -> tuple[dict[str, Any], list[str]]: + if isinstance(value, _TensorObjectBufferPayload): + total_bytes = self._transport.put_tensor_object_buffer(key, value, config) + return { + "key": key, + "bytes": total_bytes, + "chunks": [{"key": key, "bytes": total_bytes}], + }, [key] + if transfer_policy.copy_mode != "copy": + if transfer_policy.copy_mode == "zero_copy": + raise ValueError( + "zero-copy structured tensor fields require a BufferPool tensor-object buffer" + ) + tensor_spec = self._transport.put_tensor_payload_direct(key, value, config) + if tensor_spec is not None: + return tensor_spec, [key] + try: + total_bytes = self._transport.put_tensor_payload_from_pool(key, value, config) + return { + "key": key, + "bytes": total_bytes, + "chunks": [{"key": key, "bytes": total_bytes}], + }, [key] + except RuntimeError: + pass + if not _has_tensor_codec_helpers() or value.tensor.dim() == 0: + payload = _torch_save_payload_bytes(value.tensor) + payload_spec, payload_keys = self._put_payload( + key, + memoryview(payload), + len(payload) or 1, + transfer_policy, + config=config, + ) + payload_spec["format"] = "torch_save" + return payload_spec, payload_keys + payload, metadata_bytes = _tensor_payload_bytes(value) + payload_spec, payload_keys = self._put_payload( + key, + memoryview(payload), + len(payload) or 1, + transfer_policy, + pre_registered=False, + config=config, ) + payload_spec["metadata_bytes"] = metadata_bytes + return payload_spec, payload_keys def _put_payload( self, @@ -475,18 +3685,20 @@ def _put_payload( value: memoryview, chunk_bytes: int, transfer_policy: BundleTransferPolicy, - pre_registered: bool, + config: Any = None, ) -> tuple[dict[str, Any], list[str]]: + if len(value) == 0: + return {"key": key, "bytes": 0, "chunks": []}, [] chunks = _split_view(value, chunk_bytes) chunk_keys = [ key if len(chunks) == 1 else f"{key}/chunk/{index}" for index in range(len(chunks)) ] - written_keys = self._transport.put_payload_chunks( + self._transport.put_multi_buffer_payload_chunks( chunk_keys, - chunks, + [[chunk] for chunk in chunks], transfer_policy, - pre_registered=pre_registered, + config, ) payload_spec = { "key": key, @@ -496,7 +3708,172 @@ def _put_payload( for chunk_key, chunk in zip(chunk_keys, chunks) ], } - return payload_spec, written_keys + return payload_spec, list(chunk_keys) + + def _put_multi_buffer_payload( + self, + key: str, + value: _MultiBufferPayload, + chunk_bytes: int, + transfer_policy: BundleTransferPolicy, + config: Any = None, + ) -> tuple[dict[str, Any], list[str]]: + total_bytes = value.nbytes + if total_bytes == 0: + return {"key": key, "bytes": 0, "chunks": []}, [] + if len(value.buffers) == 1: + return self._put_payload( + key, value.buffers[0], chunk_bytes, transfer_policy, config=config + ) + chunk_groups = _split_multi_buffer_payload(value.buffers, chunk_bytes) + chunk_keys = [ + key if len(chunk_groups) == 1 else f"{key}/chunk/{index}" + for index in range(len(chunk_groups)) + ] + self._transport.put_multi_buffer_payload_chunks( + chunk_keys, + chunk_groups, + transfer_policy, + config, + ) + payload_spec = { + "key": key, + "bytes": total_bytes, + "chunks": [ + {"key": chunk_key, "bytes": sum(len(part) for part in group)} + for chunk_key, group in zip(chunk_keys, chunk_groups) + ], + } + return payload_spec, list(chunk_keys) + + + def _put_direct_copy_payload( + self, + key: str, + value: _DirectCopyPayload, + chunk_bytes: int, + transfer_policy: BundleTransferPolicy, + config: Any = None, + ) -> tuple[dict[str, Any], list[str]]: + if transfer_policy.copy_mode == "zero_copy": + raise RuntimeError("zero-copy put requires tensor-object buffers") + total_bytes = value.total_bytes + if total_bytes == 0: + return {"key": key, "bytes": 0, "chunks": []}, [] + if len(value.arrays) == 1: + buf = memoryview(value.arrays[0].data).cast("B") + return self._put_payload( + key, + buf, + chunk_bytes, + transfer_policy, + config=config, + ) + arrays = value.arrays + transport = self._transport + pool = transport._ensure_buffer_pool() + batch_put_from = transport._batch_put_from + if pool is None or not callable(batch_put_from): + buf = memoryview(np.concatenate( + [a.ravel().view(np.uint8) for a in arrays] + ).data).cast("B") + return self._put_payload( + key, + buf, + chunk_bytes, + transfer_policy, + config=config, + ) + # Group arrays into chunk-sized batches + n = len(arrays) + chunk_batches: list[tuple[int, int, int]] = [] # (start, count, bytes) + batch_start = 0 + batch_bytes = 0 + fallback = False + for i in range(n): + ab = arrays[i].nbytes + if ab > chunk_bytes: + fallback = True + break + if batch_bytes + ab > chunk_bytes and batch_bytes > 0: + chunk_batches.append((batch_start, i - batch_start, batch_bytes)) + batch_start = i + batch_bytes = 0 + batch_bytes += ab + if fallback: + buf = memoryview(np.concatenate( + [a.ravel().view(np.uint8) for a in arrays] + ).data).cast("B") + return self._put_payload( + key, + buf, + chunk_bytes, + transfer_policy, + config=config, + ) + if batch_bytes > 0: + chunk_batches.append((batch_start, n - batch_start, batch_bytes)) + num_chunks = len(chunk_batches) + chunk_keys = [ + key if num_chunks == 1 else f"{key}/chunk/{idx}" + for idx in range(num_chunks) + ] + + def _put_chunk_batch(keys, batches): + for ck, (start, count, size) in zip(keys, batches): + lease = pool.acquire(size) + try: + copied = _concat_arrays_into(arrays, lease.ptr, size, start, count) + if copied != size: + raise RuntimeError( + f"native fast-copy wrote {copied} bytes, expected {size}" + ) + results = _batch_put_from_with_optional_config( + batch_put_from, [ck], [lease.ptr], [size], config + ) + transport._check_batch_put_results(results, [ck], "batch_put_from") + finally: + lease.release() + + max_inflight = transfer_policy.max_inflight_put + use_parallel = ( + transfer_policy.put_mode != "batch" + and max_inflight > 1 + and num_chunks >= AUTO_PARALLEL_MIN_CHUNKS + and total_bytes >= AUTO_PARALLEL_MIN_BYTES + ) + futures: list = [] + try: + if not use_parallel: + _put_chunk_batch(chunk_keys, chunk_batches) + else: + group_count = max(1, min(max_inflight, num_chunks)) + group_size = (num_chunks + group_count - 1) // group_count + groups = [ + (chunk_keys[s:s + group_size], chunk_batches[s:s + group_size]) + for s in range(0, num_chunks, group_size) + ] + with ThreadPoolExecutor(max_workers=len(groups)) as executor: + futures = [ + executor.submit(_put_chunk_batch, gk, gb) + for gk, gb in groups + ] + for f in as_completed(futures): + f.result() + except Exception: + for f in futures: + f.cancel() + _cleanup_keys(self._store, chunk_keys, strict=False) + raise + payload_spec = { + "key": key, + "bytes": total_bytes, + "chunks": [ + {"key": ck, "bytes": cb[2]} + for ck, cb in zip(chunk_keys, chunk_batches) + ], + } + return payload_spec, list(chunk_keys) def _policy( self, @@ -506,12 +3883,16 @@ def _policy( result = policy or BundleTransferPolicy() if max_inflight_put is not None: result = BundleTransferPolicy( - max_inflight_put=max_inflight_put, put_mode=result.put_mode + max_inflight_put=max_inflight_put, + put_mode=result.put_mode, + copy_mode=result.copy_mode, ) if result.max_inflight_put < 1: raise ValueError("max_inflight_put must be positive") if result.put_mode not in {"auto", "batch", "parallel"}: raise ValueError(f"unsupported put_mode: {result.put_mode}") + if result.copy_mode not in {"auto", "zero_copy", "copy"}: + raise ValueError(f"unsupported copy_mode: {result.copy_mode}") return result def _validate_manifest(self, manifest: Mapping[str, Any]) -> None: @@ -521,20 +3902,48 @@ def _validate_manifest(self, manifest: Mapping[str, Any]) -> None: if not isinstance(object_id, str): raise ValueError("bundle manifest object_id must be a string") base_key = f"{self._key_prefix}/{object_id}" - self._validate_payload_spec(manifest.get("meta"), base_key) + buffer_object_ids = manifest.get("buffer_object_ids", [object_id]) + if not isinstance(buffer_object_ids, list) or not all( + isinstance(item, str) for item in buffer_object_ids + ): + raise ValueError("bundle manifest buffer_object_ids must be a list of strings") + allowed_buffer_base_keys = [ + f"{self._key_prefix}/{buffer_object_id}" + for buffer_object_id in buffer_object_ids + ] + self._validate_payload_spec(manifest.get("meta"), base_keys=[base_key]) + cleanup_keys = manifest.get("cleanup_keys", []) + if not isinstance(cleanup_keys, list) or not all( + self._is_allowed_cleanup_key(item, allowed_buffer_base_keys) + for item in cleanup_keys + ): + raise ValueError("bundle manifest cleanup_keys are outside the bundle namespace") buffers = manifest.get("buffers") if not isinstance(buffers, dict): raise ValueError("bundle manifest buffers must be a dict") for name, payload_spec in buffers.items(): _validate_key_segment(name, "buffer name") - self._validate_payload_spec(payload_spec, base_key) + self._validate_payload_spec( + payload_spec, base_keys=allowed_buffer_base_keys + ) + + def _is_allowed_cleanup_key(self, key: Any, base_keys: Sequence[str]) -> bool: + if not isinstance(key, str): + return False + return any( + key == f"{base_key}/manifest" or key.startswith(f"{base_key}/") + for base_key in base_keys + ) - def _validate_payload_spec(self, payload_spec: Any, base_key: str) -> None: + def _validate_payload_spec( + self, payload_spec: Any, base_keys: Sequence[str] | None = None + ) -> None: if not isinstance(payload_spec, dict): raise ValueError("bundle payload spec must be a dict") payload_key = payload_spec.get("key") - if not isinstance(payload_key, str) or not payload_key.startswith( - f"{base_key}/" + if not isinstance(payload_key, str) or ( + base_keys is not None + and not any(payload_key.startswith(f"{base_key}/") for base_key in base_keys) ): raise ValueError("bundle payload key is outside the bundle namespace") expected_bytes = int(payload_spec.get("bytes", -1)) @@ -591,35 +4000,71 @@ def _payload_keys(self, manifest: Mapping[str, Any]) -> list[str]: class _MooncakePayloadTransport: """Move payload bytes through Mooncake, preferring fast-path APIs and falling back to generic store calls.""" - def __init__(self, store: BundleStore) -> None: + def __init__(self, store: BundleStore, buffer_pool: Any = None) -> None: self._store = store + self._buffer_pool = buffer_pool self._batch_put_from = getattr(store, "batch_put_from", None) + self._put_tensor_from = getattr(store, "put_tensor_from", None) self._batch_get_into = getattr(store, "batch_get_into", None) self._get_into = getattr(store, "get_into", None) self._get_into_ranges = getattr(store, "get_into_ranges", None) self._register_buffer = getattr(store, "register_buffer", None) self._unregister_buffer = getattr(store, "unregister_buffer", None) - def put_payload_chunks( + def _store_can_auto_create_buffer_pool(self) -> bool: + get_capsule = getattr(self._store, "_get_pyclient_capsule", None) + if not callable(get_capsule): + return False + try: + return get_capsule() is not None + except Exception: + return False + + def _ensure_buffer_pool(self) -> Any: + if self._buffer_pool is None: + if not self._store_can_auto_create_buffer_pool(): + return None + try: + from mooncake.buffer_pool import BufferPool + + self._buffer_pool = BufferPool(self._store) + except (ImportError, AttributeError, RuntimeError, TypeError): + return None + return self._buffer_pool + + def put_multi_buffer_payload_chunks( self, chunk_keys: Sequence[str], - chunks: Sequence[memoryview], + chunk_groups: Sequence[Sequence[memoryview]], transfer_policy: BundleTransferPolicy, - pre_registered: bool, + config: Any = None, ) -> list[str]: - if not self._has_batch_put_support(): - return self._put_chunks_direct(chunk_keys, chunks) - put_mode = self._resolve_put_mode(chunks, transfer_policy) - if put_mode == "batch": - self.batch_put_chunks_from( - chunk_keys, chunks, pre_registered=pre_registered + def fallback_to_direct_put() -> list[str]: + chunks = [memoryview(b"".join(group)) for group in chunk_groups] + return self._put_chunks_direct(chunk_keys, chunks, config) + + if transfer_policy.copy_mode == "zero_copy": + raise RuntimeError( + "zero-copy put requires tensor-object buffers" ) + if transfer_policy.copy_mode == "copy" or self._ensure_buffer_pool() is None: + return fallback_to_direct_put() + if all(len(group) == 1 for group in chunk_groups): + self.batch_put_buffer_groups_from( + chunk_keys, [[group[0]] for group in chunk_groups], config + ) + return list(chunk_keys) + if not callable(self._batch_put_from): + return fallback_to_direct_put() + put_mode = self._resolve_buffer_group_put_mode(chunk_groups, transfer_policy) + if put_mode == "batch": + self.batch_put_buffer_groups_from(chunk_keys, chunk_groups, config) return list(chunk_keys) - return self._put_chunks_parallel( + return self._put_buffer_groups_parallel( list(chunk_keys), - list(chunks), + [list(group) for group in chunk_groups], transfer_policy.max_inflight_put, - pre_registered=pre_registered, + config, ) def read_payload(self, payload_spec: Mapping[str, Any]) -> bytes: @@ -630,6 +4075,13 @@ def read_payload(self, payload_spec: Mapping[str, Any]) -> bytes: self.read_payload_into(payload_spec, data) return bytes(data) + def read_tensor_payload(self, payload_spec: Mapping[str, Any]) -> Any: + get_tensor = getattr(self._store, "get_tensor", None) + key = payload_spec["key"] + if not callable(get_tensor): + raise RuntimeError("structured tensor payload does not support get_tensor") + return get_tensor(key) + def read_payload_into( self, payload_spec: Mapping[str, Any], destination: bytearray | np.ndarray ) -> None: @@ -640,6 +4092,61 @@ def read_payload_into( for offset, chunk in zip(offsets, chunks): self._read_chunk_with_get(chunk, destination, offset) + def read_tensor_payload_into( + self, + payload_spec: Mapping[str, Any], + destination: _TensorObjectBufferPayload, + ) -> Any: + if payload_spec.get("kind") == "tensor": + get_tensor_into = getattr(self._store, "get_tensor_into", None) + if not callable(get_tensor_into): + raise RuntimeError( + "structured tensor payload does not support get_tensor_into" + ) + result = get_tensor_into( + payload_spec["key"], destination.ptr, destination.size + ) + _ = destination.owner + return result + expected_bytes = int(payload_spec["bytes"]) + if destination.size < expected_bytes: + raise ValueError( + f"tensor destination has {destination.size} bytes, expected at least {expected_bytes}" + ) + chunks = payload_spec["chunks"] + if not self._read_payload_range_into_raw_destination( + chunks, destination.ptr, 0, expected_bytes, allow_get_into=False + ): + offset = 0 + for chunk in chunks: + data = self._store.get(chunk["key"]) + size = int(chunk["bytes"]) + if len(data) != size: + raise RuntimeError( + f"get failed for {chunk['key']}: expected {size} bytes, got {len(data)}" + ) + ctypes.memmove(destination.ptr + offset, data, size) + offset += size + _ = destination.owner + return None + + def read_payload_range( + self, + payload_spec: Mapping[str, Any], + byte_offset: int, + byte_length: int, + ) -> bytes: + if byte_length == 0: + return b"" + data = bytearray(byte_length) + if not self._read_payload_range_into_registered_destination( + payload_spec["chunks"], data, byte_offset, byte_length, False + ): + self._copy_payload_range_into_bytearray( + payload_spec["chunks"], data, byte_offset, byte_length + ) + return bytes(data) + def read_payload_range_into_destination( self, payload_spec: Mapping[str, Any], @@ -660,60 +4167,212 @@ def read_payload_range_into_destination( chunks, destination, byte_offset, byte_length ) - def batch_put_chunks_from( + def read_payload_range_into_raw_destination( + self, + payload_spec: Mapping[str, Any], + destination_ptr: int, + destination_offset: int, + byte_offset: int, + byte_length: int, + ) -> bool: + if byte_length == 0: + return True + return self._read_payload_range_into_raw_destination( + payload_spec["chunks"], + destination_ptr, + byte_offset, + byte_length, + allow_get_into=False, + destination_offset=destination_offset, + ) + + def read_payload_ranges_into_array( + self, + payload_spec: Mapping[str, Any], + destination: np.ndarray, + ranges: Sequence[tuple[int, int, int]], + ) -> None: + if self._read_payload_ranges_into_registered_destination( + payload_spec["chunks"], destination, ranges + ): + return + self._copy_payload_ranges_into_destination( + payload_spec["chunks"], destination.view(np.uint8).reshape(-1), ranges + ) + + def read_payload_ranges_into_bytearray( + self, + payload_spec: Mapping[str, Any], + destination: bytearray, + ranges: Sequence[tuple[int, int, int]], + ) -> None: + if self._read_payload_ranges_into_registered_destination( + payload_spec["chunks"], destination, ranges + ): + return + self._copy_payload_ranges_into_destination( + payload_spec["chunks"], destination, ranges + ) + + def read_payload_ranges_into_raw_destination( + self, + payload_spec: Mapping[str, Any], + destination_ptr: int, + ranges: Sequence[tuple[int, int, int]], + ) -> bool: + return self._read_payload_ranges_into_raw_destination( + payload_spec["chunks"], destination_ptr, ranges + ) + + def put_tensor_object_buffer( + self, + key: str, + value: _TensorObjectBufferPayload, + config: Any = None, + ) -> int: + put_tensor_from = self._put_tensor_from + put_from = getattr(self._store, "put_from", None) + if not callable(put_tensor_from) and not callable(put_from): + raise RuntimeError("put_from is unavailable") + _check_status( + _put_from_with_optional_config(self._store, key, value.ptr, value.size, config), + "put_from", + key, + ) + _ = value.owner + return value.size + + def put_tensor_payload_direct( + self, + key: str, + value: _TensorPayload, + config: Any = None, + ) -> dict[str, Any] | None: + if config is not None: + return None + put_tensor = getattr(self._store, "put_tensor", None) + if not callable(put_tensor): + return None + tensor = value.tensor + nbytes = int(getattr(tensor, "nbytes", 0)) + metadata_bytes = _tensor_metadata_size() + total_bytes = metadata_bytes + nbytes + _check_status(put_tensor(key, tensor), "put_tensor", key) + return { + "key": key, + "kind": "tensor", + "bytes": total_bytes, + "dtype": str(getattr(tensor, "dtype", "")), + "shape": list(getattr(tensor, "shape", ())), + "chunks": [{"key": key, "bytes": total_bytes}], + "metadata_bytes": metadata_bytes, + } + + def put_tensor_payload_from_pool( + self, + key: str, + value: _TensorPayload, + config: Any = None, + ) -> int: + if self._buffer_pool is None: + raise RuntimeError("structured tensor zero-copy requires a BufferPool") + put_tensor_from = self._put_tensor_from + put_from = getattr(self._store, "put_from", None) + if not callable(put_tensor_from) and not callable(put_from): + raise RuntimeError("put_from is unavailable") + metadata, data_ptr, tensor_nbytes, owner = _tensor_payload_parts(value) + total_bytes = len(metadata) + tensor_nbytes + lease = self._buffer_pool.acquire(total_bytes) + view = None + try: + view = lease.buffer + view[: len(metadata)] = metadata + if tensor_nbytes: + ctypes.memmove(lease.ptr + len(metadata), data_ptr, tensor_nbytes) + view.release() + view = None + _check_status( + _put_from_with_optional_config(self._store, key, lease.ptr, total_bytes, config), + "put_from", + key, + ) + _ = owner + return total_bytes + finally: + if view is not None: + view.release() + lease.release() + + def batch_put_buffer_groups_from( self, chunk_keys: Sequence[str], - chunks: Sequence[memoryview], - pre_registered: bool, + chunk_groups: Sequence[Sequence[memoryview]], + config: Any = None, ) -> None: batch_put_from = self._batch_put_from if not callable(batch_put_from): raise RuntimeError("batch_put_from is unavailable") if not chunk_keys: return - prepared_chunks = [_prepare_chunk_source_buffer(chunk) for chunk in chunks] - buffer_ptrs = [ptr for _owner, ptr, _size in prepared_chunks] - sizes = [size for _owner, _ptr, size in prepared_chunks] - registered_ptrs = self._register_buffers( - buffer_ptrs, sizes, pre_registered, "bundle source payload" - ) - try: - results = batch_put_from(list(chunk_keys), buffer_ptrs, sizes) - if len(results) != len(chunk_keys): - raise RuntimeError( - f"batch_put_from returned {len(results)} results for {len(chunk_keys)} chunks" + sizes = [_buffer_group_nbytes(group) for group in chunk_groups] + pool = self._buffer_pool + # Stream one chunk at a time: acquire -> memcpy -> put -> release. + # This keeps pool pressure minimal and allows RDMA transfers to pipeline. + for chunk_key, group, size in zip(chunk_keys, chunk_groups, sizes): + lease = pool.acquire(size) + try: + _copy_memoryviews_to_lease(group, lease) + results = _batch_put_from_with_optional_config( + batch_put_from, [chunk_key], [lease.ptr], [size], config ) - for chunk_key, status in zip(chunk_keys, results): - _check_status(status, "batch_put_from", chunk_key) - except Exception: - _cleanup_keys(self._store, chunk_keys, strict=False) - raise - finally: - self._unregister_buffers(registered_ptrs, "bundle source payload") + self._check_batch_put_results(results, [chunk_key], "batch_put_from") + except Exception: + _cleanup_keys(self._store, chunk_keys, strict=False) + raise + finally: + lease.release() + + @staticmethod + def _check_batch_put_results( + results: Sequence[int], chunk_keys: Sequence[str], operation: str + ) -> None: + if len(results) != len(chunk_keys): + raise RuntimeError( + f"{operation} returned {len(results)} results for {len(chunk_keys)} chunks" + ) + for chunk_key, status in zip(chunk_keys, results): + _check_status(status, operation, chunk_key) def _put_chunks_direct( self, chunk_keys: Sequence[str], chunks: Sequence[memoryview], + config: Any = None, ) -> list[str]: written_keys: list[str] = [] try: for chunk_key, chunk in zip(chunk_keys, chunks): - _check_status(self._store.put(chunk_key, chunk), "put", chunk_key) + _check_status( + _put_with_optional_config(self._store, chunk_key, chunk, config), + "put", + chunk_key, + ) written_keys.append(chunk_key) except Exception: _cleanup_keys(self._store, written_keys, strict=False) raise return list(chunk_keys) - def _put_chunks_parallel( + def _put_buffer_groups_parallel( self, chunk_keys: list[str], - chunks: list[memoryview], + chunk_groups: list[Sequence[memoryview]], max_inflight_put: int, - pre_registered: bool, + config: Any = None, ) -> list[str]: - groups = self._group_chunk_ranges(chunk_keys, chunks, max_inflight_put) + groups = self._group_buffer_group_ranges( + chunk_keys, chunk_groups, max_inflight_put + ) futures: list[Future[None]] = [] try: with ThreadPoolExecutor( @@ -721,10 +4380,10 @@ def _put_chunks_parallel( ) as executor: futures = [ executor.submit( - self.batch_put_chunks_from, + self.batch_put_buffer_groups_from, group_keys, group_chunks, - pre_registered, + config, ) for group_keys, group_chunks in groups ] @@ -743,22 +4402,22 @@ def _put_chunks_parallel( raise return chunk_keys - def _group_chunk_ranges( + def _group_buffer_group_ranges( self, chunk_keys: Sequence[str], - chunks: Sequence[memoryview], + chunk_groups: Sequence[Sequence[memoryview]], max_inflight_put: int, - ) -> list[tuple[list[str], list[memoryview]]]: + ) -> list[tuple[list[str], list[Sequence[memoryview]]]]: if not chunk_keys: return [] - group_count = max(1, min(max_inflight_put, len(chunks))) - group_size = (len(chunks) + group_count - 1) // group_count + group_count = max(1, min(max_inflight_put, len(chunk_groups))) + group_size = (len(chunk_groups) + group_count - 1) // group_count return [ ( list(chunk_keys[start : start + group_size]), - list(chunks[start : start + group_size]), + list(chunk_groups[start : start + group_size]), ) - for start in range(0, len(chunks), group_size) + for start in range(0, len(chunk_groups), group_size) ] def _read_chunks_with_batch_get_into( @@ -777,99 +4436,219 @@ def _read_chunks_with_batch_get_into( read_sizes = batch_get_into(keys, ptrs, sizes) if len(read_sizes) != len(keys): raise RuntimeError( - f"batch_get_into returned {len(read_sizes)} results for {len(keys)} chunks" + f"batch_get_into returned {len(read_sizes)} results for {len(keys)} chunks" + ) + for key, expected_size, actual_size in zip(keys, sizes, read_sizes): + if actual_size != expected_size: + raise RuntimeError( + f"batch_get_into failed for {key}: expected {expected_size}, got {actual_size}" + ) + return True + + def _read_chunk_with_get( + self, + chunk: Mapping[str, Any], + destination: bytearray | np.ndarray, + offset: int, + ) -> None: + chunk_bytes = int(chunk["bytes"]) + if chunk_bytes == 0: + return + data = self._store.get(chunk["key"]) + if len(data) != chunk_bytes: + raise RuntimeError( + f"get failed for {chunk['key']}: expected {chunk_bytes} bytes, got {len(data)}" + ) + destination[offset : offset + chunk_bytes] = data + + def _read_payload_range_into_registered_destination( + self, + chunks: Sequence[Mapping[str, Any]], + destination: np.ndarray, + byte_offset: int, + byte_length: int, + destination_pre_registered: bool, + ) -> bool: + if not self._has_buffer_registration_support(): + return False + with self._registered_buffer( + destination, + "structured ndarray payload", + pre_registered=destination_pre_registered, + ) as base_ptr: + return self._read_payload_range_into_raw_destination( + chunks, base_ptr, byte_offset, byte_length + ) + + def _read_payload_ranges_into_registered_destination( + self, + chunks: Sequence[Mapping[str, Any]], + destination: bytearray | np.ndarray, + ranges: Sequence[tuple[int, int, int]], + ) -> bool: + if not self._has_buffer_registration_support(): + return False + with self._registered_buffer(destination, "structured ranged payload") as base_ptr: + return self._read_payload_ranges_into_raw_destination( + chunks, base_ptr, ranges + ) + + def _read_payload_ranges_into_raw_destination( + self, + chunks: Sequence[Mapping[str, Any]], + base_ptr: int, + ranges: Sequence[tuple[int, int, int]], + ) -> bool: + get_into_ranges = self._get_into_ranges + if not callable(get_into_ranges): + return False + fragments = [] + for source_offset, destination_offset, byte_length in ranges: + fragments.extend( + _payload_range_fragments(chunks, source_offset, byte_length, destination_offset) + ) + if not fragments: + return True + keys = [ + key + for key, _chunk_size, _destination_offset, _source_offset, _size in fragments + ] + dst_offsets = [ + [destination_offset] + for _key, _chunk_size, destination_offset, _source_offset, _size in fragments + ] + src_offsets = [ + [source_offset] + for _key, _chunk_size, _destination_offset, source_offset, _size in fragments + ] + sizes = [ + [size] + for _key, _chunk_size, _destination_offset, _source_offset, size in fragments + ] + results = get_into_ranges( + [base_ptr], [keys], [dst_offsets], [src_offsets], [sizes] + ) + if len(results) != 1 or len(results[0]) != len(keys): + raise RuntimeError( + f"get_into_ranges returned invalid ranged result shape for {len(keys)} chunks" + ) + for key, expected_sizes, actual_sizes in zip(keys, sizes, results[0]): + if len(actual_sizes) != len(expected_sizes): + raise RuntimeError( + f"get_into_ranges returned invalid ranged fragment count for {key}" ) - for key, expected_size, actual_size in zip(keys, sizes, read_sizes): + for expected_size, actual_size in zip(expected_sizes, actual_sizes): if actual_size != expected_size: raise RuntimeError( - f"batch_get_into failed for {key}: expected {expected_size}, got {actual_size}" + f"get_into_ranges failed for {key}: expected {expected_size}, got {actual_size}" ) return True - def _read_chunk_with_get( + def read_payload_range_into_pool_array( self, - chunk: Mapping[str, Any], - destination: bytearray | np.ndarray, - offset: int, - ) -> None: - chunk_bytes = int(chunk["bytes"]) - if chunk_bytes == 0: - return - data = self._store.get(chunk["key"]) - if len(data) != chunk_bytes: - raise RuntimeError( - f"get failed for {chunk['key']}: expected {chunk_bytes} bytes, got {len(data)}" - ) - destination[offset : offset + chunk_bytes] = data + payload_spec: Mapping[str, Any], + dtype: np.dtype[Any], + shape: tuple[int, ...], + byte_offset: int, + ) -> np.ndarray | None: + nbytes = ( + int(np.prod(shape, dtype=np.int64)) * dtype.itemsize + if shape + else dtype.itemsize + ) + if nbytes == 0: + return np.empty(shape, dtype=dtype) + pool = self._ensure_buffer_pool() + if pool is None: + return None + try: + lease = pool.acquire(nbytes) + except RuntimeError: + return None + owner = _PoolLeaseOwner(lease) + try: + if not self._read_payload_range_into_raw_destination( + payload_spec["chunks"], + lease.ptr, + byte_offset, + nbytes, + allow_get_into=False, + ): + owner.release() + return None + return _PoolBackedNdarray(owner, dtype, shape) + except Exception: + owner.release() + raise - def _read_payload_range_into_registered_destination( + def _read_payload_range_into_raw_destination( self, chunks: Sequence[Mapping[str, Any]], - destination: np.ndarray, + base_ptr: int, byte_offset: int, byte_length: int, - destination_pre_registered: bool, + allow_get_into: bool = True, + destination_offset: int = 0, ) -> bool: get_into = self._get_into - can_use_get_into = len(chunks) == 1 and byte_offset == 0 and callable(get_into) - get_into_ranges = self._get_into_ranges - can_use_get_into_ranges = callable(get_into_ranges) - if not self._has_buffer_registration_support(): - return False - if not can_use_get_into and not can_use_get_into_ranges: - return False - with self._registered_buffer( - destination, - "structured ndarray payload", - pre_registered=destination_pre_registered, - ) as base_ptr: - if can_use_get_into and int(chunks[0]["bytes"]) == byte_length: - expected_size = int(chunks[0]["bytes"]) - read_size = get_into(chunks[0]["key"], base_ptr, expected_size) + if ( + allow_get_into + and len(chunks) == 1 + and byte_offset == 0 + and callable(get_into) + ): + expected_size = int(chunks[0]["bytes"]) + if expected_size == byte_length: + read_size = get_into( + chunks[0]["key"], base_ptr + destination_offset, expected_size + ) if read_size != expected_size: raise RuntimeError( f"get_into failed for {chunks[0]['key']}: expected {expected_size}, got {read_size}" ) return True - if not can_use_get_into_ranges: - return False - fragments = _payload_range_fragments(chunks, byte_offset, byte_length) - if not fragments: - return True - keys = [ - key - for key, _chunk_size, _destination_offset, _source_offset, _size in fragments - ] - dst_offsets = [ - [destination_offset] - for _key, _chunk_size, destination_offset, _source_offset, _size in fragments - ] - src_offsets = [ - [source_offset] - for _key, _chunk_size, _destination_offset, source_offset, _size in fragments - ] - sizes = [ - [size] - for _key, _chunk_size, _destination_offset, _source_offset, size in fragments - ] - results = get_into_ranges( - [base_ptr], [keys], [dst_offsets], [src_offsets], [sizes] + get_into_ranges = self._get_into_ranges + if not callable(get_into_ranges): + return False + fragments = _payload_range_fragments( + chunks, byte_offset, byte_length, destination_offset + ) + if not fragments: + return True + keys = [ + key + for key, _chunk_size, _destination_offset, _source_offset, _size in fragments + ] + dst_offsets = [ + [fragment_destination_offset] + for _key, _chunk_size, fragment_destination_offset, _source_offset, _size in fragments + ] + src_offsets = [ + [source_offset] + for _key, _chunk_size, _destination_offset, source_offset, _size in fragments + ] + sizes = [ + [size] + for _key, _chunk_size, _destination_offset, _source_offset, size in fragments + ] + results = get_into_ranges( + [base_ptr], [keys], [dst_offsets], [src_offsets], [sizes] + ) + if len(results) != 1 or len(results[0]) != len(keys): + raise RuntimeError( + f"get_into_ranges returned invalid ranged result shape for {len(keys)} chunks" ) - if len(results) != 1 or len(results[0]) != len(keys): + for key, expected_sizes, actual_sizes in zip(keys, sizes, results[0]): + if len(actual_sizes) != len(expected_sizes): raise RuntimeError( - f"get_into_ranges returned invalid ranged result shape for {len(keys)} chunks" + f"get_into_ranges returned invalid ranged fragment count for {key}" ) - for key, expected_sizes, actual_sizes in zip(keys, sizes, results[0]): - if len(actual_sizes) != len(expected_sizes): + for expected_size, actual_size in zip(expected_sizes, actual_sizes): + if actual_size != expected_size: raise RuntimeError( - f"get_into_ranges returned invalid ranged fragment count for {key}" + f"get_into_ranges failed for {key}: expected {expected_size}, got {actual_size}" ) - for expected_size, actual_size in zip(expected_sizes, actual_sizes): - if actual_size != expected_size: - raise RuntimeError( - f"get_into_ranges failed for {key}: expected {expected_size}, got {actual_size}" - ) - return True + return True def _copy_payload_range_into_destination( self, @@ -877,6 +4656,42 @@ def _copy_payload_range_into_destination( destination: np.ndarray, byte_offset: int, byte_length: int, + ) -> None: + data = bytearray(byte_length) + self._copy_payload_range_into_bytearray(chunks, data, byte_offset, byte_length) + destination[:byte_length] = np.frombuffer(data, dtype=np.uint8) + + def _copy_payload_ranges_into_destination( + self, + chunks: Sequence[Mapping[str, Any]], + destination: bytearray | np.ndarray, + ranges: Sequence[tuple[int, int, int]], + ) -> None: + for source_offset, destination_offset, byte_length in ranges: + for ( + key, + chunk_size, + fragment_destination_offset, + fragment_source_offset, + size, + ) in _payload_range_fragments( + chunks, source_offset, byte_length, destination_offset + ): + data = self._store.get(key) + if len(data) != chunk_size: + raise RuntimeError( + f"get failed for {key}: expected {chunk_size} bytes, got {len(data)}" + ) + destination[ + fragment_destination_offset : fragment_destination_offset + size + ] = data[fragment_source_offset : fragment_source_offset + size] + + def _copy_payload_range_into_bytearray( + self, + chunks: Sequence[Mapping[str, Any]], + destination: bytearray, + byte_offset: int, + byte_length: int, ) -> None: for ( key, @@ -884,21 +4699,15 @@ def _copy_payload_range_into_destination( destination_offset, source_offset, size, - ) in _payload_range_fragments( - chunks, - byte_offset, - byte_length, - ): + ) in _payload_range_fragments(chunks, byte_offset, byte_length): data = self._store.get(key) if len(data) != chunk_size: raise RuntimeError( f"get failed for {key}: expected {chunk_size} bytes, got {len(data)}" ) - fragment = memoryview(data)[source_offset : source_offset + size] - target_end = destination_offset + size - destination[destination_offset:target_end] = np.frombuffer( - fragment, dtype=np.uint8 - ) + destination[destination_offset : destination_offset + size] = data[ + source_offset : source_offset + size + ] def _register_buffers( self, @@ -916,6 +4725,8 @@ def _register_buffers( registered_ptrs: list[int] = [] try: for ptr, size in zip(buffer_ptrs, sizes): + if size == 0: + continue register_status = register_buffer(ptr, size) if register_status == 0: registered_ptrs.append(ptr) @@ -962,9 +4773,9 @@ def _registered_buffer( if succeeded: raise - def _resolve_put_mode( + def _resolve_buffer_group_put_mode( self, - chunks: Sequence[memoryview], + chunk_groups: Sequence[Sequence[memoryview]], transfer_policy: BundleTransferPolicy, ) -> Literal["batch", "parallel"]: if transfer_policy.put_mode == "parallel": @@ -973,19 +4784,15 @@ def _resolve_put_mode( return "batch" if transfer_policy.max_inflight_put <= 1: return "batch" - if len(chunks) < AUTO_PARALLEL_MIN_CHUNKS: + if len(chunk_groups) < AUTO_PARALLEL_MIN_CHUNKS: return "batch" - if sum(len(chunk) for chunk in chunks) < AUTO_PARALLEL_MIN_BYTES: + total_bytes = sum(_buffer_group_nbytes(group) for group in chunk_groups) + if total_bytes < AUTO_PARALLEL_MIN_BYTES: return "batch" - if min(transfer_policy.max_inflight_put, len(chunks)) < 2: + if min(transfer_policy.max_inflight_put, len(chunk_groups)) < 2: return "batch" return "parallel" - def _has_batch_put_support(self) -> bool: - return ( - callable(self._batch_put_from) and self._has_buffer_registration_support() - ) - def _has_buffer_registration_support(self) -> bool: return callable(self._register_buffer) and callable(self._unregister_buffer) @@ -998,9 +4805,19 @@ def _resolve_ndarray_destination( ) -> np.ndarray: if destination is None: return np.empty(shape, dtype=dtype) + if isinstance(destination, _RawDestinationBuffer): + nbytes = int(np.prod(shape, dtype=np.int64)) * dtype.itemsize + if destination.size < nbytes: + raise ValueError( + f"raw destination has {destination.size} bytes, expected at least {nbytes}" + ) + _ = destination.owner + return np.ctypeslib.as_array( + (ctypes.c_uint8 * nbytes).from_address(destination.ptr) + ).view(dtype).reshape(shape) if not isinstance(destination, np.ndarray): raise TypeError( - f"structured ndarray field {name} destination must be a numpy.ndarray" + f"structured ndarray field {name} destination must be a numpy.ndarray or raw_destination" ) if destination.dtype != dtype: raise ValueError( @@ -1092,6 +4909,9 @@ def _bytes_view(value: Any, name: str) -> memoryview: def _prepare_chunk_source_buffer(chunk: memoryview) -> tuple[Any, int, int]: + if len(chunk) == 0: + copied = ctypes.create_string_buffer(0) + return copied, ctypes.addressof(copied), 0 if chunk.c_contiguous and not chunk.readonly: return chunk, ctypes.addressof(ctypes.c_char.from_buffer(chunk)), len(chunk) copied = ctypes.create_string_buffer(bytes(chunk)) @@ -1116,6 +4936,59 @@ def _split_view(view: memoryview, chunk_bytes: int) -> list[memoryview]: ] +def _split_multi_buffer_payload( + buffers: Sequence[memoryview], chunk_bytes: int +) -> list[list[memoryview]]: + if len(buffers) == 1: + return [[chunk] for chunk in _split_view(buffers[0], chunk_bytes)] + groups: list[list[memoryview]] = [] + current: list[memoryview] = [] + current_bytes = 0 + for buffer in buffers: + offset = 0 + while offset < len(buffer): + remaining = chunk_bytes - current_bytes + part = buffer[offset : offset + remaining] + current.append(part) + current_bytes += len(part) + offset += len(part) + if current_bytes == chunk_bytes: + groups.append(current) + current = [] + current_bytes = 0 + if current: + groups.append(current) + return groups + + +def _buffer_group_nbytes(buffers: Sequence[memoryview]) -> int: + return sum(len(buffer) for buffer in buffers) + + +def _copy_memoryviews(buffers: Sequence[memoryview], destination: memoryview) -> None: + destination_bytes = destination.cast("B") + offset = 0 + for buffer in buffers: + source_bytes = buffer.cast("B") + n = len(source_bytes) + if n == 0: + continue + destination_bytes[offset : offset + n] = source_bytes + offset += n + + +def _copy_memoryviews_to_lease(buffers: Sequence[memoryview], lease: Any) -> None: + copy_from_buffers = getattr(lease, "copy_from_buffers", None) + if callable(copy_from_buffers): + copy_from_buffers(buffers) + return + view = lease.buffer + try: + _copy_memoryviews(buffers, view) + finally: + view.release() + + def _chunk_offsets(chunks: Sequence[Mapping[str, Any]]) -> list[int]: offsets = [0] for chunk in chunks[:-1]: @@ -1127,6 +5000,7 @@ def _payload_range_fragments( chunks: Sequence[Mapping[str, Any]], byte_offset: int, byte_length: int, + destination_offset: int = 0, ) -> list[tuple[str, int, int, int, int]]: fragments: list[tuple[str, int, int, int, int]] = [] read_end = byte_offset + byte_length @@ -1141,7 +5015,7 @@ def _payload_range_fragments( ( chunk["key"], chunk_size, - overlap_start - byte_offset, + destination_offset + overlap_start - byte_offset, overlap_start - chunk_offset, overlap_end - overlap_start, ) @@ -1212,7 +5086,37 @@ def _encode_structured_fields( return normalized_metadata, encoded_fields +def tensor_object_buffer( + ptr: int, size: int, owner: Any = None, batch_size: int | None = None +) -> _TensorObjectBufferPayload: + return _TensorObjectBufferPayload( + ptr=int(ptr), size=int(size), owner=owner, batch_size=batch_size + ) + + +def raw_destination( + ptr: int, size: int, owner: Any = None, *, pre_registered: bool = False +) -> _RawDestinationBuffer: + return _RawDestinationBuffer( + ptr=int(ptr), size=int(size), owner=owner, pre_registered=pre_registered + ) + + def _encode_structured_field(value: Any) -> tuple[dict[str, Any], Any]: + if isinstance(value, _TensorObjectBufferPayload): + return {"encoding": "torch_tensor"}, value + if isinstance(value, _TensorPayload): + return _encode_torch_tensor_field(value.tensor) + if isinstance(value, (_DirectCopyPayload, _MultiBufferPayload)): + if value.dtype is not None and value.shape is not None: + return { + "encoding": "ndarray", + "dtype": value.dtype, + "shape": list(value.shape), + }, value + return {"encoding": "bytes"}, value + if _torch is not None and isinstance(value, _torch.Tensor): + return _encode_torch_tensor_field(value) if isinstance(value, np.ndarray): array = np.ascontiguousarray(value) return { @@ -1223,6 +5127,329 @@ def _encode_structured_field(value: Any) -> tuple[dict[str, Any], Any]: return {"encoding": "bytes"}, value +def _encode_torch_tensor_field(value: Any) -> tuple[dict[str, Any], Any]: + return { + "encoding": "torch_tensor", + "dtype": str(value.dtype), + "shape": list(value.shape), + "element_size": int(value.element_size()), + }, _TensorPayload(tensor=value) + + +def _tensor_payload_parts(value: _TensorPayload) -> tuple[bytes, int, int, Any]: + metadata, data_ptr, tensor_nbytes, owner = _tensor_codec_helper( + "_serialize_tensor" + )(value.tensor) + return bytes(metadata), int(data_ptr), int(tensor_nbytes), owner + + +def _encode_recursive_if_structured( + path: str, values: list[Any] +) -> "_EncodedStructuredLeaf | None": + leaves: list[_InferredLeaf] = [] + nodes: list[_InferredNode] = [] + infer_structure(path, values, leaves, nodes) + if not nodes or not any( + (leaf.decision.codec == "typed_ragged" and leaf.decision.metadata.get("recursive_source") == "ndarray") + or leaf.decision.codec in {"ragged_tensor", "bytes_ragged", "media_bytes", "media_list_ragged"} + for leaf in leaves + ): + return None + payload: dict[str, Any] = {} + node_specs: list[dict[str, Any]] = [] + for node_id, node in enumerate(nodes): + spec: dict[str, Any] = { + "id": node_id, + "path": node.path, + "node_type": node.node_type, + "children": list(node.children), + } + missing_payload_name = f"node.{node_id}.missing" + payload[missing_payload_name] = np.asarray( + [_lookup_structured_path(value, path, node.path) is MISSING for value in values], + dtype=np.bool_, + ) + spec["missing_payload"] = missing_payload_name + if node.row_mask is not None: + payload_name = f"node.{node_id}.row_mask" + payload[payload_name] = np.asarray(node.row_mask, dtype=np.bool_) + spec["row_mask_payload"] = payload_name + if node.lengths is not None: + payload_name = f"node.{node_id}.lengths" + payload[payload_name] = np.asarray(node.lengths, dtype=np.int64) + spec["lengths_payload"] = payload_name + node_specs.append(spec) + + leaf_specs: list[dict[str, Any]] = [] + for leaf_id, leaf in enumerate(leaves): + missing = np.asarray( + [isinstance(value, _Missing) for value in leaf.values], dtype=np.bool_ + ) + codec_values = [ + None if is_missing else value + for is_missing, value in zip(missing, leaf.values) + ] + decision = leaf.decision + if not decision.accepted and all( + value is None or isinstance(value, _Missing) for value in leaf.values + ): + decision = _CodecDecision(True, "json_ragged", "all rows are null or missing", "json") + encoded = _encode_with_schema( + leaf.path, + np.asarray( + _normalize_values_for_fallback_codec(codec_values, decision.codec), + dtype=object, + ), + FieldSchema(codec=decision.codec, metadata=decision.metadata), + ) + leaf_payload_members: dict[str, str] = {} + for payload_name, payload_value in encoded.payload.items(): + recursive_payload_name = f"leaf.{leaf_id}.{payload_name}" + payload[recursive_payload_name] = payload_value + leaf_payload_members[payload_name] = recursive_payload_name + missing_payload_name = f"leaf.{leaf_id}.missing" + payload[missing_payload_name] = missing + leaf_payload_members["missing"] = missing_payload_name + leaf_specs.append( + { + "id": leaf_id, + "path": leaf.path, + "codec": encoded.codec, + "rows": encoded.rows, + "metadata": encoded.metadata, + "payload_members": leaf_payload_members, + } + ) + return _EncodedStructuredLeaf( + codec="structured_recursive", + rows=len(values), + payload=payload, + metadata={ + "schema_source": "inferred_from_runtime_values", + "structure_version": 1, + "root_path": path, + "nodes": node_specs, + "leaves": leaf_specs, + }, + ) + + + +def _encode_with_fallback(path: str, value: np.ndarray) -> "_EncodedStructuredLeaf": + """Encode using safe type-based fallbacks, with recursive expansion for structured rows.""" + values = list(value) + errors: list[str] = [] + try: + recursive = _encode_recursive_if_structured(path, values) + if recursive is not None: + return recursive + except _ENCODING_FALLBACK_ERRORS as error: + errors.append(f"structured_recursive: {error}") + decision = _choose_leaf_codec(values) + codecs = [decision.codec] + for codec in ("msgpack_ragged", "json_ragged"): + if codec not in codecs: + codecs.append(codec) + for codec in codecs: + metadata = dict(decision.metadata or {}) if codec == decision.codec else {} + codec_values = _normalize_values_for_fallback_codec(values, codec) + codec_array = np.asarray(codec_values, dtype=object) + try: + return _encode_with_schema(path, codec_array, FieldSchema(codec=codec, metadata=metadata)) + except _ENCODING_FALLBACK_ERRORS as error: + errors.append(f"{codec}: {error}") + raise ValueError( + f"unable to infer a safe codec for structured non-tensor field {path!r}; " + f"tried {errors}" + ) + + + +def _normalize_values_for_fallback_codec(values: list[Any], codec: str) -> list[Any]: + if codec in {"msgpack_ragged", "json_ragged"}: + return [_normalize_structured_scalar(value) for value in values] + return values + + + +def _normalize_structured_scalar(value: Any) -> Any: + if isinstance(value, np.ndarray): + if value.dtype == object: + # Object ndarrays do not have a stable contiguous typed representation. + # For generic fallbacks, materialize their Python shape once so msgpack/json + # can preserve the logical nested values instead of treating the object array + # as a numeric ragged buffer. + return _normalize_structured_scalar(value.tolist()) + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Mapping): + return {key: _normalize_structured_scalar(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + return [_normalize_structured_scalar(item) for item in value] + return value + + + + + + + +def _encode_msgpack_ragged_values( + path: str, values: list[Any] +) -> tuple[dict[str, Any], dict[str, Any]]: + row_count = len(values) + offsets = np.empty(row_count + 1, dtype=np.int64) + nulls = np.empty(row_count, dtype=np.bool_) + offsets[0] = 0 + packer = _msgpack.Packer(use_bin_type=True, strict_types=True) + buf = bytearray() + try: + for i, value in enumerate(values): + if value is None: + nulls[i] = True + offsets[i + 1] = offsets[i] + else: + buf.extend(packer.pack(value)) + nulls[i] = False + offsets[i + 1] = len(buf) + except TypeError as exc: + raise ValueError( + f"unsupported structured non-tensor field {path}: " + "msgpack codec cannot encode value" + ) from exc + return ( + {"data": buf, "offsets": offsets, "nulls": nulls}, + {}, + ) + + + +def _decode_msgpack_ragged_values(payload: dict[str, Any], rows: int) -> list[Any]: + data = payload["data"] + offsets = payload["offsets"] + nulls = payload["nulls"] + if len(offsets) != rows + 1: + raise ValueError( + f"msgpack_ragged offsets length {len(offsets)} does not match rows {rows}" + ) + if len(nulls) != rows: + raise ValueError( + f"msgpack_ragged nulls length {len(nulls)} does not match rows {rows}" + ) + raw_data = bytes(data) if not isinstance(data, bytes) else data + values = [] + for row, is_null in enumerate(nulls): + if bool(is_null): + values.append(None) + else: + begin = int(offsets[row]) + end = int(offsets[row + 1]) + values.append(_msgpack.unpackb(raw_data[begin:end], raw=False)) + return values + + +class _OwnerBackedList(list): + def __init__(self, values: Sequence[Any], owner: Any) -> None: + super().__init__(values) + self._mooncake_pool_owner = owner + + + +class _OwnerBackedObjectArray(np.ndarray): + def __array_finalize__(self, obj: Any) -> None: + if obj is not None: + self._mooncake_pool_owner = getattr(obj, "_mooncake_pool_owner", None) + + def tolist(self) -> list[Any]: + values = super().tolist() + owner = getattr(self, "_mooncake_pool_owner", None) + if owner is None: + return values + return _OwnerBackedList(values, owner) + + + +def _object_array_from_decoded_values(values: list[Any]) -> np.ndarray: + array = np.empty(len(values), dtype=object) + array[:] = values + owner = getattr(values, "_mooncake_pool_owner", None) or next( + (getattr(v, "_mooncake_pool_owner", None) for v in values if v is not None), None + ) + if owner is None: + return array + result = array.view(_OwnerBackedObjectArray) + result._mooncake_pool_owner = owner + return result + + + +def _has_tensor_codec_helpers() -> bool: + return ( + _mooncake_store is not None + and callable(getattr(_mooncake_store, "_serialize_tensor", None)) + and callable(getattr(_mooncake_store, "_deserialize_tensor", None)) + ) + + +def _tensor_metadata_size() -> int: + helper = ( + None + if _mooncake_store is None + else getattr(_mooncake_store, "_tensor_metadata_size", None) + ) + if callable(helper): + return int(helper()) + return 0 + + +def _torch_save_payload_bytes(value: Any) -> bytes: + buffer = io.BytesIO() + _torch.save(value, buffer) + return buffer.getvalue() + + +def _deserialize_torch_save_payload(payload: bytes) -> Any: + return _torch.load(io.BytesIO(payload), weights_only=True) + + +def _slice_tensor_metadata( + metadata: bytes, shape: Sequence[int], data_bytes: int +) -> bytes: + patched = bytearray(metadata) + ctypes.c_uint64.from_buffer(patched, 32).value = int(data_bytes) + global_shape_offset = 40 + local_shape_offset = global_shape_offset + 8 * 8 + for index, dim in enumerate(shape): + ctypes.c_int64.from_buffer( + patched, global_shape_offset + index * 8 + ).value = int(dim) + ctypes.c_int64.from_buffer(patched, local_shape_offset + index * 8).value = int( + dim + ) + return bytes(patched) + + +def _tensor_payload_bytes(value: _TensorPayload) -> tuple[bytes, int]: + metadata, data_ptr, tensor_nbytes, _owner = _tensor_payload_parts(value) + if tensor_nbytes == 0: + return metadata, len(metadata) + return metadata + ctypes.string_at(data_ptr, tensor_nbytes), len(metadata) + + +def _deserialize_tensor_payload(payload: bytes) -> Any: + return _tensor_codec_helper("_deserialize_tensor")(payload) + + +def _tensor_codec_helper(name: str) -> Any: + helper = None if _mooncake_store is None else getattr(_mooncake_store, name, None) + if not callable(helper): + raise RuntimeError( + "mooncake.store tensor serialization helpers are required for structured tensor fields" + ) + return helper + + def _structured_field_specs(metadata: Mapping[str, Any]) -> dict[str, Any]: field_specs = metadata.get(STRUCTURED_FIELD_SPECS_KEY, {}) if not isinstance(field_specs, dict): @@ -1230,6 +5457,25 @@ def _structured_field_specs(metadata: Mapping[str, Any]) -> dict[str, Any]: return field_specs +def _merge_structured_stage_metadata( + old_metadata: Mapping[str, Any], new_metadata: Mapping[str, Any] +) -> dict[str, Any]: + old_dataproto = old_metadata.get("dataproto") + new_dataproto = new_metadata.get("dataproto") + if old_dataproto != new_dataproto: + raise ValueError("DataProto stage metadata mismatch during manifest merge") + merged = dict(old_metadata) + old_specs = dict(_structured_field_specs(old_metadata)) + new_specs = dict(_structured_field_specs(new_metadata)) + collisions = sorted(set(old_specs) & set(new_specs)) + if collisions: + raise ValueError(f"structured members already exist: {collisions}") + old_specs.update(new_specs) + if old_specs: + merged[STRUCTURED_FIELD_SPECS_KEY] = old_specs + return merged + + def _encode_structured_metadata(metadata: Mapping[str, Any]) -> bytes: return _encode_json_dict(metadata, "structured metadata") @@ -1252,6 +5498,43 @@ def _decode_json_dict(payload: bytes, label: str) -> dict[str, Any]: return value +def _call_write_with_optional_config(fn: Any, *args: Any, config: Any = None) -> Any: + if config is None: + return fn(*args) + return fn(*args, config=config) + + +def _put_with_optional_config( + store: BundleStore, key: str, value: Any, config: Any = None +) -> int: + return _call_write_with_optional_config(store.put, key, value, config=config) + + +def _put_from_with_optional_config( + store: BundleStore, key: str, ptr: int, size: int, config: Any = None +) -> int: + put_tensor_from = getattr(store, "put_tensor_from", None) + if config is None and callable(put_tensor_from): + return put_tensor_from(key, ptr, size) + put_from = getattr(store, "put_from", None) + if callable(put_from): + return _call_write_with_optional_config(put_from, key, ptr, size, config=config) + raise RuntimeError("put_from is unavailable") + + +def _batch_put_from_with_optional_config( + batch_put_from: Any, + keys: Sequence[str], + ptrs: Sequence[int], + sizes: Sequence[int], + config: Any = None, +) -> Sequence[int]: + return _call_write_with_optional_config( + batch_put_from, list(keys), list(ptrs), list(sizes), config=config + ) + + + def _check_status(status: Any, operation: str, key: str) -> None: if status not in (None, 0): raise RuntimeError(f"{operation} failed for {key}: {status}") @@ -1321,6 +5604,8 @@ def _cleanup_keys(store: BundleStore, keys: Sequence[str], strict: bool) -> None _torch = None # type: ignore[assignment] + + @dataclass class _CodecDecision: accepted: bool @@ -1337,12 +5622,23 @@ class _InferredLeaf: decision: _CodecDecision +@dataclass +class _EncodedStructuredLeaf: + codec: str + rows: int + payload: dict[str, Any] + metadata: dict[str, Any] + + class _Missing: """Sentinel for dict keys absent from a row (distinct from value-is-None).""" + __slots__ = () + def __repr__(self) -> str: return "" + MISSING = _Missing() @@ -1369,10 +5665,27 @@ def _is_bytes_like(value: Any) -> bool: def _is_media_list(value: Any) -> bool: - return ( - isinstance(value, (list, tuple)) - and len(value) > 0 - and all(_is_pil_image(item) or _is_bytes_like(item) for item in value) + return isinstance(value, (list, tuple)) and all( + _is_pil_image(item) or _is_bytes_like(item) for item in value + ) + + +def _can_media_list(values: list[Any]) -> _CodecDecision: + nn = _non_null(values) + if not nn: + return _CodecDecision( + False, "media_list_ragged", "all rows are null", "media list" + ) + if not all(_is_media_list(v) for v in nn): + return _CodecDecision( + False, "media_list_ragged", "not all rows are media list", "media list" + ) + if not any(len(v) > 0 for v in nn): + return _CodecDecision( + False, "media_list_ragged", "all media lists are empty", "media list" + ) + return _CodecDecision( + True, "media_list_ragged", "all non-null rows are media list", "media list" ) @@ -1386,31 +5699,53 @@ def _check_all( if not nn: return _CodecDecision(False, codec, "all rows are null", normalized_type) if not all(predicate(v) for v in nn): - return _CodecDecision(False, codec, f"not all rows are {normalized_type}", normalized_type) - return _CodecDecision(True, codec, f"all non-null rows are {normalized_type}", normalized_type) + return _CodecDecision( + False, codec, f"not all rows are {normalized_type}", normalized_type + ) + return _CodecDecision( + True, codec, f"all non-null rows are {normalized_type}", normalized_type + ) def _can_tensor(values: list[Any]) -> _CodecDecision: if _torch is None: - return _CodecDecision(False, "ragged_tensor", "torch is not available", "torch.Tensor") + return _CodecDecision( + False, "ragged_tensor", "torch is not available", "torch.Tensor" + ) nn = _non_null(values) if not nn: - return _CodecDecision(False, "ragged_tensor", "all rows are null", "torch.Tensor") + return _CodecDecision( + False, "ragged_tensor", "all rows are null", "torch.Tensor" + ) if not all(isinstance(v, _torch.Tensor) for v in nn): - return _CodecDecision(False, "ragged_tensor", "not all rows are Tensor", "torch.Tensor") + return _CodecDecision( + False, "ragged_tensor", "not all rows are Tensor", "torch.Tensor" + ) dtypes = sorted({str(v.dtype) for v in nn}) if len(dtypes) != 1: - return _CodecDecision(False, "ragged_tensor", f"mixed dtypes: {dtypes}", "torch.Tensor") + return _CodecDecision( + False, "ragged_tensor", f"mixed dtypes: {dtypes}", "torch.Tensor" + ) ndims = {v.ndim for v in nn} if len(ndims) != 1: - return _CodecDecision(False, "ragged_tensor", f"mixed dimensions: {ndims}", "torch.Tensor") - return _CodecDecision(True, "ragged_tensor", "all non-null rows are Tensor", "torch.Tensor", {"dtype": dtypes[0]}) + return _CodecDecision( + False, "ragged_tensor", f"mixed dimensions: {ndims}", "torch.Tensor" + ) + return _CodecDecision( + True, + "ragged_tensor", + "all non-null rows are Tensor", + "torch.Tensor", + {"dtype": dtypes[0]}, + ) def _can_numeric_sequence(values: list[Any]) -> _CodecDecision: nn = _non_null(values) if not nn: - return _CodecDecision(False, "typed_ragged", "all rows are null", "numeric sequence") + return _CodecDecision( + False, "typed_ragged", "all rows are null", "numeric sequence" + ) dtypes = [] for v in nn: if isinstance(v, np.ndarray): @@ -1419,17 +5754,37 @@ def _can_numeric_sequence(values: list[Any]) -> _CodecDecision: try: arr = np.asarray(v) except (ValueError, TypeError): - return _CodecDecision(False, "typed_ragged", "row cannot be converted to ndarray", "numeric sequence") + return _CodecDecision( + False, + "typed_ragged", + "row cannot be converted to ndarray", + "numeric sequence", + ) else: - return _CodecDecision(False, "typed_ragged", "row is not array-like", "numeric sequence") + return _CodecDecision( + False, "typed_ragged", "row is not array-like", "numeric sequence" + ) if arr.dtype == object or not np.issubdtype(arr.dtype, np.number): - return _CodecDecision(False, "typed_ragged", f"non-numeric dtype: {arr.dtype}", "numeric sequence") + return _CodecDecision( + False, + "typed_ragged", + f"non-numeric dtype: {arr.dtype}", + "numeric sequence", + ) dtypes.append(arr.dtype) try: dtype = np.result_type(*dtypes) except (TypeError, ValueError, OverflowError): - return _CodecDecision(False, "typed_ragged", "cannot determine common dtype", "numeric sequence") - return _CodecDecision(True, "typed_ragged", "all rows promote to common numeric dtype", "numeric sequence", {"dtype": str(dtype)}) + return _CodecDecision( + False, "typed_ragged", "cannot determine common dtype", "numeric sequence" + ) + return _CodecDecision( + True, + "typed_ragged", + "all rows promote to common numeric dtype", + "numeric sequence", + {"dtype": str(dtype)}, + ) def _can_numeric_scalar(values: list[Any]) -> _CodecDecision: @@ -1437,14 +5792,53 @@ def _can_numeric_scalar(values: list[Any]) -> _CodecDecision: if not nn: return _CodecDecision(False, "ndarray", "all rows are null", "numeric scalar") if not all(isinstance(v, (bool, int, float, np.number)) for v in nn): - return _CodecDecision(False, "ndarray", "not all rows are numeric scalar", "numeric scalar") + return _CodecDecision( + False, "ndarray", "not all rows are numeric scalar", "numeric scalar" + ) try: dtype = np.result_type(*nn) except (TypeError, ValueError, OverflowError): - return _CodecDecision(False, "ndarray", "cannot determine common dtype", "numeric scalar") + return _CodecDecision( + False, "ndarray", "cannot determine common dtype", "numeric scalar" + ) if not np.issubdtype(dtype, np.number) and not np.issubdtype(dtype, np.bool_): - return _CodecDecision(False, "ndarray", f"non-numeric dtype: {dtype}", "numeric scalar") - return _CodecDecision(True, "ndarray", "all rows are numeric scalar", "numeric scalar", {"dtype": str(dtype)}) + return _CodecDecision( + False, "ndarray", f"non-numeric dtype: {dtype}", "numeric scalar" + ) + return _CodecDecision( + True, + "ndarray", + "all rows are numeric scalar", + "numeric scalar", + {"dtype": str(dtype)}, + ) + + +def _can_msgpack(values: list[Any]) -> _CodecDecision: + nn = _non_null(values) + if not nn: + return _CodecDecision(False, "msgpack_ragged", "all rows are null", "msgpack") + if not all(isinstance(v, (dict, list, tuple)) for v in nn): + return _CodecDecision( + False, "msgpack_ragged", "not all rows are structured objects", "msgpack" + ) + sampled_bytes = 0 + for i, v in enumerate(nn): + try: + encoded = _msgpack.packb(v, use_bin_type=True, strict_types=True) + except (TypeError, ValueError, OverflowError): + return _CodecDecision( + False, "msgpack_ragged", "serialization failed", "msgpack" + ) + if i < _INFER_MAX_SAMPLE_ROWS: + sampled_bytes += len(encoded) + if sampled_bytes > _INFER_MAX_JSON_BYTES: + return _CodecDecision( + False, "msgpack_ragged", "sampled payload too large", "msgpack" + ) + return _CodecDecision( + True, "msgpack_ragged", "all rows pass msgpack serialization", "msgpack" + ) def _can_json(values: list[Any]) -> _CodecDecision: @@ -1454,34 +5848,67 @@ def _can_json(values: list[Any]) -> _CodecDecision: sampled_bytes = 0 for i, v in enumerate(nn): try: - encoded = json.dumps(v, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + encoded = json.dumps(v, ensure_ascii=False, separators=(",", ":")).encode( + "utf-8" + ) except (TypeError, ValueError, OverflowError, RecursionError): return _CodecDecision(False, "json_ragged", "serialization failed", "json") if i < _INFER_MAX_SAMPLE_ROWS: sampled_bytes += len(encoded) if sampled_bytes > _INFER_MAX_JSON_BYTES: return _CodecDecision(False, "json_ragged", "sampled payload too large", "json") - return _CodecDecision(True, "json_ragged", "all rows pass JSON serialization", "json") - - -_CODEC_PREDICATES: tuple[Any, ...] = ( - _can_tensor, - lambda v: _check_all(v, _is_media_list, "media_list_ragged", "media list"), - _can_numeric_sequence, - _can_numeric_scalar, - lambda v: _check_all(v, _is_bytes_like, "bytes_ragged", "bytes-like"), - lambda v: _check_all(v, _is_pil_image, "media_bytes", "media"), - lambda v: _check_all(v, lambda x: isinstance(x, str), "utf8_ragged", "str"), - _can_json, -) + return _CodecDecision( + True, "json_ragged", "all rows pass JSON serialization", "json" + ) def _choose_leaf_codec(values: list[Any]) -> _CodecDecision: - for predicate in _CODEC_PREDICATES: - decision = predicate(values) - if decision.accepted: - return decision - return _CodecDecision(False, "pickle_ragged_fallback", "no optimized codec matched", "python object") + nn = _non_null(values) + if not nn: + return _CodecDecision(False, "msgpack_ragged", "all rows are null", "msgpack") + if _torch is not None and all(isinstance(value, _torch.Tensor) for value in nn): + dtypes = {value.dtype for value in nn} + if len(dtypes) != 1: + return _CodecDecision( + False, + "ragged_tensor", + f"mixed tensor dtype: {sorted(str(dtype) for dtype in dtypes)}", + "torch.Tensor", + ) + dtype = str(nn[0].dtype) + return _CodecDecision(True, "ragged_tensor", "all non-null rows are Tensor", "torch.Tensor", {"dtype": dtype}) + if all(_is_media_list(value) for value in nn): + return _CodecDecision(True, "media_list_ragged", "all rows are media lists", "media list") + if all(isinstance(value, np.ndarray) for value in nn): + if all(np.issubdtype(value.dtype, np.number) for value in nn): + return _CodecDecision( + True, + "typed_ragged", + "all rows are numeric ndarray", + "numeric sequence", + {"recursive_source": "ndarray"}, + ) + return _CodecDecision(False, "msgpack_ragged", "object ndarray rows", "python object") + if all(isinstance(value, (list, tuple)) for value in nn): + try: + dtypes = [np.asarray(value).dtype for value in nn] + dtype = np.result_type(*dtypes) + except (TypeError, ValueError, OverflowError): + dtype = None + if dtype is not None and np.issubdtype(dtype, np.number): + return _CodecDecision(True, "typed_ragged", "all rows are numeric sequences", "numeric sequence", {"dtype": str(dtype)}) + if all(isinstance(value, (bool, int, float, np.number)) for value in nn): + dtype = np.result_type(*nn) + return _CodecDecision(True, "ndarray", "all rows are numeric scalar", "numeric scalar", {"dtype": str(dtype)}) + if all(_is_bytes_like(value) for value in nn): + return _CodecDecision(True, "bytes_ragged", "all rows are bytes-like", "bytes-like") + if all(_is_pil_image(value) for value in nn): + return _CodecDecision(True, "media_bytes", "all rows are media", "media") + if all(isinstance(value, str) for value in nn): + return _CodecDecision(True, "utf8_ragged", "all rows are str", "str") + if all(isinstance(value, (dict, list, tuple)) for value in nn): + return _CodecDecision(True, "msgpack_ragged", "all rows are msgpack-like", "msgpack") + return _CodecDecision(False, "msgpack_ragged", "no codec matched", "python object") def _try_expand_dict(values: list[Any]) -> list[str] | None: @@ -1501,7 +5928,9 @@ def _try_expand_list(values: list[Any]) -> tuple[int, list[int]] | None: max_len = max(len(v) for v in nn) if max_len > _INFER_MAX_LIST_LEN: return None - if not all(item is None or isinstance(item, (dict, list, tuple)) for v in nn for item in v): + if not all( + item is None or isinstance(item, (dict, list, tuple)) for v in nn for item in v + ): return None lengths = [len(v) if isinstance(v, (list, tuple)) else 0 for v in values] return max_len, lengths @@ -1527,24 +5956,35 @@ def infer_structure( ``row_mask`` that records which rows had a real parent container. """ if _depth > _INFER_MAX_DEPTH: - raise ValueError(f"infer_structure exceeded max depth {_INFER_MAX_DEPTH} at {path!r}") + raise ValueError( + f"infer_structure exceeded max depth {_INFER_MAX_DEPTH} at {path!r}" + ) dict_keys = _try_expand_dict(values) if dict_keys is not None: row_mask = [isinstance(v, dict) for v in values] nodes.append(_InferredNode(path, "dict", dict_keys, row_mask=row_mask)) for key in dict_keys: - child = [v.get(key, MISSING) if isinstance(v, dict) else None for v in values] - infer_structure(f"{path}.{_escape_key(key)}", child, leaves, nodes, _depth=_depth + 1) + child = [ + v.get(key, MISSING) if isinstance(v, dict) else None for v in values + ] + infer_structure( + f"{path}.{_escape_key(key)}", child, leaves, nodes, _depth=_depth + 1 + ) return list_result = _try_expand_list(values) if list_result is not None: max_len, lengths = list_result row_mask = [isinstance(v, (list, tuple)) for v in values] - nodes.append(_InferredNode(path, "list", list(range(max_len)), lengths, row_mask=row_mask)) + nodes.append( + _InferredNode( + path, "list", list(range(max_len)), lengths, row_mask=row_mask + ) + ) for index in range(max_len): child = [ - v[index] if isinstance(v, (list, tuple)) and index < len(v) else - (MISSING if isinstance(v, (list, tuple)) else None) + v[index] + if isinstance(v, (list, tuple)) and index < len(v) + else (MISSING if isinstance(v, (list, tuple)) else None) for v in values ] infer_structure(f"{path}[{index}]", child, leaves, nodes, _depth=_depth + 1) diff --git a/mooncake-wheel/pyproject.toml b/mooncake-wheel/pyproject.toml index 7a3af227d0..87deee582f 100644 --- a/mooncake-wheel/pyproject.toml +++ b/mooncake-wheel/pyproject.toml @@ -1,29 +1,48 @@ +# mooncake-wheel/pyproject.toml +# +# This file is built from this directory; scripts/build_wheel.sh `cd`s +# here before invoking setuptools. Do not move this file or its parent +# build script out of sync. +# +# NOTE: the `keywords` and `classifiers` arrays below MUST stay on a +# single line. They are sed-substituted by scripts/build_wheel.sh for +# each build variant (non-cuda, cuda13, npu, efa, efa-cuda13, efa-non-cuda, +# musa); +# reordering or wrapping will silently break the variant build matrix. +# +# `readme = "README.md"` points at a per-build copy of the repo-root +# README.md, materialized by scripts/build_wheel.sh (which `cp`s +# ../README.md into this directory before invoking the build backend, and +# `rm`s it after). This avoids modern setuptools' path-sandbox check +# (rejects `../`-traversal) while keeping the root README.md the single +# source of truth for the long description. + [build-system] -requires = ["setuptools>=61.0.0", "wheel>=0.37.0"] +requires = ["setuptools>=61.0.0", "wheel>=0.37.0", "numpy"] build-backend = "setuptools.build_meta" [project] name = "mooncake-transfer-engine" -version = "0.3.11.post1" -description = "Python binding of a Mooncake library using pybind11" +version = "0.3.12.post1" +description = "A KVCache-centric Disaggregated Architecture for large-scale LLM inference and training." authors = [ { name = "Mooncake Authors" } ] -requires-python = ">=3.8" -dependencies = ["aiohttp", "requests"] -keywords = ["mooncake", "data transfer", "kv cache", "llm inference"] -classifiers = [ - "Programming Language :: Python :: 3", - "Programming Language :: C++", - "Operating System :: POSIX :: Linux", -] -readme = "" +requires-python = ">=3.10" +readme = "README.md" # local copy of ../README.md, materialized at build time by scripts/build_wheel.sh +dependencies = ["aiohttp", "requests", "msgpack"] +# Single-line arrays: sed-substituted by scripts/build_wheel.sh for each variant. +keywords = ["mooncake", "transfer engine", "kv cache", "llm inference", "rdma"] +classifiers = ["Development Status :: 5 - Production/Stable", "Environment :: GPU :: NVIDIA CUDA", "Intended Audience :: Developers", "Intended Audience :: Science/Research", "License :: OSI Approved :: Apache Software License", "Operating System :: POSIX :: Linux", "Programming Language :: C++", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Programming Language :: Python :: Implementation :: CPython", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: System :: Distributed Computing"] [project.urls] Homepage = "https://github.com/kvcache-ai/Mooncake" Documentation = "https://kvcache-ai.github.io/Mooncake" Source = "https://github.com/kvcache-ai/Mooncake" Issues = "https://github.com/kvcache-ai/Mooncake/issues" +Blog = "https://kvcache.ai/blog" +Changelog = "https://github.com/kvcache-ai/Mooncake/releases" +Slack = "https://join.slack.com/t/mooncake-project/shared_invite/zt-3qx4x35ea-zSSTqTHItHJs9SCoXLOSPA" [project.scripts] mooncake_master = "mooncake.cli:main" diff --git a/mooncake-wheel/setup.py b/mooncake-wheel/setup.py index 8891f2f343..56cb1c5a32 100644 --- a/mooncake-wheel/setup.py +++ b/mooncake-wheel/setup.py @@ -1,6 +1,8 @@ import sys import platform -from setuptools import setup, Distribution + +import numpy as np +from setuptools import setup, Distribution, Extension from wheel.bdist_wheel import bdist_wheel # --------------------------------------------------------------------------- @@ -162,10 +164,24 @@ def finalize_options(self): self.plat_name = get_platform() + +# --------------------------------------------------------------------------- +# C extensions +# --------------------------------------------------------------------------- +_fast_copy_ext = Extension( + "mooncake._fast_copy", + sources=["mooncake/_fast_copy.c"], + include_dirs=[np.get_include()], + extra_compile_args=["-O2", "-pthread"], + extra_link_args=["-lpthread"], +) + + # --------------------------------------------------------------------------- # setup() # --------------------------------------------------------------------------- setup( distclass=BinaryDistribution, cmdclass={"bdist_wheel": CustomBdistWheel}, + ext_modules=[_fast_copy_ext], ) diff --git a/mooncake-wheel/tests/ep_test_utils.py b/mooncake-wheel/tests/ep_test_utils.py index 6504ab9524..fa6b97acd3 100644 --- a/mooncake-wheel/tests/ep_test_utils.py +++ b/mooncake-wheel/tests/ep_test_utils.py @@ -1,3 +1,4 @@ +import importlib import os import sys import numpy as np @@ -5,7 +6,8 @@ import torch.distributed as dist from typing import Optional -import mooncake.pg +# Side-effect import: registers the mooncake process-group backend. +importlib.import_module("mooncake.pg") def init_dist(local_rank: int, num_local_ranks: int): diff --git a/mooncake-wheel/tests/test_cli.py b/mooncake-wheel/tests/test_cli.py index de8ed0fe16..885411de24 100644 --- a/mooncake-wheel/tests/test_cli.py +++ b/mooncake-wheel/tests/test_cli.py @@ -3,7 +3,6 @@ Test script to verify that the mooncake_master entry point works correctly. """ -import os import sys import subprocess import time @@ -55,9 +54,21 @@ def test_entry_point_installed(): def test_run_master_and_client(): """Test running the master service through the entry point.""" try: - # Run mooncake_master with a non-default port to avoid conflicts + # Use non-default ports to avoid collisions with processes started + # earlier in the CI test-wheel-ubuntu job: + # * --port: gRPC port. The test client connects via + # --master_server_address (direct gRPC), so it does not need the + # embedded HTTP metadata server on the default 8080 — the job runs a + # shared Python mooncake_http_metadata_server on 8080 for the whole + # run_tests.sh suite, so enabling it here would collide and + # LOG(FATAL) (master.cpp:1449). + # * --metrics_port: the master admin server always binds this port + # (master.cpp:1485, created unconditionally and cannot be disabled), + # and the default 9003 is still held by a lingering master from the + # prior CXL test step when this smoke test runs. process = subprocess.Popen( - ["mooncake_master", "--port=61351", "--max_threads=2", "--enable_http_metadata_server=true"], + ["mooncake_master", "--port=61351", "--max_threads=2", + "--metrics_port=19003"], stdout=subprocess.PIPE, stderr=subprocess.PIPE ) diff --git a/mooncake-wheel/tests/test_distributed_object_store.py b/mooncake-wheel/tests/test_distributed_object_store.py index 9e5d9e96bd..880a3b532c 100644 --- a/mooncake-wheel/tests/test_distributed_object_store.py +++ b/mooncake-wheel/tests/test_distributed_object_store.py @@ -2,7 +2,6 @@ import os import time import threading -import random from mooncake.store import MooncakeDistributedStore # The lease time of the kv object, should be set equal to @@ -303,7 +302,7 @@ def test_get_into_ranges_operations(self): self.assertEqual(bytes(buffer1[16:20]), data1[10:14]) mismatch_results = self.store.get_into_ranges( - [buffer_ptr0], [[key1, key2]], [[[0], []]], [[[0, 1], []]], [[[4, 4], []]] + [buffer_ptr0], [[key1, key2]], [[[0], []]], [[[0, 1], []]], [[[], []]] ) self.assertEqual(len(mismatch_results), 1) self.assertEqual(len(mismatch_results[0]), 2) diff --git a/mooncake-wheel/tests/test_distributed_object_store_cxl.py b/mooncake-wheel/tests/test_distributed_object_store_cxl.py index d883fc90cf..a08382f5dc 100644 --- a/mooncake-wheel/tests/test_distributed_object_store_cxl.py +++ b/mooncake-wheel/tests/test_distributed_object_store_cxl.py @@ -2,7 +2,6 @@ import os import time import threading -import random import tempfile from mooncake.store import MooncakeDistributedStore diff --git a/mooncake-wheel/tests/test_dummy_client.py b/mooncake-wheel/tests/test_dummy_client.py index d3ead6482e..c35c9d9492 100644 --- a/mooncake-wheel/tests/test_dummy_client.py +++ b/mooncake-wheel/tests/test_dummy_client.py @@ -2,7 +2,12 @@ import os import time import threading -import random + +try: + import torch +except ImportError: + torch = None + from mooncake.store import MooncakeDistributedStore # The lease time of the kv object, should be set equal to @@ -322,6 +327,234 @@ def test_batch_put_from_operations(self): for key in keys: self.assertEqual(self.store.remove(key), 0) + def test_tensor_operations(self): + """Test tensor wrappers through dummy-client shared memory.""" + import ctypes + + if torch is None: + self.skipTest("PyTorch is not available") + + key = f"test_dummy_tensor_{os.getpid()}" + key_from = f"{key}_from" + batch_keys = [f"{key}_batch_{i}" for i in range(2)] + upsert_key = f"{key}_upsert" + pub_key = f"{key}_pub" + tp_key = f"{key}_tp" + batch_tp_key = f"{key}_batch_tp" + parallel_key = f"{key}_parallel" + parallel_batch_keys = [f"{key}_parallel_batch_{i}" for i in range(2)] + parallel_from_key = f"{key}_parallel_from" + parallel_upsert_key = f"{key}_parallel_upsert" + parallel_upsert_from_key = f"{key}_parallel_upsert_from" + cleanup_keys = [key, key_from, *batch_keys, upsert_key, pub_key] + cleanup_keys.extend( + [ + parallel_key, + *parallel_batch_keys, + parallel_from_key, + parallel_upsert_key, + parallel_upsert_from_key, + ] + ) + cleanup_keys.extend(f"{tp_key}_tp_{rank}" for rank in range(2)) + cleanup_keys.extend(f"{batch_tp_key}_tp_{rank}" for rank in range(2)) + tensor = torch.arange(12, dtype=torch.float32).reshape(3, 4) + + self.assertEqual(self.store.put_tensor(key, tensor), 0) + retrieved = self.store.get_tensor(key) + self.assertIsNotNone(retrieved) + self.assertEqual(retrieved.dtype, tensor.dtype) + self.assertEqual(tuple(retrieved.shape), tuple(tensor.shape)) + self.assertTrue(torch.equal(retrieved, tensor)) + + buffer_size = self.store.get_size(key) + self.assertGreater(buffer_size, 0) + buffer_ptr = self.store.alloc_from_mem_pool(buffer_size) + self.assertNotEqual(buffer_ptr, 0) + self.assertEqual(self.store.register_buffer(buffer_ptr, buffer_size), 0) + try: + ctypes.memset(buffer_ptr, 0, buffer_size) + into = self.store.get_tensor_into(key, buffer_ptr, buffer_size) + self.assertIsNotNone(into) + self.assertTrue(torch.equal(into, tensor)) + + mismatch = self.store.batch_get_tensor_into( + [key, f"{key}_missing_arg"], [buffer_ptr], [buffer_size] + ) + self.assertEqual(len(mismatch), 2) + self.assertTrue(all(result < 0 for result in mismatch)) + + self.assertEqual( + self.store.put_tensor_from(key_from, buffer_ptr, buffer_size), + 0, + ) + from_tensor = self.store.get_tensor(key_from) + self.assertIsNotNone(from_tensor) + self.assertTrue(torch.equal(from_tensor, tensor)) + + batch_tensors = [tensor, tensor + 1] + batch_results = self.store.batch_put_tensor(batch_keys, batch_tensors) + self.assertEqual(list(batch_results), [0, 0]) + batch_retrieved = self.store.batch_get_tensor(batch_keys) + for expected, actual in zip(batch_tensors, batch_retrieved): + self.assertIsNotNone(actual) + self.assertTrue(torch.equal(actual, expected)) + + updated = tensor + 2 + self.assertEqual(self.store.upsert_tensor(upsert_key, updated), 0) + self.assertTrue(torch.equal(self.store.get_tensor(upsert_key), updated)) + + self.assertEqual(self.store.pub_tensor(pub_key, tensor + 3), 0) + self.assertTrue(torch.equal(self.store.get_tensor(pub_key), tensor + 3)) + + self.assertEqual( + self.store.put_tensor_with_tp( + tp_key, tensor, tp_size=2, split_dim=1 + ), + 0, + ) + tp_shards = torch.chunk(tensor, 2, dim=1) + for rank, expected in enumerate(tp_shards): + actual = self.store.get_tensor_with_tp( + tp_key, tp_rank=rank, tp_size=2 + ) + self.assertIsNotNone(actual) + self.assertTrue(torch.equal(actual, expected.contiguous())) + + batch_tp_results = self.store.batch_put_tensor_with_tp( + [batch_tp_key], [tensor], tp_size=2, split_dim=1 + ) + self.assertEqual(list(batch_tp_results), [0]) + + self.assertEqual( + self.store.put_tensor_with_parallelism(parallel_key, tensor), + 0, + ) + self.assertTrue(torch.equal(self.store.get_tensor(parallel_key), tensor)) + + parallel_batch = [tensor + 4, tensor + 5] + parallel_batch_results = self.store.batch_put_tensor_with_parallelism( + parallel_batch_keys, parallel_batch + ) + self.assertEqual(list(parallel_batch_results), [0, 0]) + parallel_batch_retrieved = self.store.batch_get_tensor( + parallel_batch_keys + ) + for expected, actual in zip(parallel_batch, parallel_batch_retrieved): + self.assertIsNotNone(actual) + self.assertTrue(torch.equal(actual, expected)) + + self.assertEqual( + self.store.put_tensor_with_parallelism_from( + parallel_from_key, buffer_ptr, buffer_size + ), + 0, + ) + self.assertTrue( + torch.equal(self.store.get_tensor(parallel_from_key), tensor) + ) + + parallel_update = tensor + 6 + self.assertEqual( + self.store.upsert_tensor_with_parallelism( + parallel_upsert_key, parallel_update + ), + 0, + ) + self.assertTrue( + torch.equal( + self.store.get_tensor(parallel_upsert_key), parallel_update + ) + ) + + parallel_upsert_from_results = ( + self.store.batch_upsert_tensor_with_parallelism_from( + [parallel_upsert_from_key], [buffer_ptr], [buffer_size] + ) + ) + self.assertEqual(list(parallel_upsert_from_results), [0]) + self.assertTrue( + torch.equal( + self.store.get_tensor(parallel_upsert_from_key), tensor + ) + ) + finally: + self.store.unregister_buffer(buffer_ptr) + for cleanup_key in cleanup_keys: + self.store.remove(cleanup_key) + + def test_00_mixed_put_and_put_tensor_concurrency(self): + """Regression test for regular put and tensor put sharing dummy SHM.""" + if torch is None: + self.skipTest("PyTorch is not available") + + prefix = f"test_dummy_mixed_put_tensor_{os.getpid()}" + iterations = 16 + regular_keys = [f"{prefix}_regular_{i}" for i in range(iterations)] + tensor_keys = [f"{prefix}_tensor_{i}" for i in range(iterations)] + regular_values = [ + (f"regular-payload-{i}-".encode() * 4096) + for i in range(iterations) + ] + tensors = [ + (torch.arange(4096, dtype=torch.float32) + i).reshape(64, 64) + for i in range(iterations) + ] + + start = threading.Event() + errors = [] + errors_lock = threading.Lock() + + def record_error(message): + with errors_lock: + errors.append(message) + + def put_regular_values(): + try: + start.wait() + for key, value in zip(regular_keys, regular_values): + ret = self.store.put(key, value) + if ret != 0: + record_error(f"put({key}) failed with {ret}") + except Exception as exc: + record_error(f"put thread raised {exc!r}") + + def put_tensors(): + try: + start.wait() + for key, tensor in zip(tensor_keys, tensors): + ret = self.store.put_tensor(key, tensor) + if ret != 0: + record_error(f"put_tensor({key}) failed with {ret}") + except Exception as exc: + record_error(f"put_tensor thread raised {exc!r}") + + threads = [ + threading.Thread(target=put_regular_values), + threading.Thread(target=put_tensors), + ] + try: + for thread in threads: + thread.start() + start.set() + for thread in threads: + thread.join() + + self.assertEqual(errors, []) + + for key, expected in zip(regular_keys, regular_values): + self.assertEqual(self.store.get(key), expected) + for key, expected in zip(tensor_keys, tensors): + actual = self.store.get_tensor(key) + self.assertIsNotNone(actual) + self.assertTrue(torch.equal(actual, expected)) + finally: + for thread in threads: + if thread.is_alive(): + thread.join() + for key in [*regular_keys, *tensor_keys]: + self.store.remove(key, force=True) + # Mark this test as zzz_ so that it is the last test to run def zzz_test_dict_fuzz_e2e(self): """End-to-end fuzz test comparing distributed store behavior with dict. diff --git a/mooncake-wheel/tests/test_fast_copy.py b/mooncake-wheel/tests/test_fast_copy.py new file mode 100644 index 0000000000..3172d14c90 --- /dev/null +++ b/mooncake-wheel/tests/test_fast_copy.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import ctypes +import unittest + +import numpy as np + +try: + from mooncake._fast_copy import concat_arrays_into +except Exception as error: # pragma: no cover - depends on built extension + concat_arrays_into = None + _FAST_COPY_IMPORT_ERROR = error +else: + _FAST_COPY_IMPORT_ERROR = None + + +@unittest.skipIf( + concat_arrays_into is None, + f"native fast-copy extension is unavailable: {_FAST_COPY_IMPORT_ERROR}", +) +class FastCopyTests(unittest.TestCase): + def test_concat_arrays_into_copies_selected_range(self) -> None: + arrays = [ + np.arange(3, dtype=np.uint8), + np.arange(4, dtype=np.uint8) + 10, + np.arange(2, dtype=np.uint8) + 20, + ] + expected = np.concatenate(arrays[1:]) + destination = ctypes.create_string_buffer(expected.nbytes) + + copied = concat_arrays_into( + arrays, ctypes.addressof(destination), len(destination), 1, 2 + ) + + self.assertEqual(copied, expected.nbytes) + self.assertEqual(destination.raw[:copied], expected.tobytes()) + + def test_concat_arrays_into_rejects_small_destination(self) -> None: + arrays = [np.arange(4, dtype=np.uint8), np.arange(4, dtype=np.uint8)] + destination = ctypes.create_string_buffer(7) + + with self.assertRaisesRegex(ValueError, "destination buffer too small"): + concat_arrays_into(arrays, ctypes.addressof(destination), len(destination)) + + def test_concat_arrays_into_rejects_non_contiguous_source(self) -> None: + arrays = [np.arange(8, dtype=np.uint8)[::2]] + destination = ctypes.create_string_buffer(arrays[0].nbytes) + + with self.assertRaisesRegex(ValueError, "C-contiguous"): + concat_arrays_into(arrays, ctypes.addressof(destination), len(destination)) + + def test_concat_arrays_into_empty_range_returns_zero(self) -> None: + arrays = [np.arange(4, dtype=np.uint8)] + destination = ctypes.create_string_buffer(1) + + copied = concat_arrays_into(arrays, ctypes.addressof(destination), 0, 0, 0) + + self.assertEqual(copied, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/mooncake-wheel/tests/test_http_metadata_server.py b/mooncake-wheel/tests/test_http_metadata_server.py new file mode 100644 index 0000000000..a494d5f203 --- /dev/null +++ b/mooncake-wheel/tests/test_http_metadata_server.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +import sys +import unittest +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from mooncake.http_metadata_server import KVBootstrapServer + + +class FakeRequest: + def __init__(self, method, key=None, body=b""): + self.method = method + self.query = {} if key is None else {"key": key} + self.body = body + + async def read(self): + return self.body + + +class HttpMetadataServerTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.server = KVBootstrapServer(port=0) + + async def test_missing_metadata_key_is_rejected_for_all_methods(self): + for method in ("GET", "PUT", "DELETE"): + with self.subTest(method=method): + response = await self.server._handle_metadata( + FakeRequest(method, body=b"value") + ) + + self.assertEqual(response.status, 400) + self.assertEqual(response.content_type, "application/json") + self.assertNotIn("", self.server.store) + + async def test_empty_metadata_key_is_rejected(self): + response = await self.server._handle_metadata( + FakeRequest("PUT", key="", body=b"value") + ) + + self.assertEqual(response.status, 400) + self.assertEqual(response.content_type, "application/json") + self.assertNotIn("", self.server.store) + + async def test_blank_metadata_key_is_rejected(self): + response = await self.server._handle_metadata( + FakeRequest("PUT", key=" ", body=b"value") + ) + + self.assertEqual(response.status, 400) + self.assertEqual(response.content_type, "application/json") + self.assertNotIn(" ", self.server.store) + + async def test_metadata_key_is_stripped_before_operations(self): + put_response = await self.server._handle_metadata( + FakeRequest("PUT", key=" valid ", body=b"value") + ) + get_response = await self.server._handle_metadata( + FakeRequest("GET", key=" valid ") + ) + + self.assertEqual(put_response.status, 200) + self.assertEqual(get_response.status, 200) + self.assertEqual(get_response.body, b"value") + self.assertIn("valid", self.server.store) + self.assertNotIn(" valid ", self.server.store) + + async def test_valid_metadata_key_still_round_trips(self): + put_response = await self.server._handle_metadata( + FakeRequest("PUT", key="valid", body=b"value") + ) + get_response = await self.server._handle_metadata( + FakeRequest("GET", key="valid") + ) + + self.assertEqual(put_response.status, 200) + self.assertEqual(get_response.status, 200) + self.assertEqual(get_response.body, b"value") + + +if __name__ == "__main__": + unittest.main() diff --git a/mooncake-wheel/tests/test_mooncake_config.py b/mooncake-wheel/tests/test_mooncake_config.py index 39b57beaef..b160fe03f7 100644 --- a/mooncake-wheel/tests/test_mooncake_config.py +++ b/mooncake-wheel/tests/test_mooncake_config.py @@ -3,6 +3,7 @@ import tempfile import unittest +from mooncake import mooncake_config as _cfg_mod from mooncake.mooncake_config import ( MooncakeConfig, DEFAULT_GLOBAL_SEGMENT_SIZE, @@ -27,7 +28,10 @@ def setUp(self): "protocol": "tcp", "device_name": "eth0", "enable_ssd_offload": True, - "ssd_offload_path": "/nvme/mooncake_offload" + "ssd_offload_path": "/nvme/mooncake_offload", + "tenant_id": "tenant-a", + "enable_client_http_server": True, + "client_http_port": 19300, } def tearDown(self): @@ -35,7 +39,7 @@ def tearDown(self): def write_config(self, config_data): """Write configuration to file""" - with open(self.config_file, 'w') as f: + with open(self.config_file, "w") as f: json.dump(config_data, f) def test_load_valid_config(self): @@ -52,13 +56,16 @@ def test_load_valid_config(self): self.assertEqual(config.device_name, "eth0") self.assertEqual(config.enable_ssd_offload, True) self.assertEqual(config.ssd_offload_path, "/nvme/mooncake_offload") + self.assertEqual(config.tenant_id, "tenant-a") + self.assertEqual(config.enable_client_http_server, True) + self.assertEqual(config.client_http_port, 19300) def test_load_with_default_values(self): """Test loading configuration with default values""" minimal_config = { "local_hostname": "localhost", "metadata_server": "localhost:8080", - "master_server_address": "localhost:8081" + "master_server_address": "localhost:8081", } self.write_config(minimal_config) config = MooncakeConfig.from_file(self.config_file) @@ -69,6 +76,93 @@ def test_load_with_default_values(self): self.assertEqual(config.device_name, "") self.assertEqual(config.enable_ssd_offload, False) self.assertEqual(config.ssd_offload_path, "") + self.assertEqual(config.tenant_id, "default") + self.assertEqual(config.enable_client_http_server, False) + self.assertEqual(config.client_http_port, 9300) + + def test_load_tenant_id_from_file(self): + """Test loading tenant_id from configuration file""" + self.write_config({**self.valid_config, "tenant_id": "tenant-from-file"}) + config = MooncakeConfig.from_file(self.config_file) + + self.assertEqual(config.tenant_id, "tenant-from-file") + + def test_tenant_id_defaults(self): + """Test tenant_id defaults to default when omitted""" + minimal_config = { + "local_hostname": "localhost", + "metadata_server": "localhost:8080", + "master_server_address": "localhost:8081", + } + self.write_config(minimal_config) + config = MooncakeConfig.from_file(self.config_file) + + self.assertEqual(config.tenant_id, "default") + + def test_tenant_id_null_defaults(self): + """Test tenant_id defaults to default when explicitly null""" + self.write_config({**self.valid_config, "tenant_id": None}) + config = MooncakeConfig.from_file(self.config_file) + + self.assertEqual(config.tenant_id, "default") + + def test_ssd_offload_path_null_defaults(self): + """Test ssd_offload_path defaults to empty when explicitly null""" + self.write_config({**self.valid_config, "ssd_offload_path": None}) + config = MooncakeConfig.from_file(self.config_file) + + self.assertEqual(config.ssd_offload_path, "") + + def test_enable_ssd_offload_string_values(self): + """from_file must parse string booleans like load_from_env, and reject typos. + + from_file used bool(value), so any non-empty string (including "false") + turned SSD offload on, disagreeing with load_from_env which parses the + string. _parse_bool now understands the common textual forms and raises + on anything unrecognized instead of silently defaulting to False. + """ + required = { + "local_hostname": "localhost", + "metadata_server": "localhost:8080", + "master_server_address": "localhost:8081", + } + cases = [ + ("false", False), + ("False", False), + ("0", False), + ("no", False), + ("off", False), + ("true", True), + ("1", True), + ("yes", True), + ("on", True), + (True, True), + (False, False), + ] + for raw, expected in cases: + with self.subTest(raw=raw): + self.write_config({**required, "enable_ssd_offload": raw}) + config = MooncakeConfig.from_file(self.config_file) + self.assertEqual(config.enable_ssd_offload, expected) + + # An unrecognized string is a config error, not a silent disable. + self.write_config({**required, "enable_ssd_offload": "notabool"}) + with self.assertRaises(ValueError): + MooncakeConfig.from_file(self.config_file) + + def test_client_http_config_from_file(self): + """Test loading client HTTP metrics settings from configuration file""" + self.write_config( + { + **self.valid_config, + "enable_client_http_server": "enable", + "client_http_port": "19444", + } + ) + config = MooncakeConfig.from_file(self.config_file) + + self.assertEqual(config.enable_client_http_server, True) + self.assertEqual(config.client_http_port, 19444) def test_missing_required_field(self): """Test missing required field""" @@ -80,61 +174,145 @@ def test_missing_required_field(self): with self.assertRaises(ValueError) as cm: MooncakeConfig.from_file(self.config_file) - self.assertIn(f"Missing required config field: {field}", str(cm.exception)) + self.assertIn( + f"Missing required config field: {field}", str(cm.exception) + ) def test_load_from_config_path_env(self): """Test loading configuration from environment variable MOONCAKE_CONFIG_PATH""" self.write_config(self.valid_config) # Set environment variable - os.environ['MOONCAKE_CONFIG_PATH'] = self.config_file + os.environ["MOONCAKE_CONFIG_PATH"] = self.config_file try: config = MooncakeConfig.load_from_env() self.assertEqual(config.local_hostname, "localhost") finally: # Clean up environment variable - del os.environ['MOONCAKE_CONFIG_PATH'] + del os.environ["MOONCAKE_CONFIG_PATH"] def test_load_from_config_env(self): """Test loading configuration from environment variable MOONCAKE_MASTER""" # Set environment variable - os.environ['MOONCAKE_MASTER'] = self.valid_config["master_server_address"] - os.environ['LOCAL_HOSTNAME'] = self.valid_config["local_hostname"] - os.environ['MOONCAKE_TE_META_DATA_SERVER'] = self.valid_config["metadata_server"] - os.environ['MOONCAKE_GLOBAL_SEGMENT_SIZE'] = str(self.valid_config["global_segment_size"]) - os.environ['MOONCAKE_PROTOCOL'] = self.valid_config["protocol"] - os.environ['MOONCAKE_DEVICE'] = self.valid_config["device_name"] - os.environ['MOONCAKE_OFFLOAD_ENABLED'] = str(self.valid_config["enable_ssd_offload"]) - os.environ['MOONCAKE_OFFLOAD_FILE_STORAGE_PATH'] = self.valid_config["ssd_offload_path"] + os.environ["MOONCAKE_MASTER"] = self.valid_config["master_server_address"] + os.environ["LOCAL_HOSTNAME"] = self.valid_config["local_hostname"] + os.environ["MOONCAKE_TE_META_DATA_SERVER"] = self.valid_config[ + "metadata_server" + ] + os.environ["MOONCAKE_GLOBAL_SEGMENT_SIZE"] = str( + self.valid_config["global_segment_size"] + ) + os.environ["MOONCAKE_PROTOCOL"] = self.valid_config["protocol"] + os.environ["MOONCAKE_DEVICE"] = self.valid_config["device_name"] + os.environ["MOONCAKE_OFFLOAD_ENABLED"] = str( + self.valid_config["enable_ssd_offload"] + ) + os.environ["MOONCAKE_OFFLOAD_FILE_STORAGE_PATH"] = self.valid_config[ + "ssd_offload_path" + ] + os.environ["MOONCAKE_TENANT_ID"] = self.valid_config["tenant_id"] + os.environ["MOONCAKE_ENABLE_CLIENT_HTTP_SERVER"] = str( + self.valid_config["enable_client_http_server"] + ) + os.environ["MOONCAKE_CLIENT_HTTP_PORT"] = str( + self.valid_config["client_http_port"] + ) try: config = MooncakeConfig.load_from_env() - self.assertEqual(config.master_server_address, self.valid_config["master_server_address"]) - self.assertEqual(config.metadata_server, self.valid_config["metadata_server"]) + self.assertEqual( + config.master_server_address, self.valid_config["master_server_address"] + ) + self.assertEqual( + config.metadata_server, self.valid_config["metadata_server"] + ) self.assertEqual(config.local_hostname, self.valid_config["local_hostname"]) - self.assertEqual(config.global_segment_size, self.valid_config["global_segment_size"]) + self.assertEqual( + config.global_segment_size, self.valid_config["global_segment_size"] + ) self.assertEqual(config.protocol, self.valid_config["protocol"]) self.assertEqual(config.device_name, self.valid_config["device_name"]) - self.assertEqual(config.enable_ssd_offload, self.valid_config["enable_ssd_offload"]) - self.assertEqual(config.ssd_offload_path, self.valid_config["ssd_offload_path"]) + self.assertEqual( + config.enable_ssd_offload, self.valid_config["enable_ssd_offload"] + ) + self.assertEqual( + config.ssd_offload_path, self.valid_config["ssd_offload_path"] + ) + self.assertEqual(config.tenant_id, self.valid_config["tenant_id"]) + self.assertEqual( + config.enable_client_http_server, + self.valid_config["enable_client_http_server"], + ) + self.assertEqual( + config.client_http_port, self.valid_config["client_http_port"] + ) finally: # Clean up environment variable - del os.environ['MOONCAKE_MASTER'] - del os.environ['LOCAL_HOSTNAME'] - del os.environ['MOONCAKE_TE_META_DATA_SERVER'] - del os.environ['MOONCAKE_GLOBAL_SEGMENT_SIZE'] - del os.environ['MOONCAKE_PROTOCOL'] - del os.environ['MOONCAKE_DEVICE'] - del os.environ['MOONCAKE_OFFLOAD_ENABLED'] - del os.environ['MOONCAKE_OFFLOAD_FILE_STORAGE_PATH'] + del os.environ["MOONCAKE_MASTER"] + del os.environ["LOCAL_HOSTNAME"] + del os.environ["MOONCAKE_TE_META_DATA_SERVER"] + del os.environ["MOONCAKE_GLOBAL_SEGMENT_SIZE"] + del os.environ["MOONCAKE_PROTOCOL"] + del os.environ["MOONCAKE_DEVICE"] + del os.environ["MOONCAKE_OFFLOAD_ENABLED"] + del os.environ["MOONCAKE_OFFLOAD_FILE_STORAGE_PATH"] + del os.environ["MOONCAKE_TENANT_ID"] + del os.environ["MOONCAKE_ENABLE_CLIENT_HTTP_SERVER"] + del os.environ["MOONCAKE_CLIENT_HTTP_PORT"] + + def test_tenant_id_from_env(self): + """Test loading tenant_id from MOONCAKE_TENANT_ID""" + previous_config_path = os.environ.pop("MOONCAKE_CONFIG_PATH", None) + previous_master = os.environ.pop("MOONCAKE_MASTER", None) + previous_tenant_id = os.environ.pop("MOONCAKE_TENANT_ID", None) + + os.environ["MOONCAKE_MASTER"] = self.valid_config["master_server_address"] + os.environ["MOONCAKE_TENANT_ID"] = "tenant-from-env" + + try: + config = MooncakeConfig.load_from_env() + self.assertEqual(config.tenant_id, "tenant-from-env") + finally: + os.environ.pop("MOONCAKE_MASTER", None) + os.environ.pop("MOONCAKE_TENANT_ID", None) + if previous_config_path is not None: + os.environ["MOONCAKE_CONFIG_PATH"] = previous_config_path + if previous_master is not None: + os.environ["MOONCAKE_MASTER"] = previous_master + if previous_tenant_id is not None: + os.environ["MOONCAKE_TENANT_ID"] = previous_tenant_id + + def test_tenant_id_env_defaults(self): + """Test tenant_id defaults to default when MOONCAKE_TENANT_ID is omitted""" + previous_config_path = os.environ.pop("MOONCAKE_CONFIG_PATH", None) + previous_master = os.environ.pop("MOONCAKE_MASTER", None) + previous_tenant_id = os.environ.pop("MOONCAKE_TENANT_ID", None) + + os.environ["MOONCAKE_MASTER"] = self.valid_config["master_server_address"] + + try: + config = MooncakeConfig.load_from_env() + self.assertEqual(config.tenant_id, "default") + finally: + os.environ.pop("MOONCAKE_MASTER", None) + os.environ.pop("MOONCAKE_TENANT_ID", None) + if previous_config_path is not None: + os.environ["MOONCAKE_CONFIG_PATH"] = previous_config_path + if previous_master is not None: + os.environ["MOONCAKE_MASTER"] = previous_master + if previous_tenant_id is not None: + os.environ["MOONCAKE_TENANT_ID"] = previous_tenant_id def test_load_from_env_missing(self): """Test loading configuration from environment variable when not set""" with self.assertRaises(ValueError) as cm: MooncakeConfig.load_from_env() - self.assertIn("Neither the environment variable 'MOONCAKE_CONFIG_PATH' nor 'MOONCAKE_MASTER' is set.", str(cm.exception)) + self.assertIn( + "Neither the environment variable 'MOONCAKE_CONFIG_PATH' nor 'MOONCAKE_MASTER' is set.", + str(cm.exception), + ) class TestParseSegmentSize(unittest.TestCase): @@ -156,20 +334,20 @@ def test_kb_suffix(self): self.assertEqual(_parse_segment_size("1.5kb"), int(1.5 * 1024)) def test_mb_suffix(self): - self.assertEqual(_parse_segment_size("1mb"), 1024 ** 2) - self.assertEqual(_parse_segment_size("512MB"), 512 * 1024 ** 2) - self.assertEqual(_parse_segment_size("1m"), 1024 ** 2) + self.assertEqual(_parse_segment_size("1mb"), 1024**2) + self.assertEqual(_parse_segment_size("512MB"), 512 * 1024**2) + self.assertEqual(_parse_segment_size("1m"), 1024**2) def test_gb_suffix(self): - self.assertEqual(_parse_segment_size("1gb"), 1024 ** 3) - self.assertEqual(_parse_segment_size("3GB"), 3 * 1024 ** 3) - self.assertEqual(_parse_segment_size("1g"), 1024 ** 3) - self.assertEqual(_parse_segment_size("1.5gb"), int(1.5 * 1024 ** 3)) + self.assertEqual(_parse_segment_size("1gb"), 1024**3) + self.assertEqual(_parse_segment_size("3GB"), 3 * 1024**3) + self.assertEqual(_parse_segment_size("1g"), 1024**3) + self.assertEqual(_parse_segment_size("1.5gb"), int(1.5 * 1024**3)) def test_tb_suffix(self): - self.assertEqual(_parse_segment_size("1tb"), 1024 ** 4) - self.assertEqual(_parse_segment_size("1TB"), 1024 ** 4) - self.assertEqual(_parse_segment_size("1t"), 1024 ** 4) + self.assertEqual(_parse_segment_size("1tb"), 1024**4) + self.assertEqual(_parse_segment_size("1TB"), 1024**4) + self.assertEqual(_parse_segment_size("1t"), 1024**4) def test_b_suffix(self): self.assertEqual(_parse_segment_size("4096b"), 4096) @@ -196,8 +374,124 @@ def test_invalid_string_raises(self): _parse_segment_size("abc") def test_whitespace_handling(self): - self.assertEqual(_parse_segment_size(" 3 gb "), 3 * 1024 ** 3) + self.assertEqual(_parse_segment_size(" 3 gb "), 3 * 1024**3) + + +class TestMooncakeConfigValidation(unittest.TestCase): + """Tests for the field validation performed in MooncakeConfig.__post_init__.""" + + BASE_KWARGS = dict( + local_hostname="localhost", + metadata_server="localhost:8080", + global_segment_size=DEFAULT_GLOBAL_SEGMENT_SIZE, + local_buffer_size=DEFAULT_LOCAL_BUFFER_SIZE, + protocol="tcp", + device_name="", + master_server_address="localhost:8081", + ) + + def make(self, **overrides): + kwargs = dict(self.BASE_KWARGS) + kwargs.update(overrides) + return MooncakeConfig(**kwargs) + + def test_known_protocols_normalized_to_lowercase(self): + # Mixed-case input is canonicalised to lowercase so it matches the + # case-sensitive C++ engine; surrounding whitespace is trimmed. Known + # protocols (including Store modes the old hard allowlist rejected) do + # not warn. + cases = { + "tcp": "tcp", + "rdma": "rdma", + "efa": "efa", + "RDMA": "rdma", + "Tcp": "tcp", + " rdma ": "rdma", + "cxl": "cxl", + "ascend": "ascend", + "nvlink_intra": "nvlink_intra", + "ub": "ub", + "ubshmem": "ubshmem", + "maca": "maca", + "sunrise_link": "sunrise_link", + "rpc_only": "rpc_only", + } + for given, expected in cases.items(): + with self.subTest(protocol=given): + config = self.make(protocol=given) + self.assertEqual(config.protocol, expected) + + def test_empty_protocol_raises(self): + # An empty or whitespace-only protocol is a caller error and still + # fails fast. + for protocol in ["", " "]: + with self.subTest(protocol=protocol): + with self.assertRaises(ValueError) as cm: + self.make(protocol=protocol) + self.assertIn("Invalid protocol", str(cm.exception)) + + def test_unknown_protocol_warns_but_is_passed_through(self): + # Unknown-but-non-empty values are no longer rejected in Python: + # MooncakeConfig drives both Transfer Engine and Store paths and the C++ + # layer is the source of truth. We warn and pass the lowercased value + # through. + for given, expected in [("rmda", "rmda"), ("udp", "udp"), ("Foo", "foo")]: + with self.subTest(protocol=given): + with self.assertLogs(_cfg_mod.logger, level="WARNING") as cm: + config = self.make(protocol=given) + self.assertEqual(config.protocol, expected) + self.assertTrue(any("Unrecognised protocol" in m for m in cm.output)) + + def test_non_string_protocol_raises(self): + with self.assertRaises(ValueError): + self.make(protocol=None) + + def test_zero_sizes_allowed(self): + # 0 is a documented sentinel (e.g. global_segment_size == 0 disables the + # store), so it must remain valid. + config = self.make(global_segment_size=0, local_buffer_size=0) + self.assertEqual(config.global_segment_size, 0) + self.assertEqual(config.local_buffer_size, 0) + + def test_negative_sizes_raise(self): + with self.assertRaises(ValueError) as cm: + self.make(global_segment_size=-1) + self.assertIn("global_segment_size", str(cm.exception)) + with self.assertRaises(ValueError) as cm: + self.make(local_buffer_size=-1024) + self.assertIn("local_buffer_size", str(cm.exception)) + + def test_empty_required_field_raises(self): + for field in ["local_hostname", "metadata_server", "master_server_address"]: + for bad in ["", " "]: + with self.subTest(field=field, value=bad): + with self.assertRaises(ValueError) as cm: + self.make(**{field: bad}) + self.assertIn(field, str(cm.exception)) + + def test_from_file_warns_on_unknown_protocol(self): + with open(self.config_path, "w") as f: + json.dump( + { + "local_hostname": "localhost", + "metadata_server": "localhost:8080", + "master_server_address": "localhost:8081", + "protocol": "rmda", # typo -> unknown, warned not rejected + }, + f, + ) + with self.assertLogs(_cfg_mod.logger, level="WARNING") as cm: + config = MooncakeConfig.from_file(self.config_path) + self.assertEqual(config.protocol, "rmda") + self.assertTrue(any("Unrecognised protocol" in m for m in cm.output)) + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.config_path = os.path.join(self._tmp.name, "config.json") + + def tearDown(self): + self._tmp.cleanup() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/mooncake-wheel/tests/test_mooncake_ep.py b/mooncake-wheel/tests/test_mooncake_ep.py index fa7636fa4e..0dd9630faf 100644 --- a/mooncake-wheel/tests/test_mooncake_ep.py +++ b/mooncake-wheel/tests/test_mooncake_ep.py @@ -1,4 +1,5 @@ import random +import os import torch import torch.distributed as dist from functools import partial @@ -6,6 +7,11 @@ from mooncake.mooncake_ep_buffer import Buffer from ep_test_utils import init_dist, bench, bench_kineto, calc_diff, hash_tensor, per_token_cast_back +_USE_MACA = ( + os.getenv("MOONCAKE_EP_USE_MACA", "").upper() in {"1", "ON", "TRUE", "YES"} + or bool(getattr(torch.version, "maca", None)) +) + def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int, rank: int, num_ranks: int, group: dist.ProcessGroup, cpu_group: dist.ProcessGroup, buffer: Buffer, seed: int = 0): @@ -34,7 +40,7 @@ def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int, hash_value, num_times = 0, 0 active_ranks = torch.ones((num_tokens, ), dtype=torch.int32, device='cuda') for return_recv_hook in (False, True): - for dispatch_use_fp8 in (False, True): + for dispatch_use_fp8 in ([False] if _USE_MACA else [False, True]): num_times += 1 for i in range((num_times % 2) + 1): packed_recv_x, packed_recv_count, handle, event, hook = \ @@ -88,16 +94,6 @@ def test_main(num_tokens: int, hidden: int, num_experts: int, num_topk: int, assert diff < 1e-5, f'Error: {diff=}, {zero_copy=}' hash_value ^= hash_tensor(combined_x) - def create_test_cast_with_outliers(num_outliers): - tmp = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') - tmp /= tmp.abs().amax(dim=1).view(-1, 1) - assert tmp.abs().amax().item() <= 1 - - # Create some amax outliers - for i in range(num_outliers): - tmp[random.randint(0, num_tokens - 1)] *= 1e3 - return tmp - # noinspection PyShadowingNames def large_gemm_with_hook(hook): mat_0 = torch.randn((8192, 8192), dtype=torch.float) diff --git a/mooncake-wheel/tests/test_mooncake_store_service_api.py b/mooncake-wheel/tests/test_mooncake_store_service_api.py index 58a2940d9b..d5b118c265 100644 --- a/mooncake-wheel/tests/test_mooncake_store_service_api.py +++ b/mooncake-wheel/tests/test_mooncake_store_service_api.py @@ -1,11 +1,17 @@ #!/usr/bin/env python3 import asyncio import json +import logging +import signal import sys +import tempfile +import time import types import unittest from pathlib import Path from types import SimpleNamespace +from unittest import mock +from unittest.mock import patch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -16,9 +22,10 @@ web_module = types.ModuleType("aiohttp.web") class Response: - def __init__(self, status=200, text="", content_type=None): + def __init__(self, status=200, text="", body=None, content_type=None): self.status = status self.text = text + self.body = body self.content_type = content_type web_module.Response = Response @@ -37,7 +44,12 @@ class MooncakeDistributedStore: store_module.MooncakeDistributedStore = MooncakeDistributedStore sys.modules["mooncake.store"] = store_module -from mooncake.mooncake_store_service import MooncakeStoreService, _shm_name_to_path +from mooncake.mooncake_store_service import ( + MooncakeStoreService, + _install_shutdown_signal_handlers, + _shm_name_to_path, + main as store_service_main, +) class FakeStore: @@ -49,6 +61,11 @@ def __init__(self): self.unmount_failures = set() self.allocated_mount_calls = [] self.free_unmount_calls = [] + self.setup_calls = [] + + def setup(self, *args): + self.setup_calls.append(args) + return 0 def mount_segment(self, path, size, offset, protocol, location): self.mount_calls.append((path, size, offset, protocol, location)) @@ -106,6 +123,69 @@ async def asyncSetUp(self): self.service.last_mount_info = {} self.service._state_lock = asyncio.Lock() + async def test_start_store_service_passes_tenant_id_to_setup(self): + fake_store = FakeStore() + self.service.config = SimpleNamespace( + local_hostname="localhost", + metadata_server="P2PHANDSHAKE", + global_segment_size=1024, + local_buffer_size=2048, + protocol="tcp", + device_name="", + master_server_address="127.0.0.1:50051", + enable_ssd_offload=False, + ssd_offload_path="", + tenant_id="tenant-a", + enable_client_http_server=False, + client_http_port=9300, + ) + + with patch( + "mooncake.mooncake_store_service.MooncakeDistributedStore", + return_value=fake_store, + ): + result = await self.service.start_store_service(max_wait_time=1) + + self.assertTrue(result) + self.assertEqual( + fake_store.setup_calls, + [ + ( + { + "local_hostname": "localhost", + "metadata_server": "P2PHANDSHAKE", + "global_segment_size": 1024, + "local_buffer_size": 2048, + "protocol": "tcp", + "rdma_devices": "", + "master_server_addr": "127.0.0.1:50051", + "enable_ssd_offload": False, + "ssd_offload_path": "", + "tenant_id": "tenant-a", + "enable_client_http_server": False, + "client_http_port": 9300, + }, + ) + ], + ) + + async def test_cli_config_can_override_tenant_id(self): + config = { + "local_hostname": "localhost", + "metadata_server": "P2PHANDSHAKE", + "master_server_address": "127.0.0.1:50051", + "tenant_id": "tenant-from-file", + } + + with tempfile.TemporaryDirectory() as tmpdir: + config_path = Path(tmpdir) / "config.json" + config_path.write_text(json.dumps(config)) + service = MooncakeStoreService( + str(config_path), {"tenant_id": "tenant-from-cli"} + ) + + self.assertEqual(service.config.tenant_id, "tenant-from-cli") + async def test_mount_shm_then_unmount_shm_api(self): mount_resp = await self.service.handle_mount_shm( FakeRequest( @@ -178,6 +258,30 @@ async def test_unmount_shm_updates_state_for_partial_success(self): self.assertEqual(self.service.current_mode, "decode") async def test_reconfigure_decode_mount_failure_rolls_back_to_prefill(self): + # Fresh prefill -> decode mount that fails: there are no previously + # serving segments to preserve, so the node still rolls back to prefill. + self.service.current_mode = "prefill" + self.service.mounted_segment_ids = [] + self.service.last_mount_info = {} + self.fake_store.fail_mount = True + + resp = await self.service.handle_reconfigure( + FakeRequest({"mode": "decode", "path": "/dev/shm/new", "size": 4096}) + ) + + self.assertEqual(resp.status, 500) + body = json.loads(resp.text) + self.assertEqual(body["mode"], "prefill") + self.assertIn("rolled back to prefill", body["error"]) + self.assertEqual(self.fake_store.unmount_calls, []) + self.assertEqual(self.service.mounted_segment_ids, []) + self.assertEqual(self.service.current_mode, "prefill") + self.assertEqual(self.service.last_mount_info, {}) + + async def test_reconfigure_decode_remount_failure_keeps_previous_segments(self): + # A remount to a DIFFERENT path that fails to mount must not destroy the + # still-healthy previous segments: the node keeps serving from them and + # stays in decode mode (make-before-break). old_id = "00000000-0000-0000-0000-000000000001" self.service.current_mode = "decode" self.service.mounted_segment_ids = [old_id] @@ -196,12 +300,132 @@ async def test_reconfigure_decode_mount_failure_rolls_back_to_prefill(self): self.assertEqual(resp.status, 500) body = json.loads(resp.text) - self.assertEqual(body["mode"], "prefill") - self.assertIn("rolled back to prefill", body["error"]) + self.assertEqual(body["mode"], "decode") + self.assertIn("keeping previous decode segments", body["error"]) + # Nothing was unmounted, capacity preserved, mode unchanged. + self.assertEqual(self.fake_store.unmount_calls, []) + self.assertEqual(self.service.mounted_segment_ids, [old_id]) + self.assertEqual(self.service.current_mode, "decode") + # last_mount_info still points at the previous (working) path so a + # subsequent same-path detection keeps working. + self.assertEqual(self.service.last_mount_info["path"], "/dev/shm/old") + + async def test_reconfigure_decode_same_path_remount_failure_keeps_previous_segments(self): + # A remount to the SAME path that fails to mount must not destroy the + # still-healthy previous segments: the node keeps serving from them and + # stays in decode mode (make-before-break). Same-path MBB is safe here + # because /api/reconfigure binds through MasterClient::MountSegment, + # which mints a fresh UUID per mount and does not enter the NoF + # te_endpoint-dedup path, so old and new cannot collide on the same path. + old_id = "00000000-0000-0000-0000-000000000001" + self.service.current_mode = "decode" + self.service.mounted_segment_ids = [old_id] + self.service.last_mount_info = { + "path": "/dev/shm/same", + "offset": 0, + "size": 4096, + "protocol": "tcp", + "location": "", + } + self.fake_store.fail_mount = True + + resp = await self.service.handle_reconfigure( + FakeRequest({"mode": "decode", "path": "/dev/shm/same", "size": 4096}) + ) + + self.assertEqual(resp.status, 500) + body = json.loads(resp.text) + self.assertEqual(body["mode"], "decode") + self.assertIn("keeping previous decode segments", body["error"]) + # Nothing was unmounted, capacity preserved, mode unchanged. + self.assertEqual(self.fake_store.unmount_calls, []) + self.assertEqual(self.service.mounted_segment_ids, [old_id]) + self.assertEqual(self.service.current_mode, "decode") + # last_mount_info still points at the previous (working) same path so a + # subsequent remount keeps working. + self.assertEqual(self.service.last_mount_info["path"], "/dev/shm/same") + + async def test_reconfigure_decode_remount_success_is_make_before_break(self): + # A successful remount to a DIFFERENT path must mount the new segment + # BEFORE retiring the old one (make-before-break), then leave only the + # new segment serving. Distinct ids let old and new be told apart. + old_id = "00000000-0000-0000-0000-000000000001" + new_id = "00000000-0000-0000-0000-000000000002" + self.service.current_mode = "decode" + self.service.mounted_segment_ids = [old_id] + self.service.last_mount_info = { + "path": "/dev/shm/old", + "offset": 0, + "size": 4096, + "protocol": "tcp", + "location": "", + } + + order = {"old_still_mounted_at_new_mount": None} + + def mount_new(path, size, offset, protocol, location): + # The old segment must still be serving when the new one is mounted. + order["old_still_mounted_at_new_mount"] = ( + old_id in self.service.mounted_segment_ids + ) + return {"ret": 0, "segment_ids": [new_id]} + + self.fake_store.mount_segment = mount_new + + resp = await self.service.handle_reconfigure( + FakeRequest({"mode": "decode", "path": "/dev/shm/new", "size": 8192}) + ) + + self.assertEqual(resp.status, 200) + body = json.loads(resp.text) + self.assertEqual(body["mode"], "decode") + # Make-before-break: the old segment was still mounted when the new one + # was created, and it is retired only after the new mount succeeds. + self.assertTrue(order["old_still_mounted_at_new_mount"]) self.assertEqual(self.fake_store.unmount_calls, [([old_id], 0)]) - self.assertEqual(self.service.mounted_segment_ids, []) - self.assertEqual(self.service.current_mode, "prefill") - self.assertEqual(self.service.last_mount_info, {}) + # Only the new segment is left serving; the old one is gone. + self.assertEqual(self.service.mounted_segment_ids, [new_id]) + self.assertEqual(self.service.current_mode, "decode") + self.assertEqual(self.service.last_mount_info["path"], "/dev/shm/new") + + async def test_reconfigure_decode_partial_unmount_failure_keeps_only_failed_ids(self): + # New mount succeeds, but retiring the previous segments only PARTIALLY + # fails. unmount_segment reports the first error for a batch, so the old + # ids must be unmounted individually; mounted_segment_ids must then hold + # the new id plus ONLY the id whose cleanup actually failed -- not the + # whole previous set (which would retain a stale, already-freed id) and + # not none of it (which would silently leak the still-live old segment). + old_ok = "00000000-0000-0000-0000-0000000000a1" + old_fail = "00000000-0000-0000-0000-0000000000a2" + new_id = "00000000-0000-0000-0000-0000000000b1" + self.service.current_mode = "decode" + self.service.mounted_segment_ids = [old_ok, old_fail] + self.service.last_mount_info = { + "path": "/dev/shm/old", + "offset": 0, + "size": 4096, + "protocol": "tcp", + "location": "", + } + self.fake_store.unmount_failures = {old_fail} + + def mount_new(path, size, offset, protocol, location): + return {"ret": 0, "segment_ids": [new_id]} + + self.fake_store.mount_segment = mount_new + + resp = await self.service.handle_reconfigure( + FakeRequest({"mode": "decode", "path": "/dev/shm/new", "size": 8192}) + ) + + self.assertEqual(resp.status, 200) + # Each previous id was unmounted on its own, not as a single batch. + self.assertEqual( + self.fake_store.unmount_calls, [([old_ok], 0), ([old_fail], 0)] + ) + # New id serves; the freed id is dropped, the un-freed id stays tracked. + self.assertEqual(self.service.mounted_segment_ids, [new_id, old_fail]) + self.assertEqual(self.service.current_mode, "decode") async def test_mount_allocates_and_frees_on_unmount(self): mount_resp = await self.service.handle_mount( @@ -306,6 +530,7 @@ async def test_handle_put_empty_key(self): # ==================== /api/get/{key} tests ==================== async def test_handle_get_success(self): + self.fake_store.is_exist = lambda key: True self.fake_store.get = lambda key: b"payload_bytes" request = FakeRequest({}) request.match_info = {"key": "my_key"} @@ -314,7 +539,8 @@ async def test_handle_get_success(self): self.assertEqual(resp.body, b"payload_bytes") async def test_handle_get_not_found(self): - self.fake_store.get = lambda key: None + self.fake_store.is_exist = lambda key: 0 + self.fake_store.get = lambda key: b"" request = FakeRequest({}) request.match_info = {"key": "missing_key"} resp = await self.service.handle_get(request) @@ -322,7 +548,28 @@ async def test_handle_get_not_found(self): body = json.loads(resp.text) self.assertIn("Key not found", body["error"]) + async def test_handle_get_exist_check_failure(self): + self.fake_store.is_exist = lambda key: -1 + self.fake_store.get = lambda key: b"" + request = FakeRequest({}) + request.match_info = {"key": "error_key"} + resp = await self.service.handle_get(request) + self.assertEqual(resp.status, 500) + body = json.loads(resp.text) + self.assertIn("Exist check failed", body["error"]) + async def test_handle_get_empty_bytes(self): + self.fake_store.is_exist = lambda key: True + self.fake_store.get = lambda key: b"" + request = FakeRequest({}) + request.match_info = {"key": "empty_value_key"} + resp = await self.service.handle_get(request) + self.assertEqual(resp.status, 200) + self.assertEqual(resp.body, b"") + + async def test_handle_get_empty_bytes_rechecks_existence(self): + existence_results = iter([1, 0]) + self.fake_store.is_exist = lambda key: next(existence_results) self.fake_store.get = lambda key: b"" request = FakeRequest({}) request.match_info = {"key": "empty_value_key"} @@ -331,10 +578,21 @@ async def test_handle_get_empty_bytes(self): body = json.loads(resp.text) self.assertIn("Key not found", body["error"]) + async def test_handle_get_none_value_is_store_failure(self): + self.fake_store.is_exist = lambda key: True + self.fake_store.get = lambda key: None + request = FakeRequest({}) + request.match_info = {"key": "empty_value_key"} + resp = await self.service.handle_get(request) + self.assertEqual(resp.status, 500) + body = json.loads(resp.text) + self.assertIn("GET operation failed", body["error"]) + async def test_handle_get_store_exception(self): def raise_error(key): raise RuntimeError("store crashed") + self.fake_store.is_exist = lambda key: True self.fake_store.get = raise_error request = FakeRequest({}) request.match_info = {"key": "crash_key"} @@ -375,6 +633,19 @@ def raise_error(key): body = json.loads(resp.text) self.assertIn("exist check crashed", body["error"]) + async def test_handle_exist_store_error(self): + # is_exist returns -1 when the store is unhealthy; this should surface + # as HTTP 500, not as HTTP 200 {"exists": true} (bool(-1) == True). + self.fake_store.is_exist = lambda key: -1 + request = FakeRequest({}) + request.match_info = {"key": "some_key"} + resp = await self.service.handle_exist(request) + self.assertEqual(resp.status, 500) + body = json.loads(resp.text) + # Assert the specific message so this pins the exists < 0 branch rather + # than any 500 (the except path returns {"error": str(e)} too). + self.assertEqual(body["error"], "Exist check failed") + # ==================== /api/remove/{key} tests ==================== async def test_handle_remove_success(self): @@ -637,6 +908,157 @@ async def test_unmount_shm_empty_list(self): self.assertEqual(resp.status, 400) +class StoreServiceShutdownTest(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.service = MooncakeStoreService.__new__(MooncakeStoreService) + self.service.store = None + self.service.config = SimpleNamespace( + local_hostname="localhost", + metadata_server="P2PHANDSHAKE", + global_segment_size=1, + local_buffer_size=0, + protocol="tcp", + device_name="", + master_server_address="localhost:50051", + enable_ssd_offload=False, + ssd_offload_path="", + tenant_id="", + ) + + async def test_shutdown_event_stops_startup_retry_sleep(self): + class FailingStore: + def setup(self, *_args): + raise RuntimeError("setup failed") + + shutdown_event = asyncio.Event() + + async def trigger_shutdown(): + await asyncio.sleep(0) + shutdown_event.set() + + with mock.patch( + "mooncake.mooncake_store_service.MooncakeDistributedStore", + FailingStore, + ): + shutdown_task = asyncio.create_task(trigger_shutdown()) + start_time = time.perf_counter() + with self.assertLogs(level=logging.WARNING): + result = await self.service.start_store_service( + max_wait_time=2, shutdown_event=shutdown_event + ) + elapsed = time.perf_counter() - start_time + await shutdown_task + + self.assertFalse(result) + self.assertLess(elapsed, 0.5) + + async def test_shutdown_during_setup_is_observed_before_success(self): + shutdown_event = asyncio.Event() + loop = asyncio.get_running_loop() + + class SchedulingStore(FakeStore): + def __init__(self): + super().__init__() + self.close_calls = 0 + + def setup(self, *args): + ret = super().setup(*args) + loop.call_soon(shutdown_event.set) + return ret + + def close(self): + self.close_calls += 1 + return 0 + + store = SchedulingStore() + with mock.patch( + "mooncake.mooncake_store_service.MooncakeDistributedStore", + return_value=store, + ): + result = await self.service.start_store_service( + max_wait_time=1, shutdown_event=shutdown_event + ) + + self.assertFalse(result) + self.assertEqual(store.close_calls, 1) + self.assertIsNone(self.service.store) + + async def test_stop_logs_nonzero_close_return(self): + store = mock.Mock() + store.close.return_value = 7 + self.service.store = store + + with self.assertLogs(level=logging.WARNING) as logs: + await self.service.stop() + + self.assertIsNone(self.service.store) + self.assertTrue( + any("close returned 7" in message for message in logs.output) + ) + await self.service.stop() + store.close.assert_called_once_with() + + async def test_signal_handler_requests_shutdown(self): + loop = mock.Mock() + shutdown_event = asyncio.Event() + + _install_shutdown_signal_handlers(loop, shutdown_event) + + sigterm_call = next( + call + for call in loop.add_signal_handler.call_args_list + if call.args[0] == signal.SIGTERM + ) + sigterm_call.args[1](sigterm_call.args[2]) + self.assertTrue(shutdown_event.is_set()) + + async def test_main_closes_store_when_shutdown_requested_during_startup(self): + args = SimpleNamespace( + config=None, + define=[], + max_wait_time=60, + port=8080, + ) + service = mock.Mock() + service.start_store_service = mock.AsyncMock(return_value=False) + service.start_http_service = mock.AsyncMock(return_value=True) + service.stop = mock.AsyncMock() + + startup_calls = [] + + def request_shutdown(_loop, shutdown_event): + startup_calls.append("install") + shutdown_event.set() + + def unblock_shutdown_signals(): + startup_calls.append("unblock") + + with ( + mock.patch( + "mooncake.mooncake_store_service.parse_arguments", + return_value=args, + ), + mock.patch( + "mooncake.mooncake_store_service.MooncakeStoreService", + return_value=service, + ), + mock.patch( + "mooncake.mooncake_store_service._unblock_shutdown_signals", + side_effect=unblock_shutdown_signals, + ), + mock.patch( + "mooncake.mooncake_store_service._install_shutdown_signal_handlers", + side_effect=request_shutdown, + ), + ): + await store_service_main() + + service.start_store_service.assert_awaited_once() + service.start_http_service.assert_not_awaited() + service.stop.assert_awaited_once() + self.assertEqual(startup_calls, ["install", "unblock"]) + + class ShmNameToPathTest(unittest.TestCase): def test_valid_simple_name(self): self.assertEqual(_shm_name_to_path("my-segment"), "/dev/shm/my-segment") diff --git a/mooncake-wheel/tests/test_put_get_tensor.py b/mooncake-wheel/tests/test_put_get_tensor.py index 210036ae6b..edf8d8c8c9 100644 --- a/mooncake-wheel/tests/test_put_get_tensor.py +++ b/mooncake-wheel/tests/test_put_get_tensor.py @@ -3,8 +3,12 @@ import unittest import os import time -import threading -import random + +try: + import torch as _torch +except Exception: + _torch = None + from mooncake.store import MooncakeDistributedStore # The lease time of the kv object, should be set equal to @@ -13,6 +17,9 @@ # Use environment variable if set, otherwise use default default_kv_lease_ttl = int(os.getenv("DEFAULT_KV_LEASE_TTL", DEFAULT_DEFAULT_KV_LEASE_TTL)) +def cuda_available(): + return _torch is not None and _torch.cuda.is_available() + # Define a test class for serialization class TestClass: def __init__(self, version=1, shape=(1, 2, 3)): @@ -119,6 +126,32 @@ def test_put_get_tensor(self): self.store.remove(key_bool) self.store.remove(key_rand) + @unittest.skipUnless(cuda_available(), "CUDA is not available") + def test_cuda_local_copy_paths(self): + """Test CUDA source writes and CUDA destination reads.""" + import torch + + prefix = f"test_cuda_local_copy_{os.getpid()}" + put_key = f"{prefix}_put" + upsert_key = f"{prefix}_upsert" + raw_key = f"{prefix}_raw" + + tensor = torch.arange(16, dtype=torch.float32, device="cuda") + self.assertEqual(self.store.put_tensor(put_key, tensor), 0) + self.assertEqual(self.store.upsert_tensor(upsert_key, tensor), 0) + + raw = bytes(range(32)) + self.assertEqual(self.store.put(raw_key, raw), 0) + + dst = torch.empty(len(raw), dtype=torch.uint8, device="cuda") + self.assertEqual(self.store.get_into(raw_key, dst.data_ptr(), len(raw)), len(raw)) + expected = torch.tensor(list(raw), dtype=torch.uint8, device="cuda") + self.assertTrue(torch.equal(dst, expected)) + + self.store.remove(put_key) + self.store.remove(upsert_key) + self.store.remove(raw_key) + def test_put_get_tensor_with_metadata(self): """Test storing and retrieving PyTorch tensors with metadata using put_tensor_with_metadata/get_tensor_with_metadata.""" import torch @@ -251,15 +284,17 @@ def test_tensor(self): self.assertTrue(d.accepted) self.assertEqual(d.codec, "ragged_tensor") - def test_tensor_mixed_dtype_rejected(self): + def test_tensor_mixed_dtype(self): import torch d = _choose_leaf_codec([torch.tensor([1], dtype=torch.float32), torch.tensor([1], dtype=torch.int64)]) self.assertFalse(d.accepted) + self.assertEqual(d.codec, "ragged_tensor") - def test_tensor_mixed_ndim_rejected(self): + def test_tensor_mixed_ndim(self): import torch d = _choose_leaf_codec([torch.tensor([1]), torch.tensor([[1, 2]])]) - self.assertFalse(d.accepted) + self.assertTrue(d.accepted) + self.assertEqual(d.codec, "ragged_tensor") def test_numeric_sequence(self): d = _choose_leaf_codec([[1, 2, 3], [4, 5]]) @@ -279,7 +314,7 @@ def test_text(self): def test_json(self): d = _choose_leaf_codec([{"a": 1}, {"b": 2}]) self.assertTrue(d.accepted) - self.assertEqual(d.codec, "json_ragged") + self.assertEqual(d.codec, "msgpack_ragged") def test_json_rejects_late_non_serializable(self): values = [{"ok": i} for i in range(200)] + [object()] @@ -294,7 +329,7 @@ def test_scalar(self): def test_fallback(self): d = _choose_leaf_codec([object(), object()]) self.assertFalse(d.accepted) - self.assertEqual(d.codec, "pickle_ragged_fallback") + self.assertEqual(d.codec, "msgpack_ragged") def test_with_nulls(self): d = _choose_leaf_codec(["hello", None, "world"]) @@ -365,6 +400,46 @@ def test_dict_missing_key_vs_none_value(self): self.assertIsInstance(x_leaf.values[1], type(MISSING)) self.assertIsNone(x_leaf.values[2]) + def test_infer_dict_of_tensors(self): + import torch + + leaves, nodes = [], [] + infer_structure( + "r", + [ + {"tokens": torch.arange(2, dtype=torch.int64), "score": 1.0}, + {"tokens": torch.arange(3, dtype=torch.int64), "score": 2.0}, + {"tokens": torch.arange(1, dtype=torch.int64), "score": None}, + ], + leaves, + nodes, + ) + self.assertEqual(len(nodes), 1) + self.assertEqual(nodes[0].node_type, "dict") + by_path = {leaf.path: leaf for leaf in leaves} + self.assertEqual(by_path["r.tokens"].decision.codec, "ragged_tensor") + self.assertEqual(by_path["r.score"].decision.codec, "ndarray") + + def test_infer_dict_of_tensors_missing_keys_and_null_rows(self): + import torch + + leaves, nodes = [], [] + rows = [ + {"tokens": torch.arange(2, dtype=torch.float32), "label": None}, + None, + {"label": 3}, + {"tokens": None, "label": 4}, + ] + infer_structure("r", rows, leaves, nodes) + self.assertEqual(nodes[0].row_mask, [True, False, True, True]) + by_path = {leaf.path: leaf for leaf in leaves} + tokens = by_path["r.tokens"] + self.assertTrue(torch.equal(tokens.values[0], rows[0]["tokens"])) + self.assertIsNone(tokens.values[1]) + self.assertIsInstance(tokens.values[2], type(MISSING)) + self.assertIsNone(tokens.values[3]) + self.assertEqual(tokens.decision.codec, "ragged_tensor") + def test_escape_key_in_path(self): self.assertEqual(_escape_key("simple"), "simple") self.assertEqual(_escape_key("a.b"), "a\\.b") diff --git a/mooncake-wheel/tests/test_regmr_overhead.py b/mooncake-wheel/tests/test_regmr_overhead.py index 506cf17a28..424a701e16 100644 --- a/mooncake-wheel/tests/test_regmr_overhead.py +++ b/mooncake-wheel/tests/test_regmr_overhead.py @@ -23,7 +23,6 @@ import time import numpy as np import unittest -import ctypes # Add project root to path for importing mooncake module sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..')) diff --git a/mooncake-wheel/tests/test_release_wheel_tags.py b/mooncake-wheel/tests/test_release_wheel_tags.py new file mode 100644 index 0000000000..9aa05e6527 --- /dev/null +++ b/mooncake-wheel/tests/test_release_wheel_tags.py @@ -0,0 +1,117 @@ +"""Guard the glibc floor of CI and release wheels. + +Distributable wheels must take their manylinux platform tag from a pinned +container image, never from the GitHub runner. ``scripts/build_wheel.sh`` derives +``PLATFORM_TAG`` from the *build host's* glibc (``detect_glibc_version``), so a +job running on a bare runner silently re-tags itself whenever GitHub bumps the +runner image. That is not hypothetical: the published aarch64 floor moved from +``manylinux_2_35`` (0.3.9) to ``manylinux_2_39`` (0.3.10) with no code change, +which stops ARM Ubuntu 22.04 from resolving any wheel and contradicts the +"OS: Ubuntu 22.04 LTS+" contract in docs/source/getting_started/build.md. + +Running CI and release builds through the same manylinux workflow makes the +detected glibc a constant of the image instead of a property of the runner. See +#2858. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +yaml = pytest.importorskip("yaml") + +# Jobs delegating to this reusable workflow produce the standard CUDA and +# non-CUDA wheel artifacts. Hardware-specific release-{efa,efa-non-cuda,musa, +# npu}.yaml workflows do not call it and are out of scope. +SHARED_BUILD_WORKFLOW = "_build-wheel.yaml" + +# The floor must come from a named manylinux image, matched to the runner's +# architecture: an aarch64 image on an x86 runner (or vice versa) is a +# misconfiguration, not a pinned floor. Deliberately not a digest assertion; +# pinning digests is worth doing for both arches at once, not here. +ARCH_CONTAINER = { + "aarch64": re.compile(r"^pytorch/manylinuxaarch64-builder:cuda\d+\.\d+$"), + "x86_64": re.compile(r"^pytorch/manylinux2_28-builder:cuda\d+\.\d+$"), +} + + +def _runner_arch(runner: str) -> str: + return "aarch64" if runner.endswith("-arm") else "x86_64" + + +def _find_workflows_dir() -> Path | None: + """Locate .github/workflows, or None when the suite has been detached. + + scripts/test_installation.sh copies this directory into a scratch tree, so + the workflow sources are not always reachable from __file__. + """ + for parent in Path(__file__).resolve().parents: + candidate = parent / ".github" / "workflows" + if candidate.is_dir(): + return candidate + return None + + +WORKFLOWS_DIR = _find_workflows_dir() + +if WORKFLOWS_DIR is None: + pytest.skip("workflow sources not available", allow_module_level=True) + + +def _collect_build_jobs() -> list: + jobs = [] + # Both extensions: the repo uses .yml for most workflows and .yaml for the + # release set, so globbing one silently ignores callers named in the other. + for path in sorted(WORKFLOWS_DIR.glob("*.y*ml")): + workflow = yaml.safe_load(path.read_text()) or {} + for job_name, job in (workflow.get("jobs") or {}).items(): + if SHARED_BUILD_WORKFLOW in str(job.get("uses", "")): + jobs.append( + pytest.param(job.get("with") or {}, id=f"{path.name}:{job_name}") + ) + return jobs + + +BUILD_JOBS = _collect_build_jobs() + + +def test_shared_build_workflow_still_has_callers() -> None: + """Fail loudly rather than skip if the marker stops matching. + + Without this, renaming _build-wheel.yaml would empty BUILD_JOBS and leave + every assertion below silently unenforced. + """ + assert BUILD_JOBS, ( + f"no job delegates to {SHARED_BUILD_WORKFLOW}; it was renamed or the " + "wheel workflows stopped using it, so this guard is disarmed" + ) + + +@pytest.mark.parametrize("with_block", BUILD_JOBS) +def test_wheel_build_pins_glibc_floor_to_a_container(with_block: dict) -> None: + runner = str(with_block.get("runner", "")) + container = with_block.get("container", "") + arch = _runner_arch(runner) + expected = ARCH_CONTAINER[arch] + + assert container, ( + "job builds wheels on a bare runner, so its manylinux tag " + "follows the runner's glibc and drifts when GitHub bumps the image; " + f"set container to a {expected.pattern} image" + ) + assert expected.match(container), ( + f"container {container!r} is not the pinned manylinux builder image " + f"for {arch} (runner {runner!r}); expected one matching " + f"{expected.pattern}" + ) + # sbsa-* makes _build-wheel.yaml install the CUDA toolkit with apt, which + # the AlmaLinux-based manylinux image does not have. Any other value + # ('container', 'none') keeps the build inside the image. + cuda = str(with_block.get("cuda", "")) + assert not cuda.startswith("sbsa-"), ( + f"cuda is {cuda!r}; a job pinned to a manylinux container cannot " + "install CUDA via the host package manager" + ) diff --git a/mooncake-wheel/tests/test_replicated_distributed_object_store.py b/mooncake-wheel/tests/test_replicated_distributed_object_store.py index 1e46d80e1f..0ad91d3e6a 100644 --- a/mooncake-wheel/tests/test_replicated_distributed_object_store.py +++ b/mooncake-wheel/tests/test_replicated_distributed_object_store.py @@ -94,13 +94,13 @@ def test_put_with_config_parameter(self): self.assertEqual(self.store.remove(key2), 0) with self.assertRaises(TypeError): - result = self.store.put(key_arg_name_error=key, value=test_data, config=config) + self.store.put(key_arg_name_error=key, value=test_data, config=config) with self.assertRaises(TypeError): - result = self.store.put(key=key, value_arg_name_error=test_data, config=config) + self.store.put(key=key, value_arg_name_error=test_data, config=config) with self.assertRaises(TypeError): - result = self.store.put(key=key, value=test_data, config_arg_name_error=config) + self.store.put(key=key, value=test_data, config_arg_name_error=config) def test_put_batch_with_config_parameter(self): """Test put_batch method with config parameter.""" diff --git a/mooncake-wheel/tests/test_ssd_offload_in_evict.py b/mooncake-wheel/tests/test_ssd_offload_in_evict.py index 389fee6010..cc85c7ab90 100644 --- a/mooncake-wheel/tests/test_ssd_offload_in_evict.py +++ b/mooncake-wheel/tests/test_ssd_offload_in_evict.py @@ -2,10 +2,7 @@ import os import time import threading -import random from mooncake.store import MooncakeDistributedStore -import statistics -import math from collections import defaultdict # The lease time of the kv object, should be set equal to @@ -255,7 +252,6 @@ def test_put_get_in_evict_operations(self): # -------------------------- get_stats.start_timer() index = 0 - count = 0 while index < MAX_REQUESTS: key = "k_" + str(index) @@ -524,8 +520,6 @@ def test_batch_get_in_evict_operations(self): self.assertEqual(len(results), len(batch_keys), "Should return result for each key") success_counter = 0 - failed_counter = 0 - count=0 for i, result in enumerate(results): current_key = batch_keys[i] expected = reference.get(current_key) diff --git a/mooncake-wheel/tests/test_structured_object_store.py b/mooncake-wheel/tests/test_structured_object_store.py index 117c160fba..ad0060d1d0 100644 --- a/mooncake-wheel/tests/test_structured_object_store.py +++ b/mooncake-wheel/tests/test_structured_object_store.py @@ -2,26 +2,55 @@ import ctypes import json +import os import threading import time import numpy as np import pytest +import mooncake.structured_object_store as sos from mooncake.structured_object_store import ( BundleTransferPolicy, + FieldSchema, MooncakeBundleTransfer, RemoteBundleRef, StructuredObjectPayload, StructuredObjectReadSpec, + export_dataproto_ref, + export_ref, + import_dataproto_ref, + import_ref, + is_dataproto_ref_handle, + raw_destination, + tensor_object_buffer, ) +class SimpleDataProto: + def __init__(self, batch=None, non_tensor_batch=None, meta_info=None) -> None: + self.batch = {} if batch is None else batch + self.non_tensor_batch = {} if non_tensor_batch is None else non_tensor_batch + self.meta_info = {} if meta_info is None else meta_info + + @classmethod + def from_dict(cls, batch, non_tensor_batch=None, meta_info=None): + return cls(batch=batch, non_tensor_batch=non_tensor_batch, meta_info=meta_info) + + +class BadDataProto: + def __init__(self) -> None: + pass + + class InMemoryStore: def __init__(self) -> None: self.objects: dict[str, bytes] = {} + self.tensor_objects: dict[str, object] = {} self.lock = threading.Lock() self.registered: set[int] = set() + self.register_buffer_calls = 0 + self.unregister_buffer_calls = 0 self.max_active_puts = 0 self.max_active_gets = 0 self.active_puts = 0 @@ -29,7 +58,13 @@ def __init__(self) -> None: self.get_into_calls = 0 self.get_into_ranges_calls = 0 self.batch_get_into_calls = 0 + self.batch_put_from_calls = 0 + self.put_tensor_from_calls = 0 self.batch_remove_calls = 0 + self.put_tensor_calls = 0 + self.get_tensor_calls = 0 + self.put_configs: list[object] = [] + self.batch_put_from_configs: list[object] = [] def _enter_put(self) -> None: with self.lock: @@ -49,7 +84,8 @@ def _exit_get(self, count: int = 1) -> None: with self.lock: self.active_gets -= count - def put(self, key: str, value) -> int: + def put(self, key: str, value, config=None) -> int: + self.put_configs.append(config) self._enter_put() try: time.sleep(0.01) @@ -71,8 +107,20 @@ def get(self, key: str) -> bytes: def remove(self, key: str, force: bool = False) -> int: with self.lock: self.objects.pop(key, None) + self.tensor_objects.pop(key, None) + return 0 + + def put_tensor(self, key: str, value) -> int: + self.put_tensor_calls += 1 + with self.lock: + self.tensor_objects[key] = value.detach().clone() return 0 + def get_tensor(self, key: str): + self.get_tensor_calls += 1 + with self.lock: + return self.tensor_objects[key].clone() + def batch_remove(self, keys: list[str], force: bool = False) -> list[int]: self.batch_remove_calls += 1 for key in keys: @@ -80,19 +128,29 @@ def batch_remove(self, keys: list[str], force: bool = False) -> list[int]: return [0 for _key in keys] def batch_put_from( - self, keys: list[str], buffer_ptrs: list[int], sizes: list[int] + self, keys: list[str], buffer_ptrs: list[int], sizes: list[int], config=None ) -> list[int]: + self.batch_put_from_calls += 1 + self.batch_put_from_configs.append(config) results: list[int] = [] for key, ptr, size in zip(keys, buffer_ptrs, sizes): data = ctypes.string_at(ptr, size) - results.append(self.put(key, data)) + results.append(self.put(key, data, config=config)) return results + def put_tensor_from(self, key: str, buffer_ptr: int, size: int) -> int: + self.put_tensor_from_calls += 1 + if buffer_ptr not in self.registered: + return -1 + return self.put(key, ctypes.string_at(buffer_ptr, size)) + def register_buffer(self, buffer_ptr: int, size: int) -> int: + self.register_buffer_calls += 1 self.registered.add(buffer_ptr) return 0 def unregister_buffer(self, buffer_ptr: int) -> int: + self.unregister_buffer_calls += 1 self.registered.remove(buffer_ptr) return 0 @@ -172,6 +230,53 @@ def batch_get_into( self._exit_get(len(keys)) +class FakeLease: + def __init__(self, pool: "FakeBufferPool", size: int) -> None: + self.pool = pool + self.size = size + self._buffer = ctypes.create_string_buffer(size) + self.ptr = ctypes.addressof(self._buffer) + self.released = False + + @property + def buffer(self): + return memoryview(self._buffer) + + def release(self) -> None: + if not self.released: + self.pool.release_count += 1 + self.released = True + + +class FakeBufferPool: + def __init__(self) -> None: + self.acquire_count = 0 + self.release_count = 0 + self.acquire_sizes: list[int] = [] + + def acquire(self, size: int) -> FakeLease: + self.acquire_count += 1 + self.acquire_sizes.append(size) + return FakeLease(self, size) + + +class FailingBatchPutStore(InMemoryStore): + def __init__(self, *, fail_on_call: int) -> None: + super().__init__() + self.fail_on_call = fail_on_call + + def batch_put_from( + self, keys: list[str], buffer_ptrs: list[int], sizes: list[int] + ) -> list[int]: + self.batch_put_from_calls += 1 + if self.batch_put_from_calls == self.fail_on_call: + return [-1 for _key in keys] + results: list[int] = [] + for key, ptr, size in zip(keys, buffer_ptrs, sizes): + results.append(self.put(key, ctypes.string_at(ptr, size))) + return results + + class GetOnlyStore(InMemoryStore): batch_get_into = None @@ -193,6 +298,15 @@ class PlainStore(MinimalStore): get_into_ranges = None +class PutTensorOnlyStore(InMemoryStore): + get_tensor = None + + +class NoTensorFastPathStore(InMemoryStore): + put_tensor = None + get_tensor = None + + class FailingPutStore(InMemoryStore): def __init__(self, fail_on_put: int) -> None: super().__init__() @@ -280,14 +394,32 @@ def make_transfer( *, key_prefix: str = "test", default_chunk_bytes: int | None = None, + buffer_pool=None, ) -> tuple[InMemoryStore, MooncakeBundleTransfer]: current_store = store or InMemoryStore() - kwargs = {"key_prefix": key_prefix} + kwargs = {"key_prefix": key_prefix, "buffer_pool": buffer_pool} if default_chunk_bytes is not None: kwargs["default_chunk_bytes"] = default_chunk_bytes return current_store, MooncakeBundleTransfer(current_store, **kwargs) +def real_transfer(key_prefix: str) -> tuple[object, MooncakeBundleTransfer]: + mooncake_store = pytest.importorskip("mooncake.store") + store = mooncake_store.MooncakeDistributedStore() + rc = store.setup( + os.getenv("LOCAL_HOSTNAME", "localhost"), + os.getenv("MC_METADATA_SERVER", "P2PHANDSHAKE"), + 16 * 1024 * 1024, + 4 * 1024 * 1024, + os.getenv("PROTOCOL", "tcp"), + os.getenv("DEVICE_NAME", ""), + os.getenv("MASTER_SERVER", "127.0.0.1:50051"), + ) + if rc != 0: + pytest.skip(f"MooncakeDistributedStore setup failed: {rc}") + return store, MooncakeBundleTransfer(store, key_prefix=key_prefix) + + def structured_payload( metadata: dict[str, object] | None = None, **buffers: object ) -> StructuredObjectPayload: @@ -302,6 +434,79 @@ def write_manifest( ) +def test_put_object_roundtrips_numpy_and_torch_tensor_fields() -> None: + torch = pytest.importorskip("torch") + store, transfer = make_transfer() + array = np.arange(6, dtype=np.float32).reshape(2, 3) + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + scalar_tensor = torch.tensor(1.5, dtype=torch.float32) + bool_tensor = torch.tensor([True, False], dtype=torch.bool) + + result = transfer.get_object( + transfer.put_object( + { + "array": array, + "tensor": tensor, + "scalar": scalar_tensor, + "bool_tensor": bool_tensor, + } + ) + ) + + assert np.array_equal(result["array"], array) + assert torch.equal(result["tensor"], tensor) + assert torch.equal(result["scalar"], scalar_tensor) + assert torch.equal(result["bool_tensor"], bool_tensor) + assert store.put_tensor_calls == 3 + assert store.get_tensor_calls == 3 + + +def test_put_object_roundtrips_wrapped_numpy_and_torch_values() -> None: + torch = pytest.importorskip("torch") + store, transfer = make_transfer() + array = np.arange(6, dtype=np.float32).reshape(2, 3) + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + + assert np.array_equal(transfer.get_object(transfer.put_object(array)), array) + assert torch.equal(transfer.get_object(transfer.put_object(tensor)), tensor) + assert store.put_tensor_calls == 1 + assert store.get_tensor_calls == 1 + + +def test_get_object_requires_get_tensor_for_tensor_payload() -> None: + torch = pytest.importorskip("torch") + store, transfer = make_transfer(PutTensorOnlyStore()) + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + + ref = transfer.put_object({"tensor": tensor}) + + with pytest.raises(RuntimeError, match="does not support get_tensor"): + transfer.get_object(ref) + + +def test_put_object_torch_tensor_raw_fallback_roundtrip() -> None: + torch = pytest.importorskip("torch") + store, transfer = make_transfer(NoTensorFastPathStore()) + tensor = torch.arange(6, dtype=torch.float32).reshape(2, 3) + scalar_tensor = torch.tensor(1.5, dtype=torch.float32) + bool_tensor = torch.tensor([True, False], dtype=torch.bool) + + result = transfer.get_object( + transfer.put_object( + { + "tensor": tensor, + "scalar": scalar_tensor, + "bool_tensor": bool_tensor, + } + ) + ) + + assert torch.equal(result["tensor"], tensor) + assert torch.equal(result["scalar"], scalar_tensor) + assert torch.equal(result["bool_tensor"], bool_tensor) + assert store.batch_get_into_calls > 0 + + def test_bundle_read_spec_full_read_is_partial_special_case() -> None: store, transfer = make_transfer() array = np.arange(16, dtype=np.int32).reshape(4, 4) @@ -366,6 +571,34 @@ def test_bundle_uses_configurable_default_chunk_size() -> None: assert ref.manifest["buffers"]["payload"]["chunks"][0]["bytes"] == 17 +def test_bundle_copy_mode_forces_store_put() -> None: + store, transfer = make_transfer() + payload = bytes(range(128)) + + ref = transfer.put_structured_object( + structured_payload(payload=payload), + chunk_bytes=17, + policy=BundleTransferPolicy(copy_mode="copy"), + ) + + assert store.batch_put_from_calls == 0 + assert store.register_buffer_calls == 0 + assert store.unregister_buffer_calls == 0 + + result = transfer.materialize(transfer.read_spec(ref)) + assert result.objects["payload"] == payload + + +def test_bundle_zero_copy_mode_requires_batch_put_support() -> None: + _store, transfer = make_transfer(MinimalStore()) + + with pytest.raises(RuntimeError, match="zero-copy put"): + transfer.put_structured_object( + structured_payload(payload=b"data"), + policy=BundleTransferPolicy(copy_mode="zero_copy"), + ) + + def test_structured_object_roundtrip() -> None: store, transfer = make_transfer() array = np.arange(12, dtype=np.int16).reshape(3, 4) @@ -428,6 +661,98 @@ def test_structured_object_multichunk_ndarray_uses_range_gather() -> None: assert store.batch_get_into_calls == batch_get_into_calls + 1 +def test_structured_object_ndarray_read_can_use_buffer_pool() -> None: + pool = FakeBufferPool() + store, transfer = make_transfer(buffer_pool=pool) + array = np.arange(64, dtype=np.int16).reshape(8, 8) + payload = structured_payload(weights=array) + + ref = transfer.put_structured_object(payload, chunk_bytes=32) + pool.acquire_sizes.clear() + before_release_count = pool.release_count + before_range_reads = store.get_into_ranges_calls + result = transfer.materialize(transfer.read_spec(ref)) + + assert np.array_equal(result.objects["weights"], array) + assert hasattr(result.objects["weights"], "_mooncake_pool_owner") + assert store.get_into_ranges_calls == before_range_reads + 1 + assert pool.acquire_sizes == [array.nbytes] + + MooncakeBundleTransfer.release_result(result.objects) + assert pool.release_count == before_release_count + 1 + MooncakeBundleTransfer.release_result(result.objects) + assert pool.release_count == before_release_count + 1 + + +def test_structured_object_multi_buffer_payload_uses_pool_batch_put() -> None: + pool = FakeBufferPool() + store, transfer = make_transfer(buffer_pool=pool) + parts = [bytearray(b"ab"), bytearray(b"cd"), bytearray(b"ef")] + payload = StructuredObjectPayload( + metadata={}, + buffers={ + "raw": sos._MultiBufferPayload( + buffers=tuple(memoryview(part) for part in parts), + owners=tuple(parts), + ) + }, + ) + + ref = transfer.put_structured_object(payload, chunk_bytes=4) + result = transfer.materialize(transfer.read_spec(ref)) + + raw = result.objects["raw"] + raw_bytes = raw if isinstance(raw, bytes) else raw.tobytes() + assert raw_bytes == b"abcdef" + assert pool.acquire_sizes == [4, 2, 6] + assert pool.release_count == 2 + assert store.batch_put_from_calls == 2 + + MooncakeBundleTransfer.release_result(result.objects) + assert pool.release_count == 3 + + +def test_structured_object_multi_buffer_put_cleans_all_chunk_keys_on_failure() -> None: + pool = FakeBufferPool() + store = FailingBatchPutStore(fail_on_call=2) + transfer = MooncakeBundleTransfer(store, key_prefix="test", buffer_pool=pool) + parts = [bytearray(b"ab"), bytearray(b"cd"), bytearray(b"ef")] + payload = StructuredObjectPayload( + metadata={}, + buffers={ + "raw": sos._MultiBufferPayload( + buffers=tuple(memoryview(part) for part in parts), + owners=tuple(parts), + ) + }, + ) + + with pytest.raises(RuntimeError): + transfer.put_structured_object(payload, chunk_bytes=2) + + assert store.objects == {} + assert pool.release_count == 2 + + +def test_structured_object_multi_buffer_zero_copy_requires_batch_put_support() -> None: + _store, transfer = make_transfer(MinimalStore()) + parts = [bytearray(b"ab"), bytearray(b"cd")] + payload = StructuredObjectPayload( + metadata={}, + buffers={ + "raw": sos._MultiBufferPayload( + buffers=tuple(memoryview(part) for part in parts), + owners=tuple(parts), + ) + }, + ) + + with pytest.raises(RuntimeError, match="zero-copy put"): + transfer.put_structured_object( + payload, policy=BundleTransferPolicy(copy_mode="zero_copy") + ) + + def test_structured_object_slice_member_uses_partial_range_reads() -> None: store, transfer = make_transfer() array = np.arange(96, dtype=np.int16).reshape(12, 8) @@ -609,48 +934,6 @@ def test_bundle_remove_recovers_after_transient_batch_failure() -> None: assert store.batch_remove_calls == 1 -def test_bundle_concurrent_put_and_read_spec_full_read() -> None: - store, transfer = make_transfer(GetOnlyStore()) - payload = bytes(range(128)) - - ref = transfer.put_structured_object( - structured_payload(payload=payload), - chunk_bytes=8, - policy=BundleTransferPolicy(max_inflight_put=4, put_mode="parallel"), - ) - result = transfer.materialize(transfer.read_spec(ref)) - - assert result.objects["payload"] == payload - assert store.max_active_puts > 1 - assert store.max_active_gets >= 1 - - -def test_bundle_duplicate_source_registration_is_tolerated() -> None: - store, transfer = make_transfer(StrictRegisterStore()) - payload = np.arange(64, dtype=np.uint8).reshape(8, 8) - payload_ptr = ctypes.addressof(ctypes.c_char.from_buffer(payload)) - assert store.register_buffer(payload_ptr, int(payload.nbytes)) == 0 - - ref = transfer.put_structured_object(structured_payload(payload=payload)) - - assert ref.manifest["buffers"]["payload"]["bytes"] == int(payload.nbytes) - assert payload_ptr in store.registered - store.unregister_buffer(payload_ptr) - - -def test_bundle_partial_register_failure_unwinds_registered_buffers() -> None: - store, transfer = make_transfer(FailingRegisterStore(fail_on_register=2)) - payload = bytes(range(64)) - - with pytest.raises(RuntimeError, match="register_buffer"): - transfer.put_structured_object( - structured_payload(payload=payload), chunk_bytes=32 - ) - - assert store.registered == set() - assert store.objects == {} - - def test_bundle_batch_get_failure_unregisters_buffer() -> None: store, transfer = make_transfer(FailingBatchGetStore()) ref = transfer.put_structured_object( @@ -706,6 +989,12 @@ def test_bundle_invalid_policy_and_chunk_size_raise() -> None: ) with pytest.raises(ValueError, match="chunk_bytes"): transfer.put_bundle(b"meta", {"payload": b"data"}, chunk_bytes=0) + with pytest.raises(ValueError, match="copy_mode"): + transfer.put_bundle( + b"meta", + {"payload": b"data"}, + policy=BundleTransferPolicy(copy_mode="invalid"), + ) def test_bundle_invalid_name_and_prefix_raise() -> None: @@ -796,3 +1085,2181 @@ def test_bundle_rejects_tampered_manifest() -> None: transfer.remove_bundle( RemoteBundleRef(manifest_key="test/other/manifest", manifest=ref.manifest) ) + + +def test_structured_object_torch_tensor_falls_back_without_buffer_pool() -> None: + torch = pytest.importorskip("torch") + mooncake_store = pytest.importorskip("mooncake.store") + if not hasattr(mooncake_store, "_serialize_tensor") or not hasattr( + mooncake_store, "_deserialize_tensor" + ): + pytest.skip("built mooncake.store lacks tensor serialization helpers") + store, transfer = make_transfer() + tensor = torch.arange(12, dtype=torch.int64).reshape(3, 4) + + ref = transfer.put_structured_object( + StructuredObjectPayload(buffers={"tensor": tensor}) + ) + result = transfer.materialize(transfer.read_spec(ref)) + + assert store.put_tensor_from_calls == 0 + assert store.batch_put_from_calls == 0 + assert torch.equal(result.objects["tensor"], tensor) + + +def test_structured_object_direct_torch_tensor_materialize_into_uses_real_store() -> ( + None +): + torch = pytest.importorskip("torch") + mooncake_store = pytest.importorskip("mooncake.store") + if not hasattr(mooncake_store, "get_tensor_into") and not hasattr( + mooncake_store.MooncakeDistributedStore, "get_tensor_into" + ): + pytest.skip("built mooncake.store lacks get_tensor_into") + store = mooncake_store.MooncakeDistributedStore() + rc = store.setup( + os.getenv("LOCAL_HOSTNAME", "localhost"), + os.getenv("MC_METADATA_SERVER", "P2PHANDSHAKE"), + 16 * 1024 * 1024, + 4 * 1024 * 1024, + os.getenv("PROTOCOL", "tcp"), + os.getenv("DEVICE_NAME", ""), + os.getenv("MASTER_SERVER", "127.0.0.1:50051"), + ) + if rc != 0: + pytest.skip(f"MooncakeDistributedStore setup failed: {rc}") + transfer = MooncakeBundleTransfer(store, key_prefix="structured-test-direct") + tensor = torch.arange(12, dtype=torch.int64).reshape(3, 4) + ref = transfer.put_structured_object( + StructuredObjectPayload(buffers={"tensor": tensor}) + ) + total_bytes = int(ref.manifest["buffers"]["tensor"]["bytes"]) + destination = ctypes.create_string_buffer(total_bytes) + + try: + result = transfer.materialize_into( + transfer.read_spec(ref), + { + "tensor": tensor_object_buffer( + ctypes.addressof(destination), + total_bytes, + destination, + batch_size=3, + ) + }, + ) + assert torch.equal(result.objects["tensor"], tensor) + finally: + transfer.remove_bundle(ref) + + +def test_structured_object_torch_tensor_zero_copy_rejects_plain_tensor() -> None: + torch = pytest.importorskip("torch") + mooncake_store = pytest.importorskip("mooncake.store") + if not hasattr(mooncake_store, "_serialize_tensor"): + pytest.skip("built mooncake.store lacks tensor serialization helpers") + _store, transfer = make_transfer() + tensor = torch.arange(12, dtype=torch.int64).reshape(3, 4) + + with pytest.raises(ValueError, match="BufferPool tensor-object buffer"): + transfer.put_structured_object( + StructuredObjectPayload(buffers={"tensor": tensor}), + policy=BundleTransferPolicy(copy_mode="zero_copy"), + ) + + +def test_structured_object_torch_tensor_slice_uses_range_reads() -> None: + torch = pytest.importorskip("torch") + mooncake_store = pytest.importorskip("mooncake.store") + if not hasattr(mooncake_store, "_serialize_tensor") or not hasattr( + mooncake_store, "_deserialize_tensor" + ): + pytest.skip("built mooncake.store lacks tensor serialization helpers") + store, transfer = make_transfer(NoTensorFastPathStore()) + tensor = torch.arange(48, dtype=torch.int64).reshape(12, 4) + ref = transfer.put_structured_object( + StructuredObjectPayload(buffers={"tensor": tensor}), chunk_bytes=40 + ) + + result = transfer.materialize( + transfer.read_spec(ref) + .select_members(["tensor"]) + .slice_member("tensor", axis=0, start=3, end=9) + ) + + assert torch.equal(result.objects["tensor"], tensor[3:9]) + assert store.get_into_ranges_calls >= 2 + + +def test_structured_object_direct_torch_tensor_slice_uses_real_store_ranges() -> None: + torch = pytest.importorskip("torch") + mooncake_store = pytest.importorskip("mooncake.store") + if not hasattr(mooncake_store, "_serialize_tensor") or not hasattr( + mooncake_store, "_deserialize_tensor" + ): + pytest.skip("built mooncake.store lacks tensor serialization helpers") + _store, transfer = real_transfer("structured-test-slice") + tensor = torch.arange(48, dtype=torch.int64).reshape(12, 4) + ref = transfer.put_structured_object(StructuredObjectPayload(buffers={"tensor": tensor})) + + try: + result = transfer.materialize( + transfer.read_spec(ref) + .select_members(["tensor"]) + .slice_member("tensor", axis=0, start=3, end=9) + ) + assert torch.equal(result.objects["tensor"], tensor[3:9]) + finally: + transfer.remove_bundle(ref) + + +def test_structured_object_tensor_object_buffer_uses_put_tensor_from() -> None: + store, transfer = make_transfer() + source = ctypes.create_string_buffer(b"tensor-payload") + store.register_buffer(ctypes.addressof(source), len(source.raw)) + + ref = transfer.put_structured_object( + StructuredObjectPayload( + buffers={ + "tensor": tensor_object_buffer( + ctypes.addressof(source), len(source.raw), source, batch_size=1 + ) + } + ), + policy=BundleTransferPolicy(copy_mode="zero_copy"), + ) + + payload = ref.manifest["buffers"]["tensor"] + assert store.put_tensor_from_calls == 1 + assert store.batch_put_from_calls == 0 + assert payload["bytes"] == len(source.raw) + assert payload["chunks"] == [{"key": payload["key"], "bytes": len(source.raw)}] + + +def test_structured_object_tensor_object_buffer_materialize_into_uses_ranges() -> None: + store, transfer = make_transfer() + source = ctypes.create_string_buffer(b"tensor-payload") + store.register_buffer(ctypes.addressof(source), len(source.raw)) + ref = transfer.put_structured_object( + StructuredObjectPayload( + buffers={ + "tensor": tensor_object_buffer( + ctypes.addressof(source), len(source.raw), source, batch_size=1 + ) + } + ), + policy=BundleTransferPolicy(copy_mode="zero_copy"), + ) + destination = ctypes.create_string_buffer(len(source.raw)) + + result = transfer.materialize_into( + transfer.read_spec(ref), + { + "tensor": tensor_object_buffer( + ctypes.addressof(destination), + len(destination.raw), + destination, + batch_size=1, + ) + }, + ) + + assert result.objects["tensor"].ptr == ctypes.addressof(destination) + assert destination.raw == source.raw + assert store.get_into_ranges_calls == 1 + + +def test_structured_object_torch_tensor_zero_copy_uses_real_buffer_pool() -> None: + torch = pytest.importorskip("torch") + mooncake_store = pytest.importorskip("mooncake.store") + if ( + not hasattr(mooncake_store, "BufferPool") + or not hasattr(mooncake_store, "_serialize_tensor") + or not hasattr(mooncake_store, "_deserialize_tensor") + ): + pytest.skip("built mooncake.store lacks BufferPool tensor helpers") + store = mooncake_store.MooncakeDistributedStore() + rc = store.setup( + os.getenv("LOCAL_HOSTNAME", "localhost"), + os.getenv("MC_METADATA_SERVER", "P2PHANDSHAKE"), + 16 * 1024 * 1024, + 4 * 1024 * 1024, + os.getenv("PROTOCOL", "tcp"), + os.getenv("DEVICE_NAME", ""), + os.getenv("MASTER_SERVER", "127.0.0.1:50051"), + ) + if rc != 0: + pytest.skip(f"MooncakeDistributedStore setup failed: {rc}") + pool = mooncake_store.BufferPool(store, min_size_class=4096, alignment=4096) + transfer = MooncakeBundleTransfer( + store, key_prefix="structured-test", buffer_pool=pool + ) + tensor = torch.arange(12, dtype=torch.int64).reshape(3, 4) + metadata, data_ptr, tensor_nbytes, owner = mooncake_store._serialize_tensor(tensor) + total_bytes = len(metadata) + int(tensor_nbytes) + lease = pool.acquire(total_bytes) + view = lease.buffer + view[: len(metadata)] = metadata + ctypes.memmove(lease.ptr + len(metadata), int(data_ptr), int(tensor_nbytes)) + view.release() + + try: + ref = transfer.put_structured_object( + StructuredObjectPayload( + buffers={ + "tensor": tensor_object_buffer( + lease.ptr, total_bytes, lease, batch_size=3 + ) + } + ), + policy=BundleTransferPolicy(copy_mode="zero_copy"), + ) + result = transfer.materialize(transfer.read_spec(ref)) + assert torch.equal(result.objects["tensor"], tensor) + _ = owner + finally: + lease.release() + pool.close() + + +def test_dataproto_helper_requires_batch_size_for_tensor_object_buffer() -> None: + _store, transfer = make_transfer() + data = SimpleDataProto(batch={"tensor": tensor_object_buffer(1, 128, object())}) + + with pytest.raises(TypeError, match="batch_size"): + transfer.put_dataproto(data, policy=BundleTransferPolicy(copy_mode="zero_copy")) + + +def test_dataproto_helper_roundtrip_uses_structured_object() -> None: + store, transfer = make_transfer() + data = SimpleDataProto( + batch={"input_ids": np.arange(12, dtype=np.int64).reshape(4, 3)}, + non_tensor_batch={ + "reward": np.asarray([1.0, 0.0, 0.5, -1.0], dtype=np.float32) + }, + meta_info={"step": 7, "source": "unit"}, + ) + + ref = transfer.put_dataproto( + data, namespace="roll", partition="train", stage="rollout" + ) + result = transfer.get_dataproto(ref, data_cls=SimpleDataProto) + + assert ref.batch_size == 4 + assert set(ref.stage_refs) == {"rollout"} + assert set(ref.field_index) == {"input_ids", "reward"} + assert np.array_equal(result.batch["input_ids"], data.batch["input_ids"]) + assert np.array_equal( + result.non_tensor_batch["reward"], data.non_tensor_batch["reward"] + ) + assert result.meta_info == data.meta_info + stage_metadata = transfer.materialize( + transfer.read_spec(ref.stage_refs["rollout"]).select_members( + ["batch.input_ids"] + ) + ).metadata + assert stage_metadata["dataproto"]["stage"] == "rollout" + + +def test_dataproto_manifest_view_reuses_structured_manifest_specs() -> None: + _store, transfer = make_transfer() + data = SimpleDataProto( + batch={"input_ids": np.arange(12, dtype=np.int64).reshape(4, 3)}, + non_tensor_batch={ + "reward": np.asarray([1.0, 0.0, 0.5, -1.0], dtype=np.float32) + }, + meta_info={"step": 7}, + ) + + ref = transfer.put_dataproto( + data, namespace="roll", partition="train", stage="rollout" + ) + view = transfer.dataproto_manifest_view(ref) + + assert view["namespace"] == "roll" + assert view["partition"] == "train" + assert view["batch_size"] == 4 + assert view["stages"]["rollout"]["dataproto"]["stage"] == "rollout" + assert view["batch_fields"]["input_ids"]["member"] == "batch.input_ids" + assert view["batch_fields"]["input_ids"]["spec"]["encoding"] == "ndarray" + assert view["batch_fields"]["input_ids"]["spec"]["shape"] == [4, 3] + assert view["non_tensor_fields"]["reward"]["spec"]["dtype"] == " None: + _store, transfer = make_transfer() + data = SimpleDataProto( + batch={ + "input_ids": np.arange(8, dtype=np.int64).reshape(4, 2), + "attention_mask": np.ones((4, 2), dtype=np.int32), + }, + non_tensor_batch={ + "reward": np.asarray([1.0, 0.0, 0.5, -1.0], dtype=np.float32), + "uid": np.arange(4, dtype=np.int64), + }, + meta_info={"step": 7, "source": "unit"}, + ) + ref = transfer.put_dataproto(data) + + result = transfer.get_dataproto( + ref, + fields=["input_ids", "reward"], + meta_info_keys=["step"], + ) + batch_only = transfer.get_dataproto(ref, batch_fields=["input_ids"]) + non_tensor_only = transfer.get_dataproto(ref, non_tensor_fields=["reward"]) + + assert set(result["batch"]) == {"input_ids"} + assert set(result["non_tensor_batch"]) == {"reward"} + assert result["meta_info"] == {"step": 7} + assert np.array_equal(result["batch"]["input_ids"], data.batch["input_ids"]) + assert np.array_equal( + result["non_tensor_batch"]["reward"], data.non_tensor_batch["reward"] + ) + assert set(batch_only["batch"]) == {"input_ids"} + assert batch_only["non_tensor_batch"] == {} + assert np.array_equal(batch_only["batch"]["input_ids"], data.batch["input_ids"]) + assert non_tensor_only["batch"] == {} + assert set(non_tensor_only["non_tensor_batch"]) == {"reward"} + assert np.array_equal( + non_tensor_only["non_tensor_batch"]["reward"], data.non_tensor_batch["reward"] + ) + + +def test_dataproto_helper_reads_rows_with_real_store_ranges() -> None: + torch = pytest.importorskip("torch") + mooncake_store = pytest.importorskip("mooncake.store") + if not hasattr(mooncake_store, "_serialize_tensor") or not hasattr( + mooncake_store, "_deserialize_tensor" + ): + pytest.skip("built mooncake.store lacks tensor serialization helpers") + _store, transfer = real_transfer("structured-test-dataproto-rows") + tensor = torch.arange(24, dtype=torch.int64).reshape(6, 4) + data = SimpleDataProto( + batch={ + "tensor": tensor, + "array": np.arange(18, dtype=np.int64).reshape(6, 3), + }, + non_tensor_batch={ + "text": np.asarray( + ["a", "bb", "ccc", "dddd", "eeeee", "ffffff"], dtype=object + ), + "json": np.asarray( + [{"i": 0}, {"i": 1}, {"i": 2}, {"i": 3}, {"i": 4}, {"i": 5}], + dtype=object, + ), + "ragged": np.asarray( + [ + torch.arange(1, dtype=torch.float32), + torch.arange(2, dtype=torch.float32), + torch.arange(3, dtype=torch.float32), + torch.arange(4, dtype=torch.float32), + torch.arange(5, dtype=torch.float32), + torch.arange(6, dtype=torch.float32), + ], + dtype=object, + ), + }, + ) + ref = transfer.put_dataproto(data) + + try: + sliced = transfer.get_dataproto(ref, rows=slice(2, 5)) + gathered = transfer.get_dataproto(ref, rows=[4, 1, 3]) + array_dst = np.empty((3, 3), dtype=np.int64) + into = transfer.get_dataproto( + ref, + batch_fields=["array"], + rows=[4, 1, 3], + destinations={"array": array_dst}, + ) + pool = mooncake_store.BufferPool(_store, 1024 * 1024) + raw_array = pool.acquire(array_dst.nbytes) + raw_array_result = transfer.get_dataproto( + ref, + batch_fields=["array"], + rows=[4, 1, 3], + destinations={ + "array": raw_destination( + raw_array.ptr, + raw_array.size, + raw_array, + pre_registered=True, + ) + }, + ) + tensor_payload_bytes = int( + ref.stage_refs["default"].manifest["buffers"]["batch.tensor"]["metadata_bytes"] + ) + tensor[[4, 1, 3]].numel() * tensor.element_size() + raw_tensor = pool.acquire(tensor_payload_bytes) + raw_tensor_result = transfer.get_dataproto( + ref, + batch_fields=["tensor"], + rows=[4, 1, 3], + destinations={ + "tensor": raw_destination( + raw_tensor.ptr, + raw_tensor.size, + raw_tensor, + pre_registered=True, + ) + }, + ) + + assert torch.equal(sliced["batch"]["tensor"], tensor[2:5]) + assert np.array_equal(sliced["batch"]["array"], data.batch["array"][2:5]) + assert sliced["non_tensor_batch"]["text"].tolist() == [ + "ccc", + "dddd", + "eeeee", + ] + assert sliced["non_tensor_batch"]["json"].tolist() == [ + {"i": 2}, + {"i": 3}, + {"i": 4}, + ] + actual_ragged = sliced["non_tensor_batch"]["ragged"] + assert torch.equal(actual_ragged[0], data.non_tensor_batch["ragged"][2]) + assert torch.equal(actual_ragged[1], data.non_tensor_batch["ragged"][3]) + assert torch.equal(actual_ragged[2], data.non_tensor_batch["ragged"][4]) + assert torch.equal(gathered["batch"]["tensor"], tensor[[4, 1, 3]]) + assert np.array_equal(gathered["batch"]["array"], data.batch["array"][[4, 1, 3]]) + assert gathered["non_tensor_batch"]["text"].tolist() == ["eeeee", "bb", "dddd"] + assert into["batch"]["array"] is array_dst + assert np.array_equal(array_dst, data.batch["array"][[4, 1, 3]]) + assert np.array_equal( + raw_array_result["batch"]["array"], data.batch["array"][[4, 1, 3]] + ) + assert raw_tensor_result["batch"]["tensor"].ptr == raw_tensor.ptr + decoded_tensor = mooncake_store._deserialize_tensor( + bytes(raw_tensor.buffer[:tensor_payload_bytes]) + ) + assert torch.equal(decoded_tensor, tensor[[4, 1, 3]]) + raw_array.release() + raw_tensor.release() + pool.close() + finally: + transfer.cleanup_dataproto(ref) + + +def test_dataproto_helper_supports_dict_cls_and_reports_bad_cls() -> None: + _store, transfer = make_transfer() + ref = transfer.put_dataproto( + SimpleDataProto(batch={"input_ids": np.arange(4)}, meta_info={"step": 1}) + ) + + result = transfer.get_dataproto(ref, data_cls=dict) + + assert result["meta_info"] == {"step": 1} + assert np.array_equal(result["batch"]["input_ids"], np.arange(4)) + with pytest.raises(TypeError, match="cannot be constructed"): + transfer.get_dataproto(ref, data_cls=BadDataProto) + + +def test_dataproto_helper_accepts_legacy_dict_inputs() -> None: + _store, transfer = make_transfer() + plain_ref = transfer.put_dataproto({"input_ids": np.arange(4)}) + envelope_ref = transfer.put_dataproto( + { + "batch": {"tokens": np.arange(6).reshape(3, 2)}, + "non_tensor_batch": {"uid": np.asarray(["a", "b", "c"], dtype=object)}, + "meta_info": {"step": 3}, + } + ) + + plain = transfer.get_dataproto(plain_ref) + envelope = transfer.get_dataproto(envelope_ref) + + assert np.array_equal(plain["batch"]["input_ids"], np.arange(4)) + assert np.array_equal(envelope["batch"]["tokens"], np.arange(6).reshape(3, 2)) + assert envelope["non_tensor_batch"]["uid"].tolist() == ["a", "b", "c"] + assert envelope["meta_info"] == {"step": 3} + +def _schema_test_data(field: str, values: np.ndarray) -> dict[str, object]: + return { + "batch": {"input_ids": np.arange(len(values))}, + "non_tensor_batch": {field: values}, + } + + +def test_dataproto_field_schema_encodes_typed_ragged_non_tensor_field() -> None: + _store, transfer = make_transfer() + values = np.empty(3, dtype=object) + values[:] = [ + np.asarray([1, 2], dtype=np.int32), + None, + np.asarray([3], dtype=np.int32), + ] + + ref = transfer.put_dataproto( + _schema_test_data("tokens", values), + field_schemas={ + "tokens": FieldSchema( + codec="typed_ragged", + metadata={"section": "non_tensor_batch", "dtype": "int32"}, + ) + }, + ) + result = transfer.get_dataproto(ref)["non_tensor_batch"]["tokens"] + + assert result[0] == [1, 2] + assert result[1] is None + assert result[2] == [3] + + bad_text = np.asarray([object()], dtype=object) + with pytest.raises(AttributeError, match="failed to encode.*'text'.*utf8_ragged"): + transfer.put_dataproto( + _schema_test_data("text", bad_text), + field_schemas={ + "text": FieldSchema(codec="utf8_ragged") + }, + ) + + +def test_dataproto_field_schema_validates_schema_errors() -> None: + _store, transfer = make_transfer() + + with pytest.raises(ValueError, match="declares section"): + transfer.put_dataproto( + {"batch": {"input_ids": np.arange(3)}}, + field_schemas={ + "input_ids": FieldSchema( + codec="ndarray", metadata={"section": "non_tensor_batch"} + ) + }, + ) + + rows = np.empty(1, dtype=object) + rows[:] = [{"image.data": object()}] + with pytest.raises(ValueError, match="must not contain"): + transfer.put_dataproto( + _schema_test_data("mm", rows), + field_schemas={ + "mm": FieldSchema( + codec="ragged_tensor_dict", + metadata={"section": "non_tensor_batch"}, + ) + }, + ) + for bad_key in ["", "image/data", "image\\data"]: + rows[:] = [{bad_key: object()}] + with pytest.raises(ValueError, match="must not"): + transfer.put_dataproto( + _schema_test_data("mm", rows), + field_schemas={ + "mm": FieldSchema( + codec="ragged_tensor_dict", + metadata={"section": "non_tensor_batch"}, + ) + }, + ) + + values = np.empty(2, dtype=object) + values[:] = ["ok", None] + with pytest.raises(ValueError, match="not nullable"): + transfer.put_dataproto( + _schema_test_data("text", values), + field_schemas={ + "text": FieldSchema( + codec="auto", + nullable=False, + metadata={"section": "non_tensor_batch"}, + ) + }, + ) + schema = { + "mm": FieldSchema( + codec="ragged_tensor_dict", + metadata={"section": "non_tensor_batch", "keys": ["image"]}, + ) + } + for row_value, error_type, message in [ + ({"image": object(), "audio": object()}, ValueError, "keys not declared"), + ({"image": None}, TypeError, "explicit None"), + ]: + rows[:] = [row_value] + with pytest.raises(error_type, match=message): + transfer.put_dataproto(_schema_test_data("mm", rows), field_schemas=schema) + + +def test_dataproto_field_schema_allows_same_named_meta_info() -> None: + _store, transfer = make_transfer() + values = np.asarray(["a", "b"], dtype=object) + + ref = transfer.put_dataproto( + { + "batch": {"input_ids": np.arange(2)}, + "non_tensor_batch": {"label": values}, + "meta_info": {"label": "metadata-label"}, + }, + field_schemas={ + "label": FieldSchema( + codec="utf8_ragged", metadata={"section": "non_tensor_batch"} + ) + }, + ) + + result = transfer.get_dataproto(ref) + + assert result["non_tensor_batch"]["label"].tolist() == ["a", "b"] + assert result["meta_info"]["label"] == "metadata-label" + + +def test_dataproto_field_schema_does_not_apply_meta_info_schema_to_non_tensor() -> None: + _store, transfer = make_transfer() + values = np.asarray([{"k": 1}, {"k": 2}], dtype=object) + + ref = transfer.put_dataproto( + { + "batch": {"input_ids": np.arange(2)}, + "non_tensor_batch": {"label": values}, + "meta_info": {"label": "metadata-label"}, + }, + field_schemas={ + "label": FieldSchema( + codec="utf8_ragged", metadata={"section": "meta_info"} + ) + }, + ) + + result = transfer.get_dataproto(ref) + assert result["non_tensor_batch"]["label"].tolist() == [{"k": 1}, {"k": 2}] + assert result["meta_info"]["label"] == "metadata-label" + + +def test_dataproto_field_schema_encodes_ragged_tensor_dict() -> None: + torch = pytest.importorskip("torch", exc_type=ImportError) + _store, transfer = make_transfer() + rows = np.empty(4, dtype=object) + rows[:] = [ + {"image": torch.arange(4, dtype=torch.float32).reshape(2, 2)}, + None, + {}, + {"image": torch.arange(2, dtype=torch.float32)}, + ] + + ref = transfer.put_dataproto( + _schema_test_data("multi_modal_inputs", rows), + field_schemas={ + "multi_modal_inputs": FieldSchema( + codec="ragged_tensor_dict", + metadata={"section": "non_tensor_batch", "keys": {"image": None}}, + ) + }, + ) + + actual = transfer.get_dataproto(ref)["non_tensor_batch"]["multi_modal_inputs"] + assert ref.encoded_non_tensor["multi_modal_inputs"]["codec"] == "ragged_tensor_dict" + assert torch.equal(actual[0]["image"], rows[0]["image"]) + assert actual[1] is None + assert actual[2] == {} + assert torch.equal(actual[3]["image"], rows[3]["image"]) + + sliced = transfer.get_dataproto(ref, rows=slice(1, 4))["non_tensor_batch"][ + "multi_modal_inputs" + ] + indexed = transfer.get_dataproto(ref, rows=[3, 0])["non_tensor_batch"][ + "multi_modal_inputs" + ] + assert sliced[0] is None + assert sliced[1] == {} + assert torch.equal(sliced[2]["image"], rows[3]["image"]) + assert torch.equal(indexed[0]["image"], rows[3]["image"]) + assert torch.equal(indexed[1]["image"], rows[0]["image"]) + all_null = np.empty(3, dtype=object) + all_null[:] = [None, None, None] + all_null_ref = transfer.put_dataproto( + _schema_test_data("multi_modal_inputs", all_null), + field_schemas={ + "multi_modal_inputs": FieldSchema( + codec="ragged_tensor_dict", + metadata={"section": "non_tensor_batch", "keys": ["image"]}, + ) + }, + ) + encoded = all_null_ref.encoded_non_tensor["multi_modal_inputs"] + assert encoded["metadata"]["keys"] == ["image"] + assert "image.data" in encoded["payload_members"] + assert transfer.get_dataproto(all_null_ref)["non_tensor_batch"][ + "multi_modal_inputs" + ].tolist() == [None, None, None] + assert transfer.get_dataproto(all_null_ref, rows=slice(1, 3))[ + "non_tensor_batch" + ]["multi_modal_inputs"].tolist() == [None, None] + assert transfer.get_dataproto(all_null_ref, rows=[2, 0])["non_tensor_batch"][ + "multi_modal_inputs" + ].tolist() == [None, None] + + +def test_unified_put_get_roundtrips_flat_dict() -> None: + _store, transfer = make_transfer() + data = { + "input_ids": np.arange(6, dtype=np.int64).reshape(3, 2), + "tokens": [np.asarray([1, 2]), None, np.asarray([3])], + "uid": ["a", "b", "c"], + "tags": ["science", "math"], + "step": 7, + } + + ref = transfer.put(data, type="dict") + result = transfer.get(ref, type="dict") + + assert np.array_equal(result["input_ids"], data["input_ids"]) + assert result["tokens"][1] is None + assert np.array_equal(result["tokens"][0], np.asarray([1, 2])) + assert np.array_equal(result["tokens"][2], np.asarray([3])) + assert result["uid"] == ["a", "b", "c"] + assert result["tags"] == ["science", "math"] + assert result["step"] == 7 + + +def test_unified_put_rejects_unknown_type() -> None: + _store, transfer = make_transfer() + + with pytest.raises(ValueError, match="unsupported Mooncake payload type"): + transfer.put({}, type="unknown") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("policy", "config_attr", "use_pool"), + [ + (BundleTransferPolicy(copy_mode="copy"), "put_configs", False), + (None, "batch_put_from_configs", True), + ], +) +def test_unified_dict_put_passes_store_config_to_writes(policy, config_attr, use_pool) -> None: + store, transfer = make_transfer(buffer_pool=FakeBufferPool() if use_pool else None) + config = object() + data = {"input_ids": np.arange(12, dtype=np.int64).reshape(4, 3)} + kwargs = {"config": config} + if policy is not None: + kwargs["policy"] = policy + + ref = transfer.put(data, type="dict", **kwargs) + result = transfer.get(ref, type="dict") + + assert np.array_equal(result["input_ids"], data["input_ids"]) + configs = getattr(store, config_attr) + assert configs + assert all(item is config for item in configs) + + +def test_unified_dict_put_accepts_field_schemas() -> None: + def schema(codec: str, section: str, dtype: str | None = None) -> FieldSchema: + metadata = {"section": section} + if dtype is not None: + metadata["dtype"] = dtype + return FieldSchema(codec=codec, metadata=metadata) + + _store, transfer = make_transfer() + values = np.empty(2, dtype=object) + values[:] = [np.asarray([1, 2], dtype=np.int32), None] + + ref = transfer.put( + { + "input_ids": np.arange(2), + "tokens": values, + "partition": [0, 1], + "response_lengths": [2, 0], + "global_batch_sizes": [2], + "num_microbatches": 1, + "step": 7, + }, + field_schemas={ + "tokens": schema("typed_ragged", "non_tensor_batch", "int32"), + "partition": schema("ndarray", "non_tensor_batch", "int64"), + "response_lengths": schema("ndarray", "non_tensor_batch", "int64"), + "global_batch_sizes": schema("auto", "meta_info"), + "num_microbatches": schema("auto", "meta_info"), + }, + type="dict", + ) + + result = transfer.get(ref, type="dict") + + assert np.array_equal(result["input_ids"], np.arange(2)) + assert result["tokens"] == [[1, 2], None] + assert result["partition"] == [0, 1] + assert result["response_lengths"] == [2, 0] + assert result["global_batch_sizes"] == [2] + assert result["num_microbatches"] == 1 + assert result["step"] == 7 + + +def test_release_result_recurses_into_ragged_row_views() -> None: + class FakeOwner: + def __init__(self) -> None: + self.release_count = 0 + + def release(self) -> None: + self.release_count += 1 + + class OwnedArray(np.ndarray): + def __new__(cls, data, owner): + array = np.asarray(data).view(cls) + array._mooncake_pool_owner = owner + return array + + def __array_finalize__(self, obj) -> None: + if obj is not None: + self._mooncake_pool_owner = getattr( + obj, "_mooncake_pool_owner", None + ) + + owner = FakeOwner() + other_owner = FakeOwner() + base = OwnedArray(np.arange(6, dtype=np.int64), owner) + other = OwnedArray(np.arange(2, dtype=np.int64), other_owner) + object_items = np.empty(1, dtype=object) + object_items[0] = base[4:5] + + MooncakeBundleTransfer.release_result( + {"values": [None, 0, base[:2], {"nested": (base[2:4], object_items)}, other[:1]]} + ) + + assert owner.release_count == 1 + assert other_owner.release_count == 1 + + +def test_ref_aliases_match_dataproto_ref_helpers() -> None: + assert export_ref is export_dataproto_ref + assert import_ref is import_dataproto_ref + + +def test_release_result_is_available_for_split_api_compatibility() -> None: + result = {"tokens": [np.asarray([1, 2]), None]} + + assert MooncakeBundleTransfer.release_result(result) is None + assert result["tokens"][0].tolist() == [1, 2] + + +def test_unified_dict_get_rejects_flattened_key_collision() -> None: + _store, transfer = make_transfer() + ref = transfer.put_dataproto( + { + "batch": {"uid": np.arange(3)}, + "meta_info": {"uid": "meta-value"}, + } + ) + + with pytest.raises(ValueError, match="Duplicate keys"): + transfer.get(ref, type="dict") + + +def test_dataproto_helper_treats_reserved_plain_dict_keys_as_batch_fields() -> None: + _store, transfer = make_transfer() + ref = transfer.put_dataproto({"batch": np.arange(4), "meta_info": np.arange(4)}) + + result = transfer.get_dataproto(ref) + + assert np.array_equal(result["batch"]["batch"], np.arange(4)) + assert np.array_equal(result["batch"]["meta_info"], np.arange(4)) + assert result["meta_info"] == {} + + +def test_dataproto_helper_rejects_cross_section_field_name_collision() -> None: + _store, transfer = make_transfer() + + with pytest.raises(ValueError, match="overlap"): + transfer.put_dataproto( + { + "batch": {"uid": np.arange(4)}, + "non_tensor_batch": {"uid": np.asarray(["a", "b", "c", "d"])}, + } + ) + + +def test_dataproto_ref_handle_roundtrip_materializes_and_cleans_up() -> None: + store, transfer = make_transfer() + input_ids = np.arange(6, dtype=np.int64).reshape(3, 2) + ref = transfer.put_dataproto( + SimpleDataProto( + batch={"input_ids": input_ids}, + non_tensor_batch={"uid": np.asarray(["a", "b", "c"], dtype=object)}, + meta_info={"step": 3, "tags": ["rollout"]}, + ), + namespace="roll", + partition="train", + stage="rollout", + ) + handle = export_dataproto_ref(ref) + + assert is_dataproto_ref_handle(handle) + assert handle["stage_refs"] == { + "rollout": {"manifest_key": ref.stage_refs["rollout"].manifest_key} + } + imported = import_dataproto_ref(handle) + assert imported.stage_refs["rollout"].manifest == {} + + result = transfer.get_dataproto(handle) + assert np.array_equal(result["batch"]["input_ids"], input_ids) + assert result["non_tensor_batch"]["uid"].tolist() == ["a", "b", "c"] + assert result["meta_info"] == {"step": 3, "tags": ["rollout"]} + assert transfer.dataproto_manifest_view(handle)["meta_info_keys"] == [ + "step", + "tags", + ] + + transfer.cleanup_dataproto(handle) + assert store.objects == {} + + +def test_dataproto_ref_handle_exports_numpy_meta_and_validates_indexes() -> None: + _store, transfer = make_transfer() + ref = transfer.put_dataproto( + SimpleDataProto( + batch={"input_ids": np.arange(4)}, + meta_info={"scores": np.asarray([1, 2], dtype=np.int64)}, + ) + ) + ref.global_indexes = [0, 2] + + handle = export_dataproto_ref(ref) + imported = import_dataproto_ref(handle) + + assert handle["meta_info"] == {"scores": [1, 2]} + assert imported.global_indexes == [0, 2] + bad_indexes = dict(handle) + bad_indexes["global_indexes"] = "0,2" + with pytest.raises(ValueError, match="global_indexes"): + import_dataproto_ref(bad_indexes) + + +def test_dataproto_ref_handle_rejects_unknown_field_stage() -> None: + _store, transfer = make_transfer() + ref = transfer.put_dataproto(SimpleDataProto(batch={"input_ids": np.arange(4)})) + handle = export_dataproto_ref(ref) + handle["field_index"]["input_ids"] = dict(handle["field_index"]["input_ids"]) + handle["field_index"]["input_ids"]["stage"] = "missing" + + with pytest.raises(ValueError, match="unknown stage"): + import_dataproto_ref(handle) + + +def test_dataproto_helper_appends_stage_fields_without_rewriting_existing() -> None: + store, transfer = make_transfer() + rollout = SimpleDataProto( + batch={"input_ids": np.arange(8, dtype=np.int64).reshape(4, 2)}, + meta_info={"step": 1}, + ) + ref = transfer.put_dataproto(rollout, stage="rollout") + initial_object_count = len(store.objects) + + old_log_prob = SimpleDataProto( + batch={"old_log_probs": np.arange(4, dtype=np.float32)}, + meta_info={"stage": "old_log_prob"}, + ) + ref = transfer.append_dataproto_fields(ref, old_log_prob, stage="old_log_prob") + result = transfer.get_dataproto(ref) + + assert set(ref.stage_refs) == {"rollout", "old_log_prob"} + assert len(store.objects) > initial_object_count + assert np.array_equal(result["batch"]["input_ids"], rollout.batch["input_ids"]) + assert np.array_equal( + result["batch"]["old_log_probs"], old_log_prob.batch["old_log_probs"] + ) + assert result["meta_info"] == {"step": 1, "stage": "old_log_prob"} + + with pytest.raises(ValueError, match="already exist"): + transfer.append_dataproto_fields(ref, rollout, stage="duplicate") + + +def test_dataproto_helper_same_stage_append_reads_rows_with_real_store() -> None: + pytest.importorskip("mooncake.store") + _store, transfer = real_transfer("structured-test-dataproto-same-stage") + input_ids = np.arange(24, dtype=np.int64).reshape(6, 4) + old_log_probs = np.linspace(0.0, 1.0, 6, dtype=np.float32) + text = np.asarray(["", "bb", "ccc", "dddd", "eeeee", "ffffff"], dtype=object) + ref = transfer.put_dataproto( + SimpleDataProto(batch={"input_ids": input_ids}, meta_info={"step": 1}), + stage="rollout", + ) + old_ref = ref + old_manifest_key = ref.stage_refs["rollout"].manifest_key + ref = transfer.append_dataproto_fields( + ref, + SimpleDataProto( + batch={"old_log_probs": old_log_probs}, + non_tensor_batch={"text": text}, + meta_info={"stage": "rollout-extra"}, + ), + stage="rollout", + ) + + try: + old_result = transfer.get_dataproto(old_ref) + result = transfer.get_dataproto(ref) + sliced = transfer.get_dataproto(ref, rows=slice(2, 5)) + gathered = transfer.get_dataproto( + ref, fields=["input_ids", "text"], rows=[4, 1, 3] + ) + view = transfer.dataproto_manifest_view(ref) + + assert set(ref.stage_refs) == {"rollout"} + assert ref.stage_refs["rollout"].manifest_key != old_manifest_key + assert np.array_equal(old_result["batch"]["input_ids"], input_ids) + assert np.array_equal(result["batch"]["input_ids"], input_ids) + assert np.array_equal(result["batch"]["old_log_probs"], old_log_probs) + assert result["non_tensor_batch"]["text"].tolist() == text.tolist() + assert result["meta_info"] == {"step": 1, "stage": "rollout-extra"} + assert np.array_equal(sliced["batch"]["input_ids"], input_ids[2:5]) + assert np.array_equal(sliced["batch"]["old_log_probs"], old_log_probs[2:5]) + assert sliced["non_tensor_batch"]["text"].tolist() == text[2:5].tolist() + assert np.array_equal(gathered["batch"]["input_ids"], input_ids[[4, 1, 3]]) + assert gathered["non_tensor_batch"]["text"].tolist() == text[[4, 1, 3]].tolist() + assert view["batch_fields"]["input_ids"]["stage"] == "rollout" + assert view["batch_fields"]["old_log_probs"]["stage"] == "rollout" + assert view["non_tensor_fields"]["text"]["stage"] == "rollout" + assert set(ref.encoded_non_tensor) == {"text"} + finally: + transfer.cleanup_dataproto(ref) + + +def test_dataproto_helper_reads_rollout_transfer_data_with_real_store() -> None: + pytest.importorskip("mooncake.store") + _store, transfer = real_transfer("structured-test-rollout-transfer") + row_ids = [f"rollout-row-{index}" for index in range(6)] + input_ids = np.arange(48, dtype=np.int64).reshape(6, 8) + attention_mask = (input_ids % 3 != 0).astype(np.int32) + position_ids = np.tile(np.arange(8, dtype=np.int64), (6, 1)) + responses = (input_ids[:, -3:] + 100).astype(np.int64) + response_mask = np.ones((6, 3), dtype=np.int32) + prompts = np.asarray( + ["", "prompt-b", "prompt-c", "prompt-d", "prompt-e", "prompt-f"], + dtype=object, + ) + sample_meta = np.asarray( + [{"row": index, "tag": row_id} for index, row_id in enumerate(row_ids)], + dtype=object, + ) + old_log_probs = np.linspace(-0.5, 0.5, 18, dtype=np.float32).reshape(6, 3) + ref_log_probs = old_log_probs + np.float32(0.25) + rewards = np.linspace(0.0, 1.0, 6, dtype=np.float32) + advantages = rewards + np.float32(10.0) + returns = rewards + np.float32(20.0) + values = rewards + np.float32(30.0) + rollout_data = { + "batch": { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "responses": responses, + "response_mask": response_mask, + }, + "non_tensor_batch": { + "prompts": prompts, + "sample_meta": sample_meta, + }, + "meta_info": {"roll_row_ids": row_ids, "global_step": 7}, + } + logprob_data = { + "batch": { + "old_log_probs": old_log_probs, + "ref_log_probs": ref_log_probs, + }, + "non_tensor_batch": {}, + "meta_info": {"logprob_stage": "complete"}, + } + value_data = { + "batch": { + "rewards": rewards, + "advantages": advantages, + "returns": returns, + "values": values, + }, + "non_tensor_batch": {}, + "meta_info": {"critic_stage": "complete"}, + } + ref = transfer.put_dataproto( + rollout_data, + namespace="roll", + partition="remote-batch", + stage="rollout", + ) + ref = transfer.append_dataproto_fields(ref, logprob_data, stage="logprob") + ref = transfer.append_dataproto_fields(ref, value_data, stage="rollout") + handle = export_dataproto_ref(ref) + + try: + full = transfer.get_dataproto(handle) + selected = transfer.get_dataproto( + handle, + fields=["input_ids", "old_log_probs", "advantages", "prompts"], + meta_info_keys=["roll_row_ids"], + ) + sliced = transfer.get_dataproto( + handle, + fields=["responses", "ref_log_probs", "returns", "sample_meta"], + rows=slice(1, 5), + ) + gathered = transfer.get_dataproto( + handle, + fields=["attention_mask", "values", "prompts", "sample_meta"], + rows=[5, 0, 3], + ) + imported = import_dataproto_ref(handle) + imported_selected = transfer.get_dataproto( + imported, + batch_fields=["position_ids", "rewards"], + rows=[2, 4], + ) + view = transfer.dataproto_manifest_view(handle) + + assert full["meta_info"] == { + "roll_row_ids": row_ids, + "global_step": 7, + "logprob_stage": "complete", + "critic_stage": "complete", + } + assert np.array_equal(full["batch"]["input_ids"], input_ids) + assert np.array_equal(full["batch"]["attention_mask"], attention_mask) + assert np.array_equal(full["batch"]["position_ids"], position_ids) + assert np.array_equal(full["batch"]["responses"], responses) + assert np.array_equal(full["batch"]["response_mask"], response_mask) + assert np.array_equal(full["batch"]["old_log_probs"], old_log_probs) + assert np.array_equal(full["batch"]["ref_log_probs"], ref_log_probs) + assert np.array_equal(full["batch"]["rewards"], rewards) + assert np.array_equal(full["batch"]["advantages"], advantages) + assert np.array_equal(full["batch"]["returns"], returns) + assert np.array_equal(full["batch"]["values"], values) + assert full["non_tensor_batch"]["prompts"].tolist() == prompts.tolist() + assert full["non_tensor_batch"]["sample_meta"].tolist() == sample_meta.tolist() + + assert set(selected["batch"]) == {"input_ids", "old_log_probs", "advantages"} + assert set(selected["non_tensor_batch"]) == {"prompts"} + assert selected["meta_info"] == {"roll_row_ids": row_ids} + assert np.array_equal(selected["batch"]["input_ids"], input_ids) + assert np.array_equal(selected["batch"]["old_log_probs"], old_log_probs) + assert np.array_equal(selected["batch"]["advantages"], advantages) + assert selected["non_tensor_batch"]["prompts"].tolist() == prompts.tolist() + + assert np.array_equal(sliced["batch"]["responses"], responses[1:5]) + assert np.array_equal(sliced["batch"]["ref_log_probs"], ref_log_probs[1:5]) + assert np.array_equal(sliced["batch"]["returns"], returns[1:5]) + assert sliced["non_tensor_batch"]["sample_meta"].tolist() == sample_meta[1:5].tolist() + + assert np.array_equal(gathered["batch"]["attention_mask"], attention_mask[[5, 0, 3]]) + assert np.array_equal(gathered["batch"]["values"], values[[5, 0, 3]]) + assert gathered["non_tensor_batch"]["prompts"].tolist() == prompts[[5, 0, 3]].tolist() + assert gathered["non_tensor_batch"]["sample_meta"].tolist() == sample_meta[[5, 0, 3]].tolist() + + assert np.array_equal(imported_selected["batch"]["position_ids"], position_ids[[2, 4]]) + assert np.array_equal(imported_selected["batch"]["rewards"], rewards[[2, 4]]) + assert view["batch_fields"]["input_ids"]["stage"] == "rollout" + assert view["batch_fields"]["old_log_probs"]["stage"] == "logprob" + assert view["batch_fields"]["advantages"]["stage"] == "rollout" + assert view["non_tensor_fields"]["prompts"]["stage"] == "rollout" + finally: + transfer.cleanup_dataproto(ref) + + +def test_dataproto_helper_reads_large_rollout_transfer_data_with_real_store() -> None: + pytest.importorskip("mooncake.store") + _store, transfer = real_transfer("structured-test-large-rollout-transfer") + batch_size = 96 + prompt_len = 128 + response_len = 64 + total_len = prompt_len + response_len + row_ids = [f"large-rollout-row-{index}" for index in range(batch_size)] + token_grid = np.arange(batch_size * total_len, dtype=np.int64).reshape( + batch_size, total_len + ) + input_ids = token_grid + 1000 + attention_mask = (token_grid % 7 != 0).astype(np.int32) + position_ids = np.tile(np.arange(total_len, dtype=np.int64), (batch_size, 1)) + responses = input_ids[:, prompt_len:] + response_mask = np.ones((batch_size, response_len), dtype=np.int32) + action_log_probs = np.linspace( + -3.0, 3.0, batch_size * response_len, dtype=np.float32 + ).reshape(batch_size, response_len) + ref_log_probs = action_log_probs + np.float32(0.125) + values = np.linspace(0.0, 1.0, batch_size, dtype=np.float32) + rewards = values + np.float32(1.0) + advantages = values + np.float32(2.0) + returns = values + np.float32(3.0) + prompts = np.asarray( + ["" if index % 17 == 0 else f"prompt-{index}-" + "x" * (index % 23) for index in range(batch_size)], + dtype=object, + ) + sample_meta = np.asarray( + [ + {"row": index, "row_id": row_id, "prompt_len": prompt_len, "response_len": response_len} + for index, row_id in enumerate(row_ids) + ], + dtype=object, + ) + ref = transfer.put_dataproto( + { + "batch": { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "responses": responses, + "response_mask": response_mask, + }, + "non_tensor_batch": { + "prompts": prompts, + "sample_meta": sample_meta, + }, + "meta_info": {"roll_row_ids": row_ids, "global_step": 11}, + }, + namespace="roll", + partition="large-remote-batch", + stage="rollout", + ) + ref = transfer.append_dataproto_fields( + ref, + { + "batch": { + "action_log_probs": action_log_probs, + "ref_log_probs": ref_log_probs, + }, + "non_tensor_batch": {}, + "meta_info": {"logprob_stage": "complete"}, + }, + stage="logprob", + ) + ref = transfer.append_dataproto_fields( + ref, + { + "batch": { + "values": values, + "rewards": rewards, + "advantages": advantages, + "returns": returns, + }, + "non_tensor_batch": {}, + "meta_info": {"critic_stage": "complete"}, + }, + stage="rollout", + ) + handle = export_dataproto_ref(ref) + gathered_rows = [95, 0, 63, 17, 42] + + try: + selected = transfer.get_dataproto( + handle, + fields=["input_ids", "action_log_probs", "values", "prompts"], + rows=slice(12, 76), + ) + gathered = transfer.get_dataproto( + handle, + fields=["responses", "ref_log_probs", "returns", "sample_meta"], + rows=gathered_rows, + ) + destination = np.empty((len(gathered_rows), total_len), dtype=np.int64) + into = transfer.get_dataproto( + handle, + batch_fields=["position_ids"], + rows=gathered_rows, + destinations={"position_ids": destination}, + ) + meta_only = transfer.get_dataproto(handle, fields=[], meta_info_keys=["roll_row_ids"]) + full_tail = transfer.get_dataproto( + handle, + fields=["attention_mask", "advantages"], + rows=slice(batch_size - 8, batch_size), + ) + + assert np.array_equal(selected["batch"]["input_ids"], input_ids[12:76]) + assert np.array_equal( + selected["batch"]["action_log_probs"], action_log_probs[12:76] + ) + assert np.array_equal(selected["batch"]["values"], values[12:76]) + assert selected["non_tensor_batch"]["prompts"].tolist() == prompts[12:76].tolist() + assert np.array_equal(gathered["batch"]["responses"], responses[gathered_rows]) + assert np.array_equal( + gathered["batch"]["ref_log_probs"], ref_log_probs[gathered_rows] + ) + assert np.array_equal(gathered["batch"]["returns"], returns[gathered_rows]) + assert gathered["non_tensor_batch"]["sample_meta"].tolist() == sample_meta[gathered_rows].tolist() + assert into["batch"]["position_ids"] is destination + assert np.array_equal(destination, position_ids[gathered_rows]) + assert meta_only["batch"] == {} + assert meta_only["non_tensor_batch"] == {} + assert meta_only["meta_info"] == {"roll_row_ids": row_ids} + assert np.array_equal( + full_tail["batch"]["attention_mask"], attention_mask[batch_size - 8 :] + ) + assert np.array_equal(full_tail["batch"]["advantages"], advantages[batch_size - 8 :]) + finally: + transfer.cleanup_dataproto(ref) + + +def test_dataproto_helper_selection_errors_and_destinations_with_real_store() -> None: + pytest.importorskip("mooncake.store") + _store, transfer = real_transfer("structured-test-selection-destinations") + batch_size = 12 + input_ids = np.arange(batch_size * 12, dtype=np.int64).reshape(batch_size, 12) + scores = np.linspace(0.0, 1.0, batch_size, dtype=np.float32) + prompts = np.asarray([f"prompt-{index}" for index in range(batch_size)], dtype=object) + ref = transfer.put_dataproto( + { + "batch": {"input_ids": input_ids, "scores": scores}, + "non_tensor_batch": {"prompts": prompts}, + "meta_info": {"step": 1}, + }, + namespace="roll", + partition="selection-destinations", + stage="rollout", + ) + + try: + with pytest.raises(ValueError, match="fields cannot be combined"): + transfer.get_dataproto(ref, fields=["input_ids"], batch_fields=["scores"]) + with pytest.raises(KeyError, match="unknown DataProto fields"): + transfer.get_dataproto(ref, fields=["missing"]) + with pytest.raises(IndexError, match="out of range"): + transfer.get_dataproto(ref, fields=["input_ids"], rows=[batch_size]) + with pytest.raises(TypeError, match="row indices"): + transfer.get_dataproto(ref, fields=["input_ids"], rows=["0"]) + with pytest.raises(ValueError, match="step must be positive"): + transfer.get_dataproto(ref, fields=["input_ids"], rows=slice(None, None, -1)) + with pytest.raises(ValueError, match="destination shape mismatch"): + transfer.get_dataproto( + ref, + batch_fields=["input_ids"], + rows=[0, 1, 2], + destinations={"input_ids": np.empty((2, 12), dtype=np.int64)}, + ) + + scores_destination = np.empty((4,), dtype=np.float32) + result = transfer.get_dataproto( + ref, + batch_fields=["scores"], + rows=[11, 3, 3, 0], + destinations={"scores": scores_destination}, + ) + non_tensor_only = transfer.get_dataproto( + ref, + non_tensor_fields=["prompts"], + rows=[2, 4, 6], + ) + + assert result["batch"]["scores"] is scores_destination + assert np.array_equal(scores_destination, scores[[11, 3, 3, 0]]) + assert result["non_tensor_batch"] == {} + assert non_tensor_only["batch"] == {} + assert non_tensor_only["non_tensor_batch"]["prompts"].tolist() == prompts[[2, 4, 6]].tolist() + finally: + transfer.cleanup_dataproto(ref) + + +def test_dataproto_helper_imported_handle_same_stage_append_with_real_store() -> None: + pytest.importorskip("mooncake.store") + _store, transfer = real_transfer("structured-test-imported-handle-append") + batch_size = 8 + input_ids = np.arange(batch_size * 6, dtype=np.int64).reshape(batch_size, 6) + masks = np.ones((batch_size, 6), dtype=np.int32) + values = np.linspace(1.0, 2.0, batch_size, dtype=np.float32) + rewards = values + np.float32(5.0) + ref = transfer.put_dataproto( + { + "batch": {"input_ids": input_ids, "masks": masks}, + "non_tensor_batch": {}, + "meta_info": {"step": 2}, + }, + namespace="roll", + partition="imported-handle-append", + stage="rollout", + ) + imported = import_dataproto_ref(export_dataproto_ref(ref)) + appended = transfer.append_dataproto_fields( + imported, + { + "batch": {"values": values, "rewards": rewards}, + "non_tensor_batch": {}, + "meta_info": {"critic": True}, + }, + stage="rollout", + ) + handle = export_dataproto_ref(appended) + + try: + result = transfer.get_dataproto(handle, fields=["input_ids", "values", "rewards"], rows=[7, 1, 4]) + assert np.array_equal(result["batch"]["input_ids"], input_ids[[7, 1, 4]]) + assert np.array_equal(result["batch"]["values"], values[[7, 1, 4]]) + assert np.array_equal(result["batch"]["rewards"], rewards[[7, 1, 4]]) + assert result["meta_info"] == {"step": 2, "critic": True} + finally: + transfer.cleanup_dataproto(appended) + + +def test_bundle_manifest_rejects_tampered_cleanup_keys() -> None: + store, transfer = make_transfer() + ref = transfer.put_bundle(b"meta", {"payload": b"abcdef"}) + tampered = dict(ref.manifest) + tampered["cleanup_keys"] = ["other/object"] + write_manifest(store, ref.manifest_key, tampered) + + with pytest.raises(ValueError, match="cleanup_keys"): + transfer.remove_bundle({"manifest_key": ref.manifest_key}) + + +def test_dataproto_helper_rollout_edge_cases_with_real_store() -> None: + pytest.importorskip("mooncake.store") + _store, transfer = real_transfer("structured-test-rollout-edge-cases") + batch_size = 10 + row_ids = [f"edge-row-{index}" for index in range(batch_size)] + input_ids = np.arange(batch_size * 16, dtype=np.int64).reshape(batch_size, 16) + attention_mask = (input_ids % 2 == 0).astype(np.int32) + prompts = np.asarray( + ["", "short", None, "三", "bytes", "long-" + "x" * 64, "7", "8", "9", "tail"], + dtype=object, + ) + scores = np.linspace(-1.0, 1.0, batch_size, dtype=np.float32) + values = scores + np.float32(2.0) + replacement_scores = scores + np.float32(10.0) + ref = transfer.put_dataproto( + { + "batch": {"input_ids": input_ids, "attention_mask": attention_mask}, + "non_tensor_batch": {"prompts": prompts}, + "meta_info": {"roll_row_ids": row_ids}, + }, + namespace="roll", + partition="edge-remote-batch", + stage="rollout", + ) + ref = transfer.append_dataproto_fields( + ref, + {"batch": {"scores": scores}, "non_tensor_batch": {}, "meta_info": {}}, + stage="scores", + ) + + try: + with pytest.raises(ValueError, match="already exist"): + transfer.append_dataproto_fields( + ref, + {"batch": {"scores": scores}, "non_tensor_batch": {}, "meta_info": {}}, + stage="other_scores", + ) + with pytest.raises(ValueError, match="batch size"): + transfer.append_dataproto_fields( + ref, + { + "batch": {"bad": np.arange(batch_size - 1, dtype=np.int64)}, + "non_tensor_batch": {}, + "meta_info": {}, + }, + stage="bad_size", + ) + + empty = transfer.get_dataproto( + ref, + fields=["input_ids", "scores", "prompts"], + rows=[], + ) + tail_and_repeat = transfer.get_dataproto( + ref, + fields=["input_ids", "attention_mask", "scores", "prompts"], + rows=[-1, 0, -1, 3], + ) + step_slice = transfer.get_dataproto( + ref, + batch_fields=["input_ids", "scores"], + rows=slice(1, 9, 2), + ) + ref = transfer.append_dataproto_fields( + ref, + { + "batch": {"values": values}, + "non_tensor_batch": {}, + "meta_info": {"same_stage": True}, + }, + stage="rollout", + ) + same_stage = transfer.get_dataproto( + ref, + fields=["input_ids", "values", "prompts"], + rows=[2, 5, 9], + ) + ref = transfer.append_dataproto_fields( + ref, + { + "batch": {"scores": replacement_scores}, + "non_tensor_batch": {}, + "meta_info": {"scores_overwritten": True}, + }, + stage="scores", + overwrite=True, + ) + overwritten = transfer.get_dataproto( + ref, + fields=["scores", "values"], + rows=[0, 4, 9], + ) + + assert empty["batch"]["input_ids"].shape == (0, 16) + assert empty["batch"]["scores"].shape == (0,) + assert empty["non_tensor_batch"]["prompts"].tolist() == [] + assert np.array_equal(tail_and_repeat["batch"]["input_ids"], input_ids[[9, 0, 9, 3]]) + assert np.array_equal( + tail_and_repeat["batch"]["attention_mask"], attention_mask[[9, 0, 9, 3]] + ) + assert np.array_equal(tail_and_repeat["batch"]["scores"], scores[[9, 0, 9, 3]]) + assert tail_and_repeat["non_tensor_batch"]["prompts"].tolist() == prompts[[9, 0, 9, 3]].tolist() + assert np.array_equal(step_slice["batch"]["input_ids"], input_ids[1:9:2]) + assert np.array_equal(step_slice["batch"]["scores"], scores[1:9:2]) + assert np.array_equal(same_stage["batch"]["input_ids"], input_ids[[2, 5, 9]]) + assert np.array_equal(same_stage["batch"]["values"], values[[2, 5, 9]]) + assert same_stage["non_tensor_batch"]["prompts"].tolist() == prompts[[2, 5, 9]].tolist() + assert np.array_equal(overwritten["batch"]["scores"], replacement_scores[[0, 4, 9]]) + assert np.array_equal(overwritten["batch"]["values"], values[[0, 4, 9]]) + assert overwritten["meta_info"]["same_stage"] is True + assert overwritten["meta_info"]["scores_overwritten"] is True + finally: + transfer.cleanup_dataproto(ref) + + with pytest.raises(Exception): + transfer.get_dataproto(ref, fields=["input_ids"]) + + +def test_structured_object_zero_byte_payload_skips_store_put() -> None: + store, transfer = make_transfer() + ref = transfer.put_structured_object( + StructuredObjectPayload(buffers={"empty": np.empty((4, 0), dtype=np.float32)}) + ) + result = transfer.materialize(transfer.read_spec(ref)) + + assert result.objects["empty"].shape == (4, 0) + assert ref.manifest["buffers"]["empty"]["bytes"] == 0 + assert ref.manifest["buffers"]["empty"]["chunks"] == [] + assert store.objects == {ref.manifest_key: store.objects[ref.manifest_key], ref.manifest["meta"]["key"]: store.objects[ref.manifest["meta"]["key"]]} + + +def test_dataproto_helper_multidim_boundary_reads_with_real_store() -> None: + pytest.importorskip("mooncake.store") + _store, transfer = real_transfer("structured-test-multidim-boundaries") + batch_size = 18 + image_features = np.arange(batch_size * 3 * 16 * 16, dtype=np.float32).reshape( + batch_size, 3, 16, 16 + ) + logits = np.linspace(-2.0, 2.0, batch_size * 4 * 5 * 6, dtype=np.float32).reshape( + batch_size, 4, 5, 6 + ) + token_matrix = np.arange(batch_size * 7 * 9, dtype=np.int64).reshape( + batch_size, 7, 9 + ) + zero_width = np.empty((batch_size, 0), dtype=np.float32) + row_scores = np.linspace(10.0, 20.0, batch_size, dtype=np.float32) + byte_blobs = np.asarray( + [ + b"", + b"alpha", + bytearray(b"beta"), + memoryview(b"gamma"), + *[f"blob-{index}".encode() for index in range(4, batch_size)], + ], + dtype=object, + ) + json_meta = np.asarray( + [ + None if index % 5 == 0 else {"row": index, "shape": [3, 16, 16]} + for index in range(batch_size) + ], + dtype=object, + ) + ref = transfer.put_dataproto( + { + "batch": { + "image_features": image_features, + "logits": logits, + "token_matrix": token_matrix, + "zero_width": zero_width, + }, + "non_tensor_batch": { + "byte_blobs": byte_blobs, + "json_meta": json_meta, + }, + "meta_info": {"batch_size": batch_size, "layout": "multidim"}, + }, + namespace="roll", + partition="multidim-boundaries", + stage="rollout", + ) + ref = transfer.append_dataproto_fields( + ref, + { + "batch": {"row_scores": row_scores}, + "non_tensor_batch": {}, + "meta_info": {"score_stage": "same-stage"}, + }, + stage="rollout", + ) + handle = export_dataproto_ref(ref) + rows = [17, 0, 5, 17, 9] + + try: + full = transfer.get_dataproto( + handle, + fields=["image_features", "zero_width", "byte_blobs", "json_meta"], + ) + empty_slice = transfer.get_dataproto( + handle, + batch_fields=["image_features", "zero_width", "row_scores"], + rows=slice(4, 4), + ) + duplicate_gather = transfer.get_dataproto( + handle, + fields=["logits", "token_matrix", "row_scores", "byte_blobs", "json_meta"], + rows=rows, + ) + image_destination = np.empty((len(rows), 3, 16, 16), dtype=np.float32) + into = transfer.get_dataproto( + handle, + batch_fields=["image_features"], + rows=rows, + destinations={"image_features": image_destination}, + ) + step_batch = transfer.get_dataproto( + handle, + batch_fields=["logits", "token_matrix"], + rows=slice(2, 17, 3), + ) + tail_non_tensor = transfer.get_dataproto( + handle, + non_tensor_fields=["byte_blobs", "json_meta"], + rows=slice(batch_size - 3, batch_size), + ) + meta_selected = transfer.get_dataproto( + handle, + fields=[], + meta_info_keys=["layout", "score_stage"], + ) + + assert np.array_equal(full["batch"]["image_features"], image_features) + assert full["batch"]["zero_width"].shape == (batch_size, 0) + assert full["non_tensor_batch"]["byte_blobs"].tolist() == [ + bytes(value) for value in byte_blobs.tolist() + ] + assert full["non_tensor_batch"]["json_meta"].tolist() == json_meta.tolist() + assert empty_slice["batch"]["image_features"].shape == (0, 3, 16, 16) + assert empty_slice["batch"]["zero_width"].shape == (0, 0) + assert empty_slice["batch"]["row_scores"].shape == (0,) + assert np.array_equal(duplicate_gather["batch"]["logits"], logits[rows]) + assert np.array_equal( + duplicate_gather["batch"]["token_matrix"], token_matrix[rows] + ) + assert np.array_equal(duplicate_gather["batch"]["row_scores"], row_scores[rows]) + assert duplicate_gather["non_tensor_batch"]["byte_blobs"].tolist() == [ + bytes(value) for value in byte_blobs[rows].tolist() + ] + assert duplicate_gather["non_tensor_batch"]["json_meta"].tolist() == json_meta[rows].tolist() + assert into["batch"]["image_features"] is image_destination + assert np.array_equal(image_destination, image_features[rows]) + assert np.array_equal(step_batch["batch"]["logits"], logits[2:17:3]) + assert np.array_equal(step_batch["batch"]["token_matrix"], token_matrix[2:17:3]) + assert tail_non_tensor["batch"] == {} + assert tail_non_tensor["non_tensor_batch"]["byte_blobs"].tolist() == [ + bytes(value) for value in byte_blobs[batch_size - 3 :].tolist() + ] + assert tail_non_tensor["non_tensor_batch"]["json_meta"].tolist() == json_meta[batch_size - 3 :].tolist() + assert meta_selected["batch"] == {} + assert meta_selected["non_tensor_batch"] == {} + assert meta_selected["meta_info"] == { + "layout": "multidim", + "score_stage": "same-stage", + } + finally: + transfer.cleanup_dataproto(ref) + + +def test_dataproto_helper_rejects_inconsistent_batch_sizes() -> None: + _store, transfer = make_transfer() + data = SimpleDataProto( + batch={ + "input_ids": np.arange(4, dtype=np.int64), + "attention_mask": np.arange(3, dtype=np.int64), + } + ) + + with pytest.raises(ValueError, match="inconsistent batch sizes"): + transfer.put_dataproto(data) + + +def test_dataproto_helper_overwrite_rejects_partial_stage_replacement() -> None: + _store, transfer = make_transfer() + ref = transfer.put_dataproto( + SimpleDataProto( + batch={"input_ids": np.arange(4), "attention_mask": np.arange(4)} + ), + stage="rollout", + ) + + with pytest.raises(ValueError, match="must include existing fields"): + transfer.append_dataproto_fields( + ref, + SimpleDataProto(batch={"input_ids": np.arange(4)}), + stage="rollout", + overwrite=True, + ) + + +def test_dataproto_helper_overwrite_replaces_encoded_metadata() -> None: + store, transfer = make_transfer() + ref = transfer.put_dataproto( + SimpleDataProto( + non_tensor_batch={"text": np.asarray(["a", None], dtype=object)} + ), + stage="meta", + ) + old_manifest_key = ref.stage_refs["meta"].manifest_key + ref = transfer.append_dataproto_fields( + ref, + SimpleDataProto(non_tensor_batch={"text": np.asarray([1, 2], dtype=np.int64)}), + stage="meta", + overwrite=True, + ) + result = transfer.get_dataproto(ref) + + assert "text" not in ref.encoded_non_tensor + assert old_manifest_key not in store.objects + assert ref.stage_refs["meta"].manifest_key in store.objects + assert np.array_equal(result["non_tensor_batch"]["text"], np.asarray([1, 2])) + + +def test_dataproto_helper_cleanup_removes_all_stage_objects() -> None: + store, transfer = make_transfer() + ref = transfer.put_dataproto( + SimpleDataProto(batch={"input_ids": np.arange(4, dtype=np.int64)}), + stage="rollout", + ) + ref = transfer.append_dataproto_fields( + ref, + SimpleDataProto(batch={"values": np.arange(4, dtype=np.float32)}), + stage="critic", + ) + + assert store.objects + transfer.cleanup_dataproto(ref) + + assert store.objects == {} + + +def test_dataproto_helper_object_non_tensor_codecs_roundtrip() -> None: + _store, transfer = make_transfer() + data = SimpleDataProto( + batch={"input_ids": np.arange(4, dtype=np.int64)}, + non_tensor_batch={ + "text": np.asarray(["hello", None, "world", "moon"], dtype=object), + "json": np.asarray( + [{"a": 1}, None, {"b": [2, 3]}, {"c": "x"}], dtype=object + ), + "blob": np.asarray( + [b"a", None, bytearray(b"bc"), memoryview(b"def")], dtype=object + ), + "nullable_int": np.asarray([1, None, 3, 4], dtype=object), + }, + ) + + ref = transfer.put_dataproto(data) + result = transfer.get_dataproto(ref) + + assert set(ref.encoded_non_tensor) == {"text", "json", "blob", "nullable_int"} + assert ref.encoded_non_tensor["json"]["codec"] == "msgpack_ragged" + assert result["non_tensor_batch"]["text"].tolist() == [ + "hello", + None, + "world", + "moon", + ] + assert result["non_tensor_batch"]["json"].tolist() == [ + {"a": 1}, + None, + {"b": [2, 3]}, + {"c": "x"}, + ] + assert result["non_tensor_batch"]["blob"].tolist() == [b"a", None, b"bc", b"def"] + assert result["non_tensor_batch"]["nullable_int"].tolist() == [1, None, 3, 4] + + +def test_msgpack_ragged_decode_validates_row_count() -> None: + payload, _metadata = sos._encode_msgpack_ragged_values( + "json", [{"a": 1}, {"b": 2}] + ) + + with pytest.raises(ValueError, match="offsets length"): + sos._decode_msgpack_ragged_values(payload, 1) + + +def test_dataproto_helper_typed_ragged_uses_multi_buffer_put() -> None: + pool = FakeBufferPool() + _store, transfer = make_transfer(buffer_pool=pool) + rows = np.empty(3, dtype=object) + rows[:] = [ + np.asarray([1, 2], dtype=np.int32), + None, + np.asarray([3, 4, 5], dtype=np.int32), + ] + data = SimpleDataProto(non_tensor_batch={"tokens": rows}) + + ref = transfer.put_dataproto(data) + result = transfer.get_dataproto(ref) + + assert ref.encoded_non_tensor["tokens"]["codec"] == "typed_ragged" + tokens = result["non_tensor_batch"]["tokens"].tolist() + assert tokens == [[1, 2], None, [3, 4, 5]] + assert rows[0].nbytes + rows[2].nbytes in pool.acquire_sizes + + +def test_typed_ragged_uses_direct_copy_payload_when_fast_copy_available() -> None: + if sos._concat_arrays_into is None: + pytest.skip("native fast-copy extension is unavailable") + rows = [ + np.arange(2, dtype=np.int32), + np.arange(3, dtype=np.int32), + ] + + payload, metadata = sos._encode_typed_ragged_values(rows, dtype_hint=np.int32) + + assert metadata["shape_policy"] == "ragged" + assert isinstance(payload["data"], sos._DirectCopyPayload) + assert payload["data"].shape == (5,) + + +def test_dataproto_helper_typed_ragged_fast_copy_put() -> None: + pool = FakeBufferPool() + store, transfer = make_transfer(buffer_pool=pool) + rows = np.empty(3, dtype=object) + rows[:] = [ + np.asarray([1, 2], dtype=np.int32), + None, + np.asarray([3, 4, 5], dtype=np.int32), + ] + data = SimpleDataProto(non_tensor_batch={"tokens": rows}) + + ref = transfer.put_dataproto(data) + result = transfer.get_dataproto(ref) + + tokens = result["non_tensor_batch"]["tokens"].tolist() + assert tokens == [[1, 2], None, [3, 4, 5]] + assert store.batch_put_from_calls > 0 + assert rows[0].nbytes + rows[2].nbytes in pool.acquire_sizes + + +def test_direct_copy_put_rolls_back_all_chunks_on_late_failure() -> None: + if sos._concat_arrays_into is None: + pytest.skip("native fast-copy extension is unavailable") + store = FailingBatchPutStore(fail_on_call=2) + pool = FakeBufferPool() + _, transfer = make_transfer(store=store, buffer_pool=pool) + arrays = [ + np.asarray([1, 2], dtype=np.int32), + np.asarray([3, 4], dtype=np.int32), + np.asarray([5, 6], dtype=np.int32), + ] + payload = sos._DirectCopyPayload.from_flat_arrays( + arrays, np.dtype(np.int32), total_elems=6 + ) + + with pytest.raises(RuntimeError, match="batch_put_from"): + transfer._bundle_store._put_direct_copy_payload( + "payload", + payload, + chunk_bytes=8, + transfer_policy=BundleTransferPolicy(), + ) + + assert store.objects == {} + assert pool.acquire_count == pool.release_count == 2 + + +def test_dataproto_helper_typed_ragged_rejects_short_fast_copy(monkeypatch) -> None: + pool = FakeBufferPool() + _store, transfer = make_transfer(buffer_pool=pool) + rows = np.empty(2, dtype=object) + rows[:] = [np.asarray([1, 2], dtype=np.int32), np.asarray([3], dtype=np.int32)] + + monkeypatch.setattr(sos, "_concat_arrays_into", lambda *args: args[2] - 1) + + with pytest.raises(RuntimeError, match="native fast-copy wrote"): + transfer.put_dataproto(SimpleDataProto(non_tensor_batch={"tokens": rows})) + + +def test_dataproto_helper_typed_ragged_zero_copy_rejects_source_buffers() -> None: + _store, transfer = make_transfer(buffer_pool=FakeBufferPool()) + rows = np.empty(2, dtype=object) + rows[:] = [ + np.asarray([1, 2], dtype=np.int32), + np.asarray([3], dtype=np.int32), + ] + + with pytest.raises(RuntimeError, match="zero-copy put"): + transfer.put_dataproto( + SimpleDataProto(non_tensor_batch={"tokens": rows}), + policy=BundleTransferPolicy(copy_mode="zero_copy"), + ) + + +def test_dataproto_helper_rejects_unsupported_object_non_tensor() -> None: + _store, transfer = make_transfer() + data = SimpleDataProto( + non_tensor_batch={"fallback": np.asarray([object(), None], dtype=object)} + ) + + with pytest.raises(ValueError, match="unsupported structured non-tensor field"): + transfer.put_dataproto(data) + with pytest.raises(ValueError, match="unable to infer a safe codec"): + sos._encode_with_schema( + "non_tensor_batch.fallback", + data.non_tensor_batch["fallback"], + FieldSchema(codec="auto"), + ) + + +def test_dataproto_helper_ragged_tensor_non_tensor_roundtrip() -> None: + torch = pytest.importorskip("torch") + _store, transfer = make_transfer() + ragged = np.asarray( + [ + torch.arange(1, dtype=torch.float32), + None, + torch.arange(3, dtype=torch.float32), + torch.arange(2, dtype=torch.float32), + ], + dtype=object, + ) + data = SimpleDataProto(non_tensor_batch={"ragged": ragged}) + + ref = transfer.put_dataproto(data) + result = transfer.get_dataproto(ref) + + assert ref.encoded_non_tensor["ragged"]["codec"] == "ragged_tensor" + actual = result["non_tensor_batch"]["ragged"] + assert torch.equal(actual[0], ragged[0]) + assert actual[1] is None + assert torch.equal(actual[2], ragged[2]) + assert torch.equal(actual[3], ragged[3]) + + mixed = np.empty(2, dtype=object) + mixed[:] = [ + torch.arange(1, dtype=torch.float32), + torch.arange(1, dtype=torch.int64), + ] + with pytest.raises(ValueError, match="mixed tensor dtype"): + transfer.put_dataproto(SimpleDataProto(non_tensor_batch={"ragged": mixed})) + + +def _assert_tensor_object_equal(actual, expected) -> None: + torch = pytest.importorskip("torch") + if expected is None: + assert actual is None + return + if isinstance(expected, dict): + assert isinstance(actual, dict) + assert set(actual) == set(expected) + for key, value in expected.items(): + _assert_tensor_object_equal(actual[key], value) + return + if isinstance(expected, list): + assert isinstance(actual, list) + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected): + _assert_tensor_object_equal(actual_item, expected_item) + return + if isinstance(expected, torch.Tensor): + assert torch.equal(actual, expected) + return + assert actual == expected + + +def test_dataproto_helper_dict_of_tensors_object_array_roundtrip() -> None: + torch = pytest.importorskip("torch") + _store, transfer = make_transfer() + samples = np.asarray( + [ + {"tokens": torch.arange(2, dtype=torch.int64), "reward": 1.0}, + {"tokens": torch.arange(3, dtype=torch.int64), "reward": 2.0}, + {"tokens": torch.arange(1, dtype=torch.int64), "reward": None}, + {"tokens": None, "reward": 4.0}, + ], + dtype=object, + ) + + ref = transfer.put_dataproto( + SimpleDataProto( + batch={"input_ids": np.arange(8, dtype=np.int64).reshape(4, 2)}, + non_tensor_batch={"samples": samples}, + meta_info={"source": "dict-of-tensors"}, + ) + ) + result = transfer.get_dataproto(ref) + view = transfer.dataproto_manifest_view(ref) + + assert ref.encoded_non_tensor["samples"]["codec"] == "structured_recursive" + assert view["non_tensor_fields"]["samples"]["spec"]["codec"] == "structured_recursive" + assert result["meta_info"] == {"source": "dict-of-tensors"} + actual = result["non_tensor_batch"]["samples"] + for row, expected in enumerate(samples): + _assert_tensor_object_equal(actual[row], expected) + + +def test_dataproto_helper_dict_of_tensors_distinguishes_missing_keys_and_nulls() -> None: + torch = pytest.importorskip("torch") + _store, transfer = make_transfer() + samples = np.asarray( + [ + {"tokens": torch.arange(2), "label": None}, + {"label": 1}, + None, + {"tokens": None, "label": 3}, + ], + dtype=object, + ) + + ref = transfer.put_dataproto(SimpleDataProto(non_tensor_batch={"samples": samples})) + actual = transfer.get_dataproto(ref)["non_tensor_batch"]["samples"] + + assert "tokens" in actual[0] + assert actual[0]["label"] is None + assert "tokens" not in actual[1] + assert actual[1]["label"] == 1 + assert actual[2] is None + assert "tokens" in actual[3] + assert actual[3]["tokens"] is None + assert actual[3]["label"] == 3 + + +def test_dataproto_helper_nested_tensor_object_array_rows() -> None: + torch = pytest.importorskip("torch") + _store, transfer = make_transfer() + samples = np.asarray( + [ + {"images": [{"pixels": torch.arange(2)}], "meta": {"rank": 0}}, + {"images": [{"pixels": torch.arange(3)}, {"pixels": None}], "meta": {}}, + {"images": [], "meta": {"rank": None}}, + {"images": [], "meta": None}, + ], + dtype=object, + ) + + ref = transfer.put_dataproto(SimpleDataProto(non_tensor_batch={"samples": samples})) + actual = transfer.get_dataproto(ref)["non_tensor_batch"]["samples"] + + for row, expected in enumerate(samples): + _assert_tensor_object_equal(actual[row], expected) + + +def test_dataproto_helper_dict_of_native_object_leaves_uses_recursive_codec() -> None: + _store, transfer = make_transfer() + samples = np.asarray( + [ + { + "media": [b"a", b"bc"], + "blob": b"payload-0", + "scores": np.asarray([1, 2], dtype=np.int64), + "label": "x", + }, + { + "media": [], + "blob": b"payload-1", + "scores": np.asarray([3], dtype=np.int64), + "label": "y", + }, + {"label": "missing-native"}, + None, + ], + dtype=object, + ) + + ref = transfer.put_dataproto(SimpleDataProto(non_tensor_batch={"samples": samples})) + result = transfer.get_dataproto(ref) + actual = result["non_tensor_batch"]["samples"] + + assert ref.encoded_non_tensor["samples"]["codec"] == "structured_recursive" + assert actual[0]["media"] == [b"a", b"bc"] + assert actual[0]["blob"] == b"payload-0" + assert actual[0]["scores"] == [1, 2] + assert actual[0]["label"] == "x" + assert actual[1]["media"] == [] + assert actual[1]["blob"] == b"payload-1" + assert actual[1]["scores"] == [3] + assert actual[1]["label"] == "y" + assert actual[2] == {"label": "missing-native"} + assert actual[3] is None + + +def test_dataproto_helper_reads_dict_of_tensors_rows_and_selected_field() -> None: + torch = pytest.importorskip("torch") + _store, transfer = make_transfer() + samples = np.asarray( + [ + {"tokens": torch.arange(i + 1, dtype=torch.int64), "row": i} + for i in range(6) + ], + dtype=object, + ) + samples[2] = {"tokens": None, "row": 2} + samples[4] = None + input_ids = np.arange(12, dtype=np.int64).reshape(6, 2) + ref = transfer.put_dataproto( + SimpleDataProto( + batch={"input_ids": input_ids}, + non_tensor_batch={"samples": samples}, + meta_info={"kind": "dict-tensor"}, + ) + ) + + sliced = transfer.get_dataproto(ref, fields=["samples"], rows=slice(1, 5)) + gathered = transfer.get_dataproto(ref, fields=["input_ids", "samples"], rows=[5, 0, 2, 4]) + + assert sliced["batch"] == {} + assert sliced["meta_info"] == {"kind": "dict-tensor"} + for row, expected in enumerate(samples[1:5]): + _assert_tensor_object_equal(sliced["non_tensor_batch"]["samples"][row], expected) + assert np.array_equal(gathered["batch"]["input_ids"], input_ids[[5, 0, 2, 4]]) + for row, expected_index in enumerate([5, 0, 2, 4]): + _assert_tensor_object_equal( + gathered["non_tensor_batch"]["samples"][row], samples[expected_index] + ) + + +def test_dataproto_helper_recursive_manifest_export_append_and_overwrite() -> None: + torch = pytest.importorskip("torch") + store, transfer = make_transfer() + input_ids = np.arange(4, dtype=np.int64) + ref = transfer.put_dataproto(SimpleDataProto(batch={"input_ids": input_ids}), stage="rollout") + samples = np.asarray( + [ + {"tokens": torch.arange(1, dtype=torch.float32), "score": 0.0}, + {"tokens": torch.arange(2, dtype=torch.float32), "score": 1.0}, + {"tokens": torch.arange(3, dtype=torch.float32), "score": None}, + {"tokens": None, "score": 3.0}, + ], + dtype=object, + ) + + ref = transfer.append_dataproto_fields( + ref, + SimpleDataProto(non_tensor_batch={"samples": samples}, meta_info={"stage": "samples"}), + stage="rollout", + ) + handle = export_dataproto_ref(ref) + json.dumps(handle) + imported = import_dataproto_ref(handle) + result = transfer.get_dataproto(imported) + view = transfer.dataproto_manifest_view(imported) + + assert np.array_equal(result["batch"]["input_ids"], input_ids) + _assert_tensor_object_equal(result["non_tensor_batch"]["samples"][1], samples[1]) + spec = view["non_tensor_fields"]["samples"]["spec"] + assert spec["codec"] == "structured_recursive" + assert spec["metadata"]["nodes"] + assert spec["metadata"]["leaves"] + assert any(name.endswith(".missing") for name in spec["payload_members"]) + + old_manifest_key = ref.stage_refs["rollout"].manifest_key + ref = transfer.append_dataproto_fields( + ref, + SimpleDataProto( + batch={"input_ids": input_ids}, + non_tensor_batch={"samples": np.asarray(["a", "b", "c", "d"], dtype=object)}, + ), + stage="rollout", + overwrite=True, + ) + overwritten = transfer.get_dataproto(ref) + assert ref.encoded_non_tensor["samples"]["codec"] == "utf8_ragged" + assert overwritten["non_tensor_batch"]["samples"].tolist() == ["a", "b", "c", "d"] + assert old_manifest_key not in store.objects diff --git a/mooncake-wheel/tests/test_transfer_on_cuda.py b/mooncake-wheel/tests/test_transfer_on_cuda.py index 9cd04028c2..c5af5f2efe 100644 --- a/mooncake-wheel/tests/test_transfer_on_cuda.py +++ b/mooncake-wheel/tests/test_transfer_on_cuda.py @@ -1,4 +1,3 @@ -import os import socket import unittest import torch diff --git a/requirements-dev.txt b/requirements.txt similarity index 83% rename from requirements-dev.txt rename to requirements.txt index a7e38e18ce..11f5e714c9 100644 --- a/requirements-dev.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ # Development dependencies for Mooncake -# Install with: pip install -r requirements-dev.txt +# Install with: pip install -r requirements.txt # Then enable hooks: pre-commit install pre-commit==3.7.1 ruff==0.6.9 diff --git a/scripts/ascend/dependencies_ascend.sh b/scripts/ascend/dependencies_ascend.sh index 101aaaf49a..357b490391 100644 --- a/scripts/ascend/dependencies_ascend.sh +++ b/scripts/ascend/dependencies_ascend.sh @@ -55,7 +55,6 @@ if command -v apt-get &> /dev/null; then wget \ libibverbs-dev \ libgoogle-glog-dev \ - libgtest-dev \ libjsoncpp-dev \ libunwind-dev \ libnuma-dev \ @@ -83,8 +82,6 @@ elif command -v yum &> /dev/null; then glog-devel \ libibverbs-devel \ numactl-devel \ - gtest \ - gtest-devel \ boost-devel \ openssl-devel \ hiredis-devel \ diff --git a/scripts/ascend/dependencies_ascend_installation.sh b/scripts/ascend/dependencies_ascend_installation.sh index eb9266dc1c..ff5944bb10 100644 --- a/scripts/ascend/dependencies_ascend_installation.sh +++ b/scripts/ascend/dependencies_ascend_installation.sh @@ -65,7 +65,6 @@ if command -v apt-get &> /dev/null; then wget \ libibverbs-dev \ libgoogle-glog-dev \ - libgtest-dev \ libjsoncpp-dev \ libunwind-dev \ libnuma-dev \ @@ -95,8 +94,6 @@ elif command -v yum &> /dev/null; then glog-devel \ libibverbs-devel \ numactl-devel \ - gtest \ - gtest-devel \ boost-devel \ openssl-devel \ hiredis-devel \ diff --git a/scripts/ascend/dependencies_openeuler.sh b/scripts/ascend/dependencies_openeuler.sh index 0d6b83f957..e2016395c8 100644 --- a/scripts/ascend/dependencies_openeuler.sh +++ b/scripts/ascend/dependencies_openeuler.sh @@ -151,7 +151,7 @@ PKG_MGR="$(detect_pkg_manager)" CORE_PACKAGES=( gcc gcc-c++ make cmake ninja-build git wget unzip gflags-devel glog-devel libibverbs-devel numactl-devel - gtest gtest-devel boost-devel openssl-devel hiredis-devel + boost-devel openssl-devel hiredis-devel libcurl-devel jsoncpp-devel libunwind-devel python3-devel zstd-devel xxhash-devel pkgconf pkgconf-pkg-config patchelf mpich mpich-devel glibc glibc-common diff --git a/scripts/ascend/perf/llmdatadist_bandwidth_test_cross_machine_demo.py b/scripts/ascend/perf/llmdatadist_bandwidth_test_cross_machine_demo.py index 6b73acbe26..7cc1d0909d 100644 --- a/scripts/ascend/perf/llmdatadist_bandwidth_test_cross_machine_demo.py +++ b/scripts/ascend/perf/llmdatadist_bandwidth_test_cross_machine_demo.py @@ -29,9 +29,13 @@ import time from llm_datadist import LLMDataDist, LLMRole, LLMConfig, CacheDesc, Cache, DataType, RegisterMemStatus, BlocksCacheKey, \ Placement +import importlib + import torch -import torch_npu -import torchair + +# Side-effect imports: register Ascend NPU backends used by tensor.npu(). +importlib.import_module("torch_npu") +importlib.import_module("torchair") PROMPT_IP_LIST = ['192.168.1.1', '192.168.1.2', '192.168.1.3', '192.168.1.4', '192.168.1.5', '192.168.1.6', '192.168.1.7', '192.168.1.8'] @@ -168,7 +172,7 @@ def run_prompt_sample(datadist, device_id: int): placement=Placement.DEVICE) tensor = torch.ones(BLOCK_NUM, BLOCK_SIZE // 4, dtype=torch.float).npu() addr = int(tensor.data_ptr()) - cache = cache_manager.register_blocks_cache(cache_desc, [addr], BlocksCacheKey(1, 0)) + cache_manager.register_blocks_cache(cache_desc, [addr], BlocksCacheKey(1, 0)) logging.info('[register_blocks_cache] success') comm_id = link(datadist, device_id, TARGRT_DEVICE_ID) diff --git a/scripts/ascend/perf/llmdatadist_bandwidth_test_single_machine_demo.py b/scripts/ascend/perf/llmdatadist_bandwidth_test_single_machine_demo.py index 5701839c98..0139a6f7dd 100644 --- a/scripts/ascend/perf/llmdatadist_bandwidth_test_single_machine_demo.py +++ b/scripts/ascend/perf/llmdatadist_bandwidth_test_single_machine_demo.py @@ -29,9 +29,13 @@ import time from llm_datadist import LLMDataDist, LLMRole, LLMConfig, CacheDesc, Cache, DataType, RegisterMemStatus, BlocksCacheKey, \ Placement +import importlib + import torch -import torch_npu -import torchair + +# Side-effect imports: register Ascend NPU backends used by tensor.npu(). +importlib.import_module("torch_npu") +importlib.import_module("torchair") NPU_IP_LIST = ['192.168.1.1', '192.168.1.2', '192.168.1.3', '192.168.1.4', '192.168.1.5', '192.168.1.6', '192.168.1.7', '192.168.1.8'] @@ -160,7 +164,7 @@ def run_prompt_sample(datadist, device_id: int): placement=Placement.DEVICE) tensor = torch.ones(BLOCK_NUM, BLOCK_SIZE // 4, dtype=torch.float).npu() addr = int(tensor.data_ptr()) - cache = cache_manager.register_blocks_cache(cache_desc, [addr], BlocksCacheKey(1, 0)) + cache_manager.register_blocks_cache(cache_desc, [addr], BlocksCacheKey(1, 0)) logging.info('[register_blocks_cache] success') comm_id = link(datadist, device_id, TARGRT_DEVICE_ID) diff --git a/scripts/build_wheel.sh b/scripts/build_wheel.sh index dc414b367a..bfbcc3c3b1 100755 --- a/scripts/build_wheel.sh +++ b/scripts/build_wheel.sh @@ -17,7 +17,7 @@ BUILD_DIR_ABS="$(pwd)/${BUILD_DIR}" echo "Building wheel for Python ${PYTHON_VERSION} with output directory ${OUTPUT_DIR}" # Ensure LD_LIBRARY_PATH includes /usr/local/lib -export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${BUILD_DIR_ABS}/mooncake-common:/usr/local/lib +export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${BUILD_DIR_ABS}/mooncake-common:${BUILD_DIR_ABS}/mooncake-common/etcd:${BUILD_DIR_ABS}/mooncake-common/k8s-lease:/usr/local/lib echo "Cleaning wheel-build directory" rm -rf mooncake-wheel/mooncake_transfer_engine* @@ -32,6 +32,12 @@ cp mooncake-integration/fabric_allocator_utils.py mooncake-wheel/mooncake/fabric # Copy engine.so to mooncake directory (will be imported by transfer module) cp ${BUILD_DIR}/mooncake-integration/engine.*.so mooncake-wheel/mooncake/engine.so +# Copy the host-only PG core to mooncake directory +if [ -f "${BUILD_DIR}/mooncake-pg/src/libmooncake_pg.so" ]; then + echo "Copying libmooncake_pg.so..." + cp "${BUILD_DIR}/mooncake-pg/src/libmooncake_pg.so" mooncake-wheel/mooncake/libmooncake_pg.so +fi + # Copy libasio.so to mooncake directory (runtime dependency of engine.so) cp ${BUILD_DIR}/mooncake-common/libasio.so mooncake-wheel/mooncake/libasio.so @@ -68,6 +74,12 @@ if [ -f ${BUILD_DIR}/mooncake-common/etcd/libetcd_wrapper.so ]; then cp ${BUILD_DIR}/mooncake-common/etcd/libetcd_wrapper.so mooncake-wheel/mooncake/libetcd_wrapper.so fi +# Copy libk8s_lease_wrapper.so to mooncake directory (only when STORE_USE_K8S_LEASE is set) +if [ -f ${BUILD_DIR}/mooncake-common/k8s-lease/libk8s_lease_wrapper.so ]; then + echo "Copying libk8s_lease_wrapper.so..." + cp ${BUILD_DIR}/mooncake-common/k8s-lease/libk8s_lease_wrapper.so mooncake-wheel/mooncake/libk8s_lease_wrapper.so +fi + # Copy libtransfer_engine.so to mooncake directory (only when BUILD_SHARED_LIBS is set) if [ -f ${BUILD_DIR}/mooncake-transfer-engine/src/libtransfer_engine.so ]; then echo "Copying libtransfer_engine.so..." @@ -154,7 +166,20 @@ echo "Building wheel package..." # Build the wheel package cd mooncake-wheel -BUILD_VARIANTS="NON_CUDA_BUILD CU13_BUILD NPU_BUILD" +# Materialize a local copy of the root README.md so that +# `readme = "README.md"` resolves inside this directory. Modern +# setuptools rejects `../`-traversal in the readme path. The file is +# removed by the EXIT trap below. +cp ../README.md README.md + +WHEEL_DIR="$(pwd)" +cleanup_wheel_metadata_state() { + [[ -f "${WHEEL_DIR}/pyproject.toml.backup" ]] && mv "${WHEEL_DIR}/pyproject.toml.backup" "${WHEEL_DIR}/pyproject.toml" + rm -f "${WHEEL_DIR}/README.md" +} +trap cleanup_wheel_metadata_state EXIT + +BUILD_VARIANTS="NON_CUDA_BUILD CU13_BUILD NPU_BUILD EFA_BUILD EFA_CU13_BUILD EFA_NON_CUDA_BUILD MUSA_BUILD HIP_BUILD" BUILD_VARIANT_COUNT=0 for build_variant in $BUILD_VARIANTS; do if [ "${!build_variant}" = "1" ]; then @@ -166,6 +191,15 @@ if [ "$BUILD_VARIANT_COUNT" -gt 1 ]; then exit 1 fi +# If a previous run was interrupted before the trailing restore (line ~481), +# pyproject.toml is left in a renamed state and pyproject.toml.backup holds the +# pristine original. Restore it first so the variant rename below always starts +# from the clean file and the backup is never overwritten with modified content. +if [ -f pyproject.toml.backup ]; then + echo "Restoring pyproject.toml from leftover backup of a previous run" + mv pyproject.toml.backup pyproject.toml +fi + # Handle package name modification for release build variants if [ "$NON_CUDA_BUILD" = "1" ]; then echo "Modifying package name for non-CUDA build" @@ -173,8 +207,9 @@ if [ "$NON_CUDA_BUILD" = "1" ]; then cp pyproject.toml pyproject.toml.backup # Replace package name and description sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-non-cuda"/' pyproject.toml - sed -i 's/description = "Python binding of a Mooncake library using pybind11"/description = "Python binding of a Mooncake library using pybind11 (Non-CUDA version)"/' pyproject.toml - sed -i 's/keywords = \["mooncake", "data transfer", "kv cache", "llm inference"\]/keywords = ["mooncake", "data transfer", "kv cache", "llm inference", "non-cuda"]/' pyproject.toml + sed -i 's/^description = "\(.*\)"$/description = "\1 (Non-CUDA version)"/' pyproject.toml + sed -i 's/^keywords = \[\(.*\)\]$/keywords = [\1, "non-cuda"]/' pyproject.toml + sed -i 's|"Environment :: GPU :: NVIDIA CUDA", ||' pyproject.toml echo "Package name modified to: mooncake-transfer-engine-non-cuda" elif [ "$CU13_BUILD" = "1" ]; then echo "Modifying package name for CU13 build" @@ -182,8 +217,8 @@ elif [ "$CU13_BUILD" = "1" ]; then cp pyproject.toml pyproject.toml.backup # Replace package name and description sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-cuda13"/' pyproject.toml - sed -i 's/description = "Python binding of a Mooncake library using pybind11"/description = "Python binding of a Mooncake library using pybind11 (CUDA 13 version)"/' pyproject.toml - sed -i 's/keywords = \["mooncake", "data transfer", "kv cache", "llm inference"\]/keywords = ["mooncake", "data transfer", "kv cache", "llm inference", "cuda13"]/' pyproject.toml + sed -i 's/^description = "\(.*\)"$/description = "\1 (CUDA 13 version)"/' pyproject.toml + sed -i 's/^keywords = \[\(.*\)\]$/keywords = [\1, "cuda13"]/' pyproject.toml echo "Package name modified to: mooncake-transfer-engine-cuda13" elif [ "$NPU_BUILD" = "1" ]; then echo "Modifying package name for Ascend NPU build" @@ -191,9 +226,62 @@ elif [ "$NPU_BUILD" = "1" ]; then cp pyproject.toml pyproject.toml.backup # Replace package name and description sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-npu"/' pyproject.toml - sed -i 's/description = "Python binding of a Mooncake library using pybind11"/description = "Python binding of a Mooncake library using pybind11 (Ascend NPU version)"/' pyproject.toml - sed -i 's/keywords = \["mooncake", "data transfer", "kv cache", "llm inference"\]/keywords = ["mooncake", "data transfer", "kv cache", "llm inference", "ascend", "npu"]/' pyproject.toml + sed -i 's/^description = "\(.*\)"$/description = "\1 (Ascend NPU version)"/' pyproject.toml + sed -i 's/^keywords = \[\(.*\)\]$/keywords = [\1, "ascend", "npu"]/' pyproject.toml + sed -i 's/^requires-python = ">=3.10"$/requires-python = ">=3.9"/' pyproject.toml + sed -i 's|"Environment :: GPU :: NVIDIA CUDA"|"Environment :: GPU"|' pyproject.toml + sed -i 's|"Programming Language :: Python :: 3.10"|"Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10"|' pyproject.toml echo "Package name modified to: mooncake-transfer-engine-npu" +elif [ "$EFA_BUILD" = "1" ]; then + echo "Modifying package name for AWS EFA build (CUDA)" + # Backup original pyproject.toml + cp pyproject.toml pyproject.toml.backup + # Replace package name and description + sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-efa"/' pyproject.toml + sed -i 's/^description = "\(.*\)"$/description = "\1 (AWS EFA, CUDA version)"/' pyproject.toml + sed -i 's/^keywords = \[\(.*\)\]$/keywords = [\1, "aws", "efa", "libfabric", "cuda"]/' pyproject.toml + echo "Package name modified to: mooncake-transfer-engine-efa" +elif [ "$EFA_CU13_BUILD" = "1" ]; then + echo "Modifying package name for AWS EFA build (CUDA 13)" + # Backup original pyproject.toml + cp pyproject.toml pyproject.toml.backup + # Replace package name and description + sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-efa-cuda13"/' pyproject.toml + sed -i 's/^description = "\(.*\)"$/description = "\1 (AWS EFA, CUDA 13 version)"/' pyproject.toml + sed -i 's/^keywords = \[\(.*\)\]$/keywords = [\1, "aws", "efa", "libfabric", "cuda13"]/' pyproject.toml + echo "Package name modified to: mooncake-transfer-engine-efa-cuda13" +elif [ "$EFA_NON_CUDA_BUILD" = "1" ]; then + echo "Modifying package name for AWS EFA build (non-CUDA)" + # Backup original pyproject.toml + cp pyproject.toml pyproject.toml.backup + # Replace package name and description + sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-efa-non-cuda"/' pyproject.toml + sed -i 's/^description = "\(.*\)"$/description = "\1 (AWS EFA, Non-CUDA version)"/' pyproject.toml + sed -i 's/^keywords = \[\(.*\)\]$/keywords = [\1, "aws", "efa", "libfabric", "non-cuda"]/' pyproject.toml + sed -i 's|"Environment :: GPU :: NVIDIA CUDA", ||' pyproject.toml + echo "Package name modified to: mooncake-transfer-engine-efa-non-cuda" +elif [ "$MUSA_BUILD" = "1" ]; then + echo "Modifying package name for MUSA build" + # Backup original pyproject.toml + cp pyproject.toml pyproject.toml.backup + # Replace package name and description + sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-musa"/' pyproject.toml + sed -i 's/^description = "\(.*\)"$/description = "\1 (MUSA version)"/' pyproject.toml + sed -i 's/^keywords = \[\(.*\)\]$/keywords = [\1, "musa", "moore-threads"]/' pyproject.toml + sed -i 's/^requires-python = ">=3.10"$/requires-python = ">=3.9"/' pyproject.toml + sed -i 's|"Environment :: GPU :: NVIDIA CUDA"|"Environment :: GPU"|' pyproject.toml + sed -i 's|"Programming Language :: Python :: 3.10"|"Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10"|' pyproject.toml + echo "Package name modified to: mooncake-transfer-engine-musa" +elif [ "$HIP_BUILD" = "1" ]; then + echo "Modifying package name for AMD ROCm/HIP build" + # Backup original pyproject.toml + cp pyproject.toml pyproject.toml.backup + # Replace package name and description + sed -i 's/name = "mooncake-transfer-engine"/name = "mooncake-transfer-engine-rocm"/' pyproject.toml + sed -i 's/^description = "\(.*\)"$/description = "\1 (AMD ROCm version)"/' pyproject.toml + sed -i 's/^keywords = \[\(.*\)\]$/keywords = [\1, "rocm", "amd", "hip"]/' pyproject.toml + sed -i 's|"Environment :: GPU :: NVIDIA CUDA"|"Environment :: GPU"|' pyproject.toml + echo "Package name modified to: mooncake-transfer-engine-rocm" else echo "Using standard package name: mooncake-transfer-engine" fi @@ -212,7 +300,7 @@ if [ "$NPU_BUILD" = "1" ]; then max_attempts=3 attempt=1 while [ $attempt -le $max_attempts ]; do - if "$PYTHON_CMD" -m pip install --upgrade pip build setuptools wheel auditwheel; then + if "$PYTHON_CMD" -m pip install --upgrade pip build setuptools wheel auditwheel numpy; then break fi echo "pip install attempt $attempt/$max_attempts failed, retrying in 5s..." @@ -344,6 +432,9 @@ ${AUDITWHEEL_CMD} repair ${OUTPUT_DIR}/*.whl \ --exclude libffi.so* \ --exclude libcuda.so* \ --exclude libcudart.so* \ + --exclude libmooncake_pg_device.so* \ + --exclude libmusa.so* \ + --exclude libmusart.so* \ --exclude libamdhip64.so* \ --exclude libhsa-runtime64.so* \ --exclude librocprofiler-register.so* \ @@ -460,6 +551,4 @@ mv ${REPAIRED_DIR}/*.whl ${OUTPUT_DIR}/ cd .. -[[ -f mooncake-wheel/pyproject.toml.backup ]] && mv mooncake-wheel/pyproject.toml.backup mooncake-wheel/pyproject.toml - -echo "Wheel package built and repaired successfully!" \ No newline at end of file +echo "Wheel package built and repaired successfully!" diff --git a/scripts/ci/run_store_go_integration.sh b/scripts/ci/run_store_go_integration.sh new file mode 100755 index 0000000000..482f93cf6c --- /dev/null +++ b/scripts/ci/run_store_go_integration.sh @@ -0,0 +1,90 @@ +#!/usr/bin/env bash + +set -e -o pipefail + +: "${GITHUB_WORKSPACE:?GITHUB_WORKSPACE must be set}" +: "${MOONCAKE_STORE_CLUSTER_ID:?MOONCAKE_STORE_CLUSTER_ID must be set}" + +case "${MOONCAKE_STORE_GO_SANITIZED:-0}" in + 0) sanitized=false ;; + 1) sanitized=true ;; + *) + echo "MOONCAKE_STORE_GO_SANITIZED must be 0 or 1" >&2 + exit 2 + ;; +esac + +case "${MOONCAKE_STORE_GO_LINK_COMMON:-0}" in + 0) link_common=false ;; + 1) link_common=true ;; + *) + echo "MOONCAKE_STORE_GO_LINK_COMMON must be 0 or 1" >&2 + exit 2 + ;; +esac + +"$GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master" \ + --eviction_high_watermark_ratio=0.95 \ + --cluster_id="$MOONCAKE_STORE_CLUSTER_ID" \ + --port 50051 & +master_pid=$! +sleep 3 + +cd "$GITHUB_WORKSPACE/mooncake-store/go" +export LD_LIBRARY_PATH="$GITHUB_WORKSPACE/build/mooncake-common:$GITHUB_WORKSPACE/build/mooncake-store/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src:$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base:$GITHUB_WORKSPACE/build/mooncake-common/etcd" +export CGO_ENABLED=1 +export CGO_CFLAGS="-I$GITHUB_WORKSPACE/mooncake-store/include -I$GITHUB_WORKSPACE/mooncake-transfer-engine/include" + +linker_flags=( + "-L$GITHUB_WORKSPACE/build/mooncake-store/src" + "-L$GITHUB_WORKSPACE/build/mooncake-store/src/cachelib_memory_allocator" + "-L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src" + "-L$GITHUB_WORKSPACE/build/mooncake-transfer-engine/src/common/base" + "-L$GITHUB_WORKSPACE/build/mooncake-common" +) +if $link_common; then + linker_flags+=("-L$GITHUB_WORKSPACE/build/mooncake-common/src") +fi +linker_flags+=("-L$GITHUB_WORKSPACE/build/mooncake-common/etcd") +if $link_common; then + linker_flags+=( + -Wl,--start-group + -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase + -lmooncake_common + -Wl,--end-group + ) +else + linker_flags+=( + -lmooncake_store -lcachelib_memory_allocator -ltransfer_engine -lbase + ) +fi +linker_flags+=( + -lasio -letcd_wrapper -lstdc++ -lnuma -lglog -lgflags -libverbs -lmlx5 + -ljsoncpp -lzstd -lcurl -luring +) +if $sanitized; then + linker_flags+=(-lasan) +fi +linker_flags+=(-lm) +if $sanitized; then + linker_flags+=(-lgcov) +fi +linker_flags+=(-lxxhash -lyaml-cpp) +export CGO_LDFLAGS="${linker_flags[*]}" + +# Link cudart if CUDA is available (needed for D2H staging in mooncake_store). +if [ -d /usr/local/cuda/lib64 ]; then + export CGO_LDFLAGS="$CGO_LDFLAGS -L/usr/local/cuda/lib64 -lcudart" +fi +# The KV events publisher is optional and linked when libzmq is installed. +if ldconfig -p 2>/dev/null | grep -q libzmq; then + export CGO_LDFLAGS="$CGO_LDFLAGS -lzmq" +fi + +test_env=(MC_METADATA_SERVER=http://127.0.0.1:8080/metadata) +if $sanitized; then + test_env=(ASAN_OPTIONS=detect_leaks=0:verify_asan_link_order=0 "${test_env[@]}") +fi +env "${test_env[@]}" go test -v ./tests/... + +kill "$master_pid" 2>/dev/null || true diff --git a/scripts/ci/run_store_rust_smoke.sh b/scripts/ci/run_store_rust_smoke.sh new file mode 100755 index 0000000000..a2de303818 --- /dev/null +++ b/scripts/ci/run_store_rust_smoke.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash + +set -e -o pipefail + +: "${GITHUB_WORKSPACE:?GITHUB_WORKSPACE must be set}" +: "${MOONCAKE_STORE_CLUSTER_ID:?MOONCAKE_STORE_CLUSTER_ID must be set}" + +# ASan-enabled CI libraries require the Rust package to link libasan first. +# Release-style nightly builds intentionally leave the ASan runtime unlinked. +case "${MOONCAKE_STORE_RUST_LINK_ASAN:-0}" in + 0) unset MOONCAKE_LINK_ASAN ;; + 1) export MOONCAKE_LINK_ASAN=1 ;; + *) + echo "MOONCAKE_STORE_RUST_LINK_ASAN must be 0 or 1" >&2 + exit 2 + ;; +esac + +"$GITHUB_WORKSPACE/build/mooncake-store/src/mooncake_master" \ + --eviction_high_watermark_ratio=0.95 \ + --cluster_id="$MOONCAKE_STORE_CLUSTER_ID" \ + --port 50051 & +master_pid=$! +sleep 3 + +cd "$GITHUB_WORKSPACE/mooncake-store/rust" +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:${LD_LIBRARY_PATH:-}" +export MOONCAKE_BUILD_DIR="$GITHUB_WORKSPACE/build" +export MOONCAKE_STORE_LIB_DIR="$GITHUB_WORKSPACE/build/mooncake-store/src" +export MOONCAKE_STORE_INCLUDE_DIR="$GITHUB_WORKSPACE/mooncake-store/include" +export MC_METADATA_SERVER=http://127.0.0.1:8080/metadata +export MC_RUST_STORE_RUN_INTEGRATION=true +export MC_RUST_STORE_MASTER_ADDR=127.0.0.1:50051 +export MC_RUST_STORE_LOCAL_HOSTNAME=127.0.0.1 +export MC_RUST_STORE_PROTOCOL=tcp +export MC_RUST_STORE_DEVICE_NAME= + +cargo test --test minimal_smoke -- --nocapture +MC_RUST_BENCH_ITERATIONS=4 \ + MC_RUST_BENCH_VALUE_SIZE=4096 \ + MC_RUST_BENCH_WARMUP=1 \ + cargo run --release --example store_benchmark + +kill "$master_pid" 2>/dev/null || true diff --git a/scripts/code_format.sh b/scripts/code_format.sh index a65f797477..a64ed89ce6 100755 --- a/scripts/code_format.sh +++ b/scripts/code_format.sh @@ -13,6 +13,8 @@ # -a, --all Format all C/C++ files in the project # -b, --base Base branch to compare against (default: origin/main) # -c, --check Check mode: only report files that need formatting +# --staged Format only added/modified lines staged for commit +# --changed-lines Check lines changed relative to base (requires --check) # -h, --help Show this help message # # Examples: @@ -20,6 +22,8 @@ # ./scripts/code_format.sh --all # Format all C/C++ files # ./scripts/code_format.sh -b origin/dev # Compare against origin/dev # ./scripts/code_format.sh --check # Check without modifying files +# ./scripts/code_format.sh --staged file.cpp # Format staged lines (pre-commit) +# ./scripts/code_format.sh --changed-lines --check # Check PR-changed lines # ============================================================================= set -e @@ -28,11 +32,15 @@ set -e BASE_BRANCH="origin/main" CHECK_MODE=false ALL_MODE=false +STAGED_MODE=false +CHANGED_LINES_MODE=false +INPUT_FILES=() SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" PROJECT_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" # File extensions to format FILE_EXTENSIONS="\.(h|hpp|cpp|cu|cuh|c|cc|cxx)$" +GIT_CLANG_FORMAT_EXTENSIONS="h,hpp,cpp,cu,cuh,c,cc,cxx" # Directories to exclude (add more patterns as needed) EXCLUDE_DIRS=( @@ -100,6 +108,26 @@ find_clang_format() { return 1 } +# Find the LLVM git integration helper. The versioned binary is preferred, but +# an unversioned helper is safe because clang-format 20 is passed via --binary. +find_git_clang_format() { + local candidates=("git-clang-format-20" "git-clang-format") + local candidate + for candidate in "${candidates[@]}"; do + if command -v "${candidate}" &> /dev/null; then + echo "${candidate}" + return 0 + fi + done + + { + print_error "A compatible git-clang-format helper was not found." + print_info "Install the clang-format 20 package, which provides git-clang-format-20." + echo " sudo apt-get install -y clang-format-20" + } >&2 + return 1 +} + # Parse command line arguments parse_args() { while [[ $# -gt 0 ]]; do @@ -116,15 +144,139 @@ parse_args() { CHECK_MODE=true shift ;; + --staged) + STAGED_MODE=true + shift + ;; + --changed-lines) + CHANGED_LINES_MODE=true + shift + ;; -h|--help) usage ;; - *) + --) + shift + INPUT_FILES+=("$@") + break + ;; + -*) print_error "Unknown option: $1" usage ;; + *) + INPUT_FILES+=("$1") + shift + ;; esac done + + if ${ALL_MODE} && ${STAGED_MODE}; then + print_error "--all and --staged cannot be used together." + exit 1 + fi + if ${ALL_MODE} && ${CHANGED_LINES_MODE}; then + print_error "--all and --changed-lines cannot be used together." + exit 1 + fi + if ${STAGED_MODE} && ${CHANGED_LINES_MODE}; then + print_error "--staged and --changed-lines cannot be used together." + exit 1 + fi + if ${CHANGED_LINES_MODE} && ! ${CHECK_MODE}; then + print_error "--changed-lines requires --check." + exit 1 + fi +} + +# Delegate changed-line selection and index handling to LLVM's git integration. +# A strict guard is retained because git-clang-format only checks for unstaged +# changes when formatting would actually rewrite a file. +format_selected_lines() { + local clang_format="$1" + local git_clang_format="$2" + shift 2 + + # With no explicit filenames, Git pathspecs select supported C/C++ files. + # Exclude pathspecs keep vendored sources out in both automatic and explicit + # modes without discovering and filtering files in this script. + if [[ $# -eq 0 ]]; then + set -- '*.c' '*.cc' '*.cpp' '*.cxx' '*.cu' '*.cuh' '*.h' '*.hpp' + fi + local paths=( + "$@" + ':(exclude,glob)**/cachelib_memory_allocator/**' + ':(exclude,glob)**/thirdparty/**' + ) + + print_info "Using $(${clang_format} --version)" + if ${STAGED_MODE}; then + print_info "Formatting only staged C/C++ line ranges" + else + print_info "Formatting only C/C++ line ranges changed against ${BASE_BRANCH}" + fi + echo "" + + if ${STAGED_MODE}; then + local git_status + if ! git_status=$(git -C "${PROJECT_ROOT}" status --porcelain=v1 \ + --untracked-files=no -- "${paths[@]}"); then + print_error "Could not inspect staged and unstaged changes." + return 1 + fi + # Both porcelain status columns are populated when the same file has + # changes in the index and the working tree. + if printf '%s\n' "${git_status}" | grep -q '^[^ ][^ ]'; then + print_error "Cannot process staged lines because a staged file also has unstaged changes." + print_info "Stage or stash the unstaged changes, then retry." + return 1 + fi + else + if ! git -C "${PROJECT_ROOT}" diff --quiet --no-ext-diff HEAD -- \ + "${paths[@]}"; then + print_error "Cannot process committed lines because a selected file has uncommitted changes." + print_info "Commit or stash the local changes, then retry." + return 1 + fi + fi + + local command=( + "${git_clang_format}" + --binary "${clang_format}" + --style file + --extensions "${GIT_CLANG_FORMAT_EXTENSIONS}" + ) + if ${CHECK_MODE}; then + command+=(--diff) + fi + if ${STAGED_MODE}; then + command+=(--staged) + else + # LLVM 20 implements --diff_from_common_commit with a triple-dot + # revision passed to diff-tree, which can produce an empty patch. + # Resolve the merge base first and pass two concrete commits instead. + local base_commit + if ! base_commit=$(git -C "${PROJECT_ROOT}" merge-base \ + "${BASE_BRANCH}" HEAD); then + print_error "Could not determine a merge base for '${BASE_BRANCH}' and HEAD." + return 1 + fi + command+=("${base_commit}" HEAD) + fi + command+=(-- "${paths[@]}") + + local status=0 + ( + cd "${PROJECT_ROOT}" + "${command[@]}" + ) || status=$? + + # git-clang-format returns 1 after successfully applying a rewrite. Treat + # that as success in apply mode; in check mode it means a diff was found. + if ! ${CHECK_MODE} && [[ ${status} -eq 1 ]]; then + return 0 + fi + return "${status}" } # Get list of all C/C++ files in the project @@ -268,6 +420,22 @@ main() { exit 1 fi + if ${STAGED_MODE} || ${CHANGED_LINES_MODE}; then + if ${CHANGED_LINES_MODE} && + ! git -C "${PROJECT_ROOT}" rev-parse --verify "${BASE_BRANCH}" &> /dev/null; then + print_error "Base branch '${BASE_BRANCH}' not found." + exit 1 + fi + + local git_clang_format + if ! git_clang_format="$(find_git_clang_format)"; then + exit 1 + fi + format_selected_lines "${clang_format}" "${git_clang_format}" \ + "${INPUT_FILES[@]}" + return + fi + # Get files to format local files if ${ALL_MODE}; then diff --git a/scripts/launch_cxi_transfer_test.sh b/scripts/launch_cxi_transfer_test.sh new file mode 100644 index 0000000000..b57f8537ba --- /dev/null +++ b/scripts/launch_cxi_transfer_test.sh @@ -0,0 +1,25 @@ +#!/bin/bash + +ip=$(hostname -I | awk '{print $1}') +echo $ip + +MC_MIN_PRC_PORT=25565 MC_MAX_PRC_PORT=25565 ./build/mooncake-transfer-engine/tests/cxi_transfer_test \ + --mode target \ + --use_device \ + --server $ip \ + --num_bufs 4 \ + --buf_size_gb 1 & +TARGET_PID=$! + +sleep 10 + +./build/mooncake-transfer-engine/tests/cxi_transfer_test --mode initiator \ + --server 127.0.0.1:12346 \ + --target $ip:25565 \ + --transfer_mb 128 \ + --threads 4 \ + --use_device \ + --warmup 0 + +kill -INT $TARGET_PID +wait $TARGET_PID \ No newline at end of file diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh index 3aaf2c7b18..3c89735870 100755 --- a/scripts/run_tests.sh +++ b/scripts/run_tests.sh @@ -7,6 +7,116 @@ set -e # Exit immediately if a command exits with a non-zero status # Ensure LD_LIBRARY_PATH includes /usr/local/lib export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib +METADATA_SERVER_PID="" +RUN_TESTS_METADATA_SERVER_MODE="${RUN_TESTS_METADATA_SERVER_MODE:-managed}" + +metadata_port_is_listening() { + ss -H -ltn 'sport = :8080' | grep -q . +} + +start_metadata_server() { + local attempt + + echo "Starting standalone HTTP metadata server..." + mooncake_http_metadata_server --port 8080 & + METADATA_SERVER_PID=$! + + for attempt in {1..50}; do + if ! kill -0 "$METADATA_SERVER_PID" 2>/dev/null; then + echo "ERROR: metadata server exited before listening on port 8080" + wait "$METADATA_SERVER_PID" 2>/dev/null || true + METADATA_SERVER_PID="" + return 1 + fi + + if metadata_port_is_listening; then + echo "Standalone HTTP metadata server is ready (pid=$METADATA_SERVER_PID)" + return 0 + fi + sleep 0.1 + done + + echo "ERROR: metadata server did not listen on port 8080 within 5 seconds" + ss -ltnp | grep ':8080' || true + return 1 +} + +use_external_metadata_server() { + local attempt + + echo "Using caller-managed HTTP metadata server on port 8080..." + for attempt in {1..50}; do + if metadata_port_is_listening; then + echo "Caller-managed HTTP metadata server is ready" + return 0 + fi + sleep 0.1 + done + + echo "ERROR: caller-managed metadata server is not listening on port 8080" + return 1 +} + +stop_metadata_server() { + local attempt + local pid="${METADATA_SERVER_PID:-}" + + if [ -z "$pid" ]; then + return 0 + fi + + echo "Stopping standalone HTTP metadata server (pid=$pid)..." + kill -TERM "$pid" 2>/dev/null || true + + for attempt in {1..50}; do + if ! metadata_port_is_listening; then + wait "$pid" 2>/dev/null || true + METADATA_SERVER_PID="" + echo "Standalone HTTP metadata server stopped" + return 0 + fi + sleep 0.1 + done + + echo "WARNING: metadata server did not release port 8080 after SIGTERM; sending SIGKILL" + ps -fp "$pid" || true + ss -ltnp | grep ':8080' || true + kill -KILL "$pid" 2>/dev/null || true + wait "$pid" 2>/dev/null || true + METADATA_SERVER_PID="" + + for attempt in {1..50}; do + if ! metadata_port_is_listening; then + echo "Standalone HTTP metadata server was force-stopped" + return 0 + fi + sleep 0.1 + done + + echo "ERROR: port 8080 is still occupied after stopping metadata server" + ps -ef | grep '[m]ooncake' || true + ss -ltnp | grep ':8080' || true + return 1 +} + +cleanup_metadata_server() { + stop_metadata_server || true +} + +case "$RUN_TESTS_METADATA_SERVER_MODE" in + managed) + trap cleanup_metadata_server EXIT + start_metadata_server + ;; + external) + use_external_metadata_server + ;; + *) + echo "ERROR: RUN_TESTS_METADATA_SERVER_MODE must be 'managed' or 'external'" + exit 1 + ;; +esac + echo "Running transfer_engine tests..." cd mooncake-wheel/tests MC_METADATA_SERVER=http://127.0.0.1:8080/metadata MC_FORCE_TCP=true python transfer_engine_target.py & @@ -104,30 +214,42 @@ else echo "Skipping test: TEST_PROMOTION_ON_HIT environment variable is not set" fi -echo "Running CXL protocol test (test_distributed_object_store_cxl.py)..." -killall mooncake_master || true -sleep 2 +if [ -n "$TEST_CXL" ]; then + echo "Running CXL protocol test (test_distributed_object_store_cxl.py)..." + killall mooncake_master || true + sleep 2 -echo "Starting Mooncake Master with CXL enabled (--enable_cxl=true)..." -mooncake_master \ - --default_kv_lease_ttl=500 \ - --enable_cxl=true \ - & -CXL_MASTER_PID=$! -sleep 3 -MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_distributed_object_store_cxl.py -kill $CXL_MASTER_PID || true -wait $CXL_MASTER_PID 2>/dev/null || true -sleep 2 -echo "CXL protocol test completed successfully!" + echo "Starting Mooncake Master with CXL enabled (--enable_cxl=true)..." + mooncake_master \ + --default_kv_lease_ttl=500 \ + --enable_cxl=true \ + & + CXL_MASTER_PID=$! + sleep 3 + MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_distributed_object_store_cxl.py + kill $CXL_MASTER_PID || true + wait $CXL_MASTER_PID 2>/dev/null || true + sleep 2 + echo "CXL protocol test completed successfully!" +else + echo "Skipping CXL protocol test: TEST_CXL is not set" +fi echo "Running CLI entry point tests..." python test_cli.py -killall mooncake_http_metadata_server || true +ENABLE_EMBEDDED_METADATA_SERVER=true +if [ "$RUN_TESTS_METADATA_SERVER_MODE" = "managed" ]; then + stop_metadata_server +else + ENABLE_EMBEDDED_METADATA_SERVER=false + echo "Leaving caller-managed HTTP metadata server running" +fi killall mooncake_master || true killall mooncake_client || true -mooncake_master --default_kv_lease_ttl=500 --enable_http_metadata_server=true & +mooncake_master \ + --default_kv_lease_ttl=500 \ + --enable_http_metadata_server="$ENABLE_EMBEDDED_METADATA_SERVER" & MASTER_PID=$! sleep 1 MC_METADATA_SERVER=http://127.0.0.1:8080/metadata DEFAULT_KV_LEASE_TTL=500 python test_distributed_object_store.py diff --git a/scripts/test_async_store.py b/scripts/test_async_store.py index b1b4046651..fea5c69407 100644 --- a/scripts/test_async_store.py +++ b/scripts/test_async_store.py @@ -9,7 +9,6 @@ import asyncio from mooncake.mooncake_config import MooncakeConfig -from dataclasses import dataclass try: diff --git a/scripts/test_installation.sh b/scripts/test_installation.sh index 9e7f07761d..68cfd181d6 100755 --- a/scripts/test_installation.sh +++ b/scripts/test_installation.sh @@ -38,6 +38,8 @@ echo "Running import structure test..." cp -r mooncake-wheel/tests test_env/ cd test_env pip install torch==2.11.0 numpy +python -c "import mooncake._fast_copy" +python tests/test_fast_copy.py python tests/test_import_structure.py echo "Running mooncake config test..." @@ -46,6 +48,7 @@ python tests/test_mooncake_config.py echo "Verifying mooncake_master entry point..." # Check if the mooncake_master entry point is installed and executable which mooncake_master || { echo "ERROR: mooncake_master entry point not found!"; exit 1; } +mooncake_master --version echo "Success: mooncake_master entry point found" # Check if the mooncake_client entry point is installed and executable diff --git a/scripts/test_tensor_api.py b/scripts/test_tensor_api.py index bdffc43732..8a9e260594 100644 --- a/scripts/test_tensor_api.py +++ b/scripts/test_tensor_api.py @@ -1,7 +1,6 @@ import ctypes import os import sys -import json import time import argparse import unittest @@ -1549,7 +1548,6 @@ def test_32_batch_unified_parallelism_full_reconstruction(self): make_deterministic_tensor((12, 12)), ] keys = [f"func_unified_tp_full_batch_{i}" for i in range(len(tensors))] - parallelisms = [build_tp_parallelism(tp_size, split_dim, rank=0) for _ in tensors] results = batch_put_full_tensors_with_unified_tp(self.store, keys, tensors, tp_size, split_dim) self.assertTrue(all(r == 0 for r in results), f"batch put failed: {results}") diff --git a/scripts/test_upsert_api.py b/scripts/test_upsert_api.py index 11aa2abc5a..6e119980db 100644 --- a/scripts/test_upsert_api.py +++ b/scripts/test_upsert_api.py @@ -28,8 +28,6 @@ import ctypes import os -import sys -import time import unittest from dataclasses import dataclass diff --git a/scripts/tone_tests/scripts/common.sh b/scripts/tone_tests/scripts/common.sh index 1c422bf2e3..68e5518ea0 100755 --- a/scripts/tone_tests/scripts/common.sh +++ b/scripts/tone_tests/scripts/common.sh @@ -231,7 +231,7 @@ get_whl(){ echo "get whl file from github action" rm -f "$whls_path/mooncake.zip" - rm -f "$whls_path/*.whl" + rm -f "$whls_path"/*.whl local max_retries=5 local base_delay=5 # seconds @@ -401,6 +401,139 @@ cleanup_test_env() { echo "Cleanup completed" } +# Wait until GPU memory on the local host drains below a threshold. +# Returns 0 once drained, 1 if it times out. +wait_gpu_idle() { + local max_seconds=${1:-90} + local threshold_mb=${2:-1024} + + if ! command -v nvidia-smi >/dev/null 2>&1; then + echo "ERROR: nvidia-smi not available; cannot verify GPU drain" >&2 + return 1 + fi + + echo "Waiting for GPU memory to drain (threshold ${threshold_mb}MB, timeout ${max_seconds}s)..." + local elapsed=0 + local max_used=0 + while [ $elapsed -lt $max_seconds ]; do + max_used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | sort -n | tail -n 1) + if ! [[ "$max_used" =~ ^[0-9]+$ ]]; then + echo "ERROR: nvidia-smi query failed; cannot verify GPU drain" >&2 + return 1 + fi + if [ "$max_used" -le "$threshold_mb" ]; then + echo "GPU memory drained (max used ${max_used}MB)" + return 0 + fi + sleep 3 + elapsed=$((elapsed + 3)) + done + echo "GPU memory not drained within ${max_seconds}s (max used ${max_used}MB)" + return 1 +} + +# Kill only GPU processes whose cgroup still identifies them as belonging to +# the reused test container. Unknown or unrelated processes must be left alone; +# the caller will quarantine the environment if GPU memory remains occupied. +gpu_pid_belongs_to_container() { + local pid=$1 + local container_id=$2 + [ -r "/proc/${pid}/cgroup" ] && grep -Fq "$container_id" "/proc/${pid}/cgroup" +} + +force_kill_container_gpu_procs() { + command -v nvidia-smi >/dev/null 2>&1 || return 1 + + local container_id + container_id=$(docker inspect --format '{{.Id}}' "${CONTAINER_NAME}" 2>/dev/null) || { + echo "ERROR: Cannot inspect container ${CONTAINER_NAME}; refusing to kill host GPU processes" >&2 + return 1 + } + if [ -z "$container_id" ]; then + echo "ERROR: Empty container ID for ${CONTAINER_NAME}; refusing to kill host GPU processes" >&2 + return 1 + fi + + local gpu_pids + gpu_pids=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | tr -cd '0-9\n' | grep -E '^[0-9]+$' | sort -u) + [ -z "$gpu_pids" ] && return 0 + + local container_pids="" + local pid + for pid in $gpu_pids; do + if gpu_pid_belongs_to_container "$pid" "$container_id"; then + container_pids="${container_pids} ${pid}" + else + echo "WARNING: GPU PID ${pid} is not owned by ${CONTAINER_NAME}; leaving it untouched" >&2 + fi + done + + if [ -z "$container_pids" ]; then + echo "No container-owned GPU processes are safe to kill" + return 0 + fi + + echo "Force-killing container-owned GPU PIDs:${container_pids}" + for pid in $container_pids; do + if ! kill -9 "$pid" 2>/dev/null; then + echo "ERROR: Failed to kill container-owned GPU PID ${pid}" >&2 + return 1 + fi + done + sleep 3 +} + +# Fully clear GPU memory on the current node: restart the container first, then +# kill only container-owned GPU processes that survived the restart. +# Restart the reused container to reset all in-container state (processes, GPU +# memory, ERDMA queue-pairs / RDMA contexts). 'docker restart' keeps the +# writable layer, so the mooncake wheel and ERDMA drivers are NOT reinstalled. +drain_gpu_local() { + echo "Restarting container ${CONTAINER_NAME} to reset GPU/ERDMA state..." + if ! docker restart "${CONTAINER_NAME}" >/dev/null 2>&1; then + echo "ERROR: Failed to restart container ${CONTAINER_NAME}; environment is unhealthy" >&2 + return 1 + fi + + if ! wait_gpu_idle 60; then + echo "GPU memory remains occupied; checking for container-owned processes..." + force_kill_container_gpu_procs || return 1 + if ! wait_gpu_idle 45; then + echo "ERROR: GPU still occupied after bounded cleanup; environment is unhealthy" >&2 + return 1 + fi + fi + + return 0 +} + +# Between test cases in run-all the container is reused; reset in-container +# state on both the local and (for double-machine runs) remote nodes via a +# lightweight container restart (no wheel / ERDMA driver reinstall). +drain_gpu_between_tests() { + echo "===== Resetting environment between test cases =====" + local reset_failed=false + if ! drain_gpu_local; then + echo "ERROR: Failed to reset the local test environment" >&2 + reset_failed=true + fi + + if [ -n "$REMOTE_IP" ]; then + echo "Resetting environment on remote node $REMOTE_IP..." + if ! ${SSH_CMD} "$REMOTE_IP" " + source ${REMOTE_TEST_DIR}/run/.shrc && \ + source ${REMOTE_TEST_DIR}/scripts/common.sh && \ + drain_gpu_local + "; then + echo "ERROR: Failed to reset the remote test environment on ${REMOTE_IP}" >&2 + reset_failed=true + fi + fi + + $reset_failed && return 1 + return 0 +} + setup_node_env() { local registry_addr=$1 echo "===== Setting up docker environment =====" @@ -416,6 +549,7 @@ setup_node_env() { fi local extra_args="" + extra_args="$extra_args -e NCCL_GIN_TYPE=0 " extra_args="$extra_args --device=/dev/infiniband/uverbs0 --device=/dev/infiniband/uverbs1 --device=/dev/infiniband/rdma_cm " if [ "${USE_HUGGINGFACE_MIRROR}" = "true" ]; then extra_args="$extra_args -e HF_ENDPOINT=${HUGGINGFACE_MIRROR} -e HF_HUB_ENABLE_HF_TRANSFER=1" @@ -840,6 +974,19 @@ collect_and_validate_model_results() { fi } +# Echo an offline env prefix when the given model already exists in the +# container's HuggingFace cache, so servers use the local snapshot instead of +# querying the hub (skips downloads and avoids hf-mirror 429 rate limiting). +# Models are pre-cached under MODEL_CACHE (mounted at /root/.cache) on both nodes. +hf_offline_prefix() { + local model_name=$1 + [ -z "$model_name" ] && return 0 + local cache_dir="models--$(echo "$model_name" | sed 's#/#--#g')" + if ${docker_exec} "ls /root/.cache/huggingface/hub/${cache_dir}/snapshots/*/config.json >/dev/null 2>&1"; then + echo "HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 " + fi +} + launch_sglang_server() { local model_path=$1 local host=$2 @@ -855,7 +1002,8 @@ launch_sglang_server() { return 1 fi - local sglang_cmd="${docker_exec} \"python -m sglang.launch_server --model-path ${model_path} --host ${host} --port ${port}" + local offline_prefix=$(hf_offline_prefix "$model_path") + local sglang_cmd="${docker_exec} \"${offline_prefix}python -m sglang.launch_server --model-path ${model_path} --host ${host} --port ${port}" if [ -n "$extra_args" ]; then sglang_cmd="${sglang_cmd} ${extra_args}" fi @@ -905,6 +1053,7 @@ launch_vllm_server() { if [ -n "$env_vars" ]; then env_prefix="${env_vars} " fi + env_prefix="${env_prefix}$(hf_offline_prefix "$model_path")" local vllm_cmd="${docker_exec} \"${env_prefix}python3 -m vllm.entrypoints.openai.api_server --model '${model_path}' --host '${host}' --port ${port}" diff --git a/scripts/tone_tests/scripts/run_test.sh b/scripts/tone_tests/scripts/run_test.sh old mode 100644 new mode 100755 index a0db0c1887..30ff521bbe --- a/scripts/tone_tests/scripts/run_test.sh +++ b/scripts/tone_tests/scripts/run_test.sh @@ -232,7 +232,10 @@ run_single_test(){ source "$RUN_DIR/.shrc" cd "$TONE_TESTS_DIR/scripts" source "./$test_name" - + + local log_dir="${BASE_DIR}/run/logs/$(basename "$test_name" .sh)" + setup_log_directory "$log_dir" + local exit_code=0 run_test "$@" || exit_code=1 @@ -263,7 +266,10 @@ run_all_tests(){ cd "$TONE_TESTS_DIR/scripts" local all_passed=true + local test_index=0 + local test_count=${#tests[@]} for test_name in "${tests[@]}"; do + test_index=$((test_index + 1)) if [[ ! -f "./$test_name" ]]; then echo "WARNING: test case $test_name is not found, skipping" continue @@ -282,6 +288,16 @@ run_all_tests(){ fi [ $exit_code -ne 0 ] && all_passed=false + + # Container is shared across cases in run-all. Do not schedule another + # case unless both nodes have been reset successfully. + if [ "$test_index" -lt "$test_count" ] && ! drain_gpu_between_tests; then + local remaining_count=$((test_count - test_index)) + echo "ERROR: Shared test environment is unhealthy after ${test_name}." >&2 + echo "ERROR: Skipping ${remaining_count} remaining test case(s); node intervention is required." >&2 + all_passed=false + break + fi done cleanup_test_env "double" @@ -330,4 +346,4 @@ case "$1" in *) show_help ;; -esac \ No newline at end of file +esac diff --git a/scripts/tone_tests/scripts/test_disaggregation_different_tp.sh b/scripts/tone_tests/scripts/test_disaggregation_different_tp.sh index fd1ff7e047..bf3f189508 100644 --- a/scripts/tone_tests/scripts/test_disaggregation_different_tp.sh +++ b/scripts/tone_tests/scripts/test_disaggregation_different_tp.sh @@ -1,7 +1,7 @@ #!/bin/bash # This test case is adapted from https://github.com/sgl-project/sglang/blob/main/test/registered/distributed/test_disaggregation_different_tp.py -# The original test has been updated to use deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct as DEFAULT_MODEL_NAME_FOR_TEST_MLA +# The original test has been updated to use deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct as DEFAULT_MODEL_NAME_FOR_TEST_MLA # and meta-llama/Llama-3.2-3B-Instruct as DEFAULT_MODEL_NAME_FOR_TEST test_case_name="test_disaggregation_different_tp" @@ -11,17 +11,44 @@ BASE_DIR=${BASE_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && cd .. && pwd)} . ${BASE_DIR}/scripts/common.sh run_test() -{ +{ echo "===== Running pytest tests =====" local log_file="${BASE_DIR}/${TEST_CASE_RESULT_PATH}/${test_case_name}.log" echo "Running tests in container and saving output to: $log_file" - ${docker_exec} "\ - cd /sgl-workspace/sglang/test/registered/distributed && \ - sed -i '0,/^class /s|^class |DEFAULT_MODEL_NAME_FOR_TEST_MLA = \"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct\"\nDEFAULT_MODEL_NAME_FOR_TEST = \"meta-llama/Llama-3.2-3B-Instruct\"\n&|' test_disaggregation_different_tp.py && \ - echo 'Model override applied successfully' && \ - python3 -m pytest test_disaggregation_different_tp.py -v -s --tb=long" | tee "$log_file" - + + local offline_prefix="" + local cache_diagnostic="" + if [ "${USE_MODELSCOPE}" = "true" ]; then + cache_diagnostic="Model cache policy: online (USE_MODELSCOPE=true)" + else + local off_mla + local off_llama + local off_qwen + off_mla=$(hf_offline_prefix "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct") + off_llama=$(hf_offline_prefix "meta-llama/Llama-3.2-3B-Instruct") + off_qwen=$(hf_offline_prefix "Qwen/Qwen3.5-4B") + if [ -z "$off_mla" ]; then + cache_diagnostic="Model cache policy: online (missing cache for deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct)" + elif [ -z "$off_llama" ]; then + cache_diagnostic="Model cache policy: online (missing cache for meta-llama/Llama-3.2-3B-Instruct)" + elif [ -z "$off_qwen" ]; then + cache_diagnostic="Model cache policy: online (missing cache for Qwen/Qwen3.5-4B)" + else + offline_prefix="$off_mla" + cache_diagnostic="Model cache policy: offline (all required models are cached)" + fi + fi + + { + echo "$cache_diagnostic" + ${docker_exec} "\ + cd /sgl-workspace/sglang/test/registered/disaggregation && \ + sed -i '0,/^class /s|^class |DEFAULT_MODEL_NAME_FOR_TEST_MLA = \"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct\"\nDEFAULT_MODEL_NAME_FOR_TEST = \"meta-llama/Llama-3.2-3B-Instruct\"\n&|' test_disaggregation_different_tp.py && \ + echo 'Model override applied successfully' && \ + ${offline_prefix}python3 -m pytest test_disaggregation_different_tp.py -v -s --tb=long" + } 2>&1 | tee "$log_file" + return ${PIPESTATUS[0]} } @@ -46,7 +73,7 @@ if [ "${BASH_SOURCE[0]}" == "${0}" ]; then if ! run_test; then exit_code=1 fi - + parse $exit_code exit $? -fi \ No newline at end of file +fi diff --git a/scripts/tone_tests/scripts/test_moe_mooncake.sh b/scripts/tone_tests/scripts/test_moe_mooncake.sh index 2e8573c536..477e3cfdc0 100644 --- a/scripts/tone_tests/scripts/test_moe_mooncake.sh +++ b/scripts/tone_tests/scripts/test_moe_mooncake.sh @@ -13,11 +13,14 @@ run_test() local log_file="${BASE_DIR}/${TEST_CASE_RESULT_PATH}/${test_case_name}.log" echo "Running tests in container and saving output to: $log_file" - + + # Use local HF cache (offline) when the model is already downloaded. + local offline_prefix=$(hf_offline_prefix "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct") + ${docker_exec} "\ export PYTHONPATH=/sgl-workspace/sglang:\$PYTHONPATH && \ cd /test_run/python && \ - python3 -m pytest test_moe_mooncake.py -v -s --tb=long" | tee "$log_file" + ${offline_prefix}python3 -m pytest test_moe_mooncake.py -v -s --tb=long" | tee "$log_file" return ${PIPESTATUS[0]} } diff --git a/scripts/tone_tests/scripts/test_vllm_1p1d_erdma.sh b/scripts/tone_tests/scripts/test_vllm_1p1d_erdma.sh index 8c856b23d5..150aa98ecb 100755 --- a/scripts/tone_tests/scripts/test_vllm_1p1d_erdma.sh +++ b/scripts/tone_tests/scripts/test_vllm_1p1d_erdma.sh @@ -34,9 +34,9 @@ start_server() local kv_config_json="{\\\"kv_connector\\\":\\\"MooncakeConnector\\\",\\\"kv_role\\\":\\\"$kv_role\\\"}" - local extra_args="--tensor-parallel-size 2 --max-model-len 32768 --no-enable-prefix-caching --kv-transfer-config '$kv_config_json'" + local extra_args="--tensor-parallel-size 2 --max-model-len 32768 --gpu-memory-utilization 0.85 --no-enable-prefix-caching --kv-transfer-config '$kv_config_json'" - local env_vars="CUDA_VISIBLE_DEVICES=0,1" + local env_vars="CUDA_VISIBLE_DEVICES=6,7" if ! launch_vllm_server "$model_name" "$host" "$port" "$vllm_server_log_path" "$kv_role" "$extra_args" "$env_vars"; then return 1 @@ -62,7 +62,7 @@ run_proxy() local use_health_check=0 if python3 -c "from packaging import version; import sys; sys.exit(0 if version.parse('$vllm_version') >= version.parse('0.16.0') else 1)" 2>/dev/null; then echo "Using Mooncake Connector Proxy (vLLM >= 0.16.0)" - proxy_script="python3 -u /vllm-workspace/examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py --prefill http://$REMOTE_IP:8010 --decode http://$LOCAL_IP:8020 --host 0.0.0.0 --port 8000" + proxy_script="python3 -u /vllm-workspace/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py --prefill http://$REMOTE_IP:8010 --decode http://$LOCAL_IP:8020 --host 0.0.0.0 --port 8000" ready_pattern="All prefiller instances are ready." else echo "Using NIXL Proxy (vLLM < 0.16.0)"